class_docstr: Policy that chooses an action based on a categorical distribution.
  
  The distribution is specified by a set of weights for each action and an
  inverse temperature. The unnormalized probability distribution is given by
  `exp(weight * inv_temp)`. Weights and inverse temperature are typically
  maintained as `Variable`s, and are updated by an `Agent`.
  
  Note that this policy does not make use of `time_step.observation` at all.
  That is, it is a non-contextual bandit policy.
class_docstr_tok: ["Policy","that","chooses","an","action","based","on","a","categorical","distribution",".","The","distribution","is","specified","by","a","set","of","weights","for","each","action","and","an","inverse","temperature",".","The","unnormalized","probability","distribution","is","given","by","`","exp","(","weight","*","inv_temp",")","`",".","Weights","and","inverse","temperature","are","typically","maintained","as","`","Variable","`","s",",","and","are","updated","by","an","`","Agent","`",".","Note","that","this","policy","does","not","make","use","of","`","time_step.observation","`","at","all",".","That","is",",","it","is","a","non-contextual","bandit","policy","."]
docstr: Implementation of `distribution`. Returns a `Categorical` distribution.
  
  The returned `Categorical` distribution has (unnormalized) probabilities
  `exp(inverse_temperature * weights)`.
  
  Args:
    time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
    policy_state: Unused in `CategoricalPolicy`. It is simply passed through.
  
  Returns:
    A `PolicyStep` named tuple containing:
      `action`: A (optionally nested) of tfp.distribution.Distribution
        capturing the distribution of next actions.
      `state`: A policy state tensor for the next call to distribution.
      `info`: Optional side information such as action log probabilities.
docstr_tok: ["Implementation","of","`","distribution","`",".","Returns","a","`","Categorical","`","distribution",".","The","returned","`","Categorical","`","distribution","has","(","unnormalized",")","probabilities","`","exp","(","inverse_temperature","*","weights",")","`",".","Args",":","time_step",":","A","`","TimeStep","`","tuple","corresponding","to","`","time_step_spec","(",")","`",".","policy_state",":","Unused","in","`","CategoricalPolicy","`",".","It","is","simply","passed","through",".","Returns",":","A","`","PolicyStep","`","named","tuple","containing",":","`","action","`",":","A","(","optionally","nested",")","of","tfp.distribution.Distribution","capturing","the","distribution","of","next","actions",".","`","state","`",":","A","policy","state","tensor","for","the","next","call","to","distribution",".","`","info","`",":","Optional","side","information","such","as","action","log","probabilities","."]
code: def _distribution(self, time_step, policy_state):
      """Implementation of `distribution`. Returns a `Categorical` distribution.
  
      The returned `Categorical` distribution has (unnormalized) probabilities
      `exp(inverse_temperature * weights)`.
  
      Args:
        time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
        policy_state: Unused in `CategoricalPolicy`. It is simply passed through.
  
      Returns:
        A `PolicyStep` named tuple containing:
          `action`: A (optionally nested) of tfp.distribution.Distribution
            capturing the distribution of next actions.
          `state`: A policy state tensor for the next call to distribution.
          `info`: Optional side information such as action log probabilities.
      """
      outer_shape = nest_utils.get_outer_shape(
          time_step, self._time_step_spec
      )
      logits = self._inverse_temperature * common.replicate(
          self._weights, outer_shape
      )
      action_distribution = tfd.Independent(
          tfd.Categorical(
              logits=logits,
              dtype=tf.nest.flatten(self.action_spec)[0].dtype,
          )
      )
      return policy_step.PolicyStep(action_distribution, policy_state)
  
file_path: agents/tf_agents/bandits/policies/categorical_policy.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies import tf_policy
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 3 for calling the CategoricalPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step