class_docstr: A JAX implementation of the DQN agent.
class_docstr_tok: ["A","JAX","implementation","of","the","DQN","agent","."]
docstr: Signals the end of the episode to the agent.
  
  We store the observation of the current time step, which is the last
  observation of the episode.
  
  Args:
    reward: float, the last reward from the environment.
    terminal: bool, whether the last state-action led to a terminal state.
docstr_tok: ["Signals","the","end","of","the","episode","to","the","agent",".","We","store","the","observation","of","the","current","time","step",",","which","is","the","last","observation","of","the","episode",".","Args",":","reward",":","float",",","the","last","reward","from","the","environment",".","terminal",":","bool",",","whether","the","last","state-action","led","to","a","terminal","state","."]
code: def end_episode(self, reward, terminal=True):
      """Signals the end of the episode to the agent.
  
      We store the observation of the current time step, which is the last
      observation of the episode.
  
      Args:
        reward: float, the last reward from the environment.
        terminal: bool, whether the last state-action led to a terminal state.
      """
      if not self.eval_mode:
          argspec = inspect.getfullargspec(self._store_transition)
          if (
              "episode_end" in argspec.args
              or "episode_end" in argspec.kwonlyargs
          ):
              self._store_transition(
                  self._observation,
                  self.action,
                  reward,
                  terminal,
                  episode_end=True,
              )
          else:
              logging.warning(
                  "_store_transition function doesn't have episode_end arg."
              )
              self._store_transition(
                  self._observation, self.action, reward, terminal
              )
  
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 10 for calling the JaxDQNAgent obj's underlying member method code with example usage: obj.end_episode(reward, terminal) without return types