chore: import upstream snapshot with attribution

rllib/examples/algorithms/appo/cartpole_appo.py

@@ -0,0 +1,31 @@
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=450.0,
+    default_timesteps=2000000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+
+config = (
+    APPOConfig()
+    .environment("CartPole-v1")
+    .training(
+        circular_buffer_iterations_per_batch=2,
+        vf_loss_coeff=0.05,
+        entropy_coeff=0.0,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(vf_share_layers=True),
+    )
+)
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/appo/multi_agent_cartpole_appo.py

@@ -0,0 +1,49 @@
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(default_timesteps=2000000)
+parser.set_defaults(
+    num_agents=2,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("env", lambda cfg: MultiAgentCartPole(config=cfg))
+
+
+config = (
+    APPOConfig()
+    .environment("env", env_config={"num_agents": args.num_agents})
+    .training(
+        vf_loss_coeff=0.005,
+        entropy_coeff=0.0,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(vf_share_layers=True),
+    )
+    .multi_agent(
+        policy_mapping_fn=(lambda agent_id, episode, **kwargs: f"p{agent_id}"),
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 350.0 * args.num_agents,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps,
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/appo/multi_agent_footsies_appo.py

@@ -0,0 +1,254 @@
+"""
+Multi-agent RLlib Footsies Example (APPO)
+
+About:
+    - Example is based on the Footsies environment (https://github.com/chasemcd/FootsiesGym).
+    - Footsies is a two-player fighting game where each player controls a character and tries to hit the opponent while avoiding being hit.
+    - Footsies is a zero-sum game, when one player wins (+1 reward) the other loses (-1 reward).
+
+Summary:
+    - Main policy is an LSTM-based policy.
+    - Training algorithm is APPO.
+
+Training:
+    - Training is governed by adding new, more complex opponents to the mix as the main policy reaches a certain win rate threshold against the current opponent.
+    - Current opponent is always the newest opponent added to the mix.
+    - Training starts with a very simple opponent: "noop" (does nothing), then progresses to "back" (only moves backwards). These are the fixed (very simple) policies that are used to kick off the training.
+    - New opponents are frozen copies of the main policy at different training stages. They will be added to the mix as "lstm_v0", "lstm_v1", etc.
+    - In this way - after kick-starting the training with fixed simple opponents - the main policy will play against a version of itself from an earlier training stage.
+    - The main policy has to achieve the win rate threshold against the current opponent to add a new opponent to the mix.
+    - Training concludes when the target mix size is reached.
+
+Evaluation:
+    - Evaluation is performed against the current (newest) opponent.
+    - Evaluation runs for a fixed number of episodes at the end of each training iteration.
+
+"""
+import functools
+from pathlib import Path
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.core.rl_module import MultiRLModuleSpec, RLModuleSpec
+from ray.rllib.env.multi_agent_env_runner import MultiAgentEnvRunner
+from ray.rllib.examples.envs.classes.multi_agent.footsies.fixed_rlmodules import (
+    BackFixedRLModule,
+    NoopFixedRLModule,
+)
+from ray.rllib.examples.envs.classes.multi_agent.footsies.footsies_env import (
+    env_creator,
+)
+from ray.rllib.examples.envs.classes.multi_agent.footsies.utils import (
+    Matchmaker,
+    Matchup,
+    MetricsLoggerCallback,
+    MixManagerCallback,
+    platform_for_binary_to_download,
+)
+from ray.rllib.examples.rl_modules.classes.lstm_containing_rlm import (
+    LSTMContainingRLModuleWithTargetNetwork,
+)
+from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME
+from ray.rllib.utils.test_utils import (
+    add_rllib_example_script_args,
+)
+from ray.tune.registry import register_env
+from ray.tune.result import TRAINING_ITERATION
+
+parser = add_rllib_example_script_args(
+    default_iters=500,
+    default_timesteps=5_000_000,
+)
+
+parser.add_argument(
+    "--train-start-port",
+    type=int,
+    default=45001,
+    help="First port number for the Footsies training environment server (default: 45001). Each server gets its own port.",
+)
+parser.add_argument(
+    "--eval-start-port",
+    type=int,
+    default=55001,
+    help="First port number for the Footsies evaluation environment server (default: 55001) Each server gets its own port.",
+)
+parser.add_argument(
+    "--binary-download-dir",
+    type=Path,
+    default="/tmp/ray/binaries/footsies",
+    help="Directory to download Footsies binaries (default: /tmp/ray/binaries/footsies)",
+)
+parser.add_argument(
+    "--binary-extract-dir",
+    type=Path,
+    default="/tmp/ray/binaries/footsies",
+    help="Directory to extract Footsies binaries (default: /tmp/ray/binaries/footsies)",
+)
+parser.add_argument(
+    "--win-rate-threshold",
+    type=float,
+    default=0.8,
+    help="The main policy should have at least 'win-rate-threshold' win rate against the "
+    "other policy to advance to the next level. Moving to the next level "
+    "means adding a new policy to the mix.",
+)
+parser.add_argument(
+    "--target-mix-size",
+    type=int,
+    default=5,
+    help="Target number of policies (RLModules) in the mix to consider the test passed. "
+    "The initial mix size is 2: 'main policy' vs. 'other'. "
+    "`--target-mix-size=5` means that 3 new policies will be added to the mix. "
+    "Whether to add new policy is decided by checking the '--win-rate-threshold' condition. ",
+)
+parser.add_argument(
+    "--rollout-fragment-length",
+    type=int,
+    default=256,
+    help="The length of each rollout fragment to be collected by the EnvRunners when sampling.",
+)
+parser.add_argument(
+    "--log-unity-output",
+    action="store_true",
+    help="Whether to log Unity output (from the game engine). Default is False.",
+    default=False,
+)
+parser.add_argument(
+    "--render",
+    action="store_true",
+    default=False,
+    help="Whether to render the Footsies environment. Default is False.",
+)
+
+main_policy = "lstm"
+args = parser.parse_args()
+register_env(name="FootsiesEnv", env_creator=env_creator)
+
+# Detect platform and choose appropriate binary
+binary_to_download = platform_for_binary_to_download(args.render)
+
+config = (
+    APPOConfig()
+    .reporting(
+        min_time_s_per_iteration=30,
+    )
+    .environment(
+        env="FootsiesEnv",
+        env_config={
+            "max_t": 1000,
+            "frame_skip": 4,
+            "observation_delay": 16,
+            "train_start_port": args.train_start_port,
+            "eval_start_port": args.eval_start_port,
+            "host": "localhost",
+            "binary_download_dir": args.binary_download_dir,
+            "binary_extract_dir": args.binary_extract_dir,
+            "binary_to_download": binary_to_download,
+            "log_unity_output": args.log_unity_output,
+        },
+    )
+    .learners(
+        num_learners=1,
+        num_cpus_per_learner=1,
+        num_gpus_per_learner=0,
+        num_aggregator_actors_per_learner=2,
+    )
+    .env_runners(
+        env_runner_cls=MultiAgentEnvRunner,
+        num_env_runners=args.num_env_runners or 1,
+        num_cpus_per_env_runner=1,
+        num_envs_per_env_runner=1,
+        batch_mode="truncate_episodes",
+        rollout_fragment_length=args.rollout_fragment_length,
+        episodes_to_numpy=True,
+        create_env_on_local_worker=False,
+    )
+    .training(
+        train_batch_size_per_learner=4096 * (args.num_env_runners or 1),
+        lr=1e-4,
+        entropy_coeff=0.01,
+    )
+    .multi_agent(
+        policies={
+            main_policy,
+            "noop",
+            "back",
+        },
+        # this is a starting policy_mapping_fn
+        # It will be updated by the MixManagerCallback during training.
+        policy_mapping_fn=Matchmaker(
+            [Matchup(main_policy, "noop", 1.0)]
+        ).agent_to_module_mapping_fn,
+        # we only train the main policy, this doesn't change during training.
+        policies_to_train=[main_policy],
+    )
+    .rl_module(
+        rl_module_spec=MultiRLModuleSpec(
+            rl_module_specs={
+                main_policy: RLModuleSpec(
+                    module_class=LSTMContainingRLModuleWithTargetNetwork,
+                    model_config={
+                        "lstm_cell_size": 128,
+                        "dense_layers": [128, 128],
+                        "max_seq_len": 64,
+                    },
+                ),
+                # for simplicity, all fixed RLModules are added to the config at the start.
+                # However, only "noop" is used at the start of training,
+                # the others are added to the mix later by the MixManagerCallback.
+                "noop": RLModuleSpec(module_class=NoopFixedRLModule),
+                "back": RLModuleSpec(module_class=BackFixedRLModule),
+            },
+        )
+    )
+    .evaluation(
+        evaluation_num_env_runners=args.evaluation_num_env_runners or 1,
+        evaluation_sample_timeout_s=120,
+        evaluation_interval=1,
+        evaluation_duration=10,  # 10 episodes is enough to get a good win rate estimate
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=False,
+        # we may add new RLModules to the mix at the end of the evaluation stage.
+        # Running evaluation in parallel may result in training for one more iteration on the old mix.
+        evaluation_force_reset_envs_before_iteration=True,
+        evaluation_config={
+            "env_config": {"env-for-evaluation": True},
+        },  # evaluation_config is used to add an argument to the env creator.
+    )
+    .callbacks(
+        [
+            functools.partial(
+                MetricsLoggerCallback,
+                main_policy=main_policy,
+            ),
+            functools.partial(
+                MixManagerCallback,
+                win_rate_threshold=args.win_rate_threshold,
+                main_policy=main_policy,
+                target_mix_size=args.target_mix_size,
+                starting_modules=[main_policy, "noop"],
+                fixed_modules_progression_sequence=(
+                    "noop",
+                    "back",
+                ),
+            ),
+        ]
+    )
+)
+
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    TRAINING_ITERATION: args.stop_iters,
+    "mix_size": args.target_mix_size,
+}
+if __name__ == "__main__":
+    from ray.rllib.utils.test_utils import run_rllib_example_script_experiment
+
+    results = run_rllib_example_script_experiment(
+        base_config=config,
+        args=args,
+        stop=stop,
+        success_metric={
+            "mix_size": args.target_mix_size
+        },  # pass the success metric for RLlib's testing framework
+    )

rllib/examples/algorithms/appo/multi_agent_pong_appo.py

@@ -0,0 +1,110 @@
+import random
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule
+from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.env.multi_agent_env import make_multi_agent
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.rl_modules.classes.random_rlm import RandomRLModule
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=0.0,
+    default_timesteps=20000000,
+    default_iters=400,
+)
+parser.set_defaults(
+    env="ale_py:ALE/Pong-v5",
+    num_agents=2,
+)
+args = parser.parse_args()
+
+
+def _make_env_to_module_connector(env, spaces, device):
+    return FrameStackingEnvToModule(num_frames=4, multi_agent=True)
+
+
+def _make_learner_connector(input_observation_space, input_action_space):
+    return FrameStackingLearner(num_frames=4, multi_agent=True)
+
+
+def _env_creator(cfg):
+    return wrap_atari_for_new_api_stack(
+        gym.make(args.env, **cfg, **{"render_mode": "rgb_array"}),
+        dim=64,
+        framestack=None,
+    )
+
+
+MultiAgentPong = make_multi_agent(_env_creator)
+NUM_POLICIES = 5
+main_spec = RLModuleSpec(
+    model_config=DefaultModelConfig(
+        vf_share_layers=True,
+        conv_filters=[(16, 4, 2), (32, 4, 2), (64, 4, 2), (128, 4, 2)],
+        conv_activation="relu",
+        head_fcnet_hiddens=[256],
+    ),
+)
+
+
+config = (
+    APPOConfig()
+    .environment(
+        MultiAgentPong,
+        env_config={
+            "num_agents": args.num_agents,
+            # Make analogous to old v4 + NoFrameskip.
+            "frameskip": 1,
+            "full_action_space": False,
+            "repeat_action_probability": 0.0,
+        },
+        clip_rewards=True,
+    )
+    .env_runners(
+        env_to_module_connector=_make_env_to_module_connector,
+    )
+    .learners(
+        num_aggregator_actors_per_learner=2,
+    )
+    .training(
+        learner_connector=_make_learner_connector,
+        train_batch_size_per_learner=500,
+        target_network_update_freq=2,
+        lr=0.0005 * ((args.num_learners or 1) ** 0.5),
+        vf_loss_coeff=1.0,
+        entropy_coeff=[[0, 0.01], [3000000, 0.0]],  # <- crucial parameter to finetune
+        # Only update connector states and model weights every n training_step calls.
+        broadcast_interval=5,
+        # learner_queue_size=1,
+        circular_buffer_num_batches=4,
+        circular_buffer_iterations_per_batch=2,
+    )
+    .rl_module(
+        rl_module_spec=MultiRLModuleSpec(
+            rl_module_specs=(
+                {f"p{i}": main_spec for i in range(NUM_POLICIES)}
+                | {"random": RLModuleSpec(module_class=RandomRLModule)}
+            ),
+        ),
+    )
+    .multi_agent(
+        policies={f"p{i}" for i in range(NUM_POLICIES)} | {"random"},
+        policy_mapping_fn=lambda aid, eps, **kw: (
+            random.choice([f"p{i}" for i in range(NUM_POLICIES)] + ["random"])
+        ),
+        policies_to_train=[f"p{i}" for i in range(NUM_POLICIES)],
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/appo/multi_agent_stateless_cartpole_appo.py

@@ -0,0 +1,61 @@
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentStatelessCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(default_timesteps=2000000)
+parser.set_defaults(
+    num_agents=2,
+    num_env_runners=6,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("env", lambda cfg: MultiAgentStatelessCartPole(config=cfg))
+
+
+config = (
+    APPOConfig()
+    .environment("env", env_config={"num_agents": args.num_agents})
+    # TODO (sven): Need to fix the MeanStdFilter(). It seems to cause NaNs when
+    #  training.
+    # .env_runners(
+    #    env_to_module_connector=lambda env, spaces, device: MeanStdFilter(multi_agent=True),
+    # )
+    .training(
+        train_batch_size_per_learner=600,
+        lr=0.0005 * ((args.num_learners or 1) ** 0.5),
+        num_epochs=1,
+        vf_loss_coeff=0.05,
+        entropy_coeff=0.005,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            vf_share_layers=True,
+            use_lstm=True,
+            max_seq_len=20,
+        ),
+    )
+    .multi_agent(
+        policy_mapping_fn=(lambda agent_id, episode, **kwargs: f"p{agent_id}"),
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 150.0 * args.num_agents,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/appo/pendulum_appo.py

@@ -0,0 +1,48 @@
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=-300.0,
+    default_timesteps=100000000,
+)
+parser.set_defaults(
+    num_env_runners=4,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+
+config = (
+    APPOConfig()
+    .environment("Pendulum-v1")
+    .env_runners(
+        num_envs_per_env_runner=20,
+    )
+    .learners(num_learners=1)
+    .training(
+        train_batch_size_per_learner=500,
+        circular_buffer_num_batches=16,
+        circular_buffer_iterations_per_batch=10,
+        target_network_update_freq=2,
+        clip_param=0.4,
+        lr=0.0003,
+        gamma=0.95,
+        lambda_=0.5,
+        entropy_coeff=0.0,
+        use_kl_loss=True,
+        kl_coeff=1.0,
+        kl_target=0.04,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(fcnet_activation="relu"),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/appo/pong_appo.py

@@ -0,0 +1,87 @@
+import gymnasium as gym
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule
+from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(
+    default_reward=20.0,
+    default_timesteps=10_000_000,
+)
+parser.set_defaults(
+    env="ale_py:ALE/Pong-v5",
+)
+args = parser.parse_args()
+
+
+def _make_env_to_module_connector(env, spaces, device):
+    return FrameStackingEnvToModule(num_frames=4)
+
+
+def _make_learner_connector(input_observation_space, input_action_space):
+    return FrameStackingLearner(num_frames=4)
+
+
+def _env_creator(cfg):
+    return wrap_atari_for_new_api_stack(
+        gym.make(args.env, **cfg, **{"render_mode": "rgb_array"}),
+        dim=64,
+        framestack=None,
+    )
+
+
+register_env("env", _env_creator)
+
+
+config = (
+    APPOConfig()
+    .environment(
+        "env",
+        env_config={
+            # Make analogous to old v4 + NoFrameskip.
+            "frameskip": 1,
+            "full_action_space": False,
+            "repeat_action_probability": 0.0,
+        },
+        clip_rewards=True,
+    )
+    .env_runners(
+        env_to_module_connector=_make_env_to_module_connector,
+        num_envs_per_env_runner=2,
+    )
+    .learners(
+        num_aggregator_actors_per_learner=2,
+    )
+    .training(
+        learner_connector=_make_learner_connector,
+        train_batch_size_per_learner=500,
+        target_network_update_freq=2,
+        lr=0.0005 * ((args.num_learners or 1) ** 0.5),
+        vf_loss_coeff=1.0,
+        entropy_coeff=[[0, 0.01], [3000000, 0.0]],  # <- crucial parameter to finetune
+        # Only update connector states and model weights every n training_step calls.
+        broadcast_interval=5,
+        # learner_queue_size=1,
+        circular_buffer_num_batches=4,
+        circular_buffer_iterations_per_batch=2,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            vf_share_layers=True,
+            conv_filters=[(16, 4, 2), (32, 4, 2), (64, 4, 2), (128, 4, 2)],
+            conv_activation="relu",
+            head_fcnet_hiddens=[256],
+        )
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/appo/stateless_cartpole_appo.py

@@ -0,0 +1,47 @@
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=2000000,
+    default_reward=300.0,
+)
+parser.set_defaults(
+    num_env_runners=3,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+
+config = (
+    APPOConfig()
+    .environment(StatelessCartPole)
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+    )
+    .training(
+        lr=0.0005 * ((args.num_learners or 1) ** 0.5),
+        num_epochs=1,
+        vf_loss_coeff=0.05,
+        entropy_coeff=0.005,
+        use_circular_buffer=False,
+        broadcast_interval=10,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            vf_share_layers=True,
+            use_lstm=True,
+            max_seq_len=20,
+        ),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/appo_custom_algorithm_w_shared_data_actor.py

@@ -0,0 +1,183 @@
+"""Example of how to write a custom APPO that uses a global shared data actor.
+
+The actor is custom code and its remote APIs can be designed as the user requires.
+It is created inside the Algorithm's `setup` method and then shared through its
+reference with all of the Algorithm's other actors, like EnvRunners, Learners, and
+aggregator actors.
+
+During sampling and through using callbacks, each EnvRunner assigns a unique ID
+to each sampled episode chunk, then sends manipulated reward data for each sampled
+episode chunk to the shared data actor. In particular, the manipulation consists of
+each individual reward being multiplied by the EnvRunner's index (from 1 to ...).
+Note that the actual reward in the episode is not altered and thus the metrics
+reporting continues to show the original reward.
+
+In the learner connector, which creates the train batch from episode data, a custom
+connector piece then gets the manipulated rewards from the shared data actor using
+the episode chunk's unique ID (see above) and uses the manipulated reward for training.
+Note that because of this, different EnvRunners provide different reward signals, which
+should make it slightly harder for the value function to learn consistently.
+Nevertheless, because the default config here only uses 2 EnvRunners, each multiplying
+their rewards by 1 and 2, respectively, this effect is negligible here and the example
+should learn how to solve the CartPole-1 env either way.
+
+This example shows:
+
+    - how to write a custom, global shared data actor class with a custom remote API.
+    - how an instance of this shared data actor is created upon algorithm
+    initialization.
+    - how to distribute the actor reference of the shared actor to all other actors
+    in the Algorithm, for example EnvRunners, AggregatorActors, and Learners
+    - how to subclass an existing algorithm class (APPO) to implement a custom
+    Algorithm, overriding the `setup` method to control, which additional actors
+    should be created (and shared) by the algo, the `get_state/set_state` methods
+    to include the state of the new actor.
+    - how - through custom callbacks - the new actor can be written to and queried
+    from anywhere within the algorithm, for example its EnvRunner actors or Learners.
+
+
+How to run this script
+----------------------
+`python [script file name].py`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+The experiment should work regardless of whether you are using aggregator
+actors or not. By default, the experiment provides one agg. actor per Learner,
+but you can set `--num-aggregator-actors-per-learner=0` to have the learner
+connector pipeline work directly inside the Learner actor(s).
+
++-------------------------------------------------+------------+--------+
+| Trial name                                      | status     |   iter |
+|                                                 |            |        |
+|-------------------------------------------------+------------+--------+
+| APPOWithSharedDataActor_CartPole-v1_4e860_00000 | TERMINATED |      7 |
++-------------------------------------------------+------------+--------+
++------------------+------------------------+
+|   total time (s) |    episode_return_mean |
+|                  |                        |
+|------------------+------------------------+
+|          70.0315 |                 468.42 |
++------------------+------------------------+
+"""
+
+import uuid
+
+import ray
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.connectors.connector_v2 import ConnectorV2
+from ray.rllib.core import Columns
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.rllib.examples.algorithms.classes.appo_w_shared_data_actor import (
+    APPOWithSharedDataActor,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=450.0,
+    default_iters=200,
+    default_timesteps=2000000,
+)
+parser.set_defaults(
+    num_aggregator_actors_per_learner=1,
+)
+
+SPECIAL_REWARDS_KEY = "special_(double)_rewards"
+ENV_RUNNER_IDX_KEY = "env_runner_index"
+UNIQUE_EPISODE_CHUNK_KEY = "unique_eps_chunk"
+
+
+# Define 2 simple EnvRunner-based callbacks:
+
+
+def on_episode_step(*, episode, env_runner, **kwargs):
+    # Multiplies the received reward by the env runner index.
+    if SPECIAL_REWARDS_KEY not in episode.custom_data:
+        episode.custom_data[SPECIAL_REWARDS_KEY] = []
+    episode.custom_data[SPECIAL_REWARDS_KEY].append(
+        episode.get_rewards(-1) * env_runner.worker_index
+    )
+
+
+def on_sample_end(*, samples, env_runner, **kwargs):
+    # Sends the (manipulated) reward sequence to the shared data actor for "pickup" by
+    # a Learner. Alternatively, one could also just store the information in the
+    # `custom_data` property.
+    for episode in samples:
+        # Provide a unique key for both episode AND record in the shared
+        # data actor.
+        unique_key = str(uuid.uuid4())
+
+        # Store the EnvRunner index and unique key in the episode.
+        episode.custom_data[ENV_RUNNER_IDX_KEY] = env_runner.worker_index
+        episode.custom_data[UNIQUE_EPISODE_CHUNK_KEY] = unique_key
+
+        # Get the manipulated rewards from the episode ..
+        special_rewards = episode.custom_data.pop(SPECIAL_REWARDS_KEY)
+        # .. and send them under the unique key to the shared data actor.
+        env_runner._shared_data_actor.put.remote(
+            key=unique_key,
+            value=special_rewards,
+        )
+
+
+class ManipulatedRewardConnector(ConnectorV2):
+    def __call__(self, *, episodes, batch, metrics, **kwargs):
+        if not isinstance(episodes[0], SingleAgentEpisode):
+            raise ValueError("This connector only works on `SingleAgentEpisodes`.")
+        # Get the manipulated rewards from the shared actor and add them to the train
+        # batch.
+        for sa_episode in self.single_agent_episode_iterator(episodes):
+            unique_key = sa_episode.custom_data[UNIQUE_EPISODE_CHUNK_KEY]
+            special_rewards = ray.get(
+                self._shared_data_actor.get.remote(unique_key, delete=True)
+            )
+            if special_rewards is None:
+                continue
+
+            assert int(special_rewards[0]) == sa_episode.custom_data[ENV_RUNNER_IDX_KEY]
+
+            # Add one more fake reward, b/c all episodes will be extended
+            # (in PPO-style algos) by one artificial timestep for GAE/v-trace
+            # computation purposes.
+            special_rewards += [0.0]
+            self.add_n_batch_items(
+                batch=batch,
+                column=Columns.REWARDS,
+                items_to_add=special_rewards[-len(sa_episode) :],
+                num_items=len(sa_episode),
+                single_agent_episode=sa_episode,
+            )
+
+        return batch
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    base_config = (
+        APPOConfig(algo_class=APPOWithSharedDataActor)
+        .environment("CartPole-v1")
+        .callbacks(
+            on_episode_step=on_episode_step,
+            on_sample_end=on_sample_end,
+        )
+        .training(
+            learner_connector=(lambda obs_sp, act_sp: ManipulatedRewardConnector()),
+        )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)

rllib/examples/algorithms/bc/benchmark_rlunplugged_atari_pong_bc.py

@@ -0,0 +1,356 @@
+"""
+schema={
+      a_t: int64,
+      r_t: float,
+      episode_return: float,
+      o_tp1: list<item: binary>,
+      episode_id: int64,
+      a_tp1: int64,
+      o_t: list<item: binary>,
+      d_t: float
+   }
+"""
+import os
+from typing import Optional
+
+import cv2
+import gymnasium as gym
+import numpy as np
+import wandb
+
+from ray import tune
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.connectors.connector_v2 import ConnectorV2
+from ray.rllib.core import ALL_MODULES
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.utils import add_rllib_example_script_args
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_TRAINED_LIFETIME,
+)
+from ray.rllib.utils.test_utils import should_stop
+
+
+# Define a `ConnectorV2` to decode stacked encoded Atari frames.
+class DecodeObservations(ConnectorV2):
+    def __init__(
+        self,
+        input_observation_space: Optional[gym.Space] = None,
+        input_action_space: Optional[gym.Space] = None,
+        *,
+        multi_agent: bool = False,
+        as_learner_connector: bool = True,
+        **kwargs,
+    ):
+        """Decodes observation from PNG to numpy array.
+
+        Note, `rl_unplugged`'s stored observations are framestacked with
+        four frames per observation. This connector returns therefore
+        decoded observations of shape `(64, 64, 4)`.
+
+        Args:
+            multi_agent: Whether this is a connector operating on a multi-agent
+                observation space mapping AgentIDs to individual agents' observations.
+            as_learner_connector: Whether this connector is part of a Learner connector
+                pipeline, as opposed to an env-to-module pipeline.
+        """
+        super().__init__(
+            input_observation_space=input_observation_space,
+            input_action_space=input_action_space,
+            **kwargs,
+        )
+
+        self._multi_agent = multi_agent
+        self._as_learner_connector = as_learner_connector
+
+    @override(ConnectorV2)
+    def recompute_output_observation_space(
+        self, input_observation_space, input_action_space
+    ):
+        return gym.spaces.Box(
+            -1.0, 1.0, (64, 64, 4), float
+        )  # <- to keep it simple hardcoded to a fixed space
+
+    @override(ConnectorV2)
+    def __call__(
+        self,
+        *,
+        rl_module,
+        data,
+        episodes,
+        explore=None,
+        shared_data=None,
+        **kwargs,
+    ):
+
+        for sa_episode in self.single_agent_episode_iterator(
+            episodes, agents_that_stepped_only=False
+        ):
+            # Map encoded PNGs into arrays of shape (64, 64, 4).
+            def _map_fn(s):
+                # Preallocate the result array with shape (64, 64, 4)
+                result = np.empty((64, 64, 4), dtype=np.uint8)
+                for i in range(4):
+                    # Convert byte data to a numpy array of uint8
+                    nparr = np.frombuffer(s[i], np.uint8)
+                    # Decode the image as grayscale
+                    img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+                    # Resize the image to 64x64 using an efficient interpolation method
+                    resized = cv2.resize(img, (64, 64), interpolation=cv2.INTER_AREA)
+                    result[:, :, i] = resized
+                return (result.astype(np.float32) / 128.0) - 1.0
+
+            # Add the observations for t.
+            self.add_n_batch_items(
+                batch=data,
+                column=Columns.OBS,
+                # Ensure, we pass in a list, otherwise it is considered
+                # an already batched array.
+                items_to_add=[
+                    _map_fn(
+                        sa_episode.get_observations(slice(0, len(sa_episode)))[0],
+                    )
+                ],
+                num_items=len(sa_episode),
+                single_agent_episode=sa_episode,
+            )
+            # Add the observations for t+1.
+            self.add_n_batch_items(
+                batch=data,
+                column=Columns.NEXT_OBS,
+                items_to_add=[
+                    _map_fn(
+                        sa_episode.get_observations(slice(1, len(sa_episode) + 1))[0],
+                    )
+                ],
+                num_items=len(sa_episode),
+                single_agent_episode=sa_episode,
+            )
+
+        return data
+
+
+# Make the learner connector.
+def _make_learner_connector(observation_space, action_space):
+    return DecodeObservations()
+
+
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values toset up `config` below.
+parser = add_rllib_example_script_args(
+    default_reward=21.0,
+    default_timesteps=3000000000,
+    default_iters=100000000000,
+)
+args = parser.parse_args()
+
+# If multiple learners are requested define a scheduling
+# strategy with best data locality.
+if args.num_learners and args.num_learners > 1:
+    import ray
+
+    ray.init()
+    # Check, if we have a multi-node cluster.
+    nodes = ray.nodes()
+    ray.shutdown()
+    print(f"Number of nodes in cluster: {len(nodes)}")
+    # If we have a multi-node cluster spread learners.
+    if len(nodes) > 1:
+        os.environ["TRAIN_ENABLE_WORKER_SPREAD_ENV"] = "1"
+        print(
+            "Multi-node cluster and multi-learner setup. "
+            "Using a 'SPREAD' scheduling strategy for learners"
+            "to support data locality."
+        )
+    # Otherwise pack the learners on the single node.
+    else:
+        print(
+            "Single-node cluster and multi-learner setup. "
+            "Using a 'PACK' scheduling strategy for learners"
+            "to support data locality."
+        )
+
+# Wrap the environment used in evalaution into `RLlib`'s Atari Wrapper
+# that automatically stacks frames and converts to the dimension used
+# in the collection of the `rl_unplugged` data.
+def _env_creator(cfg):
+    return wrap_atari_for_new_api_stack(
+        gym.make("ale_py:ALE/Pong-v5", **cfg),
+        framestack=4,
+        dim=64,
+    )
+
+
+# Register the wrapped environment to `tune`. Note, environment registration
+# to Ray Tune must happen after checking the number of nodes, otherwise the
+# registration is removed.
+tune.register_env("WrappedALE/Pong-v5", _env_creator)
+
+# Anyscale RLUnplugged storage bucket. The bucket contains from the
+# original `RLUnplugged` bucket only the first `atari/Pong` run.
+# TODO (simon, artur): Create an extra bucket for the data and do not
+# use the `ANYSCALE_ARTIFACT_STORAGE`.
+anyscale_storage_bucket = os.environ["ANYSCALE_ARTIFACT_STORAGE"]
+anyscale_rlunplugged_atari_path = anyscale_storage_bucket + "/rllib/rl_unplugged/atari"
+
+# We only use the Atari game `Pong` here. Users can choose other Atari
+# games and set here the name.
+game = "Pong"
+
+# Path to the directory with all runs from Atari Pong.
+anyscale_rlunplugged_atari_pong_path = anyscale_rlunplugged_atari_path + f"/{game}"
+print(
+    "Streaming RLUnplugged Atari Pong data from path: "
+    f"{anyscale_rlunplugged_atari_pong_path}"
+)
+
+# Define the config for Behavior Cloning.
+config = (
+    BCConfig()
+    .environment(
+        env="WrappedALE/Pong-v5",
+        clip_rewards=True,
+        env_config={
+            # Make analogous to old v4 + NoFrameskip.
+            "frameskip": 4,
+            "full_action_space": False,
+            "repeat_action_probability": 0.0,
+        },
+    )
+    # Use the new API stack that makes directly use of `ray.data`.
+    .api_stack(
+        enable_rl_module_and_learner=True,
+        enable_env_runner_and_connector_v2=True,
+    )
+    # Evaluate in the actual environment online.
+    .evaluation(
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_parallel_to_training=True,
+        evaluation_config=BCConfig.overrides(exploration=False),
+    )
+    .learners(
+        num_learners=args.num_learners,
+        num_gpus_per_learner=args.num_gpus_per_learner,
+    )
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[anyscale_rlunplugged_atari_pong_path],
+        # `rl_unplugged`'s data schema is different from the one used
+        # internally in `RLlib`. Define the schema here so it can be used
+        # when transforming column data to episodes.
+        input_read_schema={
+            Columns.EPS_ID: "episode_id",
+            Columns.OBS: "o_t",
+            Columns.ACTIONS: "a_t",
+            Columns.REWARDS: "r_t",
+            Columns.NEXT_OBS: "o_tp1",
+            Columns.TERMINATEDS: "d_t",
+        },
+        # Do not materialize data, instead stream the data from Anyscale's
+        # S3 bucket (note, streaming data is an Anyscale-platform-only feature).
+        materialize_data=False,
+        materialize_mapped_data=False,
+        # Increase the parallelism in transforming batches, such that while
+        # training, new batches are transformed while others are used in updating.
+        map_batches_kwargs={
+            "concurrency": 40 * (max(args.num_learners, 1) or 1),
+            "num_cpus": 1,
+        },
+        # When iterating over batches in the dataset, prefetch at least 4
+        # batches per learner.
+        iter_batches_kwargs={
+            "prefetch_batches": 10,
+        },
+        # Iterate over 200 batches per RLlib iteration if multiple learners
+        # are used.
+        dataset_num_iters_per_learner=200,
+    )
+    .training(
+        # To increase learning speed with multiple learners,
+        # increase the learning rate correspondingly.
+        lr=0.0001
+        * max(
+            1,
+            (args.num_learners if args.num_learners and args.num_learners > 1 else 1)
+            ** 0.5,
+        ),
+        train_batch_size_per_learner=2048,
+        # Use the defined learner connector above, to decode observations.
+        learner_connector=_make_learner_connector,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            conv_filters=[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]],
+            conv_activation="relu",
+            head_fcnet_hiddens=[256],
+        ),
+    )
+    .debugging(
+        log_level="ERROR",
+    )
+)
+
+# Stop, if either the maximum point in Pong is reached (21.0) or 10 million steps
+# were trained.
+stop = {
+    f"{EVALUATION_RESULTS} / {ENV_RUNNER_RESULTS} / {EPISODE_RETURN_MEAN}": args.stop_reward,
+    f"{LEARNER_RESULTS} / {ALL_MODULES} / {NUM_ENV_STEPS_TRAINED_LIFETIME}": args.stop_timesteps,
+}
+
+# Build the algorithm.
+algo = config.build()
+
+# Shall we use wandb for logging results?
+if args.wandb_key:
+    # Login to wandb.
+    wandb.login(
+        key=args.wandb_key,
+        verify=True,
+        relogin=True,
+        force=True,
+    )
+
+    # Initialize wandb.
+    wandb.init(project=args.wandb_project)
+    # Clean results to log seemlessly to wandb.
+    from ray.air.integrations.wandb import _clean_log
+
+i = 0
+while True:
+    print("---------------------------------------------------------------")
+    print(f"Iteration {i + 1}")
+    results = algo.train()
+    print(results)
+
+    if args.wandb_key:
+        # Log results to wandb.
+        wandb.log(data=_clean_log(results), step=i)
+
+    if stop:
+        if should_stop(stop, results):
+            algo.cleanup()
+            break
+    i += 1
+
+print("------------------------------------------------")
+print()
+print("Training finished:\n")
+print(
+    f"Mean Episode Return in Evaluation: {results[EVALUATION_RESULTS][ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
+)
+print(
+    f"Number of Environment Steps trained: {results[LEARNER_RESULTS][ALL_MODULES][NUM_ENV_STEPS_TRAINED_LIFETIME]}"
+)
+print("================================================")

rllib/examples/algorithms/bc/cartpole_bc.py

@@ -0,0 +1,91 @@
+import warnings
+from pathlib import Path
+
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune.result import TRAINING_ITERATION
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "CartPole-v1" or args.env is None
+), "This tuned example works only with `CartPole-v1`."
+
+# Define the data paths.
+data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+base_path = Path(__file__).parents[3]
+data_path = "local://" / base_path / data_path
+print(f"data_path={data_path}")
+
+# Define the BC config.
+config = (
+    BCConfig()
+    .environment("CartPole-v1")
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[data_path.as_posix()],
+        # Concurrency defines the number of processes that run the
+        # `map_batches` transformations. This should be aligned with the
+        # 'prefetch_batches' argument in 'iter_batches_kwargs'.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 1},
+        # This data set is small so do not prefetch too many batches and use no
+        # local shuffle.
+        iter_batches_kwargs={"prefetch_batches": 1},
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration.
+        dataset_num_iters_per_learner=5,
+    )
+    .training(
+        train_batch_size_per_learner=1024,
+        # To increase learning speed with multiple learners,
+        # increase the learning rate correspondingly.
+        lr=0.0008 * (args.num_learners or 1) ** 0.5,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+        ),
+    )
+    .evaluation(
+        evaluation_interval=3,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_parallel_to_training=True,
+        evaluation_config=BCConfig.overrides(explore=False),
+    )
+)
+
+if not args.no_tune:
+    warnings.warn(
+        "You are running the example with Ray Tune. Offline RL uses "
+        "Ray Data, which doesn't interact seamlessly with Ray Tune. "
+        "If you encounter difficulties try to run the example without "
+        "Ray Tune using `--no-tune`."
+    )
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 350.0,
+    TRAINING_ITERATION: 350,
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/bc/cartpole_bc_with_offline_evaluation.py

@@ -0,0 +1,123 @@
+import warnings
+from pathlib import Path
+
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune.result import TRAINING_ITERATION
+
+parser = add_rllib_example_script_args()
+
+parser.add_argument(
+    "--offline-evaluation-interval",
+    type=int,
+    default=1,
+    help=(
+        "The interval in which offline evaluation should run in relation "
+        "to training iterations, e.g. if 1 offline evaluation runs in each "
+        "iteration, if 3 it runs each 3rd training iteration."
+    ),
+)
+parser.add_argument(
+    "--num-offline-eval-runners",
+    type=int,
+    default=2,
+    help=("The number of offline evaluation runners to be used in offline evaluation."),
+)
+parser.add_argument(
+    "--num-gpus-per-offline-eval-runner",
+    type=float,
+    default=0.0,
+    help=(
+        "The number of GPUs to be used in offline evaluation per offline "
+        "evaluation runner. Can be fractional."
+    ),
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "CartPole-v1" or args.env is None
+), "This tuned example works only with `CartPole-v1`."
+
+# Define the data paths.
+data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+base_path = Path(__file__).parents[3]
+print(f"base_path={base_path}")
+data_path = "local://" / base_path / data_path
+print(f"data_path={data_path}")
+
+# Define the BC config.
+config = (
+    BCConfig()
+    .environment(
+        "CartPole-v1",
+    )
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[data_path.as_posix()],
+        # Concurrency defines the number of processes that run the
+        # `map_batches` transformations. This should be aligned with the
+        # 'prefetch_batches' argument in 'iter_batches_kwargs'.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 1},
+        # This data set is small so do not prefetch too many batches and use no
+        # local shuffle.
+        iter_batches_kwargs={"prefetch_batches": 1},
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration.
+        dataset_num_iters_per_learner=5,
+    )
+    .training(
+        train_batch_size_per_learner=1024,
+        # To increase learning speed with multiple learners,
+        # increase the learning rate correspondingly.
+        lr=0.0008 * (args.num_learners or 1) ** 0.5,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+        ),
+    )
+    .evaluation(
+        evaluation_interval=1,
+        evaluation_parallel_to_training=False,
+        evaluation_config=BCConfig.overrides(explore=False),
+        offline_evaluation_interval=1,
+        offline_evaluation_type="eval_loss",
+        num_offline_eval_runners=args.num_offline_eval_runners,
+        num_gpus_per_offline_eval_runner=args.num_gpus_per_offline_eval_runner,
+        offline_eval_batch_size_per_runner=128,
+    )
+)
+
+if not args.no_tune:
+    warnings.warn(
+        "You are running the example with Ray Tune. Offline RL uses "
+        "Ray Data, which doesn't interact seamlessly with Ray Tune. "
+        "If you encounter difficulties try to run the example without "
+        "Ray Tune using `--no-tune`."
+    )
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 350.0,
+    TRAINING_ITERATION: 350,
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/bc/pendulum_bc.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune.result import TRAINING_ITERATION
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "Pendulum-v1" or args.env is None
+), "This tuned example works only with `Pendulum-v1`."
+
+# Define the data paths.
+data_path = "offline/tests/data/pendulum/pendulum-v1_large"
+base_path = Path(__file__).parents[3]
+print(f"base_path={base_path}")
+data_path = "local://" / base_path / data_path
+print(f"data_path={data_path}")
+
+# Define the BC config.
+config = (
+    BCConfig()
+    .environment(env="Pendulum-v1")
+    .api_stack(
+        enable_rl_module_and_learner=True,
+        enable_env_runner_and_connector_v2=True,
+    )
+    .evaluation(
+        evaluation_interval=3,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_parallel_to_training=True,
+        evaluation_config=BCConfig.overrides(explore=False),
+    )
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[data_path.as_posix()],
+        # Concurrency defines the number of processes that run the
+        # `map_batches` transformations. This should be aligned with the
+        # 'prefetch_batches' argument in 'iter_batches_kwargs'.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 2},
+        # This data set is small so do not prefetch too many batches and use no
+        # local shuffle.
+        iter_batches_kwargs={
+            "prefetch_batches": 1,
+        },
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration. For single-learner mode, 1 is the only option.
+        dataset_num_iters_per_learner=1 if not args.num_learners else None,
+    )
+    .training(
+        # To increase learning speed with multiple learners,
+        # increase the learning rate correspondingly.
+        lr=0.0008 * (args.num_learners or 1) ** 0.5,
+        train_batch_size_per_learner=1024,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+        ),
+    )
+)
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -200.0,
+    TRAINING_ITERATION: 350,
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/classes/appo_w_shared_data_actor.py

