class_docstr: A trainer for the DQN algorithm and its variants.
  
  Configures optimizers, builds the losses and train functions for the
  Q-learning based algorithm variants. Supports MaxQ and SARSA style
  TD-learning under both, standard and batch-constrained q-learning.
  During training, updates CPE metrics estimators.
  
  Attributes:
      double_q_learning: a boolean flag whether to use double-q learning.
      minibatch_size: an int number of samples per minibatch.
      minibatches_per_step: an int number of minibatch updates per step.
      q_network: a network object mapping states to q-values of all actions.
      q_network_target: a copy of q-network for training stability used in
          estimating q-values targets.
      q_network_optimizer: an optimizer object for training q-network.
      bcq: a config file for batch-constrained q-learning (BCQ).
      bcq_imitator: if using batch-constrained q-learning, the behavior
          policy used for BCQ training.
      all_action_scores: a torch tensor containing q-network
              predictions from the current states.
docstr: Builds loss functions for updating q- and reward networks.
  
  Args:
      training_batch: a training batch data object.
      batch_idx: an integer batch index.
  
  Yields:
      If a calc_cpe_in_training flag is True, yields a tuple
      (td_loss, reward_loss, cpe_metric_loss, soft_update_loss)
      for updating q_network, reward_network, q_network_cpe and target
      networks respectively. If calc_cpe_in_training is False, yields
      a tuple (td_loss, soft_update_loss) for updating q_network and
      q_network_target.
docstr_tok: ["Builds","loss","functions","for","updating","q-","and","reward","networks",".","Args",":","training_batch",":","a","training","batch","data","object",".","batch_idx",":","an","integer","batch","index",".","Yields",":","If","a","calc_cpe_in_training","flag","is","True",",","yields","a","tuple","(","td_loss",",","reward_loss",",","cpe_metric_loss",",","soft_update_loss",")","for","updating","q_network",",","reward_network",",","q_network_cpe","and","target","networks","respectively",".","If","calc_cpe_in_training","is","False",",","yields","a","tuple","(","td_loss",",","soft_update_loss",")","for","updating","q_network","and","q_network_target","."]
code: def train_step_gen(
      self, training_batch: rlt.DiscreteDqnInput, batch_idx: int
  ):
      """Builds loss functions for updating q- and reward networks.
  
      Args:
          training_batch: a training batch data object.
          batch_idx: an integer batch index.
  
      Yields:
          If a calc_cpe_in_training flag is True, yields a tuple
          (td_loss, reward_loss, cpe_metric_loss, soft_update_loss)
          for updating q_network, reward_network, q_network_cpe and target
          networks respectively. If calc_cpe_in_training is False, yields
          a tuple (td_loss, soft_update_loss) for updating q_network and
          q_network_target.
      """
      # TODO: calls to _maybe_run_optimizer removed, should be replaced with Trainer parameter
      self._check_input(training_batch)
  
      rewards = self.boost_rewards(
          training_batch.reward, training_batch.action
      )
      not_done_mask = training_batch.not_terminal.float()
  
      discount_tensor = self.compute_discount_tensor(
          training_batch, rewards
      )
      td_loss = self.compute_td_loss(
          training_batch, rewards, discount_tensor
      )
      yield td_loss
      td_loss = td_loss.detach()
  
      # Get Q-values of next states, used in computing cpe
      all_next_action_scores = self.q_network(
          training_batch.next_state
      ).detach()
      logged_action_idxs = torch.argmax(
          training_batch.action, dim=1, keepdim=True
      )
  
      yield from self._calculate_cpes(
          training_batch,
          training_batch.state,
          training_batch.next_state,
          self.all_action_scores,
          all_next_action_scores,
          logged_action_idxs,
          discount_tensor,
          not_done_mask,
      )
  
      if self.maxq_learning:
          possible_actions_mask = training_batch.possible_actions_mask
  
      if self.bcq:
          action_on_policy = get_valid_actions_from_imitator(
              self.bcq_imitator,
              training_batch.state,
              self.bcq_drop_threshold,
          )
          possible_actions_mask *= action_on_policy
  
      # Do we ever use model_action_idxs computed below?
      model_action_idxs = self.get_max_q_values(
          self.all_action_scores,
          # pyre-fixme[61]: `possible_actions_mask` is undefined, or not always
          #  defined.
          possible_actions_mask
          if self.maxq_learning
          else training_batch.action,
      )[1]
  
      self._log_dqn(
          td_loss,
          logged_action_idxs,
          training_batch,
          rewards,
          model_action_idxs,
      )
  
      # Use the soft update rule to update target network
      yield self.soft_update_result()
  
imports: import logging
  from typing import List, Optional, Tuple
  import reagent.core.types
  import torch
  from reagent.core.configuration import resolve_defaults
  from reagent.core.dataclasses import dataclass, field
  from reagent.core.parameters import EvaluationParameters, RLParameters
  from reagent.optimizer import Optimizer__Union, SoftUpdate
  from reagent.training.dqn_trainer_base import DQNTrainerBaseLightning
  from reagent.training.imitator_training import get_valid_actions_from_imitator
label_desc: Use image node_id 6 for calling the DQNTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types