docstr: Check that saving and loading learner group state works.
docstr_tok: ["Check","that","saving","and","loading","learner","group","state","works","."]
code: def test_save_load_state(self):
      """Check that saving and loading learner group state works."""
      fws = ["torch", "tf2"]
      # this is expanded to more scaling modes on the release ci.
      scaling_modes = ["multi-gpu-ddp", "local-cpu"]
      test_iterator = itertools.product(fws, scaling_modes)
      batch = SampleBatch(FAKE_BATCH)
      for fw, scaling_mode in test_iterator:
          print(
              f"Testing framework: {fw}, scaling mode: {scaling_mode}."
          )
          env = gym.make("CartPole-v1")
  
          config_overrides = REMOTE_CONFIGS.get(
              scaling_mode
          ) or LOCAL_CONFIGS.get(scaling_mode)
          config = BaseTestingAlgorithmConfig().update_from_dict(
              config_overrides
          )
          initial_learner_group = config.build_learner_group(env=env)
  
          # checkpoint the initial learner state for later comparison
          initial_learner_checkpoint_dir = (
              tempfile.TemporaryDirectory().name
          )
          initial_learner_group.save_state(
              initial_learner_checkpoint_dir
          )
          initial_learner_group_weights = (
              initial_learner_group.get_weights()
          )
  
          # do a single update
          initial_learner_group.update_from_batch(
              batch=batch.as_multi_agent(), reduce_fn=None
          )
  
          # checkpoint the learner state after 1 update for later comparison
          learner_after_1_update_checkpoint_dir = (
              tempfile.TemporaryDirectory().name
          )
          initial_learner_group.save_state(
              learner_after_1_update_checkpoint_dir
          )
  
          # remove that learner, construct a new one, and load the state of the old
          # learner into the new one
          initial_learner_group.shutdown()
          del initial_learner_group
          new_learner_group = config.build_learner_group(env=env)
          new_learner_group.load_state(
              learner_after_1_update_checkpoint_dir
          )
  
          # do another update
          results_with_break = new_learner_group.update_from_batch(
              batch=batch.as_multi_agent(), reduce_fn=None
          )
          weights_after_1_update_with_break = (
              new_learner_group.get_weights()
          )
          new_learner_group.shutdown()
          del new_learner_group
  
          # construct a new learner group and load the initial state of the learner
          learner_group = config.build_learner_group(env=env)
          learner_group.load_state(initial_learner_checkpoint_dir)
          check(
              learner_group.get_weights(), initial_learner_group_weights
          )
          learner_group.update_from_batch(
              batch.as_multi_agent(), reduce_fn=None
          )
          results_without_break = learner_group.update_from_batch(
              batch=batch.as_multi_agent(), reduce_fn=None
          )
          weights_after_1_update_without_break = (
              learner_group.get_weights()
          )
          learner_group.shutdown()
          del learner_group
  
          # compare the results of the two updates
          check(results_with_break, results_without_break)
          check(
              weights_after_1_update_with_break,
              weights_after_1_update_without_break,
          )
  
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 TestLearnerGroupSaveLoadState obj's underlying member method code with example usage: obj.test_save_load_state() without return types