code: def test__calculate_cpes(self):
      inp = DiscreteDqnInput(
          state=FeatureData(
              float_features=torch.rand(self.batch_size, self.state_dim)
          ),
          next_state=FeatureData(
              float_features=torch.rand(self.batch_size, self.state_dim)
          ),
          reward=torch.ones(self.batch_size, 1),
          time_diff=torch.ones(self.batch_size, 1) * 2,
          step=torch.ones(self.batch_size, 1) * 2,
          not_terminal=torch.ones(self.batch_size, 1),
          action=torch.tensor([[0, 1], [1, 0], [0, 1]]),
          next_action=torch.tensor([[1, 0], [0, 1], [1, 0]]),
          possible_actions_mask=torch.ones(
              self.batch_size, self.action_dim
          ),
          possible_next_actions_mask=torch.ones(
              self.batch_size, self.action_dim
          ),
          extras=ExtraData(),
      )
  
      mse_backward_type = type(
          torch.nn.functional.mse_loss(
              torch.tensor([1.0, 0.0], requires_grad=True),
              torch.zeros(1),
          ).grad_fn
      )
      trainer = self._construct_trainer()
      trainer._initialize_cpe(
          self.reward_network,
          self.q_network_cpe,
          self.q_network_cpe_target,
          self.params.optimizer,
      )
      trainer.reward_network = self.reward_network
      not_done_mask = inp.not_terminal.float()
      discount_tensor = torch.ones(inp.reward.shape)
      all_q_values = self.q_network(inp.state)
      all_action_scores = all_q_values.detach()
      all_next_action_scores = self.q_network(inp.next_state).detach()
      logged_action_idxs = torch.tensor([[1], [0], [1]])
  
      cpes_gen = trainer._calculate_cpes(
          inp,
          inp.state,
          inp.next_state,
          all_action_scores,
          all_next_action_scores,
          logged_action_idxs,
          discount_tensor,
          not_done_mask,
      )
      cpes = list(cpes_gen)
  
      ## expected reward_loss
  
      reward_target = inp.reward
      reward_estimate = trainer.reward_network(inp.state)
      # rewards at offset + logged action idx
      reward_estimate = reward_estimate.gather(
          1, torch.tensor([[1], [0], [1]])
      )
      mse_reward_loss = torch.nn.functional.mse_loss(
          reward_estimate, reward_target
      )
  
      ## expected metric_q_value_loss
  
      # assuming masked_softmax is tested elsewhere,
      # we can treat this as expected ground truth value
      model_propensities_next_states = masked_softmax(
          all_next_action_scores,
          inp.possible_next_actions_mask
          if trainer.maxq_learning
          else inp.next_action,
          trainer.rl_temperature,
      )
      metric_q_values = trainer.q_network_cpe(inp.state).gather(
          1, torch.tensor([[1], [0], [1]])
      )
      metrics_target_q_values = trainer.q_network_cpe_target(
          inp.next_state
      )
      per_metric_next_q_values = torch.sum(
          metrics_target_q_values * model_propensities_next_states,
          1,
          keepdim=True,
      )
      per_metric_next_q_values *= not_done_mask
      metrics_target_q_values = (
          reward_target + discount_tensor * per_metric_next_q_values
      )
      metric_q_value_loss = trainer.q_network_loss(
          metric_q_values, metrics_target_q_values
      )
      self.assertEqual(len(cpes), 2)
      self.assertEqual(type(cpes[0].grad_fn), mse_backward_type)
      self.assertEqual(type(cpes[1].grad_fn), mse_backward_type)
      self.assertEqual(cpes[0], mse_reward_loss)
      self.assertEqual(cpes[1], metric_q_value_loss)
  
imports: import unittest
  from typing import List, Optional, Tuple
  import torch
  from reagent.core.parameters import EvaluationParameters, RLParameters
  from reagent.core.torch_utils import masked_softmax
  from reagent.core.types import DiscreteDqnInput, ExtraData, FeatureData
  from reagent.evaluation.evaluation_data_page import EvaluationDataPage
  from reagent.evaluation.evaluator import get_metrics_to_score
  from reagent.models.dqn import FullyConnectedDQN
  from reagent.training.dqn_trainer_base import DQNTrainerBaseLightning
  from reagent.training.parameters import DQNTrainerParameters
  from reagent.workflow.types import RewardOptions
label_desc: Use image node_id 8 for calling the TestDQNTrainerBaseLightning obj's underlying member method code with example usage: obj.test__calculate_cpes() without return types