def __init__(self,
               grammar_str: str,
               restrictions: Optional[Dict[str, List[str]]] = None):
    r"""Context Free Grammar for the NQ environment.

    Args:
      grammar_str: A grammar string in the form required by nltk.CFG.fromstring.
        E.g. "A -> B \n B -> 'c'"
      restrictions: A dictionary that maps from the string representation of the
        last two terminals to a list of identifiers for the allowed
        vocabularies. E.g. {'[pos][contents]': ['Vap', 'Vat'],}

    Returns:
      Context Free Grammar for the NQ environment.
    """

    self.nonterminals = set(
        [production.lhs() for production in self.productions()])
    self.production_lookup = collections.defaultdict(list)
    # Maps a rule to the corresponding action.
    self.rule_to_action = {}
    # Maps an action back to the rule.
    self.action_to_rule = {}
    # Maps a terminal (word piece) to its (corresponding rule's) action.
    self.terminal_to_action = {}

    # Associate each production with an integer which corresponds to the id of
    # an agent action.
    for i, production in enumerate(self.productions()):
      self.production_lookup[str(production.lhs())].append((i, production))
      self.rule_to_action[str(production)] = i
      self.action_to_rule[str(i)] = str(production)
      if production.is_lexical() and len(production.rhs()) == 1:
        self.terminal_to_action[production.rhs()[0]] = i

    self.min_steps_to_terminal = {
        non_terminal: self._min_steps_to_terminal(non_terminal)
        for non_terminal in self.nonterminals
    }
    self.restrictions = restrictions
Context Free Grammar for the NQ environment.

    Args:
      grammar_str: A grammar string in the form required by nltk.CFG.fromstring.
        E.g. "A -> B \n B -> 'c'"
      restrictions: A dictionary that maps from the string representation of the
        last two terminals to a list of identifiers for the allowed
        vocabularies. E.g. {'[pos][contents]': ['Vap', 'Vat'],}

    Returns:
      Context Free Grammar for the NQ environment.
    
language/search_agents/muzero/state_tree.py