class_docstr: Class to build categorical Q-policies.
class_docstr_tok: ["Class","to","build","categorical","Q-policies","."]
docstr: Generates the distribution over next actions given the time_step.
  
  Args:
    time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
    policy_state: A Tensor, or a nested dict, list or tuple of Tensors
      representing the previous policy_state.
  
  Returns:
    A tfp.distributions.Categorical capturing the distribution of next
      actions.
    A policy_state Tensor, or a nested dict, list or tuple of Tensors,
      representing the new policy state.
docstr_tok: ["Generates","the","distribution","over","next","actions","given","the","time_step",".","Args",":","time_step",":","A","`","TimeStep","`","tuple","corresponding","to","`","time_step_spec","(",")","`",".","policy_state",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","representing","the","previous","policy_state",".","Returns",":","A","tfp.distributions.Categorical","capturing","the","distribution","of","next","actions",".","A","policy_state","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors",",","representing","the","new","policy","state","."]
code: def _distribution(self, time_step, policy_state):
      """Generates the distribution over next actions given the time_step.
  
      Args:
        time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
        policy_state: A Tensor, or a nested dict, list or tuple of Tensors
          representing the previous policy_state.
  
      Returns:
        A tfp.distributions.Categorical capturing the distribution of next
          actions.
        A policy_state Tensor, or a nested dict, list or tuple of Tensors,
          representing the new policy state.
      """
      network_observation = time_step.observation
      observation_and_action_constraint_splitter = (
          self.observation_and_action_constraint_splitter
      )
  
      if observation_and_action_constraint_splitter is not None:
          (
              network_observation,
              mask,
          ) = observation_and_action_constraint_splitter(
              network_observation
          )
  
      q_logits, policy_state = self._q_network(
          network_observation,
          step_type=time_step.step_type,
          network_state=policy_state,
      )
      q_logits.shape.assert_has_rank(3)
      q_values = common.convert_q_logits_to_values(
          q_logits, self._support
      )
  
      logits = q_values
  
      if observation_and_action_constraint_splitter is not None:
          # Overwrite the logits for invalid actions to -inf.
          neg_inf = tf.constant(-np.inf, dtype=logits.dtype)
          logits = tf.compat.v2.where(
              tf.cast(mask, tf.bool), logits, neg_inf
          )
  
      action_spec = cast(tf.TensorSpec, self.action_spec)
      dist = tfp.distributions.Categorical(
          logits=logits, dtype=action_spec.dtype
      )
  
      return policy_step.PolicyStep(dist, policy_state)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, cast
  import gin
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 3 for calling the CategoricalQPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step