class_docstr: Extension of standard TorchModelV2 to provide dueling-Q functionality.
class_docstr_tok: ["Extension","of","standard","TorchModelV2","to","provide","dueling-Q","functionality","."]
docstr: Returns distributional values for Q(s, a) given a state embedding.
  
  Override this in your custom model to customize the Q output head.
  
  Args:
      model_out: Embedding from the model layers.
  
  Returns:
      (action_scores, logits, dist) if num_atoms == 1, otherwise
      (action_scores, z, support_logits_per_action, logits, dist)
docstr_tok: ["Returns","distributional","values","for","Q","(","s",",","a",")","given","a","state","embedding",".","Override","this","in","your","custom","model","to","customize","the","Q","output","head",".","Args",":","model_out",":","Embedding","from","the","model","layers",".","Returns",":","(","action_scores",",","logits",",","dist",")","if","num_atoms","==","1",",","otherwise","(","action_scores",",","z",",","support_logits_per_action",",","logits",",","dist",")"]
code: def get_q_value_distributions(self, model_out):
      """Returns distributional values for Q(s, a) given a state embedding.
  
      Override this in your custom model to customize the Q output head.
  
      Args:
          model_out: Embedding from the model layers.
  
      Returns:
          (action_scores, logits, dist) if num_atoms == 1, otherwise
          (action_scores, z, support_logits_per_action, logits, dist)
      """
      action_scores = self.advantage_module(model_out)
  
      if self.num_atoms > 1:
          # Distributional Q-learning uses a discrete support z
          # to represent the action value distribution
          z = torch.arange(0.0, self.num_atoms, dtype=torch.float32).to(
              action_scores.device
          )
          z = self.v_min + z * (self.v_max - self.v_min) / float(
              self.num_atoms - 1
          )
  
          support_logits_per_action = torch.reshape(
              action_scores,
              shape=(-1, self.action_space.n, self.num_atoms),
          )
          support_prob_per_action = nn.functional.softmax(
              support_logits_per_action, dim=-1
          )
          action_scores = torch.sum(z * support_prob_per_action, dim=-1)
          logits = support_logits_per_action
          probs = support_prob_per_action
          return (
              action_scores,
              z,
              support_logits_per_action,
              logits,
              probs,
          )
      else:
          logits = torch.unsqueeze(torch.ones_like(action_scores), -1)
          return action_scores, logits, logits
  
file_path: ray/python/ray/rllib/algorithms/dqn/dqn_torch_model.py
imports: from typing import Sequence
  import gymnasium
  from ray.rllib.models.torch.misc import SlimFC
  from ray.rllib.models.torch.modules.noisy_layer import NoisyLayer
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.typing import ModelConfigDict
label_desc: Use image node_id 2 for calling the DQNTorchModel obj's underlying member method code with example usage: obj.get_q_value_distributions(model_out) and returns: action_scores,  z,  support_logits_per_action,  logits,  probs, action_scores,  logits,  logits