class_docstr: A compact implementation of a simplified Rainbow agent.
class_docstr_tok: ["A","compact","implementation","of","a","simplified","Rainbow","agent","."]
docstr: Initializes the agent and constructs the components of its graph.
  
  Args:
    sess: `tf.compat.v1.Session`, for executing ops.
    num_actions: int, number of actions the agent can take at any state.
    observation_shape: tuple of ints or an int. If single int, the observation
      is assumed to be a 2D square.
    observation_dtype: tf.DType, specifies the type of the observations. Note
      that if your inputs are continuous, you should set this to tf.float32.
    stack_size: int, number of frames to use in state stack.
    network: tf.Keras.Model, expects four parameters:
      (num_actions, num_atoms, support, network_type).  This class is used to
      generate network instances that are used by the agent. Each
      instantiation would have different set of variables. See
      dopamine.discrete_domains.atari_lib.RainbowNetwork as an example.
    num_atoms: int, the number of buckets of the value function distribution.
    vmin: float, the value distribution support is [vmin, vmax]. If None, we
      set it to be -vmax.
    vmax: float, the value distribution support is [vmin, vmax].
    gamma: float, discount factor with the usual RL meaning.
    update_horizon: int, horizon at which updates are performed, the 'n' in
      n-step update.
    min_replay_history: int, number of transitions that should be experienced
      before the agent begins training its value function.
    update_period: int, period between DQN updates.
    target_update_period: int, update period for the target network.
    epsilon_fn: function expecting 4 parameters:
      (decay_period, step, warmup_steps, epsilon). This function should return
      the epsilon value used for exploration during training.
    epsilon_train: float, the value to which the agent's epsilon is eventually
      decayed during training.
    epsilon_eval: float, epsilon used when evaluating the agent.
    epsilon_decay_period: int, length of the epsilon decay schedule.
    replay_scheme: str, 'prioritized' or 'uniform', the sampling scheme of the
      replay memory.
    tf_device: str, Tensorflow device on which the agent's graph is executed.
    use_staging: bool, when True use a staging area to prefetch the next
      training batch, speeding training up by about 30%.
    optimizer: `tf.compat.v1.train.Optimizer`, for training the value
      function.
    summary_writer: SummaryWriter object for outputting training statistics.
      Summary writing disabled if set to None.
    summary_writing_frequency: int, frequency with which summaries will be
      written. Lower values will result in slower training.
