chore: import upstream snapshot with attribution

rllib/examples/_docs/rllib_on_rllib_readme.py

@@ -0,0 +1,103 @@
+import gymnasium as gym
+import numpy as np
+import torch
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+)
+
+
+# Define your problem using python and Farama-Foundation's gymnasium API:
+class ParrotEnv(gym.Env):
+    """Environment in which an agent must learn to repeat the seen observations.
+
+    Observations are float numbers indicating the to-be-repeated values,
+    e.g. -1.0, 5.1, or 3.2.
+
+    The action space is always the same as the observation space.
+
+    Rewards are r=-abs(observation - action), for all steps.
+    """
+
+    def __init__(self, config):
+        # Make the space (for actions and observations) configurable.
+        self.action_space = config.get(
+            "parrot_shriek_range", gym.spaces.Box(-1.0, 1.0, (1,), np.float32)
+        )
+        # Since actions should repeat observations, their spaces must be the
+        # same.
+        self.observation_space = self.action_space
+        self.cur_obs = None
+        self.episode_len = 0
+
+    def reset(self, *, seed=None, options=None):
+        """Resets the episode and returns the initial observation of the new one."""
+        # Reset the episode len.
+        self.episode_len = 0
+        # Sample a random number from our observation space.
+        self.cur_obs = self.observation_space.sample()
+        # Return initial observation.
+        return self.cur_obs, {}
+
+    def step(self, action):
+        """Takes a single step in the episode given `action`
+
+        Returns: New observation, reward, done-flag, info-dict (empty).
+        """
+        # Set `terminated` and `truncated` flags to True after 10 steps.
+        self.episode_len += 1
+        terminated = truncated = self.episode_len >= 10
+        # r = -abs(obs - action)
+        reward = -sum(abs(self.cur_obs - action))
+        # Set a new observation (random sample).
+        self.cur_obs = self.observation_space.sample()
+        return self.cur_obs, reward, terminated, truncated, {}
+
+
+# Create an RLlib Algorithm instance from a PPOConfig to learn how to
+# act in the above environment.
+config = (
+    PPOConfig().environment(
+        # Env class to use (your gym.Env subclass from above).
+        env=ParrotEnv,
+        # Config dict to be passed to your custom env's constructor.
+        env_config={"parrot_shriek_range": gym.spaces.Box(-5.0, 5.0, (1,))},
+    )
+    # Parallelize environment rollouts.
+    .env_runners(num_env_runners=3)
+)
+algo = config.build()
+
+# Train for n iterations and report results (mean episode rewards).
+# Since we have to guess 10 times and the optimal reward is 0.0
+# (exact match between observation and action value),
+# we can expect to reach an optimal episode reward of 0.0.
+for i in range(5):
+    results = algo.train()
+    print(f"Iter: {i}; avg. reward={results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}")
+
+# Perform inference (action computations) based on given env observations.
+# Note that we are using a slightly simpler env here (-3.0 to 3.0, instead
+# of -5.0 to 5.0!), however, this should still work as the agent has
+# (hopefully) learned to "just always repeat the observation!".
+env = ParrotEnv({"parrot_shriek_range": gym.spaces.Box(-3.0, 3.0, (1,))})
+# Get the initial observation (some value between -10.0 and 10.0).
+obs, info = env.reset()
+done = False
+total_reward = 0.0
+# Play one episode.
+while not done:
+    # Compute a single action, given the current observation
+    # from the environment.
+    model_outputs = algo.env_runner.module.forward_inference(
+        {"obs": torch.from_numpy(obs)}
+    )
+    action = model_outputs["action_dist_inputs"][0].numpy()
+    # Apply the computed action in the environment.
+    obs, reward, done, truncated, info = env.step(action)
+    # Sum up rewards for reporting purposes.
+    total_reward += reward
+# Report results.
+print(f"Played 1 episode; total-reward={total_reward}")

rllib/examples/_old_api_stack/algorithms/atari-dist-dqn.yaml

@@ -0,0 +1,35 @@
+# @OldAPIStack
+atari-dist-dqn:
+    env:
+        grid_search:
+            - ale_py:ALE/Breakout-v5
+            - ale_py:ALE/BeamRider-v5
+            - ale_py:ALE/Qbert-v5
+            - ale_py:ALE/SpaceInvaders-v5
+    run: DQN
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        double_q: false
+        dueling: false
+        num_atoms: 51
+        noisy: false
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+          capacity: 1000000
+        num_steps_sampled_before_learning_starts: 20000
+        n_step: 1
+        target_network_update_freq: 8000
+        lr: .0000625
+        adam_epsilon: .00015
+        hiddens: [512]
+        rollout_fragment_length: 4
+        train_batch_size: 32
+        exploration_config:
+          epsilon_timesteps: 200000
+          final_epsilon: 0.01
+        num_gpus: 0.2
+        min_sample_timesteps_per_iteration: 10000

rllib/examples/_old_api_stack/algorithms/atari-dqn.yaml

@@ -0,0 +1,39 @@
+# @OldAPIStack
+# Runs on a single g3.4xl node
+# See https://github.com/ray-project/rl-experiments for results
+atari-basic-dqn:
+    env:
+        grid_search:
+            - ale_py:ALE/Breakout-v5
+            - ale_py:ALE/BeamRider-v5
+            - ale_py:ALE/Qbert-v5
+            - ale_py:ALE/SpaceInvaders-v5
+    run: DQN
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        double_q: false
+        dueling: false
+        num_atoms: 1
+        noisy: false
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+          capacity: 1000000
+        num_steps_sampled_before_learning_starts: 20000
+        n_step: 1
+        target_network_update_freq: 8000
+        lr: .0000625
+        adam_epsilon: .00015
+        hiddens: [512]
+        rollout_fragment_length: 4
+        train_batch_size: 32
+        exploration_config:
+          epsilon_timesteps: 200000
+          final_epsilon: 0.01
+        num_gpus: 0.2
+        min_sample_timesteps_per_iteration: 10000

rllib/examples/_old_api_stack/algorithms/atari-duel-ddqn.yaml

@@ -0,0 +1,39 @@
+# @OldAPIStack
+# Runs on a single g3.4xl node
+# See https://github.com/ray-project/rl-experiments for results
+dueling-ddqn:
+    env:
+        grid_search:
+            - ale_py:ALE/Breakout-v5
+            - ale_py:ALE/BeamRider-v5
+            - ale_py:ALE/Qbert-v5
+            - ale_py:ALE/SpaceInvaders-v5
+    run: DQN
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        double_q: true
+        dueling: true
+        num_atoms: 1
+        noisy: false
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+          capacity: 1000000
+        num_steps_sampled_before_learning_starts: 20000
+        n_step: 1
+        target_network_update_freq: 8000
+        lr: .0000625
+        adam_epsilon: .00015
+        hiddens: [512]
+        rollout_fragment_length: 4
+        train_batch_size: 32
+        exploration_config:
+          epsilon_timesteps: 200000
+          final_epsilon: 0.01
+        num_gpus: 0.2
+        min_sample_timesteps_per_iteration: 10000

rllib/examples/_old_api_stack/algorithms/atari-impala-large.yaml

@@ -0,0 +1,28 @@
+# @OldAPIStack
+# Runs on a g3.16xl node with 5 m5.24xl workers
+# Takes roughly 10 minutes.
+atari-impala:
+    env:
+        grid_search:
+            - ale_py:ALE/Breakout-v5
+            - ale_py:ALE/BeamRider-v5
+            - ale_py:ALE/Qbert-v5
+            - ale_py:ALE/SpaceInvaders-v5
+    run: IMPALA
+    stop:
+        timesteps_total: 3000000
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 500
+        num_env_runners: 128
+        num_envs_per_env_runner: 5
+        clip_rewards: True
+        lr_schedule: [
+            [0, 0.0005],
+            [20000000, 0.000000000001],
+        ]

rllib/examples/_old_api_stack/algorithms/atari-impala-multi-gpu.yaml

@@ -0,0 +1,25 @@
+# @OldAPIStack
+# Runs on a p2.8xlarge single head node machine.
+# Should reach ~400 reward in about 1h and after 15-20M ts.
+atari-impala:
+    env: ale_py:ALE/Breakout-v5
+    run: IMPALA
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 4000
+        num_gpus: 4
+        num_env_runners: 31
+        num_gpus_per_env_runner: 0  # works also for partial GPUs (<1.0) per worker
+        num_envs_per_env_runner: 5
+        clip_rewards: True
+        lr_schedule: [
+            [0, 0.0005],
+            [20000000, 0.000000000001],
+        ]

rllib/examples/_old_api_stack/algorithms/atari-impala.yaml

@@ -0,0 +1,26 @@
+# @OldAPIStack
+# Runs on a g3.16xl node with 3 m4.16xl workers
+# See https://github.com/ray-project/rl-experiments for results
+atari-impala:
+    env:
+        grid_search:
+            - ale_py:ALE/Breakout-v5
+            - ale_py:ALE/BeamRider-v5
+            - ale_py:ALE/Qbert-v5
+            - ale_py:ALE/SpaceInvaders-v5
+    run: IMPALA
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 500
+        num_env_runners: 32
+        num_envs_per_env_runner: 5
+        clip_rewards: True
+        lr_schedule: [
+            [0, 0.0005],
+            [20000000, 0.000000000001],
+        ]

rllib/examples/_old_api_stack/algorithms/atari-sac.yaml

@@ -0,0 +1,55 @@
+# @OldAPIStack
+# Run e.g. on a g3.16xlarge (4 GPUs) with `num_gpus=1` (1 for each trial;
+# MsPacman torch + tf; Pong torch + tf).
+# Uses the hyperparameters published in [2] (see rllib/algorithms/sac/README.md).
+atari-sac-tf-and-torch:
+    env:
+        grid_search:
+            - ale_py:ALE/MsPacman-v5
+            - ale_py:ALE/Pong-v5
+    run: SAC
+    stop:
+        timesteps_total: 20000000
+    config:
+        # Works for both torch and tf.
+        framework:
+            grid_search: [tf, torch]
+        env_config:
+            frameskip: 1  # no frameskip
+        gamma: 0.99
+        q_model_config:
+            hidden_activation: relu
+            hidden_layer_sizes: [512]
+        policy_model_config:
+            hidden_activation: relu
+            hidden_layer_sizes: [512]
+        # Do hard syncs.
+        # Soft-syncs seem to work less reliably for discrete action spaces.
+        tau: 1.0
+        target_network_update_freq: 8000
+        # auto = 0.98 * -log(1/|A|)
+        target_entropy: auto
+        clip_rewards: 1.0
+        n_step: 1
+        rollout_fragment_length: 1
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+          capacity: 1000000
+          # How many steps of the model to sample before learning starts.
+          # If True prioritized replay buffer will be used.
+          prioritized_replay_alpha: 0.6
+          prioritized_replay_beta: 0.4
+          prioritized_replay_eps: 1e-6
+        num_steps_sampled_before_learning_starts: 100000
+        train_batch_size: 64
+        min_sample_timesteps_per_iteration: 4
+        # Paper uses 20k random timesteps, which is not exactly the same, but
+        # seems to work nevertheless. We use 100k here for the longer Atari
+        # runs (DQN style: filling up the buffer a bit before learning).
+        optimization:
+            actor_learning_rate: 0.0003
+            critic_learning_rate: 0.0003
+            entropy_learning_rate: 0.0003
+        num_env_runners: 0
+        num_gpus: 1
+        metrics_num_episodes_for_smoothing: 5

rllib/examples/_old_api_stack/algorithms/cartpole-appo-separate-losses.py

@@ -0,0 +1,44 @@
+# @OldAPIStack
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 400,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000,
+}
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment("CartPole-v1")
+    # Switch on >1 loss/optimizer API for TFPolicy and EagerTFPolicy.
+    .experimental(_tf_policy_handles_more_than_one_loss=True)
+    .training(
+        # APPO will produce two separate loss terms: policy loss + value function loss.
+        _separate_vf_optimizer=True,
+        # Separate learning rate (and schedule) for the value function branch.
+        _lr_vf=tune.grid_search([0.00075, [[0, 0.00075], [100000, 0.0003]]]),
+        num_epochs=6,
+        # `vf_loss_coeff` will be ignored anyways as we use separate loss terms.
+        vf_loss_coeff=0.01,
+        vtrace=True,
+        model={
+            # Make sure we really have completely separate branches.
+            "vf_share_layers": False,
+        },
+    )
+    .env_runners(
+        num_envs_per_env_runner=5,
+        num_env_runners=1,
+        observation_filter="MeanStdFilter",
+    )
+    .resources(num_gpus=0)
+)

rllib/examples/_old_api_stack/algorithms/cartpole-bc.yaml

@@ -0,0 +1,26 @@
+# @OldAPIStack
+# To generate training data, first run:
+# $ ./train.py --run=PPO --env=CartPole-v1 \
+#      --stop='{"timesteps_total": 50000}' \
+#      --config='{"output": "dataset", "output_config": {"format": "json", "path": "/tmp/out", "max_num_samples_per_file": 1}, "batch_mode": "complete_episodes"}'
+cartpole-bc:
+    env: CartPole-v1
+    run: BC
+    stop:
+        timesteps_total: 500000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        enable_rl_module_and_learner: false
+        enable_env_runner_and_connector_v2: false
+        # In order to evaluate on an actual environment, use these following
+        # settings:
+        evaluation_num_env_runners: 1
+        evaluation_interval: 1
+        evaluation_config:
+            input: sampler
+        # The historic (offline) data file from the PPO run (at the top).
+        input: dataset
+        input_config:
+            format: json
+            paths: /tmp/out

rllib/examples/_old_api_stack/algorithms/cartpole-crashing-and-stalling-recreate-workers-appo.py

@@ -0,0 +1,82 @@
+# @OldAPIStack
+"""
+Tests, whether APPO can learn in a fault-tolerant fashion.
+
+Workers will be configured to automatically get recreated upon failures (here: within
+the environment).
+The environment we use here is configured to crash with a certain probability on each
+`step()` and/or `reset()` call. Additionally, the environment is configured to stall
+with a configured probability on each `step()` call for a certain amount of time.
+"""
+from gymnasium.wrappers import TimeLimit
+
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.cartpole_crashing import CartPoleCrashing
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.register_env(
+    "env",
+    lambda cfg: TimeLimit(CartPoleCrashing(cfg), max_episode_steps=500),
+)
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment(
+        "env",
+        env_config={
+            "p_crash": 0.0001,  # prob to crash during step()
+            "p_crash_reset": 0.001,  # prob to crash during reset()
+            "crash_on_worker_indices": [1, 2],
+            "init_time_s": 2.0,
+            "p_stall": 0.0005,  # prob to stall during step()
+            "p_stall_reset": 0.001,  # prob to stall during reset()
+            "stall_time_sec": (2, 5),  # stall between 2 and 10sec.
+            "stall_on_worker_indices": [2, 3],
+        },
+    )
+    .env_runners(
+        num_env_runners=1,
+        num_envs_per_env_runner=1,
+    )
+    # Switch on resiliency (recreate any failed worker).
+    .fault_tolerance(
+        restart_failed_env_runners=True,
+    )
+    .evaluation(
+        evaluation_num_env_runners=4,
+        evaluation_interval=1,
+        evaluation_duration=25,
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=True,
+        evaluation_config=APPOConfig.overrides(
+            explore=False,
+            env_config={
+                # Make eval workers solid.
+                # This test is to prove that we can learn with crashing envs,
+                # not evaluate with crashing envs.
+                "p_crash": 0.0,
+                "p_crash_reset": 0.0,
+                "init_time_s": 0.0,
+                "p_stall": 0.01,
+                "stall_time_sec": 300,  # stall for 5min.
+                "p_stall_reset": 0.0,
+                "stall_on_worker_indices": [1, 2],
+            },
+        ),
+    )
+)
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 500.0,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 2000000,
+}

rllib/examples/_old_api_stack/algorithms/cartpole-crashing-recreate-workers-appo.py

@@ -0,0 +1,69 @@
+# @OldAPIStack
+"""
+Tests, whether APPO can learn in a fault-tolerant fashion.
+
+Workers will be configured to automatically get recreated upon failures (here: within
+the environment).
+The environment we use here is configured to crash with a certain probability on each
+`step()` and/or `reset()` call.
+"""
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.cartpole_crashing import CartPoleCrashing
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.register_env("env", lambda cfg: CartPoleCrashing(cfg))
+
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 400.0,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 250000,
+}
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment(
+        "env",
+        env_config={
+            # Crash roughly every 500 ts.
+            "p_crash": 0.0005,  # prob to crash during step()
+            "p_crash_reset": 0.005,  # prob to crash during reset()
+            "crash_on_worker_indices": [1, 2],
+        },
+    )
+    .env_runners(
+        num_env_runners=3,
+        num_envs_per_env_runner=1,
+    )
+    # Switch on resiliency (recreate any failed worker).
+    .fault_tolerance(
+        restart_failed_env_runners=True,
+    )
+    .evaluation(
+        evaluation_num_env_runners=1,
+        evaluation_interval=1,
+        evaluation_duration=25,
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=True,
+        evaluation_config=APPOConfig.overrides(
+            explore=False,
+            env_config={
+                # Make eval workers solid.
+                # This test is to prove that we can learn with crashing envs,
+                # not evaluate with crashing envs.
+                "p_crash": 0.0,
+                "p_crash_reset": 0.0,
+                "init_time_s": 0.0,
+            },
+        ),
+    )
+)

rllib/examples/_old_api_stack/algorithms/cartpole-dqn-fake-gpus.yaml

@@ -0,0 +1,20 @@
+# @OldAPIStack
+cartpole-dqn-fake-gpus:
+    env: CartPole-v1
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 150
+        training_iteration: 400
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        model:
+            fcnet_hiddens: [64]
+            fcnet_activation: linear
+        n_step: 3
+        # Double batch size (2 GPUs).
+        train_batch_size: 64
+
+        # Fake 2 GPUs.
+        num_gpus: 2
+        _fake_gpus: true

rllib/examples/_old_api_stack/algorithms/cartpole-dqn-param-noise.yaml

@@ -0,0 +1,20 @@
+# @OldAPIStack
+cartpole-dqn-w-param-noise:
+    env: CartPole-v1
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 150
+        timesteps_total: 300000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        exploration_config:
+            type: ParameterNoise
+            random_timesteps: 10000
+            initial_stddev: 1.0
+        batch_mode: complete_episodes
+        lr: 0.0008
+        num_env_runners: 0
+        model:
+            fcnet_hiddens: [32, 32]
+            fcnet_activation: tanh

rllib/examples/_old_api_stack/algorithms/cartpole-dqn-softq.yaml

@@ -0,0 +1,17 @@
+# @OldAPIStack
+cartpole-dqn:
+    env: CartPole-v1
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 150
+        timesteps_total: 100000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        model:
+            fcnet_hiddens: [64]
+            fcnet_activation: linear
+        n_step: 3
+        exploration_config:
+            type: SoftQ
+            temperature: 0.5

rllib/examples/_old_api_stack/algorithms/cartpole-dqn.yaml

@@ -0,0 +1,14 @@
+# @OldAPIStack
+cartpole-dqn:
+    env: CartPole-v1
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 100
+        timesteps_total: 100000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        model:
+            fcnet_hiddens: [64]
+            fcnet_activation: linear
+        n_step: 3

rllib/examples/_old_api_stack/algorithms/cartpole-impala-separate-losses.py

@@ -0,0 +1,44 @@
+# @OldAPIStack
+from ray.rllib.algorithms.impala import IMPALAConfig
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 150,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000,
+}
+
+config = (
+    IMPALAConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment("CartPole-v1")
+    # Switch on >1 loss/optimizer API for TFPolicy and EagerTFPolicy.
+    .experimental(_tf_policy_handles_more_than_one_loss=True)
+    .training(
+        # IMPALA will produce two separate loss terms: policy loss + value function
+        # loss.
+        _separate_vf_optimizer=True,
+        # Separate learning rate for the value function branch.
+        _lr_vf=0.00075,
+        num_epochs=6,
+        # `vf_loss_coeff` will be ignored anyways as we use separate loss terms.
+        vf_loss_coeff=0.01,
+        vtrace=True,
+        model={
+            # Make sure we really have completely separate branches.
+            "vf_share_layers": False,
+        },
+    )
+    .env_runners(
+        num_envs_per_env_runner=5,
+        num_env_runners=1,
+        observation_filter="MeanStdFilter",
+    )
+    .resources(num_gpus=0)
+)

rllib/examples/_old_api_stack/algorithms/cartpole-marwil.yaml

@@ -0,0 +1,22 @@
+# @OldAPIStack
+# To generate training data, first run:
+# $ ./train.py --run=PPO --env=CartPole-v1 \
+#      --stop='{"timesteps_total": 50000}' \
+#      --config='{"output": "/tmp/out", "batch_mode": "complete_episodes"}'
+cartpole-marwil:
+    env: CartPole-v1
+    run: MARWIL
+    stop:
+        timesteps_total: 500000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # In order to evaluate on an actual environment, use these following
+        # settings:
+        evaluation_num_env_runners: 1
+        evaluation_interval: 1
+        evaluation_config:
+            input: sampler
+        beta: 1.0  # Compare to behavior cloning (beta=0.0).
+        # The historic (offline) data file from the PPO run (at the top).
+        input: /tmp/out

rllib/examples/_old_api_stack/algorithms/cartpole-sac.yaml

@@ -0,0 +1,22 @@
+# @OldAPIStack
+cartpole-sac:
+    env: CartPole-v1
+    run: SAC
+    stop:
+        env_runners/episode_return_mean: 150.0
+        timesteps_total: 100000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        gamma: 0.95
+        target_network_update_freq: 32
+        tau: 1.0
+        # initial_alpha: 0.5
+        train_batch_size: 32
+        optimization:
+            actor_learning_rate: 0.005
+            critic_learning_rate: 0.005
+            entropy_learning_rate: 0.0001
+        # grad_norm_clipping: 40.0
+        # evaluation_config:
+        #     explore: true

rllib/examples/_old_api_stack/algorithms/frozenlake-appo-vtrace.yaml

@@ -0,0 +1,33 @@
+# @OldAPIStack
+frozenlake-appo-vtrace:
+    env: FrozenLake-v1
+    run: APPO
+    stop:
+        env_runners/episode_return_mean: 0.99
+        timesteps_total: 1000000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+
+        # Sparse reward environment (short horizon).
+        env_config:
+            desc:
+                - SFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFF
+                - FFFFFFFG
+            is_slippery: false
+        horizon: 20
+        rollout_fragment_length: 10
+        batch_mode: complete_episodes
+        vtrace: true
+
+        num_envs_per_env_runner: 5
+        num_env_runners: 4
+        num_gpus: 0
+        num_epochs: 1
+        vf_loss_coeff: 0.01

rllib/examples/_old_api_stack/algorithms/halfcheetah-bc.yaml

