chore: import upstream snapshot with attribution

rllib/algorithms/tests/run_regression_tests.py

@@ -0,0 +1,353 @@
+#!/usr/bin/env python
+
+# @OldAPIStack
+
+import argparse
+import importlib
+import json
+import os
+import re
+import sys
+import uuid
+from pathlib import Path
+
+import yaml
+
+import ray
+from ray import air
+from ray._common.deprecation import deprecation_warning
+from ray.air.integrations.wandb import WandbLoggerCallback
+from ray.rllib import _register_all
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune import run_experiments
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--framework",
+    type=str,
+    choices=["torch", "tf2", "tf"],
+    default=None,
+    help="The deep learning framework to use. If not provided, try using the one "
+    "specified in the file, otherwise, use RLlib's default: `torch`.",
+)
+parser.add_argument(
+    "--dir",
+    type=str,
+    required=True,
+    help="The directory or file in which to find all tests.",
+)
+parser.add_argument(
+    "--env",
+    type=str,
+    default=None,
+    help="An optional env override setting. If not provided, try using the one "
+    "specified in the file.",
+)
+parser.add_argument("--num-cpus", type=int, default=None)
+parser.add_argument(
+    "--local-mode",
+    action="store_true",
+    help=argparse.SUPPRESS,  # Deprecated.
+)
+parser.add_argument(
+    "--num-samples",
+    type=int,
+    default=1,
+    help="The number of seeds/samples to run with the given experiment config.",
+)
+parser.add_argument(
+    "--override-mean-reward",
+    type=float,
+    default=0.0,
+    help=(
+        "Override the mean reward specified by the yaml file in the stopping criteria. "
+        "This is particularly useful for timed tests."
+    ),
+)
+parser.add_argument(
+    "--verbose",
+    type=int,
+    default=2,
+    help="The verbosity level for the main `tune.run_experiments()` call.",
+)
+parser.add_argument(
+    "--wandb-key",
+    type=str,
+    default=None,
+    help="The WandB API key to use for uploading results.",
+)
+parser.add_argument(
+    "--wandb-project",
+    type=str,
+    default=None,
+    help="The WandB project name to use.",
+)
+parser.add_argument(
+    "--wandb-run-name",
+    type=str,
+    default=None,
+    help="The WandB run name to use.",
+)
+parser.add_argument(
+    "--checkpoint-freq",
+    type=int,
+    default=0,
+    help=(
+        "The frequency (in training iterations) with which to create checkpoints. "
+        "Note that if --wandb-key is provided, these checkpoints will automatically "
+        "be uploaded to WandB."
+    ),
+)
+
+# Obsoleted arg, use --dir instead.
+parser.add_argument("--yaml-dir", type=str, default="")
+
+
+def _load_experiments_from_file(
+    config_file: str,
+    file_type: str,
+    stop=None,
+    checkpoint_config=None,
+) -> dict:
+    # Yaml file.
+    if file_type == "yaml":
+        with open(config_file) as f:
+            experiments = yaml.safe_load(f)
+            if stop is not None and stop != "{}":
+                raise ValueError("`stop` criteria only supported for python files.")
+        # Make sure yaml experiments are always old API stack.
+        for experiment in experiments.values():
+            experiment["config"]["enable_rl_module_and_learner"] = False
+            experiment["config"]["enable_env_runner_and_connector_v2"] = False
+    # Python file case (ensured by file type enum)
+    else:
+        module_name = os.path.basename(config_file).replace(".py", "")
+        spec = importlib.util.spec_from_file_location(module_name, config_file)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+
+        if not hasattr(module, "config"):
+            raise ValueError(
+                "Your Python file must contain a 'config' variable "
+                "that is an AlgorithmConfig object."
+            )
+        algo_config = module.config
+        if stop is None:
+            stop = getattr(module, "stop", {})
+        else:
+            stop = json.loads(stop)
+
+        # Note: we do this gymnastics to support the old format that
+        # "_run_rllib_experiments" expects. Ideally, we'd just build the config and
+        # run the algo.
+        config = algo_config.to_dict()
+        experiments = {
+            f"default_{uuid.uuid4().hex}": {
+                "run": algo_config.algo_class,
+                "env": config.get("env"),
+                "config": config,
+                "stop": stop,
+            }
+        }
+
+    for key, val in experiments.items():
+        experiments[key]["checkpoint_config"] = checkpoint_config or {}
+
+    return experiments
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.yaml_dir != "":
+        deprecation_warning(old="--yaml-dir", new="--dir", error=True)
+
+    # Bazel regression test mode: Get path to look for yaml files.
+    # Get the path or single file to use.
+    rllib_dir = Path(__file__).parent.parent
+    print(f"rllib dir={rllib_dir}")
+
+    abs_path = os.path.join(rllib_dir, args.dir)
+    # Single file given.
+    if os.path.isfile(abs_path):
+        files = [abs_path]
+    # Path given -> Get all yaml files in there via rglob.
+    elif os.path.isdir(abs_path):
+        files = []
+        for type_ in ["yaml", "yml", "py"]:
+            files += list(rllib_dir.rglob(args.dir + f"/*.{type_}"))
+        files = sorted(map(lambda path: str(path.absolute()), files), reverse=True)
+    # Given path/file does not exist.
+    else:
+        raise ValueError(f"--dir ({args.dir}) not found!")
+
+    print("Will run the following regression tests:")
+    for file in files:
+        print("->", file)
+
+    # Loop through all collected files.
+    for file in files:
+        config_is_python = False
+        # For python files, need to make sure, we only deliver the module name into the
+        # `_load_experiments_from_file` function (everything from "/ray/rllib" on).
+        if file.endswith(".py"):
+            if file.endswith("__init__.py"):  # weird CI learning test (BAZEL) case
+                continue
+            experiments = _load_experiments_from_file(file, "py")
+            config_is_python = True
+        else:
+            experiments = _load_experiments_from_file(file, "yaml")
+
+        assert (
+            len(experiments) == 1
+        ), "Error, can only run a single experiment per file!"
+
+        exp = list(experiments.values())[0]
+        exp_name = list(experiments.keys())[0]
+
+        # Set the number of samples to run.
+        exp["num_samples"] = args.num_samples
+
+        # Make sure there is a config and a stopping criterium.
+        exp["config"] = exp.get("config", {})
+        exp["stop"] = exp.get("stop", {})
+
+        # Override framework setting with the command line one, if provided.
+        # Otherwise, will use framework setting in file (or default: torch).
+        if args.framework is not None:
+            exp["config"]["framework"] = args.framework
+        # Override env setting if given on command line.
+        if args.env is not None:
+            exp["config"]["env"] = args.env
+        else:
+            exp["config"]["env"] = exp["env"]
+
+        # Override the mean reward if specified. This is used by the ray ci
+        # for overriding the episode reward mean for tf2 tests for off policy
+        # long learning tests such as sac and ddpg on the pendulum environment.
+        if args.override_mean_reward != 0.0:
+            exp["stop"][
+                f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+            ] = args.override_mean_reward
+
+        # Checkpoint settings.
+        exp["checkpoint_config"] = air.CheckpointConfig(
+            checkpoint_frequency=args.checkpoint_freq,
+            checkpoint_at_end=args.checkpoint_freq > 0,
+        )
+
+        # Always run with eager-tracing when framework=tf2, if not in local-mode
+        # and unless the yaml explicitly tells us to disable eager tracing.
+        if (
+            (args.framework == "tf2" or exp["config"].get("framework") == "tf2")
+            # Note: This check will always fail for python configs, b/c normally,
+            # algorithm configs have `self.eager_tracing=False` by default.
+            # Thus, you'd have to set `eager_tracing` to True explicitly in your python
+            # config to make sure we are indeed using eager tracing.
+            and exp["config"].get("eager_tracing") is not False
+        ):
+            exp["config"]["eager_tracing"] = True
+
+        # Print out the actual config (not for py files as yaml.dump weirdly fails).
+        if not config_is_python:
+            print("== Test config ==")
+            print(yaml.dump(experiments))
+
+        callbacks = None
+        if args.wandb_key is not None:
+            project = args.wandb_project or (
+                exp["run"].lower()
+                + "-"
+                + re.sub("\\W+", "-", exp["config"]["env"].lower())
+                if config_is_python
+                else list(experiments.keys())[0]
+            )
+            callbacks = [
+                WandbLoggerCallback(
+                    api_key=args.wandb_key,
+                    project=project,
+                    upload_checkpoints=True,
+                    **({"name": args.wandb_run_name} if args.wandb_run_name else {}),
+                )
+            ]
+
+        if args.local_mode:
+            raise ValueError("`--local-mode` is no longer supported.")
+
+        # Try running each test 3 times and make sure it reaches the given
+        # reward.
+        passed = False
+        for i in range(3):
+            # Try starting a new ray cluster.
+            try:
+                ray.init(num_cpus=args.num_cpus)
+            # Allow running this script on existing cluster as well.
+            except ConnectionError:
+                ray.init()
+            else:
+                try:
+                    trials = run_experiments(
+                        experiments,
+                        resume=False,
+                        verbose=args.verbose,
+                        callbacks=callbacks,
+                    )
+                finally:
+                    ray.shutdown()
+                    _register_all()
+
+            for t in trials:
+                # If we have evaluation workers, use their rewards.
+                # This is useful for offline learning tests, where
+                # we evaluate against an actual environment.
+                check_eval = bool(exp["config"].get("evaluation_interval"))
+                reward_mean = (
+                    t.last_result[EVALUATION_RESULTS][ENV_RUNNER_RESULTS][
+                        EPISODE_RETURN_MEAN
+                    ]
+                    if check_eval
+                    else (
+                        # Some algos don't store sampler results under `env_runners`
+                        # e.g. ARS. Need to keep this logic around for now.
+                        t.last_result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
+                        if ENV_RUNNER_RESULTS in t.last_result
+                        else t.last_result[EPISODE_RETURN_MEAN]
+                    )
+                )
+
+                # If we are using evaluation workers, we may have
+                # a stopping criterion under the "evaluation/" scope. If
+                # not, use `episode_return_mean`.
+                if check_eval:
+                    min_reward = t.stopping_criterion.get(
+                        f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/"
+                        f"{EPISODE_RETURN_MEAN}",
+                        t.stopping_criterion.get(
+                            f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+                        ),
+                    )
+                # Otherwise, expect `env_runners/episode_return_mean` to be set.
+                else:
+                    min_reward = t.stopping_criterion.get(
+                        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+                    )
+
+                # If min reward not defined, always pass.
+                if min_reward is None or reward_mean >= min_reward:
+                    passed = True
+                    break
+
+            if passed:
+                print("Regression test PASSED")
+                break
+            else:
+                print("Regression test FAILED on attempt {}".format(i + 1))
+
+        if not passed:
+            print("Overall regression FAILED: Exiting with Error.")
+            sys.exit(1)

rllib/algorithms/tests/test_algorithm.py

