docstr: Loss function.
  
  Args:
      model: The ModelV2 to run foward pass on.
      dist_class: The distribution of this policy.
      train_batch: Training SampleBatch.
          Keys and shapes: {
              OBS: [batch_size, max_seq_len, obs_dim],
              ACTIONS: [batch_size, max_seq_len, act_dim],
              RETURNS_TO_GO: [batch_size, max_seq_len + 1, 1],
              TIMESTEPS: [batch_size, max_seq_len],
              ATTENTION_MASKS: [batch_size, max_seq_len],
          }
  Returns:
      Loss scalar tensor.
docstr_tok: ["Loss","function",".","Args",":","model",":","The","ModelV2","to","run","foward","pass","on",".","dist_class",":","The","distribution","of","this","policy",".","train_batch",":","Training","SampleBatch",".","Keys","and","shapes",":","{","OBS",":","[","batch_size",",","max_seq_len",",","obs_dim","]",",","ACTIONS",":","[","batch_size",",","max_seq_len",",","act_dim","]",",","RETURNS_TO_GO",":","[","batch_size",",","max_seq_len","+","1",",","1","]",",","TIMESTEPS",":","[","batch_size",",","max_seq_len","]",",","ATTENTION_MASKS",":","[","batch_size",",","max_seq_len","]",",","}","Returns",":","Loss","scalar","tensor","."]
code: def loss(
      self,
      model: ModelV2,
      dist_class: Type[TorchDistributionWrapper],
      train_batch: SampleBatch,
  ) -> Union[TensorType, List[TensorType]]:
      """Loss function.
  
      Args:
          model: The ModelV2 to run foward pass on.
          dist_class: The distribution of this policy.
          train_batch: Training SampleBatch.
              Keys and shapes: {
                  OBS: [batch_size, max_seq_len, obs_dim],
                  ACTIONS: [batch_size, max_seq_len, act_dim],
                  RETURNS_TO_GO: [batch_size, max_seq_len + 1, 1],
                  TIMESTEPS: [batch_size, max_seq_len],
                  ATTENTION_MASKS: [batch_size, max_seq_len],
              }
      Returns:
          Loss scalar tensor.
      """
      train_batch = self._lazy_tensor_dict(train_batch)
  
      # policy forward and get prediction
      model_out, _ = self.model(train_batch)  # noop, just returns obs.
      preds = self.model.get_prediction(model_out, train_batch)
  
      # get the regression targets
      targets = self.model.get_targets(model_out, train_batch)
  
      # get the attention masks for masked-loss
      masks = train_batch[SampleBatch.ATTENTION_MASKS]
  
      # compute loss
      loss = self._masked_loss(preds, targets, masks)
  
      self.log("cur_lr", torch.tensor(self.cur_lr))
  
      return loss
  
imports: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Box, Discrete
  from rllib_dt.dt.dt_torch_model import DTTorchModel
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.mingpt import configure_gpt_optimizer
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDeterministic, TorchDistributionWrapper
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI, override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.threading import with_lock
  from ray.rllib.utils.torch_utils import apply_grad_clipping
  from ray.rllib.utils.typing import TensorShape, TensorStructType, TensorType, TrainerConfigDict
label_desc: Use image node_id 11 for calling the DTTorchPolicy obj's underlying member method code with example usage: obj.loss(model, dist_class, train_batch) and returns: loss