class_docstr: A compact implementation of the full Rainbow agent.
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()
      states = self.preprocess_fn(self.replay_elements["state"])
      next_states = self.preprocess_fn(
          self.replay_elements["next_state"]
      )
  
      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:
          # 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,
      ) = 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":
          # 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,
              )
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_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 4 for calling the JaxFullRainbowAgent obj's underlying member method code with example usage: obj._training_step_update() without return types