class_docstr: A compact implementation of agents for Atari 100k.
docstr: Creates the Rainbow-based agent for the Atari 100k benchmark.
  
  On Atari 100k, an agent is evaluated after 100k environment steps, which
  corresponds to 2-3 hours of game play, for training.
  Args:
    num_actions: int, number of actions the agent can take at any state.
    data_augmentation: bool, whether to use data augmentation.
    mse_loss: bool, mse loss function.
    summary_writer: SummaryWriter object, for outputting training statistics.
    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.
    seed: int, a seed for Jax RNG and initialization.
docstr_tok: ["Creates","the","Rainbow-based","agent","for","the","Atari","100k","benchmark",".","On","Atari","100k",",","an","agent","is","evaluated","after","100k","environment","steps",",","which","corresponds","to","2-3","hours","of","game","play",",","for","training",".","Args",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","data_augmentation",":","bool",",","whether","to","use","data","augmentation",".","mse_loss",":","bool",",","mse","loss","function",".","summary_writer",":","SummaryWriter","object",",","for","outputting","training","statistics",".","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",".","seed",":","int",",","a","seed","for","Jax","RNG","and","initialization","."]
code: def __init__(
      self,
      num_actions,
      data_augmentation=False,
      mse_loss=False,
      summary_writer=None,
      network=networks.FullRainbowNetwork,
      seed=None,
  ):
      """Creates the Rainbow-based agent for the Atari 100k benchmark.
  
      On Atari 100k, an agent is evaluated after 100k environment steps, which
      corresponds to 2-3 hours of game play, for training.
      Args:
        num_actions: int, number of actions the agent can take at any state.
        data_augmentation: bool, whether to use data augmentation.
        mse_loss: bool, mse loss function.
        summary_writer: SummaryWriter object, for outputting training statistics.
        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.
        seed: int, a seed for Jax RNG and initialization.
      """
      super().__init__(
          num_actions=num_actions,
          preprocess_fn=preprocess_inputs_with_augmentation,
          summary_writer=summary_writer,
          network=network,
          seed=seed,
      )
      logging.info("\t data_augmentation: %s", data_augmentation)
      self._data_augmentation = data_augmentation
      self._mse_loss = mse_loss
      logging.info("\t data_augmentation: %s", data_augmentation)
      # Preprocessing during training and evaluation can be possibly different,
      # for example, when using data augmentation during training.
      self.train_preprocess_fn = functools.partial(
          preprocess_inputs_with_augmentation,
          data_augmentation=data_augmentation,
      )
      self.state_shape = self.state.shape
  
imports: import copy
  import functools
  from absl import logging
  from dopamine.jax import networks
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 1 to create a new Atari100kRainbowAgent object from inherited base classes: full_rainbow_agent with example: obj = Atari100kRainbowAgent(num_actions, data_augmentation, mse_loss, summary_writer, network, seed)