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: Initializes CPE networks, optimizers and an evaluator object.
  
  Given the number of discrete actions n and e.g. a set of reward
  metrics `a`, `b`, and `c`, the reward_network output has shape
  (batch_size, n * 3), where the [:, :n] sub-tensor contains the metric
  `a` q-values estimates for each action, the [:, n: 2 * n] sub-tensor
  contains the metric `b` q-values estimates, etc. In addition to
  initializing the reward and cpe networks, this function computes
  the offsets of each metric in the reward network output tensor,
  i.e. the torch.tensor([0, n, 2 * n]).
  
  Args:
      reward_network: a network object mapping states to rewards and
          extra reward metrics.
      q_network_cpe: a network object mapping states to CPE metrics.
          The network output has the same shape as the reward_network
          output: (batch_size, n * k), where n is a number of discrete
          actions, k is a number of reward metrics.
      q_network_cpe_target: a copy of q_network_cpe for training stability.
      optimizer: an optimizer object for training q_network_cpe.
docstr_tok: ["Initializes","CPE","networks",",","optimizers","and","an","evaluator","object",".","Given","the","number","of","discrete","actions","n","and","e.g",".","a","set","of","reward","metrics","`","a","`",",","`","b","`",",","and","`","c","`",",","the","reward_network","output","has","shape","(","batch_size",",","n","*","3",")",",","where","the","[",":",",",":","n","]","sub-tensor","contains","the","metric","`","a","`","q-values","estimates","for","each","action",",","the","[",":",",","n",":","2","*","n","]","sub-tensor","contains","the","metric","`","b","`","q-values","estimates",",","etc",".","In","addition","to","initializing","the","reward","and","cpe","networks",",","this","function","computes","the","offsets","of","each","metric","in","the","reward","network","output","tensor",",","i.e",".","the","torch.tensor","(","[","0",",","n",",","2","*","n","]",")",".","Args",":","reward_network",":","a","network","object","mapping","states","to","rewards","and","extra","reward","metrics",".","q_network_cpe",":","a","network","object","mapping","states","to","CPE","metrics",".","The","network","output","has","the","same","shape","as","the","reward_network","output",":","(","batch_size",",","n","*","k",")",",","where","n","is","a","number","of","discrete","actions",",","k","is","a","number","of","reward","metrics",".","q_network_cpe_target",":","a","copy","of","q_network_cpe","for","training","stability",".","optimizer",":","an","optimizer","object","for","training","q_network_cpe","."]
code: def _initialize_cpe(
      self,
      reward_network,
      q_network_cpe,
      q_network_cpe_target,
      optimizer: Optimizer__Union,
  ) -> None:
      """Initializes CPE networks, optimizers and an evaluator object.
  
      Given the number of discrete actions n and e.g. a set of reward
      metrics `a`, `b`, and `c`, the reward_network output has shape
      (batch_size, n * 3), where the [:, :n] sub-tensor contains the metric
      `a` q-values estimates for each action, the [:, n: 2 * n] sub-tensor
      contains the metric `b` q-values estimates, etc. In addition to
      initializing the reward and cpe networks, this function computes
      the offsets of each metric in the reward network output tensor,
      i.e. the torch.tensor([0, n, 2 * n]).
  
      Args:
          reward_network: a network object mapping states to rewards and
              extra reward metrics.
          q_network_cpe: a network object mapping states to CPE metrics.
              The network output has the same shape as the reward_network
              output: (batch_size, n * k), where n is a number of discrete
              actions, k is a number of reward metrics.
          q_network_cpe_target: a copy of q_network_cpe for training stability.
          optimizer: an optimizer object for training q_network_cpe.
      """
      if not self.calc_cpe_in_training:
          self.reward_network = None
          return
  
      assert (
          reward_network is not None
      ), "reward_network is required for CPE"
      self.reward_network = reward_network
      self.reward_network_optimizer = optimizer
      assert (
          q_network_cpe is not None and q_network_cpe_target is not None
      ), "q_network_cpe and q_network_cpe_target are required for CPE"
      self.q_network_cpe = q_network_cpe
      self.q_network_cpe_target = q_network_cpe_target
      self.q_network_cpe_optimizer = optimizer
      num_output_nodes = len(self.metrics_to_score) * self.num_actions
      reward_idx_offsets = torch.arange(
          0,
          num_output_nodes,
          self.num_actions,
          dtype=torch.long,
      )
      self.register_buffer("reward_idx_offsets", reward_idx_offsets)
  
      reward_stripped_metrics_to_score = (
          self.metrics_to_score[:-1]
          if len(self.metrics_to_score) > 1
          else None
      )
      self.evaluator = Evaluator(
          self._actions,
          self.rl_parameters.gamma,
          self,
          metrics_to_score=reward_stripped_metrics_to_score,
      )
  
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 7 for calling the DQNTrainerBaseLightning obj's underlying member method code with example usage: obj._initialize_cpe(reward_network, q_network_cpe, q_network_cpe_target, optimizer) without return types