class_docstr: Class to build Q-Policies.
class_docstr_tok: ["Class","to","build","Q-Policies","."]
code: def _distribution(self, time_step, policy_state):
      # In DQN, we always either take a uniformly random action, or the action
      # with the highest Q-value. However, to support more complicated policies,
      # we expose all Q-values as a categorical distribution with Q-values as
      # logits, and apply the GreedyPolicy wrapper in dqn_agent.py to select the
      # action with the highest Q-value.
      observation_and_action_constraint_splitter = (
          self.observation_and_action_constraint_splitter
      )
      network_observation = time_step.observation
  
      if observation_and_action_constraint_splitter is not None:
          (
              network_observation,
              mask,
          ) = observation_and_action_constraint_splitter(
              network_observation
          )
  
      q_values, policy_state = self._q_network(
          network_observation,
          network_state=policy_state,
          step_type=time_step.step_type,
      )
  
      # use action values
      logits = q_values[1] if isinstance(q_values, tuple) else q_values
  
      if observation_and_action_constraint_splitter is not None:
          # Overwrite the logits for invalid actions to logits.dtype.min.
          almost_neg_inf = tf.constant(
              logits.dtype.min, dtype=logits.dtype
          )
          logits = tf.compat.v2.where(
              tf.cast(mask, tf.bool), logits, almost_neg_inf
          )
  
      if self._flat_action_spec.minimum != 0:
          distribution = shifted_categorical.ShiftedCategorical(
              logits=logits,
              dtype=self._flat_action_spec.dtype,
              shift=self._flat_action_spec.minimum,
          )
      else:
          distribution = tfp.distributions.Categorical(
              logits=logits, dtype=self._flat_action_spec.dtype
          )
  
      distribution = tf.nest.pack_sequence_as(
          self._action_spec, [distribution]
      )
      return policy_step.PolicyStep(distribution, policy_state)
  
file_path: agents/tf_agents/policies/q_policy.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text, cast
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import shifted_categorical
  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
label_desc: Use image node_id 3 for calling the QPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step