@@ -0,0 +1,50 @@
+# @OldAPIStack
+halfcheetah_bc:
+    env:
+        grid_search:
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_random
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_medium
+            - ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_expert
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_medium_replay
+    run: CQL
+    config:
+        # SAC Configs
+        #input: d4rl.halfcheetah-random-v0
+        #input: d4rl.halfcheetah-medium-v0
+        input: d4rl.halfcheetah-expert-v0
+        #input: d4rl.halfcheetah-medium-replay-v0
+        framework: torch
+        q_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        policy_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 1
+        rollout_fragment_length: 1
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+        num_steps_sampled_before_learning_starts: 10
+        train_batch_size: 256
+        target_network_update_freq: 0
+        min_train_timesteps_per_iteration: 1000
+        optimization:
+          actor_learning_rate: 0.0001
+          critic_learning_rate: 0.0003
+          entropy_learning_rate: 0.0001
+        num_env_runners: 0
+        num_gpus: 1
+        clip_actions: false
+        normalize_actions: true
+        evaluation_interval: 1
+        metrics_num_episodes_for_smoothing: 5
+        # CQL Configs
+        min_q_weight: 5.0
+        bc_iters: 200000000
+        temperature: 1.0
+        num_actions: 10
+        lagrangian: False
+        evaluation_config:
+            input: sampler

rllib/examples/_old_api_stack/algorithms/halfcheetah-cql.yaml

@@ -0,0 +1,51 @@
+# @OldAPIStack
+halfcheetah_cql:
+    env:
+        grid_search:
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_random
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_medium
+            - ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_expert
+            #- ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_medium_replay
+    run: CQL
+    config:
+        # SAC Configs
+        #input: d4rl.halfcheetah-random-v0
+        #input: d4rl.halfcheetah-medium-v0
+        input: d4rl.halfcheetah-expert-v0
+        #input: d4rl.halfcheetah-medium-replay-v0
+
+        # Works for both torch and tf.
+        framework: torch
+        q_model_config:
+            fcnet_activation: relu
+            fcnet_hiddens: [256, 256, 256]
+        policy_model_config:
+            fcnet_activation: relu
+            fcnet_hiddens: [256, 256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 3
+        rollout_fragment_length: 1
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+        num_steps_sampled_before_learning_starts: 256
+        train_batch_size: 256
+        target_network_update_freq: 0
+        min_train_timesteps_per_iteration: 1000
+        optimization:
+            actor_learning_rate: 0.0001
+            critic_learning_rate: 0.0003
+            entropy_learning_rate: 0.0001
+        num_env_runners: 0
+        num_gpus: 1
+        metrics_num_episodes_for_smoothing: 5
+        # CQL Configs
+        min_q_weight: 5.0
+        bc_iters: 20000
+        temperature: 1.0
+        num_actions: 10
+        lagrangian: False
+
+        evaluation_interval: 3
+        evaluation_config:
+            input: sampler

rllib/examples/_old_api_stack/algorithms/halfcheetah-ppo.yaml

@@ -0,0 +1,26 @@
+# @OldAPIStack
+halfcheetah-ppo:
+    env: HalfCheetah-v2
+    run: PPO
+    stop:
+        env_runners/episode_return_mean: 9800
+        time_total_s: 10800
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        gamma: 0.99
+        lambda: 0.95
+        kl_coeff: 1.0
+        num_epochs: 32
+        lr: .0003
+        vf_loss_coeff: 0.5
+        clip_param: 0.2
+        minibatch_size: 4096
+        train_batch_size: 65536
+        num_env_runners: 16
+        num_gpus: 1
+        grad_clip: 0.5
+        num_envs_per_env_runner:
+            grid_search: [16, 32]
+        batch_mode: truncate_episodes
+        observation_filter: MeanStdFilter

rllib/examples/_old_api_stack/algorithms/hopper-bc.yaml

@@ -0,0 +1,50 @@
+# @OldAPIStack
+hopper_bc:
+    env:
+        grid_search:
+            - ray.rllib.examples.envs.classes.d4rl_env.hopper_random
+            #- ray.rllib.examples.envs.classes..d4rl_env.hopper_medium
+            #- ray.rllib.examples.envs.classes..d4rl_env.hopper_expert
+            #- ray.rllib.examples.envs.classes..d4rl_env.hopper_medium_replay
+    run: CQL
+    config:
+        # SAC Configs
+        input: d4rl.hopper-random-v0
+        #input: d4rl.hopper-medium-v0
+        #input: d4rl.hopper-expert-v0
+        #input: d4rl.hopper-medium-replay-v0
+        framework: torch
+        q_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        policy_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 1
+        rollout_fragment_length: 1
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+        num_steps_sampled_before_learning_starts: 10
+        train_batch_size: 256
+        target_network_update_freq: 0
+        min_train_timesteps_per_iteration: 1000
+        optimization:
+          actor_learning_rate: 0.0001
+          critic_learning_rate: 0.0003
+          entropy_learning_rate: 0.0001
+        num_env_runners: 0
+        num_gpus: 1
+        clip_actions: false
+        normalize_actions: true
+        evaluation_interval: 1
+        metrics_num_episodes_for_smoothing: 5
+        # CQL Configs
+        min_q_weight: 5.0
+        bc_iters: 200000000
+        temperature: 1.0
+        num_actions: 10
+        lagrangian: False
+        evaluation_config:
+            input: sampler

rllib/examples/_old_api_stack/algorithms/hopper-cql.yaml

@@ -0,0 +1,50 @@
+# @OldAPIStack
+hopper_cql:
+    env:
+        grid_search:
+            #- ray.rllib.examples.envs.classes.d4rl_env.hopper_random
+            - ray.rllib.examples.envs.classes.d4rl_env.hopper_medium
+            #- ray.rllib.examples.envs.classes.d4rl_env.hopper_expert
+            #- ray.rllib.examples.envs.classes.d4rl_env.hopper_medium_replay
+    run: CQL
+    config:
+        # SAC Configs
+        #input: d4rl.hopper-random-v0
+        input: d4rl.hopper-medium-v0
+        #input: d4rl.hopper-expert-v0
+        #input: d4rl.hopper-medium-replay-v0
+        framework: torch
+        q_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        policy_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 1
+        rollout_fragment_length: 1
+        replay_buffer_config:
+          type: MultiAgentReplayBuffer
+        num_steps_sampled_before_learning_starts: 10
+        train_batch_size: 256
+        target_network_update_freq: 0
+        min_train_timesteps_per_iteration: 1000
+        optimization:
+          actor_learning_rate: 0.0001
+          critic_learning_rate: 0.0003
+          entropy_learning_rate: 0.0001
+        num_env_runners: 0
+        num_gpus: 1
+        clip_actions: false
+        normalize_actions: true
+        evaluation_interval: 1
+        metrics_num_episodes_for_smoothing: 5
+        # CQL Configs
+        min_q_weight: 5.0
+        bc_iters: 20000
+        temperature: 1.0
+        num_actions: 10
+        lagrangian: False
+        evaluation_config:
+            input: sampler

rllib/examples/_old_api_stack/algorithms/hopper-ppo.yaml

@@ -0,0 +1,17 @@
+# @OldAPIStack
+hopper-ppo:
+    env: Hopper-v1
+    run: PPO
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        gamma: 0.995
+        kl_coeff: 1.0
+        num_epochs: 20
+        lr: .0001
+        minibatch_size: 32768
+        train_batch_size: 160000
+        num_env_runners: 64
+        num_gpus: 4
+        batch_mode: complete_episodes
+        observation_filter: MeanStdFilter

rllib/examples/_old_api_stack/algorithms/humanoid-ppo-gae.yaml

@@ -0,0 +1,24 @@
+# @OldAPIStack
+humanoid-ppo-gae:
+    env: Humanoid-v1
+    run: PPO
+    stop:
+        env_runners/episode_return_mean: 6000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        gamma: 0.995
+        lambda: 0.95
+        clip_param: 0.2
+        kl_coeff: 1.0
+        num_epochs: 20
+        lr: .0001
+        minibatch_size: 32768
+        horizon: 5000
+        train_batch_size: 320000
+        model:
+            free_log_std: true
+        num_env_runners: 64
+        num_gpus: 4
+        batch_mode: complete_episodes
+        observation_filter: MeanStdFilter

rllib/examples/_old_api_stack/algorithms/humanoid-ppo.yaml

@@ -0,0 +1,22 @@
+# @OldAPIStack
+humanoid-ppo:
+    env: Humanoid-v1
+    run: PPO
+    stop:
+        env_runners/episode_return_mean: 6000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        gamma: 0.995
+        kl_coeff: 1.0
+        num_epochs: 20
+        lr: .0001
+        minibatch_size: 32768
+        train_batch_size: 320000
+        model:
+            free_log_std: true
+        use_gae: false
+        num_env_runners: 64
+        num_gpus: 4
+        batch_mode: complete_episodes
+        observation_filter: MeanStdFilter

rllib/examples/_old_api_stack/algorithms/memory-leak-test-appo.yaml

@@ -0,0 +1,15 @@
+# @OldAPIStack
+memory-leak-test-appo:
+    env:
+        ray.rllib.examples.envs.classes.random_env.RandomLargeObsSpaceEnv
+    run: APPO
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Switch off np.random, which is known to have memory leaks.
+        env_config:
+            config:
+                static_samples: true
+        num_env_runners: 4
+        num_envs_per_env_runner: 5
+        rollout_fragment_length: 20

rllib/examples/_old_api_stack/algorithms/memory-leak-test-dqn.yaml

@@ -0,0 +1,14 @@
+# @OldAPIStack
+memory-leak-test-dqn:
+    env:
+        ray.rllib.examples.envs.classes.random_env.RandomLargeObsSpaceEnv
+    run: DQN
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Switch off np.random, which is known to have memory leaks.
+        env_config:
+            config:
+                static_samples: true
+        replay_buffer_config:
+          capacity: 500  # use small buffer to catch memory leaks

rllib/examples/_old_api_stack/algorithms/memory-leak-test-ppo.yaml

@@ -0,0 +1,17 @@
+# @OldAPIStack
+memory-leak-test-ppo:
+    env:
+        ray.rllib.examples.envs.classes.random_env.RandomLargeObsSpaceEnv
+    run: PPO
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Switch off np.random, which is known to have memory leaks.
+        env_config:
+            config:
+                static_samples: true
+        num_env_runners: 4
+        num_envs_per_env_runner: 5
+        train_batch_size: 500
+        minibatch_size: 256
+        num_epochs: 5

rllib/examples/_old_api_stack/algorithms/memory-leak-test-sac.yaml

@@ -0,0 +1,14 @@
+# @OldAPIStack
+memory-leak-test-sac:
+    env:
+        ray.rllib.examples.envs.classes.random_env.RandomLargeObsSpaceEnvContActions
+    run: SAC
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Switch off np.random, which is known to have memory leaks.
+        env_config:
+            config:
+                static_samples: true
+        replay_buffer_config:
+          capacity: 500  # use small buffer to catch memory leaks

rllib/examples/_old_api_stack/algorithms/mspacman-sac.yaml

@@ -0,0 +1,45 @@
+# @OldAPIStack
+# Our implementation of SAC discrete can reach up
+# to ~750 reward in 40k timesteps. Run e.g. on a g3.4xlarge with `num_gpus=1`.
+# Uses the hyperparameters published in [2] (see rllib/algorithms/sac/README.md).
+mspacman-sac-tf:
+    env: ale_py:ALE/MsPacman-v5
+    run: SAC
+    stop:
+        env_runners/episode_return_mean: 800
+        timesteps_total: 100000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        env_config:
+            frameskip: 1  # no frameskip
+        gamma: 0.99
+        q_model_config:
+            fcnet_hiddens: [512]
+            fcnet_activation: relu
+        policy_model_config:
+            fcnet_hiddens: [512]
+            fcnet_activation: relu
+        # Do hard syncs.
+        # Soft-syncs seem to work less reliably for discrete action spaces.
+        tau: 1.0
+        target_network_update_freq: 8000
+        # paper uses: 0.98 * -log(1/|A|)
+        target_entropy: 1.755
+        clip_rewards: 1.0
+        n_step: 1
+        rollout_fragment_length: 1
+        train_batch_size: 64
+        min_sample_timesteps_per_iteration: 4
+        # Paper uses 20k random timesteps, which is not exactly the same, but
+        # seems to work nevertheless.
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+        num_steps_sampled_before_learning_starts: 20000
+        optimization:
+            actor_learning_rate: 0.0003
+            critic_learning_rate: 0.0003
+            entropy_learning_rate: 0.0003
+        num_env_runners: 0
+        num_gpus: 0
+        metrics_num_episodes_for_smoothing: 5

rllib/examples/_old_api_stack/algorithms/multi-agent-cartpole-crashing-and-stalling-recreate-workers-appo.py

@@ -0,0 +1,77 @@
+# @OldAPIStack
+"""
+Tests, whether APPO can learn in a fault-tolerant fashion in a
+multi-agent setting.
+
+Workers will be configured to automatically get recreated upon failures (here: within
+the environment).
+The environment we use here is configured to crash with a certain probability on each
+`step()` and/or `reset()` call.
+"""
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.cartpole_crashing import MultiAgentCartPoleCrashing
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.register_env("ma_env", lambda cfg: MultiAgentCartPoleCrashing(cfg))
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 800.0,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 250000,
+}
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment(
+        "ma_env",
+        env_config={
+            "num_agents": 2,
+            # Crash roughly every 300 ts. This should be ok to measure 180.0
+            # reward (episodes are 200 ts long).
+            "p_crash": 0.00005,  # prob to crash during step()
+            "p_crash_reset": 0.0005,  # prob to crash during reset()
+            "init_time_s": 2.0,
+            "p_stall": 0.001,  # prob to stall during step()
+            "p_stall_reset": 0.001,  # prob to stall during reset()
+            "stall_time_sec": (2, 5),  # stall between 2 and 10sec.
+            "stall_on_worker_indices": [2, 3],
+        },
+    )
+    .env_runners(
+        num_env_runners=3,
+        num_envs_per_env_runner=1,
+    )
+    # Switch on resiliency (recreate any failed worker).
+    .fault_tolerance(
+        restart_failed_env_runners=True,
+    )
+    .evaluation(
+        evaluation_num_env_runners=1,
+        evaluation_interval=1,
+        evaluation_duration=25,
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=True,
+        evaluation_config=APPOConfig.overrides(
+            explore=False,
+            env_config={
+                # Make eval workers solid.
+                # This test is to prove that we can learn with crashing envs,
+                # not evaluate with crashing envs.
+                "p_crash": 0.0,
+                "p_crash_reset": 0.0,
+                "init_time_s": 0.0,
+                "p_stall": 0.0,
+                "p_stall_reset": 0.0,
+            },
+        ),
+    )
+)

rllib/examples/_old_api_stack/algorithms/multi-agent-cartpole-crashing-recreate-workers-appo.py

@@ -0,0 +1,70 @@
+# @OldAPIStack
+"""
+Tests, whether APPO can learn in a fault-tolerant fashion in a
+multi-agent setting.
+
+Workers will be configured to automatically get recreated upon failures (here: within
+the environment).
+The environment we use here is configured to crash with a certain probability on each
+`step()` and/or `reset()` call.
+"""
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.cartpole_crashing import MultiAgentCartPoleCrashing
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.register_env("ma_env", lambda cfg: MultiAgentCartPoleCrashing(cfg))
+
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 800.0,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 250000,
+}
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment(
+        "ma_env",
+        env_config={
+            "num_agents": 2,
+            # Crash roughly every 300 ts. This should be ok to measure 180.0
+            # reward (episodes are 200 ts long).
+            "p_crash": 0.0005,  # prob to crash during step()
+            "p_crash_reset": 0.005,  # prob to crash during reset()
+        },
+    )
+    .env_runners(
+        num_env_runners=4,
+        num_envs_per_env_runner=1,
+    )
+    # Switch on resiliency (recreate any failed worker).
+    .fault_tolerance(
+        restart_failed_env_runners=True,
+    )
+    .evaluation(
+        evaluation_num_env_runners=1,
+        evaluation_interval=1,
+        evaluation_duration=25,
+        evaluation_duration_unit="episodes",
+        evaluation_parallel_to_training=True,
+        evaluation_config=APPOConfig.overrides(
+            explore=False,
+            env_config={
+                # Make eval workers solid.
+                # This test is to prove that we can learn with crashing envs,
+                # not evaluate with crashing envs.
+                "p_crash": 0.0,
+                "p_crash_reset": 0.0,
+                "init_time_s": 0.0,
+            },
+        ),
+    )
+)

rllib/examples/_old_api_stack/algorithms/multi-agent-cartpole-w-100-policies-appo.py

@@ -0,0 +1,83 @@
+# @OldAPIStack
+import numpy as np
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EVALUATION_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+
+register_env("multi_cartpole", lambda _: MultiAgentCartPole({"num_agents": 2}))
+
+# Number of policies overall in the PolicyMap.
+num_policies = 20
+# Number of those policies that should be trained. These are a subset of `num_policies`.
+num_trainable = 10
+
+num_envs_per_env_runner = 5
+
+# Define the config as an APPOConfig object.
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment("multi_cartpole")
+    .env_runners(
+        num_env_runners=4,
+        num_envs_per_env_runner=num_envs_per_env_runner,
+        observation_filter="MeanStdFilter",
+    )
+    .training(
+        model={
+            "fcnet_hiddens": [32],
+            "fcnet_activation": "linear",
+            "vf_share_layers": True,
+        },
+        num_epochs=1,
+        vf_loss_coeff=0.005,
+        vtrace=True,
+    )
+    .multi_agent(
+        # 2 agents per sub-env.
+        # This is to avoid excessive swapping during an episode rollout, since
+        # Policies are only re-picked at the beginning of each episode.
+        policy_map_capacity=2 * num_envs_per_env_runner,
+        policy_states_are_swappable=True,
+        policies={f"pol{i}" for i in range(num_policies)},
+        # Train only the first n policies.
+        policies_to_train=[f"pol{i}" for i in range(num_trainable)],
+        # Pick one trainable and one non-trainable policy per episode.
+        policy_mapping_fn=(
+            lambda aid, eps, worker, **kw: "pol"
+            + str(
+                np.random.randint(0, num_trainable)
+                if aid == 0
+                else np.random.randint(num_trainable, num_policies)
+            )
+        ),
+    )
+    # On the eval track, always let policy 0 play so we get its results in each results
+    # dict.
+    .evaluation(
+        evaluation_config=APPOConfig.overrides(
+            policy_mapping_fn=(
+                lambda aid, eps, worker, **kw: "pol"
+                + str(0 if aid == 0 else np.random.randint(num_trainable, num_policies))
+            ),
+        ),
+        evaluation_num_env_runners=2,
+        evaluation_interval=1,
+        evaluation_parallel_to_training=True,
+    )
+)
+
+# Define some stopping criteria.
+stop = {
+    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/policy_reward_mean/pol0": 50.0,
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 500000,
+}

rllib/examples/_old_api_stack/algorithms/multi_agent_cartpole_appo_old_api_stack.py

@@ -0,0 +1,46 @@
+# @OldAPIStack
+from ray import tune
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.registry.register_env("env", lambda cfg: MultiAgentCartPole(config=cfg))
+
+
+config = (
+    APPOConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment("env", env_config={"num_agents": 4})
+    .env_runners(
+        num_envs_per_env_runner=5,
+        num_env_runners=4,
+        observation_filter="MeanStdFilter",
+    )
+    .resources(num_gpus=1, _fake_gpus=True)
+    .multi_agent(
+        policies=["p0", "p1", "p2", "p3"],
+        policy_mapping_fn=(lambda agent_id, episode, worker, **kwargs: f"p{agent_id}"),
+    )
+    .training(
+        num_epochs=1,
+        vf_loss_coeff=0.005,
+        vtrace=True,
+        model={
+            "fcnet_hiddens": [32],
+            "fcnet_activation": "linear",
+            "vf_share_layers": True,
+        },
+    )
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 600,  # 600 / 4 (==num_agents) = 150
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000,
+}

rllib/examples/_old_api_stack/algorithms/multi_agent_cartpole_impala_old_api_stack.py

@@ -0,0 +1,47 @@
+# @OldAPIStack
+from ray import tune
+from ray.rllib.algorithms.impala import IMPALAConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+tune.registry.register_env("env", lambda cfg: MultiAgentCartPole(config=cfg))
+
+
+config = (
+    IMPALAConfig()
+    .api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    .environment("env", env_config={"num_agents": 4})
+    .env_runners(
+        num_envs_per_env_runner=5,
+        num_env_runners=4,
+        observation_filter="MeanStdFilter",
+    )
+    .resources(num_gpus=1, _fake_gpus=True)
+    .multi_agent(
+        policies=["p0", "p1", "p2", "p3"],
+        policy_mapping_fn=(lambda agent_id, episode, worker, **kwargs: f"p{agent_id}"),
+    )
+    .training(
+        num_epochs=1,
+        vf_loss_coeff=0.005,
+        vtrace=True,
+        model={
+            "fcnet_hiddens": [32],
+            "fcnet_activation": "linear",
+            "vf_share_layers": True,
+        },
+        replay_proportion=0.0,
+    )
+)
+
+stop = {
+    f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 600,  # 600 / 4 (==num_agents) = 150
+    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000,
+}

rllib/examples/_old_api_stack/algorithms/pendulum-cql.yaml

@@ -0,0 +1,44 @@
+# @OldAPIStack
+# Given a SAC-generated offline file generated via:
+# rllib train -f examples/algorithms/sac/pendulum-sac.yaml --no-ray-ui
+
+# Pendulum CQL can attain ~ -300 reward in 10k from that file.
+pendulum-cql:
+    env: Pendulum-v1
+    run: CQL
+    stop:
+        evaluation/env_runners/episode_return_mean: -700
+        timesteps_total: 800000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+
+        # Set seed.
+        seed: 0
+
+        # Use one or more offline files or "input: sampler" for online learning.
+        input: 'dataset'
+        input_config:
+            paths: ["offline/tests/data/pendulum/enormous.zip"]
+            format: 'json'
+        # Our input file above comes from an SAC run. Actions in there
+        # are already normalized (produced by SquashedGaussian).
+        actions_in_input_normalized: true
+        clip_actions: true
+
+        twin_q: true
+        train_batch_size: 2000
+        bc_iters: 100
+        num_env_runners: 2
+        min_time_s_per_iteration: 10
+
+        metrics_num_episodes_for_smoothing: 5
+
+        # Evaluate in an actual environment.
+        evaluation_interval: 1
+        evaluation_num_env_runners: 2
+        evaluation_duration: 10
+        evaluation_parallel_to_training: true
+        evaluation_config:
+            input: sampler
+            explore: False

rllib/examples/_old_api_stack/algorithms/pendulum-sac.yaml

@@ -0,0 +1,35 @@
+# @OldAPIStack
+# Pendulum SAC can attain -150+ reward in 6-7k
+# Configurations are the similar to original softlearning/sac codebase
+pendulum-sac:
+    env: Pendulum-v1
+    run: SAC
+    stop:
+        env_runners/episode_return_mean: -250
+        timesteps_total: 10000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        q_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256]
+        policy_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 1
+        rollout_fragment_length: 1
+        train_batch_size: 256
+        target_network_update_freq: 1
+        min_sample_timesteps_per_iteration: 1000
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+        num_steps_sampled_before_learning_starts: 256
+        optimization:
+          actor_learning_rate: 0.0003
+          critic_learning_rate: 0.0003
+          entropy_learning_rate: 0.0003
+        num_env_runners: 0
+        num_gpus: 0
+        metrics_num_episodes_for_smoothing: 5

rllib/examples/_old_api_stack/algorithms/pendulum-transformed-actions-ppo.yaml