@@ -0,0 +1,88 @@
+from typing import List
+
+import ray
+from ray.rllib.algorithms import AlgorithmConfig
+from ray.rllib.algorithms.appo import APPO
+from ray.rllib.env.env_runner_group import EnvRunnerGroup
+
+
+@ray.remote
+class SharedDataActor:
+    """Simple example of an actor that's accessible from all other actors of an algo.
+
+    Exposes remote APIs `put` and `get` to other actors for storing and retrieving
+    arbitrary data.
+    """
+
+    def __init__(self):
+        self.storage = {}
+
+    def get(self, key, delete: bool = False):
+        value = self.storage.get(key)
+        if delete and key in self.storage:
+            del self.storage[key]
+        return value
+
+    def put(self, key, value):
+        self.storage[key] = value
+
+    def get_state(self):
+        return self.storage
+
+    def set_state(self, state):
+        self.storage = state
+
+
+class APPOWithSharedDataActor(APPO):
+    def setup(self, config: AlgorithmConfig):
+        # Call to parent `setup`.
+        super().setup(config)
+
+        # Create shared data actor.
+        self.shared_data_actor = SharedDataActor.remote()
+
+        # Share the actor with all other relevant actors.
+
+        def _share(actor, shared_act=self.shared_data_actor):
+            actor._shared_data_actor = shared_act
+            # Also add shared actor reference to all the learner connector pieces,
+            # if applicable.
+            if hasattr(actor, "_learner_connector") and actor._learner_connector:
+                for conn in actor._learner_connector:
+                    conn._shared_data_actor = shared_act
+
+        self.env_runner_group.foreach_env_runner(func=_share)
+        if self.eval_env_runner_group:
+            self.eval_env_runner_group.foreach_env_runner(func=_share)
+        self.learner_group.foreach_learner(func=_share)
+        if self._aggregator_actor_manager:
+            self._aggregator_actor_manager.foreach_actor(func=_share)
+
+    def get_state(self, *args, **kwargs):
+        state = super().get_state(*args, **kwargs)
+        # Add shared actor's state.
+        state["shared_data_actor"] = ray.get(self.shared_data_actor.get_state.remote())
+        return state
+
+    def set_state(self, state, *args, **kwargs):
+        super().set_state(state, *args, **kwargs)
+        # Set shared actor's state.
+        if "shared_data_actor" in state:
+            self.shared_data_actor.set_state.remote(state["shared_data_actor"])
+
+    def restore_env_runners(self, env_runner_group: EnvRunnerGroup) -> List[int]:
+        restored = super().restore_env_runners(env_runner_group)
+
+        # For the restored EnvRunners, send them the latest shared, global state
+        # from the `SharedDataActor`.
+        for restored_idx in restored:
+            state_ref = self.shared_data_actor.get.remote(
+                key=f"EnvRunner_{restored_idx}"
+            )
+            env_runner_group.foreach_env_runner(
+                lambda env_runner, state=state_ref: env_runner._global_state,
+                remote_worker_ids=[restored_idx],
+                timeout_seconds=0.0,
+            )
+
+        return restored