@@ -0,0 +1,685 @@
+import os
+import time
+import unittest
+from pathlib import Path
+from random import choice
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+import ray.rllib.algorithms.dqn as dqn
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.examples.evaluation.evaluation_parallel_to_training import (
+    AssertEvalCallback,
+)
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.framework import convert_to_tensor
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    LEARNER_RESULTS,
+)
+from ray.rllib.utils.metrics.learner_info import LEARNER_INFO
+from ray.tune import register_env
+
+
+class TestAlgorithm(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+        register_env("multi_cart", lambda cfg: MultiAgentCartPole(cfg))
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_add_module_and_remove_module(self):
+        config = (
+            ppo.PPOConfig()
+            .environment(
+                env="multi_cart",
+                env_config={"num_agents": 4},
+            )
+            .env_runners(num_cpus_per_env_runner=0.1)
+            .training(
+                train_batch_size=100,
+                minibatch_size=50,
+                num_epochs=1,
+            )
+            .rl_module(
+                model_config=DefaultModelConfig(
+                    fcnet_hiddens=[5], fcnet_activation="linear"
+                ),
+            )
+            .multi_agent(
+                # Start with a single policy.
+                policies={"p0"},
+                policy_mapping_fn=lambda *a, **kw: "p0",
+                # TODO (sven): Support object store caching on new API stack.
+                # # And only two policies that can be stored in memory at a
+                # # time.
+                # policy_map_capacity=2,
+            )
+            .evaluation(
+                evaluation_num_env_runners=1,
+                evaluation_config=ppo.PPOConfig.overrides(num_cpus_per_env_runner=0.1),
+            )
+        )
+
+        # Construct the Algorithm with a single policy in it.
+        algo = config.build()
+        mod0 = algo.get_module("p0")
+        r = algo.train()
+        self.assertTrue("p0" in r[LEARNER_RESULTS])
+        for i in range(1, 3):
+
+            def new_mapping_fn(agent_id, episode, i=i, **kwargs):
+                return f"p{choice([i, i - 1])}"
+
+            # Add a new RLModule by class (and options).
+            mid = f"p{i}"
+            print(f"Adding new RLModule {mid} ...")
+            new_marl_spec = algo.add_module(
+                module_id=mid,
+                module_spec=RLModuleSpec.from_module(mod0),
+                # Test changing the mapping fn.
+                new_agent_to_module_mapping_fn=new_mapping_fn,
+                # Change the list of modules to train.
+                new_should_module_be_updated=[f"p{i}", f"p{i-1}"],
+            )
+            new_module = algo.get_module(mid)
+            self._assert_modules_added(
+                algo=algo,
+                marl_spec=new_marl_spec,
+                mids=[0, i],
+                trainable=[i, i - 1],
+                mapped=[i, i - 1],
+                not_mapped=[i - 2],
+            )
+
+            # Assert new policy is part of local worker (eval worker set does NOT
+            # have a local worker, only the main EnvRunnerGroup does).
+            multi_rl_module = algo.env_runner.module
+            self.assertTrue(new_module is not mod0)
+            for j in range(i + 1):
+                self.assertTrue(f"p{j}" in multi_rl_module)
+            self.assertTrue(len(multi_rl_module) == i + 1)
+            algo.train()
+            checkpoint = algo.save_to_path()
+
+            # Test restoring from the checkpoint (which has more policies
+            # than what's defined in the config dict).
+            test = Algorithm.from_checkpoint(checkpoint)
+            self._assert_modules_added(
+                algo=test,
+                marl_spec=None,
+                mids=[0, i - 1, i],
+                trainable=[i - 1, i],
+                mapped=[i - 1, i],
+                not_mapped=[i - 2],
+            )
+            # Make sure algorithm can continue training the restored policy.
+            test.train()
+            # Test creating an inference action with the added (and restored) RLModule.
+            mod0 = test.get_module("p0")
+            out = mod0.forward_inference(
+                {
+                    Columns.OBS: convert_to_tensor(
+                        np.expand_dims(mod0.config.observation_space.sample(), 0),
+                        framework=mod0.framework,
+                    ),
+                },
+            )
+            action_dist_inputs = out[Columns.ACTION_DIST_INPUTS]
+            self.assertTrue(action_dist_inputs.shape == (1, 2))
+            test.stop()
+
+            # After having added 2 Modules, try to restore the Algorithm,
+            # but only with 1 of the originally added Modules (plus the initial
+            # p0).
+            if i == 2:
+
+                def new_mapping_fn(agent_id, episode, **kwargs):
+                    return f"p{choice([0, 2])}"
+
+                test2 = Algorithm.from_checkpoint(path=checkpoint)
+                test2.remove_module(
+                    module_id="p1",
+                    new_agent_to_module_mapping_fn=new_mapping_fn,
+                    new_should_module_be_updated=["p0"],
+                )
+                self._assert_modules_added(
+                    algo=test2,
+                    marl_spec=None,
+                    mids=[0, 2],
+                    trainable=[0],
+                    mapped=[0, 2],
+                    not_mapped=[1, 4, 5, 6],
+                )
+                # Make sure algorithm can continue training the restored policy.
+                mod2 = test2.get_module("p2")
+                test2.train()
+                # Test creating an inference action with the added (and restored)
+                # RLModule.
+                out = mod2.forward_exploration(
+                    {
+                        Columns.OBS: convert_to_tensor(
+                            np.expand_dims(mod0.config.observation_space.sample(), 0),
+                            framework=mod0.framework,
+                        ),
+                    },
+                )
+                action_dist_inputs = out[Columns.ACTION_DIST_INPUTS]
+                self.assertTrue(action_dist_inputs.shape == (1, 2))
+                test2.stop()
+
+        # Delete all added modules again from Algorithm.
+        for i in range(2, 0, -1):
+            mid = f"p{i}"
+            marl_spec = algo.remove_module(
+                mid,
+                # Note that the complete signature of a policy_mapping_fn
+                # is: `agent_id, episode, worker, **kwargs`.
+                new_agent_to_module_mapping_fn=(
+                    lambda agent_id, episode, i=i, **kwargs: f"p{i - 1}"
+                ),
+                # Update list of policies to train.
+                new_should_module_be_updated=[f"p{i - 1}"],
+            )
+            self._assert_modules_added(
+                algo=algo,
+                marl_spec=marl_spec,
+                mids=[0, i - 1],
+                trainable=[i - 1],
+                mapped=[i - 1],
+                not_mapped=[i, i + 1],
+            )
+
+        algo.stop()
+
+    @OldAPIStack
+    def test_add_policy_and_remove_policy(self):
+        config = (
+            ppo.PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .environment(
+                env=MultiAgentCartPole,
+                env_config={
+                    "config": {
+                        "num_agents": 4,
+                    },
+                },
+            )
+            .env_runners(num_cpus_per_env_runner=0.1)
+            .training(
+                train_batch_size=100,
+                minibatch_size=50,
+                num_epochs=1,
+                model={
+                    "fcnet_hiddens": [5],
+                    "fcnet_activation": "linear",
+                },
+            )
+            .multi_agent(
+                # Start with a single policy.
+                policies={"p0"},
+                policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: "p0",
+                # And only two policies that can be stored in memory at a
+                # time.
+                policy_map_capacity=2,
+            )
+            .evaluation(
+                evaluation_num_env_runners=1,
+                evaluation_config=ppo.PPOConfig.overrides(num_cpus_per_env_runner=0.1),
+            )
+        )
+
+        obs_space = gym.spaces.Box(-2.0, 2.0, (4,))
+        act_space = gym.spaces.Discrete(2)
+
+        # Pre-generate a policy instance to test adding these directly to an
+        # existing algorithm.
+        policy_obj = ppo.PPOTorchPolicy(obs_space, act_space, config.to_dict())
+
+        # Construct the Algorithm with a single policy in it.
+        algo = config.build()
+        pol0 = algo.get_policy("p0")
+        r = algo.train()
+        self.assertTrue("p0" in r["info"][LEARNER_INFO])
+        for i in range(1, 3):
+
+            def new_mapping_fn(agent_id, episode, worker, i=i, **kwargs):
+                return f"p{choice([i, i - 1])}"
+
+            # Add a new policy either by class (and options) or by instance.
+            pid = f"p{i}"
+            print(f"Adding policy {pid} ...")
+            # By (already instantiated) instance.
+            if i == 2:
+                new_pol = algo.add_policy(
+                    pid,
+                    # Pass in an already existing policy instance.
+                    policy=policy_obj,
+                    # Test changing the mapping fn.
+                    policy_mapping_fn=new_mapping_fn,
+                    # Change the list of policies to train.
+                    policies_to_train=[f"p{i}", f"p{i - 1}"],
+                )
+            # By class (and options).
+            else:
+                new_pol = algo.add_policy(
+                    pid,
+                    algo.get_default_policy_class(config),
+                    observation_space=obs_space,
+                    action_space=act_space,
+                    # Test changing the mapping fn.
+                    policy_mapping_fn=new_mapping_fn,
+                    # Change the list of policies to train.
+                    policies_to_train=[f"p{i}", f"p{i-1}"],
+                )
+
+            # Make sure new policy is part of remote workers in the
+            # worker set and the eval worker set.
+            self.assertTrue(
+                all(
+                    algo.env_runner_group.foreach_env_runner(
+                        func=lambda w, pid=pid: pid in w.policy_map
+                    )
+                )
+            )
+            self.assertTrue(
+                all(
+                    algo.eval_env_runner_group.foreach_env_runner(
+                        func=lambda w, pid=pid: pid in w.policy_map
+                    )
+                )
+            )
+
+            # Assert new policy is part of local worker (eval worker set does NOT
+            # have a local worker, only the main EnvRunnerGroup does).
+            pol_map = algo.env_runner.policy_map
+            self.assertTrue(new_pol is not pol0)
+            for j in range(i + 1):
+                self.assertTrue(f"p{j}" in pol_map)
+            self.assertTrue(len(pol_map) == i + 1)
+            algo.train()
+            checkpoint = algo.save().checkpoint
+
+            # Test restoring from the checkpoint (which has more policies
+            # than what's defined in the config dict).
+            test = ppo.PPO.from_checkpoint(checkpoint)
+
+            # Make sure evaluation worker also got the restored, added policy.
+            def _has_policies(w, pid=pid):
+                return w.get_policy("p0") is not None and w.get_policy(pid) is not None
+
+            self.assertTrue(
+                all(test.eval_env_runner_group.foreach_env_runner(_has_policies))
+            )
+
+            # Make sure algorithm can continue training the restored policy.
+            pol0 = test.get_policy("p0")
+            test.train()
+            # Test creating an action with the added (and restored) policy.
+            a = test.compute_single_action(
+                np.zeros_like(pol0.observation_space.sample()), policy_id=pid
+            )
+            self.assertTrue(pol0.action_space.contains(a))
+            test.stop()
+
+            # After having added 2 policies, try to restore the Algorithm,
+            # but only with 1 of the originally added policies (plus the initial
+            # p0).
+            if i == 2:
+
+                def new_mapping_fn(agent_id, episode, worker, **kwargs):
+                    return f"p{choice([0, 2])}"
+
+                test2 = ppo.PPO.from_checkpoint(
+                    path=checkpoint,
+                    policy_ids=["p0", "p2"],
+                    policy_mapping_fn=new_mapping_fn,
+                    policies_to_train=["p0"],
+                )
+
+                # Make sure evaluation workers have the same policies.
+                def _has_policies(w):
+                    return (
+                        w.get_policy("p0") is not None
+                        and w.get_policy("p2") is not None
+                        and w.get_policy("p1") is None
+                    )
+
+                self.assertTrue(
+                    all(test2.eval_env_runner_group.foreach_env_runner(_has_policies))
+                )
+
+                # Make sure algorithm can continue training the restored policy.
+                pol2 = test2.get_policy("p2")
+                test2.train()
+                # Test creating an action with the added (and restored) policy.
+                a = test2.compute_single_action(
+                    np.zeros_like(pol2.observation_space.sample()), policy_id=pid
+                )
+                self.assertTrue(pol2.action_space.contains(a))
+                test2.stop()
+
+        # Delete all added policies again from Algorithm.
+        for i in range(2, 0, -1):
+            pid = f"p{i}"
+            algo.remove_policy(
+                pid,
+                # Note that the complete signature of a policy_mapping_fn
+                # is: `agent_id, episode, worker, **kwargs`.
+                policy_mapping_fn=(
+                    lambda agent_id, episode, worker, i=i, **kwargs: f"p{i - 1}"
+                ),
+                # Update list of policies to train.
+                policies_to_train=[f"p{i - 1}"],
+            )
+            # Make sure removed policy is no longer part of remote workers in the
+            # worker set and the eval worker set.
+            self.assertTrue(
+                algo.env_runner_group.foreach_env_runner(
+                    func=lambda w, pid=pid: pid not in w.policy_map
+                )[0]
+            )
+            self.assertTrue(
+                algo.eval_env_runner_group.foreach_env_runner(
+                    func=lambda w, pid=pid: pid not in w.policy_map
+                )[0]
+            )
+            # Assert removed policy is no longer part of local worker
+            # (eval worker set does NOT have a local worker, only the main
+            # EnvRunnerGroup does).
+            pol_map = algo.env_runner.policy_map
+            self.assertTrue(pid not in pol_map)
+            self.assertTrue(len(pol_map) == i)
+
+        algo.stop()
+
+    def test_evaluation_option(self):
+        # Use a custom callback that asserts that we are running the
+        # configured exact number of episodes per evaluation.
+        config = (
+            dqn.DQNConfig()
+            .environment(env="CartPole-v1")
+            .evaluation(
+                evaluation_interval=2,
+                evaluation_duration=2,
+                evaluation_duration_unit="episodes",
+                evaluation_config=dqn.DQNConfig.overrides(gamma=0.98),
+            )
+            .callbacks(callbacks_class=AssertEvalCallback)
+        )
+
+        algo = config.build()
+        # Given evaluation_interval=2, r0, r2 should not contain
+        # evaluation metrics, while r1, r3 should.
+        r0 = algo.train()
+        print(r0)
+        r1 = algo.train()
+        print(r1)
+        r2 = algo.train()
+        print(r2)
+        r3 = algo.train()
+        print(r3)
+        algo.stop()
+
+        # No eval results yet in first iteration (eval has not run yet).
+        self.assertFalse(EVALUATION_RESULTS in r0)
+        self.assertTrue(EVALUATION_RESULTS in r1)
+        self.assertTrue(EVALUATION_RESULTS in r2)
+        self.assertTrue(EVALUATION_RESULTS in r3)
+        self.assertTrue(ENV_RUNNER_RESULTS in r1[EVALUATION_RESULTS])
+        self.assertTrue(
+            EPISODE_RETURN_MEAN in r1[EVALUATION_RESULTS][ENV_RUNNER_RESULTS]
+        )
+        self.assertNotEqual(r1[EVALUATION_RESULTS], r3[EVALUATION_RESULTS])
+
+    def test_evaluation_option_always_attach_eval_metrics(self):
+        # Use a custom callback that asserts that we are running the
+        # configured exact number of episodes per evaluation.
+        config = (
+            dqn.DQNConfig()
+            .environment("CartPole-v1")
+            .evaluation(
+                evaluation_interval=2,
+                evaluation_duration=2,
+                evaluation_duration_unit="episodes",
+                evaluation_config=dqn.DQNConfig.overrides(gamma=0.98),
+            )
+            .reporting(min_sample_timesteps_per_iteration=100)
+            .callbacks(callbacks_class=AssertEvalCallback)
+        )
+        algo = config.build()
+        # Should only see eval results, when eval actually ran.
+        r0 = algo.train()
+        r1 = algo.train()
+        r2 = algo.train()
+        r3 = algo.train()
+        algo.stop()
+
+        # Eval results are not available at step 0.
+        self.assertTrue(EVALUATION_RESULTS not in r0)
+        # But step 3 should still have it, even though no eval was
+        # run during that step (b/c the new API stack always attaches eval
+        # results, after the very first evaluation).
+        self.assertTrue(EVALUATION_RESULTS in r1)
+        self.assertTrue(EVALUATION_RESULTS in r2)
+        self.assertTrue(EVALUATION_RESULTS in r3)
+
+    def test_evaluation_wo_eval_env_runner_group(self):
+        # Use a custom callback that asserts that we are running the
+        # configured exact number of episodes per evaluation.
+        config = (
+            ppo.PPOConfig()
+            .environment(env="CartPole-v1")
+            .callbacks(callbacks_class=AssertEvalCallback)
+        )
+
+        # Setup algorithm w/o evaluation worker set and still call
+        # evaluate() -> Expect error.
+        algo_wo_env_on_local_worker = config.build()
+        self.assertRaisesRegex(
+            ValueError,
+            "doesn't have an env!",
+            algo_wo_env_on_local_worker.evaluate,
+        )
+        algo_wo_env_on_local_worker.stop()
+
+        # Try again using `create_local_env_runner=True`.
+        # This force-adds the env on the local-worker, so this Algorithm
+        # can `evaluate` even though it doesn't have an evaluation-worker
+        # set.
+        config.create_env_on_local_worker = True
+        algo_w_env_on_local_worker = config.build()
+        results = algo_w_env_on_local_worker.evaluate()
+        assert (
+            ENV_RUNNER_RESULTS in results
+            and EPISODE_RETURN_MEAN in results[ENV_RUNNER_RESULTS]
+        )
+        algo_w_env_on_local_worker.stop()
+
+    def test_no_env_but_eval_workers_do_have_env(self):
+        """Tests whether no env on workers, but env on eval workers works ok."""
+        script_path = Path(__file__)
+        input_file = os.path.join(
+            script_path.parent.parent.parent, "offline/tests/data/cartpole/small.json"
+        )
+
+        env = gym.make("CartPole-v1")
+
+        offline_rl_config = (
+            BCConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment(
+                observation_space=env.observation_space,
+                action_space=env.action_space,
+            )
+            .evaluation(
+                evaluation_interval=1,
+                evaluation_num_env_runners=1,
+                evaluation_config=BCConfig.overrides(
+                    env="CartPole-v1",
+                    input_="sampler",
+                    observation_space=None,  # Test, whether this is inferred.
+                    action_space=None,  # Test, whether this is inferred.
+                ),
+            )
+            .offline_data(input_=[input_file])
+        )
+
+        bc = offline_rl_config.build()
+        bc.train()
+        bc.stop()
+
+    def test_counters_after_checkpoint(self):
+        # We expect algorithm to no start counters from zero after loading a
+        # checkpoint on a fresh Algorithm instance
+        config = (
+            ppo.PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment(env="CartPole-v1")
+        )
+        algo = config.build()
+
+        self.assertTrue(all(c == 0 for c in algo._counters.values()))
+        algo.step()
+        self.assertTrue((all(c != 0 for c in algo._counters.values())))
+        counter_values = list(algo._counters.values())
+        state = algo.__getstate__()
+        algo.stop()
+
+        algo2 = config.build()
+        self.assertTrue(all(c == 0 for c in algo2._counters.values()))
+        algo2.__setstate__(state)
+        counter_values2 = list(algo2._counters.values())
+        self.assertEqual(counter_values, counter_values2)
+
+    def _assert_modules_added(
+        self,
+        *,
+        algo,
+        marl_spec,
+        mids,
+        trainable,
+        mapped,
+        not_mapped,
+    ):
+        # Make sure Learner has the correct `should_module_be_updated` list.
+        self.assertEqual(
+            set(algo.learner_group._learner.config.policies_to_train),
+            {f"p{i}" for i in trainable},
+        )
+        # Make sure mids are all in marl_spec.
+        if marl_spec is not None:
+            self.assertTrue(all(f"p{m}" in marl_spec for m in mids))
+        # Make sure module is part of remote EnvRunners in the
+        # EnvRunnerGroup and the eval EnvRunnerGroup.
+        self.assertTrue(
+            all(
+                algo.env_runner_group.foreach_env_runner(
+                    lambda w, mids=mids: all(f"p{i}" in w.module for i in mids)
+                )
+            )
+        )
+        self.assertTrue(
+            all(
+                algo.eval_env_runner_group.foreach_env_runner(
+                    lambda w, mids=mids: all(f"p{i}" in w.module for i in mids)
+                )
+            )
+        )
+        # Make sure that EnvRunners have received the correct mapping fn.
+        mapped_pols = [
+            algo.env_runner.config.policy_mapping_fn(0, None) for _ in range(100)
+        ]
+        self.assertTrue(all(f"p{i}" in mapped_pols for i in mapped))
+        self.assertTrue(not any(f"p{i}" in mapped_pols for i in not_mapped))
+
+    def test_evaluation_in_parallel_to_training(self):
+        SECONDS_TO_SLEEP = 2
+
+        class SluggishEnv(gym.Env):
+            def __init__(self, config):
+                self.action_space = gym.spaces.Discrete(2)
+                self.observation_space = gym.spaces.Box(-1, 1, dtype=np.float32)
+
+            def step(self, action):
+                time.sleep(SECONDS_TO_SLEEP)
+                return self.observation_space.sample(), 1, True, False, {}
+
+            def reset(self, *, seed=None, options=None):
+                super().reset(seed=seed)
+                return self.observation_space.sample(), {}
+
+        config = (
+            ppo.PPOConfig()
+            .environment(env=SluggishEnv)
+            .evaluation(
+                evaluation_parallel_to_training=True,
+                evaluation_interval=1,
+                evaluation_num_env_runners=1,
+                evaluation_duration=1,
+                evaluation_duration_unit="timesteps",
+            )
+            .training(train_batch_size=1, minibatch_size=1)  # Speed things up
+        )
+        algo = config.build()
+        metrics = algo.train()
+        # This can only be true if we do not execute training and evaluation in sequence
+        assert metrics["time_this_iter_s"] < SECONDS_TO_SLEEP * 2
+        assert metrics["time_this_iter_s"] > SECONDS_TO_SLEEP
+        algo.stop()
+
+        config.evaluation(evaluation_parallel_to_training=False)
+        algo_2 = config.build()
+        metrics_2 = algo_2.train()
+        # This must be true if we execute training and evaluation in sequence
+        assert metrics_2["time_this_iter_s"] > SECONDS_TO_SLEEP * 2
+        algo_2.stop()
+
+    def test_custom_eval_function_falsy_results(self):
+        """Test that custom eval function can return ({}, 0, 0)."""
+        config = (
+            ppo.PPOConfig()
+            .environment("CartPole-v1")
+            .evaluation(
+                custom_evaluation_function=lambda algo, eval_workers: ({}, 0, 0),
+                evaluation_interval=1,
+                evaluation_num_env_runners=1,
+                evaluation_duration=1,
+                evaluation_duration_unit="episodes",
+            )
+            .training(train_batch_size=50, minibatch_size=25, num_epochs=1)
+        )
+        algo = config.build()
+        metrics = algo.train()
+        self.assertIn(EVALUATION_RESULTS, metrics)
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

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__]))

