chore: import upstream snapshot with attribution

rllib/algorithms/tests/test_algorithm_config.py

@@ -0,0 +1,525 @@
+import unittest
+from typing import Type
+
+import gymnasium as gym
+
+import ray
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.ppo import PPO, PPOConfig
+from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner
+from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
+from ray.rllib.core.rl_module.multi_rl_module import (
+    MultiRLModule,
+    MultiRLModuleSpec,
+)
+from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec
+from ray.rllib.utils.test_utils import check
+
+
+class TestAlgorithmConfig(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_running_specific_algo_with_generic_config(self):
+        """Tests, whether some algo can be run with the generic AlgorithmConfig."""
+        config = (
+            AlgorithmConfig(algo_class=PPO)
+            .environment("CartPole-v0")
+            .training(lr=0.12345, train_batch_size=3000, minibatch_size=300)
+        )
+        algo = config.build()
+        self.assertTrue(algo.config.lr == 0.12345)
+        self.assertTrue(algo.config.train_batch_size == 3000)
+        algo.train()
+        algo.stop()
+
+    def test_freezing_of_algo_config(self):
+        """Tests, whether freezing an AlgorithmConfig actually works as expected."""
+        config = (
+            AlgorithmConfig()
+            .environment("CartPole-v0")
+            .training(lr=0.12345, train_batch_size=3000)
+            .multi_agent(
+                policies={
+                    "pol1": (None, None, None, AlgorithmConfig.overrides(lr=0.001))
+                },
+                policy_mapping_fn=lambda agent_id, episode, worker, **kw: "pol1",
+            )
+        )
+        config.freeze()
+
+        def set_lr(config):
+            config.lr = 0.01
+
+        self.assertRaisesRegex(
+            AttributeError,
+            "Cannot set attribute.+of an already frozen AlgorithmConfig",
+            lambda: set_lr(config),
+        )
+
+        # TODO: Figure out, whether we should convert all nested structures into
+        #  frozen ones (set -> frozenset; dict -> frozendict; list -> tuple).
+
+        def set_one_policy(config):
+            config.policies["pol1"] = (None, None, None, {"lr": 0.123})
+
+        # self.assertRaisesRegex(
+        #    AttributeError,
+        #    "Cannot set attribute.+of an already frozen AlgorithmConfig",
+        #    lambda: set_one_policy(config),
+        # )
+
+    def test_rollout_fragment_length(self):
+        """Tests the proper auto-computation of the `rollout_fragment_length`."""
+        config = (
+            AlgorithmConfig()
+            .env_runners(
+                num_env_runners=4,
+                num_envs_per_env_runner=3,
+                rollout_fragment_length="auto",
+            )
+            .training(train_batch_size=2456)
+        )
+        # 2456 / (3 * 4) -> 204.666 -> 204 or 205 (depending on worker index).
+        # Actual train batch size: 2457 (off by only 1).
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=0) == 205)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=1) == 205)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=2) == 205)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=3) == 205)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=4) == 204)
+
+        config = (
+            AlgorithmConfig()
+            .env_runners(
+                num_env_runners=3,
+                num_envs_per_env_runner=2,
+                rollout_fragment_length="auto",
+            )
+            .training(train_batch_size=4000)
+        )
+        # 4000 / 6 -> 666.66 -> 666 or 667 (depending on worker index)
+        # Actual train batch size: 4000 (perfect match)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=0) == 667)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=1) == 667)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=2) == 667)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=3) == 666)
+
+        config = (
+            AlgorithmConfig()
+            .env_runners(
+                num_env_runners=12,
+                rollout_fragment_length="auto",
+            )
+            .training(train_batch_size=1342)
+        )
+        # 1342 / 12 -> 111.83 -> 111 or 112 (depending on worker index)
+        # Actual train batch size: 1342 (perfect match)
+        for i in range(11):
+            self.assertTrue(config.get_rollout_fragment_length(worker_index=i) == 112)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=11) == 111)
+        self.assertTrue(config.get_rollout_fragment_length(worker_index=12) == 111)
+
+    def test_detect_atari_env(self):
+        """Tests that we can properly detect Atari envs."""
+        config = AlgorithmConfig().environment(
+            env="ale_py:ALE/Breakout-v5", env_config={"frameskip": 1}
+        )
+        self.assertTrue(config.is_atari)
+
+        config = AlgorithmConfig().environment(env="ale_py:ALE/Pong-v5")
+        self.assertTrue(config.is_atari)
+
+        config = AlgorithmConfig().environment(env="CartPole-v1")
+        # We do not auto-detect callable env makers for Atari envs.
+        self.assertFalse(config.is_atari)
+
+        config = AlgorithmConfig().environment(
+            env=lambda ctx: gym.make(
+                "ale_py:ALE/Breakout-v5",
+                frameskip=1,
+            )
+        )
+        # We do not auto-detect callable env makers for Atari envs.
+        self.assertFalse(config.is_atari)
+
+        config = AlgorithmConfig().environment(env="NotAtari")
+        self.assertFalse(config.is_atari)
+
+    def test_rl_module_api(self):
+        config = PPOConfig().environment("CartPole-v1").framework("torch")
+
+        self.assertEqual(config.rl_module_spec.module_class, PPOTorchRLModule)
+
+        class A:
+            pass
+
+        config = config.rl_module(rl_module_spec=RLModuleSpec(A))
+        self.assertEqual(config.rl_module_spec.module_class, A)
+
+    def test_config_per_module(self):
+        """Tests, whether per-module config overrides (multi-agent) work as expected."""
+
+        # Compile individual agents' PPO configs from a config object.
+        config = (
+            PPOConfig()
+            .training(kl_coeff=0.5)
+            .multi_agent(
+                policies={"module_1", "module_2", "module_3"},
+                # Override config settings fro `module_1` and `module_2`.
+                algorithm_config_overrides_per_module={
+                    "module_1": PPOConfig.overrides(lr=0.01, kl_coeff=0.1),
+                    "module_2": PPOConfig.overrides(grad_clip=100.0),
+                },
+            )
+        )
+
+        # Check default config.
+        check(config.lr, 0.00005)
+        check(config.grad_clip, None)
+        check(config.grad_clip_by, "global_norm")
+        check(config.kl_coeff, 0.5)
+
+        # `module_1` overrides.
+        config_1 = config.get_config_for_module("module_1")
+        check(config_1.lr, 0.01)
+        check(config_1.grad_clip, None)
+        check(config_1.grad_clip_by, "global_norm")
+        check(config_1.kl_coeff, 0.1)
+
+        # `module_2` overrides.
+        config_2 = config.get_config_for_module("module_2")
+        check(config_2.lr, 0.00005)
+        check(config_2.grad_clip, 100.0)
+        check(config_2.grad_clip_by, "global_norm")
+        check(config_2.kl_coeff, 0.5)
+
+        # No `module_3` overrides (b/c module_3 uses the top-level config
+        # object directly).
+        self.assertTrue("module_3" not in config._per_module_overrides)
+        config_3 = config.get_config_for_module("module_3")
+        self.assertTrue(config_3 is config)
+
+    def test_learner_api(self):
+        config = PPOConfig().environment("CartPole-v1")
+
+        self.assertEqual(config.learner_class, PPOTorchLearner)
+
+    def _assertEqualMARLSpecs(self, spec1, spec2):
+        self.assertEqual(spec1.multi_rl_module_class, spec2.multi_rl_module_class)
+
+        self.assertEqual(set(spec1.module_specs.keys()), set(spec2.module_specs.keys()))
+        for k, module_spec1 in spec1.module_specs.items():
+            module_spec2 = spec2.module_specs[k]
+
+            self.assertEqual(module_spec1.module_class, module_spec2.module_class)
+            self.assertEqual(
+                module_spec1.observation_space, module_spec2.observation_space
+            )
+            self.assertEqual(module_spec1.action_space, module_spec2.action_space)
+            self.assertEqual(
+                module_spec1.model_config_dict, module_spec2.model_config_dict
+            )
+
+    def _get_expected_marl_spec(
+        self,
+        config: AlgorithmConfig,
+        expected_module_class: Type[RLModule],
+        passed_module_class: Type[RLModule] = None,
+        expected_multi_rl_module_class: Type[MultiRLModule] = None,
+    ):
+        """This is a utility function that retrieves the expected marl specs.
+
+        Args:
+            config: The algorithm config.
+            expected_module_class: This is the expected RLModule class that is going to
+                be reference in the RLModuleSpec parts of the MultiLModuleSpec.
+            passed_module_class: This is the RLModule class that is passed into the
+                module_spec argument of get_multi_rl_module_spec. The function is
+                designed so that it will use the passed in module_spec for the
+                RLModuleSpec parts of the MultiRLModuleSpec.
+            expected_multi_rl_module_class: This is the expected MultiRLModule class
+                that is going to be reference in the MultiRLModuleSpec.
+
+        Returns:
+            Tuple of the returned MultiRLModuleSpec from config.
+            get_multi_rl_module_spec() and the expected MultiRLModuleSpec.
+        """
+        from ray.rllib.policy.policy import PolicySpec
+
+        if expected_multi_rl_module_class is None:
+            expected_multi_rl_module_class = MultiRLModule
+
+        env = gym.make("CartPole-v1")
+        policy_spec_ph = PolicySpec(
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            config=AlgorithmConfig(),
+        )
+
+        marl_spec = config.get_multi_rl_module_spec(
+            policy_dict={"p1": policy_spec_ph, "p2": policy_spec_ph},
+            single_agent_rl_module_spec=RLModuleSpec(module_class=passed_module_class)
+            if passed_module_class
+            else None,
+        )
+
+        expected_marl_spec = MultiRLModuleSpec(
+            multi_rl_module_class=expected_multi_rl_module_class,
+            rl_module_specs={
+                "p1": RLModuleSpec(
+                    module_class=expected_module_class,
+                    observation_space=env.observation_space,
+                    action_space=env.action_space,
+                ),
+                "p2": RLModuleSpec(
+                    module_class=expected_module_class,
+                    observation_space=env.observation_space,
+                    action_space=env.action_space,
+                ),
+            },
+        )
+
+        return marl_spec, expected_marl_spec
+
+    def test_get_multi_rl_module_spec(self):
+        """Tests whether the get_multi_rl_module_spec() method works properly."""
+        from ray.rllib.examples.rl_modules.classes.vpg_torch_rlm import VPGTorchRLModule
+
+        class CustomRLModule1(VPGTorchRLModule):
+            pass
+
+        class CustomRLModule2(VPGTorchRLModule):
+            pass
+
+        class CustomRLModule3(VPGTorchRLModule):
+            pass
+
+        class CustomMultiRLModule1(MultiRLModule):
+            pass
+
+        ########################################
+        # single agent
+        class SingleAgentAlgoConfig(AlgorithmConfig):
+            def get_default_rl_module_spec(self):
+                return RLModuleSpec(module_class=VPGTorchRLModule)
+
+        # multi-agent
+        class MultiAgentAlgoConfigWithNoSingleAgentSpec(AlgorithmConfig):
+            def get_default_rl_module_spec(self):
+                return MultiRLModuleSpec(multi_rl_module_class=CustomMultiRLModule1)
+
+        ########################################
+        # This is the simplest case where we have to construct the MultiRLModule based
+        # on the default specs only.
+        config = SingleAgentAlgoConfig().api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+
+        spec, expected = self._get_expected_marl_spec(config, VPGTorchRLModule)
+        self._assertEqualMARLSpecs(spec, expected)
+
+        # expected module should become the passed module if we pass it in.
+        spec, expected = self._get_expected_marl_spec(
+            config, CustomRLModule2, passed_module_class=CustomRLModule2
+        )
+        self._assertEqualMARLSpecs(spec, expected)
+
+        ########################################
+        # This is the case where we pass in a `MultiRLModuleSpec` that asks the
+        # algorithm to assign a specific type of RLModule class to certain module_ids.
+        config = (
+            SingleAgentAlgoConfig()
+            .api_stack(
+                enable_rl_module_and_learner=True,
+                enable_env_runner_and_connector_v2=True,
+            )
+            .rl_module(
+                rl_module_spec=MultiRLModuleSpec(
+                    rl_module_specs={
+                        "p1": RLModuleSpec(module_class=CustomRLModule1),
+                        "p2": RLModuleSpec(module_class=CustomRLModule1),
+                    },
+                ),
+            )
+        )
+
+        spec, expected = self._get_expected_marl_spec(config, CustomRLModule1)
+        self._assertEqualMARLSpecs(spec, expected)
+
+        ########################################
+        # This is the case where we ask the algorithm to assign a specific type of
+        # RLModule class to ALL module_ids.
+        config = (
+            SingleAgentAlgoConfig()
+            .api_stack(
+                enable_rl_module_and_learner=True,
+                enable_env_runner_and_connector_v2=True,
+            )
+            .rl_module(
+                rl_module_spec=RLModuleSpec(module_class=CustomRLModule1),
+            )
+        )
+
+        spec, expected = self._get_expected_marl_spec(config, CustomRLModule1)
+        self._assertEqualMARLSpecs(spec, expected)
+
+        # expected module should become the passed module if we pass it in.
+        spec, expected = self._get_expected_marl_spec(
+            config, CustomRLModule2, passed_module_class=CustomRLModule2
+        )
+        self._assertEqualMARLSpecs(spec, expected)
+
+        ########################################
+        # This is not only assigning a specific type of RLModule class to EACH
+        # module_id, but also defining a new custom MultiRLModule class to be used
+        # in the multi-agent scenario.
+        config = (
+            SingleAgentAlgoConfig()
+            .api_stack(
+                enable_rl_module_and_learner=True,
+                enable_env_runner_and_connector_v2=True,
+            )
+            .rl_module(
+                rl_module_spec=MultiRLModuleSpec(
+                    multi_rl_module_class=CustomMultiRLModule1,
+                    rl_module_specs={
+                        "p1": RLModuleSpec(module_class=CustomRLModule1),
+                        "p2": RLModuleSpec(module_class=CustomRLModule1),
+                    },
+                ),
+            )
+        )
+
+        spec, expected = self._get_expected_marl_spec(
+            config, CustomRLModule1, expected_multi_rl_module_class=CustomMultiRLModule1
+        )
+        self._assertEqualMARLSpecs(spec, expected)
+
+        # This is expected to return CustomRLModule1 instead of CustomRLModule3 which
+        # is passed in. Because the default for p1, p2 is to use CustomRLModule1. The
+        # passed module_spec only sets a default to fall back onto in case the
+        # module_id is not specified in the original MultiRLModuleSpec. Since P1
+        # and P2 are both assigned to CustomeRLModule1, the passed module_spec will not
+        # be used. This is the expected behavior for adding a new modules to a
+        # `MultiRLModule` that is not defined in the original MultiRLModuleSpec.
+        spec, expected = self._get_expected_marl_spec(
+            config,
+            CustomRLModule1,
+            passed_module_class=CustomRLModule3,
+            expected_multi_rl_module_class=CustomMultiRLModule1,
+        )
+        self._assertEqualMARLSpecs(spec, expected)
+
+        ########################################
+        # This is the case where we ask the algorithm to use its default
+        # MultiRLModuleSpec, but the MultiRLModuleSpec has not defined its
+        # RLModuleSpecs.
+        config = MultiAgentAlgoConfigWithNoSingleAgentSpec().api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+
+        self.assertRaisesRegex(
+            ValueError,
+            "Module_specs cannot be None",
+            lambda: config.rl_module_spec,
+        )
+
+    def test_rollout_fragment_length_with_small_batch_and_multiple_learners(self):
+        """Test that get_rollout_fragment_length doesn't return 0 when train_batch_size=1 and num_learners > 1."""
+        for num_env_runners in [1, 2, 3, 4]:
+            config = (
+                AlgorithmConfig()
+                .env_runners(
+                    rollout_fragment_length="auto",
+                    num_env_runners=num_env_runners,
+                )
+                .learners(
+                    num_learners=2
+                )  # Multiple learners with train_batch_size=1 causes the issue
+                .training(
+                    train_batch_size=1
+                )  # Small batch size with multiple learners causes integer division to 0
+            )
+
+            # This should not return 0
+            rollout_fragment_length = config.get_rollout_fragment_length(0)
+            self.assertEqual(
+                rollout_fragment_length,
+                1,
+            )
+
+    def test_to_dict_roundtrip_new_api_stack(self):
+        """Tests that to_dict() round-trips New API stack batch sizes.
+
+        `to_dict()` does NOT eagerly resolve the effective batch size (that stays
+        lazy via the `total_train_batch_size` property). It only serializes the raw
+        fields, which is what makes it safe to call on an as-yet-unresolved config
+        (e.g. one carrying Tune search spaces).
+        """
+        from ray.rllib.algorithms.ppo import PPOConfig
+
+        # 1. Create a config on the New API Stack
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=True,
+                enable_env_runner_and_connector_v2=True,
+            )
+            .training(train_batch_size_per_learner=123)
+        )
+
+        # 2. Export to dictionary
+        config_dict = config.to_dict()
+
+        # to_dict() does not inject computed properties (would break round-trip).
+        self.assertNotIn("total_train_batch_size", config_dict)
+        self.assertNotIn("train_batch_size_per_learner", config_dict)
+
+        # 3. Roundtrip: Create a new config and update from the dictionary, and
+        # verify the per-learner batch size (and the total derived from it) survives.
+        new_config = PPOConfig().update_from_dict(config_dict)
+        self.assertEqual(new_config.train_batch_size_per_learner, 123)
+        self.assertEqual(new_config.total_train_batch_size, 123)
+
+    def test_to_dict_with_tune_search_space(self):
+        """to_dict() must not eagerly resolve batch size when it's a Tune search space.
+
+        Regression test: passing an AlgorithmConfig with a search-space
+        `train_batch_size_per_learner` as Tune's `param_space` calls `to_dict()` on
+        an unresolved config. Computing `total_train_batch_size` (`Domain * int`)
+        would raise TypeError, so `to_dict()` must not attempt it.
+        """
+        from ray import tune
+        from ray.rllib.algorithms.ppo import PPOConfig
+
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=True,
+                enable_env_runner_and_connector_v2=True,
+            )
+            .training(train_batch_size_per_learner=tune.qrandint(256, 2048, 64))
+        )
+
+        # Must not raise (this is the bug: TypeError from `Domain * int`).
+        config_dict = config.to_dict()
+
+        # The unresolved search space survives serialization so Tune can sample it.
+        self.assertIsInstance(
+            config_dict["_train_batch_size_per_learner"], tune.search.sample.Domain
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))