class_docstr: distributional Q-function output for discrete action space.
  
  Args:
      q_dist: Probabilities of atoms. Its shape must be
          (batchsize, n_actions, n_atoms).
      z_values (ndarray): Values represented by atoms.
          Its shape must be (n_atoms,).
class_docstr_tok: ["distributional","Q-function","output","for","discrete","action","space",".","Args",":","q_dist",":","Probabilities","of","atoms",".","Its","shape","must","be","(","batchsize",",","n_actions",",","n_atoms",")",".","z_values","(","ndarray",")",":","Values","represented","by","atoms",".","Its","shape","must","be","(","n_atoms",",",")","."]
code: def __init__(self, q_dist, z_values, q_values_formatter=lambda x: x):
      assert isinstance(q_dist, torch.Tensor)
      assert isinstance(z_values, torch.Tensor)
      assert q_dist.ndim == 3
      assert z_values.ndim == 1
      assert q_dist.shape[2] == z_values.shape[0]
      self.z_values = z_values
      q_scaled = q_dist * self.z_values[None, None, ...]
      self.q_values = q_scaled.sum(dim=2)
      self.q_dist = q_dist
      self.n_actions = q_dist.shape[1]
      self.q_values_formatter = q_values_formatter
  
imports: import warnings
  from abc import ABCMeta, abstractmethod, abstractproperty
  import torch
  import torch.nn.functional
  from torch.distributions.utils import lazy_property
label_desc: Use image node_id 1 to create a new DistributionalDiscreteActionValue object from inherited base classes: ActionValue with example: obj = DistributionalDiscreteActionValue(q_dist, z_values, q_values_formatter)