rllib/algorithms/tests/test_algorithm_export_checkpoint.py

@@ -0,0 +1,105 @@
+import os
+import shutil
+import unittest
+
+import numpy as np
+
+import ray
+import ray._common
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.framework import try_import_torch
+from ray.tune.registry import get_trainable_cls
+
+torch, _ = try_import_torch()
+
+# Keep a set of all RLlib algos that support the RLModule API.
+# For these algos we need to disable the RLModule API in the config for the purpose of
+# this test. This test is made for the ModelV2 API which is not the same as RLModule.
+RLMODULE_SUPPORTED_ALGOS = {"PPO"}
+
+
+def save_test(alg_name, framework="tf", multi_agent=False):
+    config = (
+        get_trainable_cls(alg_name)
+        .get_default_config()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False, enable_rl_module_and_learner=False
+        )
+        .framework(framework)
+        # Switch on saving native DL-framework (tf, torch) model files.
+        .checkpointing(export_native_model_files=True)
+    )
+
+    if alg_name in RLMODULE_SUPPORTED_ALGOS:
+        config.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+
+    if multi_agent:
+        config.multi_agent(
+            policies={"pol1", "pol2"},
+            policy_mapping_fn=(
+                lambda agent_id, episode, worker, **kwargs: "pol1"
+                if agent_id == "agent1"
+                else "pol2"
+            ),
+        )
+        config.environment(MultiAgentCartPole, env_config={"num_agents": 2})
+    else:
+        config.environment("CartPole-v1")
+    algo = config.build()
+    test_obs = np.array([[0.1, 0.2, 0.3, 0.4]])
+
+    export_dir = os.path.join(
+        ray._common.utils.get_default_ray_temp_dir(),
+        "export_dir_%s" % alg_name,
+    )
+
+    algo.train()
+    print("Exporting algo checkpoint", alg_name, export_dir)
+    export_dir = algo.save(export_dir).checkpoint.path
+    model_dir = os.path.join(
+        export_dir,
+        "policies",
+        "pol1" if multi_agent else DEFAULT_POLICY_ID,
+        "model",
+    )
+
+    # Test loading exported model and perform forward pass.
+    filename = os.path.join(model_dir, "model.pt")
+    model = torch.load(filename, weights_only=False)
+    assert model
+    results = model(
+        input_dict={"obs": torch.from_numpy(test_obs)},
+        # TODO (sven): Make non-RNN models NOT expect these args at all.
+        state=[torch.tensor(0)],  # dummy value
+        seq_lens=torch.tensor(0),  # dummy value
+    )
+    assert len(results) == 2
+    assert results[0].shape == (1, 2)
+    assert results[1] == [torch.tensor(0)]  # dummy
+
+    shutil.rmtree(export_dir)
+
+
+class TestAlgorithmSave(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init(num_cpus=4)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_save_appo_multi_agent(self):
+        save_test("APPO", "torch", multi_agent=True)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_algorithm_imports.py

@@ -0,0 +1,24 @@
+import unittest
+
+import ray
+from ray.rllib.algorithms.registry import ALGORITHMS
+
+
+class TestAlgorithmImport(unittest.TestCase):
+    def setUp(self):
+        ray.init()
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def test_algo_import(self):
+        for name, func in ALGORITHMS.items():
+            func()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_algorithm_rl_module_restore.py

@@ -0,0 +1,336 @@
+import shutil
+import tempfile
+import unittest
+
+import gymnasium as gym
+import numpy as np
+import tree
+
+import ray
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.multi_rl_module import (
+    MultiRLModule,
+    MultiRLModuleSpec,
+)
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.test_utils import check
+
+NUM_AGENTS = 2
+
+
+class TestAlgorithmRLModuleRestore(unittest.TestCase):
+    """Test RLModule loading from rl module spec across a local node."""
+
+    def setUp(self) -> None:
+        ray.init()
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    @staticmethod
+    def get_ppo_config(num_agents=NUM_AGENTS):
+        def policy_mapping_fn(agent_id, episode, **kwargs):
+            # policy_id is policy_i where i is the agent id
+            pol_id = f"policy_{agent_id}"
+            return pol_id
+
+        scaling_config = {
+            "num_learners": 0,
+            "num_gpus_per_learner": 0,
+        }
+
+        policies = {f"policy_{i}" for i in range(num_agents)}
+
+        config = (
+            PPOConfig()
+            .env_runners(rollout_fragment_length=4)
+            .learners(**scaling_config)
+            .environment(MultiAgentCartPole, env_config={"num_agents": num_agents})
+            .training(num_epochs=1, train_batch_size=8, minibatch_size=8)
+            .multi_agent(policies=policies, policy_mapping_fn=policy_mapping_fn)
+        )
+        return config
+
+    def test_e2e_load_simple_multi_rl_module(self):
+        """Test if we can train a PPO algo with a checkpointed MultiRLModule e2e."""
+        config = self.get_ppo_config()
+        env = MultiAgentCartPole({"num_agents": NUM_AGENTS})
+        # create a multi_rl_module to load and save it to a checkpoint directory
+        module_specs = {}
+        for i in range(NUM_AGENTS):
+            module_specs[f"policy_{i}"] = RLModuleSpec(
+                module_class=PPOTorchRLModule,
+                observation_space=env.get_observation_space(0),
+                action_space=env.get_action_space(0),
+                # If we want to use this externally created module in the algorithm,
+                # we need to provide the same config as the algorithm.
+                model_config=DefaultModelConfig(fcnet_hiddens=[32 * (i + 1)]),
+                catalog_class=PPOCatalog,
+            )
+        multi_rl_module_spec = MultiRLModuleSpec(rl_module_specs=module_specs)
+        multi_rl_module = multi_rl_module_spec.build()
+        multi_rl_module_weights = convert_to_numpy(multi_rl_module.get_state())
+        marl_checkpoint_path = tempfile.mkdtemp()
+        multi_rl_module.save_to_path(marl_checkpoint_path)
+
+        # create a new MARL_spec with the checkpoint from the previous one
+        multi_rl_module_spec_from_checkpoint = MultiRLModuleSpec(
+            rl_module_specs=module_specs,
+            load_state_path=marl_checkpoint_path,
+        )
+        config.rl_module(rl_module_spec=multi_rl_module_spec_from_checkpoint)
+
+        # create the algorithm with multiple nodes and check if the weights
+        # are the same as the original MultiRLModule
+        algo = config.build()
+        algo_module_weights = algo.learner_group.get_weights()
+        check(algo_module_weights, multi_rl_module_weights)
+        algo.train()
+        algo.stop()
+        del algo
+        shutil.rmtree(marl_checkpoint_path)
+
+    def test_e2e_load_complex_multi_rl_module(self):
+        """Test if we can train a PPO algorithm with a cpkt MARL and RL module e2e."""
+        config = self.get_ppo_config()
+        env = MultiAgentCartPole({"num_agents": NUM_AGENTS})
+        # create a multi_rl_module to load and save it to a checkpoint directory
+        module_specs = {}
+        for i in range(NUM_AGENTS):
+            module_specs[f"policy_{i}"] = RLModuleSpec(
+                module_class=PPOTorchRLModule,
+                observation_space=env.get_observation_space(0),
+                action_space=env.get_action_space(0),
+                # If we want to use this externally created module in the algorithm,
+                # we need to provide the same config as the algorithm.
+                model_config=DefaultModelConfig(fcnet_hiddens=[32 * (i + 1)]),
+                catalog_class=PPOCatalog,
+            )
+        multi_rl_module_spec = MultiRLModuleSpec(rl_module_specs=module_specs)
+        multi_rl_module = multi_rl_module_spec.build()
+        marl_checkpoint_path = tempfile.mkdtemp()
+        multi_rl_module.save_to_path(marl_checkpoint_path)
+
+        # create a RLModule to load and override the "policy_1" module with
+        module_to_swap_in = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.get_observation_space(0),
+            action_space=env.get_action_space(0),
+            # Note, we need to pass in the default model config for the algorithm
+            # to be able to use this module later.
+            model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+            catalog_class=PPOCatalog,
+        ).build()
+
+        module_to_swap_in_path = tempfile.mkdtemp()
+        module_to_swap_in.save_to_path(module_to_swap_in_path)
+
+        # create a new MARL_spec with the checkpoint from the marl_checkpoint
+        # and the module_to_swap_in_checkpoint
+        module_specs["policy_1"] = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.get_observation_space(0),
+            action_space=env.get_action_space(0),
+            model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+            catalog_class=PPOCatalog,
+            load_state_path=module_to_swap_in_path,
+        )
+        multi_rl_module_spec_from_checkpoint = MultiRLModuleSpec(
+            rl_module_specs=module_specs,
+            load_state_path=marl_checkpoint_path,
+        )
+        config = config.rl_module(rl_module_spec=multi_rl_module_spec_from_checkpoint)
+
+        # create the algorithm with multiple nodes and check if the weights
+        # are the same as the original MultiRLModule
+        algo = config.build()
+        algo_module_weights = algo.learner_group.get_weights()
+
+        multi_rl_module_with_swapped_in_module = MultiRLModule()
+        multi_rl_module_with_swapped_in_module.add_module(
+            "policy_0", multi_rl_module["policy_0"]
+        )
+        multi_rl_module_with_swapped_in_module.add_module("policy_1", module_to_swap_in)
+
+        check(
+            algo_module_weights,
+            convert_to_numpy(multi_rl_module_with_swapped_in_module.get_state()),
+        )
+        algo.train()
+        algo.stop()
+        del algo
+        shutil.rmtree(marl_checkpoint_path)
+
+    def test_e2e_load_rl_module(self):
+        """Test if we can train a PPO algorithm with a cpkt RL module e2e."""
+        scaling_config = {
+            "num_learners": 0,
+            "num_gpus_per_learner": 0,
+        }
+
+        config = (
+            PPOConfig()
+            .env_runners(rollout_fragment_length=4)
+            .learners(**scaling_config)
+            .environment("CartPole-v1")
+            .training(num_epochs=1, train_batch_size=8, minibatch_size=8)
+        )
+        env = gym.make("CartPole-v1")
+        # create a multi_rl_module to load and save it to a checkpoint directory
+        module_spec = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            # If we want to use this externally created module in the algorithm,
+            # we need to provide the same config as the algorithm.
+            model_config=DefaultModelConfig(fcnet_hiddens=[32]),
+            catalog_class=PPOCatalog,
+        )
+        module = module_spec.build()
+
+        module_ckpt_path = tempfile.mkdtemp()
+        module.save_to_path(module_ckpt_path)
+
+        module_to_load_spec = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            model_config=DefaultModelConfig(fcnet_hiddens=[32]),
+            catalog_class=PPOCatalog,
+            load_state_path=module_ckpt_path,
+        )
+
+        config.rl_module(rl_module_spec=module_to_load_spec)
+
+        # create the algorithm with multiple nodes and check if the weights
+        # are the same as the original MultiRLModule
+        algo = config.build()
+        algo_module_weights = algo.learner_group.get_weights()
+
+        check(
+            algo_module_weights[DEFAULT_MODULE_ID],
+            convert_to_numpy(module.get_state()),
+        )
+        algo.train()
+        algo.stop()
+        del algo
+        shutil.rmtree(module_ckpt_path)
+
+    def test_e2e_load_complex_multi_rl_module_with_modules_to_load(self):
+        """Test if we can train a PPO algorithm with a cpkt MARL and RL module e2e.
+
+        Additionally, check if we can set modules to load so that we can exclude
+        a module from our ckpted MultiRLModule from being loaded.
+
+        """
+        num_agents = 3
+        config = self.get_ppo_config(num_agents=num_agents)
+        env = MultiAgentCartPole({"num_agents": num_agents})
+        # create a multi_rl_module to load and save it to a checkpoint directory
+        module_specs = {}
+        for i in range(num_agents):
+            module_specs[f"policy_{i}"] = RLModuleSpec(
+                module_class=PPOTorchRLModule,
+                observation_space=env.get_observation_space(0),
+                action_space=env.get_action_space(0),
+                # Note, we need to pass in the default model config for the
+                # algorithm to be able to use this module later.
+                model_config=DefaultModelConfig(fcnet_hiddens=[32 * (i + 1)]),
+                catalog_class=PPOCatalog,
+            )
+        multi_rl_module_spec = MultiRLModuleSpec(rl_module_specs=module_specs)
+        multi_rl_module = multi_rl_module_spec.build()
+        marl_checkpoint_path = tempfile.mkdtemp()
+        multi_rl_module.save_to_path(marl_checkpoint_path)
+
+        # create a RLModule to load and override the "policy_1" module with
+        module_to_swap_in = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.get_observation_space(0),
+            action_space=env.get_action_space(0),
+            # Note, we need to pass in the default model config for the algorithm
+            # to be able to use this module later.
+            model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+            catalog_class=PPOCatalog,
+        ).build()
+
+        module_to_swap_in_path = tempfile.mkdtemp()
+        module_to_swap_in.save_to_path(module_to_swap_in_path)
+
+        # create a new MARL_spec with the checkpoint from the marl_checkpoint
+        # and the module_to_swap_in_checkpoint
+        module_specs["policy_1"] = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.get_observation_space(0),
+            action_space=env.get_action_space(0),
+            model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+            catalog_class=PPOCatalog,
+            load_state_path=module_to_swap_in_path,
+        )
+        multi_rl_module_spec_from_checkpoint = MultiRLModuleSpec(
+            rl_module_specs=module_specs,
+            load_state_path=marl_checkpoint_path,
+            modules_to_load={
+                "policy_0",
+            },
+        )
+        config.rl_module(rl_module_spec=multi_rl_module_spec_from_checkpoint)
+
+        # create the algorithm with multiple nodes and check if the weights
+        # are the same as the original MultiRLModule
+        algo = config.build()
+        algo_module_weights = algo.learner_group.get_weights()
+
+        # weights of "policy_0" should be the same as in the loaded MultiRLModule
+        # since we specified it as being apart of the modules_to_load
+        check(
+            algo_module_weights["policy_0"],
+            convert_to_numpy(multi_rl_module["policy_0"].get_state()),
+        )
+        # weights of "policy_1" should be the same as in the module_to_swap_in since
+        # we specified its load path separately in an rl_module_spec inside of the
+        # multi_rl_module_spec_from_checkpoint
+        check(
+            algo_module_weights["policy_1"],
+            convert_to_numpy(module_to_swap_in.get_state()),
+        )
+        # weights of "policy_2" should be different from the loaded MultiRLModule
+        # since we didn't specify it as being apart of the modules_to_load
+        policy_2_algo_module_weight_sum = np.sum(
+            [
+                np.sum(s)
+                for s in tree.flatten(convert_to_numpy(algo_module_weights["policy_2"]))
+            ]
+        )
+        policy_2_multi_rl_module_weight_sum = np.sum(
+            [
+                np.sum(s)
+                for s in tree.flatten(
+                    convert_to_numpy(multi_rl_module["policy_2"].get_state())
+                )
+            ]
+        )
+        check(
+            policy_2_algo_module_weight_sum,
+            policy_2_multi_rl_module_weight_sum,
+            false=True,
+        )
+
+        algo.train()
+        algo.stop()
+        del algo
+        shutil.rmtree(marl_checkpoint_path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_algorithm_save_load_checkpoint_connectors.py

@@ -0,0 +1,232 @@
+import tempfile
+import unittest
+
+import ray
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter
+from ray.rllib.core import COMPONENT_ENV_TO_MODULE_CONNECTOR
+from ray.rllib.utils.filter import RunningStat
+from ray.rllib.utils.test_utils import check
+
+algorithms_and_configs = {
+    "PPO": (PPOConfig().training(train_batch_size=2, minibatch_size=2))
+}
+
+
+@ray.remote
+def save_train_and_get_states(
+    algo_cfg: AlgorithmConfig, num_env_runners: int, env: str, tmpdir
+):
+    """Create an algo, train for 10 iterations, then checkpoint it.
+
+    Note: This function uses a seeded algorithm that can modify the global random state.
+        Running it multiple times in the same process can affect other algorithms.
+        Making it a Ray task runs it in a separate process and prevents it from
+        affecting other algorithms' random state.
+
+    Args:
+        algo_cfg: The algorithm config to build the algo from.
+        num_env_runners: Number of environment runners to use.
+        env: The gym genvironment to train on.
+        tmpdir: The temporary directory to save the checkpoint to.
+
+    Returns:
+        The env-runner states after 10 iterations of training.
+    """
+    algo_cfg = (
+        algo_cfg.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        .environment(env)
+        .env_runners(
+            num_env_runners=num_env_runners,
+            env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+        )
+        # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+        # to make sure that we get results as soon as sampling/training is done at
+        # least once
+        .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+        .debugging(seed=10)
+    )
+    algo = algo_cfg.build()
+    for _ in range(10):
+        algo.train()
+    algo.save_to_path(tmpdir)
+    states = algo.env_runner_group.foreach_env_runner(
+        "get_state",
+        local_env_runner=False,
+    )
+    return states
+
+
+@ray.remote
+def load_and_get_states(
+    algo_cfg: AlgorithmConfig, num_env_runners: int, env: str, tmpdir
+):
+    """Loads the checkpoint saved by save_train_and_get_states and returns connector states.
+
+    Note: This function uses a seeded algorithm that can modify the global random state.
+        Running it multiple times in the same process can affect other algorithms.
+        Making it a Ray task runs it in a separate process and prevents it from
+        affecting other algorithms' random state.
+
+    Args:
+        algo_cfg: The algorithm config to build the algo from.
+        num_env_runners: Number of env-runners to use.
+        env: The gym genvironment to train on.
+        tmpdir: The temporary directory to save the checkpoint to.
+
+    Returns:
+        The connector states of remote env-runners after 10 iterations of training.
+
+    """
+    algo_cfg = (
+        algo_cfg.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        .environment(env)
+        .env_runners(
+            num_env_runners=num_env_runners,
+            env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+        )
+        # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+        # to make sure that we get results as soon as sampling/training is done at
+        # least once
+        .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+        .debugging(seed=10)
+    )
+    algo = algo_cfg.build()
+    algo.restore_from_path(tmpdir)
+    states = algo.env_runner_group.foreach_env_runner(
+        "get_state",
+        local_env_runner=False,
+    )
+
+    return states
+
+
+class TestAlgorithmWithConnectorsSaveAndRestore(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_save_and_restore_w_remote_env_runners(self):
+        num_env_runners = 2
+        for algo_name in algorithms_and_configs:
+            config = algorithms_and_configs[algo_name]
+            with tempfile.TemporaryDirectory() as tmpdir:
+                # create an algorithm, checkpoint it, then train for 2 iterations
+                connector_states_algo_1 = ray.get(
+                    save_train_and_get_states.remote(
+                        config, num_env_runners, "CartPole-v1", tmpdir
+                    )
+                )
+                # load that checkpoint into a new algorithm and check the states.
+                connector_states_algo_2 = ray.get(  # noqa
+                    load_and_get_states.remote(
+                        config, num_env_runners, "CartPole-v1", tmpdir
+                    )
+                )
+
+                # Assert that all running stats are the same.
+                self._assert_running_stats_consistency(
+                    connector_states_algo_1, connector_states_algo_2
+                )
+
+    def test_save_and_restore_w_remote_env_runners_and_wo_local_env_runner(self):
+        num_env_runners = 2
+        for algo_name in algorithms_and_configs:
+            config = algorithms_and_configs[algo_name].env_runners(
+                create_local_env_runner=False
+            )
+            with tempfile.TemporaryDirectory() as tmpdir:
+                # create an algorithm, checkpoint it, then train for 2 iterations
+                connector_states_algo_1 = ray.get(
+                    save_train_and_get_states.remote(
+                        config, num_env_runners, "CartPole-v1", tmpdir
+                    )
+                )
+                # load that checkpoint into a new algorithm and check the states.
+                connector_states_algo_2 = ray.get(  # noqa
+                    load_and_get_states.remote(
+                        config, num_env_runners, "CartPole-v1", tmpdir
+                    )
+                )
+                # Assert that all running stats are the same.
+                self._assert_running_stats_consistency(
+                    connector_states_algo_1, connector_states_algo_2
+                )
+
+    def _assert_running_stats_consistency(
+        self, connector_states_algo_1: list, connector_states_algo_2: list
+    ):
+        """
+        Asserts consistency of running stats within and between algorithms.
+        """
+
+        running_stats_states_algo_1 = [
+            state[COMPONENT_ENV_TO_MODULE_CONNECTOR]["MeanStdFilter"][None][
+                "running_stats"
+            ]
+            for state in connector_states_algo_1
+        ]
+        running_stats_states_algo_2 = [
+            state[COMPONENT_ENV_TO_MODULE_CONNECTOR]["MeanStdFilter"][None][
+                "running_stats"
+            ]
+            for state in connector_states_algo_2
+        ]
+
+        running_stats_states_algo_1 = [
+            [RunningStat.from_state(s) for s in running_stats_state]
+            for running_stats_state in running_stats_states_algo_1
+        ]
+        running_stats_states_algo_2 = [
+            [RunningStat.from_state(s) for s in running_stats_state]
+            for running_stats_state in running_stats_states_algo_2
+        ]
+
+        running_stats_states_algo_1 = [
+            (
+                running_stat[0].n,
+                running_stat[0].mean_array,
+                running_stat[0].sum_sq_diff_array,
+            )
+            for running_stat in running_stats_states_algo_1
+        ]
+        running_stats_states_algo_2 = [
+            (
+                running_stat[0].n,
+                running_stat[0].mean_array,
+                running_stat[0].sum_sq_diff_array,
+            )
+            for running_stat in running_stats_states_algo_2
+        ]
+
+        # The number of env-runners must be two for the following checks to make sense.
+        self.assertEqual(len(running_stats_states_algo_1), 2)
+        self.assertEqual(len(running_stats_states_algo_2), 2)
+
+        # Assert that all running stats in algo-1 are the same (for consistency).
+        check(running_stats_states_algo_1[0][0], running_stats_states_algo_1[1][0])
+
+        # Now ensure that the connector states on remote `EnvRunner`s were restored.
+        check(running_stats_states_algo_1[0][0], running_stats_states_algo_2[0][0])
+
+        # Ensure also that all states are the same in algo-2 (for consistency).
+        check(running_stats_states_algo_2[0][0], running_stats_states_algo_2[1][0])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_algorithm_save_load_checkpoint_learner.py

@@ -0,0 +1,131 @@
+import tempfile
+import unittest
+
+import ray
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+
+algorithms_and_configs = {
+    "PPO": (PPOConfig().training(train_batch_size=2, minibatch_size=2))
+}
+
+
+@ray.remote
+def save_and_train(algo_cfg: AlgorithmConfig, env: str, tmpdir):
+    """Create an algo, checkpoint it, then train for 2 iterations.
+
+    Note: This function uses a seeded algorithm that can modify the global random state.
+        Running it multiple times in the same process can affect other algorithms.
+        Making it a Ray task runs it in a separate process and prevents it from
+        affecting other algorithms' random state.
+
+    Args:
+        algo_cfg: The algorithm config to build the algo from.
+        env: The gym genvironment to train on.
+        tmpdir: The temporary directory to save the checkpoint to.
+
+    Returns:
+        The learner stats after 2 iterations of training.
+    """
+    algo_cfg = (
+        algo_cfg.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        .environment(env)
+        .env_runners(num_env_runners=0)
+        # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+        # to make sure that we get results as soon as sampling/training is done at
+        # least once
+        .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+        .debugging(seed=10)
+    )
+    algo = algo_cfg.build()
+
+    algo.save_to_path(tmpdir)
+    for _ in range(2):
+        results = algo.train()
+    return results[LEARNER_RESULTS][DEFAULT_MODULE_ID]
+
+
+@ray.remote
+def load_and_train(algo_cfg: AlgorithmConfig, env: str, tmpdir):
+    """Loads the checkpoint saved by save_and_train and trains for 2 iterations.
+
+    Note: This function uses a seeded algorithm that can modify the global random state.
+        Running it multiple times in the same process can affect other algorithms.
+        Making it a Ray task runs it in a separate process and prevents it from
+        affecting other algorithms' random state.
+
+    Args:
+        algo_cfg: The algorithm config to build the algo from.
+        env: The gym genvironment to train on.
+        tmpdir: The temporary directory to save the checkpoint to.
+
+    Returns:
+        The learner stats after 2 iterations of training.
+
+    """
+    algo_cfg = (
+        algo_cfg.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        .environment(env)
+        .env_runners(num_env_runners=0)
+        # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+        # to make sure that we get results as soon as sampling/training is done at
+        # least once
+        .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+        .debugging(seed=10)
+    )
+    algo = algo_cfg.build()
+    algo.restore_from_path(tmpdir)
+    for _ in range(2):
+        results = algo.train()
+    return results[LEARNER_RESULTS][DEFAULT_MODULE_ID]
+
+
+class TestAlgorithmWithLearnerSaveAndRestore(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_save_and_restore(self):
+        for algo_name in algorithms_and_configs:
+            config = algorithms_and_configs[algo_name]
+            with tempfile.TemporaryDirectory() as tmpdir:
+                # create an algorithm, checkpoint it, then train for 2 iterations
+                ray.get(save_and_train.remote(config, "CartPole-v1", tmpdir))
+                # load that checkpoint into a new algorithm and train for 2
+                # iterations
+                results_algo_2 = ray.get(  # noqa
+                    load_and_train.remote(config, "CartPole-v1", tmpdir)
+                )
+
+                # load that checkpoint into another new algorithm and train for 2
+                # iterations
+                results_algo_3 = ray.get(  # noqa
+                    load_and_train.remote(config, "CartPole-v1", tmpdir)
+                )
+
+                # check that the results are the same across loaded algorithms
+                # they won't be the same as the first algorithm since the random
+                # state that is used for each algorithm is not preserved across
+                # checkpoints.
+                # TODO (sven): Uncomment once seeding works on EnvRunners.
+                # check(results_algo_3, results_algo_2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_custom_resource.py

@@ -0,0 +1,43 @@
+import pytest
+
+import ray
+from ray import tune
+from ray.tune.registry import get_trainable_cls
+from ray.tune.result import TRAINING_ITERATION
+
+
+@pytest.mark.parametrize("algorithm", ["PPO", "IMPALA"])
+def test_custom_resource(algorithm):
+    if ray.is_initialized:
+        ray.shutdown()
+
+    ray.init(
+        resources={"custom_resource": 1},
+        include_dashboard=False,
+    )
+
+    config = (
+        get_trainable_cls(algorithm)
+        .get_default_config()
+        .environment("CartPole-v1")
+        .framework("torch")
+        .env_runners(
+            num_env_runners=1,
+            custom_resources_per_env_runner={"custom_resource": 0.01},
+        )
+        .resources(num_gpus=0)
+    )
+    stop = {TRAINING_ITERATION: 1}
+
+    tune.Tuner(
+        algorithm,
+        param_space=config,
+        run_config=tune.RunConfig(stop=stop, verbose=0),
+        tune_config=tune.TuneConfig(num_samples=1),
+    ).fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_custom_resources_per_learner.py

@@ -0,0 +1,117 @@
+import pytest
+
+import ray
+import ray._common
+from ray.cluster_utils import Cluster
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.utils import _get_learner_bundles
+
+EXPECTED_PER_NODE_OBJECT_STORE_MEMORY = 10**8
+HEAD_REDIS_PORT = 6379
+HEAD_CPUS = 4
+WORKER_CPUS = 4
+PINNED_RESOURCE = "learner_pool"
+DECOY_RESOURCE = "decoy_pool"
+
+
+def test_round_trip():
+    """Test that custom_resources_per_learner is set correctly."""
+    cfg = AlgorithmConfig().learners(
+        custom_resources_per_learner={"my_label": 0.001, "other": 1}
+    )
+    assert cfg.custom_resources_per_learner == {"my_label": 0.001, "other": 1}
+
+
+@pytest.mark.parametrize("reserved", ["CPU", "GPU"])
+def test_reserved_keys_rejected(reserved):
+    """`CPU`/`GPU` belong to `num_*_per_learner`, not custom resources."""
+    with pytest.raises(ValueError, match="CPU.*GPU"):
+        AlgorithmConfig().learners(custom_resources_per_learner={reserved: 1})
+
+
+def test_placement_group_bundles_include_custom_resources():
+    """The PG bundles built for Tune must reserve the custom learner resources;
+    otherwise learners scheduled within the PG can never satisfy their request.
+    """
+    cfg = AlgorithmConfig().learners(
+        num_learners=2,
+        num_cpus_per_learner=1,
+        custom_resources_per_learner={"learner_pool": 0.5},
+    )
+    bundles = _get_learner_bundles(cfg)
+    assert len(bundles) == 2
+    assert all(b.get("learner_pool") == 0.5 for b in bundles)
+
+
+@pytest.fixture
+def cluster():
+    """3-node fake cluster:
+
+    - head:    HEAD_CPUS CPUs, no custom resources
+    - pinned:  WORKER_CPUS CPUs, {PINNED_RESOURCE: 4}
+    - decoy:   WORKER_CPUS CPUs, {DECOY_RESOURCE: 4} (must not host learners
+               when requesting PINNED_RESOURCE)
+    """
+    assert (
+        3 * EXPECTED_PER_NODE_OBJECT_STORE_MEMORY
+        < ray._common.utils.get_system_memory() / 2
+    ), "Not enough memory on this machine to run this workload."
+
+    cluster = Cluster()
+    head = cluster.add_node(
+        redis_port=HEAD_REDIS_PORT,
+        num_cpus=HEAD_CPUS,
+        object_store_memory=EXPECTED_PER_NODE_OBJECT_STORE_MEMORY,
+        include_dashboard=True,
+    )
+    pinned = cluster.add_node(
+        num_cpus=WORKER_CPUS,
+        resources={PINNED_RESOURCE: 4},
+        object_store_memory=EXPECTED_PER_NODE_OBJECT_STORE_MEMORY,
+        include_dashboard=False,
+    )
+    decoy = cluster.add_node(
+        num_cpus=WORKER_CPUS,
+        resources={DECOY_RESOURCE: 4},
+        object_store_memory=EXPECTED_PER_NODE_OBJECT_STORE_MEMORY,
+        include_dashboard=False,
+    )
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    yield {"cluster": cluster, "head": head, "pinned": pinned, "decoy": decoy}
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+def test_algorithm_pins_learners_to_node_with_custom_resource(cluster):
+    """End-to-end: a built Algorithm places its learners on the node exposing
+    PINNED_RESOURCE, even though head+decoy together have 8 free CPUs that
+    would otherwise absorb them.
+    """
+    config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .env_runners(num_env_runners=0)
+        .learners(
+            num_learners=2,
+            num_cpus_per_learner=1,
+            custom_resources_per_learner={PINNED_RESOURCE: 1},
+        )
+    )
+    algo = config.build()
+
+    refs = algo.learner_group.foreach_learner(
+        lambda _: ray.get_runtime_context().get_node_id()
+    )
+    node_ids = [r.get() for r in refs]
+    assert len(node_ids) == 2
+    assert all(nid == cluster["pinned"].node_id for nid in node_ids)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_dependency_tf.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+if __name__ == "__main__":
+    # Do not import tf for testing purposes.
+    os.environ["RLLIB_TEST_NO_TF_IMPORT"] = "1"
+
+    # Test registering (includes importing) all Algorithms.
+    from ray.rllib import _register_all
+
+    # This should surface any dependency on tf, e.g. inside function
+    # signatures/typehints.
+    _register_all()
+
+    from ray.rllib.algorithms.ppo import PPOConfig
+
+    assert (
+        "tensorflow" not in sys.modules
+    ), "`tensorflow` initially present, when it shouldn't!"
+
+    config = (
+        PPOConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=True,
+            enable_rl_module_and_learner=True,
+        )
+        .environment("CartPole-v1")
+        .framework("torch")
+        .env_runners(num_env_runners=0)
+    )
+    # Note: No ray.init(), to test it works without Ray
+    algo = config.build()
+    algo.train()
+
+    assert (
+        "tensorflow" not in sys.modules
+    ), "`tensorflow` should not be imported after creating and training A3C!"
+
+    # Clean up.
+    del os.environ["RLLIB_TEST_NO_TF_IMPORT"]
+
+    algo.stop()
+
+    print("ok")

rllib/algorithms/tests/test_dependency_torch.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+if __name__ == "__main__":
+    # Do not import torch for testing purposes.
+    os.environ["RLLIB_TEST_NO_TORCH_IMPORT"] = "1"
+
+    # Test registering (includes importing) all Algorithms.
+    from ray.rllib import _register_all
+
+    # This should surface any dependency on torch, e.g. inside function
+    # signatures/typehints.
+    _register_all()
+
+    from ray.rllib.algorithms.ppo import PPOConfig
+
+    assert "torch" not in sys.modules, "`torch` initially present, when it shouldn't!"
+
+    # Note: No ray.init(), to test it works without Ray
+    config = (
+        PPOConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("CartPole-v1")
+        .framework("tf")
+        .env_runners(num_env_runners=0)
+    )
+
+    # Disable auto-added TBX logger callback to avoid importing torch
+    # via the tensorboardX.SummaryWriter class.
+    os.environ["TUNE_DISABLE_AUTO_CALLBACK_LOGGERS"] = "1"
+    algo = config.build()
+    algo.train()
+
+    assert (
+        "torch" not in sys.modules
+    ), "`torch` should not be imported after creating and training A3C!"
+
+    # Clean up.
+    del os.environ["RLLIB_TEST_NO_TORCH_IMPORT"]
+    del os.environ["TUNE_DISABLE_AUTO_CALLBACK_LOGGERS"]
+
+    algo.stop()
+
+    print("ok")

rllib/algorithms/tests/test_env_runner_state_server.py

@@ -0,0 +1,63 @@
+import sys
+
+import pytest
+
+from ray.rllib.env.env_runner_state_server import EnvRunnerStateServer
+from ray.rllib.utils.metrics import WEIGHTS_SEQ_NO
+
+
+def _state(seq_no, weights="WEIGHTS_OBJ_REF"):
+    # The server stores whatever StateDict it is handed. In production the `rl_module`
+    # value is a `ray.ObjectRef`, but the server never dereferences it, so a plain
+    # sentinel is sufficient to exercise the store/serve logic here.
+    return {WEIGHTS_SEQ_NO: seq_no, "rl_module": weights}
+
+
+def test_push_pull_roundtrip_preserves_state_verbatim():
+    server = EnvRunnerStateServer()
+    state = _state(3)
+    server.push(state)
+    # Returned verbatim (including the - here sentinel - weights "ObjectRef"); the
+    # server must NOT dereference or copy it.
+    assert server.pull() is state
+    assert server.pull()["rl_module"] == "WEIGHTS_OBJ_REF"
+    assert server.get_version() == 3
+
+
+def test_push_replaces_and_advances_version():
+    server = EnvRunnerStateServer()
+    server.push(_state(1))
+    server.push(_state(2))
+    assert server.get_version() == 2
+    assert server.pull()[WEIGHTS_SEQ_NO] == 2
+
+
+def test_push_rejects_state_without_seq_no():
+    # A state without WEIGHTS_SEQ_NO could never be pulled (EnvRunners version-gate via
+    # `pull_if_newer`), so the server rejects it loudly instead of holding state that no
+    # EnvRunner would ever apply.
+    server = EnvRunnerStateServer()
+    with pytest.raises(ValueError, match="WEIGHTS_SEQ_NO"):
+        server.push({"rl_module": "WEIGHTS_OBJ_REF"})
+    # Nothing was stored.
+    assert server.pull() is None
+    assert server.get_version() == -1
+
+
+def test_pull_if_newer_only_returns_strictly_newer_state():
+    server = EnvRunnerStateServer()
+    # Empty server -> nothing to return, regardless of the caller's version.
+    assert server.pull_if_newer(-1) is None
+
+    state = _state(5)
+    server.push(state)
+    # Equal-or-newer caller version -> None, so the (heavy) state is NOT transferred.
+    assert server.pull_if_newer(5) is None
+    assert server.pull_if_newer(6) is None
+    # Older caller version -> the full state, returned verbatim (same object).
+    assert server.pull_if_newer(4) is state
+    assert server.pull_if_newer(-1) is state
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_eval_workers_all_unhealthy.py

@@ -0,0 +1,258 @@
+"""Tests for evaluating when all configured remote eval EnvRunners are
+unhealthy.
+
+The behavior is controlled by two orthogonal config knobs:
+
+- ``evaluation_unhealthy_workers_timeout_s``: how long to wait for at
+  least one eval EnvRunner to recover before deciding what to do
+  (default 0: don't wait).
+- ``evaluation_error_after_n_consecutive_skips``: tolerate this many
+  consecutive evaluation iterations in which all configured remote eval
+  EnvRunners are unhealthy. On the next such iteration, ``evaluate()``
+  raises ``RuntimeError``. ``None`` (default) tolerates an unbounded
+  number of consecutive skips.
+
+Both apply identically regardless of ``evaluation_parallel_to_training``.
+"""
+import time
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray.rllib.algorithms.ppo import PPOConfig
+
+
+@pytest.fixture(params=[True, False], ids=["parallel", "sequential"])
+def parallel_to_training(request):
+    return request.param
+
+
+def _algo_with_unhealthy_eval_workers(
+    *,
+    timeout_s=0,
+    error_after_n_consecutive_skips=None,
+    parallel_to_training=True,
+    spawn_replacement=False,
+    enable_v2_stack=True,
+):
+    """Build a PPO algo and put every remote eval worker into the failed
+    state.
+
+    Always ``ray.kill(no_restart=True)``-s the original actors and sets
+    ``restart_failed_env_runners=False`` so Ray Core does not auto-resurrect
+    them. This keeps each test deterministic: the only way a worker reappears
+    is for the test to explicitly add one back (see ``spawn_replacement``).
+
+    Args:
+        spawn_replacement: If True, additionally spawn a fresh remote eval
+            EnvRunner via ``add_workers(1)`` after the kill, then flip its
+            health flag to False as well. The replacement is what a
+            subsequent ``probe_unhealthy_env_runners`` will be able to
+            ping successfully -- simulating "a new worker has come back"
+            without relying on Ray Core's restart timing.
+        enable_v2_stack: If True (default), use the new API stack (v2
+            connectors + Learner). If False, use the old API stack so
+            the recovery path's ``_sync_filters_if_needed`` branch is
+            exercised instead of ``sync_env_runner_states``.
+
+    Config is the smallest that exercises the failure path: no remote
+    training EnvRunners, 1 remote eval EnvRunner, fixed-duration eval so
+    we can call ``evaluate()`` directly without a parallel-training
+    future.
+    """
+    config = PPOConfig()
+    if not enable_v2_stack:
+        config = config.api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+    config = (
+        config.environment("CartPole-v1")
+        # Local-only training: skips remote train EnvRunner setup.
+        .env_runners(num_env_runners=0).evaluation(
+            # 1 eval worker is enough; killing it leaves 0 healthy, which
+            # is the condition we're testing.
+            evaluation_num_env_runners=1,
+            evaluation_interval=1,
+            evaluation_parallel_to_training=parallel_to_training,
+            # Fixed duration so `evaluate()` doesn't need a parallel-train
+            # future (avoids the `auto` branch's `assert future is not None`).
+            evaluation_duration=1,
+            evaluation_duration_unit="episodes",
+            evaluation_unhealthy_workers_timeout_s=timeout_s,
+            evaluation_error_after_n_consecutive_skips=(
+                error_after_n_consecutive_skips
+            ),
+        )
+        # Disable auto-restart so the tests are fully in control of when
+        # (and whether) a worker comes back.
+        .fault_tolerance(restart_failed_env_runners=False)
+    )
+    algo = config.build()
+
+    eval_grp = algo.eval_env_runner_group
+    # Hard-kill the originals; no Ray Core restart.
+    for a in list(eval_grp._worker_manager._actors.values()):
+        ray.kill(a, no_restart=True)
+
+    if spawn_replacement:
+        # The dead originals are evicted from the worker manager *before*
+        # spawning the replacement so that `add_workers` assigns the
+        # fresh worker the lowest free `worker_index` (1). Eval's
+        # `_evaluate_with_fixed_duration` builds a per-worker num-units
+        # list of length `evaluation_num_env_runners + 1` indexed by
+        # `worker.worker_index`; leaving the dead originals in the
+        # manager would push the new index past the end of that list.
+        for actor_id in list(eval_grp._worker_manager.actor_ids()):
+            eval_grp._worker_manager.remove_actor(actor_id)
+        before = set(eval_grp._worker_manager.actor_ids())
+        eval_grp.add_workers(1, validate=False)
+        new_ids = set(eval_grp._worker_manager.actor_ids()) - before
+        # `add_workers` registers the replacement with `is_healthy=True`;
+        # flip it so the wait loop is forced to probe. The probe's
+        # `ping()` then succeeds against this alive actor and flips it
+        # back to healthy -- mirroring the production "Ray Core restarted
+        # the actor" recovery sequence.
+        for actor_id in new_ids:
+            eval_grp._worker_manager.set_actor_state(actor_id, healthy=False)
+    else:
+        # No replacement: just mark the (dead) originals unhealthy.
+        for actor_id in list(eval_grp._worker_manager.actor_ids()):
+            eval_grp._worker_manager.set_actor_state(actor_id, healthy=False)
+
+    assert eval_grp.num_healthy_remote_workers() == 0
+    return algo
+
+
+def test_default_skips_eval_silently(parallel_to_training):
+    """Default (threshold=None): evaluate() must return cleanly even if
+    every eval iteration finds 0 healthy workers, indefinitely."""
+    algo = _algo_with_unhealthy_eval_workers(parallel_to_training=parallel_to_training)
+    for _ in range(3):
+        algo.evaluate()  # must not raise on any of these
+
+
+def test_threshold_one_raises_on_first_skip(parallel_to_training):
+    """threshold=1: first failed iteration raises (strictest setting)."""
+    algo = _algo_with_unhealthy_eval_workers(
+        error_after_n_consecutive_skips=1,
+        parallel_to_training=parallel_to_training,
+    )
+    with pytest.raises(
+        RuntimeError, match="evaluation_error_after_n_consecutive_skips"
+    ):
+        algo.evaluate()
+
+
+def test_threshold_n_tolerates_n_minus_1_skips(parallel_to_training):
+    """threshold=N: tolerate N-1 consecutive skips, raise on the N-th."""
+    threshold = 3
+    algo = _algo_with_unhealthy_eval_workers(
+        error_after_n_consecutive_skips=threshold,
+        parallel_to_training=parallel_to_training,
+    )
+    # First (threshold - 1) iterations skip silently.
+    for _ in range(threshold - 1):
+        algo.evaluate()
+    # The threshold-th iteration trips the check and raises.
+    with pytest.raises(
+        RuntimeError, match="evaluation_error_after_n_consecutive_skips"
+    ):
+        algo.evaluate()
+
+
+def test_timeout_waits_then_skips_when_no_recovery(parallel_to_training):
+    """timeout_s>0 with workers that never come back: evaluate() should
+    take at least roughly that long (waiting for recovery), then skip
+    silently because the threshold defaults to None."""
+    timeout_s = 2
+    algo = _algo_with_unhealthy_eval_workers(
+        timeout_s=timeout_s, parallel_to_training=parallel_to_training
+    )
+    start = time.monotonic()
+    algo.evaluate()
+    elapsed = time.monotonic() - start
+    # Check against lower bound but also a loose upper bound to sanity check
+    assert 5 > elapsed >= timeout_s
+
+
+def test_timeout_recovers_resyncs_and_evaluates(parallel_to_training):
+    """When a fresh replacement eval worker appears during the wait
+    window, ``evaluate()`` must re-sync weights *and* connector state to
+    it and then run eval normally (no skip, no raise, counter stays at 0).
+
+    The corresponding syncs at the top of ``evaluate()`` run before the
+    wait and skip workers that are unhealthy at that moment. Without the
+    post-recovery re-syncs inside ``_maybe_wait_for_eval_env_runner_recovery``,
+    the recovered worker would run eval with stale/initial weights and
+    empty/stale connector state (e.g. no obs-normalization filter), each
+    silently producing wrong metrics for one iteration.
+    """
+    timeout_s = 30  # generous; the first probe will revive quickly.
+    algo = _algo_with_unhealthy_eval_workers(
+        timeout_s=timeout_s,
+        parallel_to_training=parallel_to_training,
+        spawn_replacement=True,
+    )
+    eval_grp = algo.eval_env_runner_group
+
+    # Spy on `sync_weights` AND `sync_env_runner_states` (v2 stack) so we
+    # can verify both re-syncs after recovery. Each is expected to be
+    # called >=2 times in one `evaluate()`:
+    #   1) the unconditional sync at the start of `evaluate()` (effectively
+    #      a no-op here because no eval worker is healthy at that moment),
+    #   2) the post-recovery re-sync inside the wait method we target here.
+    with patch.object(
+        eval_grp, "sync_weights", wraps=eval_grp.sync_weights
+    ) as weights_spy, patch.object(
+        eval_grp,
+        "sync_env_runner_states",
+        wraps=eval_grp.sync_env_runner_states,
+    ) as states_spy:
+        start = time.monotonic()
+        algo.evaluate()
+        elapsed = time.monotonic() - start
+
+    # Recovery should be near-instant
+    assert elapsed < timeout_s
+    assert algo._counters["num_consecutive_eval_no_workers_iterations"] == 0
+    assert weights_spy.call_count == 2
+    assert states_spy.call_count == 2
+
+
+def test_timeout_recovers_resyncs_and_evaluates_old_stack(parallel_to_training):
+    """Same as ``test_timeout_recovers_resyncs_and_evaluates`` but for the *old* API stack"""
+    timeout_s = 30
+    algo = _algo_with_unhealthy_eval_workers(
+        timeout_s=timeout_s,
+        parallel_to_training=parallel_to_training,
+        spawn_replacement=True,
+        enable_v2_stack=False,
+    )
+    eval_grp = algo.eval_env_runner_group
+
+    # Spy on `sync_weights` (on the eval group) AND
+    # `_sync_filters_if_needed` (on the algo itself; that's where the
+    # method lives in the old API stack).
+    with patch.object(
+        eval_grp, "sync_weights", wraps=eval_grp.sync_weights
+    ) as weights_spy, patch.object(
+        algo,
+        "_sync_filters_if_needed",
+        wraps=algo._sync_filters_if_needed,
+    ) as filters_spy:
+        start = time.monotonic()
+        algo.evaluate()
+        elapsed = time.monotonic() - start
+
+    assert elapsed < timeout_s
+    assert algo._counters["num_consecutive_eval_no_workers_iterations"] == 0
+    assert weights_spy.call_count == 2
+    assert filters_spy.call_count == 2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_get_main_process_bundle.py

@@ -0,0 +1,47 @@
+import pytest
+
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.utils import _get_main_process_bundle
+
+
+def test_get_main_process_bundle():
+    # num_learners=0, so main process gets num_cpus_per_learner.
+    config = AlgorithmConfig().learners(
+        num_learners=0, num_cpus_per_learner=4, num_gpus_per_learner=1
+    )
+    bundle = _get_main_process_bundle(config)
+    assert bundle["CPU"] == 4
+    assert bundle["GPU"] == 1
+
+    # custom_resources_for_main_process included in bundle (num_learners>0).
+    config = (
+        AlgorithmConfig()
+        .resources(custom_resources_for_main_process={"my_resource": 1})
+        .learners(num_learners=1)
+    )
+    bundle = _get_main_process_bundle(config)
+    assert bundle == {"CPU": 1, "GPU": 0, "my_resource": 1}
+
+
+def test_default_resource_request_old_stack_custom_resources():
+    """custom_resources_for_main_process must appear in the placement group
+    even when using the old API stack (enable_rl_module_and_learner=False)."""
+    config = (
+        AlgorithmConfig()
+        .api_stack(
+            enable_rl_module_and_learner=False,
+            enable_env_runner_and_connector_v2=False,
+        )
+        .resources(custom_resources_for_main_process={"my_resource": 2})
+        .env_runners(num_env_runners=0)
+    )
+    pg_factory = Algorithm.default_resource_request(config)
+    main_bundle = pg_factory.bundles[0]
+    assert main_bundle["my_resource"] == 2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_nn_framework_import_errors.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+import os
+
+import pytest
+
+import ray.rllib.algorithms.ppo as ppo
+
+
+def test_dont_import_torch_error():
+    """Check error being thrown, if torch not installed but configured."""
+    # Do not import tf for testing purposes.
+    os.environ["RLLIB_TEST_NO_TORCH_IMPORT"] = "1"
+    config = (
+        ppo.PPOConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("CartPole-v1")
+        .framework("torch")
+    )
+    with pytest.raises(ImportError, match="However, no installation was found"):
+        config.build()
+
+
+if __name__ == "__main__":
+    test_dont_import_torch_error()

rllib/algorithms/tests/test_node_failures.py

@@ -0,0 +1,230 @@
+import pytest
+
+import ray
+import ray._common
+from ray._private.test_utils import get_other_nodes
+from ray.cluster_utils import Cluster
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.algorithms.ppo import PPOConfig
+
+EXPECTED_PER_NODE_OBJECT_STORE_MEMORY = 10**8
+HEAD_REDIS_PORT = 6379
+HEAD_CPUS = 2
+WORKER_CPUS = 4
+NUM_ENV_RUNNERS = 4
+
+
+def _add_node(cluster, worker=True):
+    cluster.add_node(
+        redis_port=None if worker else HEAD_REDIS_PORT,
+        num_cpus=WORKER_CPUS if worker else HEAD_CPUS,
+        object_store_memory=EXPECTED_PER_NODE_OBJECT_STORE_MEMORY,
+        include_dashboard=not worker,
+    )
+
+
+@pytest.fixture
+def cluster():
+    """Create a 2-node fake cluster: head (2 CPUs) + worker (4 CPUs).
+
+    Head holds the algo process + local env runner.
+    Worker holds all 4 remote env runners (deterministic placement).
+    """
+    assert (
+        2 * EXPECTED_PER_NODE_OBJECT_STORE_MEMORY
+        < ray._common.utils.get_system_memory() / 2
+    ), "Not enough memory on this machine to run this workload."
+
+    cluster = Cluster()
+    _add_node(cluster, worker=False)
+    _add_node(cluster)
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    yield cluster
+
+    # Detach head_node before cluster.shutdown() and kill its processes
+    # directly afterwards. Otherwise cluster.remove_node(head) refuses if a
+    # daemon thread (e.g. APPO/IMPALA's learner thread) called an
+    # auto-init-wrapped Ray API and re-established global_worker.node after
+    # our ray.shutdown(), leaving port 8265 bound for the next test.
+    ray.shutdown()
+    head = cluster.head_node
+    cluster.head_node = None
+    cluster.shutdown()
+    head.kill_all_processes(check_alive=False, allow_graceful=False, wait=True)
+
+
+def _kill_worker_node(cluster):
+    others = get_other_nodes(cluster, exclude_head=True)
+    if others:
+        cluster.remove_node(others[0])
+        return True
+    return False
+
+
+def _train(cluster, algo, config, iters, preempt_freq):
+    """Train loop with periodic node kill/restore and health tracking."""
+    num_runners = config.num_env_runners
+    saw_healthy_drop = False
+    saw_recovery = False
+
+    for i in range(iters):
+        algo.train()
+
+        assert algo.env_runner_group.num_remote_env_runners() == num_runners
+        healthy = algo.env_runner_group.num_healthy_remote_workers()
+        assert 0 <= healthy <= num_runners
+
+        if healthy < num_runners:
+            saw_healthy_drop = True
+        if saw_healthy_drop and healthy == num_runners:
+            saw_recovery = True
+
+        print(
+            f"ITER={i}, healthy={healthy}/{num_runners}, "
+            f"saw_drop={saw_healthy_drop}, saw_recovery={saw_recovery}"
+        )
+
+        # Shut down one node every preempt_freq iterations.
+        if i % preempt_freq == 0:
+            _kill_worker_node(cluster)
+
+        # Bring back a previously failed node.
+        elif (i - 1) % preempt_freq == 0:
+            _add_node(cluster)
+
+    # Workers must have gone down at some point.
+    assert saw_healthy_drop, (
+        "Expected healthy worker count to drop after node kill, " "but it never did."
+    )
+    # If restart is enabled, workers must have come back.
+    if config.restart_failed_env_runners:
+        assert saw_recovery, (
+            "Expected workers to recover after node restore "
+            "(restart_failed_env_runners=True), but they never did."
+        )
+    # If restart is disabled, workers must NOT have come back.
+    if not config.restart_failed_env_runners:
+        assert (
+            not saw_recovery
+        ), "Workers recovered despite restart_failed_env_runners=False."
+
+
+def test_node_failure_ignore(cluster):
+    """restart=False, ignore=True: workers die and stay dead, training
+    continues with fewer workers."""
+    config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .env_runners(
+            num_env_runners=NUM_ENV_RUNNERS,
+            sample_timeout_s=5.0,
+        )
+        .training(
+            train_batch_size=500,
+            num_epochs=1,
+            minibatch_size=500,
+        )
+        .reporting(min_train_timesteps_per_iteration=1)
+        .fault_tolerance(
+            ignore_env_runner_failures=True,
+            restart_failed_env_runners=False,
+            env_runner_health_probe_timeout_s=20.0,
+        )
+    )
+
+    algo = config.build()
+    _train(cluster, algo, config, iters=10, preempt_freq=3)
+
+
+def test_node_failure_recreate_appo(cluster):
+    """restart=True with APPO (async): workers die, get auto-restarted by Ray,
+    and restore_env_runners() syncs their state."""
+    config = (
+        APPOConfig()
+        .environment("CartPole-v1")
+        .learners(num_learners=0)
+        .experimental(_validate_config=False)
+        .env_runners(
+            num_env_runners=NUM_ENV_RUNNERS,
+        )
+        .reporting(
+            # Must be >= 2s so APPO's async mechanism has time to detect
+            # worker death within a single iteration.
+            min_time_s_per_iteration=2,
+            min_train_timesteps_per_iteration=1,
+        )
+        .fault_tolerance(
+            restart_failed_env_runners=True,
+            env_runner_health_probe_timeout_s=20.0,
+        )
+    )
+
+    algo = config.build()
+    _train(cluster, algo, config, iters=10, preempt_freq=7)
+
+
+def test_node_failure_recreate_ppo(cluster):
+    """restart=True with PPO (sync): workers die, get auto-restarted by Ray,
+    and restore_env_runners() syncs their state."""
+    config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .learners(num_learners=0)
+        .env_runners(
+            num_env_runners=NUM_ENV_RUNNERS,
+            sample_timeout_s=5.0,
+        )
+        .training(
+            train_batch_size=500,
+            num_epochs=1,
+            minibatch_size=500,
+        )
+        .reporting(
+            min_time_s_per_iteration=2,
+            min_train_timesteps_per_iteration=1,
+        )
+        .fault_tolerance(
+            restart_failed_env_runners=True,
+            env_runner_health_probe_timeout_s=20.0,
+        )
+    )
+
+    algo = config.build()
+    _train(cluster, algo, config, iters=10, preempt_freq=7)
+
+
+def test_node_failure_no_recovery(cluster):
+    """restart=False, ignore=False: dead worker RayErrors propagate and
+    crash training. Verify the crash happens."""
+    config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .env_runners(
+            num_env_runners=NUM_ENV_RUNNERS,
+            sample_timeout_s=5.0,
+        )
+        .training(
+            train_batch_size=500,
+            num_epochs=1,
+            minibatch_size=500,
+        )
+        .reporting(min_train_timesteps_per_iteration=1)
+        .fault_tolerance(
+            restart_failed_env_runners=False,
+            env_runner_health_probe_timeout_s=20.0,
+        )
+    )
+
+    algo = config.build()
+    # _train will crash with an ActorDiedError when dead workers are detected
+    # (ignore=False, restart=False → errors propagate).
+    with pytest.raises(ray.exceptions.ActorDiedError, match="actor died unexpectedly"):
+        _train(cluster, algo, config, iters=10, preempt_freq=3)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_placement_groups.py

