class_docstr: A JAX implementation of the DQN agent.
class_docstr_tok: ["A","JAX","implementation","of","the","DQN","agent","."]
docstr: Runs a single training step.
  
  Runs training if both:
    (1) A minimum number of frames have been added to the replay buffer.
    (2) `training_steps` is a multiple of `update_period`.
  
  Also, syncs weights from online_params to target_network_params if training
  steps is a multiple of target update period.
docstr_tok: ["Runs","a","single","training","step",".","Runs","training","if","both",":","(","1",")","A","minimum","number","of","frames","have","been","added","to","the","replay","buffer",".","(","2",")","`","training_steps","`","is","a","multiple","of","`","update_period","`",".","Also",",","syncs","weights","from","online_params","to","target_network_params","if","training","steps","is","a","multiple","of","target","update","period","."]
code: def _train_step(self):
      """Runs a single training step.
  
      Runs training if both:
        (1) A minimum number of frames have been added to the replay buffer.
        (2) `training_steps` is a multiple of `update_period`.
  
      Also, syncs weights from online_params to target_network_params if training
      steps is a multiple of target update period.
      """
      # Run a train op at the rate of self.update_period if enough training steps
      # have been run. This matches the Nature DQN behaviour.
      if self._replay.add_count > self.min_replay_history:
          if self.training_steps % self.update_period == 0:
              self._sample_from_replay_buffer()
              states = self.preprocess_fn(self.replay_elements["state"])
              next_states = self.preprocess_fn(
                  self.replay_elements["next_state"]
              )
              self.optimizer_state, self.online_params, loss = train(
                  self.network_def,
                  self.online_params,
                  self.target_network_params,
                  self.optimizer,
                  self.optimizer_state,
                  states,
                  self.replay_elements["action"],
                  next_states,
                  self.replay_elements["reward"],
                  self.replay_elements["terminal"],
                  self.cumulative_gamma,
                  self._loss_type,
              )
              if (
                  self.summary_writer is not None
                  and self.training_steps > 0
                  and self.training_steps
                  % self.summary_writing_frequency
                  == 0
              ):
                  with self.summary_writer.as_default():
                      tf.summary.scalar(
                          "HuberLoss", loss, step=self.training_steps
                      )
                  self.summary_writer.flush()
                  if hasattr(self, "collector_dispatcher"):
                      self.collector_dispatcher.write(
                          [
                              statistics_instance.StatisticsInstance(
                                  "Loss",
                                  onp.asarray(loss),
                                  step=self.training_steps,
                              ),
                          ],
                          collector_allowlist=self._collector_allowlist,
                      )
          if self.training_steps % self.target_update_period == 0:
              self._sync_weights()
  
      self.training_steps += 1
  
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 11 for calling the JaxDQNAgent obj's underlying member method code with example usage: obj._train_step() without return types