docstr: Test saving and loading from checkpoint after adding / removing modules.
docstr_tok: ["Test","saving","and","loading","from","checkpoint","after","adding","\/","removing","modules","."]
code: def test_save_to_from_checkpoint(self):
      """Test saving and loading from checkpoint after adding / removing modules."""
      env_class = make_multi_agent("CartPole-v0")
      env = env_class({"num_agents": 2})
      module = DiscreteBCTorchModule(
          config=RLModuleConfig(
              env.observation_space,
              env.action_space,
              model_config_dict={"fcnet_hiddens": [32]},
          )
      ).as_multi_agent()
  
      module.add_module(
          "test",
          DiscreteBCTorchModule(
              config=RLModuleConfig(
                  env.observation_space,
                  env.action_space,
                  model_config_dict={"fcnet_hiddens": [32]},
              )
          ),
      )
      module.add_module(
          "test2",
          DiscreteBCTorchModule(
              config=RLModuleConfig(
                  env.observation_space,
                  env.action_space,
                  model_config_dict={"fcnet_hiddens": [128]},
              )
          ),
      )
  
      with tempfile.TemporaryDirectory() as tmpdir:
          module.save_to_checkpoint(tmpdir)
          module2 = MultiAgentRLModule.from_checkpoint(tmpdir)
          check(module.get_state(), module2.get_state())
          self.assertEqual(module.keys(), module2.keys())
          self.assertEqual(
              module.keys(), {"test", "test2", DEFAULT_POLICY_ID}
          )
          self.assertNotEqual(id(module), id(module2))
  
      module.remove_module("test")
  
      # check that after removing a module, the checkpoint is correct
      with tempfile.TemporaryDirectory() as tmpdir:
          module.save_to_checkpoint(tmpdir)
          module2 = MultiAgentRLModule.from_checkpoint(tmpdir)
          check(module.get_state(), module2.get_state())
          self.assertEqual(module.keys(), module2.keys())
          self.assertEqual(module.keys(), {"test2", DEFAULT_POLICY_ID})
          self.assertNotEqual(id(module), id(module2))
  
      # check that after adding a new module, the checkpoint is correct
      module.add_module(
          "test3",
          DiscreteBCTorchModule(
              config=RLModuleConfig(
                  env.observation_space,
                  env.action_space,
                  model_config_dict={"fcnet_hiddens": [120]},
              )
          ),
      )
      # check that after adding a module, the checkpoint is correct
      with tempfile.TemporaryDirectory() as tmpdir:
          tmpdir = "/tmp/test_marl_module"
          module.save_to_checkpoint(tmpdir)
          module2 = MultiAgentRLModule.from_checkpoint(tmpdir)
          check(module.get_state(), module2.get_state())
          self.assertEqual(module.keys(), module2.keys())
          self.assertEqual(
              module.keys(), {"test2", "test3", DEFAULT_POLICY_ID}
          )
          self.assertNotEqual(id(module), id(module2))
  
file_path: ray/rllib/core/rl_module/tests/test_marl_module.py
imports: import tempfile
  import unittest
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec, RLModuleConfig
  from ray.rllib.core.rl_module.marl_module import MultiAgentRLModule, MultiAgentRLModuleConfig
  from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule
  from ray.rllib.core.testing.utils import DEFAULT_POLICY_ID
  from ray.rllib.env.multi_agent_env import make_multi_agent
  from ray.rllib.utils.test_utils import check
label_desc: Use image node_id 5 for calling the TestMARLModule obj's underlying member method code with example usage: obj.test_save_to_from_checkpoint() without return types