class_docstr: A JAX implementation of the DQN agent.
class_docstr_tok: ["A","JAX","implementation","of","the","DQN","agent","."]
docstr: Stores a transition when in training mode.
  
  Stores the following tuple in the replay buffer (last_observation, action,
  reward, is_terminal, priority).
  
  Args:
    last_observation: Last observation, type determined via observation_type
      parameter in the replay_memory constructor.
    action: An integer, the action taken.
    reward: A float, the reward.
    is_terminal: Boolean indicating if the current state is a terminal state.
    *args: Any, other items to be added to the replay buffer.
    priority: Float. Priority of sampling the transition. If None, the default
      priority will be used. If replay scheme is uniform, the default priority
      is 1. If the replay scheme is prioritized, the default priority is the
      maximum ever seen [Schaul et al., 2015].
    episode_end: bool, whether this transition is the last for the episode.
      This can be different than terminal when ending the episode because
      of a timeout, for example.
docstr_tok: ["Stores","a","transition","when","in","training","mode",".","Stores","the","following","tuple","in","the","replay","buffer","(","last_observation",",","action",",","reward",",","is_terminal",",","priority",")",".","Args",":","last_observation",":","Last","observation",",","type","determined","via","observation_type","parameter","in","the","replay_memory","constructor",".","action",":","An","integer",",","the","action","taken",".","reward",":","A","float",",","the","reward",".","is_terminal",":","Boolean","indicating","if","the","current","state","is","a","terminal","state",".","*","args",":","Any",",","other","items","to","be","added","to","the","replay","buffer",".","priority",":","Float",".","Priority","of","sampling","the","transition",".","If","None",",","the","default","priority","will","be","used",".","If","replay","scheme","is","uniform",",","the","default","priority","is","1",".","If","the","replay","scheme","is","prioritized",",","the","default","priority","is","the","maximum","ever","seen","[","Schaul","et","al.",",","2015","]",".","episode_end",":","bool",",","whether","this","transition","is","the","last","for","the","episode",".","This","can","be","different","than","terminal","when","ending","the","episode","because","of","a","timeout",",","for","example","."]
code: def _store_transition(
      self,
      last_observation,
      action,
      reward,
      is_terminal,
      *args,
      priority=None,
      episode_end=False,
  ):
      """Stores a transition when in training mode.
  
      Stores the following tuple in the replay buffer (last_observation, action,
      reward, is_terminal, priority).
  
      Args:
        last_observation: Last observation, type determined via observation_type
          parameter in the replay_memory constructor.
        action: An integer, the action taken.
        reward: A float, the reward.
        is_terminal: Boolean indicating if the current state is a terminal state.
        *args: Any, other items to be added to the replay buffer.
        priority: Float. Priority of sampling the transition. If None, the default
          priority will be used. If replay scheme is uniform, the default priority
          is 1. If the replay scheme is prioritized, the default priority is the
          maximum ever seen [Schaul et al., 2015].
        episode_end: bool, whether this transition is the last for the episode.
          This can be different than terminal when ending the episode because
          of a timeout, for example.
      """
      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 collections
  import functools
  import inspect
  import math
  import time
  from absl import logging
  from dopamine.agents.dqn import dqn_agent
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import circular_replay_buffer
  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 12 for calling the JaxDQNAgent obj's underlying member method code with example usage: obj._store_transition(last_observation, action, reward, is_terminal) without return types