@@ -0,0 +1,133 @@
+import os
+import unittest
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.ppo import PPO, PPOConfig
+from ray.tune import Callback
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Trial
+from ray.tune.result import TRAINING_ITERATION
+
+trial_executor = None
+
+
+class _TestCallback(Callback):
+    def on_step_end(self, iteration, trials, **info):
+        num_running = len([t for t in trials if t.status == Trial.RUNNING])
+
+        # All 3 trials (3 different learning rates) should be scheduled.
+        assert 3 == min(3, len(trials))
+        # Cannot run more than 2 at a time
+        # (due to different resource restrictions in the test cases).
+        assert num_running <= 2
+
+
+class TestPlacementGroups(unittest.TestCase):
+    def setUp(self) -> None:
+        os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
+        ray.init(num_cpus=6)
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_overriding_default_resource_request(self):
+        # 3 Trials: Can only run 2 at a time (num_cpus=6; needed: 3).
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .training(
+                model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
+            )
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=2)
+            .framework("tf")
+        )
+
+        # Create an Algorithm with an overridden default_resource_request
+        # method that returns a PlacementGroupFactory.
+
+        class MyAlgo(PPO):
+            @classmethod
+            def default_resource_request(cls, config):
+                head_bundle = {"CPU": 1, "GPU": 0}
+                child_bundle = {"CPU": 1}
+                return PlacementGroupFactory(
+                    [head_bundle, child_bundle, child_bundle],
+                    strategy=config["placement_strategy"],
+                )
+
+        tune.register_trainable("my_trainable", MyAlgo)
+
+        tune.Tuner(
+            "my_trainable",
+            param_space=config,
+            run_config=tune.RunConfig(
+                stop={TRAINING_ITERATION: 2},
+                verbose=2,
+                callbacks=[_TestCallback()],
+            ),
+        ).fit()
+
+    def test_default_resource_request(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .resources(placement_strategy="SPREAD")
+            .env_runners(
+                num_env_runners=2,
+                num_cpus_per_env_runner=2,
+            )
+            .training(
+                model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
+            )
+            .environment("CartPole-v1")
+            .framework("torch")
+        )
+        # 3 Trials: Can only run 1 at a time (num_cpus=6; needed: 5).
+
+        tune.Tuner(
+            PPO,
+            param_space=config,
+            run_config=tune.RunConfig(
+                stop={TRAINING_ITERATION: 2},
+                verbose=2,
+                callbacks=[_TestCallback()],
+            ),
+            tune_config=tune.TuneConfig(reuse_actors=False),
+        ).fit()
+
+    def test_default_resource_request_plus_manual_leads_to_error(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .training(model={"fcnet_hiddens": [10]})
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=0)
+        )
+
+        try:
+            tune.Tuner(
+                tune.with_resources(PPO, PlacementGroupFactory([{"CPU": 1}])),
+                param_space=config,
+                run_config=tune.RunConfig(stop={TRAINING_ITERATION: 2}, verbose=2),
+            ).fit()
+        except ValueError as e:
+            assert "have been automatically set to" in e.args[0]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_ray_client.py

