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: Computes q_network td loss.
  
  Computes a temporal difference loss for training the q-network based
  on a corresponding Bellman update. Supports maxQ and SARSA
  style updates.
  
  Args:
      batch: a training batch object.
      boosted_rewards: a (batch_size, 1) shaped torch tensor with
          boosted rewards values.
      discount_tensor: a (batch_size, 1) torch tensor containing the
          discount to apply.
  
  Returns:
      A temporal difference loss object for training the q-network.
docstr_tok: ["Computes","q_network","td","loss",".","Computes","a","temporal","difference","loss","for","training","the","q-network","based","on","a","corresponding","Bellman","update",".","Supports","maxQ","and","SARSA","style","updates",".","Args",":","batch",":","a","training","batch","object",".","boosted_rewards",":","a","(","batch_size",",","1",")","shaped","torch","tensor","with","boosted","rewards","values",".","discount_tensor",":","a","(","batch_size",",","1",")","torch","tensor","containing","the","discount","to","apply",".","Returns",":","A","temporal","difference","loss","object","for","training","the","q-network","."]
code: def compute_td_loss(
      self,
      batch: rlt.DiscreteDqnInput,
      boosted_rewards: torch.Tensor,
      discount_tensor: torch.Tensor,
  ):
      """Computes q_network td loss.
  
      Computes a temporal difference loss for training the q-network based
      on a corresponding Bellman update. Supports maxQ and SARSA
      style updates.
  
      Args:
          batch: a training batch object.
          boosted_rewards: a (batch_size, 1) shaped torch tensor with
              boosted rewards values.
          discount_tensor: a (batch_size, 1) torch tensor containing the
              discount to apply.
  
      Returns:
          A temporal difference loss object for training the q-network.
      """
      not_done_mask = batch.not_terminal.float()
      (
          all_next_q_values,
          all_next_q_values_target,
      ) = self.get_detached_model_outputs(batch.next_state)
  
      if self.maxq_learning:
          # Compute max a' Q(s', a') over all possible actions using target network
          possible_next_actions_mask = (
              batch.possible_next_actions_mask.float()
          )
          if self.bcq:
              action_on_policy = get_valid_actions_from_imitator(
                  self.bcq_imitator,
                  batch.next_state,
                  self.bcq_drop_threshold,
              )
              possible_next_actions_mask *= action_on_policy
          (
              next_q_values,
              max_q_action_idxs,
          ) = self.get_max_q_values_with_target(
              all_next_q_values,
              all_next_q_values_target,
              possible_next_actions_mask,
          )
      else:
          # SARSA
          (
              next_q_values,
              max_q_action_idxs,
          ) = self.get_max_q_values_with_target(
              all_next_q_values,
              all_next_q_values_target,
              batch.next_action,
          )
  
      filtered_next_q_vals = next_q_values * not_done_mask
  
      target_q_values = boosted_rewards + (
          discount_tensor * filtered_next_q_vals
      )
  
      # Get Q-value of action taken
      all_q_values = self.q_network(batch.state)
      self.all_action_scores = all_q_values.detach()
      q_values = torch.sum(all_q_values * batch.action, 1, keepdim=True)
      td_loss = self.q_network_loss(q_values, target_q_values.detach())
      return td_loss
  
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 5 for calling the DQNTrainer obj's underlying member method code with example usage: obj.compute_td_loss(batch, boosted_rewards, discount_tensor) and returns: td_loss