@@ -0,0 +1,31 @@
+# @OldAPIStack
+# Can expect improvement to -140 reward in ~300-500k timesteps.
+pendulum-ppo:
+    env: ray.rllib.examples.envs.classes.transformed_action_space_env.TransformedActionPendulum
+    run: PPO
+    stop:
+        env_runners/episode_return_mean: -500
+        timesteps_total: 400000
+    config:
+        # Works for both torch and tf.
+        framework: torch
+
+        # Test, whether PPO is able to learn in "distorted" action spaces.
+        env_config:
+            config:
+                low: 300.0
+                high: 500.0
+
+        normalize_actions: true
+        clip_actions: false
+        vf_clip_param: 10.0
+        num_envs_per_env_runner: 20
+        lambda: 0.1
+        gamma: 0.95
+        lr: 0.0003
+        train_batch_size: 512
+        minibatch_size: 64
+        num_epochs: 6
+        observation_filter: MeanStdFilter
+        model:
+            fcnet_activation: relu

rllib/examples/_old_api_stack/algorithms/pendulum-transformed-actions-sac.yaml

@@ -0,0 +1,44 @@
+# @OldAPIStack
+# TransformedActionPendulum SAC can attain -150+ reward in 6-7k
+# Configurations are the similar to original softlearning/sac codebase
+transformed-actions-pendulum-sac-dummy-torch:
+    env: ray.rllib.examples.envs.classes.transformed_action_space_env.TransformedActionPendulum
+    run: SAC
+    stop:
+        env_runners/episode_return_mean: -200
+        timesteps_total: 10000
+    config:
+        # Works for both torch and tf.
+        seed: 42
+        framework: torch
+
+        # Test, whether SAC is able to learn in "distorted" action spaces.
+        env_config:
+          config:
+            low: 300.0
+            high: 500.0
+
+        horizon: 200
+        q_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256]
+        policy_model_config:
+          fcnet_activation: relu
+          fcnet_hiddens: [256, 256]
+        tau: 0.005
+        target_entropy: auto
+        n_step: 1
+        rollout_fragment_length: 1
+        train_batch_size: 256
+        target_network_update_freq: 1
+        min_sample_timesteps_per_iteration: 1000
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+        num_steps_sampled_before_learning_starts: 256
+        optimization:
+          actor_learning_rate: 0.0003
+          critic_learning_rate: 0.0003
+          entropy_learning_rate: 0.0003
+        num_env_runners: 0
+        num_gpus: 0
+        metrics_num_episodes_for_smoothing: 5

rllib/examples/_old_api_stack/algorithms/pong-dqn.yaml

@@ -0,0 +1,35 @@
+# @OldAPIStack
+# You can expect ~20 reward within 1.1m timesteps / 2.1 hours on a K80 GPU
+pong-deterministic-dqn:
+    env: ale_py:ALE/Pong-v5
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 20
+        time_total_s: 7200
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        num_gpus: 1
+        gamma: 0.99
+        lr: .0001
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+          capacity: 50000
+        num_steps_sampled_before_learning_starts: 10000
+        rollout_fragment_length: 4
+        train_batch_size: 32
+        exploration_config:
+          epsilon_timesteps: 200000
+          final_epsilon: .01
+        model:
+          grayscale: True
+          zero_mean: False
+          dim: 42
+        # we should set compress_observations to True because few machines
+        # would be able to contain the replay buffers in memory otherwise
+        compress_observations: True

rllib/examples/_old_api_stack/algorithms/pong-impala-fast.yaml

@@ -0,0 +1,25 @@
+# @OldAPIStack
+# This can reach 18-19 reward in ~3 minutes on p3.16xl head w/m4.16xl workers
+#   128 workers -> 3 minutes (best case)
+#    64 workers -> 4 minutes
+#    32 workers -> 7 minutes
+# See also: pong-impala.yaml, pong-impala-vectorized.yaml
+pong-impala-fast:
+    env: ale_py:ALE/Pong-v5
+    run: IMPALA
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 1000
+        num_env_runners: 128
+        num_envs_per_env_runner: 5
+        broadcast_interval: 5
+        max_sample_requests_in_flight_per_worker: 1
+        num_multi_gpu_tower_stacks: 4
+        num_gpus: 2
+        model:
+          dim: 42

rllib/examples/_old_api_stack/algorithms/pong-impala-vectorized.yaml

@@ -0,0 +1,17 @@
+# @OldAPIStack
+# This can reach 18-19 reward within 10 minutes on a Tesla M60 GPU (e.g., G3 EC2 node)
+# with 32 workers and 10 envs per worker. This is more efficient than the non-vectorized
+# configuration which requires 128 workers to achieve the same performance.
+pong-impala-vectorized:
+    env: ale_py:ALE/Pong-v5
+    run: IMPALA
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 500
+        num_env_runners: 32
+        num_envs_per_env_runner: 10

rllib/examples/_old_api_stack/algorithms/pong-impala.yaml

@@ -0,0 +1,19 @@
+# @OldAPIStack
+# This can reach 18-19 reward within 10 minutes on a Tesla M60 GPU (e.g., G3 EC2 node):
+#   128 workers -> 8 minutes
+#    32 workers -> 17 minutes
+#    16 workers -> 40 min+
+# See also: pong-impala-fast.yaml, pong-impala-vectorized.yaml
+pong-impala:
+    env: ale_py:ALE/Pong-v5
+    run: IMPALA
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        rollout_fragment_length: 50
+        train_batch_size: 500
+        num_env_runners: 128
+        num_envs_per_env_runner: 1

rllib/examples/_old_api_stack/algorithms/pong-rainbow.yaml

@@ -0,0 +1,37 @@
+# @OldAPIStack
+pong-deterministic-rainbow:
+    env: ale_py:ALE/Pong-v5
+    run: DQN
+    stop:
+        env_runners/episode_return_mean: 20
+    config:
+        # Make analogous to old v4 + NoFrameskip.
+        env_config:
+            frameskip: 1
+            full_action_space: false
+            repeat_action_probability: 0.0
+        num_atoms: 51
+        noisy: True
+        gamma: 0.99
+        lr: .0001
+        hiddens: [512]
+        rollout_fragment_length: 4
+        train_batch_size: 32
+        exploration_config:
+          epsilon_timesteps: 2
+          final_epsilon: 0.0
+        target_network_update_freq: 500
+        replay_buffer_config:
+          type: MultiAgentPrioritizedReplayBuffer
+          prioritized_replay_alpha: 0.5
+          capacity: 50000
+        num_steps_sampled_before_learning_starts: 10000
+        n_step: 3
+        gpu: True
+        model:
+          grayscale: True
+          zero_mean: False
+          dim: 42
+        # we should set compress_observations to True because few machines
+        # would be able to contain the replay buffers in memory otherwise
+        compress_observations: True

rllib/examples/_old_api_stack/algorithms/walker2d-ppo.yaml

@@ -0,0 +1,16 @@
+# @OldAPIStack
+walker2d-v1-ppo:
+    env: Walker2d-v1
+    run: PPO
+    config:
+        # Works for both torch and tf.
+        framework: torch
+        kl_coeff: 1.0
+        num_epochs: 20
+        lr: .0001
+        minibatch_size: 32768
+        train_batch_size: 320000
+        num_env_runners: 64
+        num_gpus: 4
+        batch_mode: complete_episodes
+        observation_filter: MeanStdFilter

rllib/examples/_old_api_stack/attention_net_supervised.py

@@ -0,0 +1,78 @@
+# @OldAPIStack
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+from rllib.models.tf.attention_net import TrXLNet
+
+from ray.rllib.utils.framework import try_import_tf
+
+tf1, tf, tfv = try_import_tf()
+
+
+def bit_shift_generator(seq_length, shift, batch_size):
+    while True:
+        values = np.array([0.0, 1.0], dtype=np.float32)
+        seq = np.random.choice(values, (batch_size, seq_length, 1))
+        targets = np.squeeze(np.roll(seq, shift, axis=1).astype(np.int32))
+        targets[:, :shift] = 0
+        yield seq, targets
+
+
+def train_loss(targets, outputs):
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        labels=targets, logits=outputs
+    )
+    return tf.reduce_mean(loss)
+
+
+def train_bit_shift(seq_length, num_iterations, print_every_n):
+
+    optimizer = tf.keras.optimizers.Adam(1e-3)
+
+    model = TrXLNet(
+        observation_space=Box(low=0, high=1, shape=(1,), dtype=np.int32),
+        action_space=Discrete(2),
+        num_outputs=2,
+        model_config={"max_seq_len": seq_length},
+        name="trxl",
+        num_transformer_units=1,
+        attention_dim=10,
+        num_heads=5,
+        head_dim=20,
+        position_wise_mlp_dim=20,
+    )
+
+    shift = 10
+    train_batch = 10
+    test_batch = 100
+    data_gen = bit_shift_generator(seq_length, shift=shift, batch_size=train_batch)
+    test_gen = bit_shift_generator(seq_length, shift=shift, batch_size=test_batch)
+
+    @tf.function
+    def update_step(inputs, targets):
+        model_out = model(
+            {"obs": inputs},
+            state=[tf.reshape(inputs, [-1, seq_length, 1])],
+            seq_lens=np.full(shape=(train_batch,), fill_value=seq_length),
+        )
+        optimizer.minimize(
+            lambda: train_loss(targets, model_out), lambda: model.trainable_variables
+        )
+
+    for i, (inputs, targets) in zip(range(num_iterations), data_gen):
+        inputs_in = np.reshape(inputs, [-1, 1])
+        targets_in = np.reshape(targets, [-1])
+        update_step(tf.convert_to_tensor(inputs_in), tf.convert_to_tensor(targets_in))
+
+        if i % print_every_n == 0:
+            test_inputs, test_targets = next(test_gen)
+            print(i, train_loss(test_targets, model(test_inputs)))
+
+
+if __name__ == "__main__":
+    tf.enable_eager_execution()
+    train_bit_shift(
+        seq_length=20,
+        num_iterations=2000,
+        print_every_n=200,
+    )

rllib/examples/_old_api_stack/catalogs/mobilenet_v2_encoder.py

@@ -0,0 +1,78 @@
+# @OldAPIStack
+
+"""
+This example shows two modifications:
+- How to write a custom Encoder (using MobileNet v2)
+- How to enhance Catalogs with this custom Encoder
+
+With the pattern shown in this example, we can enhance Catalogs such that they extend
+to new observation- or action spaces while retaining their original functionality.
+"""
+# __sphinx_doc_begin__
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.ppo.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples._old_api_stack.models.mobilenet_v2_encoder import (
+    MOBILENET_INPUT_SHAPE,
+    MobileNetV2EncoderConfig,
+)
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+
+
+# Define a PPO Catalog that we can use to inject our MobileNetV2 Encoder into RLlib's
+# decision tree of what model to choose
+class MobileNetEnhancedPPOCatalog(PPOCatalog):
+    @classmethod
+    def _get_encoder_config(
+        cls,
+        observation_space: gym.Space,
+        **kwargs,
+    ):
+        if (
+            isinstance(observation_space, gym.spaces.Box)
+            and observation_space.shape == MOBILENET_INPUT_SHAPE
+        ):
+            # Inject our custom encoder here, only if the observation space fits it
+            return MobileNetV2EncoderConfig()
+        else:
+            return super()._get_encoder_config(observation_space, **kwargs)
+
+
+# Create a generic config with our enhanced Catalog
+ppo_config = (
+    PPOConfig()
+    .rl_module(rl_module_spec=RLModuleSpec(catalog_class=MobileNetEnhancedPPOCatalog))
+    .env_runners(num_env_runners=0)
+    # The following training settings make it so that a training iteration is very
+    # quick. This is just for the sake of this example. PPO will not learn properly
+    # with these settings!
+    .training(train_batch_size_per_learner=32, minibatch_size=16, num_epochs=1)
+)
+
+# CartPole's observation space is not compatible with our MobileNetV2 Encoder, so
+# this will use the default behaviour of Catalogs
+ppo_config.environment("CartPole-v1")
+results = ppo_config.build().train()
+print(results)
+
+# For this training, we use a RandomEnv with observations of shape
+# MOBILENET_INPUT_SHAPE. This will use our custom Encoder.
+ppo_config.environment(
+    RandomEnv,
+    env_config={
+        "action_space": gym.spaces.Discrete(2),
+        # Test a simple Image observation space.
+        "observation_space": gym.spaces.Box(
+            0.0,
+            1.0,
+            shape=MOBILENET_INPUT_SHAPE,
+            dtype=np.float32,
+        ),
+    },
+)
+results = ppo_config.build().train()
+print(results)
+# __sphinx_doc_end__

rllib/examples/_old_api_stack/catalogs/onnx_torch_lstm.py

@@ -0,0 +1,131 @@
+# @OldAPIStack
+
+import numpy as np
+import onnxruntime
+
+import ray
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.examples.utils import add_rllib_example_script_args, check
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, _ = try_import_torch()
+
+parser = add_rllib_example_script_args()
+parser.set_defaults(
+    num_env_runners=1,
+    # ONNX is not supported by RLModule API yet.
+    old_api_stack=True,
+)
+
+
+class ONNXCompatibleWrapper(torch.nn.Module):
+    def __init__(self, original_model):
+        super(ONNXCompatibleWrapper, self).__init__()
+        self.original_model = original_model
+
+    def forward(self, a, b0, b1, c):
+        # Convert the separate tensor inputs back into the list format
+        # expected by the original model's forward method.
+        b = [b0, b1]
+        ret = self.original_model({"obs": a}, b, c)
+        # results, state_out_0, state_out_1
+        return ret[0], ret[1][0], ret[1][1]
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    ray.init()
+
+    # Configure our PPO Algorithm.
+    config = (
+        ppo.PPOConfig()
+        .environment("CartPole-v1")
+        .env_runners(num_env_runners=args.num_env_runners)
+        .training(model={"use_lstm": True})
+    )
+
+    B = 3
+    T = 5
+    LSTM_CELL = 256
+
+    # Input data for a python inference forward call.
+    test_data_python = {
+        "obs": np.random.uniform(0, 1.0, size=(B * T, 4)).astype(np.float32),
+        "state_ins": [
+            np.random.uniform(0, 1.0, size=(B, LSTM_CELL)).astype(np.float32),
+            np.random.uniform(0, 1.0, size=(B, LSTM_CELL)).astype(np.float32),
+        ],
+        "seq_lens": np.array([T] * B, np.float32),
+    }
+    # Input data for the ONNX session.
+    test_data_onnx = {
+        "obs": test_data_python["obs"],
+        "state_in_0": test_data_python["state_ins"][0],
+        "state_in_1": test_data_python["state_ins"][1],
+        "seq_lens": test_data_python["seq_lens"],
+    }
+
+    # Input data for compiling the ONNX model.
+    test_data_onnx_input = convert_to_torch_tensor(test_data_onnx)
+
+    # Initialize a PPO Algorithm.
+    algo = config.build()
+
+    # You could train the model here
+    # algo.train()
+
+    # Let's run inference on the torch model
+    policy = algo.get_policy()
+    result_pytorch, _ = policy.model(
+        {
+            "obs": torch.tensor(test_data_python["obs"]),
+        },
+        [
+            torch.tensor(test_data_python["state_ins"][0]),
+            torch.tensor(test_data_python["state_ins"][1]),
+        ],
+        torch.tensor(test_data_python["seq_lens"]),
+    )
+
+    # Evaluate tensor to fetch numpy array
+    result_pytorch = result_pytorch.detach().numpy()
+
+    # Wrap the actual ModelV2 with the torch wrapper above to make this all work with
+    # LSTMs (extra `state` in- and outputs and `seq_lens` inputs).
+    onnx_compatible = ONNXCompatibleWrapper(policy.model)
+    exported_model_file = "model.onnx"
+    input_names = [
+        "obs",
+        "state_in_0",
+        "state_in_1",
+        "seq_lens",
+    ]
+
+    # This line will export the model to ONNX.
+    torch.onnx.export(
+        onnx_compatible,
+        tuple(test_data_onnx_input[n] for n in input_names),
+        exported_model_file,
+        export_params=True,
+        opset_version=11,
+        do_constant_folding=True,
+        input_names=input_names,
+        output_names=[
+            "output",
+            "state_out_0",
+            "state_out_1",
+        ],
+        dynamic_axes={k: {0: "batch_size"} for k in input_names},
+    )
+    # Start an inference session for the ONNX model.
+    session = onnxruntime.InferenceSession(exported_model_file, None)
+    result_onnx = session.run(["output"], test_data_onnx)
+
+    # These results should be equal!
+    print("PYTORCH", result_pytorch)
+    print("ONNX", result_onnx[0])
+
+    check(result_pytorch, result_onnx[0])
+    print("Model outputs are equal. PASSED")

rllib/examples/_old_api_stack/centralized_critic.py

@@ -0,0 +1,318 @@
+# @OldAPIStack
+
+# ***********************************************************************************
+# IMPORTANT NOTE: This script uses the old API stack and will soon be replaced by
+# `ray.rllib.examples.multi_agent.pettingzoo_shared_value_function.py`!
+# ***********************************************************************************
+
+"""An example of customizing PPO to leverage a centralized critic.
+
+Here the model and policy are hard-coded to implement a centralized critic
+for TwoStepGame, but you can adapt this for your own use cases.
+
+Compared to simply running `rllib/examples/two_step_game.py --run=PPO`,
+this centralized critic version reaches vf_explained_variance=1.0 more stably
+since it takes into account the opponent actions as well as the policy's.
+Note that this is also using two independent policies instead of weight-sharing
+with one.
+
+See also: centralized_critic_2.py for a simpler approach that instead
+modifies the environment.
+"""
+
+import argparse
+import os
+
+import numpy as np
+from gymnasium.spaces import Discrete
+
+from ray import tune
+from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig
+from ray.rllib.algorithms.ppo.ppo_tf_policy import (
+    PPOTF1Policy,
+    PPOTF2Policy,
+)
+from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy
+from ray.rllib.evaluation.postprocessing import Postprocessing, compute_advantages
+from ray.rllib.examples._old_api_stack.models.centralized_critic_models import (
+    CentralizedCriticModel,
+    TorchCentralizedCriticModel,
+)
+from ray.rllib.examples.envs.classes.multi_agent.two_step_game import TwoStepGame
+from ray.rllib.models import ModelCatalog
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.test_utils import check_learning_achieved
+from ray.rllib.utils.tf_utils import explained_variance, make_tf_callable
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+from ray.tune.result import TRAINING_ITERATION
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+OPPONENT_OBS = "opponent_obs"
+OPPONENT_ACTION = "opponent_action"
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--as-test",
+    action="store_true",
+    help="Whether this script should be run as a test: --stop-reward must "
+    "be achieved within --stop-timesteps AND --stop-iters.",
+)
+parser.add_argument(
+    "--stop-iters", type=int, default=100, help="Number of iterations to train."
+)
+parser.add_argument(
+    "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train."
+)
+parser.add_argument(
+    "--stop-reward", type=float, default=7.99, help="Reward at which we stop training."
+)
+
+
+class CentralizedValueMixin:
+    """Add method to evaluate the central value function from the model."""
+
+    def __init__(self):
+        if self.config["framework"] != "torch":
+            self.compute_central_vf = make_tf_callable(self.get_session())(
+                self.model.central_value_function
+            )
+        else:
+            self.compute_central_vf = self.model.central_value_function
+
+
+# Grabs the opponent obs/act and includes it in the experience train_batch,
+# and computes GAE using the central vf predictions.
+def centralized_critic_postprocessing(
+    policy, sample_batch, other_agent_batches=None, episode=None
+):
+    pytorch = policy.config["framework"] == "torch"
+    if (pytorch and hasattr(policy, "compute_central_vf")) or (
+        not pytorch and policy.loss_initialized()
+    ):
+        assert other_agent_batches is not None
+        [(_, _, opponent_batch)] = list(other_agent_batches.values())
+
+        # also record the opponent obs and actions in the trajectory
+        sample_batch[OPPONENT_OBS] = opponent_batch[SampleBatch.CUR_OBS]
+        sample_batch[OPPONENT_ACTION] = opponent_batch[SampleBatch.ACTIONS]
+
+        # overwrite default VF prediction with the central VF
+        if args.framework == "torch":
+            sample_batch[SampleBatch.VF_PREDS] = (
+                policy.compute_central_vf(
+                    convert_to_torch_tensor(
+                        sample_batch[SampleBatch.CUR_OBS], policy.device
+                    ),
+                    convert_to_torch_tensor(sample_batch[OPPONENT_OBS], policy.device),
+                    convert_to_torch_tensor(
+                        sample_batch[OPPONENT_ACTION], policy.device
+                    ),
+                )
+                .cpu()
+                .detach()
+                .numpy()
+            )
+        else:
+            sample_batch[SampleBatch.VF_PREDS] = convert_to_numpy(
+                policy.compute_central_vf(
+                    sample_batch[SampleBatch.CUR_OBS],
+                    sample_batch[OPPONENT_OBS],
+                    sample_batch[OPPONENT_ACTION],
+                )
+            )
+    else:
+        # Policy hasn't been initialized yet, use zeros.
+        sample_batch[OPPONENT_OBS] = np.zeros_like(sample_batch[SampleBatch.CUR_OBS])
+        sample_batch[OPPONENT_ACTION] = np.zeros_like(sample_batch[SampleBatch.ACTIONS])
+        sample_batch[SampleBatch.VF_PREDS] = np.zeros_like(
+            sample_batch[SampleBatch.REWARDS], dtype=np.float32
+        )
+
+    completed = sample_batch[SampleBatch.TERMINATEDS][-1]
+    if completed:
+        last_r = 0.0
+    else:
+        last_r = sample_batch[SampleBatch.VF_PREDS][-1]
+
+    train_batch = compute_advantages(
+        sample_batch,
+        last_r,
+        policy.config["gamma"],
+        policy.config["lambda"],
+        use_gae=policy.config["use_gae"],
+    )
+    return train_batch
+
+
+# Copied from PPO but optimizing the central value function.
+def loss_with_central_critic(policy, base_policy, model, dist_class, train_batch):
+    # Save original value function.
+    vf_saved = model.value_function
+
+    # Calculate loss with a custom value function.
+    model.value_function = lambda: policy.model.central_value_function(
+        train_batch[SampleBatch.CUR_OBS],
+        train_batch[OPPONENT_OBS],
+        train_batch[OPPONENT_ACTION],
+    )
+    policy._central_value_out = model.value_function()
+    loss = base_policy.loss(model, dist_class, train_batch)
+
+    # Restore original value function.
+    model.value_function = vf_saved
+
+    return loss
+
+
+def central_vf_stats(policy, train_batch):
+    # Report the explained variance of the central value function.
+    return {
+        "vf_explained_var": explained_variance(
+            train_batch[Postprocessing.VALUE_TARGETS], policy._central_value_out
+        )
+    }
+
+
+def get_ccppo_policy(base):
+    class CCPPOTFPolicy(CentralizedValueMixin, base):
+        def __init__(self, observation_space, action_space, config):
+            base.__init__(self, observation_space, action_space, config)
+            CentralizedValueMixin.__init__(self)
+
+        @override(base)
+        def loss(self, model, dist_class, train_batch):
+            # Use super() to get to the base PPO policy.
+            # This special loss function utilizes a shared
+            # value function defined on self, and the loss function
+            # defined on PPO policies.
+            return loss_with_central_critic(
+                self, super(), model, dist_class, train_batch
+            )
+
+        @override(base)
+        def postprocess_trajectory(
+            self, sample_batch, other_agent_batches=None, episode=None
+        ):
+            return centralized_critic_postprocessing(
+                self, sample_batch, other_agent_batches, episode
+            )
+
+        @override(base)
+        def stats_fn(self, train_batch: SampleBatch):
+            stats = super().stats_fn(train_batch)
+            stats.update(central_vf_stats(self, train_batch))
+            return stats
+
+    return CCPPOTFPolicy
+
+
+CCPPOStaticGraphTFPolicy = get_ccppo_policy(PPOTF1Policy)
+CCPPOEagerTFPolicy = get_ccppo_policy(PPOTF2Policy)
+
+
+class CCPPOTorchPolicy(CentralizedValueMixin, PPOTorchPolicy):
+    def __init__(self, observation_space, action_space, config):
+        PPOTorchPolicy.__init__(self, observation_space, action_space, config)
+        CentralizedValueMixin.__init__(self)
+
+    @override(PPOTorchPolicy)
+    def loss(self, model, dist_class, train_batch):
+        return loss_with_central_critic(self, super(), model, dist_class, train_batch)
+
+    @override(PPOTorchPolicy)
+    def postprocess_trajectory(
+        self, sample_batch, other_agent_batches=None, episode=None
+    ):
+        return centralized_critic_postprocessing(
+            self, sample_batch, other_agent_batches, episode
+        )
+
+
+class CentralizedCritic(PPO):
+    @classmethod
+    @override(PPO)
+    def get_default_policy_class(cls, config):
+        if config["framework"] == "torch":
+            return CCPPOTorchPolicy
+        elif config["framework"] == "tf":
+            return CCPPOStaticGraphTFPolicy
+        else:
+            return CCPPOEagerTFPolicy
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    ModelCatalog.register_custom_model(
+        "cc_model",
+        TorchCentralizedCriticModel
+        if args.framework == "torch"
+        else CentralizedCriticModel,
+    )
+
+    config = (
+        PPOConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment(TwoStepGame)
+        .framework(args.framework)
+        .env_runners(batch_mode="complete_episodes", num_env_runners=0)
+        .training(model={"custom_model": "cc_model"})
+        .multi_agent(
+            policies={
+                "pol1": (
+                    None,
+                    Discrete(6),
+                    TwoStepGame.action_space,
+                    # `framework` would also be ok here.
+                    PPOConfig.overrides(framework_str=args.framework),
+                ),
+                "pol2": (
+                    None,
+                    Discrete(6),
+                    TwoStepGame.action_space,
+                    # `framework` would also be ok here.
+                    PPOConfig.overrides(framework_str=args.framework),
+                ),
+            },
+            policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: "pol1"
+            if agent_id == 0
+            else "pol2",
+        )
+        # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+        .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0")))
+    )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+    }
+
+    tuner = tune.Tuner(
+        CentralizedCritic,
+        param_space=config.to_dict(),
+        run_config=tune.RunConfig(stop=stop, verbose=1),
+    )
+    results = tuner.fit()
+
+    if args.as_test:
+        check_learning_achieved(results, args.stop_reward)