@@ -0,0 +1,27 @@
+import sys
+
+import pytest
+
+from ray._private.client_mode_hook import client_mode_should_convert, enable_client_mode
+from ray.rllib.algorithms import dqn
+from ray.util.client.ray_client_helpers import ray_start_client_server
+
+
+def test_basic_dqn():
+    with ray_start_client_server():
+        # Need to enable this for client APIs to be used.
+        with enable_client_mode():
+            # Confirming mode hook is enabled.
+            assert client_mode_should_convert()
+            config = (
+                dqn.DQNConfig()
+                .environment("CartPole-v1")
+                .env_runners(num_env_runners=0, compress_observations=True)
+            )
+            trainer = config.build()
+            for i in range(2):
+                trainer.train()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

rllib/algorithms/tests/test_registry.py

@@ -0,0 +1,38 @@
+import unittest
+
+from ray.rllib.algorithms.registry import (
+    ALGORITHMS,
+    ALGORITHMS_CLASS_TO_NAME,
+    POLICIES,
+    get_policy_class,
+    get_policy_class_name,
+)
+
+
+class TestPolicies(unittest.TestCase):
+    def test_load_policies(self):
+        for name in POLICIES.keys():
+            self.assertIsNotNone(get_policy_class(name))
+
+    def test_get_eager_traced_class_name(self):
+        from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF2Policy
+
+        traced = PPOTF2Policy.with_tracing()
+        self.assertEqual(get_policy_class_name(traced), "PPOTF2Policy")
+
+    def test_registered_algorithm_names(self):
+        """All RLlib registered algorithms should have their name listed in the
+        registry dictionary."""
+
+        for class_name in ALGORITHMS_CLASS_TO_NAME.keys():
+            registered_name = ALGORITHMS_CLASS_TO_NAME[class_name]
+            algo_class, _ = ALGORITHMS[registered_name]()
+            self.assertEqual(class_name.upper(), algo_class.__name__.upper())
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/tests/test_telemetry.py

