class_docstr: A Categorical DQN Agent based on the DQN Agent.
class_docstr_tok: ["A","Categorical","DQN","Agent","based","on","the","DQN","Agent","."]
docstr: Compute the q distribution of the next state for TD error computation.
  
  Args:
    next_time_steps: A batch of next timesteps
  
  Returns:
    A [batch_size, num_atoms] tensor representing the Q-distribution for the
    next state.
docstr_tok: ["Compute","the","q","distribution","of","the","next","state","for","TD","error","computation",".","Args",":","next_time_steps",":","A","batch","of","next","timesteps","Returns",":","A","[","batch_size",",","num_atoms","]","tensor","representing","the","Q-distribution","for","the","next","state","."]
code: def _next_q_distribution(self, next_time_steps):
      """Compute the q distribution of the next state for TD error computation.
  
      Args:
        next_time_steps: A batch of next timesteps
  
      Returns:
        A [batch_size, num_atoms] tensor representing the Q-distribution for the
        next state.
      """
      network_observation = next_time_steps.observation
  
      if self._observation_and_action_constraint_splitter is not None:
          (
              network_observation,
              _,
          ) = self._observation_and_action_constraint_splitter(
              network_observation
          )
  
      next_target_logits, _ = self._target_q_network(
          network_observation,
          step_type=next_time_steps.step_type,
          training=False,
      )
      batch_size = (
          next_target_logits.shape[0] or tf.shape(next_target_logits)[0]
      )
      next_target_probabilities = tf.nn.softmax(next_target_logits)
      next_target_q_values = tf.reduce_sum(
          self._support * next_target_probabilities, axis=-1
      )
      dummy_state = self._target_greedy_policy.get_initial_state(
          batch_size
      )
      # Find the greedy actions using our target greedy policy. This ensures that
      # action constraints are respected and helps centralize the greedy logic.
      greedy_actions = self._target_greedy_policy.action(
          next_time_steps, dummy_state
      ).action
      next_qt_argmax = tf.cast(greedy_actions, tf.int32)[:, None]
      batch_indices = tf.range(
          tf.cast(tf.shape(next_target_q_values)[0], tf.int32)
      )[:, None]
      next_qt_argmax = tf.concat(
          [batch_indices, next_qt_argmax], axis=-1
      )
      return tf.gather_nd(next_target_probabilities, next_qt_argmax)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text
  import gin
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.agents.dqn import dqn_agent
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.policies import boltzmann_policy
  from tf_agents.policies import categorical_q_policy
  from tf_agents.policies import epsilon_greedy_policy
  from tf_agents.policies import greedy_policy
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tf_agents.utils import value_ops
label_desc: Use image node_id 5 for calling the CategoricalDqnAgent obj's underlying member method code with example usage: obj._next_q_distribution(next_time_steps) and returns: tf