class_docstr: Ranibow Agent with recycled parameters.
class_docstr_tok: ["Ranibow","Agent","with","recycled","parameters","."]
code: def _training_step_update(self):
      self._sample_from_replay_buffer()
      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:
          loss_weights = jnp.ones(
              self.replay_elements["state"].shape[0]
          )
  
      results = rainbow_agent.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,
      )
      (
          self.optimizer_state,
          self.online_params,
          loss,
          mean_loss,
      ) = results
      online_params = self.online_params
  
      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({"Loss": float(mean_loss)}, update_step)
  
      # Neuron/layer recycling starts if reset_mode is not None.
      # Otherwise, we log dead neurons over training for standard agent.
      is_intermediated = self.weight_recycler.is_intermediated_required(
          update_step
      )
      # get intermediate activation per layer to calculate neuron score
      intermediates = (
          self.get_intermediates(online_params)
          if is_intermediated
          else None
      )
      log_dict_neurons = self.weight_recycler.maybe_log_deadneurons(
          update_step, intermediates
      )
      # logging dead neurons.
      self._log_stats(log_dict_neurons, update_step)
      if self.is_debugging:
          log_dict_intersected = self.weight_recycler.intersected_dead_neurons_with_last_reset(
              intermediates, update_step
          )
          self._log_stats(log_dict_intersected, update_step)
  
      # Neuron/layer recyling.
      self._rng, key = jax.random.split(self._rng)
      (
          online_params,
          opt_state,
      ) = self.weight_recycler.maybe_update_weights(
          update_step,
          intermediates,
          online_params,
          key,
          self.optimizer_state,
      )
      self.optimizer_state = opt_state
      self.online_params = online_params
  
file_path: dopamine/dopamine/labs/redo/recycled_rainbow_agent.py
imports: import functools
  from dopamine.jax.agents.rainbow import 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 RecycledRainbowAgent obj's underlying member method code with example usage: obj._training_step_update() without return types