class_docstr: DQN Agent with masked parameters.
class_docstr_tok: ["DQN","Agent","with","masked","parameters","."]
code: def _training_step_update(self):
      # 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
      )
  
      self._sample_from_replay_buffer()
      batch = {
          k: (
              self.preprocess_fn(v)
              if k in ["state", "next_state"]
              else v
          )
          for k, v in self.replay_elements.items()
      }
      online_params = self.online_params
  
      loss, grad = get_gradients(
          online_params=online_params,
          apply_fn=self.network_def.apply,
          target_params=self.target_network_params,
          batch=batch,
          cumulative_gamma=self.cumulative_gamma,
      )
  
      new_online_params, self.optimizer_state = apply_updates_jitted(
          online_params, grad, self.optimizer_state, self.optimizer
      )
      self.online_params = new_online_params
      online_params = self.online_params
  
      if is_logging:
          self._log_stats({"HuberLoss": float(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
  
imports: import functools
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_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 RecycledDQNAgent obj's underlying member method code with example usage: obj._training_step_update() without return types