def legal_actions_after_sequence(self,
                                   actions_sequence: Optional[Tuple[int]]):
    """Returns the legal action space after a sequence of actions."""
    stack = self.apply_actions_to_stack(self.initial_stack, actions_sequence)
    full_tree_action_history = self.actions + actions_sequence

    original_legal_actions = self.legal_actions_for_stack(stack=stack)
    legal_actions = original_legal_actions

    def _find_word_prefix():

      # Find the latest `_W_ -> *` rule, if any.
      w_pos = None
      for pos in range(len(full_tree_action_history) - 1, -1, -1):
        rule = self.grammar.action_to_rule[str(full_tree_action_history[pos])]
        if (rule.startswith('_W_') or rule.startswith('_Wq_') or
            rule.startswith('_Wa_') or rule.startswith('_Wd_') or
            rule.startswith('_Wt_')):
          w_pos = pos
          break
      if w_pos is None:
        return ()

      # Collect all the terminals generated so far.
      word_prefix = []
      for pos in range(w_pos + 1, len(full_tree_action_history)):
        if self.grammar.productions()[
            full_tree_action_history[pos]].is_lexical():
          word_prefix.append(full_tree_action_history[pos])
      return tuple(word_prefix)

    def _compute_crossed_search_boundary():
      # We crossed a search boundary if at any point during the actions sequence
      # we have another grammar start on top of the stack.
      for i in range(len(actions_sequence)):
        intermediate_stack = self.apply_actions_to_stack(
            self.initial_stack, actions_sequence[:i + 1])
        if intermediate_stack[-1] == self.grammar.start():
          return True
      return False

    def _find_word_action():
      # Find the latest rule that indicates the type of `Word` to use.
      for pos in range(len(full_tree_action_history) - 1, -1, -1):
        rule = self.grammar.action_to_rule[str(full_tree_action_history[pos])]
        if '_W_' in rule:
          return '_W_'
        if '_Wq_' in rule:
          return '_Wq_'
        if '_Wa_' in rule:
          return '_Wa_'
        if '_Wd_' in rule:
          return '_Wd_'
        if '_Wt_' in rule:
          return '_Wt_'
      return ''

    def _find_operator():
      # Find the latest rule that indicates the type of operator to use.
      for pos in range(len(full_tree_action_history) - 1, -1, -1):
        rule = self.grammar.action_to_rule[str(full_tree_action_history[pos])]
        if '[pos]' in rule:
          return '[pos]'
        if '[neg]' in rule:
          return '[neg]'
        if '[or]' in rule:
          return '[or]'
      return ''

    def _can_construct(pieces, valid_actions) -> bool:
      if crossed_search_boundary:
        return True
      return all(piece in valid_actions for piece in pieces)

    def _get_restricted_trie(rule: str, trie: pygtrie.Trie) -> pygtrie.Trie:
      if trie:
        if '[pos]' in rule or '[or]' in rule:
          # For [pos] or [or] operators, remove words from input trie that are
          # not valid intersect_word_actions.
          valid_word_actions = self.valid_word_actions.intersect_word_actions
        elif '[neg]' in rule:
          # For [neg] rules, remove words from the input trie that are not
          # valid diff_word_actions.
          valid_word_actions = self.valid_word_actions.diff_word_actions
        else:
          valid_word_actions = None
        if valid_word_actions:
          result = copy.deepcopy(trie)
          for word in trie.iterkeys(()):
            if not _can_construct(word, valid_word_actions):
              del result[word]
          return result
      return trie

    def _log_error(error: str):
      logging.error(error)
      logging.error('  action_sequence: %s', actions_sequence)
      logging.error('  full_action_sequence: %s', full_tree_action_history)
      for action in full_tree_action_history:
        logging.error('    action: %d = [%s]', action,
                      self.grammar.action_to_rule[str(action)])
      logging.error('  word_prefix: %s', word_prefix)
      logging.error('  word_action: %s', word_action)
      logging.error('  stack: %s', stack)
      logging.error('  original_legal_actions: %s', original_legal_actions)
      for action in original_legal_actions:
        logging.error('    action: %d = [%s]', action,
                      self.grammar.action_to_rule[str(action)])
      logging.error('  known_word_tries: %s', self.known_word_tries)
      logging.error('  current_trie: %s', current_trie)
      logging.error('  valid_word_actions: %s', self.valid_word_actions)
      logging.error('  current_valid_word_actions: %s',
                    current_valid_word_actions)
      logging.error('  crossed_search: %s', crossed_search_boundary)

    def _actions_since_last_query(taken_actions, simulated_actions):
      """Sequence of actions since last issued query."""
      if self.initial_stack and self.initial_stack[-1] == self.grammar.start():
        return simulated_actions
      actions = []
      for action in taken_actions[::-1]:
        actions.append(action)
        if self.grammar.productions()[action].lhs() == self.grammar.start():
          break
      return tuple(actions[::-1] + list(simulated_actions))

    current_node = f'{stack[-1].symbol()}'
    crossed_search_boundary = _compute_crossed_search_boundary()
    word_prefix = _find_word_prefix()
    word_action = _find_word_action()

    actions_since_last_query = _actions_since_last_query(
        self.actions, actions_sequence)
    if self.restriction_trie and self.restriction_trie.has_subtrie(
        actions_since_last_query):
      restricted_legal_actions = []
      for sequence in self.restriction_trie.iterkeys(actions_since_last_query):
        if len(sequence) > len(actions_since_last_query):
          restricted_legal_actions.append(
              sequence[len(actions_since_last_query)])
      if restricted_legal_actions:
        return tuple(set(restricted_legal_actions))

    # When we use a trie, especially when splitting the `Word` types into
    # _Wq_, _Wa_, _Wd_ and _Wt_, it's possible that the associated trie is
    # empty. We need to check that it's possible to generate a term from a
    # given trie.
    if (current_node == '_Q_' and not crossed_search_boundary and
        self.known_word_tries is not None):
      legal_actions = []
      for action in original_legal_actions:
        rule = self.grammar.action_to_rule[str(action)]
        if '_Wq_' in rule:
          if self.known_word_tries.question_trie:
            legal_actions.append(action)
        elif '_Wa_' in rule:
          if _get_restricted_trie(
              rule=rule, trie=self.known_word_tries.answer_trie):
            legal_actions.append(action)
        elif '_Wd_' in rule:
          if _get_restricted_trie(
              rule=rule, trie=self.known_word_tries.document_trie):
            legal_actions.append(action)
        elif '_Wt_' in rule:
          if _get_restricted_trie(
              rule=rule, trie=self.known_word_tries.title_trie):
            legal_actions.append(action)
        elif '_W_' in rule:
          if _get_restricted_trie(
              rule=rule, trie=self.known_word_tries.local_trie):
            legal_actions.append(action)
        else:
          legal_actions.append(action)
      return tuple(legal_actions)

    current_trie = None
    if self.known_word_tries is not None:
      if crossed_search_boundary:
        current_trie = self.known_word_tries.global_trie
      else:
        operator = _find_operator()
        if word_action.startswith('_Wq_'):
          current_trie = self.known_word_tries.question_trie
        elif word_action.startswith('_Wa_'):
          current_trie = _get_restricted_trie(
              rule=operator, trie=self.known_word_tries.answer_trie)
        elif word_action.startswith('_Wd_'):
          current_trie = _get_restricted_trie(
              rule=operator, trie=self.known_word_tries.document_trie)
        elif word_action.startswith('_Wt_'):
          current_trie = _get_restricted_trie(
              rule=operator, trie=self.known_word_tries.title_trie)
        else:
          current_trie = _get_restricted_trie(
              rule=operator, trie=self.known_word_tries.local_trie)

    current_valid_word_actions = None
    if self.valid_word_actions is not None:
      if word_action.startswith('_Wq_'):
        current_valid_word_actions = self.valid_word_actions.question_word_actions
      elif word_action.startswith('_Wa_'):
        current_valid_word_actions = self.valid_word_actions.answer_word_actions
      elif word_action.startswith('_Wd_'):
        current_valid_word_actions = self.valid_word_actions.document_word_actions
      elif word_action.startswith('_Wt_'):
        current_valid_word_actions = self.valid_word_actions.title_word_actions
      else:
        current_valid_word_actions = self.valid_word_actions.all_word_actions

    if (current_node in ['_Vw_', '_Vsh_'] and
        current_valid_word_actions is not None and not crossed_search_boundary):
      # For lexical rules, we apply restrictions if we were given the set of
      # available word action as per BERT state.
      # Additionally, if a known-trie has been passed, we will further restrict
      # to only those actions that yield valid words.

      # We have not yet crossed a "query" boundary, so we will restrict words to
      # be expanded to the word actions valid for the known BERT state.
      legal_actions = [
          action for action in legal_actions
          if action in current_valid_word_actions
      ]

    if current_trie is not None:
      if current_node in ['_W_', '_Wq_', '_Wa_', '_Wd_', '_Wt_']:
        # If we use a trie, we need to control the recursion of word
        # derivations. In particular, we need to check that
        #   1)  the _W_ -> _Vw_ _W-_ is only available if _any_ multi-word
        #       sequence could be generated;
        #   2)  the _W_ -> _Vw_ is only available if at least one initial word
        #       sequence can be generated.
        has_length_1_words = False
        has_length_2_words = False
        for word in current_trie.iterkeys(()):
          if has_length_1_words and has_length_2_words:
            break
          if len(word) == 1 and _can_construct(word,
                                               current_valid_word_actions):
            has_length_1_words = True
          elif not has_length_2_words and _can_construct(
              word, current_valid_word_actions):
            has_length_2_words = True

        valid_rules = []
        if has_length_1_words:
          valid_rules.append(
              self.grammar.rule_to_action[f'{current_node} -> _Vw_'])
        if has_length_2_words:
          valid_rules.append(
              self.grammar.rule_to_action[f'{current_node} -> _Vw_ _W-_'])
        legal_actions = [
            action for action in legal_actions if action in valid_rules
        ]
      elif current_node == '_W-_':
        # If we use a trie, we need to control the continuation of recursion.
        # In particular, we need to check that
        #   1) the _W-_ -> _Vsh_ is only available if at least one valid word
        #      can be generated by terminating and not recursing;
        #   2) the _W-_ -> _Vsh_ _W-_ is only available if at least one word
        #      can be generated if another recursion happens.
        has_immediate_word = False
        has_continuing_word = False
        try:
          for word in current_trie.iterkeys(word_prefix):
            if has_immediate_word and has_continuing_word:
              break
            if not has_immediate_word and len(
                word) == len(word_prefix) + 1 and _can_construct(
                    word[-1:], current_valid_word_actions):
              has_immediate_word = True
            elif not has_continuing_word and len(
                word) > len(word_prefix) + 1 and _can_construct(
                    word[len(word_prefix):], current_valid_word_actions):
              has_continuing_word = True
        except KeyError:
          _log_error('INVALID WORD PREFIX')

        valid_rules = []
        if has_immediate_word:
          valid_rules.append(self.grammar.rule_to_action['_W-_ -> _Vsh_'])
        if has_continuing_word:
          valid_rules.append(self.grammar.rule_to_action['_W-_ -> _Vsh_ _W-_'])
        legal_actions = [
            action for action in legal_actions if action in valid_rules
        ]
      elif current_node == '_Vw_':
        # Check if the last rule ends with `_W-_` to determine if we recursed
        # or not.
        rule = self.grammar.action_to_rule[str(full_tree_action_history[-1])]
        is_recursive = rule.endswith('_W-_')
        if crossed_search_boundary:
          legal_actions = (
              self.filter_legal_actions_for_vw_node_crossed_search_boundary(
                  legal_actions, is_recursive))
        else:
          legal_actions = self.filter_legal_actions_for_vw_node(
              legal_actions, is_recursive, current_trie)
      elif current_node == '_Vsh_':
        # Check if the last rule ends with `_W-_` to determine if we recursed
        # or not.
        rule = self.grammar.action_to_rule[str(full_tree_action_history[-1])]
        is_recursive = rule.endswith('_W-_')
        if crossed_search_boundary:
          legal_actions = (
              self.filter_legal_actions_for_vsh_node_crossed_search_boundary(
                  legal_actions, is_recursive, word_prefix))
        else:
          legal_actions = self.filter_legal_actions_for_vsh_node(
              legal_actions, is_recursive, word_prefix, current_trie)

    # These should not happen, but the Trie-based constraining surfaces the odd
    # error case every once in a while, so try to dump some info whenever this
    # happens.
    if not legal_actions:
      _log_error('INVALID TRANSITION')

    return tuple(legal_actions)
Returns the legal action space after a sequence of actions.
language/search_agents/muzero/state_tree.py