class_docstr: Wrapped DqnAgent that supports recycled training.
class_docstr_tok: ["Wrapped","DqnAgent","that","supports","recycled","training","."]
docstr: Recycle the dead neurons by reinitializing its weights.
  
  Args:
    optimizer: optimizer of a network weights
    act: activation of each layer that is used to calculate the neuron score
      for some score metrics
    act_grad: gradient of the activation of each hidden layer
    model: behavior network
    target_model: target network. When there is no target model (e.g. actor
      network case), we iterate only on the behavior network layers
docstr_tok: ["Recycle","the","dead","neurons","by","reinitializing","its","weights",".","Args",":","optimizer",":","optimizer","of","a","network","weights","act",":","activation","of","each","layer","that","is","used","to","calculate","the","neuron","score","for","some","score","metrics","act_grad",":","gradient","of","the","activation","of","each","hidden","layer","model",":","behavior","network","target_model",":","target","network",".","When","there","is","no","target","model","(","e.g",".","actor","network","case",")",",","we","iterate","only","on","the","behavior","network","layers"]
code: def reset_dead_neurons(
      self, optimizer, act, act_grad, model, target_model=None
  ):
      """Recycle the dead neurons by reinitializing its weights.
  
      Args:
        optimizer: optimizer of a network weights
        act: activation of each layer that is used to calculate the neuron score
          for some score metrics
        act_grad: gradient of the activation of each hidden layer
        model: behavior network
        target_model: target network. When there is no target model (e.g. actor
          network case), we iterate only on the behavior network layers
      """
      model_layers = get_all_layers(model, filter_fn=is_dense_layer)
      target_model_layers = (
          model_layers
          if target_model is None
          else get_all_layers(target_model, filter_fn=is_dense_layer)
      )
      (
          dead_neuron_masks,
          incoming_masks,
          outgoing_masks,
      ) = self.get_mask_dead_neurons_weights(act, act_grad, model)
      # update incoming weights
      for layer, target_layer in zip(
          model_layers[:-1], target_model_layers[:-1]
      ):
          incoming_mask = incoming_masks[layer.name]
          incoming_var = layer.kernel
          incoming_target_var = target_layer.kernel
          w_initializer = type(layer.kernel_initializer)(
              seed=self.train_step_counter
          )
          new_in_weights = w_initializer(incoming_var.shape)
  
          if self.scale_recycled_weights:
              neuron_mask = dead_neuron_masks[layer.name]
              new_in_weights = self._rescale_weights(
                  neuron_mask,
                  incoming_var,
                  new_in_weights,
                  0,
                  self.recycled_incoming_scaler,
              )
  
          new_in_weights = tf.where(
              incoming_mask == 1, new_in_weights, incoming_var
          )
  
          incoming_var.assign(new_in_weights)
          if self.reset_target_models and target_model is not None:
              new_target_in_weights = tf.where(
                  incoming_mask == 1,
                  new_in_weights,
                  incoming_target_var,
              )
              incoming_target_var.assign(new_target_in_weights)
          # reset momentum for the weights that are resetted
          self.reset_momentum(optimizer, incoming_var, incoming_mask)
  
          # reset bias of dead neurons
          new_bias = tf.zeros_like(layer.bias)
          new_bias = tf.where(
              tf.math.equal(dead_neuron_masks[layer.name], 1),
              new_bias,
              layer.bias,
          )
          layer.bias.assign(new_bias)
  
      # update outgoing weights
      # model layers returns list of each layer in the model.
      # e.g. model_layers = [dense_1, dense_2, output layer]
      # updating the outgoing connections of the first hidden layer (dense_1)
      # means that we update the weights of the second hidden layer (dense_2).
      # in case of weight scaling, we identify the non dead neurons in dense_1
      # to get the average norm that will be used to scale the outgoing weights
      # of dense_1 (weights of dense_2).
      # model[1:] = [dense_2, output layer] to access weights
      # model[:-1] = [dense_1, dense_2] to access the corresponding dead neurons.
      for layer, next_layer, target_next_layer in zip(
          model_layers[:-1], model_layers[1:], target_model_layers[1:]
      ):
          outgoing_mask = outgoing_masks[next_layer.name]
          outgoing_var = next_layer.kernel
          outgoing_target_var = target_next_layer.kernel
          if self.init_method == "random":
              w_initializer = type(next_layer.kernel_initializer)(
                  seed=self.train_step_counter
              )
              new_out_weights = w_initializer(outgoing_var.shape)
              # if we are interested in scaling the ingoing only,
              # keep recycled_outgoing_scaler with the default 0.
              if (
                  self.scale_recycled_weights
                  and self.recycled_outgoing_scaler > 0
              ):
                  neuron_mask = dead_neuron_masks[layer.name]
                  new_out_weights = self._rescale_weights(
                      neuron_mask,
                      outgoing_var,
                      new_out_weights,
                      1,
                      self.recycled_outgoing_scaler,
                  )
  
          elif self.init_method == "zero":
              new_out_weights = tf.zeros_like(outgoing_var)
          else:
              raise ValueError(
                  "init_method:%s  is not valid." % self.init_method
              )
  
          new_out_weights = tf.where(
              outgoing_mask == 1, new_out_weights, outgoing_var
          )
          outgoing_var.assign(new_out_weights)
          if self.reset_target_models and target_model is not None:
              new_target_out_weights = tf.where(
                  outgoing_mask == 1,
                  new_out_weights,
                  outgoing_target_var,
              )
              outgoing_target_var.assign(new_target_out_weights)
          self.reset_momentum(optimizer, outgoing_var, outgoing_mask)
  
imports: import functools
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  import numpy
  import reverb
  from rigl.rl.tfagents import sparse_tanh_normal_projection_network
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.agents.sac import sac_agent
  from tf_agents.environments import suite_mujoco
  from tf_agents.keras_layers import inner_reshape
  from tf_agents.metrics import py_metrics
  from tf_agents.networks import nest_map
  from tf_agents.networks import sequential
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import py_tf_eager_policy
  from tf_agents.policies import random_py_policy
  from tf_agents.replay_buffers import reverb_replay_buffer
  from tf_agents.replay_buffers import reverb_utils
  from tf_agents.train import actor
  from tf_agents.train import learner
  from tf_agents.train import triggers
  from tf_agents.train.utils import spec_utils
  from tf_agents.train.utils import strategy_utils
  from tf_agents.train.utils import train_utils
  from tf_agents.utils import common
  from tf_agents.utils import object_identity
label_desc: Use image node_id 18 for calling the RecycledSacAgent obj's underlying member method code with example usage: obj.reset_dead_neurons(optimizer, act, act_grad, model, target_model) without return types