class_docstr: A JAX implementation of the DQN agent.
class_docstr_tok: ["A","JAX","implementation","of","the","DQN","agent","."]
docstr: Initializes the agent and constructs the necessary components.
  
  Note: We are using the Adam optimizer by default for JaxDQN, which differs
        from the original NatureDQN and the dopamine TensorFlow version. In
        the experiments we have ran, we have found that using Adam yields
        improved training performance.
  
  Args:
    num_actions: int, number of actions the agent can take at any state.
    observation_shape: tuple of ints describing the observation shape.
    observation_dtype: jnp.dtype, specifies the type of the observations.
    stack_size: int, number of frames to use in state stack.
    network: Jax network to use for training.
    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.
    eval_mode: bool, True for evaluation and False for training.
    optimizer: str, name of optimizer to use.
    summary_writer: SummaryWriter object for outputting training statistics.
      May also be a str specifying the base directory, in which case the
      SummaryWriter will be created by the agent.
    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).
    seed: int, a seed for DQN's internal RNG, used for initialization and
      sampling actions. If None, will use the current time in nanoseconds.
    loss_type: str, whether to use Huber or MSE loss during training.
    preprocess_fn: function expecting the input state as parameter which
      it preprocesses (such as normalizing the pixel values between 0 and 1)
      before passing it to the Q-network. Defaults to None.
    collector_allowlist: list of str, if using CollectorDispatcher, this can
      be used to specify which Collectors to log to.
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Note",":","We","are","using","the","Adam","optimizer","by","default","for","JaxDQN",",","which","differs","from","the","original","NatureDQN","and","the","dopamine","TensorFlow","version",".","In","the","experiments","we","have","ran",",","we","have","found","that","using","Adam","yields","improved","training","performance",".","Args",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","observation_shape",":","tuple","of","ints","describing","the","observation","shape",".","observation_dtype",":","jnp.dtype",",","specifies","the","type","of","the","observations",".","stack_size",":","int",",","number","of","frames","to","use","in","state","stack",".","network",":","Jax","network","to","use","for","training",".","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",".","eval_mode",":","bool",",","True","for","evaluation","and","False","for","training",".","optimizer",":","str",",","name","of","optimizer","to","use",".","summary_writer",":","SummaryWriter","object","for","outputting","training","statistics",".","May","also","be","a","str","specifying","the","base","directory",",","in","which","case","the","SummaryWriter","will","be","created","by","the","agent",".","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",")",".","seed",":","int",",","a","seed","for","DQN","'s","internal","RNG",",","used","for","initialization","and","sampling","actions",".","If","None",",","will","use","the","current","time","in","nanoseconds",".","loss_type",":","str",",","whether","to","use","Huber","or","MSE","loss","during","training",".","preprocess_fn",":","function","expecting","the","input","state","as","parameter","which","it","preprocesses","(","such","as","normalizing","the","pixel","values","between","0","and","1",")","before","passing","it","to","the","Q-network",".","Defaults","to","None",".","collector_allowlist",":","list","of","str",",","if","using","CollectorDispatcher",",","this","can","be","used","to","specify","which","Collectors","to","log","to","."]
code: def __init__(
      self,
      num_actions,
      observation_shape=NATURE_DQN_OBSERVATION_SHAPE,
      observation_dtype=NATURE_DQN_DTYPE,
      stack_size=NATURE_DQN_STACK_SIZE,
      network=networks.NatureDQNNetwork,
      gamma=0.99,
      update_horizon=1,
      min_replay_history=20000,
      update_period=4,
      target_update_period=8000,
      epsilon_fn=linearly_decaying_epsilon,
      epsilon_train=0.01,
      epsilon_eval=0.001,
      epsilon_decay_period=250000,
      eval_mode=False,
      optimizer="adam",
      summary_writer=None,
      summary_writing_frequency=500,
      allow_partial_reload=False,
      seed=None,
      loss_type="mse",
      preprocess_fn=None,
      collector_allowlist=("tensorboard",),
  ):
      """Initializes the agent and constructs the necessary components.
  
      Note: We are using the Adam optimizer by default for JaxDQN, which differs
            from the original NatureDQN and the dopamine TensorFlow version. In
            the experiments we have ran, we have found that using Adam yields
            improved training performance.
  
      Args:
        num_actions: int, number of actions the agent can take at any state.
        observation_shape: tuple of ints describing the observation shape.
        observation_dtype: jnp.dtype, specifies the type of the observations.
        stack_size: int, number of frames to use in state stack.
        network: Jax network to use for training.
        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.
        eval_mode: bool, True for evaluation and False for training.
        optimizer: str, name of optimizer to use.
        summary_writer: SummaryWriter object for outputting training statistics.
          May also be a str specifying the base directory, in which case the
          SummaryWriter will be created by the agent.
        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).
        seed: int, a seed for DQN's internal RNG, used for initialization and
          sampling actions. If None, will use the current time in nanoseconds.
        loss_type: str, whether to use Huber or MSE loss during training.
        preprocess_fn: function expecting the input state as parameter which
          it preprocesses (such as normalizing the pixel values between 0 and 1)
          before passing it to the Q-network. Defaults to None.
        collector_allowlist: list of str, if using CollectorDispatcher, this can
          be used to specify which Collectors to log to.
      """
      assert isinstance(observation_shape, tuple)
      seed = int(time.time() * 1e6) if seed is None else seed
      logging.info(
          "Creating %s agent with the following parameters:",
          self.__class__.__name__,
      )
      logging.info("\t gamma: %f", gamma)
      logging.info("\t update_horizon: %f", update_horizon)
      logging.info("\t min_replay_history: %d", min_replay_history)
      logging.info("\t update_period: %d", update_period)
      logging.info("\t target_update_period: %d", target_update_period)
      logging.info("\t epsilon_train: %f", epsilon_train)
      logging.info("\t epsilon_eval: %f", epsilon_eval)
      logging.info("\t epsilon_decay_period: %d", epsilon_decay_period)
      logging.info("\t optimizer: %s", optimizer)
      logging.info("\t seed: %d", seed)
      logging.info("\t loss_type: %s", loss_type)
      logging.info("\t preprocess_fn: %s", preprocess_fn)
      logging.info(
          "\t summary_writing_frequency: %d", summary_writing_frequency
      )
      logging.info("\t allow_partial_reload: %s", allow_partial_reload)
  
      self.num_actions = num_actions
      self.observation_shape = tuple(observation_shape)
      self.observation_dtype = observation_dtype
      self.stack_size = stack_size
      if preprocess_fn is None:
          self.network_def = network(num_actions=num_actions)
          self.preprocess_fn = networks.identity_preprocess_fn
      else:
          self.network_def = network(
              num_actions=num_actions, inputs_preprocessed=True
          )
          self.preprocess_fn = preprocess_fn
      self.gamma = gamma
      self.update_horizon = update_horizon
      self.cumulative_gamma = math.pow(gamma, update_horizon)
      self.min_replay_history = min_replay_history
      self.target_update_period = target_update_period
      self.epsilon_fn = epsilon_fn
      self.epsilon_train = epsilon_train
      self.epsilon_eval = epsilon_eval
      self.epsilon_decay_period = epsilon_decay_period
      self.update_period = update_period
      self.eval_mode = eval_mode
      self.training_steps = 0
      if isinstance(summary_writer, str):
          try:
              tf.compat.v1.enable_v2_behavior()
          except ValueError:
              pass
          self.summary_writer = tf.summary.create_file_writer(
              summary_writer
          )
      else:
          self.summary_writer = summary_writer
      self.summary_writing_frequency = summary_writing_frequency
      self.allow_partial_reload = allow_partial_reload
      self._loss_type = loss_type
      self._collector_allowlist = collector_allowlist
  
      self._rng = jax.random.PRNGKey(seed)
      state_shape = self.observation_shape + (stack_size,)
      self.state = onp.zeros(state_shape)
      self._replay = self._build_replay_buffer()
      self._optimizer_name = optimizer
      self._build_networks_and_optimizer()
  
      # Variables to be initialized by the agent once it interacts with the
      # environment.
      self._observation = None
      self._last_observation = None
  
imports: import collections
  import functools
  import inspect
  import math
  import time
  from absl import logging
  from dopamine.agents.dqn import dqn_agent
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import circular_replay_buffer
  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 JaxDQNAgent object from inherited base classes: object with example: obj = JaxDQNAgent(num_actions, observation_shape, observation_dtype, stack_size, network, gamma, update_horizon, min_replay_history, update_period, target_update_period, epsilon_fn, epsilon_train, epsilon_eval, epsilon_decay_period, eval_mode, optimizer, summary_writer, summary_writing_frequency, allow_partial_reload, seed, loss_type, preprocess_fn, collector_allowlist)