rllib/examples/_old_api_stack/checkpoints/cartpole_dqn_export.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+
+# @OldAPIStack
+
+import os
+
+import numpy as np
+
+import ray
+import ray._common
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.framework import try_import_tf
+from ray.tune.registry import get_trainable_cls
+
+tf1, tf, tfv = try_import_tf()
+
+ray.init()
+
+
+def train_and_export_policy_and_model(algo_name, num_steps, model_dir, ckpt_dir):
+    cls = get_trainable_cls(algo_name)
+    config = cls.get_default_config()
+    config.api_stack(
+        enable_rl_module_and_learner=False,
+        enable_env_runner_and_connector_v2=False,
+    )
+    # This Example is only for tf.
+    config.framework("tf")
+    # Set exporting native (DL-framework) model files to True.
+    config.export_native_model_files = True
+    config.env = "CartPole-v1"
+    alg = config.build()
+    for _ in range(num_steps):
+        alg.train()
+
+    # Export Policy checkpoint.
+    alg.export_policy_checkpoint(ckpt_dir)
+    # Export tensorflow keras Model for online serving
+    alg.export_policy_model(model_dir)
+
+
+def restore_saved_model(export_dir):
+    signature_key = (
+        tf1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+    )
+    g = tf1.Graph()
+    with g.as_default():
+        with tf1.Session(graph=g) as sess:
+            meta_graph_def = tf1.saved_model.load(
+                sess, [tf1.saved_model.tag_constants.SERVING], export_dir
+            )
+            print("Model restored!")
+            print("Signature Def Information:")
+            print(meta_graph_def.signature_def[signature_key])
+            print("You can inspect the model using TensorFlow SavedModel CLI.")
+            print("https://www.tensorflow.org/guide/saved_model")
+
+
+def restore_policy_from_checkpoint(export_dir):
+    # Load the model from the checkpoint.
+    policy = Policy.from_checkpoint(export_dir)
+    # Perform a dummy (CartPole) forward pass.
+    test_obs = np.array([0.1, 0.2, 0.3, 0.4])
+    results = policy.compute_single_action(test_obs)
+    # Check results for correctness.
+    assert len(results) == 3
+    assert results[0].shape == ()  # pure single action (int)
+    assert results[1] == []  # RNN states
+    assert results[2]["action_dist_inputs"].shape == (2,)  # categorical inputs
+
+
+if __name__ == "__main__":
+    algo = "PPO"
+    model_dir = os.path.join(
+        ray._common.utils.get_default_ray_temp_dir(), "model_export_dir"
+    )
+    ckpt_dir = os.path.join(
+        ray._common.utils.get_default_ray_temp_dir(), "ckpt_export_dir"
+    )
+    num_steps = 1
+    train_and_export_policy_and_model(algo, num_steps, model_dir, ckpt_dir)
+    restore_saved_model(model_dir)
+    restore_policy_from_checkpoint(ckpt_dir)

rllib/examples/_old_api_stack/compute_adapted_gae_on_postprocess_trajectory.py

@@ -0,0 +1,158 @@
+# @OldAPIStack
+
+"""
+Adapted (time-dependent) GAE for PPO algorithm that you can activate by setting
+use_adapted_gae=True in the policy config. Additionally, it's required that
+"callbacks" include the custom callback class in the Algorithm's config.
+Furthermore, the env must return in its info dictionary a key-value pair of
+the form "d_ts": ... where the value is the length (time) of recent agent step.
+
+This adapted, time-dependent computation of advantages may be useful in cases
+where agent's actions take various times and thus time steps are not
+equidistant (https://docdro.id/400TvlR)
+"""
+
+import numpy as np
+
+from ray.rllib.callbacks.callbacks import RLlibCallback
+from ray.rllib.evaluation.postprocessing import Postprocessing
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+
+
+class MyCallbacks(RLlibCallback):
+    @override(RLlibCallback)
+    def on_postprocess_trajectory(
+        self,
+        *,
+        worker,
+        episode,
+        agent_id,
+        policy_id,
+        policies,
+        postprocessed_batch,
+        original_batches,
+        **kwargs
+    ):
+        super().on_postprocess_trajectory(
+            worker=worker,
+            episode=episode,
+            agent_id=agent_id,
+            policy_id=policy_id,
+            policies=policies,
+            postprocessed_batch=postprocessed_batch,
+            original_batches=original_batches,
+            **kwargs
+        )
+
+        if policies[policy_id].config.get("use_adapted_gae", False):
+            policy = policies[policy_id]
+            assert policy.config[
+                "use_gae"
+            ], "Can't use adapted gae without use_gae=True!"
+
+            info_dicts = postprocessed_batch[SampleBatch.INFOS]
+            assert np.all(
+                ["d_ts" in info_dict for info_dict in info_dicts]
+            ), "Info dicts in sample batch must contain data 'd_ts' \
+                (=ts[i+1]-ts[i] length of time steps)!"
+
+            d_ts = np.array(
+                [np.float(info_dict.get("d_ts")) for info_dict in info_dicts]
+            )
+            assert np.all(
+                [e.is_integer() for e in d_ts]
+            ), "Elements of 'd_ts' (length of time steps) must be integer!"
+
+            # Trajectory is actually complete -> last r=0.0.
+            if postprocessed_batch[SampleBatch.TERMINATEDS][-1]:
+                last_r = 0.0
+            # Trajectory has been truncated -> last r=VF estimate of last obs.
+            else:
+                # Input dict is provided to us automatically via the Model's
+                # requirements. It's a single-timestep (last one in trajectory)
+                # input_dict.
+                # Create an input dict according to the Model's requirements.
+                input_dict = postprocessed_batch.get_single_step_input_dict(
+                    policy.model.view_requirements, index="last"
+                )
+                last_r = policy._value(**input_dict)
+
+            gamma = policy.config["gamma"]
+            lambda_ = policy.config["lambda"]
+
+            vpred_t = np.concatenate(
+                [postprocessed_batch[SampleBatch.VF_PREDS], np.array([last_r])]
+            )
+            delta_t = (
+                postprocessed_batch[SampleBatch.REWARDS]
+                + gamma**d_ts * vpred_t[1:]
+                - vpred_t[:-1]
+            )
+            # This formula for the advantage is an adaption of
+            # "Generalized Advantage Estimation"
+            # (https://arxiv.org/abs/1506.02438) which accounts for time steps
+            # of irregular length (see proposal here ).
+            # NOTE: last time step delta is not required
+            postprocessed_batch[
+                Postprocessing.ADVANTAGES
+            ] = generalized_discount_cumsum(delta_t, d_ts[:-1], gamma * lambda_)
+            postprocessed_batch[Postprocessing.VALUE_TARGETS] = (
+                postprocessed_batch[Postprocessing.ADVANTAGES]
+                + postprocessed_batch[SampleBatch.VF_PREDS]
+            ).astype(np.float32)
+
+            postprocessed_batch[Postprocessing.ADVANTAGES] = postprocessed_batch[
+                Postprocessing.ADVANTAGES
+            ].astype(np.float32)
+
+
+def generalized_discount_cumsum(
+    x: np.ndarray, deltas: np.ndarray, gamma: float
+) -> np.ndarray:
+    """Calculates the 'time-dependent' discounted cumulative sum over a
+    (reward) sequence `x`.
+
+    Recursive equations:
+
+    y[t] - gamma**deltas[t+1]*y[t+1] = x[t]
+
+    reversed(y)[t] - gamma**reversed(deltas)[t-1]*reversed(y)[t-1] =
+    reversed(x)[t]
+
+    Args:
+        x (np.ndarray): A sequence of rewards or one-step TD residuals.
+        deltas (np.ndarray): A sequence of time step deltas (length of time
+            steps).
+        gamma: The discount factor gamma.
+
+    Returns:
+        np.ndarray: The sequence containing the 'time-dependent' discounted
+            cumulative sums for each individual element in `x` till the end of
+            the trajectory.
+
+    .. testcode::
+        :skipif: True
+
+        x = np.array([0.0, 1.0, 2.0, 3.0])
+        deltas = np.array([1.0, 4.0, 15.0])
+        gamma = 0.9
+        generalized_discount_cumsum(x, deltas, gamma)
+
+    .. testoutput::
+
+        array([0.0 + 0.9^1.0*1.0 + 0.9^4.0*2.0 + 0.9^15.0*3.0,
+               1.0 + 0.9^4.0*2.0 + 0.9^15.0*3.0,
+               2.0 + 0.9^15.0*3.0,
+               3.0])
+    """
+    reversed_x = x[::-1]
+    reversed_deltas = deltas[::-1]
+    reversed_y = np.empty_like(x)
+    reversed_y[0] = reversed_x[0]
+    for i in range(1, x.size):
+        reversed_y[i] = (
+            reversed_x[i] + gamma ** reversed_deltas[i - 1] * reversed_y[i - 1]
+        )
+
+    return reversed_y[::-1]

rllib/examples/_old_api_stack/connectors/prepare_checkpoint.py

@@ -0,0 +1,51 @@
+# @OldAPIStack
+import random
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.algorithms.sac import SACConfig
+
+
+def create_appo_cartpole_checkpoint(output_dir, use_lstm=False):
+    config = (
+        APPOConfig()
+        .api_stack(
+            enable_rl_module_and_learner=False,
+            enable_env_runner_and_connector_v2=False,
+        )
+        .environment("CartPole-v1")
+        .training(model={"use_lstm": use_lstm})
+    )
+    # Build algorithm object.
+    algo = config.build()
+    algo.save(checkpoint_dir=output_dir)
+
+
+def create_open_spiel_checkpoint(output_dir):
+    def _policy_mapping_fn(*args, **kwargs):
+        random.choice(["main", "opponent"])
+
+    config = (
+        SACConfig()
+        .environment("open_spiel_env")
+        # Intentionally create a TF2 policy to demonstrate that we can restore
+        # and use a TF policy in a Torch training stack.
+        .framework("tf2")
+        .env_runners(
+            num_env_runners=1,
+            num_envs_per_env_runner=5,
+            # We will be restoring a TF2 policy.
+            # So tell the RolloutWorkers to enable TF eager exec as well, even if
+            # framework is set to torch.
+            enable_tf1_exec_eagerly=True,
+        )
+        .training(model={"fcnet_hiddens": [512, 512]})
+        .multi_agent(
+            policies={"main", "opponent"},
+            policy_mapping_fn=_policy_mapping_fn,
+            # Just train the "main" policy.
+            policies_to_train=["main"],
+        )
+    )
+    # Build algorithm object.
+    algo = config.build()
+    algo.save(checkpoint_dir=output_dir)

rllib/examples/_old_api_stack/connectors/run_connector_policy.py

@@ -0,0 +1,78 @@
+# @OldAPIStack
+"""This example script loads a connector enabled policy,
+and uses it in a serving or inference setting.
+"""
+
+import argparse
+import os
+import tempfile
+
+import gymnasium as gym
+
+from ray.rllib.examples._old_api_stack.connectors.prepare_checkpoint import (
+    # For demo purpose only. Would normally not need this.
+    create_appo_cartpole_checkpoint,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.policy import local_policy_inference
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--use-lstm", action="store_true", help="Add LSTM to the setup.")
+
+
+def run(checkpoint_path, policy_id):
+    # __sphinx_doc_begin__
+    # Restore policy.
+    policy = Policy.from_checkpoint(
+        checkpoint=checkpoint_path,
+        policy_ids=[policy_id],
+    )
+
+    # Run CartPole.
+    env = gym.make("CartPole-v1")
+    env_id = "env_1"
+    obs, info = env.reset()
+    # Run for 2 episodes.
+    episodes = step = 0
+    while episodes < 2:
+        # Use local_policy_inference() to run inference, so we do not have to
+        # provide policy states or extra fetch dictionaries.
+        # "env_1" and "agent_1" are dummy env and agent IDs to run connectors with.
+        policy_outputs = local_policy_inference(
+            policy, env_id, "agent_1", obs, explore=False
+        )
+        assert len(policy_outputs) == 1
+        action, _, _ = policy_outputs[0]
+        print(f"episode {episodes} step {step}", obs, action)
+
+        # Step environment forward one more step.
+        obs, _, terminated, truncated, _ = env.step(action)
+        step += 1
+
+        # If the episode is done, reset the env and our connectors and start a new
+        # episode.
+        if terminated or truncated:
+            episodes += 1
+            step = 0
+            obs, info = env.reset()
+            policy.agent_connectors.reset(env_id)
+
+    # __sphinx_doc_end__
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        policy_id = "default_policy"
+
+        # Note, this is just for demo purpose.
+        # Normally, you would use a policy checkpoint from a real training run.
+        create_appo_cartpole_checkpoint(tmpdir, args.use_lstm)
+        policy_checkpoint_path = os.path.join(
+            tmpdir,
+            "policies",
+            policy_id,
+        )
+
+        run(policy_checkpoint_path, policy_id)

rllib/examples/_old_api_stack/connectors/self_play_with_policy_checkpoint.py

@@ -0,0 +1,152 @@
+# @OldAPIStack
+"""Example showing to restore a connector enabled TF policy
+checkpoint for a new self-play PyTorch training job.
+You can train the checkpointed policy with a different algorithm too.
+"""
+
+import argparse
+import os
+import tempfile
+from functools import partial
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.sac import SACConfig
+from ray.rllib.callbacks.callbacks import RLlibCallback
+from ray.rllib.env.utils import try_import_pyspiel
+from ray.rllib.env.wrappers.open_spiel import OpenSpielEnv
+from ray.rllib.examples._old_api_stack.connectors.prepare_checkpoint import (
+    create_open_spiel_checkpoint,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    NUM_EPISODES,
+)
+from ray.tune import CLIReporter, register_env
+from ray.tune.result import TRAINING_ITERATION
+
+pyspiel = try_import_pyspiel(error=True)
+register_env(
+    "open_spiel_env", lambda _: OpenSpielEnv(pyspiel.load_game("connect_four"))
+)
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--train_iteration",
+    type=int,
+    default=10,
+    help="Number of iterations to train.",
+)
+args = parser.parse_args()
+
+
+MAIN_POLICY_ID = "main"
+OPPONENT_POLICY_ID = "opponent"
+
+
+class AddPolicyCallback(RLlibCallback):
+    def __init__(self, checkpoint_dir):
+        self._checkpoint_dir = checkpoint_dir
+        super().__init__()
+
+    def on_algorithm_init(self, *, algorithm, metrics_logger, **kwargs):
+        policy = Policy.from_checkpoint(
+            self._checkpoint_dir, policy_ids=[OPPONENT_POLICY_ID]
+        )
+
+        # Add restored policy to Algorithm.
+        # Note that this policy doesn't have to be trained with the same algorithm
+        # of the training stack. You can even mix up TF policies with a Torch stack.
+        algorithm.add_policy(
+            policy_id=OPPONENT_POLICY_ID,
+            policy=policy,
+            add_to_eval_env_runners=True,
+        )
+
+
+def policy_mapping_fn(agent_id, episode, worker, **kwargs):
+    # main policy plays against opponent policy.
+    return MAIN_POLICY_ID if episode.episode_id % 2 == agent_id else OPPONENT_POLICY_ID
+
+
+def main(checkpoint_dir):
+    config = (
+        SACConfig()
+        .environment("open_spiel_env")
+        .framework("torch")
+        .callbacks(partial(AddPolicyCallback, checkpoint_dir))
+        .env_runners(
+            num_env_runners=1,
+            num_envs_per_env_runner=5,
+            # We will be restoring a TF2 policy.
+            # So tell the RolloutWorkers to enable TF eager exec as well, even if
+            # framework is set to torch.
+            enable_tf1_exec_eagerly=True,
+        )
+        .training(model={"fcnet_hiddens": [512, 512]})
+        .multi_agent(
+            # Initial policy map: Random and PPO. This will be expanded
+            # to more policy snapshots taken from "main" against which "main"
+            # will then play (instead of "random"). This is done in the
+            # custom callback defined above (`SelfPlayCallback`).
+            # Note: We will add the "opponent" policy with callback.
+            policies={MAIN_POLICY_ID},  # Our main policy, we'd like to optimize.
+            # Assign agent 0 and 1 randomly to the "main" policy or
+            # to the opponent ("random" at first). Make sure (via episode_id)
+            # that "main" always plays against "random" (and not against
+            # another "main").
+            policy_mapping_fn=policy_mapping_fn,
+            # Always just train the "main" policy.
+            policies_to_train=[MAIN_POLICY_ID],
+        )
+    )
+
+    stop = {TRAINING_ITERATION: args.train_iteration}
+
+    # Train the "main" policy to play really well using self-play.
+    tuner = tune.Tuner(
+        "SAC",
+        param_space=config.to_dict(),
+        run_config=tune.RunConfig(
+            stop=stop,
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_at_end=True,
+                checkpoint_frequency=10,
+            ),
+            verbose=2,
+            progress_reporter=CLIReporter(
+                metric_columns={
+                    TRAINING_ITERATION: "iter",
+                    "time_total_s": "time_total_s",
+                    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": "ts",
+                    f"{ENV_RUNNER_RESULTS}/{NUM_EPISODES}": "train_episodes",
+                    (
+                        f"{ENV_RUNNER_RESULTS}/module_episode_returns_mean/main"
+                    ): "reward_main",
+                },
+                sort_by_metric=True,
+            ),
+        ),
+    )
+    tuner.fit()
+
+
+if __name__ == "__main__":
+    ray.init()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        create_open_spiel_checkpoint(tmpdir)
+
+        policy_checkpoint_path = os.path.join(
+            tmpdir,
+            "checkpoint_000000",
+            "policies",
+            OPPONENT_POLICY_ID,
+        )
+
+        main(policy_checkpoint_path)
+
+    ray.shutdown()

rllib/examples/_old_api_stack/envs/greyscale_env.py

@@ -0,0 +1,126 @@
+# @OldAPIStack
+"""
+Example of interfacing with an environment that produces 2D observations.
+
+This example shows how turning 2D observations with shape (A, B) into a 3D
+observations with shape (C, D, 1) can enable usage of RLlib's default models.
+RLlib's default Catalog class does not provide default models for 2D observation
+spaces, but it does so for 3D observations.
+Therefore, one can either write a custom model or transform the 2D observations into 3D
+observations. This enables RLlib to use one of the default CNN filters, even though the
+original observation space of the environment does not fit them.
+
+This simple example should reach rewards of 50 within 150k timesteps.
+"""
+
+import argparse
+
+from numpy import float32
+from pettingzoo.butterfly import pistonball_v6
+from supersuit import (
+    color_reduction_v0,
+    dtype_v0,
+    normalize_obs_v0,
+    reshape_v0,
+    resize_v1,
+)
+
+from ray import tune
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.env import PettingZooEnv
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import register_env
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--framework",
+    choices=["tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--as-test",
+    action="store_true",
+    help="Whether this script should be run as a compilation test.",
+)
+parser.add_argument(
+    "--stop-iters", type=int, default=150, help="Number of iterations to train."
+)
+parser.add_argument(
+    "--stop-timesteps", type=int, default=1000000, help="Number of timesteps to train."
+)
+parser.add_argument(
+    "--stop-reward", type=float, default=50, help="Reward at which we stop training."
+)
+
+args = parser.parse_args()
+
+
+# The space we down-sample and transform the greyscale pistonball images to.
+# Other spaces supported by RLlib can be chosen here.
+TRANSFORMED_OBS_SPACE = (42, 42, 1)
+
+
+def env_creator(config):
+    env = pistonball_v6.env(n_pistons=5)
+    env = dtype_v0(env, dtype=float32)
+    # This gives us greyscale images for the color red
+    env = color_reduction_v0(env, mode="R")
+    env = normalize_obs_v0(env)
+    # This gives us images that are upsampled to the number of pixels in the
+    # default CNN filter
+    env = resize_v1(
+        env, x_size=TRANSFORMED_OBS_SPACE[0], y_size=TRANSFORMED_OBS_SPACE[1]
+    )
+    # This gives us 3D images for which we have default filters
+    env = reshape_v0(env, shape=TRANSFORMED_OBS_SPACE)
+    return env
+
+
+# Register env
+register_env("pistonball", lambda config: PettingZooEnv(env_creator(config)))
+
+config = (
+    PPOConfig()
+    .environment("pistonball", env_config={"local_ratio": 0.5}, clip_rewards=True)
+    .env_runners(
+        num_env_runners=15 if not args.as_test else 2,
+        num_envs_per_env_runner=1,
+        observation_filter="NoFilter",
+        rollout_fragment_length="auto",
+    )
+    .framework("torch")
+    .training(
+        entropy_coeff=0.01,
+        vf_loss_coeff=0.1,
+        clip_param=0.1,
+        vf_clip_param=10.0,
+        num_epochs=10,
+        kl_coeff=0.5,
+        lr=0.0001,
+        grad_clip=100,
+        minibatch_size=500,
+        train_batch_size=5000 if not args.as_test else 1000,
+        model={"vf_share_layers": True},
+    )
+    .resources(num_gpus=1 if not args.as_test else 0)
+    .reporting(min_time_s_per_iteration=30)
+)
+
+tune.Tuner(
+    "PPO",
+    param_space=config.to_dict(),
+    run_config=tune.RunConfig(
+        stop={
+            TRAINING_ITERATION: args.stop_iters,
+            NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+            f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+        },
+        verbose=2,
+    ),
+).fit()

