class_docstr: Critic head that uses a categorical to represent action values.
class_docstr_tok: ["Critic","head","that","uses","a","categorical","to","represent","action","values","."]
code: def __call__(self, embedding: chex.Array) -> tfd.Distribution:
      output = self._head(embedding)
      return output
  
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 2 for calling the CategoricalCriticHead obj's underlying member method code with example usage: obj.__call__(embedding) and returns: output