rllib/examples/algorithms/classes/maml_lr_differentiable_learner.py

@@ -0,0 +1,35 @@
+from typing import TYPE_CHECKING, Any, Dict
+
+from ray.rllib.core.learner.torch.torch_differentiable_learner import (
+    TorchDifferentiableLearner,
+)
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+if TYPE_CHECKING:
+    from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+
+torch, nn = try_import_torch()
+
+
+class MAMLTorchDifferentiableLearner(TorchDifferentiableLearner):
+    """A `TorchDifferentiableLearner` to perform MAML learning.
+
+    This `TorchDifferentiableLearner`
+    - defines a funcitonal MSE loss for learning simple (here non-linear)
+    prediction.
+    """
+
+    @override(TorchDifferentiableLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: "AlgorithmConfig",
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        """Defines a simple MSE prediction loss for continuous task."""
+
+        return nn.functional.mse_loss(fwd_out["y_pred"], batch["y"])

rllib/examples/algorithms/classes/maml_lr_differentiable_rlm.py

@@ -0,0 +1,39 @@
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class DifferentiableTorchRLModule(TorchRLModule):
+    """Differentiable neural network to learn sinusoid curves.
+
+    This `TorchRLModule`:
+    - defines a simple neural network to learn sinusoid curves with two
+    feed forward layern and ReLU activations,
+    - defines a differentiable `forward` call by overriding the `_forward`
+    method (which is implicitly used by the module's `forward` method); this
+    enables `torch.func.functional_call?` to work.
+    """
+
+    def setup(self):
+        """Sets up a simple neural network
+
+        The network contains two hidden layers and ReLU activations. Note,
+        input and output are single dimensional b/c the sinusoid curve is.
+        """
+        self.net = nn.Sequential(
+            nn.Linear(1, 40), nn.ReLU(), nn.Linear(40, 40), nn.ReLU(), nn.Linear(40, 1)
+        )
+
+    def _forward(self, batch, **kwargs):
+        """Defines method to be called for general forward path.
+
+        Note, it is important that the `RLModule.forward` method contains the
+        logic to be used for training forward pass b/c otherwise the functional
+        call via `torch.func.functional_call` will not work. See for reference
+        https://pytorch.org/docs/stable/generated/torch.func.functional_call.html.
+        """
+        outs = {}
+        outs["y_pred"] = self.net(batch[Columns.OBS])
+        return outs

rllib/examples/algorithms/classes/maml_lr_meta_learner.py

@@ -0,0 +1,38 @@
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from ray.rllib.core.learner.torch.torch_meta_learner import TorchMetaLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+if TYPE_CHECKING:
+    from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+
+torch, nn = try_import_torch()
+
+
+class MAMLTorchMetaLearner(TorchMetaLearner):
+    """A `TorchMetaLearner` to perform MAML learning.
+
+    This `TorchMetaLearner`
+    - defines a MSE loss for learning simple (here non-linear) prediction.
+    """
+
+    @override(TorchMetaLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: "AlgorithmConfig",
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType],
+        others_loss_per_module: List[Dict[ModuleID, TensorType]] = None,
+    ) -> TensorType:
+        """Defines a simple MSE prediction loss for continuous task.
+
+        Note, MAML does not need the losses from the registered differentiable
+        learners (contained in `others_loss_per_module`) b/c it computes a test
+        loss on an unseen data batch.
+        """
+        # Use a simple MSE loss for the meta learning task.
+        return torch.nn.functional.mse_loss(fwd_out["y_pred"], batch["y"])

rllib/examples/algorithms/classes/vpg.py

@@ -0,0 +1,175 @@
+import tree  # pip install dm_tree
+from typing_extensions import Self
+
+from ray.rllib.algorithms import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    ENV_RUNNER_SAMPLING_TIMER,
+    LEARNER_RESULTS,
+    LEARNER_UPDATE_TIMER,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TIMERS,
+)
+
+
+class VPGConfig(AlgorithmConfig):
+    """A simple VPG (vanilla policy gradient) algorithm w/o value function support.
+
+    Use for testing purposes only!
+
+    This Algorithm should use the VPGTorchLearner and VPGTorchRLModule
+    """
+
+    # A test setting to activate metrics on mean weights.
+    report_mean_weights: bool = True
+
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or VPG)
+
+        # VPG specific settings.
+        self.num_episodes_per_train_batch = 10
+        # Note that we don't have to set this here, because we tell the EnvRunners
+        # explicitly to sample entire episodes. However, for good measure, we change
+        # this setting here either way.
+        self.batch_mode = "complete_episodes"
+
+        # VPG specific defaults (from AlgorithmConfig).
+        self.num_env_runners = 1
+
+    @override(AlgorithmConfig)
+    def training(self, *, num_episodes_per_train_batch=NotProvided, **kwargs) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            num_episodes_per_train_batch: The number of complete episodes per train
+                batch. VPG requires entire episodes to be sampled from the EnvRunners.
+                For environments with varying episode lengths, this leads to varying
+                batch sizes (in timesteps) as well possibly causing slight learning
+                instabilities. However, for simplicity reasons, we stick to collecting
+                always exactly n episodes per training update.
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+
+        if num_episodes_per_train_batch is not NotProvided:
+            self.num_episodes_per_train_batch = num_episodes_per_train_batch
+
+        return self
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self):
+        if self.framework_str == "torch":
+            from ray.rllib.examples.rl_modules.classes.vpg_torch_rlm import (
+                VPGTorchRLModule,
+            )
+
+            spec = RLModuleSpec(
+                module_class=VPGTorchRLModule,
+                model_config={"hidden_dim": 64},
+            )
+        else:
+            raise ValueError(f"Unsupported framework: {self.framework_str}")
+
+        return spec
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self):
+        if self.framework_str == "torch":
+            from ray.rllib.examples.learners.classes.vpg_torch_learner import (
+                VPGTorchLearner,
+            )
+
+            return VPGTorchLearner
+        else:
+            raise ValueError(f"Unsupported framework: {self.framework_str}")
+
+
+class VPG(Algorithm):
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> VPGConfig:
+        return VPGConfig()
+
+    @override(Algorithm)
+    def training_step(self) -> None:
+        """Override of the training_step method of `Algorithm`.
+
+        Runs the following steps per call:
+        - Sample B timesteps (B=train batch size). Note that we don't sample complete
+        episodes due to simplicity. For an actual VPG algo, due to the loss computation,
+        you should always sample only completed episodes.
+        - Send the collected episodes to the VPG LearnerGroup for model updating.
+        - Sync the weights from LearnerGroup to all EnvRunners.
+        """
+        # Sample.
+        with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)):
+            episodes, env_runner_results = self._sample_episodes()
+        # Merge results from n parallel sample calls into self's metrics logger.
+        self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS)
+
+        # Just for demonstration purposes, log the number of time steps sampled in this
+        # `training_step` round.
+        # Mean over a window of 100:
+        self.metrics.log_value(
+            "episode_timesteps_sampled_mean_win100",
+            sum(map(len, episodes)),
+            reduce="mean",
+            window=100,
+        )
+        # Exponential Moving Average (EMA) with coeff=0.1:
+        self.metrics.log_value(
+            "episode_timesteps_sampled_ema",
+            sum(map(len, episodes)),
+            ema_coeff=0.1,  # <- weight of new value; weight of old avg=1.0-ema_coeff
+        )
+
+        # Update model.
+        with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
+            learner_results = self.learner_group.update(
+                episodes=episodes,
+                timesteps={
+                    NUM_ENV_STEPS_SAMPLED_LIFETIME: (
+                        self.metrics.peek(
+                            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME)
+                        )
+                    ),
+                },
+            )
+        # Merge results from m parallel update calls into self's metrics logger.
+        self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+        # Sync weights.
+        with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)):
+            self.env_runner_group.sync_weights(
+                from_worker_or_learner_group=self.learner_group,
+                inference_only=True,
+            )
+
+    def _sample_episodes(self):
+        # How many episodes to sample from each EnvRunner?
+        num_episodes_per_env_runner = self.config.num_episodes_per_train_batch // (
+            self.config.num_env_runners or 1
+        )
+        # Send parallel remote requests to sample and get the metrics.
+        sampled_data = self.env_runner_group.foreach_env_runner(
+            # Return tuple of [episodes], [metrics] from each EnvRunner.
+            lambda env_runner: (
+                env_runner.sample(num_episodes=num_episodes_per_env_runner),
+                env_runner.get_metrics(),
+            ),
+            # Loop over remote EnvRunners' `sample()` method in parallel or use the
+            # local EnvRunner if there aren't any remote ones.
+            local_env_runner=self.env_runner_group.num_remote_workers() <= 0,
+        )
+        # Return one list of episodes and a list of metrics dicts (one per EnvRunner).
+        episodes = tree.flatten([s[0] for s in sampled_data])
+        stats_dicts = [s[1] for s in sampled_data]
+
+        return episodes, stats_dicts

rllib/examples/algorithms/cql/pendulum_cql.py

@@ -0,0 +1,103 @@
+import warnings
+from pathlib import Path
+
+from ray.rllib.algorithms.cql.cql import CQLConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "Pendulum-v1" or args.env is None
+), "This tuned example works only with `Pendulum-v1`."
+
+# Define the base path relative to this file.
+base_path = Path(__file__).parents[3]
+# Use the larger data set of Pendulum we have. Note, these are
+# parquet data, the default in `AlgorithmConfig.offline_data`.
+data_path = base_path / "offline/tests/data/pendulum/pendulum-v1_enormous"
+data_path_uri = f"local://{data_path.as_posix()}"
+print(f"data_path_uri={data_path_uri}")
+# Define the configuration.
+config = (
+    CQLConfig()
+    .environment("Pendulum-v1")
+    .offline_data(
+        input_=[data_path_uri],
+        # The `kwargs` for the `map_batches` method in which our
+        # `OfflinePreLearner` is run. 2 data workers should be run
+        # concurrently.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 1},
+        # The `kwargs` for the `iter_batches` method. Due to the small
+        # dataset we choose only a single batch to prefetch.
+        iter_batches_kwargs={"prefetch_batches": 1},
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration.
+        dataset_num_iters_per_learner=5,
+        # TODO (sven): Has this any influence in the connectors?
+        actions_in_input_normalized=True,
+    )
+    .training(
+        bc_iters=200,
+        tau=9.5e-3,
+        min_q_weight=5.0,
+        train_batch_size_per_learner=1024,
+        twin_q=True,
+        actor_lr=1.7e-3 * (args.num_learners or 1) ** 0.5,
+        critic_lr=2.5e-3 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=1e-3 * (args.num_learners or 1) ** 0.5,
+        # Set this to `None` for all `SAC`-like algorithms. These
+        # algorithms use learning rates for each optimizer.
+        lr=None,
+    )
+    .reporting(
+        min_time_s_per_iteration=10,
+        metrics_num_episodes_for_smoothing=5,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        )
+    )
+    .evaluation(
+        evaluation_interval=3,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_config={
+            "explore": False,
+        },
+    )
+)
+
+if not args.no_tune:
+    warnings.warn(
+        "You are running the example with Ray Tune. Offline RL uses "
+        "Ray Data, which doesn't does not interact seamlessly with Ray Tune. "
+        "If you encounter difficulties try to run the example without "
+        "Ray Tune using `--no-tune`."
+    )
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -700.0,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: 800000,
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/dqn/benchmark_dqn_atari.py

