class_docstr: A Categorical DQN Agent based on the DQN Agent.
class_docstr_tok: ["A","Categorical","DQN","Agent","based","on","the","DQN","Agent","."]
docstr: Creates a Categorical DQN Agent.
  
  Args:
    time_step_spec: A `TimeStep` spec of the expected time_steps.
    action_spec: A `BoundedTensorSpec` representing the actions.
    categorical_q_network: A categorical_q_network.CategoricalQNetwork that
      returns the q_distribution for each action.
    optimizer: The optimizer to use for training.
    observation_and_action_constraint_splitter: A function used to process
      observations with action constraints. These constraints can indicate,
      for example, a mask of valid/invalid actions for a given state of the
      environment. The function takes in a full observation and returns a
      tuple consisting of 1) the part of the observation intended as input to
      the network and 2) the constraint. An example
      `observation_and_action_constraint_splitter` could be as simple as: ```
      def observation_and_action_constraint_splitter(observation): return
      observation['network_input'], observation['constraint'] ``` *Note*: when
      using `observation_and_action_constraint_splitter`, make sure the
      provided `q_network` is compatible with the network-specific half of the
      output of the `observation_and_action_constraint_splitter`. In
      particular, `observation_and_action_constraint_splitter` will be called
      on the observation before passing to the network. If
      `observation_and_action_constraint_splitter` is None, action constraints
      are not applied.
    min_q_value: A float specifying the minimum Q-value, used for setting up
      the support.
    max_q_value: A float specifying the maximum Q-value, used for setting up
      the support.
    epsilon_greedy: probability of choosing a random action in the default
      epsilon-greedy collect policy (used only if a wrapper is not provided to
      the collect_policy method). Only one of epsilon_greedy and
      boltzmann_temperature should be provided.
    n_step_update: The number of steps to consider when computing TD error and
      TD loss. Defaults to single-step updates. Note that this requires the
      user to call train on Trajectory objects with a time dimension of
      `n_step_update + 1`. However, note that we do not yet support
      `n_step_update > 1` in the case of RNNs (i.e., non-empty
      `q_network.state_spec`).
    boltzmann_temperature: Temperature value to use for Boltzmann sampling of
      the actions during data collection. The closer to 0.0, the higher the
      probability of choosing the best action. Only one of epsilon_greedy and
      boltzmann_temperature should be provided.
    target_categorical_q_network: (Optional.)  A `tf_agents.network.Network`
      to be used as the target network during Q learning.  Every
      `target_update_period` train steps, the weights from
      `categorical_q_network` are copied (possibly with smoothing via
      `target_update_tau`) to `target_categorical_q_network`.  If
      `target_categorical_q_network` is not provided, it is created by making
      a copy of `categorical_q_network`, which initializes a new network with
      the same structure and its own layers and weights.  Network copying is
      performed via the `Network.copy` superclass method, and may
      inadvertently lead to the resulting network to share weights with the
      original.  This can happen if, for example, the original network
      accepted a pre-built Keras layer in its `__init__`, or accepted a Keras
      layer that wasn't built, but neglected to create a new copy.  In these
      cases, it is up to you to provide a target Network having weights that
      are not shared with the original `categorical_q_network`. If you provide
      a `target_categorical_q_network` that shares any weights with
      `categorical_q_network`, a warning will be logged but no exception is
      thrown.  Note; shallow copies of Keras layers may be built via the code:
      ```python new_layer = type(layer).from_config(layer.get_config()) ```
    target_update_tau: Factor for soft update of the target networks.
    target_update_period: Period for soft update of the target networks.
    td_errors_loss_fn: A function for computing the TD errors loss. If None, a
      default value of huber_loss is used. This function takes as input the
      target and the estimated Q values and returns the loss for each element
      of the batch.
    gamma: A discount factor for future rewards.
    reward_scale_factor: Multiplicative scale for the reward.
    gradient_clipping: Norm length to clip gradients.
    debug_summaries: A bool to gather debug summaries.
    summarize_grads_and_vars: If True, gradient and network variable summaries
      will be written during training.
    train_step_counter: An optional counter to increment every time the train
      op is run.  Defaults to the global_step.
    name: The name of this agent. All variables in this module will fall under
      that name. Defaults to the class name.
  
  Raises:
    TypeError: If the action spec contains more than one action.
    TypeError: If the q network(s) lack a `num_atoms` property.
