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 cpkt RL module e2e.
docstr_tok: ["Test","if","we","can","train","a","PPO","algorithm","with","a","cpkt","RL","module","e2e","."]
code: def test_e2e_load_rl_module(self):
      """Test if we can train a PPO algorithm with a cpkt RL module e2e."""
      scaling_config = {
          "num_learner_workers": 2,
          "num_gpus_per_learner_worker": 1,
      }
  
      config = (
          PPOConfig()
          .rollouts(rollout_fragment_length=4)
          .environment("CartPole-v1")
          .training(
              num_sgd_iter=1, train_batch_size=8, sgd_minibatch_size=8
          )
          .training(_enable_learner_api=True)
          .resources(**scaling_config)
      )
      env = gym.make("CartPole-v1")
      for fw in framework_iterator(config, frameworks=["tf2", "torch"]):
          # create a marl_module to load and save it to a checkpoint directory
          module_class = PPO_MODULES[fw]
          module_spec = SingleAgentRLModuleSpec(
              module_class=module_class,
              observation_space=env.observation_space,
              action_space=env.action_space,
              model_config_dict={"fcnet_hiddens": [32]},
              catalog_class=PPOCatalog,
          )
          module = module_spec.build()
  
          module_ckpt_path = tempfile.mkdtemp()
          module.save_to_checkpoint(module_ckpt_path)
  
          module_to_load_spec = SingleAgentRLModuleSpec(
              module_class=module_class,
              observation_space=env.observation_space,
              action_space=env.action_space,
              model_config_dict={"fcnet_hiddens": [32]},
              catalog_class=PPOCatalog,
              load_state_path=module_ckpt_path,
          )
  
          config = config.rl_module(
              rl_module_spec=module_to_load_spec,
              _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[DEFAULT_POLICY_ID],
              convert_to_numpy(module.get_state()),
          )
          algo.train()
          algo.stop()
          del algo
          shutil.rmtree(module_ckpt_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 6 for calling the TestE2ERLModuleLoad obj's underlying member method code with example usage: obj.test_e2e_load_rl_module() without return types