docstr: Tests the compute/apply_gradients with add/remove modules.
  
  Tests that if we add a module with SGD optimizer with a known lr (different
  from default), and remove the default module, with a loss that is the sum of
  all variables the updated parameters follow the SGD update rule.
docstr_tok: ["Tests","the","compute\/apply_gradients","with","add\/remove","modules",".","Tests","that","if","we","add","a","module","with","SGD","optimizer","with","a","known","lr","(","different","from","default",")",",","and","remove","the","default","module",",","with","a","loss","that","is","the","sum","of","all","variables","the","updated","parameters","follow","the","SGD","update","rule","."]
code: def test_add_remove_module(self):
      """Tests the compute/apply_gradients with add/remove modules.
  
      Tests that if we add a module with SGD optimizer with a known lr (different
      from default), and remove the default module, with a loss that is the sum of
      all variables the updated parameters follow the SGD update rule.
      """
      config = BaseTestingAlgorithmConfig().training(lr=0.0003)
  
      for fw in framework_iterator(config, frameworks=("torch", "tf2")):
          learner = config.build_learner(env=self.ENV)
          rl_module_spec = config.get_default_rl_module_spec()
          rl_module_spec.observation_space = self.ENV.observation_space
          rl_module_spec.action_space = self.ENV.action_space
          learner.add_module(
              module_id="test",
              module_spec=rl_module_spec,
          )
          learner.remove_module(DEFAULT_POLICY_ID)
  
          # only test module should be left
          self.assertEqual(set(learner.module.keys()), {"test"})
  
          # calculated the expected new params based on gradients of all ones.
          params = learner.get_parameters(learner.module["test"])
          n_steps = 100
          expected = [
              convert_to_numpy(param)
              - n_steps * learner.config.lr * np.ones(param.shape)
              for param in params
          ]
          for _ in range(n_steps):
              tape = None
              if fw == "torch":
                  loss_per_module = {
                      ALL_MODULES: sum(param.sum() for param in params)
                  }
              else:
                  with tf.GradientTape() as tape:
                      loss_per_module = {
                          ALL_MODULES: sum(
                              tf.reduce_sum(param) for param in params
                          )
                      }
              gradients = learner.compute_gradients(
                  loss_per_module, gradient_tape=tape
              )
              learner.apply_gradients(gradients)
  
          check(params, expected)
  
imports: import gymnasium
  import numpy
  import tempfile
  import unittest
  import ray
  from ray.rllib.core.learner.learner import Learner
  from ray.rllib.core.testing.testing_learner import BaseTestingAlgorithmConfig
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.test_utils import check, framework_iterator, get_cartpole_dataset_reader
  from ray.rllib.utils.metrics import ALL_MODULES
label_desc: Use image node_id 7 for calling the TestLearner obj's underlying member method code with example usage: obj.test_add_remove_module() without return types