class_docstr: An implementation of Quantile regression DQN agent.
class_docstr_tok: ["An","implementation","of","Quantile","regression","DQN","agent","."]
docstr: Initializes the agent and constructs the Graph.
  
  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, expects 3 parameters: num_actions, num_atoms,
      network_type.
    kappa: Float, Huber loss cutoff.
    num_atoms: Int, the number of buckets for the value function distribution.
    gamma: Float, exponential decay factor as commonly used in the RL
      literature.
    update_horizon: Int, horizon at which updates are performed, the 'n' in
      n-step update.
    min_replay_history: Int, number of stored transitions for training to
      start.
    update_period: Int, period between DQN updates.
    target_update_period: Int, ppdate period for the target network.
    epsilon_fn: Function expecting 4 parameters: (decay_period, step,
      warmup_steps, epsilon), and which returns the epsilon value used for
      exploration during training.
    epsilon_train: Float, final epsilon for training.
    epsilon_eval: Float, epsilon during evaluation.
    epsilon_decay_period: Int, number of steps for epsilon to decay.
    replay_scheme: String, replay memory scheme to be used. Choices are:
      uniform - Standard (DQN) replay buffer (Mnih et al., 2015)
      prioritized - Prioritized replay buffer (Schaul et al., 2015)
    optimizer: str, name of optimizer to use.
    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.
    seed: int, a seed for DQN's internal RNG, used for initialization and
      sampling actions. If None, will use the current time in nanoseconds.
    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","Graph",".","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",",","expects","3","parameters",":","num_actions",",","num_atoms",",","network_type",".","kappa",":","Float",",","Huber","loss","cutoff",".","num_atoms",":","Int",",","the","number","of","buckets","for","the","value","function","distribution",".","gamma",":","Float",",","exponential","decay","factor","as","commonly","used","in","the","RL","literature",".","update_horizon",":","Int",",","horizon","at","which","updates","are","performed",",","the","'n","'","in","n-step","update",".","min_replay_history",":","Int",",","number","of","stored","transitions","for","training","to","start",".","update_period",":","Int",",","period","between","DQN","updates",".","target_update_period",":","Int",",","ppdate","period","for","the","target","network",".","epsilon_fn",":","Function","expecting","4","parameters",":","(","decay_period",",","step",",","warmup_steps",",","epsilon",")",",","and","which","returns","the","epsilon","value","used","for","exploration","during","training",".","epsilon_train",":","Float",",","final","epsilon","for","training",".","epsilon_eval",":","Float",",","epsilon","during","evaluation",".","epsilon_decay_period",":","Int",",","number","of","steps","for","epsilon","to","decay",".","replay_scheme",":","String",",","replay","memory","scheme","to","be","used",".","Choices","are",":","uniform","-","Standard","(","DQN",")","replay","buffer","(","Mnih","et","al.",",","2015",")","prioritized","-","Prioritized","replay","buffer","(","Schaul","et","al.",",","2015",")","optimizer",":","str",",","name","of","optimizer","to","use",".","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",".","seed",":","int",",","a","seed","for","DQN","'s","internal","RNG",",","used","for","initialization","and","sampling","actions",".","If","None",",","will","use","the","current","time","in","nanoseconds",".","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.QuantileNetwork,
      kappa=1.0,
      num_atoms=200,
      gamma=0.99,
      update_horizon=1,
      min_replay_history=50000,
      update_period=4,
      target_update_period=10000,
      epsilon_fn=dqn_agent.linearly_decaying_epsilon,
      epsilon_train=0.1,
      epsilon_eval=0.05,
      epsilon_decay_period=1000000,
      replay_scheme="prioritized",
      optimizer="adam",
      summary_writer=None,
      summary_writing_frequency=500,
      seed=None,
      allow_partial_reload=False,
  ):
      """Initializes the agent and constructs the Graph.
  
      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, expects 3 parameters: num_actions, num_atoms,
          network_type.
        kappa: Float, Huber loss cutoff.
        num_atoms: Int, the number of buckets for the value function distribution.
        gamma: Float, exponential decay factor as commonly used in the RL
          literature.
        update_horizon: Int, horizon at which updates are performed, the 'n' in
          n-step update.
        min_replay_history: Int, number of stored transitions for training to
          start.
        update_period: Int, period between DQN updates.
        target_update_period: Int, ppdate period for the target network.
        epsilon_fn: Function expecting 4 parameters: (decay_period, step,
          warmup_steps, epsilon), and which returns the epsilon value used for
          exploration during training.
        epsilon_train: Float, final epsilon for training.
        epsilon_eval: Float, epsilon during evaluation.
        epsilon_decay_period: Int, number of steps for epsilon to decay.
        replay_scheme: String, replay memory scheme to be used. Choices are:
          uniform - Standard (DQN) replay buffer (Mnih et al., 2015)
          prioritized - Prioritized replay buffer (Schaul et al., 2015)
        optimizer: str, name of optimizer to use.
        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.
        seed: int, a seed for DQN's internal RNG, used for initialization and
          sampling actions. If None, will use the current time in nanoseconds.
        allow_partial_reload: bool, whether we allow reloading a partial agent
          (for instance, only the network parameters).
      """
      self._num_atoms = num_atoms
      self._kappa = kappa
      self._replay_scheme = replay_scheme
  
      super(JaxQuantileAgent, 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,
          summary_writer=summary_writer,
          summary_writing_frequency=summary_writing_frequency,
          seed=seed,
          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 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 JaxQuantileAgent object from inherited base classes: dqn_agent with example: obj = JaxQuantileAgent(num_actions, observation_shape, observation_dtype, stack_size, network, kappa, num_atoms, gamma, update_horizon, min_replay_history, update_period, target_update_period, epsilon_fn, epsilon_train, epsilon_eval, epsilon_decay_period, replay_scheme, optimizer, summary_writer, summary_writing_frequency, seed, allow_partial_reload)