class_docstr: A compact implementation of SPR in Jax.
class_docstr_tok: ["A","compact","implementation","of","SPR","in","Jax","."]
code: def _sample_from_replay_buffer(self):
      self._rng, rng = jax.random.split(self._rng)
      samples = self._replay.sample_transition_batch(rng)
      types = self._replay.get_transition_elements()
      self.replay_elements = collections.OrderedDict()
      for element, element_type in zip(samples, types):
          self.replay_elements[element_type.name] = element
  
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 4 for calling the SPRAgent obj's underlying member method code with example usage: obj._sample_from_replay_buffer() without return types