rllib/examples/_old_api_stack/inference/policy_inference_after_training_with_attention.py

@@ -0,0 +1,264 @@
+# @OldAPIStack
+"""
+Example showing how you can use your trained policy for inference
+(computing actions) in an environment.
+
+Includes options for LSTM-based models (--use-lstm), attention-net models
+(--use-attention), and plain (non-recurrent) models.
+"""
+import argparse
+import os
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import get_trainable_cls
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use."
+)
+parser.add_argument("--num-cpus", type=int, default=0)
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--prev-n-actions",
+    type=int,
+    default=0,
+    help="Feed n most recent actions to the attention net as part of its input.",
+)
+parser.add_argument(
+    "--prev-n-rewards",
+    type=int,
+    default=0,
+    help="Feed n most recent rewards to the attention net as part of its input.",
+)
+parser.add_argument(
+    "--stop-iters",
+    type=int,
+    default=200,
+    help="Number of iterations to train before we do inference.",
+)
+parser.add_argument(
+    "--stop-timesteps",
+    type=int,
+    default=100000,
+    help="Number of timesteps to train before we do inference.",
+)
+parser.add_argument(
+    "--stop-reward",
+    type=float,
+    default=150.0,
+    help="Reward at which we stop training before we do inference.",
+)
+parser.add_argument(
+    "--explore-during-inference",
+    action="store_true",
+    help="Whether the trained policy should use exploration during action "
+    "inference.",
+)
+parser.add_argument(
+    "--num-episodes-during-inference",
+    type=int,
+    default=10,
+    help="Number of episodes to do inference over after training.",
+)
+parser.add_argument(
+    "--use-onnx-for-inference",
+    action="store_true",
+    help="Whether to convert the loaded module to ONNX format and then perform "
+    "inference through this ONNX model.",
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.use_onnx_for_inference:
+        if args.explore_during_inference:
+            raise ValueError(
+                "Can't set `--explore-during-inference` and `--use-onnx-for-inference` together!"
+            )
+        import onnxruntime
+
+    ray.init(num_cpus=args.num_cpus or None)
+
+    config = (
+        get_trainable_cls(args.run)
+        .get_default_config()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("FrozenLake-v1")
+        # Run with tracing enabled for tf2?
+        .framework(args.framework)
+        .training(
+            model={
+                "use_attention": True,
+                "attention_num_transformer_units": 1,
+                "attention_use_n_prev_actions": args.prev_n_actions,
+                "attention_use_n_prev_rewards": args.prev_n_rewards,
+                "attention_dim": 32,
+                "attention_memory_inference": 10,
+                "attention_memory_training": 10,
+            },
+        )
+        # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+        .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0")))
+    )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+    }
+
+    print("Training policy until desired reward/timesteps/iterations. ...")
+    tuner = tune.Tuner(
+        args.run,
+        param_space=config,
+        run_config=tune.RunConfig(
+            stop=stop,
+            verbose=2,
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=1,
+                checkpoint_at_end=True,
+            ),
+        ),
+    )
+    results = tuner.fit()
+
+    print("Training completed. Restoring new Algorithm for action inference.")
+    # Get the last checkpoint from the above training run.
+    checkpoint = results.get_best_result().checkpoint
+    # Create new Algorithm and restore its state from the last checkpoint.
+    algo = Algorithm.from_checkpoint(checkpoint)
+    # Export ONNX model if relevant
+    if args.use_onnx_for_inference:
+        algo.get_policy().export_model(
+            "frozenlake_attention_model_onnx",
+            # ONNX opset version 12 required to support einsum operator.
+            # Requires ONNX >= 1.7 and ONNX runtime >= 1.3
+            onnx=12,
+        )
+
+    # Create the env to do inference in.
+    env = gym.make("FrozenLake-v1")
+    obs, info = env.reset()
+
+    # In case the model needs previous-reward/action inputs, keep track of
+    # these via these variables here (we'll have to pass them into the
+    # compute_actions methods below).
+    init_prev_a = prev_a = None
+    init_prev_r = prev_r = None
+
+    # Set attention net's initial internal state.
+    num_transformers = config["model"]["attention_num_transformer_units"]
+    memory_inference = config["model"]["attention_memory_inference"]
+    attention_dim = config["model"]["attention_dim"]
+    init_state = state = [
+        np.zeros([memory_inference, attention_dim], np.float32)
+        for _ in range(num_transformers)
+    ]
+    # Do we need prev-action/reward as part of the input?
+    if args.prev_n_actions:
+        init_prev_a = prev_a = np.array([0] * args.prev_n_actions)
+    if args.prev_n_rewards:
+        init_prev_r = prev_r = np.array([0.0] * args.prev_n_rewards)
+
+    num_episodes = 0
+    ort_session = None
+
+    while num_episodes < args.num_episodes_during_inference:
+
+        # Compute an action (`a`).
+        if args.use_onnx_for_inference:
+            # Prepare the ONNX runtime session.
+            if ort_session is None:
+                ort_session = onnxruntime.InferenceSession(
+                    "frozenlake_attention_model_onnx/model.onnx"
+                )
+            # Prepare the inputs dict.
+            seq_len = np.array([config["model"]["max_seq_len"]], dtype=np.int32)
+
+            # pre-process observation: obs is an integer.
+            # we need to convert it to a one-hot vector (FrozenLake-v1 space).
+            n = env.observation_space.n
+            obs_one_hot = np.zeros(n, dtype=np.float32)
+            obs_one_hot[obs] = 1.0
+            obs = obs_one_hot
+            # Add batch dimension.
+            obs = np.array(obs, dtype=np.float32)[np.newaxis, :]
+
+            state_ins = np.array(state, dtype=np.float32)
+
+            ort_inputs = {
+                "obs": obs,
+                "state_ins": state_ins,
+                "seq_lens": seq_len,
+            }
+
+            if init_prev_a is not None:
+                ort_inputs["prev_actions"] = prev_a.astype(np.int64)[np.newaxis, :]
+            if init_prev_r is not None:
+                ort_inputs["prev_rewards"] = prev_r.astype(np.float32)[np.newaxis, :]
+            # Run the ONNX model.
+            ort_outs = ort_session.run(
+                output_names=["output", "state_outs"],
+                input_feed=ort_inputs,
+            )
+            # Extract action and state-out from the ONNX model outputs.
+            dist_inputs = ort_outs[0][0]
+            # Exploration could be added here based on `dist_inputs`.
+            # This would require using the configured exploration strategy.
+            # Not implemented in this example.
+            a = np.argmax(dist_inputs)
+
+            state_out = [ort_outs[i + 1][0] for i in range(len(state))]
+        else:
+            a, state_out, _ = algo.compute_single_action(
+                observation=obs,
+                state=state,
+                prev_action=prev_a,
+                prev_reward=prev_r,
+                explore=args.explore_during_inference,
+                policy_id="default_policy",  # <- default value
+            )
+        # Send the computed action `a` to the env.
+        obs, reward, done, truncated, _ = env.step(a)
+        # Is the episode `done`? -> Reset.
+        if done:
+            obs, info = env.reset()
+            num_episodes += 1
+            state = init_state
+            prev_a = init_prev_a
+            prev_r = init_prev_r
+        # Episode is still ongoing -> Continue.
+        else:
+            # Append the just received state-out (most recent timestep) to the
+            # cascade (memory) of our state-ins and drop the oldest state-in.
+            state = [
+                np.concatenate([state[i], [state_out[i]]], axis=0)[1:]
+                for i in range(num_transformers)
+            ]
+            if init_prev_a is not None:
+                prev_a = np.concatenate([prev_a, [a]])[1:]
+            if init_prev_r is not None:
+                prev_r = np.concatenate([prev_r, [reward]])[1:]
+
+    algo.stop()
+
+    ray.shutdown()

rllib/examples/_old_api_stack/inference/policy_inference_after_training_with_lstm.py

@@ -0,0 +1,186 @@
+# @OldAPIStack
+"""
+Example showing how you can use your trained policy for inference
+(computing actions) in an environment.
+
+Includes options for LSTM-based models (--use-lstm), attention-net models
+(--use-attention), and plain (non-recurrent) models.
+"""
+import argparse
+import os
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.tune.registry import get_trainable_cls
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use."
+)
+parser.add_argument("--num-cpus", type=int, default=0)
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--prev-action",
+    action="store_true",
+    help="Feed most recent action to the LSTM as part of its input.",
+)
+parser.add_argument(
+    "--prev-reward",
+    action="store_true",
+    help="Feed most recent reward to the LSTM as part of its input.",
+)
+parser.add_argument(
+    "--stop-iters",
+    type=int,
+    default=2,
+    help="Number of iterations to train before we do inference.",
+)
+parser.add_argument(
+    "--stop-timesteps",
+    type=int,
+    default=100000,
+    help="Number of timesteps to train before we do inference.",
+)
+parser.add_argument(
+    "--stop-reward",
+    type=float,
+    default=0.8,
+    help="Reward at which we stop training before we do inference.",
+)
+parser.add_argument(
+    "--explore-during-inference",
+    action="store_true",
+    help="Whether the trained policy should use exploration during action "
+    "inference.",
+)
+parser.add_argument(
+    "--num-episodes-during-inference",
+    type=int,
+    default=10,
+    help="Number of episodes to do inference over after training.",
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    ray.init(num_cpus=args.num_cpus or None)
+
+    config = (
+        get_trainable_cls(args.run)
+        .get_default_config()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("FrozenLake-v1")
+        # Run with tracing enabled for tf2?
+        .framework(args.framework)
+        .training(
+            model={
+                "use_lstm": True,
+                "lstm_cell_size": 256,
+                "lstm_use_prev_action": args.prev_action,
+                "lstm_use_prev_reward": args.prev_reward,
+            },
+        )
+        # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+        .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0")))
+    )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+    }
+
+    print("Training policy until desired reward/timesteps/iterations. ...")
+    tuner = tune.Tuner(
+        args.run,
+        param_space=config,
+        run_config=tune.RunConfig(
+            stop=stop,
+            verbose=2,
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=1,
+                checkpoint_at_end=True,
+            ),
+        ),
+    )
+    results = tuner.fit()
+
+    print("Training completed. Restoring new Algorithm for action inference.")
+    # Get the last checkpoint from the above training run.
+    checkpoint = results.get_best_result().checkpoint
+    # Create new Algorithm from the last checkpoint.
+    algo = Algorithm.from_checkpoint(checkpoint)
+
+    # Create the env to do inference in.
+    env = gym.make("FrozenLake-v1")
+    obs, info = env.reset()
+
+    # In case the model needs previous-reward/action inputs, keep track of
+    # these via these variables here (we'll have to pass them into the
+    # compute_actions methods below).
+    init_prev_a = prev_a = None
+    init_prev_r = prev_r = None
+
+    # Set LSTM's initial internal state.
+    lstm_cell_size = config["model"]["lstm_cell_size"]
+    # range(2) b/c h- and c-states of the LSTM.
+    if algo.config.enable_rl_module_and_learner:
+        init_state = state = algo.get_policy().model.get_initial_state()
+    else:
+        init_state = state = [np.zeros([lstm_cell_size], np.float32) for _ in range(2)]
+    # Do we need prev-action/reward as part of the input?
+    if args.prev_action:
+        init_prev_a = prev_a = 0
+    if args.prev_reward:
+        init_prev_r = prev_r = 0.0
+
+    num_episodes = 0
+
+    while num_episodes < args.num_episodes_during_inference:
+        # Compute an action (`a`).
+        a, state_out, _ = algo.compute_single_action(
+            observation=obs,
+            state=state,
+            prev_action=prev_a,
+            prev_reward=prev_r,
+            explore=args.explore_during_inference,
+            policy_id="default_policy",  # <- default value
+        )
+        # Send the computed action `a` to the env.
+        obs, reward, done, truncated, info = env.step(a)
+        # Is the episode `done`? -> Reset.
+        if done:
+            obs, info = env.reset()
+            num_episodes += 1
+            state = init_state
+            prev_a = init_prev_a
+            prev_r = init_prev_r
+        # Episode is still ongoing -> Continue.
+        else:
+            state = state_out
+            if init_prev_a is not None:
+                prev_a = a
+            if init_prev_r is not None:
+                prev_r = reward
+
+    algo.stop()
+
+    ray.shutdown()

rllib/examples/_old_api_stack/models/action_mask_model.py

@@ -0,0 +1,126 @@
+# @OldAPIStack
+from gymnasium.spaces import Dict
+
+from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+from ray.rllib.utils.torch_utils import FLOAT_MIN
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class ActionMaskModel(TFModelV2):
+    """Model that handles simple discrete action masking.
+
+    This assumes the outputs are logits for a single Categorical action dist.
+    Getting this to work with a more complex output (e.g., if the action space
+    is a tuple of several distributions) is also possible but left as an
+    exercise to the reader.
+    """
+
+    def __init__(
+        self, obs_space, action_space, num_outputs, model_config, name, **kwargs
+    ):
+
+        orig_space = getattr(obs_space, "original_space", obs_space)
+        assert (
+            isinstance(orig_space, Dict)
+            and "action_mask" in orig_space.spaces
+            and "observations" in orig_space.spaces
+        )
+
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+
+        self.internal_model = FullyConnectedNetwork(
+            orig_space["observations"],
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_internal",
+        )
+
+        # disable action masking --> will likely lead to invalid actions
+        self.no_masking = model_config["custom_model_config"].get("no_masking", False)
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the unmasked logits.
+        logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]})
+
+        # If action masking is disabled, directly return unmasked logits
+        if self.no_masking:
+            return logits, state
+
+        # Convert action_mask into a [0.0 || -inf]-type mask.
+        inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)
+        masked_logits = logits + inf_mask
+
+        # Return masked logits.
+        return masked_logits, state
+
+    def value_function(self):
+        return self.internal_model.value_function()
+
+
+class TorchActionMaskModel(TorchModelV2, nn.Module):
+    """PyTorch version of above ActionMaskingModel."""
+
+    def __init__(
+        self,
+        obs_space,
+        action_space,
+        num_outputs,
+        model_config,
+        name,
+        **kwargs,
+    ):
+        orig_space = getattr(obs_space, "original_space", obs_space)
+        assert (
+            isinstance(orig_space, Dict)
+            and "action_mask" in orig_space.spaces
+            and "observations" in orig_space.spaces
+        )
+
+        TorchModelV2.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name, **kwargs
+        )
+        nn.Module.__init__(self)
+
+        self.internal_model = TorchFC(
+            orig_space["observations"],
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_internal",
+        )
+
+        # disable action masking --> will likely lead to invalid actions
+        self.no_masking = False
+        if "no_masking" in model_config["custom_model_config"]:
+            self.no_masking = model_config["custom_model_config"]["no_masking"]
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the unmasked logits.
+        logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]})
+
+        # If action masking is disabled, directly return unmasked logits
+        if self.no_masking:
+            return logits, state
+
+        # Convert action_mask into a [0.0 || -inf]-type mask.
+        inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN)
+        masked_logits = logits + inf_mask
+
+        # Return masked logits.
+        return masked_logits, state
+
+    def value_function(self):
+        return self.internal_model.value_function()

rllib/examples/_old_api_stack/models/autoregressive_action_dist.py

@@ -0,0 +1,149 @@
+# @OldAPIStack
+from ray.rllib.models.tf.tf_action_dist import ActionDistribution, Categorical
+from ray.rllib.models.torch.torch_action_dist import (
+    TorchCategorical,
+    TorchDistributionWrapper,
+)
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class BinaryAutoregressiveDistribution(ActionDistribution):
+    """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)"""
+
+    def deterministic_sample(self):
+        # First, sample a1.
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.deterministic_sample()
+
+        # Sample a2 conditioned on a1.
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.deterministic_sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # Return the action tuple.
+        return (a1, a2)
+
+    def sample(self):
+        # First, sample a1.
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.sample()
+
+        # Sample a2 conditioned on a1.
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # Return the action tuple.
+        return (a1, a2)
+
+    def logp(self, actions):
+        a1, a2 = actions[:, 0], actions[:, 1]
+        a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1)
+        a1_logits, a2_logits = self.model.action_model([self.inputs, a1_vec])
+        return Categorical(a1_logits).logp(a1) + Categorical(a2_logits).logp(a2)
+
+    def sampled_action_logp(self):
+        return self._action_logp
+
+    def entropy(self):
+        a1_dist = self._a1_distribution()
+        a2_dist = self._a2_distribution(a1_dist.sample())
+        return a1_dist.entropy() + a2_dist.entropy()
+
+    def kl(self, other):
+        a1_dist = self._a1_distribution()
+        a1_terms = a1_dist.kl(other._a1_distribution())
+
+        a1 = a1_dist.sample()
+        a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1))
+        return a1_terms + a2_terms
+
+    def _a1_distribution(self):
+        BATCH = tf.shape(self.inputs)[0]
+        a1_logits, _ = self.model.action_model([self.inputs, tf.zeros((BATCH, 1))])
+        a1_dist = Categorical(a1_logits)
+        return a1_dist
+
+    def _a2_distribution(self, a1):
+        a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1)
+        _, a2_logits = self.model.action_model([self.inputs, a1_vec])
+        a2_dist = Categorical(a2_logits)
+        return a2_dist
+
+    @staticmethod
+    def required_model_output_shape(action_space, model_config):
+        return 16  # controls model output feature vector size
+
+
+class TorchBinaryAutoregressiveDistribution(TorchDistributionWrapper):
+    """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)"""
+
+    def deterministic_sample(self):
+        # First, sample a1.
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.deterministic_sample()
+
+        # Sample a2 conditioned on a1.
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.deterministic_sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # Return the action tuple.
+        return (a1, a2)
+
+    def sample(self):
+        # First, sample a1.
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.sample()
+
+        # Sample a2 conditioned on a1.
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # Return the action tuple.
+        return (a1, a2)
+
+    def logp(self, actions):
+        a1, a2 = actions[:, 0], actions[:, 1]
+        a1_vec = torch.unsqueeze(a1.float(), 1)
+        a1_logits, a2_logits = self.model.action_module(self.inputs, a1_vec)
+        return TorchCategorical(a1_logits).logp(a1) + TorchCategorical(a2_logits).logp(
+            a2
+        )
+
+    def sampled_action_logp(self):
+        return self._action_logp
+
+    def entropy(self):
+        a1_dist = self._a1_distribution()
+        a2_dist = self._a2_distribution(a1_dist.sample())
+        return a1_dist.entropy() + a2_dist.entropy()
+
+    def kl(self, other):
+        a1_dist = self._a1_distribution()
+        a1_terms = a1_dist.kl(other._a1_distribution())
+
+        a1 = a1_dist.sample()
+        a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1))
+        return a1_terms + a2_terms
+
+    def _a1_distribution(self):
+        BATCH = self.inputs.shape[0]
+        zeros = torch.zeros((BATCH, 1)).to(self.inputs.device)
+        a1_logits, _ = self.model.action_module(self.inputs, zeros)
+        a1_dist = TorchCategorical(a1_logits)
+        return a1_dist
+
+    def _a2_distribution(self, a1):
+        a1_vec = torch.unsqueeze(a1.float(), 1)
+        _, a2_logits = self.model.action_module(self.inputs, a1_vec)
+        a2_dist = TorchCategorical(a2_logits)
+        return a2_dist
+
+    @staticmethod
+    def required_model_output_shape(action_space, model_config):
+        return 16  # controls model output feature vector size

rllib/examples/_old_api_stack/models/autoregressive_action_model.py

@@ -0,0 +1,161 @@
+# @OldAPIStack
+from gymnasium.spaces import Discrete, Tuple
+
+from ray.rllib.models.tf.misc import normc_initializer
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC, normc_initializer as normc_init_torch
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class AutoregressiveActionModel(TFModelV2):
+    """Implements the `.action_model` branch required above."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super(AutoregressiveActionModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+        if action_space != Tuple([Discrete(2), Discrete(2)]):
+            raise ValueError("This model only supports the [2, 2] action space")
+
+        # Inputs
+        obs_input = tf.keras.layers.Input(shape=obs_space.shape, name="obs_input")
+        a1_input = tf.keras.layers.Input(shape=(1,), name="a1_input")
+        ctx_input = tf.keras.layers.Input(shape=(num_outputs,), name="ctx_input")
+
+        # Output of the model (normally 'logits', but for an autoregressive
+        # dist this is more like a context/feature layer encoding the obs)
+        context = tf.keras.layers.Dense(
+            num_outputs,
+            name="hidden",
+            activation=tf.nn.tanh,
+            kernel_initializer=normc_initializer(1.0),
+        )(obs_input)
+
+        # V(s)
+        value_out = tf.keras.layers.Dense(
+            1,
+            name="value_out",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01),
+        )(context)
+
+        # P(a1 | obs)
+        a1_logits = tf.keras.layers.Dense(
+            2,
+            name="a1_logits",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01),
+        )(ctx_input)
+
+        # P(a2 | a1)
+        # --note: typically you'd want to implement P(a2 | a1, obs) as follows:
+        # a2_context = tf.keras.layers.Concatenate(axis=1)(
+        #     [ctx_input, a1_input])
+        a2_context = a1_input
+        a2_hidden = tf.keras.layers.Dense(
+            16,
+            name="a2_hidden",
+            activation=tf.nn.tanh,
+            kernel_initializer=normc_initializer(1.0),
+        )(a2_context)
+        a2_logits = tf.keras.layers.Dense(
+            2,
+            name="a2_logits",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01),
+        )(a2_hidden)
+
+        # Base layers
+        self.base_model = tf.keras.Model(obs_input, [context, value_out])
+        self.base_model.summary()
+
+        # Autoregressive action sampler
+        self.action_model = tf.keras.Model(
+            [ctx_input, a1_input], [a1_logits, a2_logits]
+        )
+        self.action_model.summary()
+
+    def forward(self, input_dict, state, seq_lens):
+        context, self._value_out = self.base_model(input_dict["obs"])
+        return context, state
+
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchAutoregressiveActionModel(TorchModelV2, nn.Module):
+    """PyTorch version of the AutoregressiveActionModel above."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        TorchModelV2.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name
+        )
+        nn.Module.__init__(self)
+
+        if action_space != Tuple([Discrete(2), Discrete(2)]):
+            raise ValueError("This model only supports the [2, 2] action space")
+
+        # Output of the model (normally 'logits', but for an autoregressive
+        # dist this is more like a context/feature layer encoding the obs)
+        self.context_layer = SlimFC(
+            in_size=obs_space.shape[0],
+            out_size=num_outputs,
+            initializer=normc_init_torch(1.0),
+            activation_fn=nn.Tanh,
+        )
+
+        # V(s)
+        self.value_branch = SlimFC(
+            in_size=num_outputs,
+            out_size=1,
+            initializer=normc_init_torch(0.01),
+            activation_fn=None,
+        )
+
+        # P(a1 | obs)
+        self.a1_logits = SlimFC(
+            in_size=num_outputs,
+            out_size=2,
+            activation_fn=None,
+            initializer=normc_init_torch(0.01),
+        )
+
+        class _ActionModel(nn.Module):
+            def __init__(self):
+                nn.Module.__init__(self)
+                self.a2_hidden = SlimFC(
+                    in_size=1,
+                    out_size=16,
+                    activation_fn=nn.Tanh,
+                    initializer=normc_init_torch(1.0),
+                )
+                self.a2_logits = SlimFC(
+                    in_size=16,
+                    out_size=2,
+                    activation_fn=None,
+                    initializer=normc_init_torch(0.01),
+                )
+
+            def forward(self_, ctx_input, a1_input):
+                a1_logits = self.a1_logits(ctx_input)
+                a2_logits = self_.a2_logits(self_.a2_hidden(a1_input))
+                return a1_logits, a2_logits
+
+        # P(a2 | a1)
+        # --note: typically you'd want to implement P(a2 | a1, obs) as follows:
+        # a2_context = tf.keras.layers.Concatenate(axis=1)(
+        #     [ctx_input, a1_input])
+        self.action_module = _ActionModel()
+
+        self._context = None
+
+    def forward(self, input_dict, state, seq_lens):
+        self._context = self.context_layer(input_dict["obs"])
+        return self._context, state
+
+    def value_function(self):
+        return torch.reshape(self.value_branch(self._context), [-1])

