def special_episode_statistics_learner(self, return_as_dict=False):
    """Everything you want logged for episodes.

    These are only used for logging.  Actual rewards for learning need to be
    returned via `step`.

    Args:
      return_as_dict:  If true, returns a "human-readable" dict, rather than a
        tensorflow-compatible tuple.

    Returns:
      Metrics/Statistics either as a tuple of numpy.ndarrays, or a dictionary.
    """
    stats = {}
    documents_list = self.get_final_document_list()
    original_documents_list = self.state.history[0].documents

    stats['ndcg_score'] = self.state.score(
        'ndcg', documents_list=documents_list, k=int(common_flags.K.value))
    stats['ndcg_score_improvement'] = stats['ndcg_score'] - self.state.score(
        'ndcg',
        documents_list=original_documents_list,
        k=int(common_flags.K.value))

    stats['em_at_1'] = self.state.score(
        'em_at_k', documents_list=documents_list, k=1)
    stats['em_at_1_improvement'] = stats['em_at_1'] - self.state.score(
        'em_at_k', documents_list=original_documents_list, k=1)

    stats[f'em_at_{int(common_flags.K.value)}'] = self.state.score(
        'em_at_k',
        documents_list=documents_list,
        k=int(common_flags.K.value))
    stats[f'em_at_{int(common_flags.K.value)}_improvement'] = stats[
        f'em_at_{int(common_flags.K.value)}'] - self.state.score(
            'em_at_k',
            documents_list=original_documents_list,
            k=int(common_flags.K.value))

    stats[f'recall_at_{int(common_flags.K.value)}'] = self.state.score(
        'recall_at_k',
        documents_list=documents_list,
        k=int(common_flags.K.value))
    stats[f'recall_at_{int(common_flags.K.value)}_improvement'] = stats[
        f'recall_at_{int(common_flags.K.value)}'] - self.state.score(
            'recall_at_k',
            documents_list=original_documents_list,
            k=int(common_flags.K.value))

    stats['recall_at_1'] = self.state.score(
        'recall_at_k', documents_list=documents_list, k=1)
    stats['recall_at_1_improvement'] = stats['recall_at_1'] - self.state.score(
        'recall_at_k', documents_list=original_documents_list, k=1)

    stats['documents_explored'] = len(
        self.state.sorted_unique_documents(
            step=self.state.num_completed_requests))

    if return_as_dict:
      return stats

    return (
        np.float32(stats['ndcg_score']),
        np.float32(stats['ndcg_score_improvement']),
        np.float32(stats['em_at_1']),
        np.float32(stats['em_at_1_improvement']),
        np.float32(stats[f'em_at_{int(common_flags.K.value)}']),
        np.float32(stats[f'em_at_{int(common_flags.K.value)}_improvement']),
        np.float32(stats[f'recall_at_{int(common_flags.K.value)}']),
        np.float32(
            stats[f'recall_at_{int(common_flags.K.value)}_improvement']),
        np.float32(stats['recall_at_1']),
        np.float32(stats['recall_at_1_improvement']),
        np.float32(stats['documents_explored']),
    )
Everything you want logged for episodes.

    These are only used for logging.  Actual rewards for learning need to be
    returned via `step`.

    Args:
      return_as_dict:  If true, returns a "human-readable" dict, rather than a
        tensorflow-compatible tuple.

    Returns:
      Metrics/Statistics either as a tuple of numpy.ndarrays, or a dictionary.
    
language/search_agents/muzero/env.py