@@ -0,0 +1,364 @@
+import gymnasium as gym
+from gymnasium.wrappers import AtariPreprocessing
+
+from ray import tune
+from ray.rllib.algorithms.dqn.dqn import DQNConfig
+from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule
+from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune import Stopper
+
+# Might need `gymnasium[atari, other]` to be installed.
+
+# See the following links for becnhmark results of other libraries:
+#   Original paper: https://arxiv.org/abs/1812.05905
+#   CleanRL: https://wandb.ai/cleanrl/cleanrl.benchmark/reports/Mujoco--VmlldzoxODE0NjE
+#   AgileRL: https://github.com/AgileRL/AgileRL?tab=readme-ov-file#benchmarks
+
+benchmark_envs = {
+    "AlienNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6022.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AmidarNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 202.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AssaultNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14491.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AsterixNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 280114.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AsteroidsNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2249.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AtlantisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 814684.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BankHeistNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 826.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BattleZoneNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 52040.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BeamRiderNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 21768.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BerzerkNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1793.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BowlingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 39.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BoxingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BreakoutNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 379.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "CentipedeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7160.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "ChopperCommandNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10916.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "CrazyClimberNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 143962.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DefenderNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 47671.3,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DemonAttackNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 109670.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DoubleDunkNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -0.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "EnduroNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2061.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FishingDerbyNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 22.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FreewayNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 29.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FrostbiteNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 4141.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "GopherNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 72595.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "GravitarNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 567.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "HeroNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 50496.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "IceHockeyNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KangarooNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10841.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KrullNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6715.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KungFuMasterNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 28999.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "MontezumaRevengeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 154.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "MsPacmanNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2570.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "NameThisGameNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11686.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PhoenixNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 103061.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PitfallNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -37.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PongNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PrivateEyeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1704.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "QbertNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 18397.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "RoadRunnerNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54261.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "RobotankNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 55.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SeaquestNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19176.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SkiingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SolarisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2860.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SpaceInvadersNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 12629.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "StarGunnerNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 123853.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SurroundNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TennisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -2.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TimePilotNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11190.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TutankhamNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 126.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "VentureNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 45.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "VideoPinballNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 506817.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "WizardOfWorNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14631.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "YarsRevengeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 93007.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "ZaxxonNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19658.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+}
+
+for env in benchmark_envs.keys():
+    tune.register_env(
+        env,
+        lambda ctx, e=env: AtariPreprocessing(
+            gym.make(e, **ctx), grayscale_newaxis=True, screen_size=84, noop_max=0
+        ),
+    )
+
+
+def _make_env_to_module_connector(env, spaces, device):
+    return FrameStackingEnvToModule(num_frames=4)
+
+
+def _make_learner_connector(input_observation_space, input_action_space):
+    return FrameStackingLearner(num_frames=4)
+
+
+# Define a `tune.Stopper` that stops the training if the benchmark is reached
+# or the maximum number of timesteps is exceeded.
+class BenchmarkStopper(Stopper):
+    def __init__(self, benchmark_envs):
+        self.benchmark_envs = benchmark_envs
+
+    def __call__(self, trial_id, result):
+        # Stop training if the mean reward is reached.
+        if (
+            result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
+            >= self.benchmark_envs[result["env"]][
+                f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+            ]
+        ):
+            return True
+        # Otherwise check, if the total number of timesteps is exceeded.
+        elif (
+            result[f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+            >= self.benchmark_envs[result["env"]][f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+        ):
+            return True
+        # Otherwise continue training.
+        else:
+            return False
+
+    # Note, this needs to implemented b/c the parent class is abstract.
+    def stop_all(self):
+        return False
+
+
+# See Table 1 in the Rainbow paper for the hyperparameters.
+config = (
+    DQNConfig()
+    .environment(
+        env=tune.grid_search(list(benchmark_envs.keys())),
+        env_config={
+            "max_episode_steps": 108000,
+            "obs_type": "grayscale",
+            # The authors actually use an action repetition of 4.
+            "repeat_action_probability": 0.25,
+        },
+        clip_rewards=True,
+    )
+    .env_runners(
+        # Every 4 agent steps a training update is performed.
+        rollout_fragment_length=4,
+        num_env_runners=1,
+        env_to_module_connector=_make_env_to_module_connector,
+    )
+    # TODO (simon): Adjust to new model_config_dict.
+    .training(
+        # Note, the paper uses also an Adam epsilon of 0.00015.
+        lr=0.0000625,
+        n_step=3,
+        tau=1.0,
+        train_batch_size=32,
+        target_network_update_freq=32000,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 1000000,
+            "alpha": 0.5,
+            # Note the paper used a linear schedule for beta.
+            "beta": 0.4,
+        },
+        # Note, these are frames.
+        num_steps_sampled_before_learning_starts=80000,
+        noisy=True,
+        num_atoms=51,
+        v_min=-10.0,
+        v_max=10.0,
+        double_q=True,
+        dueling=True,
+        model={
+            "cnn_filter_specifiers": [[32, 8, 4], [64, 4, 2], [64, 3, 1]],
+            "fcnet_activation": "tanh",
+            "post_fcnet_hiddens": [512],
+            "post_fcnet_activation": "relu",
+            "post_fcnet_weights_initializer": "orthogonal_",
+            "post_fcnet_weights_initializer_config": {"gain": 0.01},
+        },
+        learner_connector=_make_learner_connector,
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=10,
+        min_sample_timesteps_per_iteration=1000,
+    )
+    .evaluation(
+        evaluation_duration="auto",
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+        evaluation_parallel_to_training=True,
+        evaluation_config={
+            "explore": False,
+        },
+    )
+)
+
+tuner = tune.Tuner(
+    "DQN",
+    param_space=config,
+    run_config=tune.RunConfig(
+        stop=BenchmarkStopper(benchmark_envs=benchmark_envs),
+        name="benchmark_dqn_atari",
+    ),
+)
+
+tuner.fit()

rllib/examples/algorithms/dqn/benchmark_dqn_atari_rllib_preprocessing.py

@@ -0,0 +1,366 @@
+import gymnasium as gym
+
+from ray import tune
+from ray.rllib.algorithms.dqn.dqn import DQNConfig
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune import Stopper
+
+# Might need `gymnasium[atari, other]` to be installed.
+
+# See the following links for becnhmark results of other libraries:
+#   Original paper: https://arxiv.org/abs/1812.05905
+#   CleanRL: https://wandb.ai/cleanrl/cleanrl.benchmark/reports/Mujoco--VmlldzoxODE0NjE
+#   AgileRL: https://github.com/AgileRL/AgileRL?tab=readme-ov-file#benchmarks
+
+benchmark_envs = {
+    "AlienNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6022.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AmidarNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 202.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AssaultNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14491.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AsterixNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 280114.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AsteroidsNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2249.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "AtlantisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 814684.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BankHeistNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 826.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BattleZoneNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 52040.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BeamRiderNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 21768.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BerzerkNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1793.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BowlingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 39.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BoxingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "BreakoutNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 379.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "CentipedeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7160.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "ChopperCommandNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10916.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "CrazyClimberNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 143962.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DefenderNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 47671.3,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DemonAttackNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 109670.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "DoubleDunkNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -0.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "EnduroNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2061.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FishingDerbyNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 22.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FreewayNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 29.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "FrostbiteNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 4141.1,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "GopherNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 72595.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "GravitarNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 567.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "HeroNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 50496.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "IceHockeyNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KangarooNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10841.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KrullNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6715.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "KungFuMasterNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 28999.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "MontezumaRevengeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 154.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "MsPacmanNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2570.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "NameThisGameNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11686.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PhoenixNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 103061.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PitfallNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -37.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PongNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "PrivateEyeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1704.4,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "QbertNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 18397.6,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "RoadRunnerNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54261.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "RobotankNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 55.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SeaquestNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19176.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SkiingNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SolarisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2860.7,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SpaceInvadersNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 12629.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "StarGunnerNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 123853.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "SurroundNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TennisNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -2.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TimePilotNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11190.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "TutankhamNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 126.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "VentureNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 45.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "VideoPinballNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 506817.2,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "WizardOfWorNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14631.5,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "YarsRevengeNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 93007.9,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+    "ZaxxonNoFrameskip-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19658.0,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
+    },
+}
+
+for env in benchmark_envs:
+    tune.register_env(
+        env,
+        # Use the RLlib atari wrapper to squeeze images to 84x84.
+        # Note, the default of this wrapper is `framestack=4`.
+        lambda ctx, e=env: wrap_atari_for_new_api_stack(gym.make(e, **ctx), dim=84),
+    )
+
+
+# Define a `tune.Stopper` that stops the training if the benchmark is reached
+# or the maximum number of timesteps is exceeded.
+class BenchmarkStopper(Stopper):
+    def __init__(self, benchmark_envs):
+        self.benchmark_envs = benchmark_envs
+
+    def __call__(self, trial_id, result):
+        # Stop training if the mean reward is reached.
+        if (
+            result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
+            >= self.benchmark_envs[result["env"]][
+                f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+            ]
+        ):
+            return True
+        # Otherwise check, if the total number of timesteps is exceeded.
+        elif (
+            result[f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+            >= self.benchmark_envs[result["env"]][f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+        ):
+            return True
+        # Otherwise continue training.
+        else:
+            return False
+
+    # Note, this needs to implemented b/c the parent class is abstract.
+    def stop_all(self):
+        return False
+
+
+# See Table 1 in the Rainbow paper for the hyperparameters.
+config = (
+    DQNConfig()
+    .environment(
+        env=tune.grid_search(list(benchmark_envs.keys())),
+        env_config={
+            # "sticky actions" but not according to Danijar's 100k configs.
+            "repeat_action_probability": 0.0,
+            # "full action space" but not according to Danijar's 100k configs.
+            "full_action_space": False,
+            # Already done by MaxAndSkip wrapper: "action repeat" == 4.
+            "frameskip": 1,
+            # NOTE, because we use the atari wrapper of RLlib, we also have
+            # framestack: 4,
+            # dim: 84,
+            # NOTE, we do not use grayscale here, so this run will need
+            # more memory on GPU and CPU (buffer).
+        },
+        clip_rewards=True,
+    )
+    .env_runners(
+        # Every 4 agent steps a training update is performed.
+        rollout_fragment_length=4,
+        num_env_runners=1,
+    )
+    .learners(
+        # We have a train/sample ratio of 1:1 and a batch of 32.
+        num_learners=1,
+        num_gpus_per_learner=1,
+    )
+    # TODO (simon): Adjust to new model_config_dict.
+    .training(
+        # Note, the paper uses also an Adam epsilon of 0.00015.
+        lr=0.0000625,
+        n_step=1,
+        tau=1.0,
+        # TODO (simon): Activate when new model_config_dict is available.
+        # epsilon=0.01,
+        train_batch_size=32,
+        target_network_update_freq=8000,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 1000000,
+            "alpha": 0.5,
+            # Note the paper used a linear schedule for beta.
+            "beta": 0.4,
+        },
+        # Note, these are frames.
+        num_steps_sampled_before_learning_starts=20000,
+        noisy=True,
+        num_atoms=51,
+        v_min=-10.0,
+        v_max=10.0,
+        double_q=True,
+        dueling=True,
+        model={
+            "cnn_filter_specifiers": [[32, 8, 4], [64, 4, 2], [64, 3, 1]],
+            "fcnet_activation": "tanh",
+            "post_fcnet_hiddens": [512],
+            "post_fcnet_activation": "relu",
+            "post_fcnet_weights_initializer": "orthogonal_",
+            "post_fcnet_weights_initializer_config": {"gain": 0.01},
+        },
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=10,
+        min_sample_timesteps_per_iteration=1000,
+    )
+    .evaluation(
+        evaluation_duration="auto",
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+        evaluation_parallel_to_training=True,
+        evaluation_config={
+            "explore": False,
+        },
+    )
+)
+
+tuner = tune.Tuner(
+    "DQN",
+    param_space=config,
+    run_config=tune.RunConfig(
+        stop=BenchmarkStopper(benchmark_envs=benchmark_envs),
+        name="benchmark_dqn_atari_rllib_preprocessing",
+    ),
+)
+
+tuner.fit()

rllib/examples/algorithms/dqn/cartpole_dqn.py

@@ -0,0 +1,47 @@
+from ray.rllib.algorithms.dqn import DQNConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=450.0,
+    default_timesteps=200000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    DQNConfig()
+    .environment(env="CartPole-v1")
+    .training(
+        lr=0.0005 * (args.num_learners or 1) ** 0.5,
+        train_batch_size_per_learner=32,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 50000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        n_step=(2, 5),
+        double_q=True,
+        dueling=True,
+        epsilon=[(0, 1.0), (10000, 0.02)],
+    )
+    .rl_module(
+        # Settings identical to old stack.
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256],
+            fcnet_activation="tanh",
+            fcnet_bias_initializer="zeros_",
+            head_fcnet_bias_initializer="zeros_",
+            head_fcnet_hiddens=[256],
+        ),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/dqn/multi_agent_cartpole_dqn.py

@@ -0,0 +1,74 @@
+from ray.rllib.algorithms.dqn import DQNConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(
+    default_timesteps=500000,
+)
+parser.set_defaults(
+    num_agents=2,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("multi_agent_cartpole", lambda cfg: MultiAgentCartPole(config=cfg))
+
+config = (
+    DQNConfig()
+    .environment(env="multi_agent_cartpole", env_config={"num_agents": args.num_agents})
+    .training(
+        lr=0.00065 * (args.num_learners or 1) ** 0.5,
+        train_batch_size_per_learner=48,
+        replay_buffer_config={
+            "type": "MultiAgentPrioritizedEpisodeReplayBuffer",
+            "capacity": 50000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        n_step=(2, 5),
+        double_q=True,
+        num_atoms=1,
+        dueling=True,
+        epsilon=[(0, 1.0), (20000, 0.02)],
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="tanh",
+            fcnet_bias_initializer="zeros_",
+            head_fcnet_bias_initializer="zeros_",
+            head_fcnet_hiddens=[256],
+        ),
+    )
+)
+
+if args.num_agents:
+    config.multi_agent(
+        policy_mapping_fn=lambda aid, *arg, **kw: f"p{aid}",
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    # `episode_return_mean` is the sum of all agents/policies' returns.
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 150.0 * args.num_agents,
+}
+
+
+if __name__ == "__main__":
+    assert (
+        args.num_agents > 0
+    ), "The `--num-agents` arg must be > 0 for this script to work."
+
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/dqn/stateless_cartpole_dqn.py

@@ -0,0 +1,58 @@
+from ray.rllib.algorithms.dqn import DQNConfig
+from ray.rllib.connectors.env_to_module import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=2000000,
+    default_reward=350.0,
+)
+parser.set_defaults(
+    num_env_runners=3,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    DQNConfig()
+    .environment(StatelessCartPole)
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+    )
+    .training(
+        lr=0.0005,
+        train_batch_size_per_learner=32,
+        replay_buffer_config={
+            "type": "EpisodeReplayBuffer",
+            "capacity": 100000,
+        },
+        n_step=1,
+        double_q=True,
+        dueling=True,
+        num_atoms=1,
+        epsilon=[(0, 1.0), (20000, 0.02)],
+        burn_in_len=8,
+    )
+    .rl_module(
+        # Settings identical to old stack.
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256],
+            fcnet_activation="tanh",
+            fcnet_bias_initializer="zeros_",
+            head_fcnet_bias_initializer="zeros_",
+            head_fcnet_hiddens=[256],
+            head_fcnet_activation="tanh",
+            lstm_kernel_initializer="xavier_uniform_",
+            use_lstm=True,
+            max_seq_len=20,
+        ),
+    )
+)
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/dreamerv3/atari_100k_dreamerv3.py

