class_docstr: A compact implementation of agents for Atari 100k.
docstr: Stores a transition when in training mode.
docstr_tok: ["Stores","a","transition","when","in","training","mode","."]
code: def _store_transition(
      self,
      last_observation,
      action,
      reward,
      is_terminal,
      *args,
      priority=None,
      episode_end=False,
  ):
      """Stores a transition when in training mode."""
      is_prioritized = hasattr(self._replay, "sum_tree")
      if is_prioritized and priority is None:
          priority = onp.ones_like(reward)
          if self._replay_scheme == "prioritized":
              priority *= self._replay.sum_tree.max_recorded_priority
  
      to_store = (last_observation, action, reward, is_terminal, *args)
      to_store = (onp.asarray(x) for x in to_store)
      if not hasattr(self._replay, "_n_envs"):
          to_store = (onp.squeeze(x) for x in to_store)
          priority = onp.squeeze(priority)
      elif hasattr(self._replay, "_n_envs") and not reward.shape:
          to_store = (
              onp.expand_dims(x, 0)
              if not (x.shape and x.shape[0] == 1)
              else x
              for x in to_store
          )
          priority = onp.expand_dims(priority, 0)
      if not self.eval_mode:
          self._replay.add(
              *to_store, priority=priority, episode_end=episode_end
          )
  
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 12 for calling the Atari100kRainbowAgent obj's underlying member method code with example usage: obj._store_transition(last_observation, action, reward, is_terminal) without return types