docstr: Do the preparations needed for different estimators.
  Some of the calcuations are shared, so we centralize all the work here.
  
  :param evaluation_dataset_as_transitions: The evaluation dataset in the form of transitions.
  :param batch_size: The batch size to use.
  :param reward_model: A reward model to be used by DR
  :param q_network: The Q network whose its policy we evaluate.
  :param network_keys: The network keys used for feeding the neural networks.
  :return:
docstr_tok: ["Do","the","preparations","needed","for","different","estimators",".","Some","of","the","calcuations","are","shared",",","so","we","centralize","all","the","work","here",".",":","param","evaluation_dataset_as_transitions",":","The","evaluation","dataset","in","the","form","of","transitions",".",":","param","batch_size",":","The","batch","size","to","use",".",":","param","reward_model",":","A","reward","model","to","be","used","by","DR",":","param","q_network",":","The","Q","network","whose","its","policy","we","evaluate",".",":","param","network_keys",":","The","network","keys","used","for","feeding","the","neural","networks",".",":","return",":"]
code: def _prepare_ope_shared_stats(
      self,
      evaluation_dataset_as_transitions: List[Transition],
      batch_size: int,
      q_network: Architecture,
      network_keys: List,
  ) -> OpeSharedStats:
      """
      Do the preparations needed for different estimators.
      Some of the calcuations are shared, so we centralize all the work here.
  
      :param evaluation_dataset_as_transitions: The evaluation dataset in the form of transitions.
      :param batch_size: The batch size to use.
      :param reward_model: A reward model to be used by DR
      :param q_network: The Q network whose its policy we evaluate.
      :param network_keys: The network keys used for feeding the neural networks.
      :return:
      """
  
      assert self.is_gathered_static_shared_data, (
          "gather_static_shared_stats() should be called once before "
          "calling _prepare_ope_shared_stats()"
      )
      # IPS
      all_policy_probs = []
      all_v_values_reward_model_based, all_v_values_q_model_based = (
          [],
          [],
      )
  
      for i in range(
          math.ceil(len(evaluation_dataset_as_transitions) / batch_size)
      ):
          batch = evaluation_dataset_as_transitions[
              i * batch_size : (i + 1) * batch_size
          ]
          batch_for_inference = Batch(batch)
  
          # we always use the first Q head to calculate OPEs. might want to change this in the future.
          # for instance, this means that for bootstrapped dqn we always use the first QHead to calculate the OPEs.
          q_values, sm_values = q_network.predict(
              batch_for_inference.states(network_keys),
              outputs=[
                  q_network.output_heads[0].q_values,
                  q_network.output_heads[0].softmax,
              ],
          )
  
          all_policy_probs.append(sm_values)
          all_v_values_reward_model_based.append(
              np.sum(
                  all_policy_probs[-1]
                  * self.all_reward_model_rewards[i],
                  axis=1,
              )
          )
          all_v_values_q_model_based.append(
              np.sum(all_policy_probs[-1] * q_values, axis=1)
          )
  
          for j, t in enumerate(batch):
              t.update_info(
                  {
                      "q_value": q_values[j],
                      "softmax_policy_prob": all_policy_probs[-1][j],
                      "v_value_q_model_based": all_v_values_q_model_based[
                          -1
                      ][
                          j
                      ],
                  }
              )
  
      all_policy_probs = np.concatenate(all_policy_probs, axis=0)
      all_v_values_reward_model_based = np.concatenate(
          all_v_values_reward_model_based, axis=0
      )
  
      # generate model probabilities
      new_policy_prob = all_policy_probs[
          np.arange(self.all_actions.shape[0]), self.all_actions
      ]
      rho_all_dataset = new_policy_prob / self.all_old_policy_probs
  
      return OpeSharedStats(
          self.all_reward_model_rewards,
          all_policy_probs,
          all_v_values_reward_model_based,
          self.all_rewards,
          self.all_actions,
          self.all_old_policy_probs,
          new_policy_prob,
          rho_all_dataset,
      )
  
file_path: coach/rl_coach/off_policy_evaluators/ope_manager.py
imports: import math
  from collections import namedtuple
  import numpy
  from typing import List
  from rl_coach.architectures.architecture import Architecture
  from rl_coach.core_types import Episode, Batch
  from rl_coach.off_policy_evaluators.bandits.doubly_robust import DoublyRobust
  from rl_coach.off_policy_evaluators.rl.sequential_doubly_robust import SequentialDoublyRobust
  from rl_coach.core_types import Transition
  from rl_coach.off_policy_evaluators.rl.weighted_importance_sampling import WeightedImportanceSampling
label_desc: Use image node_id 2 for calling the OpeManager obj's underlying member method code with example usage: obj._prepare_ope_shared_stats(evaluation_dataset_as_transitions, batch_size, q_network, network_keys) and returns: OpeSharedStats