@@ -0,0 +1,100 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py --env ale_py:ALE/[gym ID e.g. Pong-v5]
+
+# To see all available options:
+# python [this script name].py --help
+
+import gymnasium as gym
+
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=1000000,
+    default_reward=20.0,
+    default_timesteps=100000,
+)
+parser.set_defaults(env="ale_py:ALE/Pong-v5")
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+# If we use >1 GPU and increase the batch size accordingly, we should also
+# increase the number of envs per worker.
+if args.num_envs_per_env_runner is None:
+    args.num_envs_per_env_runner = args.num_learners or 1
+
+
+# Create the DreamerV3-typical Atari setup.
+def _env_creator(cfg):
+    return wrap_atari_for_new_api_stack(
+        gym.make(args.env, **cfg, render_mode="rgb_array"),
+        # No framestacking necessary for Dreamer.
+        framestack=None,
+        # No grayscaling necessary for Dreamer.
+        grayscale=False,
+    )
+
+
+tune.register_env("env", _env_creator)
+
+default_config = DreamerV3Config()
+lr_multiplier = args.num_learners or 1
+
+config = (
+    DreamerV3Config()
+    .environment(
+        env="env",
+        # [2]: "We follow the evaluation protocol of Machado et al. (2018) with 200M
+        # environment steps, action repeat of 4, a time limit of 108,000 steps per
+        # episode that correspond to 30 minutes of game play, no access to life
+        # information, full action space, and sticky actions. Because the world model
+        # integrates information over time, DreamerV2 does not use frame stacking.
+        # The experiments use a single-task setup where a separate agent is trained
+        # for each game. Moreover, each agent uses only a single environment instance.
+        env_config={
+            # "sticky actions" but not according to Danijar's 100k configs.
+            "repeat_action_probability": 0.0,
+            # "full action space" but not according to Danijar's 100k configs.
+            "full_action_space": False,
+            # Already done by MaxAndSkip wrapper: "action repeat" == 4.
+            "frameskip": 1,
+        },
+    )
+    .env_runners(
+        remote_worker_envs=(args.num_learners and args.num_learners > 1),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(args.num_learners or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="S",
+        training_ratio=1024,
+        batch_size_B=16 * (args.num_learners or 1),
+        world_model_lr=default_config.world_model_lr * lr_multiplier,
+        actor_lr=default_config.actor_lr * lr_multiplier,
+        critic_lr=default_config.critic_lr * lr_multiplier,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/dreamerv3/atari_200M_dreamerv3.py

@@ -0,0 +1,88 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py --env ale_py:ALE/[gym ID e.g. Pong-v5]
+
+# To see all available options:
+# python [this script name].py --help
+
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=1000000,
+    default_reward=20.0,
+    default_timesteps=1000000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+# If we use >1 GPU and increase the batch size accordingly, we should also
+# increase the number of envs per worker.
+if args.num_envs_per_env_runner is None:
+    args.num_envs_per_env_runner = 8 * (args.num_learners or 1)
+
+default_config = DreamerV3Config()
+lr_multiplier = (args.num_learners or 1) ** 0.5
+
+config = (
+    DreamerV3Config()
+    .resources(
+        # For each (parallelized) env, we should provide a CPU. Lower this number
+        # if you don't have enough CPUs.
+        num_cpus_for_main_process=8
+        * (args.num_learners or 1),
+    )
+    .environment(
+        env=args.env,
+        # [2]: "We follow the evaluation protocol of Machado et al. (2018) with 200M
+        # environment steps, action repeat of 4, a time limit of 108,000 steps per
+        # episode that correspond to 30 minutes of game play, no access to life
+        # information, full action space, and sticky actions. Because the world model
+        # integrates information over time, DreamerV2 does not use frame stacking.
+        # The experiments use a single-task setup where a separate agent is trained
+        # for each game. Moreover, each agent uses only a single environment instance.
+        env_config={
+            # "sticky actions" but not according to Danijar's 100k configs.
+            "repeat_action_probability": 0.0,
+            # "full action space" but not according to Danijar's 100k configs.
+            "full_action_space": False,
+            # Already done by MaxAndSkip wrapper: "action repeat" == 4.
+            "frameskip": 1,
+        },
+    )
+    .env_runners(
+        remote_worker_envs=True,
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(args.num_learners or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="XL",
+        training_ratio=64,
+        batch_size_B=16 * (args.num_learners or 1),
+        world_model_lr=default_config.world_model_lr * lr_multiplier,
+        actor_lr=default_config.actor_lr * lr_multiplier,
+        critic_lr=default_config.critic_lr * lr_multiplier,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/dreamerv3/cartpole_dreamerv3.py

@@ -0,0 +1,22 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Run with:
+# python [this script name].py
+
+config = (
+    DreamerV3Config()
+    .environment("CartPole-v1")
+    .training(
+        model_size="XS",
+        training_ratio=1024,
+    )
+)

rllib/examples/algorithms/dreamerv3/dm_control_suite_vision_dreamerv3.py

@@ -0,0 +1,88 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py --env DMC/[task]/[domain] (e.g. DMC/cartpole/swingup)
+
+# To see all available options:
+# python [this script name].py --help
+
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.env.wrappers.dm_control_wrapper import ActionClip, DMCEnv
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=1000000,
+    default_reward=800.0,
+    default_timesteps=1000000,
+)
+parser.set_defaults(env="DMC/cartpole/swingup")
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+# If we use >1 GPU and increase the batch size accordingly, we should also
+# increase the number of envs per worker.
+if args.num_envs_per_env_runner is None:
+    args.num_envs_per_env_runner = 4 * (args.num_learners or 1)
+
+parts = args.env.split("/")
+assert len(parts) == 3, (
+    "ERROR: DMC env must be formatted as 'DMC/[task]/[domain]', e.g. "
+    f"'DMC/cartpole/swingup'! You provided '{args.env}'."
+)
+
+
+def env_creator(cfg):
+    return ActionClip(
+        DMCEnv(
+            parts[1],
+            parts[2],
+            from_pixels=True,
+            channels_first=False,
+        )
+    )
+
+
+tune.register_env("env", env_creator)
+
+default_config = DreamerV3Config()
+lr_multiplier = (args.num_learners or 1) ** 0.5
+
+config = (
+    DreamerV3Config()
+    # Use image observations.
+    .environment(env="env")
+    .env_runners(
+        remote_worker_envs=True,
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(args.num_learners or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="S",
+        training_ratio=512,
+        batch_size_B=16 * (args.num_learners or 1),
+        world_model_lr=default_config.world_model_lr * lr_multiplier,
+        actor_lr=default_config.actor_lr * lr_multiplier,
+        critic_lr=default_config.critic_lr * lr_multiplier,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/dreamerv3/flappy_bird_dreamerv3.py

@@ -0,0 +1,80 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Number of GPUs to run on.
+num_gpus = 0
+
+# DreamerV3 config and default (1 GPU) learning rates.
+config = DreamerV3Config()
+w = config.world_model_lr
+c = config.critic_lr
+
+
+def _env_creator(ctx):
+    import flappy_bird_gymnasium  # noqa doctest: +SKIP
+    import gymnasium as gym
+    from supersuit.generic_wrappers import resize_v1
+    from ray.rllib.env.wrappers.atari_wrappers import NormalizedImageEnv
+
+    return NormalizedImageEnv(
+        resize_v1(  # resize to 64x64 and normalize images
+            gym.make("FlappyBird-rgb-v0", audio_on=False), x_size=64, y_size=64
+        )
+    )
+
+
+# Register the FlappyBird-rgb-v0 env including necessary wrappers via the
+# `tune.register_env()` API.
+tune.register_env("flappy-bird", _env_creator)
+
+# Further specify the DreamerV3 config object to use.
+(
+    config.environment("flappy-bird")
+    .resources(
+        num_cpus_for_main_process=1,
+    )
+    .learners(
+        num_learners=0 if num_gpus == 1 else num_gpus,
+        num_gpus_per_learner=1 if num_gpus else 0,
+    )
+    .env_runners(
+        # If we use >1 GPU and increase the batch size accordingly, we should also
+        # increase the number of envs per worker.
+        num_envs_per_env_runner=8 * (num_gpus or 1),
+        remote_worker_envs=True,
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(num_gpus or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="M",
+        training_ratio=64,
+        batch_size_B=16 * (num_gpus or 1),
+        # Use a well established 4-GPU lr scheduling recipe:
+        # ~ 1000 training updates with 0.4x[default rates], then over a few hundred
+        # steps, increase to 4x[default rates].
+        world_model_lr=[[0, 0.4 * w], [8000, 0.4 * w], [10000, 3 * w]],
+        critic_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]],
+        actor_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]],
+    )
+)

rllib/examples/algorithms/dreamerv3/frozenlake_2x2_dreamerv3.py

@@ -0,0 +1,34 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+config = (
+    DreamerV3Config()
+    .environment(
+        "FrozenLake-v1",
+        env_config={
+            "desc": [
+                "SF",
+                "HG",
+            ],
+            "is_slippery": False,
+        },
+    )
+    .training(
+        model_size="XS",
+        training_ratio=1024,
+    )
+)

rllib/examples/algorithms/dreamerv3/frozenlake_4x4_deterministic_dreamerv3.py

@@ -0,0 +1,31 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+config = (
+    DreamerV3Config()
+    .environment(
+        "FrozenLake-v1",
+        env_config={
+            "map_name": "4x4",
+            "is_slippery": False,
+        },
+    )
+    .training(
+        model_size="nano",
+        training_ratio=1024,
+    )
+)

rllib/examples/algorithms/dreamerv3/gymnasium_robotics_dreamerv3.py

@@ -0,0 +1,68 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+try:
+    import gymnasium_robotics  # noqa
+except (ImportError, ModuleNotFoundError):
+    print("You have to `pip install gymnasium_robotics` in order to run this example!")
+
+import gymnasium as gym
+
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Number of GPUs to run on.
+num_gpus = 4
+
+# Register the gymnasium robotics env (including necessary wrappers and options) via the
+# `tune.register_env()` API.
+# Create the specific gymnasium robotics env.
+# e.g. AdroitHandHammerSparse-v1 or FrankaKitchen-v1.
+# return gym.make("FrankaKitchen-v1", tasks_to_complete=["microwave", "kettle"])
+tune.register_env("flappy-bird", lambda ctx: gym.make("AdroitHandHammer-v1"))
+
+# Define the DreamerV3 config object to use.
+config = DreamerV3Config()
+w = config.world_model_lr
+c = config.critic_lr
+# Further specify the details of our config object.
+(
+    config.resources(
+        num_cpus_for_main_process=8 * (num_gpus or 1),
+    )
+    .learners(
+        num_learners=0 if num_gpus == 1 else num_gpus,
+        num_gpus_per_learner=1 if num_gpus else 0,
+    )
+    # If we use >1 GPU and increase the batch size accordingly, we should also
+    # increase the number of envs per worker.
+    .env_runners(num_envs_per_env_runner=8 * (num_gpus or 1), remote_worker_envs=True)
+    .reporting(
+        metrics_num_episodes_for_smoothing=(num_gpus or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="XL",
+        training_ratio=64,
+        batch_size_B=16 * (num_gpus or 1),
+        world_model_lr=[[0, 0.4 * w], [50000, 0.4 * w], [100000, 3 * w]],
+        critic_lr=[[0, 0.4 * c], [50000, 0.4 * c], [100000, 3 * c]],
+        actor_lr=[[0, 0.4 * c], [50000, 0.4 * c], [100000, 3 * c]],
+    )
+)

rllib/examples/algorithms/dreamerv3/highway_env_dreamerv3.py

@@ -0,0 +1,73 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+try:
+    import highway_env  # noqa
+except (ImportError, ModuleNotFoundError):
+    print("You have to `pip install highway_env` in order to run this example!")
+
+import gymnasium as gym
+
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+# Number of GPUs to run on.
+num_gpus = 4
+
+# Register the highway env (including necessary wrappers and options) via the
+# `tune.register_env()` API.
+# Create the specific env.
+# e.g. roundabout-v0 or racetrack-v0
+tune.register_env("flappy-bird", lambda ctx: gym.make("intersection-v0", policy_freq=5))
+
+# Define the DreamerV3 config object to use.
+config = DreamerV3Config()
+w = config.world_model_lr
+c = config.critic_lr
+
+(
+    config.resources(
+        num_cpus_for_main_process=1,
+    )
+    .learners(
+        num_learners=0 if num_gpus == 1 else num_gpus,
+        num_gpus_per_learner=1 if num_gpus else 0,
+    )
+    .env_runners(
+        # If we use >1 GPU and increase the batch size accordingly, we should also
+        # increase the number of envs per worker.
+        num_envs_per_env_runner=8 * (num_gpus or 1),
+        remote_worker_envs=True,
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(num_gpus or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="M",
+        training_ratio=64,
+        batch_size_B=16 * (num_gpus or 1),
+        # Use a well established 4-GPU lr scheduling recipe:
+        # ~ 1000 training updates with 0.4x[default rates], then over a few hundred
+        # steps, increase to 4x[default rates].
+        world_model_lr=[[0, 0.4 * w], [8000, 0.4 * w], [10000, 3 * w]],
+        critic_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]],
+        actor_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]],
+    )
+)

rllib/examples/algorithms/dreamerv3/pendulum_dreamerv3.py

@@ -0,0 +1,64 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=10000,
+    default_reward=-200.0,
+    default_timesteps=100000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+# If we use >1 GPU and increase the batch size accordingly, we should also
+# increase the number of envs per worker.
+if args.num_envs_per_env_runner is None:
+    args.num_envs_per_env_runner = args.num_learners or 1
+
+# Run with:
+# python [this script name].py
+
+# To see all available options:
+# python [this script name].py --help
+
+default_config = DreamerV3Config()
+lr_multiplier = args.num_learners or 1
+
+
+config = (
+    DreamerV3Config()
+    .environment("Pendulum-v1")
+    .env_runners(
+        remote_worker_envs=(args.num_learners and args.num_learners > 1),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=(args.num_learners or 1),
+        report_images_and_videos=False,
+        report_dream_data=False,
+        report_individual_batch_item_stats=False,
+    )
+    # See Appendix A.
+    .training(
+        model_size="S",
+        training_ratio=1024,
+        batch_size_B=16 * (args.num_learners or 1),
+        world_model_lr=default_config.world_model_lr * lr_multiplier,
+        actor_lr=default_config.actor_lr * lr_multiplier,
+        critic_lr=default_config.critic_lr * lr_multiplier,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/impala/cartpole_impala.py

@@ -0,0 +1,90 @@
+"""Example showing how to run IMPALA on the CartPole environment.
+
+IMPALA (Importance Weighted Actor-Learner Architecture) is a distributed
+reinforcement learning algorithm that decouples acting from learning. It uses
+V-trace for off-policy correction, enabling efficient training across many
+distributed actors while maintaining stable learning.
+
+This example:
+    - trains on the classic CartPole-v1 control task
+    - uses gradient clipping by global norm (40.0) for training stability
+    - scales the learning rate with the square root of the number of learners
+    - shares value function layers with the policy network for parameter efficiency
+    - targets a reward of 450 (near-optimal for CartPole-v1's max of 500)
+
+How to run this script
+----------------------
+`python cartpole_impala.py [options]`
+
+To run with default settings:
+`python cartpole_impala.py`
+
+To scale up with distributed learning using multiple learners and env-runners:
+`python cartpole_impala.py --num-learners=2 --num-env-runners=8`
+
+To use a GPU-based learner add the number of GPUs per learner:
+`python cartpole_impala.py --num-learners=1 --num-gpus-per-learner=1`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0 --num-learners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+By setting `--num-learners=0` and `--num-env-runners=0` will make them run locally
+instead of as remote Ray Actors where breakpoints aren't possible.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key]
+ --wandb-project=[some project name]
+ --wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+Results to expect
+-----------------
+The algorithm should reach the reward threshold of 450 on CartPole-v1
+within 2 million timesteps (see: `default_timesteps` in the code).
+CartPole-v1 has a maximum episode reward of 500, and IMPALA should
+consistently achieve near-optimal performance on this task.
+"""
+from ray.rllib.algorithms.impala import IMPALAConfig
+from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=450.0,
+    default_timesteps=2_000_000,
+)
+parser.set_defaults(
+    num_env_runners=4,
+    num_envs_per_env_runner=16,
+    num_learners=1,
+)
+args = parser.parse_args()
+
+
+config = (
+    IMPALAConfig()
+    .environment("CartPole-v1")
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+    )
+    .training(
+        train_batch_size_per_learner=500,
+        grad_clip=40.0,
+        grad_clip_by="global_norm",
+        lr=0.0005 * ((args.num_learners or 1) ** 0.5),
+        vf_loss_coeff=0.05,
+        entropy_coeff=0.0,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            vf_share_layers=True,
+        ),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/impala/tictactoe_impala.py

@@ -0,0 +1,134 @@
+"""Example showing how to run multi-agent IMPALA on TicTacToe with self-play.
+
+This example demonstrates multi-agent reinforcement learning using IMPALA on a
+TicTacToe environment. The setup includes trainable policies that learn to play
+against each other and a frozen random policy that provides diverse opponents.
+This self-play with random opponents approach helps prevent overfitting to a
+single opponent strategy.
+
+This example:
+    - trains multiple policies on the TicTacToe multi-agent environment
+    - uses a RandomRLModule as a frozen opponent that is not trained
+    - randomly maps agents to policies (including the random policy) each episode
+    - demonstrates MultiRLModuleSpec for configuring multiple policies
+    - uses 4 env runners by default for parallel experience collection
+
+How to run this script
+----------------------
+`python tictactoe_impala.py [options]`
+
+To run with default settings (5 trainable agents):
+`python tictactoe_impala.py`
+
+To run with a different number of trainable agents:
+`python tictactoe_impala.py --num-agents=4`
+
+To scale up with distributed learning using multiple learners and env-runners:
+`python tictactoe_impala.py --num-learners=2 --num-env-runners=8`
+
+To use a GPU-based learner add the number of GPUs per learner:
+`python tictactoe_impala.py --num-learners=1 --num-gpus-per-learner=1`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0 --num-learners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+By setting `--num-learners=0` and `--num-env-runners=0` will make them run locally
+instead of as remote Ray Actors where breakpoints aren't possible.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key]
+ --wandb-project=[some project name]
+ --wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+Results to expect
+-----------------
+Four policies are trained plus a fifth random policy are randomly paired against
+each other. Training is stopped when policy 0 achieves a return of < -0.3 within
+2 million timesteps. A reward close to 0 or positive indicates
+the policies are learning to win or draw more often than they lose.
+"""
+import random
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.rllib.algorithms.impala import IMPALAConfig
+from ray.rllib.core.rl_module import MultiRLModuleSpec, RLModuleSpec
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent.tic_tac_toe import TicTacToe
+from ray.rllib.examples.rl_modules.classes.random_rlm import RandomRLModule
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_MODULE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=-0.5,
+    default_timesteps=2_000_000,
+)
+parser.set_defaults(
+    num_env_runners=4,
+    num_envs_per_env_runner=3,
+    num_learners=1,
+    num_agents=5,
+)
+args = parser.parse_args()
+
+config = (
+    IMPALAConfig()
+    .environment(TicTacToe)
+    .env_runners(
+        num_env_runners=args.num_env_runners,
+        num_envs_per_env_runner=args.num_envs_per_env_runner,
+    )
+    .learners(
+        num_learners=args.num_learners,
+    )
+    .training(
+        train_batch_size_per_learner=1000,
+        grad_clip=30.0,
+        grad_clip_by="global_norm",
+        lr=0.0005,
+        vf_loss_coeff=0.01,
+        entropy_coeff=0.0,
+    )
+    .rl_module(
+        rl_module_spec=MultiRLModuleSpec(
+            rl_module_specs=(
+                {
+                    f"p{i}": RLModuleSpec(
+                        model_config=DefaultModelConfig(vf_share_layers=True),
+                    )
+                    for i in range(args.num_agents)
+                }
+                | {"random": RLModuleSpec(module_class=RandomRLModule)}
+            ),
+        ),
+    )
+    .multi_agent(
+        policies={f"p{i}" for i in range(args.num_agents)} | {"random"},
+        policy_mapping_fn=lambda aid, eps, **kw: (
+            random.choice([f"p{i}" for i in range(args.num_agents)] + ["random"])
+        ),
+        policies_to_train=[f"p{i}" for i in range(args.num_agents)],
+    )
+)
+
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_MODULE_RETURN_MEAN}/p0": args.stop_reward,
+    f"{ENV_RUNNER_RESULTS}/{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps,
+    TRAINING_ITERATION: args.stop_iters,
+}
+success_metric = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_MODULE_RETURN_MEAN}/p0": args.stop_reward
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(
+        config, args, stop=stop, success_metric=success_metric
+    )

rllib/examples/algorithms/iql/pendulum_iql.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+
+from ray.rllib.algorithms.iql.iql import IQLConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune.result import TRAINING_ITERATION
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "Pendulum-v1" or args.env is None
+), "This tuned example works only with `Pendulum-v1`."
+
+# Define the data paths.
+data_path = "offline/tests/data/pendulum/pendulum-v1_enormous"
+base_path = Path(__file__).parents[3]
+print(f"base_path={base_path}")
+data_path = "local://" / base_path / data_path
+print(f"data_path={data_path}")
+
+# Define the IQL config.
+config = (
+    IQLConfig()
+    .environment(env="Pendulum-v1")
+    .evaluation(
+        evaluation_interval=3,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_parallel_to_training=True,
+    )
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[data_path.as_posix()],
+        # Concurrency defines the number of processes that run the
+        # `map_batches` transformations. This should be aligned with the
+        # 'prefetch_batches' argument in 'iter_batches_kwargs'.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 2},
+        # This data set is small so do not prefetch too many batches and use no
+        # local shuffle.
+        iter_batches_kwargs={
+            "prefetch_batches": 1,
+        },
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration.
+        dataset_num_iters_per_learner=5,
+    )
+    .training(
+        # To increase learning speed with multiple learners,
+        # increase the learning rates correspondingly.
+        actor_lr=2.59e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=2.14e-4 * (args.num_learners or 1) ** 0.5,
+        value_lr=3.7e-5 * (args.num_learners or 1) ** 0.5,
+        # Smooth Polyak-averaging for the target network.
+        tau=6e-4,
+        # Update the target network each training iteration.
+        target_network_update_freq=1,
+        train_batch_size_per_learner=1024,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        )
+    )
+)
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -200.0,
+    TRAINING_ITERATION: 1250,
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/maml_lr_supervised_learning.py

