class_docstr: A JAX implementation of the DQN agent.
class_docstr_tok: ["A","JAX","implementation","of","the","DQN","agent","."]
docstr: Returns a self-contained bundle of the agent's state.
  
  This is used for checkpointing. It will return a dictionary containing all
  non-TensorFlow objects (to be saved into a file by the caller), and it saves
  all TensorFlow objects into a checkpoint file.
  
  Args:
    checkpoint_dir: str, directory where TensorFlow objects will be saved.
    iteration_number: int, iteration number to use for naming the checkpoint
      file.
  
  Returns:
    A dict containing additional Python objects to be checkpointed by the
      experiment. If the checkpoint directory does not exist, returns None.
docstr_tok: ["Returns","a","self-contained","bundle","of","the","agent","'s","state",".","This","is","used","for","checkpointing",".","It","will","return","a","dictionary","containing","all","non-TensorFlow","objects","(","to","be","saved","into","a","file","by","the","caller",")",",","and","it","saves","all","TensorFlow","objects","into","a","checkpoint","file",".","Args",":","checkpoint_dir",":","str",",","directory","where","TensorFlow","objects","will","be","saved",".","iteration_number",":","int",",","iteration","number","to","use","for","naming","the","checkpoint","file",".","Returns",":","A","dict","containing","additional","Python","objects","to","be","checkpointed","by","the","experiment",".","If","the","checkpoint","directory","does","not","exist",",","returns","None","."]
code: def bundle_and_checkpoint(self, checkpoint_dir, iteration_number):
      """Returns a self-contained bundle of the agent's state.
  
      This is used for checkpointing. It will return a dictionary containing all
      non-TensorFlow objects (to be saved into a file by the caller), and it saves
      all TensorFlow objects into a checkpoint file.
  
      Args:
        checkpoint_dir: str, directory where TensorFlow objects will be saved.
        iteration_number: int, iteration number to use for naming the checkpoint
          file.
  
      Returns:
        A dict containing additional Python objects to be checkpointed by the
          experiment. If the checkpoint directory does not exist, returns None.
      """
      if not tf.io.gfile.exists(checkpoint_dir):
          return None
      # Checkpoint the out-of-graph replay buffer.
      self._replay.save(checkpoint_dir, iteration_number)
      bundle_dictionary = {
          "state": self.state,
          "training_steps": self.training_steps,
          "online_params": self.online_params,
          "optimizer_state": self.optimizer_state,
          "target_params": self.target_network_params,
      }
      return bundle_dictionary
  
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 13 for calling the JaxDQNAgent obj's underlying member method code with example usage: obj.bundle_and_checkpoint(checkpoint_dir, iteration_number) and returns: bundle_dictionary, None