class_docstr: Wrapped DqnAgent that supports recycled training.
class_docstr_tok: ["Wrapped","DqnAgent","that","supports","recycled","training","."]
code: def reset_weights(self, optimizer, model, target_model=None):
      logging.info("weights reset")
      model_layers = get_all_layers(model)
      target_model_layers = (
          model_layers
          if target_model is None
          else get_all_layers(target_model)
      )
      assert len(model_layers) == len(target_model_layers)
      for layer, target_layer in zip(model_layers, target_model_layers):
          assert len(layer.weights) == len(target_layer.weights)
          for var, target_var in zip(
              layer.weights, target_layer.weights
          ):
              abs_weights = tf.math.abs(var)
              k = tf.dtypes.cast(
                  tf.math.maximum(
                      tf.math.round(
                          tf.dtypes.cast(
                              tf.size(abs_weights), tf.float32
                          )
                          * self.reset_frac
                      ),
                      1,
                  ),
                  tf.int32,
              )
              w_initializer = layer.kernel_initializer
              logging.info("shape %s", var.shape)
              logging.info("w_initializer%s", w_initializer)
  
              b_initializer = layer.bias_initializer
              new_weights, mask = self.create_new_weights(
                  k, var, abs_weights, w_initializer, b_initializer
              )
              var.assign(new_weights)
              if target_model is not None:
                  new_target_weights = tf.where(
                      mask == 1, new_weights, target_var
                  )
                  target_var.assign(new_target_weights)
  
              # reset momentum for the weights that are resetted
              if len(var.get_shape().as_list()) > 1:
                  self.reset_momentum(optimizer, var, 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 13 for calling the RecycledSacAgent obj's underlying member method code with example usage: obj.reset_weights(optimizer, model, target_model) without return types