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 necessary components.
  
  Args:
    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: DType, specifies the type of the observations. Note
      that if your inputs are continuous, you should set this to jnp.float32.
    stack_size: int, number of frames to use in state stack.
    network: flax.linen Module that is initialized by shape in _create_network
      below. See dopamine.jax.networks.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.
    optimizer: str, name of optimizer to use.
    seed: int, a seed for Rainbow's internal RNG, used for initialization and
      sampling actions. If None, will use the current time in nanoseconds.
    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.
    allow_partial_reload: bool, whether we allow reloading a partial agent
      (for instance, only the network parameters).
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Args",":","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",":","DType",",","specifies","the","type","of","the","observations",".","Note","that","if","your","inputs","are","continuous",",","you","should","set","this","to","jnp.float32",".","stack_size",":","int",",","number","of","frames","to","use","in","state","stack",".","network",":","flax.linen","Module","that","is","initialized","by","shape","in","_create_network","below",".","See","dopamine.jax.networks.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",".","optimizer",":","str",",","name","of","optimizer","to","use",".","seed",":","int",",","a","seed","for","Rainbow","'s","internal","RNG",",","used","for","initialization","and","sampling","actions",".","If","None",",","will","use","the","current","time","in","nanoseconds",".","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",".","allow_partial_reload",":","bool",",","whether","we","allow","reloading","a","partial","agent","(","for","instance",",","only","the","network","parameters",")","."]
code: def __init__(
      self,
      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=networks.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",
      optimizer="adam",
      seed=None,
      summary_writer=None,
      summary_writing_frequency=500,
      allow_partial_reload=False,
  ):
      """Initializes the agent and constructs the necessary components.
  
      Args:
        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: DType, specifies the type of the observations. Note
          that if your inputs are continuous, you should set this to jnp.float32.
        stack_size: int, number of frames to use in state stack.
        network: flax.linen Module that is initialized by shape in _create_network
          below. See dopamine.jax.networks.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.
        optimizer: str, name of optimizer to use.
        seed: int, a seed for Rainbow's internal RNG, used for initialization and
          sampling actions. If None, will use the current time in nanoseconds.
        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.
        allow_partial_reload: bool, whether we allow reloading a partial agent
          (for instance, only the network parameters).
      """
      # 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 = jnp.linspace(vmin, vmax, num_atoms)
      self._replay_scheme = replay_scheme
  
      super(JaxRainbowAgent, self).__init__(
          num_actions=num_actions,
          observation_shape=observation_shape,
          observation_dtype=observation_dtype,
          stack_size=stack_size,
          network=functools.partial(network, num_atoms=num_atoms),
          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,
          optimizer=optimizer,
          seed=seed,
          summary_writer=summary_writer,
          summary_writing_frequency=summary_writing_frequency,
          allow_partial_reload=allow_partial_reload,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 1 to create a new JaxRainbowAgent object from inherited base classes: dqn_agent with example: obj = JaxRainbowAgent(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, optimizer, seed, summary_writer, summary_writing_frequency, allow_partial_reload)