class_docstr: A compact implementation of a simplified Rainbow agent.
class_docstr_tok: ["A","compact","implementation","of","a","simplified","Rainbow","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.
      """
      if self._replay.add_count > self.min_replay_history:
          if self.training_steps % self.update_period == 0:
              self._sample_from_replay_buffer()
  
              if self._replay_scheme == "prioritized":
                  # The original prioritized experience replay uses a linear exponent
                  # schedule 0.4 -> 1.0. Comparing the schedule to a fixed exponent of
                  # 0.5 on 5 games (Asterix, Pong, Q*Bert, Seaquest, Space Invaders)
                  # suggested a fixed exponent actually performs better, except on Pong.
                  probs = self.replay_elements["sampling_probabilities"]
                  # Weight the loss by the inverse priorities.
                  loss_weights = 1.0 / jnp.sqrt(probs + 1e-10)
                  loss_weights /= jnp.max(loss_weights)
              else:
                  loss_weights = jnp.ones(
                      self.replay_elements["state"].shape[0]
                  )
  
              (
                  self.optimizer_state,
                  self.online_params,
                  loss,
                  mean_loss,
              ) = train(
                  self.network_def,
                  self.online_params,
                  self.target_network_params,
                  self.optimizer,
                  self.optimizer_state,
                  self.preprocess_fn(self.replay_elements["state"]),
                  self.replay_elements["action"],
                  self.preprocess_fn(
                      self.replay_elements["next_state"]
                  ),
                  self.replay_elements["reward"],
                  self.replay_elements["terminal"],
                  loss_weights,
                  self._support,
                  self.cumulative_gamma,
              )
  
              if self._replay_scheme == "prioritized":
                  # Rainbow and prioritized replay are parametrized by an exponent
                  # alpha, but in both cases it is set to 0.5 - for simplicity's sake we
                  # leave it as is here, using the more direct sqrt(). Taking the square
                  # root "makes sense", as we are dealing with a squared loss.  Add a
                  # small nonzero value to the loss to avoid 0 priority items. While
                  # technically this may be okay, setting all items to 0 priority will
                  # cause troubles, and also result in 1.0 / 0.0 = NaN correction terms.
                  self._replay.set_priority(
                      self.replay_elements["indices"],
                      jnp.sqrt(loss + 1e-10),
                  )
  
              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(
                          "CrossEntropyLoss",
                          mean_loss,
                          step=self.training_steps,
                      )
                  self.summary_writer.flush()
                  if hasattr(self, "collector_dispatcher"):
                      self.collector_dispatcher.write(
                          [
                              statistics_instance.StatisticsInstance(
                                  "Loss",
                                  onp.asarray(mean_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: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_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 6 for calling the JaxRainbowAgent obj's underlying member method code with example usage: obj._train_step() without return types