docstr: Test that module state can be loaded from a checkpoint.
docstr_tok: ["Test","that","module","state","can","be","loaded","from","a","checkpoint","."]
code: def test_load_module_state(self):
      """Test that module state can be loaded from a checkpoint."""
      fws = ["torch", "tf2"]
      # this is expanded to more scaling modes on the release ci.
      scaling_modes = ["local-cpu", "multi-gpu-ddp"]
  
      test_iterator = itertools.product(fws, scaling_modes)
      for fw, scaling_mode in test_iterator:
          print(
              f"Testing framework: {fw}, scaling mode: {scaling_mode}."
          )
          # env will have agent ids 0 and 1
          env = MultiAgentCartPole({"num_agents": 2})
  
          config_overrides = REMOTE_CONFIGS.get(
              scaling_mode
          ) or LOCAL_CONFIGS.get(scaling_mode)
          config = BaseTestingAlgorithmConfig().update_from_dict(
              config_overrides
          )
          learner_group = config.build_learner_group(env=env)
          spec = config.get_marl_module_spec(env=env).module_specs[
              DEFAULT_POLICY_ID
          ]
          learner_group.add_module(module_id="0", module_spec=spec)
          learner_group.add_module(module_id="1", module_spec=spec)
          learner_group.remove_module(DEFAULT_POLICY_ID)
  
          module_0 = spec.build()
          module_1 = spec.build()
          marl_module = MultiAgentRLModule()
          marl_module.add_module(module_id="0", module=module_0)
          marl_module.add_module(module_id="1", module=module_1)
  
          # Check if we can load just the MARL Module
          with tempfile.TemporaryDirectory() as tmpdir:
              marl_module.save_to_checkpoint(tmpdir)
              old_learner_weights = learner_group.get_weights()
              learner_group.load_module_state(
                  marl_module_ckpt_dir=tmpdir
              )
              # check the weights of the module in the learner group are the
              # same as the weights of the newly created marl module
              check(
                  learner_group.get_weights(), marl_module.get_state()
              )
              learner_group.set_weights(old_learner_weights)
  
          # Check if we can load just single agent RL Modules
          with tempfile.TemporaryDirectory() as tmpdir:
              module_0.save_to_checkpoint(tmpdir)
              with tempfile.TemporaryDirectory() as tmpdir2:
                  temp_module = spec.build()
                  temp_module.save_to_checkpoint(tmpdir2)
  
                  old_learner_weights = learner_group.get_weights()
                  learner_group.load_module_state(
                      rl_module_ckpt_dirs={"0": tmpdir, "1": tmpdir2}
                  )
                  # check the weights of the module in the learner group are the
                  # same as the weights of the newly created marl module
                  new_marl_module = MultiAgentRLModule()
                  new_marl_module.add_module(
                      module_id="0", module=module_0
                  )
                  new_marl_module.add_module(
                      module_id="1", module=temp_module
                  )
                  check(
                      learner_group.get_weights(),
                      new_marl_module.get_state(),
                  )
                  learner_group.set_weights(old_learner_weights)
  
          # Check if we can load a MARL Module and single agent RL Module
          # at the same time. Check that the single agent RL Module is loaded
          # over the sub module in the MARL Module
          with tempfile.TemporaryDirectory() as tmpdir:
              module_0 = spec.build()
              marl_module = MultiAgentRLModule()
              marl_module.add_module(module_id="0", module=module_0)
              marl_module.add_module(module_id="1", module=spec.build())
              marl_module.save_to_checkpoint(tmpdir)
              with tempfile.TemporaryDirectory() as tmpdir2:
                  module_1 = spec.build()
                  module_1.save_to_checkpoint(tmpdir2)
                  learner_group.load_module_state(
                      marl_module_ckpt_dir=tmpdir,
                      rl_module_ckpt_dirs={"1": tmpdir2},
                  )
                  new_marl_module = MultiAgentRLModule()
                  new_marl_module.add_module(
                      module_id="0", module=module_0
                  )
                  new_marl_module.add_module(
                      module_id="1", module=module_1
                  )
                  check(
                      learner_group.get_weights(),
                      new_marl_module.get_state(),
                  )
          del learner_group
  
imports: import gymnasium
  import itertools
  import numpy
  from typing import Any, Dict, List
  import tempfile
  import unittest
  import ray
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.algorithms.ppo.tests.test_ppo_learner import FAKE_BATCH
  from ray.rllib.core.learner.learner import Learner
  from ray.rllib.core.rl_module.marl_module import MultiAgentRLModule
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.core.testing.tf.bc_learner import BCTfLearner
  from ray.rllib.core.testing.tf.bc_module import DiscreteBCTFModule
  from ray.rllib.core.testing.utils import add_module_to_learner_or_learner_group
  from ray.rllib.core.testing.testing_learner import BaseTestingAlgorithmConfig
  from ray.rllib.examples.env.multi_agent import MultiAgentCartPole
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch, MultiAgentBatch
  from ray.rllib.utils.test_utils import check, get_cartpole_dataset_reader
  from ray.rllib.utils.metrics import ALL_MODULES
  from ray.util.timer import _Timer
label_desc: Use image node_id 3 for calling the TestLearnerGroupCheckpointRestore obj's underlying member method code with example usage: obj.test_load_module_state() without return types