class_docstr: C51-style head.
  
  For each action, it produces the logits for a discrete distribution over
  atoms. Therefore, the returned logits represents several distributions, one
  for each action.
class_docstr_tok: ["C51-style","head",".","For","each","action",",","it","produces","the","logits","for","a","discrete","distribution","over","atoms",".","Therefore",",","the","returned","logits","represents","several","distributions",",","one","for","each","action","."]
code: def __init__(
      self,
      num_actions: int,
      head_units: int = 512,
      num_atoms: int = 51,
      v_min: float = -1.0,
      v_max: float = 1.0,
  ):
      super().__init__("DiscreteValued")
      self._num_actions = num_actions
      self._num_atoms = num_atoms
      self._atoms = jnp.linspace(v_min, v_max, self._num_atoms)
      self._network = hk.nets.MLP([head_units, num_actions * num_atoms])
  
imports: from typing import Any, List, Optional, Sequence, Union, Callable
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 1 to create a new DiscreteValued object from inherited base classes: hk with example: obj = DiscreteValued(num_actions, head_units, num_atoms, v_min, v_max)