class_docstr: Extension of standard TFModel to provide distributional Q values.
  
  It also supports options for noisy nets and parameter space noise.
  
  Data flow:
      obs -> forward() -> model_out
      model_out -> get_q_value_distributions() -> Q(s, a) atoms
      model_out -> get_state_value() -> V(s)
  
  Note that this class by itself is not a valid model unless you
  implement forward() in a subclass.
class_docstr_tok: ["Extension","of","standard","TFModel","to","provide","distributional","Q","values",".","It","also","supports","options","for","noisy","nets","and","parameter","space","noise",".","Data","flow",":","obs","-",">","forward","(",")","-",">","model_out","model_out","-",">","get_q_value_distributions","(",")","-",">","Q","(","s",",","a",")","atoms","model_out","-",">","get_state_value","(",")","-",">","V","(","s",")","Note","that","this","class","by","itself","is","not","a","valid","model","unless","you","implement","forward","(",")","in","a","subclass","."]
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: TensorType
  ) -> List[TensorType]:
      """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)
      """
      return self.q_value_head(model_out)
  
file_path: ray/rllib/algorithms/dqn/distributional_q_tf_model.py
imports: from typing import List
  import gymnasium
  from ray.rllib.models.tf.layers import NoisyLayer
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
label_desc: Use image node_id 2 for calling the DistributionalQTFModel obj's underlying member method code with example usage: obj.get_q_value_distributions(model_out) and returns: self