class_docstr: A compact implementation of SPR in Jax.
class_docstr_tok: ["A","compact","implementation","of","SPR","in","Jax","."]
docstr: Initializes the agent and constructs the necessary components.
  
  Args:
    num_actions: int, number of actions the agent can take at any state.
    jumps: int >= 0, number of SPR prediction steps to do.
    spr_weight: float, weight given to the SPR loss.
    summary_writer: SummaryWriter object, for outputting training statistics.
    seed: int, a seed for Jax RNG and initialization.
    epsilon_fn: Type of epsilon decay to use. By default, linearly_decaying
      will use e-greedy during initial data collection, matching the PyTorch
      codebase.
    network: Network class to use.
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Args",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","jumps",":","int",">","=","0",",","number","of","SPR","prediction","steps","to","do",".","spr_weight",":","float",",","weight","given","to","the","SPR","loss",".","summary_writer",":","SummaryWriter","object",",","for","outputting","training","statistics",".","seed",":","int",",","a","seed","for","Jax","RNG","and","initialization",".","epsilon_fn",":","Type","of","epsilon","decay","to","use",".","By","default",",","linearly_decaying","will","use","e-greedy","during","initial","data","collection",",","matching","the","PyTorch","codebase",".","network",":","Network","class","to","use","."]
code: def __init__(
      self,
      num_actions,
      jumps=5,
      spr_weight=5,
      summary_writer=None,
      seed=None,
      epsilon_fn=dqn_agent.linearly_decaying_epsilon,
      network=SPRNetwork,
  ):
      """Initializes the agent and constructs the necessary components.
  
      Args:
        num_actions: int, number of actions the agent can take at any state.
        jumps: int >= 0, number of SPR prediction steps to do.
        spr_weight: float, weight given to the SPR loss.
        summary_writer: SummaryWriter object, for outputting training statistics.
        seed: int, a seed for Jax RNG and initialization.
        epsilon_fn: Type of epsilon decay to use. By default, linearly_decaying
          will use e-greedy during initial data collection, matching the PyTorch
          codebase.
        network: Network class to use.
      """
      logging.info(
          "Creating %s agent with the following parameters:",
          self.__class__.__name__,
      )
      logging.info("\t spr_weight: %s", spr_weight)
      logging.info("\t jumps: %s", jumps)
      self._jumps = jumps
      self.spr_weight = spr_weight
      super().__init__(
          num_actions=num_actions,
          summary_writer=summary_writer,
          seed=seed,
          network=network,
      )
  
      # Parent class JaxFullRainbowAgent will overwrite this with the wrong value,
      # so just reverse its change.
      self.epsilon_fn = epsilon_fn
      self.start_time = time.time()
  
imports: import collections
  import copy
  import functools
  import time
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_agent
  from dopamine.labs.atari_100k import atari_100k_rainbow_agent
  from dopamine.labs.atari_100k.replay_memory import subsequence_replay_buffer
  from dopamine.labs.atari_100k.spr_networks import SPRNetwork
  import gin
  import jax
  import jax.numpy
  import optax
  import tensorflow
label_desc: Use image node_id 1 to create a new SPRAgent object from inherited base classes: atari_100k_rainbow_agent with example: obj = SPRAgent(num_actions, jumps, spr_weight, summary_writer, seed, epsilon_fn, network)