class_docstr: PyTorch policy class used with SimpleQTrainer.
docstr: Compute loss for SimpleQ.
  
  Args:
      model: The Model to calculate the loss for.
      dist_class: The action distr. class.
      train_batch: The training data.
  
  Returns:
      The SimpleQ loss tensor given the input batch.
docstr_tok: ["Compute","loss","for","SimpleQ",".","Args",":","model",":","The","Model","to","calculate","the","loss","for",".","dist_class",":","The","action","distr",".","class",".","train_batch",":","The","training","data",".","Returns",":","The","SimpleQ","loss","tensor","given","the","input","batch","."]
code: def loss(
      self,
      model: ModelV2,
      dist_class: Type[TorchDistributionWrapper],
      train_batch: SampleBatch,
  ) -> Union[TensorType, List[TensorType]]:
      """Compute loss for SimpleQ.
  
      Args:
          model: The Model to calculate the loss for.
          dist_class: The action distr. class.
          train_batch: The training data.
  
      Returns:
          The SimpleQ loss tensor given the input batch.
      """
      target_model = self.target_models[model]
  
      # q network evaluation
      q_t = self._compute_q_values(
          model, train_batch[SampleBatch.CUR_OBS], is_training=True
      )
  
      # target q network evalution
      q_tp1 = self._compute_q_values(
          target_model,
          train_batch[SampleBatch.NEXT_OBS],
          is_training=True,
      )
  
      # q scores for actions which we know were selected in the given state.
      one_hot_selection = F.one_hot(
          train_batch[SampleBatch.ACTIONS].long(), self.action_space.n
      )
      q_t_selected = torch.sum(q_t * one_hot_selection, 1)
  
      # compute estimate of best possible value starting from state at t + 1
      dones = train_batch[SampleBatch.TERMINATEDS].float()
      q_tp1_best_one_hot_selection = F.one_hot(
          torch.argmax(q_tp1, 1), self.action_space.n
      )
      q_tp1_best = torch.sum(q_tp1 * q_tp1_best_one_hot_selection, 1)
      q_tp1_best_masked = (1.0 - dones) * q_tp1_best
  
      # compute RHS of bellman equation
      q_t_selected_target = (
          train_batch[SampleBatch.REWARDS]
          + self.config["gamma"] * q_tp1_best_masked
      )
  
      # Compute the error (Square/Huber).
      td_error = q_t_selected - q_t_selected_target.detach()
      loss = torch.mean(huber_loss(td_error))
  
      # Store values for stats function in model (tower), such that for
      # multi-GPU, we do not override them during the parallel loss phase.
      model.tower_stats["loss"] = loss
      # TD-error tensor in final stats
      # will be concatenated and retrieved for each individual batch item.
      model.tower_stats["td_error"] = td_error
  
      return loss
  
imports: import logging
  from typing import Any, Dict, List, Tuple, Type, Union
  from rllib_simple_q.simple_q.utils import make_q_models
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDistributionWrapper
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule, TargetNetworkMixin
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import concat_multi_gpu_td_errors, huber_loss
  from ray.rllib.utils.typing import TensorStructType, TensorType
label_desc: Use image node_id 4 for calling the SimpleQTorchPolicy obj's underlying member method code with example usage: obj.loss(model, dist_class, train_batch) and returns: loss