class_docstr: Test RLModule loading from rl module spec across a multi node cluster.
class_docstr_tok: ["Test","RLModule","loading","from","rl","module","spec","across","a","multi","node","cluster","."]
docstr: Test if we can train a PPO algorithm with a checkpointed MARL module e2e.
docstr_tok: ["Test","if","we","can","train","a","PPO","algorithm","with","a","checkpointed","MARL","module","e2e","."]
code: def test_e2e_load_simple_marl_module(self):
      """Test if we can train a PPO algorithm with a checkpointed MARL module e2e."""
      config = self.get_ppo_config()
      env = MultiAgentCartPole({"num_agents": NUM_AGENTS})
      for fw in framework_iterator(config, frameworks=["tf2", "torch"]):
          # create a marl_module to load and save it to a checkpoint directory
          module_specs = {}
          module_class = PPO_MODULES[fw]
          for i in range(NUM_AGENTS):
              module_specs[f"policy_{i}"] = SingleAgentRLModuleSpec(
                  module_class=module_class,
                  observation_space=env.observation_space,
                  action_space=env.action_space,
                  model_config_dict={"fcnet_hiddens": [32 * (i + 1)]},
                  catalog_class=PPOCatalog,
              )
          marl_module_spec = MultiAgentRLModuleSpec(
              module_specs=module_specs
          )
          marl_module = marl_module_spec.build()
          marl_module_weights = convert_to_numpy(
              marl_module.get_state()
          )
          marl_checkpoint_path = tempfile.mkdtemp()
          marl_module.save_to_checkpoint(marl_checkpoint_path)
  
          # create a new MARL_spec with the checkpoint from the previous one
          marl_module_spec_from_checkpoint = MultiAgentRLModuleSpec(
              module_specs=module_specs,
              load_state_path=marl_checkpoint_path,
          )
          config = config.rl_module(
              rl_module_spec=marl_module_spec_from_checkpoint,
              _enable_rl_module_api=True,
          )
  
          # create the algorithm with multiple nodes and check if the weights
          # are the same as the original MARL Module
          algo = config.build()
          algo_module_weights = algo.learner_group.get_weights()
          check(algo_module_weights, marl_module_weights)
          algo.train()
          algo.stop()
          del algo
          shutil.rmtree(marl_checkpoint_path)
  
imports: import gymnasium
  import numpy
  import shutil
  import tempfile
  import tree
  import unittest
  import ray
  from ray.rllib.algorithms.ppo import PPOConfig
  from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
  from ray.rllib.algorithms.ppo.tf.ppo_tf_rl_module import PPOTfRLModule
  from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.core.rl_module.marl_module import MultiAgentRLModuleSpec, MultiAgentRLModule
  from ray.rllib.examples.env.multi_agent import MultiAgentCartPole
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils.test_utils import check, framework_iterator
  from ray.rllib.utils.numpy import convert_to_numpy
label_desc: Use image node_id 4 for calling the TestE2ERLModuleLoad obj's underlying member method code with example usage: obj.test_e2e_load_simple_marl_module() without return types