class_docstr: Atari100k Ranibow Agent with recycled parameters.
class_docstr_tok: ["Atari100k","Ranibow","Agent","with","recycled","parameters","."]
docstr: Gradient update.
docstr_tok: ["Gradient","update","."]
code: def _training_step_update(self):
      """Gradient update."""
      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._rng,
      )
  
      if self._replay_scheme == "prioritized":
          self._replay.set_priority(
              self.replay_elements["indices"], jnp.sqrt(loss + 1e-10)
          )
  
      # We are using # gradient_update_steps in our calculations and logging.
      update_step = self.gradient_step
      is_logging = (
          update_step > 0
          and update_step % self.summary_writing_frequency == 0
      )
      if is_logging:
          self._log_stats(
              {"CrossEntropyLoss": float(mean_loss)}, update_step
          )
  
      self.online_params = self._weight_recycle(
          update_step, self.online_params
      )
  
file_path: dopamine/dopamine/labs/redo/recycled_atari100k_rainbow_agent.py
imports: from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  from dopamine.labs.atari_100k import atari_100k_rainbow_agent
  from dopamine.labs.redo import networks
  from dopamine.labs.redo import weight_recyclers
  from dopamine.metrics import statistics_instance
  import gin
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 4 for calling the RecycledAtari100kRainbowAgent obj's underlying member method code with example usage: obj._training_step_update() without return types