class_docstr: A compact implementation of agents for Atari 100k.
docstr: Gradient update during every training step.
docstr_tok: ["Gradient","update","during","every","training","step","."]
code: def _training_step_update(self):
      """Gradient update during every training step."""
  
      self._sample_from_replay_buffer()
      self._rng, rng1, rng2 = jax.random.split(self._rng, num=3)
      states = self.train_preprocess_fn(
          self.replay_elements["state"], rng=rng1
      )
      next_states = self.train_preprocess_fn(
          self.replay_elements["next_state"], rng=rng2
      )
  
      if self._replay_scheme == "prioritized":
          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:
          # Uniform weights if not using prioritized replay.
          loss_weights = jnp.ones(states.shape[0])
  
      (
          self.optimizer_state,
          self.online_params,
          loss,
          mean_loss,
          self._rng,
      ) = full_rainbow_agent.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"],
          loss_weights,
          self._support,
          self.cumulative_gamma,
          self._double_dqn,
          self._distributional,
          self._mse_loss,
          self._rng,
      )
  
      if self._replay_scheme == "prioritized":
          self._replay.set_priority(
              self.replay_elements["indices"], jnp.sqrt(loss + 1e-10)
          )
  
      if self.summary_writer is not None:
          with self.summary_writer.as_default():
              tf.summary.scalar(
                  "CrossEntropyLoss",
                  mean_loss,
                  step=self.training_steps,
              )
          self.summary_writer.flush()
  
imports: import copy
  import functools
  from absl import logging
  from dopamine.jax import networks
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling the Atari100kRainbowAgent obj's underlying member method code with example usage: obj._training_step_update() without return types