def legal_actions(self):
    """Return only rules having the to-be-expanded non-terminal as lhs."""
    current_node = self._find_node_to_expand()
    if current_node is None:
      return []
    legal_actions = self.grammar.production_lookup[str(current_node.label())]
    return [idx for idx, legal_action in legal_actions]
Return only rules having the to-be-expanded non-terminal as lhs.
language/search_agents/muzero/state_tree.py