rllib/examples/_old_api_stack/models/centralized_critic_models.py

@@ -0,0 +1,182 @@
+# @OldAPIStack
+from gymnasium.spaces import Box
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class CentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized value function."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super(CentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+        # Base of the model
+        self.model = FullyConnectedNetwork(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+
+        # Central VF maps (obs, opp_obs, opp_act) -> vf_pred
+        obs = tf.keras.layers.Input(shape=(6,), name="obs")
+        opp_obs = tf.keras.layers.Input(shape=(6,), name="opp_obs")
+        opp_act = tf.keras.layers.Input(shape=(2,), name="opp_act")
+        concat_obs = tf.keras.layers.Concatenate(axis=1)([obs, opp_obs, opp_act])
+        central_vf_dense = tf.keras.layers.Dense(
+            16, activation=tf.nn.tanh, name="c_vf_dense"
+        )(concat_obs)
+        central_vf_out = tf.keras.layers.Dense(1, activation=None, name="c_vf_out")(
+            central_vf_dense
+        )
+        self.central_vf = tf.keras.Model(
+            inputs=[obs, opp_obs, opp_act], outputs=central_vf_out
+        )
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        return self.model.forward(input_dict, state, seq_lens)
+
+    def central_value_function(self, obs, opponent_obs, opponent_actions):
+        return tf.reshape(
+            self.central_vf(
+                [obs, opponent_obs, tf.one_hot(tf.cast(opponent_actions, tf.int32), 2)]
+            ),
+            [-1],
+        )
+
+    @override(ModelV2)
+    def value_function(self):
+        return self.model.value_function()  # not used
+
+
+class YetAnotherCentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized value function.
+
+    It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the
+    former of which can be used for computing actions (i.e., decentralized
+    execution), and the latter for optimization (i.e., centralized learning).
+
+    This model has two parts:
+    - An action model that looks at just 'own_obs' to compute actions
+    - A value model that also looks at the 'opponent_obs' / 'opponent_action'
+      to compute the value (it does this by using the 'obs_flat' tensor).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super(YetAnotherCentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+
+        self.action_model = FullyConnectedNetwork(
+            Box(low=0, high=1, shape=(6,)),  # one-hot encoded Discrete(6)
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_action",
+        )
+
+        self.value_model = FullyConnectedNetwork(
+            obs_space, action_space, 1, model_config, name + "_vf"
+        )
+
+    def forward(self, input_dict, state, seq_lens):
+        self._value_out, _ = self.value_model(
+            {"obs": input_dict["obs_flat"]}, state, seq_lens
+        )
+        return self.action_model({"obs": input_dict["obs"]["own_obs"]}, state, seq_lens)
+
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchCentralizedCriticModel(TorchModelV2, nn.Module):
+    """Multi-agent model that implements a centralized VF."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        TorchModelV2.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name
+        )
+        nn.Module.__init__(self)
+
+        # Base of the model
+        self.model = TorchFC(obs_space, action_space, num_outputs, model_config, name)
+
+        # Central VF maps (obs, opp_obs, opp_act) -> vf_pred
+        input_size = 6 + 6 + 2  # obs + opp_obs + opp_act
+        self.central_vf = nn.Sequential(
+            SlimFC(input_size, 16, activation_fn=nn.Tanh),
+            SlimFC(16, 1),
+        )
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        model_out, _ = self.model(input_dict, state, seq_lens)
+        return model_out, []
+
+    def central_value_function(self, obs, opponent_obs, opponent_actions):
+        input_ = torch.cat(
+            [
+                obs,
+                opponent_obs,
+                torch.nn.functional.one_hot(opponent_actions.long(), 2).float(),
+            ],
+            1,
+        )
+        return torch.reshape(self.central_vf(input_), [-1])
+
+    @override(ModelV2)
+    def value_function(self):
+        return self.model.value_function()  # not used
+
+
+class YetAnotherTorchCentralizedCriticModel(TorchModelV2, nn.Module):
+    """Multi-agent model that implements a centralized value function.
+
+    It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the
+    former of which can be used for computing actions (i.e., decentralized
+    execution), and the latter for optimization (i.e., centralized learning).
+
+    This model has two parts:
+    - An action model that looks at just 'own_obs' to compute actions
+    - A value model that also looks at the 'opponent_obs' / 'opponent_action'
+      to compute the value (it does this by using the 'obs_flat' tensor).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        TorchModelV2.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name
+        )
+        nn.Module.__init__(self)
+
+        self.action_model = TorchFC(
+            Box(low=0, high=1, shape=(6,)),  # one-hot encoded Discrete(6)
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_action",
+        )
+
+        self.value_model = TorchFC(
+            obs_space, action_space, 1, model_config, name + "_vf"
+        )
+        self._model_in = None
+
+    def forward(self, input_dict, state, seq_lens):
+        # Store model-input for possible `value_function()` call.
+        self._model_in = [input_dict["obs_flat"], state, seq_lens]
+        return self.action_model({"obs": input_dict["obs"]["own_obs"]}, state, seq_lens)
+
+    def value_function(self):
+        value_out, _ = self.value_model(
+            {"obs": self._model_in[0]}, self._model_in[1], self._model_in[2]
+        )
+        return torch.reshape(value_out, [-1])

rllib/examples/_old_api_stack/models/custom_loss_model.py

@@ -0,0 +1,137 @@
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions
+from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
+from ray.rllib.models.tf.tf_action_dist import Categorical
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.offline import JsonReader
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class CustomLossModel(TFModelV2):
+    """Custom model that adds an imitation loss on top of the policy loss."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+
+        self.fcnet = FullyConnectedNetwork(
+            self.obs_space, self.action_space, num_outputs, model_config, name="fcnet"
+        )
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        # Delegate to our FCNet.
+        return self.fcnet(input_dict, state, seq_lens)
+
+    @override(ModelV2)
+    def value_function(self):
+        # Delegate to our FCNet.
+        return self.fcnet.value_function()
+
+    @override(ModelV2)
+    def custom_loss(self, policy_loss, loss_inputs):
+        # Create a new input reader per worker.
+        reader = JsonReader(self.model_config["custom_model_config"]["input_files"])
+        input_ops = reader.tf_input_ops()
+
+        # Define a secondary loss by building a graph copy with weight sharing.
+        obs = restore_original_dimensions(
+            tf.cast(input_ops["obs"], tf.float32), self.obs_space
+        )
+        logits, _ = self.forward({"obs": obs}, [], None)
+
+        # Compute the IL loss.
+        action_dist = Categorical(logits, self.model_config)
+        self.policy_loss = policy_loss
+        self.imitation_loss = tf.reduce_mean(-action_dist.logp(input_ops["actions"]))
+        return policy_loss + 10 * self.imitation_loss
+
+    def metrics(self):
+        return {
+            "policy_loss": self.policy_loss,
+            "imitation_loss": self.imitation_loss,
+        }
+
+
+class TorchCustomLossModel(TorchModelV2, nn.Module):
+    """PyTorch version of the CustomLossModel above."""
+
+    def __init__(
+        self, obs_space, action_space, num_outputs, model_config, name, input_files
+    ):
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+        nn.Module.__init__(self)
+
+        self.input_files = input_files
+        # Create a new input reader per worker.
+        self.reader = JsonReader(self.input_files)
+        self.fcnet = TorchFC(
+            self.obs_space, self.action_space, num_outputs, model_config, name="fcnet"
+        )
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        # Delegate to our FCNet.
+        return self.fcnet(input_dict, state, seq_lens)
+
+    @override(ModelV2)
+    def value_function(self):
+        # Delegate to our FCNet.
+        return self.fcnet.value_function()
+
+    @override(ModelV2)
+    def custom_loss(self, policy_loss, loss_inputs):
+        """Calculates a custom loss on top of the given policy_loss(es).
+
+        Args:
+            policy_loss (List[TensorType]): The list of already calculated
+                policy losses (as many as there are optimizers).
+            loss_inputs: Struct of np.ndarrays holding the
+                entire train batch.
+
+        Returns:
+            List[TensorType]: The altered list of policy losses. In case the
+                custom loss should have its own optimizer, make sure the
+                returned list is one larger than the incoming policy_loss list.
+                In case you simply want to mix in the custom loss into the
+                already calculated policy losses, return a list of altered
+                policy losses (as done in this example below).
+        """
+        # Get the next batch from our input files.
+        batch = self.reader.next()
+
+        # Define a secondary loss by building a graph copy with weight sharing.
+        obs = restore_original_dimensions(
+            torch.from_numpy(batch["obs"]).float().to(policy_loss[0].device),
+            self.obs_space,
+            tensorlib="torch",
+        )
+        logits, _ = self.forward({"obs": obs}, [], None)
+
+        # Compute the IL loss.
+        action_dist = TorchCategorical(logits, self.model_config)
+        imitation_loss = torch.mean(
+            -action_dist.logp(
+                torch.from_numpy(batch["actions"]).to(policy_loss[0].device)
+            )
+        )
+        self.imitation_loss_metric = imitation_loss.item()
+        self.policy_loss_metric = np.mean([loss.item() for loss in policy_loss])
+
+        # Add the imitation loss to each already calculated policy loss term.
+        # Alternatively (if custom loss has its own optimizer):
+        # return policy_loss + [10 * self.imitation_loss]
+        return [loss_ + 10 * imitation_loss for loss_ in policy_loss]
+
+    def metrics(self):
+        return {
+            "policy_loss": self.policy_loss_metric,
+            "imitation_loss": self.imitation_loss_metric,
+        }

rllib/examples/_old_api_stack/models/fast_model.py

@@ -0,0 +1,80 @@
+# @OldAPIStack
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class FastModel(TFModelV2):
+    """An example for a non-Keras ModelV2 in tf that learns a single weight.
+
+    Defines all network architecture in `forward` (not `__init__` as it's
+    usually done for Keras-style TFModelV2s).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+        # Have we registered our vars yet (see `forward`)?
+        self._registered = False
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        with tf1.variable_scope("model", reuse=tf1.AUTO_REUSE):
+            bias = tf1.get_variable(
+                dtype=tf.float32,
+                name="bias",
+                initializer=tf.keras.initializers.Zeros(),
+                shape=(),
+            )
+            output = bias + tf.zeros([tf.shape(input_dict["obs"])[0], self.num_outputs])
+            self._value_out = tf.reduce_mean(output, -1)  # fake value
+
+        if not self._registered:
+            self.register_variables(
+                tf1.get_collection(
+                    tf1.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+"
+                )
+            )
+            self._registered = True
+
+        return output, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchFastModel(TorchModelV2, nn.Module):
+    """Torch version of FastModel (tf)."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        TorchModelV2.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name
+        )
+        nn.Module.__init__(self)
+
+        self.bias = nn.Parameter(
+            torch.tensor([0.0], dtype=torch.float32, requires_grad=True)
+        )
+
+        # Only needed to give some params to the optimizer (even though,
+        # they are never used anywhere).
+        self.dummy_layer = SlimFC(1, 1)
+        self._output = None
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        self._output = self.bias + torch.zeros(
+            size=(input_dict["obs"].shape[0], self.num_outputs)
+        ).to(self.bias.device)
+        return self._output, []
+
+    @override(ModelV2)
+    def value_function(self):
+        assert self._output is not None, "must call forward first!"
+        return torch.reshape(torch.mean(self._output, -1), [-1])

rllib/examples/_old_api_stack/models/neural_computer.py

@@ -0,0 +1,248 @@
+# @OldAPIStack
+from collections import OrderedDict
+from typing import Dict, List, Tuple, Union
+
+import gymnasium as gym
+
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModelConfigDict, TensorType
+
+try:
+    from dnc import DNC
+except ModuleNotFoundError:
+    print("dnc module not found. Did you forget to 'pip install dnc'?")
+    raise
+
+torch, nn = try_import_torch()
+
+
+class DNCMemory(TorchModelV2, nn.Module):
+    """Differentiable Neural Computer wrapper around ixaxaar's DNC implementation,
+    see https://github.com/ixaxaar/pytorch-dnc"""
+
+    DEFAULT_CONFIG = {
+        "dnc_model": DNC,
+        # Number of controller hidden layers
+        "num_hidden_layers": 1,
+        # Number of weights per controller hidden layer
+        "hidden_size": 64,
+        # Number of LSTM units
+        "num_layers": 1,
+        # Number of read heads, i.e. how many addrs are read at once
+        "read_heads": 4,
+        # Number of memory cells in the controller
+        "nr_cells": 32,
+        # Size of each cell
+        "cell_size": 16,
+        # LSTM activation function
+        "nonlinearity": "tanh",
+        # Observation goes through this torch.nn.Module before
+        # feeding to the DNC
+        "preprocessor": torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.Tanh()),
+        # Input size to the preprocessor
+        "preprocessor_input_size": 64,
+        # The output size of the preprocessor
+        # and the input size of the dnc
+        "preprocessor_output_size": 64,
+    }
+
+    MEMORY_KEYS = [
+        "memory",
+        "link_matrix",
+        "precedence",
+        "read_weights",
+        "write_weights",
+        "usage_vector",
+    ]
+
+    def __init__(
+        self,
+        obs_space: gym.spaces.Space,
+        action_space: gym.spaces.Space,
+        num_outputs: int,
+        model_config: ModelConfigDict,
+        name: str,
+        **custom_model_kwargs,
+    ):
+        nn.Module.__init__(self)
+        super(DNCMemory, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+        self.num_outputs = num_outputs
+        self.obs_dim = gym.spaces.utils.flatdim(obs_space)
+        self.act_dim = gym.spaces.utils.flatdim(action_space)
+
+        self.cfg = dict(self.DEFAULT_CONFIG, **custom_model_kwargs)
+        assert (
+            self.cfg["num_layers"] == 1
+        ), "num_layers != 1 has not been implemented yet"
+        self.cur_val = None
+
+        self.preprocessor = torch.nn.Sequential(
+            torch.nn.Linear(self.obs_dim, self.cfg["preprocessor_input_size"]),
+            self.cfg["preprocessor"],
+        )
+
+        self.logit_branch = SlimFC(
+            in_size=self.cfg["hidden_size"],
+            out_size=self.num_outputs,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_,
+        )
+
+        self.value_branch = SlimFC(
+            in_size=self.cfg["hidden_size"],
+            out_size=1,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_,
+        )
+
+        self.dnc: Union[None, DNC] = None
+
+    def get_initial_state(self) -> List[TensorType]:
+        ctrl_hidden = [
+            torch.zeros(self.cfg["num_hidden_layers"], self.cfg["hidden_size"]),
+            torch.zeros(self.cfg["num_hidden_layers"], self.cfg["hidden_size"]),
+        ]
+        m = self.cfg["nr_cells"]
+        r = self.cfg["read_heads"]
+        w = self.cfg["cell_size"]
+        memory = [
+            torch.zeros(m, w),  # memory
+            torch.zeros(1, m, m),  # link_matrix
+            torch.zeros(1, m),  # precedence
+            torch.zeros(r, m),  # read_weights
+            torch.zeros(1, m),  # write_weights
+            torch.zeros(m),  # usage_vector
+        ]
+
+        read_vecs = torch.zeros(w * r)
+
+        state = [*ctrl_hidden, read_vecs, *memory]
+        assert len(state) == 9
+        return state
+
+    def value_function(self) -> TensorType:
+        assert self.cur_val is not None, "must call forward() first"
+        return self.cur_val
+
+    def unpack_state(
+        self,
+        state: List[TensorType],
+    ) -> Tuple[List[Tuple[TensorType, TensorType]], Dict[str, TensorType], TensorType]:
+        """Given a list of tensors, reformat for self.dnc input"""
+        assert len(state) == 9, "Failed to verify unpacked state"
+        ctrl_hidden: List[Tuple[TensorType, TensorType]] = [
+            (
+                state[0].permute(1, 0, 2).contiguous(),
+                state[1].permute(1, 0, 2).contiguous(),
+            )
+        ]
+        read_vecs: TensorType = state[2]
+        memory: List[TensorType] = state[3:]
+        memory_dict: OrderedDict[str, TensorType] = OrderedDict(
+            zip(self.MEMORY_KEYS, memory)
+        )
+
+        return ctrl_hidden, memory_dict, read_vecs
+
+    def pack_state(
+        self,
+        ctrl_hidden: List[Tuple[TensorType, TensorType]],
+        memory_dict: Dict[str, TensorType],
+        read_vecs: TensorType,
+    ) -> List[TensorType]:
+        """Given the dnc output, pack it into a list of tensors
+        for rllib state. Order is ctrl_hidden, read_vecs, memory_dict"""
+        state = []
+        ctrl_hidden = [
+            ctrl_hidden[0][0].permute(1, 0, 2),
+            ctrl_hidden[0][1].permute(1, 0, 2),
+        ]
+        state += ctrl_hidden
+        assert len(state) == 2, "Failed to verify packed state"
+        state.append(read_vecs)
+        assert len(state) == 3, "Failed to verify packed state"
+        state += memory_dict.values()
+        assert len(state) == 9, "Failed to verify packed state"
+        return state
+
+    def validate_unpack(self, dnc_output, unpacked_state):
+        """Ensure the unpacked state shapes match the DNC output"""
+        s_ctrl_hidden, s_memory_dict, s_read_vecs = unpacked_state
+        ctrl_hidden, memory_dict, read_vecs = dnc_output
+
+        for i in range(len(ctrl_hidden)):
+            for j in range(len(ctrl_hidden[i])):
+                assert s_ctrl_hidden[i][j].shape == ctrl_hidden[i][j].shape, (
+                    "Controller state mismatch: got "
+                    f"{s_ctrl_hidden[i][j].shape} should be "
+                    f"{ctrl_hidden[i][j].shape}"
+                )
+
+        for k in memory_dict:
+            assert s_memory_dict[k].shape == memory_dict[k].shape, (
+                "Memory state mismatch at key "
+                f"{k}: got {s_memory_dict[k].shape} should be "
+                f"{memory_dict[k].shape}"
+            )
+
+        assert s_read_vecs.shape == read_vecs.shape, (
+            "Read state mismatch: got "
+            f"{s_read_vecs.shape} should be "
+            f"{read_vecs.shape}"
+        )
+
+    def build_dnc(self, device_idx: Union[int, None]) -> None:
+        self.dnc = self.cfg["dnc_model"](
+            input_size=self.cfg["preprocessor_output_size"],
+            hidden_size=self.cfg["hidden_size"],
+            num_layers=self.cfg["num_layers"],
+            num_hidden_layers=self.cfg["num_hidden_layers"],
+            read_heads=self.cfg["read_heads"],
+            cell_size=self.cfg["cell_size"],
+            nr_cells=self.cfg["nr_cells"],
+            nonlinearity=self.cfg["nonlinearity"],
+            gpu_id=device_idx,
+        )
+
+    def forward(
+        self,
+        input_dict: Dict[str, TensorType],
+        state: List[TensorType],
+        seq_lens: TensorType,
+    ) -> Tuple[TensorType, List[TensorType]]:
+
+        flat = input_dict["obs_flat"]
+        # Batch and Time
+        # Forward expects outputs as [B, T, logits]
+        B = len(seq_lens)
+        T = flat.shape[0] // B
+
+        # Deconstruct batch into batch and time dimensions: [B, T, feats]
+        flat = torch.reshape(flat, [-1, T] + list(flat.shape[1:]))
+
+        # First run
+        if self.dnc is None:
+            gpu_id = flat.device.index if flat.device.index is not None else -1
+            self.build_dnc(gpu_id)
+            hidden = (None, None, None)
+
+        else:
+            hidden = self.unpack_state(state)  # type: ignore
+
+        # Run thru preprocessor before DNC
+        z = self.preprocessor(flat.reshape(B * T, self.obs_dim))
+        z = z.reshape(B, T, self.cfg["preprocessor_output_size"])
+        output, hidden = self.dnc(z, hidden)
+        packed_state = self.pack_state(*hidden)
+
+        # Compute action/value from output
+        logits = self.logit_branch(output.view(B * T, -1))
+        values = self.value_branch(output.view(B * T, -1))
+
+        self.cur_val = values.squeeze(1)
+
+        return logits, packed_state

rllib/examples/_old_api_stack/models/parametric_actions_model.py

@@ -0,0 +1,201 @@
+# @OldAPIStack
+from gymnasium.spaces import Box
+
+from ray.rllib.algorithms.dqn.distributional_q_tf_model import DistributionalQTFModel
+from ray.rllib.algorithms.dqn.dqn_torch_model import DQNTorchModel
+from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+from ray.rllib.utils.torch_utils import FLOAT_MAX, FLOAT_MIN
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class ParametricActionsModel(DistributionalQTFModel):
+    """Parametric action model that handles the dot product and masking.
+
+    This assumes the outputs are logits for a single Categorical action dist.
+    Getting this to work with a more complex output (e.g., if the action space
+    is a tuple of several distributions) is also possible but left as an
+    exercise to the reader.
+    """
+
+    def __init__(
+        self,
+        obs_space,
+        action_space,
+        num_outputs,
+        model_config,
+        name,
+        true_obs_shape=(4,),
+        action_embed_size=2,
+        **kw
+    ):
+        super(ParametricActionsModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name, **kw
+        )
+        self.action_embed_model = FullyConnectedNetwork(
+            Box(-1, 1, shape=true_obs_shape),
+            action_space,
+            action_embed_size,
+            model_config,
+            name + "_action_embed",
+        )
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        avail_actions = input_dict["obs"]["avail_actions"]
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the predicted action embedding
+        action_embed, _ = self.action_embed_model({"obs": input_dict["obs"]["cart"]})
+
+        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
+        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
+        intent_vector = tf.expand_dims(action_embed, 1)
+
+        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
+        action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2)
+
+        # Mask out invalid actions (use tf.float32.min for stability)
+        inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)
+        return action_logits + inf_mask, state
+
+    def value_function(self):
+        return self.action_embed_model.value_function()
+
+
+class TorchParametricActionsModel(DQNTorchModel):
+    """PyTorch version of above ParametricActionsModel."""
+
+    def __init__(
+        self,
+        obs_space,
+        action_space,
+        num_outputs,
+        model_config,
+        name,
+        true_obs_shape=(4,),
+        action_embed_size=2,
+        **kw
+    ):
+        DQNTorchModel.__init__(
+            self, obs_space, action_space, num_outputs, model_config, name, **kw
+        )
+
+        self.action_embed_model = TorchFC(
+            Box(-1, 1, shape=true_obs_shape),
+            action_space,
+            action_embed_size,
+            model_config,
+            name + "_action_embed",
+        )
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        avail_actions = input_dict["obs"]["avail_actions"]
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the predicted action embedding
+        action_embed, _ = self.action_embed_model({"obs": input_dict["obs"]["cart"]})
+
+        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
+        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
+        intent_vector = torch.unsqueeze(action_embed, 1)
+
+        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
+        action_logits = torch.sum(avail_actions * intent_vector, dim=2)
+
+        # Mask out invalid actions (use -inf to tag invalid).
+        # These are then recognized by the EpsilonGreedy exploration component
+        # as invalid actions that are not to be chosen.
+        inf_mask = torch.clamp(torch.log(action_mask), FLOAT_MIN, FLOAT_MAX)
+
+        return action_logits + inf_mask, state
+
+    def value_function(self):
+        return self.action_embed_model.value_function()
+
+
+class ParametricActionsModelThatLearnsEmbeddings(DistributionalQTFModel):
+    """Same as the above ParametricActionsModel.
+
+    However, this version also learns the action embeddings.
+    """
+
+    def __init__(
+        self,
+        obs_space,
+        action_space,
+        num_outputs,
+        model_config,
+        name,
+        true_obs_shape=(4,),
+        action_embed_size=2,
+        **kw
+    ):
+        super(ParametricActionsModelThatLearnsEmbeddings, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name, **kw
+        )
+
+        action_ids_shifted = tf.constant(
+            list(range(1, num_outputs + 1)), dtype=tf.float32
+        )
+
+        obs_cart = tf.keras.layers.Input(shape=true_obs_shape, name="obs_cart")
+        valid_avail_actions_mask = tf.keras.layers.Input(
+            shape=(num_outputs,), name="valid_avail_actions_mask"
+        )
+
+        self.pred_action_embed_model = FullyConnectedNetwork(
+            Box(-1, 1, shape=true_obs_shape),
+            action_space,
+            action_embed_size,
+            model_config,
+            name + "_pred_action_embed",
+        )
+
+        # Compute the predicted action embedding
+        pred_action_embed, _ = self.pred_action_embed_model({"obs": obs_cart})
+        _value_out = self.pred_action_embed_model.value_function()
+
+        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
+        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
+        intent_vector = tf.expand_dims(pred_action_embed, 1)
+
+        valid_avail_actions = action_ids_shifted * valid_avail_actions_mask
+        # Embedding for valid available actions which will be learned.
+        # Embedding vector for 0 is an invalid embedding (a "dummy embedding").
+        valid_avail_actions_embed = tf.keras.layers.Embedding(
+            input_dim=num_outputs + 1,
+            output_dim=action_embed_size,
+            name="action_embed_matrix",
+        )(valid_avail_actions)
+
+        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
+        action_logits = tf.reduce_sum(valid_avail_actions_embed * intent_vector, axis=2)
+
+        # Mask out invalid actions (use tf.float32.min for stability)
+        inf_mask = tf.maximum(tf.math.log(valid_avail_actions_mask), tf.float32.min)
+
+        action_logits = action_logits + inf_mask
+
+        self.param_actions_model = tf.keras.Model(
+            inputs=[obs_cart, valid_avail_actions_mask],
+            outputs=[action_logits, _value_out],
+        )
+        self.param_actions_model.summary()
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions mask tensor from the observation.
+        valid_avail_actions_mask = input_dict["obs"]["valid_avail_actions_mask"]
+
+        action_logits, self._value_out = self.param_actions_model(
+            [input_dict["obs"]["cart"], valid_avail_actions_mask]
+        )
+
+        return action_logits, state
+
+    def value_function(self):
+        return self._value_out