@@ -0,0 +1,443 @@
+"""Example of how to use `TorchMetaLearner` and `DifferentiableLearner` for MAML.
+
+Meta-learning, or “learning to learn,” trains models to quickly adapt to new tasks
+using only a few examples. One prominent method is Model-Agnostic Meta-Learning
+(MAML), which is compatible with any model trained via gradient descent. MAML has
+been successfully applied across domains such as classification, regression, and
+reinforcement learning.
+
+In this MAML example, the goal is to train a model that can adapt to an infinite
+number of tasks, where each task corresponds to a sinusoidal function with randomly
+sampled amplitude and phase. Because each new task introduces a shift in data
+distribution, traditional learning algorithms would fail to generalize — they’d
+overfit to the training task and struggle on unseen ones. Meta-learning addresses
+this by optimizing the model parameters such that they can be fine-tuned rapidly
+for any new task.
+
+During training, a DifferentiableLearner performs an inner-loop update using the
+training error for each task. The outer-loop TorchMetaLearner then evaluates the
+model’s performance on held-out data (the task's test set) and updates the meta-
+parameters so that they lead to better generalization across all tasks. This bi-
+level optimization ensures that gradients across tasks remain close, enabling
+fast adaptation.
+
+At inference time, the trained model can adapt to a new task using just a small
+batch of examples — performing few-shot learning to adjust quickly and accurately.
+
+This example shows:
+- how to implement MAML with RLlib in just a few lines of code.
+- how to define a `TorchDifferentiableLearner` to register a custom train loss
+    function.
+- how to define a `TorchMetaLearner` class to implement a custom meta (test) train
+    loss function.
+- how to configure both learners top be used with each others via the
+    `DifferentiableAlgorithmConfig` and `DifferentiableLearnerConfig`.
+- how to update the `RLModule` in a meta-learning fashion.
+- how to fine-tune an `RLModule` with gradient descent within a few iterations with
+    only using the meta (test) loss.
+
+See :py:class:`~ray.rllib.examples.learners.classes.lr_meta_learner.LRTorchMetaLearner` # noqa
+class for details on how to override the main `TorchMetaLearner`. And see
+:py:class:`~ray.rllib.examples.learners.classes.lr_differentiable_learner.LRTorchDifferentiableLearner` # noqa
+class for an example of how to override the main `TorchDifferentiableLearner`.
+
+Note, the meta-learner needs a long-enough training (`default_iters`=~70,000) to learn
+to adapt quickly to new tasks.
+
+How to run this script
+----------------------
+`python [script file name].py --iters=70000 --meta-train-batch-size=5 --fine-tune-batch-size=5`
+
+Use the `--meta-train-batch-size` to set the training/testing batch size in meta-learning and
+the `--fine-tune-batch-size` to adjust the number of samples used in all updates during
+few-shot learning.
+
+To suppress plotting (plotting is the default) use `--no-plot` and for taking a longer
+look at the plot increase the seconds for which plotting is paused at the end of the
+script by `--pause-plot-secs`.
+
+Results to expect
+-----------------
+You should expect to see sometimes alternating test losses ("Total Loss") due to new
+(unseen) tasks during meta learning. In few-shot learning after the meta-learning the
+(few shot) loss should decrease almost monotonically. In the plot you can expect to see
+a decent adaption to the new task after fine-tuning updates of the `RLModule` weights.
+
+With `--iters=70_000`, `--meta-train-batch-size=5`, `--fine-tune-batch-size=5`,
+`--fine-tune-lr=0.01`, `--fine-tune-iters=10`, `--meta-lr=0.001`, `--noise-std=0.0`,
+and no seed defined.
+-------------------------
+
+Iteration: 68000
+Total loss: 0.013758559711277485
+-------------------------
+
+Iteration: 69000
+Total loss: 0.7246640920639038
+-------------------------
+
+Iteration: 70000
+Total loss: 3.091259002685547
+
+Few shot loss: 2.754437208175659
+Few shot loss: 2.7399725914001465
+Few shot loss: 2.499554395675659
+Few shot loss: 2.1763901710510254
+Few shot loss: 1.793503999710083
+Few shot loss: 1.4362313747406006
+Few shot loss: 1.083552598953247
+Few shot loss: 0.7845061421394348
+Few shot loss: 0.5579453110694885
+Few shot loss: 0.4087105393409729
+"""
+import gymnasium as gym
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ray.rllib.algorithms.algorithm_config import DifferentiableAlgorithmConfig
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.differentiable_learner_config import (
+    DifferentiableLearnerConfig,
+)
+from ray.rllib.core.learner.training_data import TrainingData
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_learner import (
+    MAMLTorchDifferentiableLearner,
+)
+from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_rlm import (
+    DifferentiableTorchRLModule,
+)
+from ray.rllib.examples.algorithms.classes.maml_lr_meta_learner import (
+    MAMLTorchMetaLearner,
+)
+from ray.rllib.examples.utils import add_rllib_example_script_args
+from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
+from ray.rllib.utils.framework import try_import_torch
+
+# Import torch.
+torch, _ = try_import_torch()
+
+# Implement generation of data from sinusoid curves.
+def generate_sinusoid_task(batch_size, noise_std=0.1, return_params=False):
+    """Generate a sinusoid task with random amplitude and phase.
+
+    Args:
+        batch_size: The number of data points to be generated.
+        noise_std: An optional standard deviation to be used in the sinusoid
+            data generation. Defines a linear error term added to the sine
+            curve.
+        return_params: If the sampled amplitude and phase should be returned.
+
+    Returns:
+        Torch tensors with the support data and the labels of a sinusoid
+        curve.
+    """
+    # Sample the amplitude and the phase for a task.
+    amplitude = np.random.uniform(0.1, 5.0)
+    phase = np.random.uniform(0.0, np.pi)
+
+    # Sample the support.
+    x = np.random.uniform(-5.0, 5.0, (batch_size, 1))
+
+    # Generate the labels.
+    y = amplitude * np.sin(x - phase)
+
+    # Add noise.
+    y += noise_std * np.random.random((batch_size, 1))
+
+    # If sampled parameters should be returned.
+    if return_params:
+        # Return torch tensors.
+        return (
+            torch.tensor(x, dtype=torch.float32),
+            torch.tensor(y, dtype=torch.float32),
+            amplitude,
+            phase,
+        )
+    # Otherwise, return only the sampled data.
+    else:
+        return (
+            torch.tensor(x, dtype=torch.float32),
+            torch.tensor(y, dtype=torch.float32),
+        )
+
+
+def sample_task(batch_size=10, noise_std=0.1, training_data=False, return_params=False):
+    """Samples training batches for meta learner and differentiable learner.
+
+    Args:
+        batch_size: The batch size for both meta learning and task learning.
+        noise_std: An optional standard deviation to be used in the sinusoid
+            data generation. Defines a linear error term added to the sine
+            curve.
+        training_data: Whether data should be returned as `TrainingData`.
+            Otherwise, a `MultiAgentBatch` is returned. Default is `False`.
+        return_params: If the sampled amplitude and phase should be returned.
+
+    Returns:
+        A tuple with training batches for the meta learner and the differentiable
+        learner. If `training_data` is `True`, the data is wrapped into
+        `TrainingData`, otherwise both batches are `MultiAgentBatch`es.
+
+    """
+    # Generate training data for meta learner and differentiable learner.
+    train_batch = {}
+    generated_data = generate_sinusoid_task(
+        batch_size * 2, noise_std=noise_std, return_params=return_params
+    )
+    train_batch[Columns.OBS], train_batch["y"] = generated_data[:2]
+
+    # Convert to `MultiAgentBatch`.
+    meta_train_batch = MultiAgentBatch(
+        env_steps=batch_size,
+        policy_batches={
+            DEFAULT_MODULE_ID: SampleBatch(
+                {k: train_batch[k][:batch_size] for k in train_batch}
+            )
+        },
+    )
+    task_train_batch = MultiAgentBatch(
+        env_steps=batch_size,
+        policy_batches={
+            DEFAULT_MODULE_ID: SampleBatch(
+                {k: train_batch[k][batch_size:] for k in train_batch}
+            )
+        },
+    )
+    # If necessary convert to `TrainingData`.
+    if training_data:
+        meta_train_batch = TrainingData(
+            batch=meta_train_batch,
+        )
+        task_train_batch = TrainingData(
+            batch=task_train_batch,
+        )
+
+    # If amplitude and phase should be returned add them to the return tuple.
+    if return_params:
+        return meta_train_batch, task_train_batch, *generated_data[2:]
+    # Otherwise return solely train data.
+    else:
+        return meta_train_batch, task_train_batch
+
+
+# Define arguments.
+parser = add_rllib_example_script_args(default_iters=70_000)
+
+parser.add_argument(
+    "--meta-train-batch-size",
+    type=int,
+    default=5,
+    help="The number of samples per train and test update (meta-learning).",
+)
+parser.add_argument(
+    "--meta-lr",
+    type=float,
+    default=0.001,
+    help="The learning rate to be used for meta learning (in the `MetaLearner`).",
+)
+parser.add_argument(
+    "--fine-tune-batch-size",
+    type=int,
+    default=10,
+    help="The number of samples for the fine-tuning updates.",
+)
+parser.add_argument(
+    "--noise-std",
+    type=float,
+    default=0.0,
+    help="The standard deviation for noise added to the single tasks.",
+)
+parser.add_argument(
+    "--seed",
+    type=int,
+    default=None,
+    help="An optional random seed. If not set, the experiment is not reproducable.",
+)
+parser.add_argument(
+    "--fine-tune-iters",
+    type=int,
+    default=10,
+    help="The number of updates in fine-tuning.",
+)
+parser.add_argument(
+    "--fine-tune-lr",
+    type=float,
+    default=0.01,
+    help="The learning rate to be used in fine-tuning the model in the test phase.",
+)
+parser.add_argument(
+    "--no-plot",
+    action="store_true",
+    help=(
+        "If plotting should suppressed. Otherwise user action is needed to close "
+        "the plot early."
+    ),
+)
+parser.add_argument(
+    "--pause-plot-secs",
+    type=int,
+    default=1000,
+    help=(
+        "The number of seconds to keep the plot open. Note the plot can always be "
+        "closed by the user when open."
+    ),
+)
+
+# Parse the arguments.
+args = parser.parse_args()
+
+# If a random seed is provided set it for torch and numpy.
+if args.seed:
+    torch.random.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+
+if __name__ == "__main__":
+    # Define the `RLModule`.
+    module_spec = RLModuleSpec(
+        module_class=DifferentiableTorchRLModule,
+        # Note, the spaces are needed by default but are not used.
+        observation_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32),
+        action_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32),
+    )
+    # `Learner`s work on `MultiRLModule`s.
+    multi_module_spec = MultiRLModuleSpec(
+        rl_module_specs={DEFAULT_MODULE_ID: module_spec}
+    )
+
+    # Build the `MultiRLModule`.
+    module = multi_module_spec.build()
+
+    # Configure the `DifferentiableLearner`.
+    diff_learner_config = DifferentiableLearnerConfig(
+        learner_class=MAMLTorchDifferentiableLearner,
+        minibatch_size=args.meta_train_batch_size,
+        lr=0.01,
+    )
+
+    # Configure the `TorchMetaLearner` via the `DifferentiableAlgorithmConfig`.
+    config = (
+        DifferentiableAlgorithmConfig()
+        .learners(
+            # Add the `DifferentiableLearnerConfig`s.
+            differentiable_learner_configs=[diff_learner_config],
+            num_gpus_per_learner=args.num_gpus_per_learner or 0,
+        )
+        .training(
+            lr=args.meta_lr,
+            train_batch_size=args.meta_train_batch_size,
+            # Use the full batch in a single update.
+            minibatch_size=args.meta_train_batch_size,
+        )
+    )
+
+    # Initialize the `TorchMetaLearner`.
+    meta_learner = MAMLTorchMetaLearner(config=config, module_spec=module_spec)
+    # Build the `TorchMetaLearner`.
+    meta_learner.build()
+
+    for i in range(args.stop_iters):
+        # Sample the training data.
+        meta_training_data, task_training_data = sample_task(
+            args.meta_train_batch_size, noise_std=args.noise_std, training_data=True
+        )
+
+        # Update the module.
+        outs = meta_learner.update(
+            training_data=meta_training_data,
+            num_epochs=1,
+            others_training_data=[task_training_data],
+        )
+        iter = i + 1
+        if iter % 1000 == 0:
+            total_loss = outs["default_policy"]["total_loss"].peek()
+            print("-------------------------\n")
+            print(f"Iteration: {iter}")
+            print(f"Total loss: {total_loss}")
+
+    # Generate test data.
+    test_batch, _, amplitude, phase = sample_task(
+        batch_size=args.fine_tune_batch_size,
+        noise_std=args.noise_std,
+        return_params=True,
+    )
+
+    if config.num_gpus_per_learner > 0:
+        test_batch = meta_learner._convert_batch_type(test_batch)
+
+    # Run inference and plot results.
+    with torch.no_grad():
+        # Generate a grid for the support.
+        x_grid = torch.tensor(
+            np.arange(-5.0, 5.0, 0.02), dtype=torch.float32, device=meta_learner._device
+        ).view(-1, 1)
+        # Get label prediction from the model trained by MAML.
+        y_pred = meta_learner.module[DEFAULT_MODULE_ID]({Columns.OBS: x_grid})["y_pred"]
+
+    # Plot the results if requested.
+    if not args.no_plot:
+        # Sort the data by the support.
+        x_order = np.argsort(test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0])
+        x_sorted = test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0][x_order]
+        y_sorted = test_batch[DEFAULT_MODULE_ID]["y"][:, 0][x_order]
+
+        # Plot the data.
+        def sinusoid(t):
+            return amplitude * np.sin(t - phase)
+
+        plt.ion()
+        plt.figure(figsize=(5, 3))
+        # Plot the true sinusoid curve.
+        plt.plot(x_grid, sinusoid(x_grid), "r", label="Ground Truth")
+        # Add the sampled support values.
+        plt.plot(x_sorted, y_sorted, "^", color="purple")
+        # Add the prediction made by the model after MAML training.
+        plt.plot(x_grid, y_pred, ":", label="Prediction", color="#90EE90")
+        plt.title(f"MAML Results from {args.fine_tune_iters} fine-tuning steps.")
+
+    # Fine-tune with the meta loss for just a few steps.
+    optim = meta_learner.get_optimizers_for_module(DEFAULT_MODULE_ID)[0][1]
+    # Set the learning rate to a larger value.
+    for g in optim.param_groups:
+        g["lr"] = args.fine_tune_lr
+    # Now run the fine-tune iterations and update the model via the meta-learner loss.
+    for i in range(args.fine_tune_iters):
+        # Forward pass.
+        fwd_out = {
+            DEFAULT_MODULE_ID: meta_learner.module[DEFAULT_MODULE_ID](
+                test_batch[DEFAULT_MODULE_ID]
+            )
+        }
+        # Compute the MSE prediction loss.
+        loss_per_module = meta_learner.compute_losses(fwd_out=fwd_out, batch=test_batch)
+        # Optimize parameters.
+        optim.zero_grad(set_to_none=True)
+        loss_per_module[DEFAULT_MODULE_ID].backward()
+        optim.step()
+        # Show the loss for few-shot learning (fine-tuning).
+        print(f"Few shot loss: {loss_per_module[DEFAULT_MODULE_ID].item()}")
+
+    # Run the model again after fine-tuning.
+    with torch.no_grad():
+        y_pred_fine_tuned = meta_learner.module[DEFAULT_MODULE_ID](
+            {Columns.OBS: x_grid}
+        )["y_pred"]
+
+    if not args.no_plot:
+        # Plot the predictions of the fine-tuned model.
+        plt.plot(
+            x_grid,
+            y_pred_fine_tuned,
+            "-.",
+            label="Tuned Prediction",
+            color="green",
+            mfc="gray",
+        )
+        plt.legend()
+        plt.show()
+
+        # Pause the plot until the user closes it.
+        plt.pause(args.pause_plot_secs)

rllib/examples/algorithms/marwil/cartpole_marwil.py

@@ -0,0 +1,87 @@
+import warnings
+from pathlib import Path
+
+from ray.rllib.algorithms.marwil import MARWILConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+assert (
+    args.env == "CartPole-v1" or args.env is None
+), "This tuned example works only with `CartPole-v1`."
+
+# Define the data paths.
+data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+base_path = Path(__file__).parents[3]
+print(f"base_path={base_path}")
+data_path = "local://" / base_path / data_path
+print(f"data_path={data_path}")
+
+# Define the MARWIL config.
+config = (
+    MARWILConfig()
+    .environment(env="CartPole-v1")
+    # Evaluate every 3 training iterations.
+    .evaluation(
+        evaluation_interval=3,
+        evaluation_num_env_runners=1,
+        evaluation_duration=5,
+        evaluation_parallel_to_training=True,
+        evaluation_config=MARWILConfig.overrides(explore=False),
+    )
+    # Note, the `input_` argument is the major argument for the
+    # new offline API. Via the `input_read_method_kwargs` the
+    # arguments for the `ray.data.Dataset` read method can be
+    # configured. The read method needs at least as many blocks
+    # as remote learners.
+    .offline_data(
+        input_=[data_path.as_posix()],
+        # The `kwargs` for the `map_batches` method in which our
+        # `OfflinePreLearner` is run. 2 data workers should be run
+        # concurrently.
+        map_batches_kwargs={"concurrency": 2, "num_cpus": 1},
+        # The `kwargs` for the `iter_batches` method. Due to the small
+        # dataset we choose only a single batch to prefetch.
+        iter_batches_kwargs={"prefetch_batches": 1},
+        # The number of iterations to be run per learner when in multi-learner
+        # mode in a single RLlib training iteration. Leave this to `None` to
+        # run an entire epoch on the dataset during a single RLlib training
+        # iteration.
+        dataset_num_iters_per_learner=5,
+    )
+    .training(
+        beta=1.0,
+        # To increase learning speed with multiple learners,
+        # increase the learning rate correspondingly.
+        lr=0.0008 * (args.num_learners or 1) ** 0.5,
+        train_batch_size_per_learner=1024,
+    )
+)
+
+if not args.no_tune:
+    warnings.warn(
+        "You are running the example with Ray Tune. Offline RL uses "
+        "Ray Data, which doesn't does not interact seamlessly with Ray Tune. "
+        "If you encounter difficulties try to run the example without "
+        "Ray Tune using `--no-tune`."
+    )
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 250.0,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: 500000,
+}
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/ppo/atari_ppo.py

@@ -0,0 +1,97 @@
+# These tags allow extracting portions of this script on Anyscale.
+# ws-template-imports-start
+import gymnasium as gym
+
+from ray import tune
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule
+from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+# ws-template-imports-end
+
+parser = add_rllib_example_script_args(
+    default_reward=float("inf"),
+    default_timesteps=3000000,
+    default_iters=100000000000,
+)
+parser.set_defaults(
+    env="ale_py:ALE/Pong-v5",
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+NUM_LEARNERS = args.num_learners or 1
+ENV = args.env
+
+
+# These tags allow extracting portions of this script on Anyscale.
+# ws-template-code-start
+def _make_env_to_module_connector(env, spaces, device):
+    return FrameStackingEnvToModule(num_frames=4)
+
+
+def _make_learner_connector(input_observation_space, input_action_space):
+    return FrameStackingLearner(num_frames=4)
+
+
+# Create a custom Atari setup (w/o the usual RLlib-hard-coded framestacking in it).
+# We would like our frame stacking connector to do this job.
+def _env_creator(cfg):
+    return wrap_atari_for_new_api_stack(
+        gym.make(ENV, **cfg, render_mode="rgb_array"),
+        # Perform frame-stacking through ConnectorV2 API.
+        framestack=None,
+    )
+
+
+tune.register_env("env", _env_creator)
+
+config = (
+    PPOConfig()
+    .environment(
+        "env",
+        env_config={
+            # Make analogous to old v4 + NoFrameskip.
+            "frameskip": 1,
+            "full_action_space": False,
+            "repeat_action_probability": 0.0,
+        },
+        clip_rewards=True,
+    )
+    .env_runners(
+        env_to_module_connector=_make_env_to_module_connector,
+    )
+    .training(
+        learner_connector=_make_learner_connector,
+        train_batch_size_per_learner=4000,
+        minibatch_size=128,
+        lambda_=0.95,
+        kl_coeff=0.5,
+        clip_param=0.1,
+        vf_clip_param=10.0,
+        entropy_coeff=0.01,
+        num_epochs=10,
+        lr=0.00015 * NUM_LEARNERS,
+        grad_clip=100.0,
+        grad_clip_by="global_norm",
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            conv_filters=[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]],
+            conv_activation="relu",
+            head_fcnet_hiddens=[256],
+            vf_share_layers=True,
+        ),
+    )
+)
+# ws-template-code-end
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args=args)

rllib/examples/algorithms/ppo/benchmark_ppo_mujoco.py

