class_docstr: A compact implementation of SPR in Jax.
class_docstr_tok: ["A","compact","implementation","of","SPR","in","Jax","."]
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."""
  
      inter_batch_time = time.time() - self.start_time
      self.start_time = time.time()
  
      sample_start_time = time.time()
      self._sample_from_replay_buffer()
      sample_time = time.time() - sample_start_time
  
      aug_start_time = time.time()
      # Add code for data augmentation.
      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"][:, 0], rng=rng2
      )
  
      aug_time = time.time() - aug_start_time
      train_start_time = time.time()
  
      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,
          mean_loss,
          dqn_loss,
          spr_loss,
          self._rng,
      ) = train(
          network_def=self.network_def,
          online_params=self.online_params,
          target_params=self.target_network_params,
          optimizer=self.optimizer,
          optimizer_state=self.optimizer_state,
          states=states,
          actions=self.replay_elements["action"],
          next_states=next_states,
          rewards=self.replay_elements["reward"][:, 0],
          terminals=self.replay_elements["terminal"][:, 0],
          same_traj_mask=self.replay_elements["same_trajectory"][:, 1:],
          loss_weights=loss_weights,
          support=self._support,
          cumulative_gamma=self.cumulative_gamma,
          double_dqn=self._double_dqn,
          distributional=self._distributional,
          rng=self._rng,
          spr_weight=self.spr_weight,
      )
  
      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(dqn_loss + 1e-10),
          )
  
      train_time = time.time() - train_start_time
      if (
          self.summary_writer is not None
          and self.training_steps > 0
          and self.training_steps % self.summary_writing_frequency == 0
      ):
          step = self.training_steps
          tf.summary.scalar("TotalLoss", float(mean_loss), step=step)
          tf.summary.scalar(
              "DQNLoss", float(dqn_loss.mean()), step=step
          )
          tf.summary.scalar(
              "SPRLoss", float(spr_loss.mean()), step=step
          )
          tf.summary.scalar(
              "InterbatchTime", float(inter_batch_time), step=step
          )
          tf.summary.scalar("TrainTime", float(train_time), step=step)
          tf.summary.scalar("SampleTime", float(sample_time), step=step)
          tf.summary.scalar("AugTime", float(aug_time), step=step)
          self.summary_writer.flush()
  
imports: import collections
  import copy
  import functools
  import time
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_agent
  from dopamine.labs.atari_100k import atari_100k_rainbow_agent
  from dopamine.labs.atari_100k.replay_memory import subsequence_replay_buffer
  from dopamine.labs.atari_100k.spr_networks import SPRNetwork
  import gin
  import jax
  import jax.numpy
  import optax
  import tensorflow
label_desc: Use image node_id 5 for calling the SPRAgent obj's underlying member method code with example usage: obj._training_step_update() without return types