wehub-resource-sync
526 lines
20 KiB
Python
526 lines
20 KiB
Python
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__]))
|