@@ -0,0 +1,138 @@
+from ray import tune
+from ray.rllib.algorithms.ppo.ppo import PPOConfig
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune import Stopper
+
+# Needs the following packages to be installed on Ubuntu:
+#   sudo apt-get libosmesa-dev
+#   sudo apt-get install patchelf
+#   python -m pip install "gymnasium[mujoco]"
+# Might need to be added to bashsrc:
+#   export MUJOCO_GL=osmesa"
+#   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco200/bin"
+
+# See the following links for becnhmark results of other libraries:
+#   Original paper: https://arxiv.org/pdf/1707.06347
+#   CleanRL: https://wandb.ai/openrlbenchmark/openrlbenchmark/reports"
+#       /MuJoCo-CleanRL-s-PPO--VmlldzoxODAwNjkw
+#   AgileRL: https://github.com/AgileRL/AgileRL?tab=readme-ov-file#benchmarks
+benchmark_envs = {
+    "HalfCheetah-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "Hopper-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2250,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "InvertedPendulum-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "InvertedDoublePendulum-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 8000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "Reacher-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -15,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "Swimmer-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 120,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "Walker2d-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 3500,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+}
+
+
+# Define a `tune.Stopper` that stops the training if the benchmark is reached
+# or the maximum number of timesteps is exceeded.
+class BenchmarkStopper(Stopper):
+    def __init__(self, benchmark_envs):
+        self.benchmark_envs = benchmark_envs
+
+    def __call__(self, trial_id, result):
+        # Stop training if the mean reward is reached.
+        if (
+            result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
+            >= self.benchmark_envs[result["env"]][
+                f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+            ]
+        ):
+            return True
+        # Otherwise check, if the total number of timesteps is exceeded.
+        elif (
+            result[f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+            >= self.benchmark_envs[result["env"]][f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+        ):
+            return True
+        # Otherwise continue training.
+        else:
+            return False
+
+    # Note, this needs to implemented b/c the parent class is abstract.
+    def stop_all(self):
+        return False
+
+
+config = (
+    PPOConfig()
+    .environment(env=tune.grid_search(list(benchmark_envs.keys())))
+    .env_runners(
+        # Following the paper.
+        num_env_runners=32,
+        rollout_fragment_length=512,
+    )
+    .learners(
+        # Let's start with a small number of learner workers and
+        # add later a tune grid search for these resources.
+        num_learners=1,
+        num_gpus_per_learner=1,
+    )
+    # TODO (simon): Adjust to new model_config_dict.
+    .training(
+        # Following the paper.
+        lambda_=0.95,
+        lr=0.0003,
+        num_epochs=15,
+        train_batch_size=32 * 512,
+        minibatch_size=4096,
+        vf_loss_coeff=0.01,
+        model={
+            "fcnet_hiddens": [64, 64],
+            "fcnet_activation": "tanh",
+            "vf_share_layers": True,
+        },
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+        min_sample_timesteps_per_iteration=1000,
+    )
+    .evaluation(
+        evaluation_duration="auto",
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+        evaluation_parallel_to_training=True,
+        evaluation_config={
+            "explore": True,
+        },
+    )
+)
+
+tuner = tune.Tuner(
+    "PPO",
+    param_space=config,
+    run_config=tune.RunConfig(
+        stop=BenchmarkStopper(benchmark_envs=benchmark_envs),
+        name="benchmark_ppo_mujoco",
+    ),
+)
+
+tuner.fit()

rllib/examples/algorithms/ppo/cartpole_heavy_ppo.py

@@ -0,0 +1,38 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import FlattenObservations
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.cartpole_with_large_observation_space import (
+    CartPoleWithLargeObservationSpace,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(default_reward=450.0, default_timesteps=300000)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    PPOConfig()
+    .environment(CartPoleWithLargeObservationSpace)
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: FlattenObservations(),
+        episodes_to_numpy=False,
+    )
+    .training(
+        lr=0.0003,
+        num_epochs=6,
+        vf_loss_coeff=0.01,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            use_lstm=True,
+            lstm_cell_size=1024,
+        ),
+    )
+)
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/ppo/cartpole_ppo.py

@@ -0,0 +1,32 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(default_reward=450.0, default_timesteps=300000)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    PPOConfig()
+    .environment("CartPole-v1")
+    .training(
+        lr=0.0003,
+        num_epochs=6,
+        vf_loss_coeff=0.01,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[32],
+            fcnet_activation="linear",
+            vf_share_layers=True,
+        ),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/ppo/cartpole_truncated_ppo.py

@@ -0,0 +1,55 @@
+import gymnasium as gym
+from gymnasium.wrappers import TimeLimit
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args()
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+# For training, use a time-truncated (max. 50 timestep) version of CartPole-v1.
+register_env(
+    "cartpole_truncated",
+    lambda _: TimeLimit(gym.make("CartPole-v1"), max_episode_steps=50),
+)
+
+config = (
+    PPOConfig()
+    .environment("cartpole_truncated")
+    .env_runners(num_envs_per_env_runner=10)
+    .training(
+        lr=0.0003,
+        num_epochs=6,
+        vf_loss_coeff=0.01,
+    )
+    # For evaluation, use the "real" CartPole-v1 env (up to 500 steps).
+    .evaluation(
+        evaluation_config=PPOConfig.overrides(
+            env="CartPole-v1",
+            explore=False,
+        ),
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+    )
+)
+
+stop = {
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 500000,
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 80.0,
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/ppo/memory_leak_test_ppo_new_stack.py

@@ -0,0 +1,13 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.examples.envs.classes.random_env import RandomLargeObsSpaceEnv
+
+config = (
+    PPOConfig()
+    # Switch off np.random, which is known to have memory leaks.
+    .environment(RandomLargeObsSpaceEnv, env_config={"static_samples": True})
+    .env_runners(
+        num_env_runners=4,
+        num_envs_per_env_runner=5,
+    )
+    .training(train_batch_size=500, minibatch_size=256, num_epochs=5)
+)

rllib/examples/algorithms/ppo/multi_agent_cartpole_ppo.py

@@ -0,0 +1,57 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args()
+parser.set_defaults(
+    num_agents=2,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("multi_agent_cartpole", lambda cfg: MultiAgentCartPole(config=cfg))
+
+config = (
+    PPOConfig()
+    .environment("multi_agent_cartpole", env_config={"num_agents": args.num_agents})
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[32],
+            fcnet_activation="linear",
+            vf_share_layers=True,
+        ),
+    )
+    .env_runners(
+        num_envs_per_env_runner=2,
+    )
+    .training(
+        lr=0.0003,
+        num_epochs=6,
+        vf_loss_coeff=0.01,
+    )
+    .multi_agent(
+        policy_mapping_fn=lambda aid, *arg, **kw: f"p{aid}",
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+)
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: 400000,
+    # Divide by num_agents to get actual return per agent.
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 450.0 * (args.num_agents or 1),
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/ppo/multi_agent_footsies_ppo.py

@@ -0,0 +1,264 @@
+"""
+Multi-agent RLlib Footsies Example (PPO)
+
+About:
+    - Example is based on the Footsies environment (https://github.com/chasemcd/FootsiesGym).
+    - Footsies is a two-player fighting game where each player controls a character and tries to hit the opponent while avoiding being hit.
+    - Footsies is a zero-sum game, when one player wins (+1 reward) the other loses (-1 reward).
+
+Summary:
+    - Main policy is an LSTM-based policy.
+    - Training algorithm is PPO.
+
+Training:
+    - Training is governed by adding new, more complex opponents to the mix as the main policy reaches a certain win rate threshold against the current opponent.
+    - Current opponent is always the newest opponent added to the mix.
+    - Training starts with a very simple opponent: "noop" (does nothing), then progresses to "back" (only moves backwards). These are the fixed (very simple) policies that are used to kick off the training.
+    - After "random", new opponents are frozen copies of the main policy at different training stages. They will be added to the mix as "lstm_v0", "lstm_v1", etc.
+    - In this way - after kick-starting the training with fixed simple opponents - the main policy will play against a version of itself from an earlier training stage.
+    - The main policy has to achieve the win rate threshold against the current opponent to add a new opponent to the mix.
+    - Training concludes when the target mix size is reached.
+
+Evaluation:
+    - Evaluation is performed against the current (newest) opponent.
+    - Evaluation runs for a fixed number of episodes at the end of each training iteration.
+
+"""
+import functools
+from pathlib import Path
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module import MultiRLModuleSpec, RLModuleSpec
+from ray.rllib.env.multi_agent_env_runner import MultiAgentEnvRunner
+from ray.rllib.examples.envs.classes.multi_agent.footsies.fixed_rlmodules import (
+    BackFixedRLModule,
+    NoopFixedRLModule,
+)
+from ray.rllib.examples.envs.classes.multi_agent.footsies.footsies_env import (
+    env_creator,
+)
+from ray.rllib.examples.envs.classes.multi_agent.footsies.utils import (
+    Matchmaker,
+    Matchup,
+    MetricsLoggerCallback,
+    MixManagerCallback,
+    platform_for_binary_to_download,
+)
+from ray.rllib.examples.rl_modules.classes.lstm_containing_rlm import (
+    LSTMContainingRLModule,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME
+from ray.tune.registry import register_env
+from ray.tune.result import TRAINING_ITERATION
+
+# setting two default stopping criteria:
+#    1. training_iteration (via "stop_iters")
+#    2. num_env_steps_sampled_lifetime (via "default_timesteps")
+# ...values very high to make sure that the test passes by adding
+# all required policies to the mix, not by hitting the iteration limit.
+# Our main stopping criterion is "target_mix_size" (see an argument below).
+parser = add_rllib_example_script_args(
+    default_iters=500,
+    default_timesteps=5_000_000,
+)
+
+parser.add_argument(
+    "--train-start-port",
+    type=int,
+    default=45001,
+    help="First port number for the Footsies training environment server (default: 45001). Each server gets its own port.",
+)
+parser.add_argument(
+    "--eval-start-port",
+    type=int,
+    default=55001,
+    help="First port number for the Footsies evaluation environment server (default: 55001) Each server gets its own port.",
+)
+parser.add_argument(
+    "--binary-download-dir",
+    type=Path,
+    default="/tmp/ray/binaries/footsies",
+    help="Directory to download Footsies binaries (default: /tmp/ray/binaries/footsies)",
+)
+parser.add_argument(
+    "--binary-extract-dir",
+    type=Path,
+    default="/tmp/ray/binaries/footsies",
+    help="Directory to extract Footsies binaries (default: /tmp/ray/binaries/footsies)",
+)
+parser.add_argument(
+    "--win-rate-threshold",
+    type=float,
+    default=0.8,
+    help="The main policy should have at least 'win-rate-threshold' win rate against the "
+    "other policy to advance to the next level. Moving to the next level "
+    "means adding a new policy to the mix.",
+)
+parser.add_argument(
+    "--target-mix-size",
+    type=int,
+    default=5,
+    help="Target number of policies (RLModules) in the mix to consider the test passed. "
+    "The initial mix size is 2: 'main policy' vs. 'other'. "
+    "`--target-mix-size=5` means that 3 new policies will be added to the mix. "
+    "Whether to add new policy is decided by checking the '--win-rate-threshold' condition. ",
+)
+parser.add_argument(
+    "--rollout-fragment-length",
+    type=int,
+    default=256,
+    help="The length of each rollout fragment to be collected by the EnvRunners when sampling.",
+)
+parser.add_argument(
+    "--log-unity-output",
+    action="store_true",
+    help="Whether to log Unity output (from the game engine). Default is False.",
+    default=False,
+)
+parser.add_argument(
+    "--render",
+    action="store_true",
+    default=False,
+    help="Whether to render the Footsies environment. Default is False.",
+)
+
+main_policy = "lstm"
+args = parser.parse_args()
+register_env(name="FootsiesEnv", env_creator=env_creator)
+
+# Detect platform and choose appropriate binary
+binary_to_download = platform_for_binary_to_download(args.render)
+
+config = (
+    PPOConfig()
+    .reporting(
+        min_time_s_per_iteration=30,
+    )
+    .environment(
+        env="FootsiesEnv",
+        env_config={
+            "max_t": 1000,
+            "frame_skip": 4,
+            "observation_delay": 16,
+            "train_start_port": args.train_start_port,
+            "eval_start_port": args.eval_start_port,
+            "host": "localhost",
+            "binary_download_dir": args.binary_download_dir,
+            "binary_extract_dir": args.binary_extract_dir,
+            "binary_to_download": binary_to_download,
+            "log_unity_output": args.log_unity_output,
+        },
+    )
+    .learners(
+        num_learners=1,
+        num_cpus_per_learner=1,
+        num_gpus_per_learner=0,
+        num_aggregator_actors_per_learner=0,
+    )
+    .env_runners(
+        env_runner_cls=MultiAgentEnvRunner,
+        num_env_runners=args.num_env_runners or 1,
+        num_cpus_per_env_runner=0.5,
+        num_envs_per_env_runner=1,
+        batch_mode="truncate_episodes",
+        rollout_fragment_length=args.rollout_fragment_length,
+        episodes_to_numpy=False,
+        create_env_on_local_worker=True,
+    )
+    .training(
+        train_batch_size_per_learner=args.rollout_fragment_length
+        * (args.num_env_runners or 1),
+        lr=1e-4,
+        entropy_coeff=0.01,
+        num_epochs=10,
+        minibatch_size=128,
+    )
+    .multi_agent(
+        policies={
+            main_policy,
+            "noop",
+            "back",
+        },
+        # this is a starting policy_mapping_fn
+        # It will be updated by the MixManagerCallback during training.
+        policy_mapping_fn=Matchmaker(
+            [Matchup(main_policy, "noop", 1.0)]
+        ).agent_to_module_mapping_fn,
+        # we only train the main policy, this doesn't change during training.
+        policies_to_train=[main_policy],
+    )
+    .rl_module(
+        rl_module_spec=MultiRLModuleSpec(
+            rl_module_specs={
+                main_policy: RLModuleSpec(
+                    module_class=LSTMContainingRLModule,
+                    model_config={
+                        "lstm_cell_size": 128,
+                        "dense_layers": [128, 128],
+                        "max_seq_len": 64,
+                    },
+                ),
+                # for simplicity, all fixed RLModules are added to the config at the start.
+                # However, only "noop" is used at the start of training,
+                # the others are added to the mix later by the MixManagerCallback.
+                "noop": RLModuleSpec(module_class=NoopFixedRLModule),
+                "back": RLModuleSpec(module_class=BackFixedRLModule),
+            },
+        )
+    )
+    .evaluation(
+        evaluation_num_env_runners=args.evaluation_num_env_runners or 1,
+        evaluation_sample_timeout_s=120,
+        evaluation_interval=1,
+        evaluation_duration=10,  # 10 episodes is enough to get a good win rate estimate
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=False,
+        # we may add new RLModules to the mix at the end of the evaluation stage.
+        # Running evaluation in parallel may result in training for one more iteration on the old mix.
+        evaluation_force_reset_envs_before_iteration=True,
+        evaluation_config={
+            "env_config": {"env-for-evaluation": True},
+        },  # evaluation_config is used to add an argument to the env creator.
+    )
+    .callbacks(
+        [
+            functools.partial(
+                MetricsLoggerCallback,
+                main_policy=main_policy,
+            ),
+            functools.partial(
+                MixManagerCallback,
+                win_rate_threshold=args.win_rate_threshold,
+                main_policy=main_policy,
+                target_mix_size=args.target_mix_size,
+                starting_modules=[main_policy, "noop"],
+                fixed_modules_progression_sequence=(
+                    "noop",
+                    "back",
+                ),
+            ),
+        ]
+    )
+)
+
+# stopping criteria to be passed to Ray Tune. The main stopping criterion is "mix_size".
+# "mix_size" is reported at the end of each training iteration by the MixManagerCallback.
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    TRAINING_ITERATION: args.stop_iters,
+    "mix_size": args.target_mix_size,
+}
+
+if __name__ == "__main__":
+    results = run_rllib_example_script_experiment(
+        base_config=config,
+        args=args,
+        stop=stop,
+        success_metric={
+            "mix_size": args.target_mix_size
+        },  # pass the success metric for RLlib's testing framework
+    )

rllib/examples/algorithms/ppo/multi_agent_pendulum_ppo.py

@@ -0,0 +1,58 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(default_timesteps=500000)
+parser.set_defaults(
+    num_agents=2,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("multi_agent_pendulum", lambda cfg: MultiAgentPendulum(config=cfg))
+
+config = (
+    PPOConfig()
+    .environment("multi_agent_pendulum", env_config={"num_agents": args.num_agents})
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(
+            multi_agent=True
+        ),
+    )
+    .training(
+        train_batch_size_per_learner=1024,
+        minibatch_size=128,
+        lr=0.0002 * (args.num_learners or 1) ** 0.5,
+        gamma=0.95,
+        lambda_=0.5,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(fcnet_activation="relu"),
+    )
+    .multi_agent(
+        policy_mapping_fn=lambda aid, *arg, **kw: f"p{aid}",
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+)
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    # Divide by num_agents to get actual return per agent.
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -300.0 * (args.num_agents or 1),
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/ppo/multi_agent_stateless_cartpole_ppo.py

@@ -0,0 +1,63 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentStatelessCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune import register_env
+
+parser = add_rllib_example_script_args(default_timesteps=4000000)
+parser.set_defaults(
+    num_agents=2,
+    num_env_runners=3,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env(
+    "multi_stateless_cart",
+    lambda _: MultiAgentStatelessCartPole({"num_agents": args.num_agents}),
+)
+
+config = (
+    PPOConfig()
+    .environment("multi_stateless_cart")
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(
+            multi_agent=True
+        ),
+    )
+    .training(
+        lr=0.0003 * ((args.num_learners or 1) ** 0.5),
+        num_epochs=6,
+        vf_loss_coeff=0.05,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            use_lstm=True,
+            max_seq_len=20,
+        ),
+    )
+    .multi_agent(
+        policy_mapping_fn=lambda aid, *arg, **kw: f"p{aid}",
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+)
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    # Divide by num_agents to get actual return per agent.
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 300.0 * (args.num_agents or 1),
+}
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/ppo/pendulum_ppo.py

@@ -0,0 +1,37 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(default_timesteps=400000, default_reward=-300)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    PPOConfig()
+    .environment("Pendulum-v1")
+    .env_runners(
+        num_env_runners=2,
+        num_envs_per_env_runner=20,
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+    )
+    .training(
+        train_batch_size_per_learner=1024,
+        minibatch_size=128,
+        lr=0.0002 * (args.num_learners or 1) ** 0.5,
+        gamma=0.95,
+        lambda_=0.5,
+        # num_epochs=8,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(fcnet_activation="relu"),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/ppo/stateless_cartpole_ppo.py

@@ -0,0 +1,43 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import MeanStdFilter
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=2000000,
+    default_reward=350.0,
+)
+parser.set_defaults(
+    num_env_runners=3,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    PPOConfig()
+    .environment(StatelessCartPole)
+    .env_runners(
+        env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
+    )
+    .training(
+        lr=0.0003 * ((args.num_learners or 1) ** 0.5),
+        num_epochs=6,
+        vf_loss_coeff=0.05,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            vf_share_layers=False,
+            use_lstm=True,
+            max_seq_len=20,
+        ),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/sac/benchmark_sac_mujoco.py

@@ -0,0 +1,143 @@
+from ray import tune
+from ray.rllib.algorithms.sac.sac import SACConfig
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune import Stopper
+
+# Needs the following packages to be installed on Ubuntu:
+#   sudo apt-get libosmesa-dev
+#   sudo apt-get install patchelf
+#   python -m pip install "gymnasium[mujoco]"
+# Might need to be added to bashsrc:
+#   export MUJOCO_GL=osmesa"
+#   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco200/bin"
+
+# See the following links for becnhmark results of other libraries:
+#   Original paper: https://arxiv.org/abs/1812.05905
+#   CleanRL: https://wandb.ai/cleanrl/cleanrl.benchmark/reports/Mujoco--VmlldzoxODE0NjE
+#   AgileRL: https://github.com/AgileRL/AgileRL?tab=readme-ov-file#benchmarks
+benchmark_envs = {
+    "HalfCheetah-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 15000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 3000000,
+    },
+    "Hopper-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 3500,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 1000000,
+    },
+    "Humanoid-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 8000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 10000000,
+    },
+    "Ant-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 5500,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 3000000,
+    },
+    "Walker2d-v4": {
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6000,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 3000000,
+    },
+}
+
+
+# Define a `tune.Stopper` that stops the training if the benchmark is reached
+# or the maximum number of timesteps is exceeded.
+class BenchmarkStopper(Stopper):
+    def __init__(self, benchmark_envs):
+        self.benchmark_envs = benchmark_envs
+
+    def __call__(self, trial_id, result):
+        # Stop training if the mean reward is reached.
+        if (
+            result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
+            >= self.benchmark_envs[result["env"]][
+                f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
+            ]
+        ):
+            return True
+        # Otherwise check, if the total number of timesteps is exceeded.
+        elif (
+            result[f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+            >= self.benchmark_envs[result["env"]][f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
+        ):
+            return True
+        # Otherwise continue training.
+        else:
+            return False
+
+    # Note, this needs to implemented b/c the parent class is abstract.
+    def stop_all(self):
+        return False
+
+
+config = (
+    SACConfig()
+    .environment(env=tune.grid_search(list(benchmark_envs.keys())))
+    .env_runners(
+        rollout_fragment_length=1,
+        num_env_runners=0,
+    )
+    .learners(
+        # Note, we have a sample/train ratio of 1:1 and a small train
+        # batch, so 1 learner with a single GPU should suffice.
+        num_learners=1,
+        num_gpus_per_learner=1,
+    )
+    # TODO (simon): Adjust to new model_config_dict.
+    .training(
+        initial_alpha=1.001,
+        # Choose a smaller learning rate for the actor (policy).
+        actor_lr=3e-5,
+        critic_lr=3e-4,
+        alpha_lr=1e-4,
+        target_entropy="auto",
+        n_step=1,
+        tau=0.005,
+        train_batch_size=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 1000000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        num_steps_sampled_before_learning_starts=256,
+        model={
+            "fcnet_hiddens": [256, 256],
+            "fcnet_activation": "relu",
+            "post_fcnet_hiddens": [],
+            "post_fcnet_activation": None,
+            "post_fcnet_weights_initializer": "orthogonal_",
+            "post_fcnet_weights_initializer_config": {"gain": 0.01},
+            "fusionnet_hiddens": [256, 256, 256],
+            "fusionnet_activation": "relu",
+        },
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+        min_sample_timesteps_per_iteration=1000,
+    )
+    .evaluation(
+        evaluation_duration="auto",
+        evaluation_interval=1,
+        evaluation_num_env_runners=1,
+        evaluation_parallel_to_training=True,
+        evaluation_config={
+            "explore": False,
+        },
+    )
+)
+
+tuner = tune.Tuner(
+    "SAC",
+    param_space=config,
+    run_config=tune.RunConfig(
+        stop=BenchmarkStopper(benchmark_envs=benchmark_envs),
+        name="benchmark_sac_mujoco",
+    ),
+)
+
+tuner.fit()

rllib/examples/algorithms/sac/halfcheetah_sac.py

@@ -0,0 +1,64 @@
+from torch import nn
+
+from ray.rllib.algorithms.sac.sac import SACConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=1000000,
+    default_reward=12000.0,
+    default_iters=2000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    SACConfig()
+    .environment("HalfCheetah-v4")
+    .training(
+        initial_alpha=1.001,
+        # lr=0.0006 is very high, w/ 4 GPUs -> 0.0012
+        # Might want to lower it for better stability, but it does learn well.
+        actor_lr=2e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=8e-4 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=9e-4 * (args.num_learners or 1) ** 0.5,
+        lr=None,
+        target_entropy="auto",
+        n_step=(1, 5),  # 1?
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 100000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        num_steps_sampled_before_learning_starts=10000,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        ),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+        min_sample_timesteps_per_iteration=1000,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/sac/humanoid_sac.py

@@ -0,0 +1,71 @@
+"""This is WIP.
+
+On a single-GPU machine, with the `--num-gpus-per-learner=1` command line option, this
+example should learn a episode return of >1000 in ~10h, which is still very basic, but
+does somewhat prove SAC's capabilities. Some more hyperparameter fine tuning, longer
+runs, and more scale (`--num-learners > 0` and `--num-env-runners > 0`) should help push
+this up.
+"""
+
+from torch import nn
+
+from ray.rllib.algorithms.sac.sac import SACConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=1000000,
+    default_reward=12000.0,
+    default_iters=2000,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+
+config = (
+    SACConfig()
+    .environment("Humanoid-v4")
+    .training(
+        initial_alpha=1.001,
+        actor_lr=0.00005,
+        critic_lr=0.00005,
+        alpha_lr=0.00005,
+        target_entropy="auto",
+        n_step=(1, 3),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 1000000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        num_steps_sampled_before_learning_starts=10000,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[1024, 1024],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        )
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+        min_sample_timesteps_per_iteration=1000,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/sac/mountaincar_sac.py

@@ -0,0 +1,61 @@
+from torch import nn
+
+from ray.rllib.algorithms.sac.sac import SACConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=20000,
+    default_reward=-250.0,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    SACConfig()
+    .environment("MountainCar-v0")
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        ),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+    )
+    .training(
+        initial_alpha=1.001,
+        # Use a smaller learning rate for the policy.
+        actor_lr=2e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=8e-4 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=9e-4 * (args.num_learners or 1) ** 0.5,
+        lr=None,
+        target_entropy="auto",
+        n_step=(2, 5),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 100000,
+            "alpha": 1.0,
+            "beta": 0.0,
+        },
+        num_steps_sampled_before_learning_starts=256 * (args.num_learners or 1),
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/sac/multi_agent_pendulum_sac.py

@@ -0,0 +1,87 @@
+from torch import nn
+
+from ray.rllib.algorithms.sac import SACConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+parser = add_rllib_example_script_args(
+    default_timesteps=500000,
+)
+parser.set_defaults(
+    num_agents=2,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+register_env("multi_agent_pendulum", lambda cfg: MultiAgentPendulum(config=cfg))
+
+config = (
+    SACConfig()
+    .environment("multi_agent_pendulum", env_config={"num_agents": args.num_agents})
+    .training(
+        initial_alpha=1.001,
+        # Use a smaller learning rate for the policy.
+        actor_lr=2e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=8e-4 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=9e-4 * (args.num_learners or 1) ** 0.5,
+        lr=None,
+        target_entropy="auto",
+        n_step=(2, 5),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "MultiAgentPrioritizedEpisodeReplayBuffer",
+            "capacity": 100000,
+            "alpha": 1.0,
+            "beta": 0.0,
+        },
+        num_steps_sampled_before_learning_starts=256,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer=nn.init.orthogonal_,
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        ),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+    )
+)
+
+if args.num_agents > 0:
+    config.multi_agent(
+        policy_mapping_fn=lambda aid, *arg, **kw: f"p{aid}",
+        policies={f"p{i}" for i in range(args.num_agents)},
+    )
+
+stop = {
+    NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+    # `episode_return_mean` is the sum of all agents/policies' returns.
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -450.0 * args.num_agents,
+}
+
+if __name__ == "__main__":
+    assert (
+        args.num_agents > 0
+    ), "The `--num-agents` arg must be > 0 for this script to work."
+
+    run_rllib_example_script_experiment(config, args, stop=stop)