docstr_tok: ["Initializes","the","agent","and","constructs","the","components","of","its","graph",".","Args",":","sess",":","`","tf.compat.v1.Session","`",",","for","executing","ops",".","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","observation_shape",":","tuple","of","ints","or","an","int",".","If","single","int",",","the","observation","is","assumed","to","be","a","2D","square",".","observation_dtype",":","tf.DType",",","specifies","the","type","of","the","observations",".","Note","that","if","your","inputs","are","continuous",",","you","should","set","this","to","tf.float32",".","stack_size",":","int",",","number","of","frames","to","use","in","state","stack",".","network",":","tf.Keras.Model",",","expects","four","parameters",":","(","num_actions",",","num_atoms",",","support",",","network_type",")",".","This","class","is","used","to","generate","network","instances","that","are","used","by","the","agent",".","Each","instantiation","would","have","different","set","of","variables",".","See","dopamine.discrete_domains.atari_lib.RainbowNetwork","as","an","example",".","num_atoms",":","int",",","the","number","of","buckets","of","the","value","function","distribution",".","vmin",":","float",",","the","value","distribution","support","is","[","vmin",",","vmax","]",".","If","None",",","we","set","it","to","be","-vmax",".","vmax",":","float",",","the","value","distribution","support","is","[","vmin",",","vmax","]",".","gamma",":","float",",","discount","factor","with","the","usual","RL","meaning",".","update_horizon",":","int",",","horizon","at","which","updates","are","performed",",","the","'n","'","in","n-step","update",".","min_replay_history",":","int",",","number","of","transitions","that","should","be","experienced","before","the","agent","begins","training","its","value","function",".","update_period",":","int",",","period","between","DQN","updates",".","target_update_period",":","int",",","update","period","for","the","target","network",".","epsilon_fn",":","function","expecting","4","parameters",":","(","decay_period",",","step",",","warmup_steps",",","epsilon",")",".","This","function","should","return","the","epsilon","value","used","for","exploration","during","training",".","epsilon_train",":","float",",","the","value","to","which","the","agent","'s","epsilon","is","eventually","decayed","during","training",".","epsilon_eval",":","float",",","epsilon","used","when","evaluating","the","agent",".","epsilon_decay_period",":","int",",","length","of","the","epsilon","decay","schedule",".","replay_scheme",":","str",",","'prioritized","'","or","'uniform","'",",","the","sampling","scheme","of","the","replay","memory",".","tf_device",":","str",",","Tensorflow","device","on","which","the","agent","'s","graph","is","executed",".","use_staging",":","bool",",","when","True","use","a","staging","area","to","prefetch","the","next","training","batch",",","speeding","training","up","by","about","30","%",".","optimizer",":","`","tf.compat.v1.train.Optimizer","`",",","for","training","the","value","function",".","summary_writer",":","SummaryWriter","object","for","outputting","training","statistics",".","Summary","writing","disabled","if","set","to","None",".","summary_writing_frequency",":","int",",","frequency","with","which","summaries","will","be","written",".","Lower","values","will","result","in","slower","training","."]
code: def __init__(
      self,
      sess,
      num_actions,
      observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
      observation_dtype=dqn_agent.NATURE_DQN_DTYPE,
      stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
      network=atari_lib.RainbowNetwork,
      num_atoms=51,
      vmin=None,
      vmax=10.0,
      gamma=0.99,
      update_horizon=1,
      min_replay_history=20000,
      update_period=4,
      target_update_period=8000,
      epsilon_fn=dqn_agent.linearly_decaying_epsilon,
      epsilon_train=0.01,
      epsilon_eval=0.001,
      epsilon_decay_period=250000,
      replay_scheme="prioritized",
      tf_device="/cpu:*",
      use_staging=False,
      optimizer=tf.compat.v1.train.AdamOptimizer(
          learning_rate=0.00025, epsilon=0.0003125
      ),
      summary_writer=None,
      summary_writing_frequency=500,
  ):
      """Initializes the agent and constructs the components of its graph.
  
      Args:
        sess: `tf.compat.v1.Session`, for executing ops.
        num_actions: int, number of actions the agent can take at any state.
        observation_shape: tuple of ints or an int. If single int, the observation
          is assumed to be a 2D square.
        observation_dtype: tf.DType, specifies the type of the observations. Note
          that if your inputs are continuous, you should set this to tf.float32.
        stack_size: int, number of frames to use in state stack.
        network: tf.Keras.Model, expects four parameters:
          (num_actions, num_atoms, support, network_type).  This class is used to
          generate network instances that are used by the agent. Each
          instantiation would have different set of variables. See
          dopamine.discrete_domains.atari_lib.RainbowNetwork as an example.
        num_atoms: int, the number of buckets of the value function distribution.
        vmin: float, the value distribution support is [vmin, vmax]. If None, we
          set it to be -vmax.
        vmax: float, the value distribution support is [vmin, vmax].
        gamma: float, discount factor with the usual RL meaning.
        update_horizon: int, horizon at which updates are performed, the 'n' in
          n-step update.
        min_replay_history: int, number of transitions that should be experienced
          before the agent begins training its value function.
        update_period: int, period between DQN updates.
        target_update_period: int, update period for the target network.
        epsilon_fn: function expecting 4 parameters:
          (decay_period, step, warmup_steps, epsilon). This function should return
          the epsilon value used for exploration during training.
        epsilon_train: float, the value to which the agent's epsilon is eventually
          decayed during training.
        epsilon_eval: float, epsilon used when evaluating the agent.
        epsilon_decay_period: int, length of the epsilon decay schedule.
        replay_scheme: str, 'prioritized' or 'uniform', the sampling scheme of the
          replay memory.
        tf_device: str, Tensorflow device on which the agent's graph is executed.
        use_staging: bool, when True use a staging area to prefetch the next
          training batch, speeding training up by about 30%.
        optimizer: `tf.compat.v1.train.Optimizer`, for training the value
          function.
        summary_writer: SummaryWriter object for outputting training statistics.
          Summary writing disabled if set to None.
        summary_writing_frequency: int, frequency with which summaries will be
          written. Lower values will result in slower training.
      """
      # We need this because some tools convert round floats into ints.
      vmax = float(vmax)
      self._num_atoms = num_atoms
      # If vmin is not specified, set it to -vmax similar to C51.
      vmin = vmin if vmin else -vmax
      self._support = tf.linspace(vmin, vmax, num_atoms)
      self._replay_scheme = replay_scheme
      # TODO(b/110897128): Make agent optimizer attribute private.
      self.optimizer = optimizer
  
      dqn_agent.DQNAgent.__init__(
          self,
          sess=sess,
          num_actions=num_actions,
          observation_shape=observation_shape,
          observation_dtype=observation_dtype,
          stack_size=stack_size,
          network=network,
          gamma=gamma,
          update_horizon=update_horizon,
          min_replay_history=min_replay_history,
          update_period=update_period,
          target_update_period=target_update_period,
          epsilon_fn=epsilon_fn,
          epsilon_train=epsilon_train,
          epsilon_eval=epsilon_eval,
          epsilon_decay_period=epsilon_decay_period,
          tf_device=tf_device,
          use_staging=use_staging,
          optimizer=self.optimizer,
          summary_writer=summary_writer,
          summary_writing_frequency=summary_writing_frequency,
      )
  
file_path: dopamine/dopamine/agents/rainbow/rainbow_agent.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from dopamine.agents.dqn import dqn_agent
  from dopamine.discrete_domains import atari_lib
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin.tf
  import tensorflow
label_desc: Use image node_id 1 to create a new RainbowAgent object from inherited base classes: dqn_agent with example: obj = RainbowAgent(sess, num_actions, observation_shape, observation_dtype, stack_size, network, num_atoms, vmin, vmax, gamma, update_horizon, min_replay_history, update_period, target_update_period, epsilon_fn, epsilon_train, epsilon_eval, epsilon_decay_period, replay_scheme, tf_device, use_staging, optimizer, summary_writer, summary_writing_frequency)