rllib/examples/_old_api_stack/models/shared_weights_model.py

@@ -0,0 +1,206 @@
+# @OldAPIStack
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+TF2_GLOBAL_SHARED_LAYER = None
+
+
+class TF2SharedWeightsModel(TFModelV2):
+    """Example of weight sharing between two different TFModelV2s.
+
+    NOTE: This will only work for tf2.x. When running with config.framework=tf,
+    use SharedWeightsModel1 and SharedWeightsModel2 below, instead!
+
+    The shared (single) layer is simply defined outside of the two Models,
+    then used by both Models in their forward pass.
+    """
+
+    def __init__(
+        self, observation_space, action_space, num_outputs, model_config, name
+    ):
+        super().__init__(
+            observation_space, action_space, num_outputs, model_config, name
+        )
+
+        global TF2_GLOBAL_SHARED_LAYER
+        # The global, shared layer to be used by both models.
+        if TF2_GLOBAL_SHARED_LAYER is None:
+            TF2_GLOBAL_SHARED_LAYER = tf.keras.layers.Dense(
+                units=64, activation=tf.nn.relu, name="fc1"
+            )
+
+        inputs = tf.keras.layers.Input(observation_space.shape)
+        last_layer = TF2_GLOBAL_SHARED_LAYER(inputs)
+        output = tf.keras.layers.Dense(
+            units=num_outputs, activation=None, name="fc_out"
+        )(last_layer)
+        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
+            last_layer
+        )
+        self.base_model = tf.keras.models.Model(inputs, [output, vf])
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class SharedWeightsModel1(TFModelV2):
+    """Example of weight sharing between two different TFModelV2s.
+
+    NOTE: This will only work for tf1 (static graph). When running with
+    config.framework_str=tf2, use TF2SharedWeightsModel, instead!
+
+    Here, we share the variables defined in the 'shared' variable scope
+    by entering it explicitly with tf1.AUTO_REUSE. This creates the
+    variables for the 'fc1' layer in a global scope called 'shared'
+    (outside of the Policy's normal variable scope).
+    """
+
+    def __init__(
+        self, observation_space, action_space, num_outputs, model_config, name
+    ):
+        super().__init__(
+            observation_space, action_space, num_outputs, model_config, name
+        )
+
+        inputs = tf.keras.layers.Input(observation_space.shape)
+        with tf1.variable_scope(
+            tf1.VariableScope(tf1.AUTO_REUSE, "shared"),
+            reuse=tf1.AUTO_REUSE,
+            auxiliary_name_scope=False,
+        ):
+            last_layer = tf.keras.layers.Dense(
+                units=64, activation=tf.nn.relu, name="fc1"
+            )(inputs)
+        output = tf.keras.layers.Dense(
+            units=num_outputs, activation=None, name="fc_out"
+        )(last_layer)
+        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
+            last_layer
+        )
+        self.base_model = tf.keras.models.Model(inputs, [output, vf])
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class SharedWeightsModel2(TFModelV2):
+    """The "other" TFModelV2 using the same shared space as the one above."""
+
+    def __init__(
+        self, observation_space, action_space, num_outputs, model_config, name
+    ):
+        super().__init__(
+            observation_space, action_space, num_outputs, model_config, name
+        )
+
+        inputs = tf.keras.layers.Input(observation_space.shape)
+
+        # Weights shared with SharedWeightsModel1.
+        with tf1.variable_scope(
+            tf1.VariableScope(tf1.AUTO_REUSE, "shared"),
+            reuse=tf1.AUTO_REUSE,
+            auxiliary_name_scope=False,
+        ):
+            last_layer = tf.keras.layers.Dense(
+                units=64, activation=tf.nn.relu, name="fc1"
+            )(inputs)
+        output = tf.keras.layers.Dense(
+            units=num_outputs, activation=None, name="fc_out"
+        )(last_layer)
+        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
+            last_layer
+        )
+        self.base_model = tf.keras.models.Model(inputs, [output, vf])
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+TORCH_GLOBAL_SHARED_LAYER = None
+if torch:
+    # The global, shared layer to be used by both models.
+    TORCH_GLOBAL_SHARED_LAYER = SlimFC(
+        64,
+        64,
+        activation_fn=nn.ReLU,
+        initializer=torch.nn.init.xavier_uniform_,
+    )
+
+
+class TorchSharedWeightsModel(TorchModelV2, nn.Module):
+    """Example of weight sharing between two different TorchModelV2s.
+
+    The shared (single) layer is simply defined outside of the two Models,
+    then used by both Models in their forward pass.
+    """
+
+    def __init__(
+        self, observation_space, action_space, num_outputs, model_config, name
+    ):
+        TorchModelV2.__init__(
+            self, observation_space, action_space, num_outputs, model_config, name
+        )
+        nn.Module.__init__(self)
+
+        # Non-shared initial layer.
+        self.first_layer = SlimFC(
+            int(np.prod(observation_space.shape)),
+            64,
+            activation_fn=nn.ReLU,
+            initializer=torch.nn.init.xavier_uniform_,
+        )
+
+        # Non-shared final layer.
+        self.last_layer = SlimFC(
+            64,
+            self.num_outputs,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_,
+        )
+        self.vf = SlimFC(
+            64,
+            1,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_,
+        )
+        self._global_shared_layer = TORCH_GLOBAL_SHARED_LAYER
+        self._output = None
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out = self.first_layer(input_dict["obs"])
+        self._output = self._global_shared_layer(out)
+        model_out = self.last_layer(self._output)
+        return model_out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        assert self._output is not None, "must call forward first!"
+        return torch.reshape(self.vf(self._output), [-1])

rllib/examples/_old_api_stack/models/simple_rpg_model.py

@@ -0,0 +1,65 @@
+# @OldAPIStack
+from ray.rllib.models.tf.fcnet import FullyConnectedNetwork as TFFCNet
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCNet
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf1, tf, tfv = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class CustomTorchRPGModel(TorchModelV2, nn.Module):
+    """Example of interpreting repeated observations."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+        nn.Module.__init__(self)
+        self.model = TorchFCNet(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+
+    def forward(self, input_dict, state, seq_lens):
+        # The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS:
+        # {
+        #   'items', <torch.Tensor shape=(?, M, N, 5)>,
+        #   'location', <torch.Tensor shape=(?, M, 2)>,
+        #   'status', <torch.Tensor shape=(?, M, 10)>,
+        # }
+        print("The unpacked input tensors:", input_dict["obs"])
+        print()
+        print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim())
+        print()
+        print("Fully unbatched", input_dict["obs"].unbatch_all())
+        print()
+        return self.model.forward(input_dict, state, seq_lens)
+
+    def value_function(self):
+        return self.model.value_function()
+
+
+class CustomTFRPGModel(TFModelV2):
+    """Example of interpreting repeated observations."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
+        super().__init__(obs_space, action_space, num_outputs, model_config, name)
+        self.model = TFFCNet(obs_space, action_space, num_outputs, model_config, name)
+
+    def forward(self, input_dict, state, seq_lens):
+        # The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS:
+        # {
+        #   'items', <tf.Tensor shape=(?, M, N, 5)>,
+        #   'location', <tf.Tensor shape=(?, M, 2)>,
+        #   'status', <tf.Tensor shape=(?, M, 10)>,
+        # }
+        print("The unpacked input tensors:", input_dict["obs"])
+        print()
+        print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim())
+        print()
+        if tf.executing_eagerly():
+            print("Fully unbatched", input_dict["obs"].unbatch_all())
+            print()
+        return self.model.forward(input_dict, state, seq_lens)
+
+    def value_function(self):
+        return self.model.value_function()

rllib/examples/_old_api_stack/offline_rl/custom_input_api.py

@@ -0,0 +1,134 @@
+# @OldAPIStack
+
+"""Example of creating a custom input API
+
+Custom input apis are useful when your data source is in a custom format or
+when it is necessary to use an external data loading mechanism.
+In this example, we train an rl agent on user specified input data.
+Instead of using the built in JsonReader, we will create our own custom input
+api, and show how to pass config arguments to it.
+
+To train CQL on the pendulum environment:
+$ python custom_input_api.py --input-files=../offline/tests/data/pendulum/enormous.zip
+"""
+
+import argparse
+import os
+
+import ray
+from ray import tune
+from ray.rllib.offline import InputReader, IOContext, JsonReader, ShuffledInput
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.tune.registry import get_trainable_cls, register_input
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--run", type=str, default="CQL", help="The RLlib-registered algorithm to use."
+)
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument("--stop-iters", type=int, default=100)
+parser.add_argument(
+    "--input-files",
+    type=str,
+    default=os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "../../offline/tests/data/pendulum/small.json",
+    ),
+)
+
+
+class CustomJsonReader(JsonReader):
+    """
+    Example custom InputReader implementation (extended from JsonReader).
+
+    This gets wrapped in ShuffledInput to comply with offline rl algorithms.
+    """
+
+    def __init__(self, ioctx: IOContext):
+        """
+        The constructor must take an IOContext to be used in the input config.
+        Args:
+            ioctx: use this to access the `input_config` arguments.
+        """
+        super().__init__(ioctx.input_config["input_files"], ioctx)
+
+
+def input_creator(ioctx: IOContext) -> InputReader:
+    """
+    The input creator method can be used in the input registry or set as the
+    config["input"] parameter.
+
+    Args:
+        ioctx: use this to access the `input_config` arguments.
+
+    Returns:
+        instance of ShuffledInput to work with some offline rl algorithms
+    """
+    return ShuffledInput(CustomJsonReader(ioctx))
+
+
+if __name__ == "__main__":
+    ray.init()
+    args = parser.parse_args()
+
+    # make absolute path because relative path looks in result directory
+    args.input_files = os.path.abspath(args.input_files)
+
+    # we register our custom input creator with this convenient function
+    register_input("custom_input", input_creator)
+
+    # Config modified from rllib/examples/algorithms/cql/pendulum-cql.yaml
+    default_config = get_trainable_cls(args.run).get_default_config()
+    config = (
+        default_config.environment("Pendulum-v1", clip_actions=True)
+        .framework(args.framework)
+        .offline_data(
+            # We can either use the tune registry ...
+            input_="custom_input",
+            # ... full classpath
+            # input_: "ray.rllib.examples.offline_rl.custom_input_api.CustomJsonReader"
+            # ... or a direct function to connect our input api.
+            # input_: input_creator
+            input_config={"input_files": args.input_files},  # <- passed to IOContext
+            actions_in_input_normalized=True,
+        )
+        .training(train_batch_size=2000)
+        .evaluation(
+            evaluation_interval=1,
+            evaluation_num_env_runners=2,
+            evaluation_duration=10,
+            evaluation_parallel_to_training=True,
+            evaluation_config=default_config.overrides(
+                input_="sampler",
+                explore=False,
+            ),
+        )
+        .reporting(metrics_num_episodes_for_smoothing=5)
+    )
+
+    if args.run == "CQL":
+        config.training(
+            twin_q=True,
+            num_steps_sampled_before_learning_starts=0,
+            bc_iters=100,
+        )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -600,
+    }
+
+    tuner = tune.Tuner(
+        args.run, param_space=config, run_config=tune.RunConfig(stop=stop, verbose=1)
+    )
+    tuner.fit()

rllib/examples/_old_api_stack/offline_rl/offline_rl.py

@@ -0,0 +1,168 @@
+# @OldAPIStack
+
+"""Example on how to use CQL to learn from an offline JSON file.
+
+Important node: Make sure that your offline data file contains only
+a single timestep per line to mimic the way SAC pulls samples from
+the buffer.
+
+Generate the offline json file by running an SAC algo until it reaches expert
+level on your command line. For example:
+$ cd ray
+$ rllib train -f rllib/examples/algorithms/sac/pendulum-sac.yaml --no-ray-ui
+
+Also make sure that in the above SAC yaml file (pendulum-sac.yaml),
+you specify an additional "output" key with any path on your local
+file system. In that path, the offline json files will be written to.
+
+Use the generated file(s) as "input" in the CQL config below
+(`config["input"] = [list of your json files]`), then run this script.
+"""
+
+import argparse
+
+import numpy as np
+
+from ray.rllib.algorithms import cql as cql
+from ray.rllib.execution.rollout_ops import (
+    synchronous_parallel_sample,
+)
+from ray.rllib.policy.sample_batch import convert_ma_batch_to_sample_batch
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+
+torch, _ = try_import_torch()
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--as-test",
+    action="store_true",
+    help="Whether this script should be run as a test: --stop-reward must "
+    "be achieved within --stop-timesteps AND --stop-iters.",
+)
+parser.add_argument(
+    "--stop-iters", type=int, default=5, help="Number of iterations to train."
+)
+parser.add_argument(
+    "--stop-reward", type=float, default=50.0, help="Reward at which we stop training."
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    # See rllib/examples/algorithms/cql/pendulum-cql.yaml for comparison.
+    config = (
+        cql.CQLConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .framework(framework="torch")
+        .env_runners(num_env_runners=0)
+        .training(
+            n_step=3,
+            bc_iters=0,
+            clip_actions=False,
+            tau=0.005,
+            target_entropy="auto",
+            q_model_config={
+                "fcnet_hiddens": [256, 256],
+                "fcnet_activation": "relu",
+            },
+            policy_model_config={
+                "fcnet_hiddens": [256, 256],
+                "fcnet_activation": "relu",
+            },
+            optimization_config={
+                "actor_learning_rate": 3e-4,
+                "critic_learning_rate": 3e-4,
+                "entropy_learning_rate": 3e-4,
+            },
+            train_batch_size=256,
+            target_network_update_freq=1,
+            num_steps_sampled_before_learning_starts=256,
+        )
+        .reporting(min_train_timesteps_per_iteration=1000)
+        .debugging(log_level="INFO")
+        .environment("Pendulum-v1", normalize_actions=True)
+        .offline_data(
+            input_config={
+                "paths": ["offline/tests/data/pendulum/enormous.zip"],
+                "format": "json",
+            }
+        )
+        .evaluation(
+            evaluation_num_env_runners=1,
+            evaluation_interval=1,
+            evaluation_duration=10,
+            evaluation_parallel_to_training=False,
+            evaluation_config=cql.CQLConfig.overrides(input_="sampler"),
+        )
+    )
+    # evaluation_parallel_to_training should be False b/c iterations are very long
+    # and this would cause evaluation to lag one iter behind training.
+
+    # Check, whether we can learn from the given file in `num_iterations`
+    # iterations, up to a reward of `min_reward`.
+    num_iterations = 5
+    min_reward = -300
+
+    cql_algorithm = cql.CQL(config=config)
+    learnt = False
+    for i in range(num_iterations):
+        print(f"Iter {i}")
+        eval_results = cql_algorithm.train().get(EVALUATION_RESULTS)
+        if eval_results:
+            print(
+                "... R={}".format(eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN])
+            )
+            # Learn until some reward is reached on an actual live env.
+            if eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_reward:
+                # Test passed gracefully.
+                if args.as_test:
+                    print("Test passed after {} iterations.".format(i))
+                    quit(0)
+                learnt = True
+                break
+
+    # Get policy and model.
+    cql_policy = cql_algorithm.get_policy()
+    cql_model = cql_policy.model
+
+    # If you would like to query CQL's learnt Q-function for arbitrary
+    # (cont.) actions, do the following:
+    obs_batch = torch.from_numpy(np.random.random(size=(5, 3)))
+    action_batch = torch.from_numpy(np.random.random(size=(5, 1)))
+    q_values = cql_model.get_q_values(obs_batch, action_batch)[0]
+    # If you are using the "twin_q", there'll be 2 Q-networks and
+    # we usually consider the min of the 2 outputs, like so:
+    twin_q_values = cql_model.get_twin_q_values(obs_batch, action_batch)[0]
+    final_q_values = torch.min(q_values, twin_q_values)[0]
+    print(f"final_q_values={final_q_values.detach().numpy()}")
+
+    # Example on how to do evaluation on the trained Algorithm.
+    # using the data from our buffer.
+    # Get a sample (MultiAgentBatch).
+
+    batch = synchronous_parallel_sample(worker_set=cql_algorithm.env_runner_group)
+    batch = convert_ma_batch_to_sample_batch(batch)
+    obs = torch.from_numpy(batch["obs"])
+    # Pass the observations through our model to get the
+    # features, which then to pass through the Q-head.
+    model_out, _ = cql_model({"obs": obs})
+    # The estimated Q-values from the (historic) actions in the batch.
+    q_values_old = cql_model.get_q_values(
+        model_out, torch.from_numpy(batch["actions"])
+    )[0]
+    # The estimated Q-values for the new actions computed by our policy.
+    actions_new = cql_policy.compute_actions_from_input_dict({"obs": obs})[0]
+    q_values_new = cql_model.get_q_values(model_out, torch.from_numpy(actions_new))[0]
+    print(f"Q-val batch={q_values_old.detach().numpy()}")
+    print(f"Q-val policy={q_values_new.detach().numpy()}")
+
+    cql_algorithm.stop()

rllib/examples/_old_api_stack/offline_rl/saving_experiences.py

@@ -0,0 +1,60 @@
+# @OldAPIStack
+
+"""Simple example of writing experiences to a file using JsonWriter."""
+
+# __sphinx_doc_begin__
+import os
+
+import gymnasium as gym
+import numpy as np
+
+from ray._common.utils import get_default_ray_temp_dir
+from ray.rllib.evaluation.sample_batch_builder import SampleBatchBuilder
+from ray.rllib.models.preprocessors import get_preprocessor
+from ray.rllib.offline.json_writer import JsonWriter
+
+if __name__ == "__main__":
+    batch_builder = SampleBatchBuilder()  # or MultiAgentSampleBatchBuilder
+    writer = JsonWriter(os.path.join(get_default_ray_temp_dir(), "demo-out"))
+
+    # You normally wouldn't want to manually create sample batches if a
+    # simulator is available, but let's do it anyways for example purposes:
+    env = gym.make("CartPole-v1")
+
+    # RLlib uses preprocessors to implement transforms such as one-hot encoding
+    # and flattening of tuple and dict observations. For CartPole a no-op
+    # preprocessor is used, but this may be relevant for more complex envs.
+    prep = get_preprocessor(env.observation_space)(env.observation_space)
+    print("The preprocessor is", prep)
+
+    for eps_id in range(100):
+        obs, info = env.reset()
+        prev_action = np.zeros_like(env.action_space.sample())
+        prev_reward = 0
+        terminated = truncated = False
+        t = 0
+        while not terminated and not truncated:
+            action = env.action_space.sample()
+            new_obs, rew, terminated, truncated, info = env.step(action)
+            batch_builder.add_values(
+                t=t,
+                eps_id=eps_id,
+                agent_index=0,
+                obs=prep.transform(obs),
+                actions=action,
+                action_prob=1.0,  # put the true action probability here
+                action_logp=0.0,
+                rewards=rew,
+                prev_actions=prev_action,
+                prev_rewards=prev_reward,
+                terminateds=terminated,
+                truncateds=truncated,
+                infos=info,
+                new_obs=prep.transform(new_obs),
+            )
+            obs = new_obs
+            prev_action = action
+            prev_reward = rew
+            t += 1
+        writer.write(batch_builder.build_and_reset())
+# __sphinx_doc_end__

rllib/examples/_old_api_stack/parametric_actions_cartpole.py