rllib/examples/algorithms/sac/pendulum_sac.py

@@ -0,0 +1,61 @@
+from torch import nn
+
+from ray.rllib.algorithms.sac.sac import SACConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=20000,
+    default_reward=-250.0,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    SACConfig()
+    .environment("Pendulum-v1")
+    .training(
+        initial_alpha=1.001,
+        # Use a smaller learning rate for the policy.
+        actor_lr=2e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=8e-4 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=9e-4 * (args.num_learners or 1) ** 0.5,
+        # TODO (sven): Maybe go back to making this a dict of the sub-learning rates?
+        lr=None,
+        target_entropy="auto",
+        n_step=(2, 5),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 100000,
+            "alpha": 1.0,
+            "beta": 0.0,
+        },
+        num_steps_sampled_before_learning_starts=256 * (args.num_learners or 1),
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        ),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+    )
+)
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/tqc/humanoid_tqc.py

@@ -0,0 +1,122 @@
+"""Example showing how to train TQC on the Humanoid-v4 MuJoCo environment.
+
+TQC (Truncated Quantile Critics) is an extension of SAC that uses distributional
+critics with quantile regression. By truncating the upper quantiles when computing
+target values, TQC reduces overestimation bias that can plague actor-critic methods,
+leading to more stable and efficient learning on complex continuous control tasks.
+
+This example:
+- Trains on Humanoid-v4, a challenging 17-DoF locomotion task
+- Uses truncated quantile critics with 25 quantiles and 2 critics
+- Drops the top 2 quantiles per network to reduce overestimation bias
+- Employs prioritized experience replay with capacity of 1M transitions
+- Uses a large network architecture (1024x1024) suitable for high-dimensional control
+- Applies mixed n-step returns (1 to 3 steps) for variance reduction
+- Expects to achieve episode returns >12000 with sufficient training
+
+How to run this script
+----------------------
+`python humanoid_tqc.py --num-env-runners=4`
+
+For faster training, use GPU acceleration and more parallelism:
+`python humanoid_tqc.py --num-learners=1 --num-gpus-per-learner=1 --num-env-runners=8`
+
+To scale up with distributed learning using multiple learners and env-runners:
+`python humanoid_tqc.py --num-learners=2 --num-env-runners=16`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+ --wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+Results to expect
+-----------------
+On a single-GPU machine with --num-gpus-per-learner=1, this example should learn
+an episode return of >1000 within approximately 10 hours. With more hyperparameter
+tuning, longer runs, and additional scale, returns of >12000 are achievable.
+"""
+
+from torch import nn
+
+from ray.rllib.algorithms.tqc.tqc import TQCConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=1_000_000,
+    default_reward=12_000.0,
+    default_iters=2_000,
+)
+parser.set_defaults(
+    num_env_runners=4,
+    num_envs_per_env_runner=8,
+    num_learners=1,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+
+config = (
+    TQCConfig()
+    .environment("Humanoid-v4")
+    .env_runners(
+        num_env_runners=args.num_env_runners,
+        num_envs_per_env_runner=args.num_envs_per_env_runner,
+    )
+    .learners(
+        num_learners=args.num_learners,
+        num_gpus_per_learner=1,
+        num_aggregator_actors_per_learner=2,
+    )
+    .training(
+        initial_alpha=1.001,
+        actor_lr=0.00005,
+        critic_lr=0.00005,
+        alpha_lr=0.00005,
+        target_entropy="auto",
+        n_step=(1, 3),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        # TQC-specific parameters
+        n_quantiles=25,
+        n_critics=2,
+        top_quantiles_to_drop_per_net=2,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 1000000,
+            "alpha": 0.6,
+            "beta": 0.4,
+        },
+        num_steps_sampled_before_learning_starts=10000,
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[1024, 1024],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        )
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+        min_sample_timesteps_per_iteration=1000,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/tqc/pendulum_tqc.py

@@ -0,0 +1,123 @@
+"""Example showing how to train TQC on the Pendulum-v1 classic control environment.
+
+TQC (Truncated Quantile Critics) is an extension of SAC that uses distributional
+critics with quantile regression to reduce overestimation bias. This example
+demonstrates TQC on a simple continuous control task suitable for quick experiments.
+
+This example:
+- Trains on Pendulum-v1, a classic swing-up control task with continuous actions
+- Uses truncated quantile critics with 25 quantiles and 2 critics
+- Drops the top 2 quantiles per network to reduce overestimation bias
+- Employs prioritized experience replay with 100K capacity
+- Scales learning rates based on the number of learners for distributed training
+- Uses mixed n-step returns (2 to 5 steps) for improved sample efficiency
+- Expects to achieve episode returns of approximately -250 within 20K timesteps
+
+How to run this script
+----------------------
+`python pendulum_tqc.py`
+
+To run with different configuration:
+`python pendulum_tqc.py --num-env-runners=2`
+
+To scale up with distributed learning using multiple learners and env-runners:
+`python pendulum_tqc.py --num-learners=2 --num-env-runners=8`
+
+To use a GPU-based learner add the number of GPUs per learners:
+`python pendulum_tqc.py --num-learners=1 --num-gpus-per-learner=1`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+ --wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+Results to expect
+-----------------
+With default settings, this example should achieve an episode return of around -250
+within 20,000 timesteps. The Pendulum environment has a maximum possible return of 0
+(perfect balancing), with typical good performance in the -200 to -300 range.
+"""
+
+from torch import nn
+
+from ray.rllib.algorithms.tqc.tqc import TQCConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=20000,
+    default_reward=-250.0,
+)
+parser.set_defaults(
+    num_env_runners=4,
+    num_envs_per_env_runner=8,
+    num_learners=1,
+)
+# Use `parser` to add your own custom command line options to this script
+# and (if needed) use their values to set up `config` below.
+args = parser.parse_args()
+
+config = (
+    TQCConfig()
+    .environment("Pendulum-v1")
+    .env_runners(
+        num_env_runners=args.num_env_runners,
+        num_envs_per_env_runner=args.num_envs_per_env_runner,
+    )
+    .learners(
+        num_learners=args.num_learners,
+        num_gpus_per_learner=1,
+        num_aggregator_actors_per_learner=2,
+    )
+    .training(
+        initial_alpha=1.001,
+        # Use a smaller learning rate for the policy.
+        actor_lr=2e-4 * (args.num_learners or 1) ** 0.5,
+        critic_lr=8e-4 * (args.num_learners or 1) ** 0.5,
+        alpha_lr=9e-4 * (args.num_learners or 1) ** 0.5,
+        lr=None,
+        target_entropy="auto",
+        n_step=(2, 5),
+        tau=0.005,
+        train_batch_size_per_learner=256,
+        target_network_update_freq=1,
+        # TQC-specific parameters
+        n_quantiles=25,
+        n_critics=2,
+        top_quantiles_to_drop_per_net=2,
+        replay_buffer_config={
+            "type": "PrioritizedEpisodeReplayBuffer",
+            "capacity": 100000,
+            "alpha": 1.0,
+            "beta": 0.0,
+        },
+        num_steps_sampled_before_learning_starts=256 * (args.num_learners or 1),
+    )
+    .rl_module(
+        model_config=DefaultModelConfig(
+            fcnet_hiddens=[256, 256],
+            fcnet_activation="relu",
+            fcnet_kernel_initializer=nn.init.xavier_uniform_,
+            head_fcnet_hiddens=[],
+            head_fcnet_activation=None,
+            head_fcnet_kernel_initializer="orthogonal_",
+            head_fcnet_kernel_initializer_kwargs={"gain": 0.01},
+            fusionnet_hiddens=[256, 256, 256],
+            fusionnet_activation="relu",
+        ),
+    )
+    .reporting(
+        metrics_num_episodes_for_smoothing=5,
+    )
+)
+
+
+if __name__ == "__main__":
+    run_rllib_example_script_experiment(config, args)

rllib/examples/algorithms/vpg_custom_algorithm.py

@@ -0,0 +1,115 @@
+"""Example of how to write a custom Algorithm.
+
+This is an end-to-end example for how to implement a custom Algorithm, including
+a matching AlgorithmConfig class and Learner class. There is no particular RLModule API
+needed for this algorithm, which means that any TorchRLModule returning actions
+or action distribution parameters suffices.
+
+The RK algorithm implemented here is "vanilla policy gradient" (VPG) in its simplest
+form, without a value function baseline.
+
+See the actual VPG algorithm class here:
+https://github.com/ray-project/ray/blob/master/rllib/examples/algorithms/classes/vpg.py
+
+The Learner class the algorithm uses by default (if the user doesn't specify a custom
+Learner):
+https://github.com/ray-project/ray/blob/master/rllib/examples/learners/classes/vpg_torch_learner.py  # noqa
+
+And the RLModule class the algorithm uses by default (if the user doesn't specify a
+custom RLModule):
+https://github.com/ray-project/ray/blob/master/rllib/examples/rl_modules/classes/vpg_torch_rlm.py  # noqa
+
+This example shows:
+    - how to subclass the AlgorithmConfig base class to implement a custom algorithm's.
+    config class.
+    - how to subclass the Algorithm base class to implement a custom Algorithm,
+    including its `training_step` method.
+    - how to subclass the TorchLearner base class to implement a custom Learner with
+    loss function, overriding `compute_loss_for_module` and
+    `after_gradient_based_update`.
+    - how to define a default RLModule used by the algorithm in case the user
+    doesn't bring their own custom RLModule. The VPG algorithm doesn't require any
+    specific RLModule APIs, so any RLModule returning actions or action distribution
+    inputs suffices.
+
+We compute a plain policy gradient loss without value function baseline.
+The experiment shows that even with such a simple setup, our custom algorithm is still
+able to successfully learn CartPole-v1.
+
+
+How to run this script
+----------------------
+`python [script file name].py`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+With some fine-tuning of the learning rate, the batch size, and maybe the
+number of env runners and number of envs per env runner, you should see decent
+learning behavior on the CartPole-v1 environment:
+
++-----------------------------+------------+--------+------------------+
+| Trial name                  | status     |   iter |   total time (s) |
+|                             |            |        |                  |
+|-----------------------------+------------+--------+------------------+
+| VPG_CartPole-v1_2973e_00000 | TERMINATED |    451 |          59.5184 |
++-----------------------------+------------+--------+------------------+
++-----------------------+------------------------+------------------------+
+|   episode_return_mean |   num_env_steps_sample |   ...env_steps_sampled |
+|                       |             d_lifetime |   _lifetime_throughput |
+|-----------------------+------------------------+------------------------|
+|                250.52 |                 415787 |                7428.98 |
++-----------------------+------------------------+------------------------+
+"""
+
+from ray.rllib.examples.algorithms.classes.vpg import VPGConfig
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_reward=250.0,
+    default_iters=1000,
+    default_timesteps=1_000_000,
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    base_config = (
+        VPGConfig()
+        .environment("CartPole-v1")
+        .training(
+            # The only VPG-specific setting. How many episodes per train batch?
+            num_episodes_per_train_batch=10,
+            # Set other config parameters.
+            lr=0.0005,
+            # Note that you don't have to set any specific Learner class, because
+            # our custom Algorithm already defines the default Learner class to use
+            # through its `get_default_learner_class` method, which returns
+            # `VPGTorchLearner`.
+            # learner_class=VPGTorchLearner,
+        )
+        # Increase the number of EnvRunners (default is 1 for VPG)
+        # or the number of envs per EnvRunner.
+        .env_runners(num_env_runners=2, num_envs_per_env_runner=1)
+        # Plug in your own RLModule class. VPG doesn't require any specific
+        # RLModule APIs, so any RLModule returning `actions` or `action_dist_inputs`
+        # from the forward methods works ok.
+        # .rl_module(
+        #    rl_module_spec=RLModuleSpec(module_class=...),
+        # )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)