docstr_tok: ["Creates","a","Categorical","DQN","Agent",".","Args",":","time_step_spec",":","A","`","TimeStep","`","spec","of","the","expected","time_steps",".","action_spec",":","A","`","BoundedTensorSpec","`","representing","the","actions",".","categorical_q_network",":","A","categorical_q_network.CategoricalQNetwork","that","returns","the","q_distribution","for","each","action",".","optimizer",":","The","optimizer","to","use","for","training",".","observation_and_action_constraint_splitter",":","A","function","used","to","process","observations","with","action","constraints",".","These","constraints","can","indicate",",","for","example",",","a","mask","of","valid\/invalid","actions","for","a","given","state","of","the","environment",".","The","function","takes","in","a","full","observation","and","returns","a","tuple","consisting","of","1",")","the","part","of","the","observation","intended","as","input","to","the","network","and","2",")","the","constraint",".","An","example","`","observation_and_action_constraint_splitter","`","could","be","as","simple","as",":","``","`","def","observation_and_action_constraint_splitter","(","observation",")",":","return","observation","[","'network_input","'","]",",","observation","[","'constraint","'","]","``","`","*","Note","*",":","when","using","`","observation_and_action_constraint_splitter","`",",","make","sure","the","provided","`","q_network","`","is","compatible","with","the","network-specific","half","of","the","output","of","the","`","observation_and_action_constraint_splitter","`",".","In","particular",",","`","observation_and_action_constraint_splitter","`","will","be","called","on","the","observation","before","passing","to","the","network",".","If","`","observation_and_action_constraint_splitter","`","is","None",",","action","constraints","are","not","applied",".","min_q_value",":","A","float","specifying","the","minimum","Q-value",",","used","for","setting","up","the","support",".","max_q_value",":","A","float","specifying","the","maximum","Q-value",",","used","for","setting","up","the","support",".","epsilon_greedy",":","probability","of","choosing","a","random","action","in","the","default","epsilon-greedy","collect","policy","(","used","only","if","a","wrapper","is","not","provided","to","the","collect_policy","method",")",".","Only","one","of","epsilon_greedy","and","boltzmann_temperature","should","be","provided",".","n_step_update",":","The","number","of","steps","to","consider","when","computing","TD","error","and","TD","loss",".","Defaults","to","single-step","updates",".","Note","that","this","requires","the","user","to","call","train","on","Trajectory","objects","with","a","time","dimension","of","`","n_step_update","+","1","`",".","However",",","note","that","we","do","not","yet","support","`","n_step_update",">","1","`","in","the","case","of","RNNs","(","i.e.",",","non-empty","`","q_network.state_spec","`",")",".","boltzmann_temperature",":","Temperature","value","to","use","for","Boltzmann","sampling","of","the","actions","during","data","collection",".","The","closer","to","0.0",",","the","higher","the","probability","of","choosing","the","best","action",".","Only","one","of","epsilon_greedy","and","boltzmann_temperature","should","be","provided",".","target_categorical_q_network",":","(","Optional",".",")","A","`","tf_agents.network.Network","`","to","be","used","as","the","target","network","during","Q","learning",".","Every","`","target_update_period","`","train","steps",",","the","weights","from","`","categorical_q_network","`","are","copied","(","possibly","with","smoothing","via","`","target_update_tau","`",")","to","`","target_categorical_q_network","`",".","If","`","target_categorical_q_network","`","is","not","provided",",","it","is","created","by","making","a","copy","of","`","categorical_q_network","`",",","which","initializes","a","new","network","with","the","same","structure","and","its","own","layers","and","weights",".","Network","copying","is","performed","via","the","`","Network.copy","`","superclass","method",",","and","may","inadvertently","lead","to","the","resulting","network","to","share","weights","with","the","original",".","This","can","happen","if",",","for","example",",","the","original","network","accepted","a","pre-built","Keras","layer","in","its","`","__init__","`",",","or","accepted","a","Keras","layer","that","was","n't","built",",","but","neglected","to","create","a","new","copy",".","In","these","cases",",","it","is","up","to","you","to","provide","a","target","Network","having","weights","that","are","not","shared","with","the","original","`","categorical_q_network","`",".","If","you","provide","a","`","target_categorical_q_network","`","that","shares","any","weights","with","`","categorical_q_network","`",",","a","warning","will","be","logged","but","no","exception","is","thrown",".","Note",";","shallow","copies","of","Keras","layers","may","be","built","via","the","code",":","``","`","python","new_layer","=","type","(","layer",")",".from_config","(","layer.get_config","(",")",")","``","`","target_update_tau",":","Factor","for","soft","update","of","the","target","networks",".","target_update_period",":","Period","for","soft","update","of","the","target","networks",".","td_errors_loss_fn",":","A","function","for","computing","the","TD","errors","loss",".","If","None",",","a","default","value","of","huber_loss","is","used",".","This","function","takes","as","input","the","target","and","the","estimated","Q","values","and","returns","the","loss","for","each","element","of","the","batch",".","gamma",":","A","discount","factor","for","future","rewards",".","reward_scale_factor",":","Multiplicative","scale","for","the","reward",".","gradient_clipping",":","Norm","length","to","clip","gradients",".","debug_summaries",":","A","bool","to","gather","debug","summaries",".","summarize_grads_and_vars",":","If","True",",","gradient","and","network","variable","summaries","will","be","written","during","training",".","train_step_counter",":","An","optional","counter","to","increment","every","time","the","train","op","is","run",".","Defaults","to","the","global_step",".","name",":","The","name","of","this","agent",".","All","variables","in","this","module","will","fall","under","that","name",".","Defaults","to","the","class","name",".","Raises",":","TypeError",":","If","the","action","spec","contains","more","than","one","action",".","TypeError",":","If","the","q","network","(","s",")","lack","a","`","num_atoms","`","property","."]
code: def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      categorical_q_network: network.Network,
      optimizer: types.Optimizer,
      observation_and_action_constraint_splitter: Optional[
          types.Splitter
      ] = None,
      min_q_value: types.Float = -10.0,
      max_q_value: types.Float = 10.0,
      epsilon_greedy: Optional[types.FloatOrReturningFloat] = 0.1,
      n_step_update: int = 1,
      boltzmann_temperature: Optional[
          types.FloatOrReturningFloat
      ] = None,
      # Params for target network updates
      target_categorical_q_network: Optional[network.Network] = None,
      target_update_tau: types.Float = 1.0,
      target_update_period: types.Int = 1,
      # Params for training.
      td_errors_loss_fn: Optional[types.LossFn] = None,
      gamma: types.Float = 1.0,
      reward_scale_factor: types.Float = 1.0,
      gradient_clipping: Optional[types.Float] = None,
      # Params for debugging
      debug_summaries: bool = False,
      summarize_grads_and_vars: bool = False,
      train_step_counter: Optional[tf.Variable] = None,
      name: Optional[Text] = None,
  ):
      """Creates a Categorical DQN Agent.
  
      Args:
        time_step_spec: A `TimeStep` spec of the expected time_steps.
        action_spec: A `BoundedTensorSpec` representing the actions.
        categorical_q_network: A categorical_q_network.CategoricalQNetwork that
          returns the q_distribution for each action.
        optimizer: The optimizer to use for training.
        observation_and_action_constraint_splitter: A function used to process
          observations with action constraints. These constraints can indicate,
          for example, a mask of valid/invalid actions for a given state of the
          environment. The function takes in a full observation and returns a
          tuple consisting of 1) the part of the observation intended as input to
          the network and 2) the constraint. An example
          `observation_and_action_constraint_splitter` could be as simple as: ```
          def observation_and_action_constraint_splitter(observation): return
          observation['network_input'], observation['constraint'] ``` *Note*: when
          using `observation_and_action_constraint_splitter`, make sure the
          provided `q_network` is compatible with the network-specific half of the
          output of the `observation_and_action_constraint_splitter`. In
          particular, `observation_and_action_constraint_splitter` will be called
          on the observation before passing to the network. If
          `observation_and_action_constraint_splitter` is None, action constraints
          are not applied.
        min_q_value: A float specifying the minimum Q-value, used for setting up
          the support.
        max_q_value: A float specifying the maximum Q-value, used for setting up
          the support.
        epsilon_greedy: probability of choosing a random action in the default
          epsilon-greedy collect policy (used only if a wrapper is not provided to
          the collect_policy method). Only one of epsilon_greedy and
          boltzmann_temperature should be provided.
        n_step_update: The number of steps to consider when computing TD error and
          TD loss. Defaults to single-step updates. Note that this requires the
          user to call train on Trajectory objects with a time dimension of
          `n_step_update + 1`. However, note that we do not yet support
          `n_step_update > 1` in the case of RNNs (i.e., non-empty
          `q_network.state_spec`).
        boltzmann_temperature: Temperature value to use for Boltzmann sampling of
          the actions during data collection. The closer to 0.0, the higher the
          probability of choosing the best action. Only one of epsilon_greedy and
          boltzmann_temperature should be provided.
        target_categorical_q_network: (Optional.)  A `tf_agents.network.Network`
          to be used as the target network during Q learning.  Every
          `target_update_period` train steps, the weights from
          `categorical_q_network` are copied (possibly with smoothing via
          `target_update_tau`) to `target_categorical_q_network`.  If
          `target_categorical_q_network` is not provided, it is created by making
          a copy of `categorical_q_network`, which initializes a new network with
          the same structure and its own layers and weights.  Network copying is
          performed via the `Network.copy` superclass method, and may
          inadvertently lead to the resulting network to share weights with the
          original.  This can happen if, for example, the original network
          accepted a pre-built Keras layer in its `__init__`, or accepted a Keras
          layer that wasn't built, but neglected to create a new copy.  In these
          cases, it is up to you to provide a target Network having weights that
          are not shared with the original `categorical_q_network`. If you provide
          a `target_categorical_q_network` that shares any weights with
          `categorical_q_network`, a warning will be logged but no exception is
          thrown.  Note; shallow copies of Keras layers may be built via the code:
          ```python new_layer = type(layer).from_config(layer.get_config()) ```
        target_update_tau: Factor for soft update of the target networks.
        target_update_period: Period for soft update of the target networks.
        td_errors_loss_fn: A function for computing the TD errors loss. If None, a
          default value of huber_loss is used. This function takes as input the
          target and the estimated Q values and returns the loss for each element
          of the batch.
        gamma: A discount factor for future rewards.
        reward_scale_factor: Multiplicative scale for the reward.
        gradient_clipping: Norm length to clip gradients.
        debug_summaries: A bool to gather debug summaries.
        summarize_grads_and_vars: If True, gradient and network variable summaries
          will be written during training.
        train_step_counter: An optional counter to increment every time the train
          op is run.  Defaults to the global_step.
        name: The name of this agent. All variables in this module will fall under
          that name. Defaults to the class name.
  
      Raises:
        TypeError: If the action spec contains more than one action.
        TypeError: If the q network(s) lack a `num_atoms` property.
      """
  
      def check_atoms(net, label):
          try:
              num_atoms = net.num_atoms
          except AttributeError:
              raise TypeError(
                  "Expected {} to have property `num_atoms`, but it "
                  "doesn't. (Note: you likely want to use a "
                  "CategoricalQNetwork.) Network is: {}".format(
                      label, net
                  )
              )
          return num_atoms
  
      self._num_atoms = check_atoms(
          categorical_q_network, "categorical_q_network"
      )
  
      if target_categorical_q_network is not None:
          target_num_atoms = check_atoms(
              target_categorical_q_network,
              "target_categorical_q_network",
          )
          if self._num_atoms != target_num_atoms:
              raise ValueError(
                  "categorical_q_network and target_categorical_q_network have "
                  "different numbers of atoms: {} vs. {}".format(
                      self._num_atoms, target_num_atoms
                  )
              )
  
      self._min_q_value = min_q_value
      self._max_q_value = max_q_value
      min_q_value = tf.convert_to_tensor(
          min_q_value, dtype_hint=tf.float32
      )
      max_q_value = tf.convert_to_tensor(
          max_q_value, dtype_hint=tf.float32
      )
      self._support = tf.linspace(
          min_q_value, max_q_value, self._num_atoms
      )
  
      super(CategoricalDqnAgent, self).__init__(
          time_step_spec,
          action_spec,
          categorical_q_network,
          optimizer,
          observation_and_action_constraint_splitter=(
              observation_and_action_constraint_splitter
          ),
          epsilon_greedy=epsilon_greedy,
          n_step_update=n_step_update,
          boltzmann_temperature=boltzmann_temperature,
          target_q_network=target_categorical_q_network,
          target_update_tau=target_update_tau,
          target_update_period=target_update_period,
          td_errors_loss_fn=td_errors_loss_fn,
          gamma=gamma,
          reward_scale_factor=reward_scale_factor,
          gradient_clipping=gradient_clipping,
          debug_summaries=debug_summaries,
          summarize_grads_and_vars=summarize_grads_and_vars,
          train_step_counter=train_step_counter,
          name=name,
      )
  
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 1 to create a new CategoricalDqnAgent object from inherited base classes: dqn_agent with example: obj = CategoricalDqnAgent(time_step_spec, action_spec, categorical_q_network, optimizer, observation_and_action_constraint_splitter, min_q_value, max_q_value, epsilon_greedy, n_step_update, boltzmann_temperature, target_categorical_q_network, target_update_tau, target_update_period, td_errors_loss_fn, gamma, reward_scale_factor, gradient_clipping, debug_summaries, summarize_grads_and_vars, train_step_counter, name)