class_docstr: A compact implementation of the full 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.
    noisy: bool, Whether to use noisy networks or not.
    dueling: bool, Whether to use dueling network architecture or not.
    double_dqn: bool, Whether to use Double DQN or not.
    distributional: bool, whether to use distributional RL or not.
    mse_loss: bool, mse loss function.
    num_updates_per_train_step: int, Number of gradient updates every training
      step. Defaults to 1.
    network: flax.linen Module, neural network used by the agent 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.
    vmax: float, the value distribution support is [vmin, vmax].
    vmin: float, the value distribution support is [vmin, vmax]. If vmin is
      None, it is set to -vmax.
    epsilon_fn: function expecting 4 parameters: (decay_period, step,
      warmup_steps, epsilon). This function should return the epsilon value
      used for exploration during training.
    replay_scheme: str, 'prioritized' or 'uniform', the sampling scheme of the
      replay memory.
    summary_writer: SummaryWriter object, for outputting training statistics.
    seed: int, a seed for Jax RNG and initialization.
    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.
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Args",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","noisy",":","bool",",","Whether","to","use","noisy","networks","or","not",".","dueling",":","bool",",","Whether","to","use","dueling","network","architecture","or","not",".","double_dqn",":","bool",",","Whether","to","use","Double","DQN","or","not",".","distributional",":","bool",",","whether","to","use","distributional","RL","or","not",".","mse_loss",":","bool",",","mse","loss","function",".","num_updates_per_train_step",":","int",",","Number","of","gradient","updates","every","training","step",".","Defaults","to","1.","network",":","flax.linen","Module",",","neural","network","used","by","the","agent","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",".","vmax",":","float",",","the","value","distribution","support","is","[","vmin",",","vmax","]",".","vmin",":","float",",","the","value","distribution","support","is","[","vmin",",","vmax","]",".","If","vmin","is","None",",","it","is","set","to","-vmax",".","epsilon_fn",":","function","expecting","4","parameters",":","(","decay_period",",","step",",","warmup_steps",",","epsilon",")",".","This","function","should","return","the","epsilon","value","used","for","exploration","during","training",".","replay_scheme",":","str",",","'prioritized","'","or","'uniform","'",",","the","sampling","scheme","of","the","replay","memory",".","summary_writer",":","SummaryWriter","object",",","for","outputting","training","statistics",".","seed",":","int",",","a","seed","for","Jax","RNG","and","initialization",".","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","."]
code: def __init__(
      self,
      num_actions,
      noisy=True,
      dueling=True,
      double_dqn=True,
      distributional=True,
      mse_loss=False,
      num_updates_per_train_step=1,
      network=networks.FullRainbowNetwork,
      num_atoms=51,
      vmax=10.0,
      vmin=None,
      epsilon_fn=dqn_agent.linearly_decaying_epsilon,
      replay_scheme="prioritized",
      summary_writer=None,
      seed=None,
      preprocess_fn=None,
  ):
      """Initializes the agent and constructs the necessary components.
  
      Args:
        num_actions: int, number of actions the agent can take at any state.
        noisy: bool, Whether to use noisy networks or not.
        dueling: bool, Whether to use dueling network architecture or not.
        double_dqn: bool, Whether to use Double DQN or not.
        distributional: bool, whether to use distributional RL or not.
        mse_loss: bool, mse loss function.
        num_updates_per_train_step: int, Number of gradient updates every training
          step. Defaults to 1.
        network: flax.linen Module, neural network used by the agent 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.
        vmax: float, the value distribution support is [vmin, vmax].
        vmin: float, the value distribution support is [vmin, vmax]. If vmin is
          None, it is set to -vmax.
        epsilon_fn: function expecting 4 parameters: (decay_period, step,
          warmup_steps, epsilon). This function should return the epsilon value
          used for exploration during training.
        replay_scheme: str, 'prioritized' or 'uniform', the sampling scheme of the
          replay memory.
        summary_writer: SummaryWriter object, for outputting training statistics.
        seed: int, a seed for Jax RNG and initialization.
        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.
      """
      logging.info(
          "Creating %s agent with the following parameters:",
          self.__class__.__name__,
      )
      logging.info("\t double_dqn: %s", double_dqn)
      logging.info("\t noisy_networks: %s", noisy)
      logging.info("\t dueling_dqn: %s", dueling)
      logging.info("\t distributional: %s", distributional)
      logging.info("\t mse_loss: %d", mse_loss)
      logging.info("\t num_atoms: %d", num_atoms)
      logging.info("\t replay_scheme: %s", replay_scheme)
      logging.info(
          "\t num_updates_per_train_step: %d",
          num_updates_per_train_step,
      )
      # We need this because some tools convert round floats into ints.
      vmax = float(vmax)
      self._num_atoms = num_atoms
      vmin = vmin if vmin else -vmax
      self._support = jnp.linspace(vmin, vmax, num_atoms)
      self._replay_scheme = replay_scheme
      self._double_dqn = double_dqn
      self._noisy = noisy
      self._dueling = dueling
      self._distributional = distributional
      self._mse_loss = mse_loss
      self._num_updates_per_train_step = num_updates_per_train_step
  
      super().__init__(
          num_actions=num_actions,
          network=functools.partial(
              network,
              num_atoms=num_atoms,
              noisy=self._noisy,
              dueling=self._dueling,
              distributional=self._distributional,
          ),
          epsilon_fn=zero_epsilon if self._noisy else epsilon_fn,
          summary_writer=summary_writer,
          seed=seed,
          preprocess_fn=preprocess_fn,
      )
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_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 JaxFullRainbowAgent object from inherited base classes: dqn_agent with example: obj = JaxFullRainbowAgent(num_actions, noisy, dueling, double_dqn, distributional, mse_loss, num_updates_per_train_step, network, num_atoms, vmax, vmin, epsilon_fn, replay_scheme, summary_writer, seed, preprocess_fn)