class_docstr: A compact implementation of SPR in Jax.
class_docstr_tok: ["A","compact","implementation","of","SPR","in","Jax","."]
docstr: Creates the replay buffer used by the agent.
docstr_tok: ["Creates","the","replay","buffer","used","by","the","agent","."]
code: def _build_replay_buffer(self):
      """Creates the replay buffer used by the agent."""
      if self._replay_scheme not in ["uniform", "prioritized"]:
          raise ValueError(
              "Invalid replay scheme: {}".format(self._replay_scheme)
          )
      if self._replay_scheme == "prioritized":
          return replay_buffers.PrioritizedJaxSubsequenceParallelEnvReplayBuffer(
              observation_shape=self.observation_shape,
              stack_size=self.stack_size,
              update_horizon=self.update_horizon,
              gamma=self.gamma,
              subseq_len=self._jumps + 1,
              observation_dtype=self.observation_dtype,
          )
      else:
          return replay_buffers.JaxSubsequenceParallelEnvReplayBuffer(
              observation_shape=self.observation_shape,
              stack_size=self.stack_size,
              update_horizon=self.update_horizon,
              gamma=self.gamma,
              subseq_len=self._jumps + 1,
              observation_dtype=self.observation_dtype,
          )
  
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 3 for calling the SPRAgent obj's underlying member method code with example usage: obj._build_replay_buffer() and returns: replay_buffers, replay_buffers