class_docstr: A compact implementation of the full Rainbow agent.
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 = isinstance(
          self._replay,
          prioritized_replay_buffer.OutOfGraphPrioritizedReplayBuffer,
      )
      if is_prioritized and priority is None:
          if self._replay_scheme == "uniform":
              priority = 1.0
          else:
              priority = self._replay.sum_tree.max_recorded_priority
  
      if not self.eval_mode:
          self._replay.add(
              last_observation,
              action,
              reward,
              is_terminal,
              *args,
              priority=priority,
              episode_end=episode_end,
          )
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_agent
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 5 for calling the JaxFullRainbowAgent obj's underlying member method code with example usage: obj._store_transition(last_observation, action, reward, is_terminal) without return types