@@ -0,0 +1,46 @@
+import sys
+
+import pytest
+
+import ray
+import ray._common.usage.usage_lib as ray_usage_lib
+from ray._common.test_utils import TelemetryCallsite, check_library_usage_telemetry
+
+
+@pytest.fixture
+def reset_usage_lib():
+    yield
+    ray.shutdown()
+    ray_usage_lib.reset_global_state()
+
+
+@pytest.mark.skip(reason="Usage currently marked on import.")
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_not_used_on_import(reset_usage_lib, callsite: TelemetryCallsite):
+    def _import_rllib():
+        from ray import rllib  # noqa: F401
+
+    check_library_usage_telemetry(
+        _import_rllib, callsite=callsite, expected_library_usages=[set(), {"core"}]
+    )
+
+
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_used_on_import(reset_usage_lib, callsite: TelemetryCallsite):
+    def _use_rllib():
+        # TODO(edoakes): this test currently fails if we don't call `ray.init()`
+        # prior to the import. This is a bug.
+        if callsite == TelemetryCallsite.DRIVER:
+            ray.init()
+
+        from ray import rllib  # noqa: F401
+
+    check_library_usage_telemetry(
+        _use_rllib,
+        callsite=callsite,
+        expected_library_usages=[{"rllib"}, {"core", "rllib"}],
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))