def reset(
      self,
      index: Optional[Union[int, environment_pb2.GetQueryResponse]] = None,
      documents: Optional[Sequence[environment_pb2.Document]] = None
  ) -> Tuple[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], Mapping[
      str, Any]]:
    """Resets the environment to play a query.

    Args:
      index:  If set as an integer, requests this specific query from the
        environment.  If set as a `GetQueryResponse`, resets for this query.  If
        `None`, a random query will be requested from the environment.
      documents:  If set, uses these results instead of actually querying the
        environment.  Useful during pre-training.

    Returns:
      This function changes the state of the environment, as the name suggests.
      It also returns a tuple of (Observation, InfoDict), where `Observation`
      is the observation for the agent at the beginning of the episode (after
      the original query has been issued).
    """

    if index is None or isinstance(index, int):
      query = self._get_query(index)
    else:
      query = index
    self._logging_info('Original Query [%d/%d]: %s | gold answers: %s.',
                       query.index, query.total, query.query, query.gold_answer)
    self.state = types.EnvState(original_query=query, k=common_flags.K.value)
    self.state.add_history_entry(
        self._get_env_output(
            query=utils.escape_for_lucene(query.query),
            original_query=query,
            documents=documents))

    if (common_flags.RELEVANCE_FEEDBACK_RESTRICT.value == 1 and
        self.training and query.gold_answer):
      target_query = (
          f'{utils.escape_for_lucene(query.query)} '
          f'+(contents:"{utils.escape_for_lucene(query.gold_answer[0])}")')
      self.state.target_documents = self._get_env_output(
          query=target_query, original_query=query).documents

      # Check that the target documents lead to a higher reward otherwise we
      # drop them.
      if self._target_reward() < self._initial_reward():
        self.state.target_documents = None

    # Signal we have not yet built the trie for these results.
    self.known_word_tries = None
    self.valid_word_actions = None
    self.use_rf_restrict = False

    documents = self.state.history[-1].documents
    # If we do not have documents at this point, we have to skip the episode.
    if not documents:
      raise mzcore.SkipEpisode(f'No documents for original query {query.query}')

    self.state.tree = state_tree.NQStateTree(grammar=self.grammar)

    self._logging_info('Initial result score: %s',
                       self._compute_reward(self.state.history[0].documents))

    self.action_history = []
    self.n_episode += 1
    obs = self._obs()
    info = self._info_dict(obs=obs)
    return obs.observation, info
Resets the environment to play a query.

    Args:
      index:  If set as an integer, requests this specific query from the
        environment.  If set as a `GetQueryResponse`, resets for this query.  If
        `None`, a random query will be requested from the environment.
      documents:  If set, uses these results instead of actually querying the
        environment.  Useful during pre-training.

    Returns:
      This function changes the state of the environment, as the name suggests.
      It also returns a tuple of (Observation, InfoDict), where `Observation`
      is the observation for the agent at the beginning of the episode (after
      the original query has been issued).
    
language/search_agents/muzero/env.py