class_docstr: Base trainer class for the DQN algorihtm variants.
  
  A base class for training variants of the DQN algorithm.
  Provides methods for q-values calculation, reward boosting, and also
  initializing and training counterfactual policy evaluation
  metrics estimators (see https://arxiv.org/abs/1811.00260 for detail).
  
  Attributes:
      rl_parameters: an object of a reagent.core.parameters.RLParameters
          data class.
      metrics_to_score: a list of strings representing reward metrics.
      evaluation_parameters: an object of a
          reagent.core.parameters.EvaluationParameters data class.
      q_network_loss: a loss function object, currently either
          torch.nn.functional.mse_loss or torch.nn.functional.smooth_l1_loss.
      reward_boosts: a torch tensor containing reward boost values at
          indices corresponding to respective discrete actions.
      calc_cpe_in_training: a boolean flag whether to calculate CPE metrics
          during training.
      reward_network: a network object mapping states to rewards and
          extra reward metrics. Given the number of discrete actions n and
          e.g. a set of k reward metrics `m_1`, `m_2`, ..., and `m_k`, the
          reward_network output has shape (batch_size, n * k), where
          the [:, (k - 1) * n : k * n] sub-tensor contains the metric `m_k`
          q-values estimates for each of the n actions. The order of metrics
          in the output is the following: the environment reward, followed by
          extra metrics sorted by their names in RewardOptions.metric_reward_values
          config.
      reward_network_optimizer: an optimizer object for training
          reward network.
      q_network_cpe: a network object mapping states to CPE metrics.
      q_network_cpe_target: a copy of q_network_cpe for training stability.
      q_network_cpe_optimizer: an optimizer object for training q_network_cpe.
      reward_idx_offsets: a flat torch tensor containing integer offsets of
          different reward metrics in the reward_network output.
      evaluator: an object of a reagent.evaluation.evaluator.Evaluator class.
docstr: Computes losses for the reward and the cpe q-values networks.
  
  Based on the actions taken idxs, computes the mse loss for training
  the reward network and either mse or huber loss for the cpe q-network.
  
  Args:
      training_batch: a training batch data object.
      states: a (batch_size, state_dim) torch tensor containing
          current states.
      next_states: a (batch_size, state_dim) torch tensor containing
          next states.
      all_action_scores: a torch tensor containing q-network
          predictions from the current states.
      all_next_action_scores: a torch tensor containing q-network
          predictions from the next states.
      logged_action_idxs: a (batch_size, 1) torch tensor with integer
          executed actions indices.
      discount_tensor: a (batch_size, 1) torch tensor containing the
          discount to apply, e.g. gamma ** k where k is a number of
          steps used in td-error estimation.
      not_done_mask: a (batch_size, 1) torch tensor with boolean values
          indicating whether corresponding state is terminal.
  Yields:
      A reward network and a CPE q-network loss objects.
docstr_tok: ["Computes","losses","for","the","reward","and","the","cpe","q-values","networks",".","Based","on","the","actions","taken","idxs",",","computes","the","mse","loss","for","training","the","reward","network","and","either","mse","or","huber","loss","for","the","cpe","q-network",".","Args",":","training_batch",":","a","training","batch","data","object",".","states",":","a","(","batch_size",",","state_dim",")","torch","tensor","containing","current","states",".","next_states",":","a","(","batch_size",",","state_dim",")","torch","tensor","containing","next","states",".","all_action_scores",":","a","torch","tensor","containing","q-network","predictions","from","the","current","states",".","all_next_action_scores",":","a","torch","tensor","containing","q-network","predictions","from","the","next","states",".","logged_action_idxs",":","a","(","batch_size",",","1",")","torch","tensor","with","integer","executed","actions","indices",".","discount_tensor",":","a","(","batch_size",",","1",")","torch","tensor","containing","the","discount","to","apply",",","e.g",".","gamma","*","*","k","where","k","is","a","number","of","steps","used","in","td-error","estimation",".","not_done_mask",":","a","(","batch_size",",","1",")","torch","tensor","with","boolean","values","indicating","whether","corresponding","state","is","terminal",".","Yields",":","A","reward","network","and","a","CPE","q-network","loss","objects","."]
code: def _calculate_cpes(
      self,
      training_batch,
      states,
      next_states,
      all_action_scores,
      all_next_action_scores,
      logged_action_idxs,
      discount_tensor,
      not_done_mask,
  ):
      """Computes losses for the reward and the cpe q-values networks.
  
      Based on the actions taken idxs, computes the mse loss for training
      the reward network and either mse or huber loss for the cpe q-network.
  
      Args:
          training_batch: a training batch data object.
          states: a (batch_size, state_dim) torch tensor containing
              current states.
          next_states: a (batch_size, state_dim) torch tensor containing
              next states.
          all_action_scores: a torch tensor containing q-network
              predictions from the current states.
          all_next_action_scores: a torch tensor containing q-network
              predictions from the next states.
          logged_action_idxs: a (batch_size, 1) torch tensor with integer
              executed actions indices.
          discount_tensor: a (batch_size, 1) torch tensor containing the
              discount to apply, e.g. gamma ** k where k is a number of
              steps used in td-error estimation.
          not_done_mask: a (batch_size, 1) torch tensor with boolean values
              indicating whether corresponding state is terminal.
      Yields:
          A reward network and a CPE q-network loss objects.
      """
      if not self.calc_cpe_in_training:
          return
      if training_batch.extras.metrics is None:
          metrics_reward_concat_real_vals = training_batch.reward
      else:
          metrics_reward_concat_real_vals = torch.cat(
              (training_batch.reward, training_batch.extras.metrics),
              dim=1,
          )
  
      model_propensities_next_states = masked_softmax(
          all_next_action_scores,
          training_batch.possible_next_actions_mask
          if self.maxq_learning
          else training_batch.next_action,
          self.rl_temperature,
      )
  
      ######### Train separate reward network for CPE evaluation #############
      reward_estimates = self.reward_network(states)
      reward_estimates_for_logged_actions = reward_estimates.gather(
          1, self.reward_idx_offsets + logged_action_idxs
      )
      reward_loss = F.mse_loss(
          reward_estimates_for_logged_actions,
          metrics_reward_concat_real_vals,
      )
      yield reward_loss
  
      ######### Train separate q-network for CPE evaluation #############
      metric_q_values = self.q_network_cpe(states).gather(
          1, self.reward_idx_offsets + logged_action_idxs
      )
      all_metrics_target_q_values = torch.chunk(
          self.q_network_cpe_target(next_states).detach(),
          len(self.metrics_to_score),
          dim=1,
      )
      target_metric_q_values = []
      for i, per_metric_target_q_values in enumerate(
          all_metrics_target_q_values
      ):
          per_metric_next_q_values = torch.sum(
              per_metric_target_q_values
              * model_propensities_next_states,
              1,
              keepdim=True,
          )
          per_metric_next_q_values = (
              per_metric_next_q_values * not_done_mask
          )
          per_metric_target_q_values = metrics_reward_concat_real_vals[
              :, i : i + 1
          ] + (discount_tensor * per_metric_next_q_values)
          target_metric_q_values.append(per_metric_target_q_values)
  
      target_metric_q_values = torch.cat(target_metric_q_values, dim=1)
      metric_q_value_loss = self.q_network_loss(
          metric_q_values, target_metric_q_values
      )
  
      # The model_propensities computed below are not used right now. The CPE graphs in the Outputs
      # tab use model_propensities computed in the function create_from_tensors_dqn() in evaluation_data_page.py,
      # which is called on the eval_table_sample in the gather_eval_data() function below.
      model_propensities = masked_softmax(
          all_action_scores,
          training_batch.possible_actions_mask
          if self.maxq_learning
          else training_batch.action,
          self.rl_temperature,
      )
      # Extract rewards predicted by the reward_network. The other columns will
      # give predicted values for other metrics, if such were specified.
      model_rewards = reward_estimates[
          :,
          torch.arange(
              self.reward_idx_offsets[0],
              self.reward_idx_offsets[0] + self.num_actions,
          ),
      ]
  
      self.reporter.log(
          reward_loss=reward_loss,
          model_propensities=model_propensities,
          model_rewards=model_rewards,
      )
  
      yield metric_q_value_loss
  
imports: import logging
  from abc import abstractmethod
  from typing import Dict, List, Optional
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.parameters import EvaluationParameters, RLParameters
  from reagent.core.torch_utils import masked_softmax
  from reagent.evaluation.evaluation_data_page import EvaluationDataPage
  from reagent.evaluation.evaluator import Evaluator
  from reagent.optimizer import Optimizer__Union
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.rl_trainer_pytorch import RLTrainerMixin
label_desc: Use image node_id 9 for calling the DQNTrainerBaseLightning obj's underlying member method code with example usage: obj._calculate_cpes(training_batch, states, next_states, all_action_scores, all_next_action_scores, logged_action_idxs, discount_tensor, not_done_mask) without return types