@@ -0,0 +1,121 @@
+# @OldAPIStack
+"""Example of handling variable length or parametric action spaces.
+
+This toy example demonstrates the action-embedding based approach for handling large
+discrete action spaces (potentially infinite in size), similar to this example:
+
+    https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/
+
+This example works with RLlib's policy gradient style algorithms
+(e.g., PG, PPO, IMPALA, A2C) and DQN.
+
+Note that since the model outputs now include "-inf" tf.float32.min
+values, not all algorithm options are supported. For example,
+algorithms might crash if they don't properly ignore the -inf action scores.
+Working configurations are given below.
+"""
+
+import argparse
+import os
+
+import ray
+from ray import tune
+from ray.rllib.examples._old_api_stack.models.parametric_actions_model import (
+    ParametricActionsModel,
+    TorchParametricActionsModel,
+)
+from ray.rllib.examples.envs.classes.parametric_actions_cartpole import (
+    ParametricActionsCartPole,
+)
+from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.test_utils import check_learning_achieved
+from ray.tune.registry import register_env
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use."
+)
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--as-test",
+    action="store_true",
+    help="Whether this script should be run as a test: --stop-reward must "
+    "be achieved within --stop-timesteps AND --stop-iters.",
+)
+parser.add_argument(
+    "--stop-iters", type=int, default=200, help="Number of iterations to train."
+)
+parser.add_argument(
+    "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train."
+)
+parser.add_argument(
+    "--stop-reward", type=float, default=150.0, help="Reward at which we stop training."
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init()
+
+    register_env("pa_cartpole", lambda _: ParametricActionsCartPole(10))
+    ModelCatalog.register_custom_model(
+        "pa_model",
+        TorchParametricActionsModel
+        if args.framework == "torch"
+        else ParametricActionsModel,
+    )
+
+    if args.run == "DQN":
+        cfg = {
+            # TODO(ekl) we need to set these to prevent the masked values
+            # from being further processed in DistributionalQModel, which
+            # would mess up the masking. It is possible to support these if we
+            # defined a custom DistributionalQModel that is aware of masking.
+            "hiddens": [],
+            "dueling": False,
+            "enable_rl_module_and_learner": False,
+            "enable_env_runner_and_connector_v2": False,
+        }
+    else:
+        cfg = {}
+
+    config = dict(
+        {
+            "env": "pa_cartpole",
+            "model": {
+                "custom_model": "pa_model",
+            },
+            # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+            "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")),
+            "num_env_runners": 0,
+            "framework": args.framework,
+        },
+        **cfg,
+    )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps,
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+    }
+
+    results = tune.Tuner(
+        args.run,
+        run_config=tune.RunConfig(stop=stop, verbose=1),
+        param_space=config,
+    ).fit()
+
+    if args.as_test:
+        check_learning_achieved(results, args.stop_reward)
+
+    ray.shutdown()

rllib/examples/_old_api_stack/parametric_actions_cartpole_embeddings_learnt_by_model.py

@@ -0,0 +1,107 @@
+# @OldAPIStack
+"""Example of handling variable length or parametric action spaces.
+
+This is a toy example of the action-embedding based approach for handling large
+discrete action spaces (potentially infinite in size), similar to this:
+
+    https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/
+
+This currently works with RLlib's policy gradient style algorithms
+(e.g., PG, PPO, IMPALA, A2C) and also DQN.
+
+Note that since the model outputs now include "-inf" tf.float32.min
+values, not all algorithm options are supported at the moment. For example,
+algorithms might crash if they don't properly ignore the -inf action scores.
+Working configurations are given below.
+"""
+
+import argparse
+import os
+
+import ray
+from ray import tune
+from ray.rllib.examples._old_api_stack.models.parametric_actions_model import (
+    ParametricActionsModelThatLearnsEmbeddings,
+)
+from ray.rllib.examples.envs.classes.parametric_actions_cartpole import (
+    ParametricActionsCartPoleNoEmbeddings,
+)
+from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.test_utils import check_learning_achieved
+from ray.tune.registry import register_env
+from ray.tune.result import TRAINING_ITERATION
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--run", type=str, default="PPO")
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2"],
+    default="tf",
+    help="The DL framework specifier (Torch not supported "
+    "due to the lack of a model).",
+)
+parser.add_argument("--as-test", action="store_true")
+parser.add_argument("--stop-iters", type=int, default=200)
+parser.add_argument("--stop-reward", type=float, default=150.0)
+parser.add_argument("--stop-timesteps", type=int, default=100000)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init()
+
+    register_env("pa_cartpole", lambda _: ParametricActionsCartPoleNoEmbeddings(10))
+
+    ModelCatalog.register_custom_model(
+        "pa_model", ParametricActionsModelThatLearnsEmbeddings
+    )
+
+    if args.run == "DQN":
+        cfg = {
+            # TODO(ekl) we need to set these to prevent the masked values
+            # from being further processed in DistributionalQModel, which
+            # would mess up the masking. It is possible to support these if we
+            # defined a custom DistributionalQModel that is aware of masking.
+            "hiddens": [],
+            "dueling": False,
+            "enable_rl_module_and_learner": False,
+            "enable_env_runner_and_connector_v2": False,
+        }
+    else:
+        cfg = {}
+
+    config = dict(
+        {
+            "env": "pa_cartpole",
+            "model": {
+                "custom_model": "pa_model",
+            },
+            # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+            "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")),
+            "num_env_runners": 0,
+            "framework": args.framework,
+            "action_mask_key": "valid_avail_actions_mask",
+        },
+        **cfg,
+    )
+
+    stop = {
+        TRAINING_ITERATION: args.stop_iters,
+        NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+        f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
+    }
+
+    results = tune.Tuner(
+        args.run,
+        run_config=tune.RunConfig(stop=stop, verbose=2),
+        param_space=config,
+    ).fit()
+
+    if args.as_test:
+        check_learning_achieved(results, args.stop_reward)
+
+    ray.shutdown()

rllib/examples/_old_api_stack/policy/cliff_walking_wall_policy.py

@@ -0,0 +1,101 @@
+# @OldAPIStack
+from typing import Dict, List, Optional, Tuple, Union
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.policy.policy import Policy, ViewRequirement
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.debug import update_global_seed_if_necessary
+from ray.rllib.utils.typing import AlgorithmConfigDict, TensorStructType, TensorType
+
+
+class CliffWalkingWallPolicy(Policy):
+    """Optimal RLlib policy for the CliffWalkingWallEnv environment, defined in
+    ray/rllib/examples/env/cliff_walking_wall_env.py, with epsilon-greedy exploration.
+
+    The policy takes a random action with probability epsilon, specified
+    by `config["epsilon"]`, and the optimal action with probability  1 - epsilon.
+    """
+
+    @override(Policy)
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        config: AlgorithmConfigDict,
+    ):
+        update_global_seed_if_necessary(seed=config.get("seed"))
+        super().__init__(observation_space, action_space, config)
+
+        # Known optimal action dist for each of the 48 states and 4 actions
+        self.action_dist = np.zeros((48, 4), dtype=float)
+        # Starting state: go up
+        self.action_dist[36] = (1, 0, 0, 0)
+        # Cliff + Goal: never actually used, set to random
+        self.action_dist[37:] = (0.25, 0.25, 0.25, 0.25)
+        # Row 2; always go right
+        self.action_dist[24:36] = (0, 1, 0, 0)
+        # Row 0 and Row 1; go down or go right
+        self.action_dist[0:24] = (0, 0.5, 0.5, 0)
+        # Col 11; always go down, supercedes previous values
+        self.action_dist[[11, 23, 35]] = (0, 0, 1, 0)
+        assert np.allclose(self.action_dist.sum(-1), 1)
+
+        # Epsilon-Greedy action selection
+        epsilon = config.get("epsilon", 0.0)
+        self.action_dist = self.action_dist * (1 - epsilon) + epsilon / 4
+        assert np.allclose(self.action_dist.sum(-1), 1)
+
+        # Attributes required for RLlib; note that while CliffWalkingWallPolicy
+        # inherits from Policy, it actually implements TorchPolicyV2.
+        self.view_requirements[SampleBatch.ACTION_PROB] = ViewRequirement()
+        self.device = "cpu"
+        self.model = None
+        self.dist_class = TorchCategorical
+
+    @override(Policy)
+    def compute_actions(
+        self,
+        obs_batch: Union[List[TensorStructType], TensorStructType],
+        state_batches: Optional[List[TensorType]] = None,
+        **kwargs,
+    ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]:
+        obs = np.array(obs_batch, dtype=int)
+        action_probs = self.action_dist[obs]
+        actions = np.zeros(len(obs), dtype=int)
+        for i in range(len(obs)):
+            actions[i] = np.random.choice(4, p=action_probs[i])
+        return (
+            actions,
+            [],
+            {SampleBatch.ACTION_PROB: action_probs[np.arange(len(obs)), actions]},
+        )
+
+    @override(Policy)
+    def compute_log_likelihoods(
+        self,
+        actions: Union[List[TensorType], TensorType],
+        obs_batch: Union[List[TensorType], TensorType],
+        **kwargs,
+    ) -> TensorType:
+        obs = np.array(obs_batch, dtype=int)
+        actions = np.array(actions, dtype=int)
+        # Compute action probs for all possible actions
+        action_probs = self.action_dist[obs]
+        # Take the action_probs corresponding to the specified actions
+        action_probs = action_probs[np.arange(len(obs)), actions]
+        # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0
+        with np.errstate(divide="ignore"):
+            return np.log(action_probs)
+
+    def action_distribution_fn(
+        self, model, obs_batch: TensorStructType, **kwargs
+    ) -> Tuple[TensorType, type, List[TensorType]]:
+        obs = np.array(obs_batch[SampleBatch.OBS], dtype=int)
+        action_probs = self.action_dist[obs]
+        # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0
+        with np.errstate(divide="ignore"):
+            return np.log(action_probs), TorchCategorical, None

rllib/examples/_old_api_stack/policy/random_policy.py

@@ -0,0 +1,102 @@
+# @OldAPIStack
+import random
+from typing import (
+    List,
+    Optional,
+    Union,
+)
+
+import numpy as np
+import tree  # pip install dm_tree
+from gymnasium.spaces import Box
+
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import ModelWeights, TensorStructType, TensorType
+
+
+class RandomPolicy(Policy):
+    """Hand-coded policy that returns random actions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Whether for compute_actions, the bounds given in action_space
+        # should be ignored (default: False). This is to test action-clipping
+        # and any Env's reaction to bounds breaches.
+        if self.config.get("ignore_action_bounds", False) and isinstance(
+            self.action_space, Box
+        ):
+            self.action_space_for_sampling = Box(
+                -float("inf"),
+                float("inf"),
+                shape=self.action_space.shape,
+                dtype=self.action_space.dtype,
+            )
+        else:
+            self.action_space_for_sampling = self.action_space
+
+    @override(Policy)
+    def init_view_requirements(self):
+        super().init_view_requirements()
+        # Disable for_training and action attributes for SampleBatch.INFOS column
+        # since it can not be properly batched.
+        vr = self.view_requirements[SampleBatch.INFOS]
+        vr.used_for_training = False
+        vr.used_for_compute_actions = False
+
+    @override(Policy)
+    def compute_actions(
+        self,
+        obs_batch: Union[List[TensorStructType], TensorStructType],
+        state_batches: Optional[List[TensorType]] = None,
+        prev_action_batch: Union[List[TensorStructType], TensorStructType] = None,
+        prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None,
+        **kwargs,
+    ):
+        # Alternatively, a numpy array would work here as well.
+        # e.g.: np.array([random.choice([0, 1])] * len(obs_batch))
+        obs_batch_size = len(tree.flatten(obs_batch)[0])
+        return (
+            [self.action_space_for_sampling.sample() for _ in range(obs_batch_size)],
+            [],
+            {},
+        )
+
+    @override(Policy)
+    def learn_on_batch(self, samples):
+        """No learning."""
+        return {}
+
+    @override(Policy)
+    def compute_log_likelihoods(
+        self,
+        actions,
+        obs_batch,
+        state_batches=None,
+        prev_action_batch=None,
+        prev_reward_batch=None,
+        **kwargs,
+    ):
+        return np.array([random.random()] * len(obs_batch))
+
+    @override(Policy)
+    def get_weights(self) -> ModelWeights:
+        """No weights to save."""
+        return {}
+
+    @override(Policy)
+    def set_weights(self, weights: ModelWeights) -> None:
+        """No weights to set."""
+        pass
+
+    @override(Policy)
+    def _get_dummy_batch_from_view_requirements(self, batch_size: int = 1):
+        return SampleBatch(
+            {
+                SampleBatch.OBS: tree.map_structure(
+                    lambda s: s[None], self.observation_space.sample()
+                ),
+            }
+        )

rllib/examples/_old_api_stack/replay_buffer_api.py

@@ -0,0 +1,82 @@
+# @OldAPIStack
+
+# __sphinx_doc_replay_buffer_api_example_script_begin__
+"""Simple example of how to modify replay buffer behaviour.
+
+We modify DQN to utilize prioritized replay but supplying it with the
+PrioritizedMultiAgentReplayBuffer instead of the standard MultiAgentReplayBuffer.
+This is possible because DQN uses the DQN training iteration function,
+which includes and a priority update, given that a fitting buffer is provided.
+"""
+
+import argparse
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.dqn import DQNConfig
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME
+from ray.rllib.utils.replay_buffers.replay_buffer import StorageUnit
+from ray.tune.result import TRAINING_ITERATION
+
+tf1, tf, tfv = try_import_tf()
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--num-cpus", type=int, default=0)
+parser.add_argument(
+    "--framework",
+    choices=["tf", "tf2", "torch"],
+    default="torch",
+    help="The DL framework specifier.",
+)
+parser.add_argument(
+    "--stop-iters", type=int, default=50, help="Number of iterations to train."
+)
+parser.add_argument(
+    "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train."
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    ray.init(num_cpus=args.num_cpus or None)
+
+    # This is where we add prioritized experiences replay
+    # The training iteration function that is used by DQN already includes a priority
+    # update step.
+    replay_buffer_config = {
+        "type": "MultiAgentPrioritizedReplayBuffer",
+        # Although not necessary, we can modify the default constructor args of
+        # the replay buffer here
+        "prioritized_replay_alpha": 0.5,
+        "storage_unit": StorageUnit.SEQUENCES,
+        "replay_burn_in": 20,
+        "zero_init_states": True,
+    }
+
+    config = (
+        DQNConfig()
+        .environment("CartPole-v1")
+        .framework(framework=args.framework)
+        .env_runners(num_env_runners=4)
+        .training(
+            model=dict(use_lstm=True, lstm_cell_size=64, max_seq_len=20),
+            replay_buffer_config=replay_buffer_config,
+        )
+    )
+
+    stop_config = {
+        NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps,
+        TRAINING_ITERATION: args.stop_iters,
+    }
+
+    results = tune.Tuner(
+        config.algo_class,
+        param_space=config,
+        run_config=tune.RunConfig(stop=stop_config),
+    ).fit()
+
+    ray.shutdown()
+
+# __sphinx_doc_replay_buffer_api_example_script_end__

rllib/examples/actions/autoregressive_actions.py

@@ -0,0 +1,107 @@
+"""Example on how to define and run with an RLModule with a dependent action space.
+
+This examples:
+    - Shows how to write a custom RLModule outputting autoregressive actions.
+    The RLModule class used here implements a prior distribution for the first couple
+    of actions and then uses the sampled actions to compute the parameters for and
+    sample from a posterior distribution.
+    - Shows how to configure a PPO algorithm to use the custom RLModule.
+    - Stops the training after 100k steps or when the mean episode return
+    exceeds -0.012 in evaluation, i.e. if the agent has learned to
+    synchronize its actions.
+
+For details on the environment used, take a look at the `CorrelatedActionsEnv`
+class. To receive an episode return over 100, the agent must learn how to synchronize
+its actions.
+
+
+How to run this script
+----------------------
+`python [script file name].py --num-env-runners 2`
+
+Control the number of `EnvRunner`s with the `--num-env-runners` flag. This
+will increase the sampling speed.
+
+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
+-----------------
+You should reach an episode return of better than -0.5 quickly through a simple PPO
+policy. The logic behind beating the env is roughly:
+
+OBS:  optimal a1:   r1:  optimal a2:   r2:
+-1      2            0      -1.0        0
+-0.5    1/2       -0.5   -0.5/-1.5      0
+0       1            0      -1.0        0
+0.5     0/1       -0.5   -0.5/-1.5      0
+1       0            0      -1.0        0
+
+Meaning, most of the time, you would receive a reward better than -0.5, but worse than
+0.0.
+
++--------------------------------------+------------+--------+------------------+
+| Trial name                           | status     |   iter |   total time (s) |
+|                                      |            |        |                  |
+|--------------------------------------+------------+--------+------------------+
+| PPO_CorrelatedActionsEnv_6660d_00000 | TERMINATED |     76 |          132.438 |
++--------------------------------------+------------+--------+------------------+
++------------------------+------------------------+------------------------+
+|    episode_return_mean |   num_env_steps_sample |   ...env_steps_sampled |
+|                        |             d_lifetime |   _lifetime_throughput |
+|------------------------+------------------------+------------------------|
+|                  -0.43 |                 152000 |                1283.48 |
++------------------------+------------------------+------------------------+
+"""
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.correlated_actions_env import CorrelatedActionsEnv
+from ray.rllib.examples.rl_modules.classes.autoregressive_actions_rlm import (
+    AutoregressiveActionsRLM,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=1000,
+    default_timesteps=2000000,
+    default_reward=-0.45,
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.algo != "PPO":
+        raise ValueError(
+            "This example script only runs with PPO! Set --algo=PPO on the command "
+            "line."
+        )
+
+    base_config = (
+        PPOConfig()
+        .environment(CorrelatedActionsEnv)
+        .training(
+            train_batch_size_per_learner=2000,
+            num_epochs=12,
+            minibatch_size=256,
+            entropy_coeff=0.005,
+            lr=0.0003,
+        )
+        # Specify the RLModule class to be used.
+        .rl_module(
+            rl_module_spec=RLModuleSpec(module_class=AutoregressiveActionsRLM),
+        )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)

rllib/examples/actions/custom_action_distribution.py

@@ -0,0 +1,116 @@
+"""Example on how to define and run an experiment with a custom action distribution.
+
+The example uses an additional `temperature` parameter on top of the built-in
+`TorchCategorical` class. Incoming logits (outputs from the RLModule) are divided by
+this temperature before creating the underlying
+torch.distributions.categorical.Categorical object.
+
+This examples:
+    - Shows how to write a custom RLlib action distribution class accepting an
+    additional parameter in its constructor.
+    - demonstrates how you can subclass the TorchRLModule base class and write your
+    own architecture by overriding the `setup()` method.
+    - shows how to set the attribute `self.action_dist_cls` in that same `setup()`
+    method. For an alternative way of defining action distribution classes for your
+    RLModules, see the `setup()` method implementation in the imported
+    `CustomActionDistributionRLModule` class.
+    - shows how you then configure an RLlib Algorithm such that it uses your custom
+    RLModule (instead of a default RLModule).
+
+
+How to run this script
+----------------------
+`python [script file name].py --temperature=0.8`
+
+Use the `--temperature` setting to change the temperature. Higher values (>>1.0) lead
+to almost random behavior, lower values (<<1.0) lead to always-greedy behavior. Note
+though, that both extremes hurt learning performance.
+
+Control the number of `EnvRunner`s with the `--num-env-runners` flag. This
+will increase the sampling speed.
+
+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 a --temperature setting of 0.75, learning seems to be particularly easy with the
+given other parameters:
+
++-----------------------------+------------+-----------------+--------+
+| Trial name                  | status     | loc             |   iter |
+|                             |            |                 |        |
+|-----------------------------+------------+-----------------+--------+
+| PPO_CartPole-v1_1bbe0_00000 | TERMINATED | 127.0.0.1:81594 |     22 |
++-----------------------------+------------+-----------------+--------+
++------------------+------------------------+------------------------+
+|   total time (s) |    episode_return_mean |   num_env_steps_sample |
+|                  |                        |             d_lifetime |
+|------------------+------------------------+------------------------|
+|          17.6368 |                 450.54 |                  88000 |
++------------------+------------------------+------------------------+
+"""
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.rl_modules.classes.custom_action_distribution_rlm import (
+    CustomActionDistributionRLModule,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_timesteps=200000,
+    default_reward=450.0,
+)
+parser.add_argument(
+    "--temperature",
+    type=float,
+    default=2.0,
+    help="The action distribution temperature to apply to the raw model logits. "
+    "Logits are first divided by the temperature, then an underlying torch.Categorical "
+    "distribution is created from those altered logits and used for sampling actions. "
+    "Set this to <<1.0 to approximate greedy behavior and to >>1.0 to approximate "
+    "random behavior.",
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.algo != "PPO":
+        raise ValueError(
+            "This example script only runs with PPO! Set --algo=PPO on the command "
+            "line."
+        )
+
+    base_config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .training(
+            lr=0.0003,
+            num_epochs=6,
+            vf_loss_coeff=0.01,
+        )
+        # Specify the RLModule class to be used.
+        .rl_module(
+            rl_module_spec=RLModuleSpec(
+                module_class=CustomActionDistributionRLModule,
+                model_config={
+                    "hidden_dim": 128,
+                    "action_dist_temperature": args.temperature,
+                },
+            ),
+        )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)

rllib/examples/actions/nested_action_spaces.py

@@ -0,0 +1,83 @@
+from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
+
+from ray.rllib.connectors.env_to_module import FlattenObservations
+from ray.rllib.examples.envs.classes.multi_agent import (
+    MultiAgentNestedSpaceRepeatAfterMeEnv,
+)
+from ray.rllib.examples.envs.classes.nested_space_repeat_after_me_env import (
+    NestedSpaceRepeatAfterMeEnv,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.tune.registry import get_trainable_cls, register_env
+
+# Read in common example script command line arguments.
+parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=-500.0)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    # Define env-to-module-connector pipeline for the new stack.
+    def _env_to_module_pipeline(env, spaces, device):
+        return FlattenObservations(multi_agent=args.num_agents > 0)
+
+    # Register our environment with tune.
+    if args.num_agents > 0:
+        register_env(
+            "env",
+            lambda c: MultiAgentNestedSpaceRepeatAfterMeEnv(
+                config=dict(c, **{"num_agents": args.num_agents})
+            ),
+        )
+    else:
+        register_env("env", lambda c: NestedSpaceRepeatAfterMeEnv(c))
+
+    # Define the AlgorithmConfig used.
+    base_config = (
+        get_trainable_cls(args.algo)
+        .get_default_config()
+        .environment(
+            "env",
+            env_config={
+                "space": Dict(
+                    {
+                        "a": Tuple(
+                            [Dict({"d": Box(-15.0, 3.0, ()), "e": Discrete(3)})]
+                        ),
+                        "b": Box(-10.0, 10.0, (2,)),
+                        "c": MultiDiscrete([3, 3]),
+                        "d": Discrete(2),
+                    }
+                ),
+                "episode_len": 100,
+            },
+        )
+        .env_runners(env_to_module_connector=_env_to_module_pipeline)
+        # No history in Env (bandit problem).
+        .training(
+            gamma=0.0,
+            lr=0.0005,
+        )
+    )
+
+    # Add a simple multi-agent setup.
+    if args.num_agents > 0:
+        base_config.multi_agent(
+            policies={f"p{i}" for i in range(args.num_agents)},
+            policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}",
+        )
+
+    # Fix some PPO-specific settings.
+    if args.algo == "PPO":
+        base_config.training(
+            # We don't want high entropy in this Env.
+            entropy_coeff=0.00005,
+            num_epochs=4,
+            vf_loss_coeff=0.01,
+        )
+
+    # Run everything as configured.
+    run_rllib_example_script_experiment(base_config, args)

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)