chore: import upstream snapshot with attribution

rllib/examples/envs/classes/action_mask_env.py

@@ -0,0 +1,42 @@
+import numpy as np
+from gymnasium.spaces import Box, Dict, Discrete
+
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+
+
+class ActionMaskEnv(RandomEnv):
+    """A randomly acting environment that publishes an action-mask each step."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Masking only works for Discrete actions.
+        assert isinstance(self.action_space, Discrete)
+        # Add action_mask to observations.
+        self.observation_space = Dict(
+            {
+                "action_mask": Box(0.0, 1.0, shape=(self.action_space.n,)),
+                "observations": self.observation_space,
+            }
+        )
+        self.valid_actions = None
+
+    def reset(self, *, seed=None, options=None):
+        obs, info = super().reset()
+        self._fix_action_mask(obs)
+        return obs, info
+
+    def step(self, action):
+        # Check whether action is valid.
+        if not self.valid_actions[action]:
+            raise ValueError(
+                f"Invalid action ({action}) sent to env! "
+                f"valid_actions={self.valid_actions}"
+            )
+        obs, rew, done, truncated, info = super().step(action)
+        self._fix_action_mask(obs)
+        return obs, rew, done, truncated, info
+
+    def _fix_action_mask(self, obs):
+        # Fix action-mask: Everything larger 0.5 is 1.0, everything else 0.0.
+        self.valid_actions = np.round(obs["action_mask"])
+        obs["action_mask"] = self.valid_actions

rllib/examples/envs/classes/cartpole_crashing.py

@@ -0,0 +1,183 @@
+import logging
+import time
+
+import numpy as np
+from gymnasium.envs.classic_control import CartPoleEnv
+
+from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.error import EnvError
+
+logger = logging.getLogger(__name__)
+
+
+class CartPoleCrashing(CartPoleEnv):
+    """A CartPole env that crashes (or stalls) from time to time.
+
+    Useful for testing faulty sub-env (within a vectorized env) handling by
+    EnvRunners.
+
+    After crashing, the env expects a `reset()` call next (calling `step()` will
+    result in yet another error), which may or may not take a very long time to
+    complete. This simulates the env having to reinitialize some sub-processes, e.g.
+    an external connection.
+
+    The env can also be configured to stall (and do nothing during a call to `step()`)
+    from time to time for a configurable amount of time.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        self.config = config if config is not None else {}
+
+        # Crash probability (in each `step()`).
+        self.p_crash = config.get("p_crash", 0.005)
+        # Crash probability when `reset()` is called.
+        self.p_crash_reset = config.get("p_crash_reset", 0.0)
+        # Crash exactly after every n steps. If a 2-tuple, will uniformly sample
+        # crash timesteps from in between the two given values.
+        self.crash_after_n_steps = config.get("crash_after_n_steps")
+        self._crash_after_n_steps = None
+        assert (
+            self.crash_after_n_steps is None
+            or isinstance(self.crash_after_n_steps, int)
+            or (
+                isinstance(self.crash_after_n_steps, tuple)
+                and len(self.crash_after_n_steps) == 2
+            )
+        )
+        # Only ever crash, if on certain worker indices.
+        faulty_indices = config.get("crash_on_worker_indices", None)
+        if faulty_indices and config.worker_index not in faulty_indices:
+            self.p_crash = 0.0
+            self.p_crash_reset = 0.0
+            self.crash_after_n_steps = None
+
+        # Stall probability (in each `step()`).
+        self.p_stall = config.get("p_stall", 0.0)
+        # Stall probability when `reset()` is called.
+        self.p_stall_reset = config.get("p_stall_reset", 0.0)
+        # Stall exactly after every n steps.
+        self.stall_after_n_steps = config.get("stall_after_n_steps")
+        self._stall_after_n_steps = None
+        # Amount of time to stall. If a 2-tuple, will uniformly sample from in between
+        # the two given values.
+        self.stall_time_sec = config.get("stall_time_sec")
+        assert (
+            self.stall_time_sec is None
+            or isinstance(self.stall_time_sec, (int, float))
+            or (
+                isinstance(self.stall_time_sec, tuple) and len(self.stall_time_sec) == 2
+            )
+        )
+
+        # Only ever stall, if on certain worker indices.
+        faulty_indices = config.get("stall_on_worker_indices", None)
+        if faulty_indices and config.worker_index not in faulty_indices:
+            self.p_stall = 0.0
+            self.p_stall_reset = 0.0
+            self.stall_after_n_steps = None
+
+        # Timestep counter for the ongoing episode.
+        self.timesteps = 0
+
+        # Time in seconds to initialize (in this c'tor).
+        sample = 0.0
+        if "init_time_s" in config:
+            sample = (
+                config["init_time_s"]
+                if not isinstance(config["init_time_s"], tuple)
+                else np.random.uniform(
+                    config["init_time_s"][0], config["init_time_s"][1]
+                )
+            )
+
+        print(f"Initializing crashing env (with init-delay of {sample}sec) ...")
+        time.sleep(sample)
+
+        # Make sure envs don't crash at the same time.
+        self._rng = np.random.RandomState()
+
+    @override(CartPoleEnv)
+    def reset(self, *, seed=None, options=None):
+        # Reset timestep counter for the new episode.
+        self.timesteps = 0
+        self._crash_after_n_steps = None
+
+        # Should we crash?
+        if self._should_crash(p=self.p_crash_reset):
+            raise EnvError(
+                f"Simulated env crash on worker={self.config.worker_index} "
+                f"env-idx={self.config.vector_index} during `reset()`! "
+                "Feel free to use any other exception type here instead."
+            )
+        # Should we stall for a while?
+        self._stall_if_necessary(p=self.p_stall_reset)
+
+        return super().reset()
+
+    @override(CartPoleEnv)
+    def step(self, action):
+        # Increase timestep counter for the ongoing episode.
+        self.timesteps += 1
+
+        # Should we crash?
+        if self._should_crash(p=self.p_crash):
+            raise EnvError(
+                f"Simulated env crash on worker={self.config.worker_index} "
+                f"env-idx={self.config.vector_index} during `step()`! "
+                "Feel free to use any other exception type here instead."
+            )
+        # Should we stall for a while?
+        self._stall_if_necessary(p=self.p_stall)
+
+        return super().step(action)
+
+    def _should_crash(self, p):
+        rnd = self._rng.rand()
+        if rnd < p:
+            print("Crashing due to p(crash)!")
+            return True
+        elif self.crash_after_n_steps is not None:
+            if self._crash_after_n_steps is None:
+                self._crash_after_n_steps = (
+                    self.crash_after_n_steps
+                    if not isinstance(self.crash_after_n_steps, tuple)
+                    else np.random.randint(
+                        self.crash_after_n_steps[0], self.crash_after_n_steps[1]
+                    )
+                )
+            if self._crash_after_n_steps == self.timesteps:
+                print("Crashing due to n timesteps reached!")
+                return True
+
+        return False
+
+    def _stall_if_necessary(self, p):
+        stall = False
+        if self._rng.rand() < p:
+            stall = True
+        elif self.stall_after_n_steps is not None:
+            if self._stall_after_n_steps is None:
+                self._stall_after_n_steps = (
+                    self.stall_after_n_steps
+                    if not isinstance(self.stall_after_n_steps, tuple)
+                    else np.random.randint(
+                        self.stall_after_n_steps[0], self.stall_after_n_steps[1]
+                    )
+                )
+            if self._stall_after_n_steps == self.timesteps:
+                stall = True
+
+        if stall:
+            sec = (
+                self.stall_time_sec
+                if not isinstance(self.stall_time_sec, tuple)
+                else np.random.uniform(self.stall_time_sec[0], self.stall_time_sec[1])
+            )
+            print(f" -> will stall for {sec}sec ...")
+            time.sleep(sec)
+
+
+MultiAgentCartPoleCrashing = make_multi_agent(lambda config: CartPoleCrashing(config))

rllib/examples/envs/classes/cartpole_sparse_rewards.py

@@ -0,0 +1,51 @@
+from copy import deepcopy
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Dict, Discrete
+
+
+class CartPoleSparseRewards(gym.Env):
+    """Wrapper for gym CartPole environment where reward is accumulated to the end."""
+
+    def __init__(self, config=None):
+        self.env = gym.make("CartPole-v1")
+        self.action_space = Discrete(2)
+        self.observation_space = Dict(
+            {
+                "obs": self.env.observation_space,
+                "action_mask": Box(
+                    low=0, high=1, shape=(self.action_space.n,), dtype=np.int8
+                ),
+            }
+        )
+        self.running_reward = 0
+
+    def reset(self, *, seed=None, options=None):
+        self.running_reward = 0
+        obs, infos = self.env.reset()
+        return {
+            "obs": obs,
+            "action_mask": np.array([1, 1], dtype=np.int8),
+        }, infos
+
+    def step(self, action):
+        obs, rew, terminated, truncated, info = self.env.step(action)
+        self.running_reward += rew
+        score = self.running_reward if terminated else 0
+        return (
+            {"obs": obs, "action_mask": np.array([1, 1], dtype=np.int8)},
+            score,
+            terminated,
+            truncated,
+            info,
+        )
+
+    def set_state(self, state):
+        self.running_reward = state[1]
+        self.env = deepcopy(state[0])
+        obs = np.array(list(self.env.unwrapped.state))
+        return {"obs": obs, "action_mask": np.array([1, 1], dtype=np.int8)}
+
+    def get_state(self):
+        return deepcopy(self.env), self.running_reward

rllib/examples/envs/classes/cartpole_with_dict_observation_space.py

@@ -0,0 +1,74 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.classic_control import CartPoleEnv
+
+
+class CartPoleWithDictObservationSpace(CartPoleEnv):
+    """CartPole gym environment that has a dict observation space.
+
+    However, otherwise, the information content in each observation remains the same.
+
+    https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_control/cartpole.py  # noqa
+
+    The new observation space looks as follows (a little quirky, but this is
+    for testing purposes only):
+
+    gym.spaces.Dict({
+        "x-pos": [x-pos],
+        "angular-pos": gym.spaces.Dict({"test": [angular-pos]}),
+        "velocs": gym.spaces.Tuple([x-veloc, angular-veloc]),
+    })
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        # Fix our observation-space as described above.
+        low = self.observation_space.low
+        high = self.observation_space.high
+
+        # Test as many quirks and oddities as possible: Dict, Dict inside a Dict,
+        # Tuple inside a Dict, and both (1,)-shapes as well as ()-shapes for Boxes.
+        # Also add a random discrete variable here.
+        self.observation_space = gym.spaces.Dict(
+            {
+                "x-pos": gym.spaces.Box(low[0], high[0], (1,), dtype=np.float32),
+                "angular-pos": gym.spaces.Dict(
+                    {
+                        "value": gym.spaces.Box(low[2], high[2], (), dtype=np.float32),
+                        # Add some random non-essential information.
+                        "some_random_stuff": gym.spaces.Discrete(3),
+                    }
+                ),
+                "velocs": gym.spaces.Tuple(
+                    [
+                        # x-veloc
+                        gym.spaces.Box(low[1], high[1], (1,), dtype=np.float32),
+                        # angular-veloc
+                        gym.spaces.Box(low[3], high[3], (), dtype=np.float32),
+                    ]
+                ),
+            }
+        )
+
+    def step(self, action):
+        next_obs, reward, done, truncated, info = super().step(action)
+        return self._compile_current_obs(next_obs), reward, done, truncated, info
+
+    def reset(self, *, seed=None, options=None):
+        init_obs, init_info = super().reset(seed=seed, options=options)
+        return self._compile_current_obs(init_obs), init_info
+
+    def _compile_current_obs(self, original_cartpole_obs):
+        # original_cartpole_obs is [x-pos, x-veloc, angle, angle-veloc]
+        return {
+            "x-pos": np.array([original_cartpole_obs[0]], np.float32),
+            "angular-pos": {
+                "value": np.array(original_cartpole_obs[2]),
+                "some_random_stuff": np.random.randint(3),
+            },
+            "velocs": (
+                np.array([original_cartpole_obs[1]], np.float32),
+                np.array(original_cartpole_obs[3], np.float32),
+            ),
+        }

rllib/examples/envs/classes/cartpole_with_large_observation_space.py

@@ -0,0 +1,69 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.classic_control import CartPoleEnv
+
+
+class CartPoleWithLargeObservationSpace(CartPoleEnv):
+    """CartPole gym environment that has a large dict observation space.
+
+    However, otherwise, the information content in each observation remains the same.
+
+    https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_control/cartpole.py  # noqa
+
+    The new observation space looks as follows (a little quirky, but this is
+    for testing purposes only):
+
+    gym.spaces.Dict({
+        "1": gym.spaces.Tuple((
+            gym.spaces.Discrete(100),
+            gym.spaces.Box(0, 256, shape=(30,), dtype=float32),
+        )),
+        "2": gym.spaces.Tuple((
+            gym.spaces.Discrete(100),
+            gym.spaces.Box(0, 256, shape=(30,), dtype=float32),
+        )),
+        "3": ...
+        "actual-obs": gym.spaces.Box(-inf, inf, (4,), float32),
+    })
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        # Fix our observation-space as described above.
+        low = self.observation_space.low
+        high = self.observation_space.high
+
+        # Test as many quirks and oddities as possible: Dict, Dict inside a Dict,
+        # Tuple inside a Dict, and both (1,)-shapes as well as ()-shapes for Boxes.
+        # Also add a random discrete variable here.
+        spaces = {
+            str(i): gym.spaces.Tuple(
+                (
+                    gym.spaces.Discrete(100),
+                    gym.spaces.Box(0, 256, shape=(30,), dtype=np.float32),
+                )
+            )
+            for i in range(100)
+        }
+        spaces.update(
+            {
+                "actually-useful-stuff": (
+                    gym.spaces.Box(low[0], high[0], (4,), np.float32)
+                )
+            }
+        )
+        self.observation_space = gym.spaces.Dict(spaces)
+
+    def step(self, action):
+        next_obs, reward, done, truncated, info = super().step(action)
+        return self._compile_current_obs(next_obs), reward, done, truncated, info
+
+    def reset(self, *, seed=None, options=None):
+        init_obs, init_info = super().reset(seed=seed, options=options)
+        return self._compile_current_obs(init_obs), init_info
+
+    def _compile_current_obs(self, original_cartpole_obs):
+        return {
+            str(i): self.observation_space.spaces[str(i)].sample() for i in range(100)
+        } | {"actually-useful-stuff": original_cartpole_obs}

rllib/examples/envs/classes/cartpole_with_protobuf_observation_space.py

@@ -0,0 +1,79 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.classic_control import CartPoleEnv
+
+from ray.rllib.examples.envs.classes.utils.cartpole_observations_proto import (
+    CartPoleObservation,
+)
+
+
+class CartPoleWithProtobufObservationSpace(CartPoleEnv):
+    """CartPole gym environment that has a protobuf observation space.
+
+    Sometimes, it is more performant for an environment to publish its observations
+    as a protobuf message (instead of a heavily nested Dict).
+
+    The protobuf message used here is originally defined in the
+    `./utils/cartpole_observations.proto` file. We converted this file into a python
+    importable module by compiling it with:
+
+    `protoc --python_out=. cartpole_observations.proto`
+
+    .. which yielded the `cartpole_observations_proto.py` file in the same directory
+    (we import this file's `CartPoleObservation` message here).
+
+    The new observation space is a (binary) Box(0, 255, ([len of protobuf],), uint8).
+
+    A ConnectorV2 pipeline or simpler gym.Wrapper will have to be used to convert this
+    observation format into an NN-readable (e.g. float32) 1D tensor.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        dummy_obs = self._convert_observation_to_protobuf(
+            np.array([1.0, 1.0, 1.0, 1.0])
+        )
+        bin_length = len(dummy_obs)
+        self.observation_space = gym.spaces.Box(0, 255, (bin_length,), np.uint8)
+
+    def step(self, action):
+        observation, reward, terminated, truncated, info = super().step(action)
+        proto_observation = self._convert_observation_to_protobuf(observation)
+        return proto_observation, reward, terminated, truncated, info
+
+    def reset(self, **kwargs):
+        observation, info = super().reset(**kwargs)
+        proto_observation = self._convert_observation_to_protobuf(observation)
+        return proto_observation, info
+
+    def _convert_observation_to_protobuf(self, observation):
+        x_pos, x_veloc, angle_pos, angle_veloc = observation
+
+        # Create the Protobuf message
+        cartpole_observation = CartPoleObservation()
+        cartpole_observation.x_pos = x_pos
+        cartpole_observation.x_veloc = x_veloc
+        cartpole_observation.angle_pos = angle_pos
+        cartpole_observation.angle_veloc = angle_veloc
+
+        # Serialize to binary string.
+        return np.frombuffer(cartpole_observation.SerializeToString(), np.uint8)
+
+
+if __name__ == "__main__":
+    env = CartPoleWithProtobufObservationSpace()
+    obs, info = env.reset()
+
+    # Test loading a protobuf object with data from the obs binary string
+    # (uint8 ndarray).
+    byte_str = obs.tobytes()
+    obs_protobuf = CartPoleObservation()
+    obs_protobuf.ParseFromString(byte_str)
+    print(obs_protobuf)
+
+    terminated = truncated = False
+    while not terminated and not truncated:
+        action = env.action_space.sample()
+        obs, reward, terminated, truncated, info = env.step(action)
+
+        print(obs)

rllib/examples/envs/classes/cliff_walking_wall_env.py

@@ -0,0 +1,71 @@
+import gymnasium as gym
+from gymnasium import spaces
+
+ACTION_UP = 0
+ACTION_RIGHT = 1
+ACTION_DOWN = 2
+ACTION_LEFT = 3
+
+
+class CliffWalkingWallEnv(gym.Env):
+    """Modified version of the CliffWalking environment from Farama-Foundation's
+    Gymnasium with walls instead of a cliff.
+
+    ### Description
+    The board is a 4x12 matrix, with (using NumPy matrix indexing):
+    - [3, 0] or obs==36 as the start at bottom-left
+    - [3, 11] or obs==47 as the goal at bottom-right
+    - [3, 1..10] or obs==37...46 as the cliff at bottom-center
+
+    An episode terminates when the agent reaches the goal.
+
+    ### Actions
+    There are 4 discrete deterministic actions:
+    - 0: move up
+    - 1: move right
+    - 2: move down
+    - 3: move left
+    You can also use the constants ACTION_UP, ACTION_RIGHT, ... defined above.
+
+    ### Observations
+    There are 3x12 + 2 possible states, not including the walls. If an action
+    would move an agent into one of the walls, it simply stays in the same position.
+
+    ### Reward
+    Each time step incurs -1 reward, except reaching the goal which gives +10 reward.
+    """
+
+    def __init__(self, seed=42) -> None:
+        self.observation_space = spaces.Discrete(48)
+        self.action_space = spaces.Discrete(4)
+        self.observation_space.seed(seed)
+        self.action_space.seed(seed)
+
+    def reset(self, *, seed=None, options=None):
+        self.position = 36
+        return self.position, {}
+
+    def step(self, action):
+        x = self.position // 12
+        y = self.position % 12
+        # UP
+        if action == ACTION_UP:
+            x = max(x - 1, 0)
+        # RIGHT
+        elif action == ACTION_RIGHT:
+            if self.position != 36:
+                y = min(y + 1, 11)
+        # DOWN
+        elif action == ACTION_DOWN:
+            if self.position < 25 or self.position > 34:
+                x = min(x + 1, 3)
+        # LEFT
+        elif action == ACTION_LEFT:
+            if self.position != 47:
+                y = max(y - 1, 0)
+        else:
+            raise ValueError(f"action {action} not in {self.action_space}")
+        self.position = x * 12 + y
+        done = self.position == 47
+        reward = -1 if not done else 10
+        return self.position, reward, done, False, {}

rllib/examples/envs/classes/correlated_actions_env.py

@@ -0,0 +1,79 @@
+from typing import Any, Dict, Optional
+
+import gymnasium as gym
+import numpy as np
+
+
+class CorrelatedActionsEnv(gym.Env):
+    """Environment that can only be solved through an autoregressive action model.
+
+    In each step, the agent observes a random number (between -1 and 1) and has
+    to choose two actions, a1 (discrete, 0, 1, or 2) and a2 (cont. between -1 and 1).
+
+    The reward is constructed such that actions need to be correlated to succeed. It's
+    impossible for the network to learn each action head separately.
+
+    There are two reward components:
+    The first is the negative absolute value of the delta between 1.0 and the sum of
+    obs + a1. For example, if obs is -0.3 and a1 was sampled to be 1, then the value of
+    the first reward component is:
+    r1 = -abs(1.0 - [obs+a1]) = -abs(1.0 - (-0.3 + 1)) = -abs(0.3) = -0.3
+    The second reward component is computed as the negative absolute value
+    of `obs + a1 + a2`. For example, if obs is 0.5, a1 was sampled to be 0,
+    and a2 was sampled to be -0.7, then the value of the second reward component is:
+    r2 = -abs(obs + a1 + a2) = -abs(0.5 + 0 - 0.7)) = -abs(-0.2) = -0.2
+
+    Because of this specific reward function, the agent must learn to optimally sample
+    a1 based on the observation and to optimally sample a2, based on the observation
+    AND the sampled value of a1.
+
+    One way to effectively learn this is through correlated action
+    distributions, e.g., in examples/actions/auto_regressive_actions.py
+
+    The game ends after the first step.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        # Observation space (single continuous value between -1. and 1.).
+        self.observation_space = gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32)
+
+        # Action space (discrete action a1 and continuous action a2).
+        self.action_space = gym.spaces.Tuple(
+            [gym.spaces.Discrete(3), gym.spaces.Box(-2.0, 2.0, (1,), np.float32)]
+        )
+
+        # Internal state for the environment (e.g., could represent a factor
+        # influencing the relationship)
+        self.obs = None
+
+    def reset(
+        self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None
+    ):
+        """Reset the environment to an initial state."""
+        super().reset(seed=seed, options=options)
+
+        # Randomly initialize the observation between -1 and 1.
+        self.obs = np.random.uniform(-1, 1, size=(1,))
+
+        return self.obs, {}
+
+    def step(self, action):
+        """Apply the autoregressive action and return step information."""
+
+        # Extract individual action components, a1 and a2.
+        a1, a2 = action
+        a2 = a2[0]  # dissolve shape=(1,)
+
+        # r1 depends on how well a1 is aligned to obs:
+        r1 = -abs(1.0 - (self.obs[0] + a1))
+        # r2 depends on how well a2 is aligned to both, obs and a1.
+        r2 = -abs(self.obs[0] + a1 + a2)
+
+        reward = r1 + r2
+
+        # Optionally: add some noise or complexity to the reward function
+        # reward += np.random.normal(0, 0.01)  # Small noise can be added
+
+        # Terminate after each step (no episode length in this simple example)
+        return self.obs, reward, True, False, {}

rllib/examples/envs/classes/d4rl_env.py

@@ -0,0 +1,46 @@
+"""
+8 Environments from D4RL Environment.
+Use fully qualified class-path in your configs:
+e.g. "env": "ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_random".
+"""
+
+import gymnasium as gym
+
+try:
+    import d4rl
+
+    _ = d4rl.__name__  # Fool LINTer.
+except ImportError:
+    d4rl = None
+
+
+def halfcheetah_random():
+    return gym.make("halfcheetah-random-v0")
+
+
+def halfcheetah_medium():
+    return gym.make("halfcheetah-medium-v0")
+
+
+def halfcheetah_expert():
+    return gym.make("halfcheetah-expert-v0")
+
+
+def halfcheetah_medium_replay():
+    return gym.make("halfcheetah-medium-replay-v0")
+
+
+def hopper_random():
+    return gym.make("hopper-random-v0")
+
+
+def hopper_medium():
+    return gym.make("hopper-medium-v0")
+
+
+def hopper_expert():
+    return gym.make("hopper-expert-v0")
+
+
+def hopper_medium_replay():
+    return gym.make("hopper-medium-replay-v0")

rllib/examples/envs/classes/debug_counter_env.py

@@ -0,0 +1,92 @@
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class DebugCounterEnv(gym.Env):
+    """Simple Env that yields a ts counter as observation (0-based).
+
+    Actions have no effect.
+    The episode length is always 15.
+    Reward is always: current ts % 3.
+    """
+
+    def __init__(self, config=None):
+        config = config or {}
+        self.action_space = gym.spaces.Discrete(2)
+        self.observation_space = gym.spaces.Box(0, 100, (1,), dtype=np.float32)
+        self.start_at_t = int(config.get("start_at_t", 0))
+        self.i = self.start_at_t
+
+    def reset(self, *, seed=None, options=None):
+        self.i = self.start_at_t
+        return self._get_obs(), {}
+
+    def step(self, action):
+        self.i += 1
+        terminated = False
+        truncated = self.i >= 15 + self.start_at_t
+        return self._get_obs(), float(self.i % 3), terminated, truncated, {}
+
+    def _get_obs(self):
+        return np.array([self.i], dtype=np.float32)
+
+
+class MultiAgentDebugCounterEnv(MultiAgentEnv):
+    def __init__(self, config):
+        super().__init__()
+        self.num_agents = config["num_agents"]
+        self.base_episode_len = config.get("base_episode_len", 103)
+
+        # Observation dims:
+        # 0=agent ID.
+        # 1=episode ID (0.0 for obs after reset).
+        # 2=env ID (0.0 for obs after reset).
+        # 3=ts (of the agent).
+        self.observation_space = gym.spaces.Dict(
+            {
+                aid: gym.spaces.Box(float("-inf"), float("inf"), (4,))
+                for aid in range(self.num_agents)
+            }
+        )
+
+        # Actions are always:
+        # (episodeID, envID) as floats.
+        self.action_space = gym.spaces.Dict(
+            {
+                aid: gym.spaces.Box(-float("inf"), float("inf"), shape=(2,))
+                for aid in range(self.num_agents)
+            }
+        )
+
+        self.timesteps = [0] * self.num_agents
+        self.terminateds = set()
+        self.truncateds = set()
+
+    def reset(self, *, seed=None, options=None):
+        self.timesteps = [0] * self.num_agents
+        self.terminateds = set()
+        self.truncateds = set()
+        return {
+            i: np.array([i, 0.0, 0.0, 0.0], dtype=np.float32)
+            for i in range(self.num_agents)
+        }, {}
+
+    def step(self, action_dict):
+        obs, rew, terminated, truncated = {}, {}, {}, {}
+        for i, action in action_dict.items():
+            self.timesteps[i] += 1
+            obs[i] = np.array([i, action[0], action[1], self.timesteps[i]])
+            rew[i] = self.timesteps[i] % 3
+            terminated[i] = False
+            truncated[i] = (
+                True if self.timesteps[i] > self.base_episode_len + i else False
+            )
+            if terminated[i]:
+                self.terminateds.add(i)
+            if truncated[i]:
+                self.truncateds.add(i)
+        terminated["__all__"] = len(self.terminateds) == self.num_agents
+        truncated["__all__"] = len(self.truncateds) == self.num_agents
+        return obs, rew, terminated, truncated, {}

rllib/examples/envs/classes/deterministic_envs.py

@@ -0,0 +1,13 @@
+import gymnasium as gym
+
+
+def create_cartpole_deterministic(config):
+    env = gym.make("CartPole-v1")
+    env.reset(seed=config.get("seed", 0))
+    return env
+
+
+def create_pendulum_deterministic(config):
+    env = gym.make("Pendulum-v1")
+    env.reset(seed=config.get("seed", 0))
+    return env

rllib/examples/envs/classes/dm_control_suite.py

@@ -0,0 +1,131 @@
+from ray.rllib.env.wrappers.dm_control_wrapper import DMCEnv
+
+"""
+8 Environments from Deepmind Control Suite
+"""
+
+
+def acrobot_swingup(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "acrobot",
+        "swingup",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def walker_walk(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "walker",
+        "walk",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def hopper_hop(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "hopper",
+        "hop",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def hopper_stand(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "hopper",
+        "stand",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def cheetah_run(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "cheetah",
+        "run",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def walker_run(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "walker",
+        "run",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def pendulum_swingup(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "pendulum",
+        "swingup",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def cartpole_swingup(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "cartpole",
+        "swingup",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )
+
+
+def humanoid_walk(
+    from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True
+):
+    return DMCEnv(
+        "humanoid",
+        "walk",
+        from_pixels=from_pixels,
+        height=height,
+        width=width,
+        frame_skip=frame_skip,
+        channels_first=channels_first,
+    )

rllib/examples/envs/classes/env_using_remote_actor.py

@@ -0,0 +1,63 @@
+"""
+Example of an environment that uses a named remote actor as parameter
+server.
+
+"""
+from gymnasium.envs.classic_control.cartpole import CartPoleEnv
+from gymnasium.utils import seeding
+
+import ray
+
+
+@ray.remote
+class ParameterStorage:
+    def get_params(self, rng):
+        return {
+            "MASSCART": rng.uniform(low=0.5, high=2.0),
+        }
+
+
+class CartPoleWithRemoteParamServer(CartPoleEnv):
+    """CartPoleMassEnv varies the weights of the cart and the pole."""
+
+    def __init__(self, env_config):
+        self.env_config = env_config
+        super().__init__()
+        # Get our param server (remote actor) by name.
+        self._handler = ray.get_actor(env_config.get("param_server", "param-server"))
+        self.rng_seed = None
+        self.np_random, _ = seeding.np_random(self.rng_seed)
+
+    def reset(self, *, seed=None, options=None):
+        if seed is not None:
+            self.rng_seed = int(seed)
+            self.np_random, _ = seeding.np_random(seed)
+            print(
+                f"Seeding env (worker={self.env_config.worker_index}) " f"with {seed}"
+            )
+
+        # Pass in our RNG to guarantee no race conditions.
+        # If `self._handler` had its own RNG, this may clash with other
+        # envs trying to use the same param-server.
+        params = ray.get(self._handler.get_params.remote(self.np_random))
+
+        # IMPORTANT: Advance the state of our RNG (self._rng was passed
+        # above via ray (serialized) and thus not altered locally here!).
+        # Or create a new RNG from another random number:
+        # Seed the RNG with a deterministic seed if set, otherwise, create
+        # a random one.
+        new_seed = int(
+            self.np_random.integers(0, 1000000) if not self.rng_seed else self.rng_seed
+        )
+        self.np_random, _ = seeding.np_random(new_seed)
+
+        print(
+            f"Env worker-idx={self.env_config.worker_index} "
+            f"mass={params['MASSCART']}"
+        )
+
+        self.masscart = params["MASSCART"]
+        self.total_mass = self.masspole + self.masscart
+        self.polemass_length = self.masspole * self.length
+
+        return super().reset()

rllib/examples/envs/classes/env_with_subprocess.py

@@ -0,0 +1,43 @@
+import atexit
+import os
+import subprocess
+
+import gymnasium as gym
+from gymnasium.spaces import Discrete
+
+
+class EnvWithSubprocess(gym.Env):
+    """An env that spawns a subprocess."""
+
+    # Dummy command to run as a subprocess with a unique name
+    UNIQUE_CMD = "sleep 20"
+
+    def __init__(self, config):
+        self.UNIQUE_FILE_0 = config["tmp_file1"]
+        self.UNIQUE_FILE_1 = config["tmp_file2"]
+        self.UNIQUE_FILE_2 = config["tmp_file3"]
+        self.UNIQUE_FILE_3 = config["tmp_file4"]
+
+        self.action_space = Discrete(2)
+        self.observation_space = Discrete(2)
+        # Subprocess that should be cleaned up.
+        self.subproc = subprocess.Popen(self.UNIQUE_CMD.split(" "), shell=False)
+        self.config = config
+        # Exit handler should be called.
+        atexit.register(lambda: self.subproc.kill())
+        if config.worker_index == 0:
+            atexit.register(lambda: os.unlink(self.UNIQUE_FILE_0))
+        else:
+            atexit.register(lambda: os.unlink(self.UNIQUE_FILE_1))
+
+    def close(self):
+        if self.config.worker_index == 0:
+            os.unlink(self.UNIQUE_FILE_2)
+        else:
+            os.unlink(self.UNIQUE_FILE_3)
+
+    def reset(self, *, seed=None, options=None):
+        return 0, {}
+
+    def step(self, action):
+        return 0, 0, True, False, {}

rllib/examples/envs/classes/fast_image_env.py

@@ -0,0 +1,20 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+
+class FastImageEnv(gym.Env):
+    def __init__(self, config):
+        self.zeros = np.zeros((84, 84, 4))
+        self.action_space = Discrete(2)
+        self.observation_space = Box(0.0, 1.0, shape=(84, 84, 4), dtype=np.float32)
+        self.i = 0
+
+    def reset(self, *, seed=None, options=None):
+        self.i = 0
+        return self.zeros, {}
+
+    def step(self, action):
+        self.i += 1
+        done = truncated = self.i > 1000
+        return self.zeros, 1, done, truncated, {}

rllib/examples/envs/classes/gpu_requiring_env.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import ray
+from ray.rllib.examples.envs.classes.simple_corridor import SimpleCorridor
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+class GPURequiringEnv(SimpleCorridor):
+    """A dummy env that requires a GPU in order to work.
+
+    The env here is a simple corridor env that additionally simulates a GPU
+    check in its constructor via `ray.get_gpu_ids()`. If this returns an
+    empty list, we raise an error.
+
+    To make this env work, use `num_gpus_per_env_runner > 0` (RolloutWorkers
+    requesting this many GPUs each) and - maybe - `num_gpus > 0` in case
+    your local worker/driver must have an env as well. However, this is
+    only the case if `create_local_env_runner`=True (default is False).
+    """
+
+    def __init__(self, config=None):
+        super().__init__(config)
+
+        # Fake-require some GPUs (at least one).
+        # If your local worker's env (`create_local_env_runner`=True) does not
+        # necessarily require a GPU, you can perform the below assertion only
+        # if `config.worker_index != 0`.
+        gpus_available = ray.get_gpu_ids()
+        print(f"{type(self).__name__} can see GPUs={gpus_available}")
+
+        # Create a dummy tensor on the GPU.
+        if len(gpus_available) > 0 and torch:
+            self._tensor = torch.from_numpy(np.random.random_sample(size=(42, 42))).to(
+                f"cuda:{gpus_available[0]}"
+            )

rllib/examples/envs/classes/look_and_push.py

@@ -0,0 +1,65 @@
+import gymnasium as gym
+import numpy as np
+
+
+class LookAndPush(gym.Env):
+    """Memory-requiring Env: Best sequence of actions depends on prev. states.
+
+    Optimal behavior:
+        0) a=0 -> observe next state (s'), which is the "hidden" state.
+            If a=1 here, the hidden state is not observed.
+        1) a=1 to always jump to s=2 (not matter what the prev. state was).
+        2) a=1 to move to s=3.
+        3) a=1 to move to s=4.
+        4) a=0 OR 1 depending on s' observed after 0): +10 reward and done.
+            otherwise: -10 reward and done.
+    """
+
+    def __init__(self):
+        self.action_space = gym.spaces.Discrete(2)
+        self.observation_space = gym.spaces.Discrete(5)
+        self._state = None
+        self._case = None
+
+    def reset(self, *, seed=None, options=None):
+        self._state = 2
+        self._case = np.random.choice(2)
+        return self._state, {}
+
+    def step(self, action):
+        assert self.action_space.contains(action)
+
+        if self._state == 4:
+            if action and self._case:
+                return self._state, 10.0, True, {}
+            else:
+                return self._state, -10, True, {}
+        else:
+            if action:
+                if self._state == 0:
+                    self._state = 2
+                else:
+                    self._state += 1
+            elif self._state == 2:
+                self._state = self._case
+
+        return self._state, -1, False, False, {}
+
+
+class OneHot(gym.Wrapper):
+    def __init__(self, env):
+        super(OneHot, self).__init__(env)
+        self.observation_space = gym.spaces.Box(0.0, 1.0, (env.observation_space.n,))
+
+    def reset(self, *, seed=None, options=None):
+        obs, info = self.env.reset(seed=seed, options=options)
+        return self._encode_obs(obs), info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return self._encode_obs(obs), reward, terminated, truncated, info
+
+    def _encode_obs(self, obs):
+        new_obs = np.ones(self.env.observation_space.n)
+        new_obs[obs] = 1.0
+        return new_obs

rllib/examples/envs/classes/memory_leaking_env.py

@@ -0,0 +1,35 @@
+import logging
+import uuid
+
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+from ray.rllib.utils.annotations import override
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryLeakingEnv(RandomEnv):
+    """An env that leaks very little memory.
+
+    Useful for proving that our memory-leak tests can catch the
+    slightest leaks.
+    """
+
+    def __init__(self, config=None):
+        super().__init__(config)
+        self._leak = {}
+        self._steps_after_reset = 0
+
+    @override(RandomEnv)
+    def reset(self, *, seed=None, options=None):
+        self._steps_after_reset = 0
+        return super().reset(seed=seed, options=options)
+
+    @override(RandomEnv)
+    def step(self, action):
+        self._steps_after_reset += 1
+
+        # Only leak once an episode.
+        if self._steps_after_reset == 2:
+            self._leak[uuid.uuid4().hex.upper()] = 1
+
+        return super().step(action)

rllib/examples/envs/classes/mock_env.py

@@ -0,0 +1,221 @@
+from typing import Optional
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.env.vector_env import VectorEnv
+from ray.rllib.utils.annotations import override
+
+
+class MockEnv(gym.Env):
+    """Mock environment for testing purposes.
+
+    Observation=0, reward=1.0, episode-len is configurable.
+    Actions are ignored.
+    """
+
+    def __init__(self, episode_length, config=None):
+        self.episode_length = episode_length
+        self.config = config
+        self.i = 0
+        self.observation_space = gym.spaces.Discrete(1)
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self, *, seed=None, options=None):
+        self.i = 0
+        return 0, {}
+
+    def step(self, action):
+        self.i += 1
+        terminated = truncated = self.i >= self.episode_length
+        return 0, 1.0, terminated, truncated, {}
+
+
+class MockEnv2(gym.Env):
+    """Mock environment for testing purposes.
+
+    Observation=ts (discrete space!), reward=100.0, episode-len is
+    configurable. Actions are ignored.
+    """
+
+    metadata = {
+        "render.modes": ["rgb_array"],
+    }
+    render_mode: Optional[str] = "rgb_array"
+
+    def __init__(self, episode_length):
+        self.episode_length = episode_length
+        self.i = 0
+        self.observation_space = gym.spaces.Discrete(self.episode_length + 1)
+        self.action_space = gym.spaces.Discrete(2)
+        self.rng_seed = None
+
+    def reset(self, *, seed=None, options=None):
+        self.i = 0
+        if seed is not None:
+            self.rng_seed = seed
+        return self.i, {}
+
+    def step(self, action):
+        self.i += 1
+        terminated = truncated = self.i >= self.episode_length
+        return self.i, 100.0, terminated, truncated, {}
+
+    def render(self):
+        # Just generate a random image here for demonstration purposes.
+        # Also see `gym/envs/classic_control/cartpole.py` for
+        # an example on how to use a Viewer object.
+        return np.random.randint(0, 256, size=(300, 400, 3), dtype=np.uint8)
+
+
+class MockEnv3(gym.Env):
+    """Mock environment for testing purposes.
+
+    Observation=ts (discrete space!), reward=100.0, episode-len is
+    configurable. Actions are ignored.
+    """
+
+    def __init__(self, episode_length):
+        self.episode_length = episode_length
+        self.i = 0
+        self.observation_space = gym.spaces.Discrete(100)
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self, *, seed=None, options=None):
+        self.i = 0
+        return self.i, {"timestep": 0}
+
+    def step(self, action):
+        self.i += 1
+        terminated = truncated = self.i >= self.episode_length
+        return self.i, self.i, terminated, truncated, {"timestep": self.i}
+
+
+class VectorizedMockEnv(VectorEnv):
+    """Vectorized version of the MockEnv.
+
+    Contains `num_envs` MockEnv instances, each one having its own
+    `episode_length` horizon.
+    """
+
+    def __init__(self, episode_length, num_envs):
+        super().__init__(
+            observation_space=gym.spaces.Discrete(1),
+            action_space=gym.spaces.Discrete(2),
+            num_envs=num_envs,
+        )
+        self.envs = [MockEnv(episode_length) for _ in range(num_envs)]
+
+    @override(VectorEnv)
+    def vector_reset(self, *, seeds=None, options=None):
+        seeds = seeds or [None] * self.num_envs
+        options = options or [None] * self.num_envs
+        obs_and_infos = [
+            e.reset(seed=seeds[i], options=options[i]) for i, e in enumerate(self.envs)
+        ]
+        return [oi[0] for oi in obs_and_infos], [oi[1] for oi in obs_and_infos]
+
+    @override(VectorEnv)
+    def reset_at(self, index, *, seed=None, options=None):
+        return self.envs[index].reset(seed=seed, options=options)
+
+    @override(VectorEnv)
+    def vector_step(self, actions):
+        obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch = (
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+        for i in range(len(self.envs)):
+            obs, rew, terminated, truncated, info = self.envs[i].step(actions[i])
+            obs_batch.append(obs)
+            rew_batch.append(rew)
+            terminated_batch.append(terminated)
+            truncated_batch.append(truncated)
+            info_batch.append(info)
+        return obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch
+
+    @override(VectorEnv)
+    def get_sub_environments(self):
+        return self.envs
+
+
+class MockVectorEnv(VectorEnv):
+    """A custom vector env that uses a single(!) CartPole sub-env.
+
+    However, this env pretends to be a vectorized one to illustrate how one
+    could create custom VectorEnvs w/o the need for actual vectorizations of
+    sub-envs under the hood.
+    """
+
+    def __init__(self, episode_length, mocked_num_envs):
+        self.env = gym.make("CartPole-v1")
+        super().__init__(
+            observation_space=self.env.observation_space,
+            action_space=self.env.action_space,
+            num_envs=mocked_num_envs,
+        )
+        self.episode_len = episode_length
+        self.ts = 0
+
+    @override(VectorEnv)
+    def vector_reset(self, *, seeds=None, options=None):
+        # Since we only have one underlying sub-environment, just use the first seed
+        # and the first options dict (the user of this env thinks, there are
+        # `self.num_envs` sub-environments and sends that many seeds/options).
+        seeds = seeds or [None]
+        options = options or [None]
+        obs, infos = self.env.reset(seed=seeds[0], options=options[0])
+        # Simply repeat the single obs/infos to pretend we really have
+        # `self.num_envs` sub-environments.
+        return (
+            [obs for _ in range(self.num_envs)],
+            [infos for _ in range(self.num_envs)],
+        )
+
+    @override(VectorEnv)
+    def reset_at(self, index, *, seed=None, options=None):
+        self.ts = 0
+        return self.env.reset(seed=seed, options=options)
+
+    @override(VectorEnv)
+    def vector_step(self, actions):
+        self.ts += 1
+        # Apply all actions sequentially to the same env.
+        # Whether this would make a lot of sense is debatable.
+        obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch = (
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+        for i in range(self.num_envs):
+            obs, rew, terminated, truncated, info = self.env.step(actions[i])
+            # Artificially truncate once time step limit has been reached.
+            # Note: Also terminate/truncate, when underlying CartPole is
+            # terminated/truncated.
+            if self.ts >= self.episode_len:
+                truncated = True
+            obs_batch.append(obs)
+            rew_batch.append(rew)
+            terminated_batch.append(terminated)
+            truncated_batch.append(truncated)
+            info_batch.append(info)
+            if terminated or truncated:
+                remaining = self.num_envs - (i + 1)
+                obs_batch.extend([obs for _ in range(remaining)])
+                rew_batch.extend([rew for _ in range(remaining)])
+                terminated_batch.extend([terminated for _ in range(remaining)])
+                truncated_batch.extend([truncated for _ in range(remaining)])
+                info_batch.extend([info for _ in range(remaining)])
+                break
+        return obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch
+
+    @override(VectorEnv)
+    def get_sub_environments(self):
+        # You may also leave this method as-is, in which case, it would
+        # return an empty list.
+        return [self.env for _ in range(self.num_envs)]

rllib/examples/envs/classes/multi_agent/__init__.py

@@ -0,0 +1,34 @@
+from ray.rllib.env.multi_agent_env import make_multi_agent
+from ray.rllib.examples.envs.classes.cartpole_with_dict_observation_space import (
+    CartPoleWithDictObservationSpace,
+)
+from ray.rllib.examples.envs.classes.multi_agent.guess_the_number_game import (
+    GuessTheNumberGame,
+)
+from ray.rllib.examples.envs.classes.multi_agent.two_step_game import (
+    TwoStepGame,
+    TwoStepGameWithGroupedAgents,
+)
+from ray.rllib.examples.envs.classes.nested_space_repeat_after_me_env import (
+    NestedSpaceRepeatAfterMeEnv,
+)
+from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole
+
+# Backward compatibility.
+__all__ = [
+    "GuessTheNumberGame",
+    "TwoStepGame",
+    "TwoStepGameWithGroupedAgents",
+]
+
+
+MultiAgentCartPole = make_multi_agent("CartPole-v1")
+MultiAgentMountainCar = make_multi_agent("MountainCarContinuous-v0")
+MultiAgentPendulum = make_multi_agent("Pendulum-v1")
+MultiAgentStatelessCartPole = make_multi_agent(lambda config: StatelessCartPole(config))
+MultiAgentCartPoleWithDictObservationSpace = make_multi_agent(
+    lambda config: CartPoleWithDictObservationSpace(config)
+)
+MultiAgentNestedSpaceRepeatAfterMeEnv = make_multi_agent(
+    lambda config: NestedSpaceRepeatAfterMeEnv(config)
+)

rllib/examples/envs/classes/multi_agent/bandit_envs_discrete.py

@@ -0,0 +1,207 @@
+import copy
+import random
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+
+class SimpleContextualBandit(gym.Env):
+    """Simple env w/ 2 states and 3 actions (arms): 0, 1, and 2.
+
+    Episodes last only for one timestep, possible observations are:
+    [-1.0, 1.0] and [1.0, -1.0], where the first element is the "current context".
+    The highest reward (+10.0) is received for selecting arm 0 for context=1.0
+    and arm 2 for context=-1.0. Action 1 always yields 0.0 reward.
+    """
+
+    def __init__(self, config=None):
+        self.action_space = Discrete(3)
+        self.observation_space = Box(low=-1.0, high=1.0, shape=(2,))
+        self.cur_context = None
+
+    def reset(self, *, seed=None, options=None):
+        self.cur_context = random.choice([-1.0, 1.0])
+        return np.array([self.cur_context, -self.cur_context]), {}
+
+    def step(self, action):
+        rewards_for_context = {
+            -1.0: [-10, 0, 10],
+            1.0: [10, 0, -10],
+        }
+        reward = rewards_for_context[self.cur_context][action]
+        return (
+            np.array([-self.cur_context, self.cur_context]),
+            reward,
+            True,
+            False,
+            {"regret": 10 - reward},
+        )
+
+
+class LinearDiscreteEnv(gym.Env):
+    """Samples data from linearly parameterized arms.
+
+    The reward for context X and arm i is given by X^T * theta_i, for some
+    latent set of parameters {theta_i : i = 1, ..., k}.
+    The thetas are sampled uniformly at random, the contexts are Gaussian,
+    and Gaussian noise is added to the rewards.
+    """
+
+    DEFAULT_CONFIG_LINEAR = {
+        "feature_dim": 8,
+        "num_actions": 4,
+        "reward_noise_std": 0.01,
+    }
+
+    def __init__(self, config=None):
+        self.config = copy.copy(self.DEFAULT_CONFIG_LINEAR)
+        if config is not None and type(config) is dict:
+            self.config.update(config)
+
+        self.feature_dim = self.config["feature_dim"]
+        self.num_actions = self.config["num_actions"]
+        self.sigma = self.config["reward_noise_std"]
+
+        self.action_space = Discrete(self.num_actions)
+        self.observation_space = Box(low=-10, high=10, shape=(self.feature_dim,))
+
+        self.thetas = np.random.uniform(-1, 1, (self.num_actions, self.feature_dim))
+        self.thetas /= np.linalg.norm(self.thetas, axis=1, keepdims=True)
+
+        self._elapsed_steps = 0
+        self._current_context = None
+
+    def _sample_context(self):
+        return np.random.normal(scale=1 / 3, size=(self.feature_dim,))
+
+    def reset(self, *, seed=None, options=None):
+        self._current_context = self._sample_context()
+        return self._current_context, {}
+
+    def step(self, action):
+        assert (
+            self._elapsed_steps is not None
+        ), "Cannot call env.step() beforecalling reset()"
+        assert action < self.num_actions, "Invalid action."
+
+        action = int(action)
+        context = self._current_context
+        rewards = self.thetas.dot(context)
+
+        opt_action = rewards.argmax()
+
+        regret = rewards.max() - rewards[action]
+
+        # Add Gaussian noise
+        rewards += np.random.normal(scale=self.sigma, size=rewards.shape)
+
+        reward = rewards[action]
+        self._current_context = self._sample_context()
+        return (
+            self._current_context,
+            reward,
+            True,
+            False,
+            {"regret": regret, "opt_action": opt_action},
+        )
+
+    def render(self, mode="human"):
+        raise NotImplementedError
+
+
+class WheelBanditEnv(gym.Env):
+    """Wheel bandit environment for 2D contexts
+    (see https://arxiv.org/abs/1802.09127).
+    """
+
+    DEFAULT_CONFIG_WHEEL = {
+        "delta": 0.5,
+        "mu_1": 1.2,
+        "mu_2": 1,
+        "mu_3": 50,
+        "std": 0.01,
+    }
+
+    feature_dim = 2
+    num_actions = 5
+
+    def __init__(self, config=None):
+        self.config = copy.copy(self.DEFAULT_CONFIG_WHEEL)
+        if config is not None and type(config) is dict:
+            self.config.update(config)
+
+        self.delta = self.config["delta"]
+        self.mu_1 = self.config["mu_1"]
+        self.mu_2 = self.config["mu_2"]
+        self.mu_3 = self.config["mu_3"]
+        self.std = self.config["std"]
+
+        self.action_space = Discrete(self.num_actions)
+        self.observation_space = Box(low=-1, high=1, shape=(self.feature_dim,))
+
+        self.means = [self.mu_1] + 4 * [self.mu_2]
+        self._elapsed_steps = 0
+        self._current_context = None
+
+    def _sample_context(self):
+        while True:
+            state = np.random.uniform(-1, 1, self.feature_dim)
+            if np.linalg.norm(state) <= 1:
+                return state
+
+    def reset(self, *, seed=None, options=None):
+        self._current_context = self._sample_context()
+        return self._current_context, {}
+
+    def step(self, action):
+        assert (
+            self._elapsed_steps is not None
+        ), "Cannot call env.step() before calling reset()"
+
+        action = int(action)
+        self._elapsed_steps += 1
+        rewards = [
+            np.random.normal(self.means[j], self.std) for j in range(self.num_actions)
+        ]
+        context = self._current_context
+        r_big = np.random.normal(self.mu_3, self.std)
+
+        if np.linalg.norm(context) >= self.delta:
+            if context[0] > 0:
+                if context[1] > 0:
+                    # First quadrant
+                    rewards[1] = r_big
+                    opt_action = 1
+                else:
+                    # Fourth quadrant
+                    rewards[4] = r_big
+                    opt_action = 4
+            else:
+                if context[1] > 0:
+                    # Second quadrant
+                    rewards[2] = r_big
+                    opt_action = 2
+                else:
+                    # Third quadrant
+                    rewards[3] = r_big
+                    opt_action = 3
+        else:
+            # Smaller region where action 0 is optimal
+            opt_action = 0
+
+        reward = rewards[action]
+
+        regret = rewards[opt_action] - reward
+
+        self._current_context = self._sample_context()
+        return (
+            self._current_context,
+            reward,
+            True,
+            False,
+            {"regret": regret, "opt_action": opt_action},
+        )
+
+    def render(self, mode="human"):
+        raise NotImplementedError

rllib/examples/envs/classes/multi_agent/bandit_envs_recommender_system.py

@@ -0,0 +1,228 @@
+"""Examples for recommender system simulating envs ready to be used by RLlib Algorithms.
+
+This env follows RecSim obs and action APIs.
+"""
+from typing import Optional
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.utils.numpy import softmax
+
+
+class ParametricRecSys(gym.Env):
+    """A recommendation environment which generates items with visible features
+    randomly (parametric actions).
+    The environment can be configured to be multi-user, i.e. different models
+    will be learned independently for each user, by setting num_users_in_db
+    parameter.
+    To enable slate recommendation, the `slate_size` config parameter can be
+    set as > 1.
+    """
+
+    def __init__(
+        self,
+        embedding_size: int = 20,
+        num_docs_to_select_from: int = 10,
+        slate_size: int = 1,
+        num_docs_in_db: Optional[int] = None,
+        num_users_in_db: Optional[int] = None,
+        user_time_budget: float = 60.0,
+    ):
+        """Initializes a ParametricRecSys instance.
+
+        Args:
+            embedding_size: Embedding size for both users and docs.
+                Each value in the user/doc embeddings can have values between
+                -1.0 and 1.0.
+            num_docs_to_select_from: The number of documents to present to the
+                agent each timestep. The agent will then have to pick a slate
+                out of these.
+            slate_size: The size of the slate to recommend to the user at each
+                timestep.
+            num_docs_in_db: The total number of documents in the DB. Set this
+                to None, in case you would like to resample docs from an
+                infinite pool.
+            num_users_in_db: The total number of users in the DB. Set this to
+                None, in case you would like to resample users from an infinite
+                pool.
+            user_time_budget: The total time budget a user has throughout an
+                episode. Once this time budget is used up (through engagements
+                with clicked/selected documents), the episode ends.
+        """
+        self.embedding_size = embedding_size
+        self.num_docs_to_select_from = num_docs_to_select_from
+        self.slate_size = slate_size
+
+        self.num_docs_in_db = num_docs_in_db
+        self.docs_db = None
+        self.num_users_in_db = num_users_in_db
+        self.users_db = None
+        self.current_user = None
+
+        self.user_time_budget = user_time_budget
+        self.current_user_budget = user_time_budget
+
+        self.observation_space = gym.spaces.Dict(
+            {
+                # The D docs our agent sees at each timestep.
+                # It has to select a k-slate out of these.
+                "doc": gym.spaces.Dict(
+                    {
+                        str(i): gym.spaces.Box(
+                            -1.0, 1.0, shape=(self.embedding_size,), dtype=np.float32
+                        )
+                        for i in range(self.num_docs_to_select_from)
+                    }
+                ),
+                # The user engaging in this timestep/episode.
+                "user": gym.spaces.Box(
+                    -1.0, 1.0, shape=(self.embedding_size,), dtype=np.float32
+                ),
+                # For each item in the previous slate, was it clicked?
+                # If yes, how long was it being engaged with (e.g. watched)?
+                "response": gym.spaces.Tuple(
+                    [
+                        gym.spaces.Dict(
+                            {
+                                # Clicked or not?
+                                "click": gym.spaces.Discrete(2),
+                                # Engagement time (how many minutes watched?).
+                                "engagement": gym.spaces.Box(
+                                    0.0, 100.0, shape=(), dtype=np.float32
+                                ),
+                            }
+                        )
+                        for _ in range(self.slate_size)
+                    ]
+                ),
+            }
+        )
+        # Our action space is
+        self.action_space = gym.spaces.MultiDiscrete(
+            [self.num_docs_to_select_from for _ in range(self.slate_size)]
+        )
+
+    def _get_embedding(self):
+        return np.random.uniform(-1, 1, size=(self.embedding_size,)).astype(np.float32)
+
+    def reset(self, *, seed=None, options=None):
+        # Reset the current user's time budget.
+        self.current_user_budget = self.user_time_budget
+
+        # Sample a user for the next episode/session.
+        # Pick from a only-once-sampled user DB.
+        if self.num_users_in_db is not None:
+            if self.users_db is None:
+                self.users_db = [
+                    self._get_embedding() for _ in range(self.num_users_in_db)
+                ]
+            self.current_user = self.users_db[np.random.choice(self.num_users_in_db)]
+        # Pick from an infinite pool of users.
+        else:
+            self.current_user = self._get_embedding()
+
+        return self._get_obs(), {}
+
+    def step(self, action):
+        # Action is the suggested slate (indices of the docs in the
+        # suggested ones).
+
+        # We calculate scores as the dot product between document features and user
+        # features. The softmax ensures regret<1 further down.
+        scores = softmax(
+            [np.dot(self.current_user, doc) for doc in self.currently_suggested_docs]
+        )
+        best_reward = np.max(scores)
+
+        # User choice model: User picks a doc stochastically,
+        # where probs are dot products between user- and doc feature
+        # (categories) vectors (rewards).
+        # There is also a no-click doc whose weight is 0.0.
+        user_doc_overlaps = np.array([scores[a] for a in action] + [0.0])
+        # We have to softmax again so that probabilities add up to 1
+        probabilities = softmax(user_doc_overlaps)
+        which_clicked = np.random.choice(
+            np.arange(self.slate_size + 1), p=probabilities
+        )
+
+        reward = 0.0
+        if which_clicked < self.slate_size:
+            # Reward is 1.0 - regret if clicked. 0.0 if not clicked.
+            regret = best_reward - user_doc_overlaps[which_clicked]
+            # The reward also represents the user engagement that we define to be
+            # withing the range [0...100].
+            reward = (1 - regret) * 100
+            # If anything clicked, deduct from the current user's time budget.
+            self.current_user_budget -= 1.0
+        done = truncated = self.current_user_budget <= 0.0
+
+        # Compile response.
+        response = tuple(
+            {
+                "click": int(idx == which_clicked),
+                "engagement": reward if idx == which_clicked else 0.0,
+            }
+            for idx in range(len(user_doc_overlaps) - 1)
+        )
+
+        return self._get_obs(response=response), reward, done, truncated, {}
+
+    def _get_obs(self, response=None):
+        # Sample D docs from infinity or our pre-existing docs.
+        # Pick from a only-once-sampled docs DB.
+        if self.num_docs_in_db is not None:
+            if self.docs_db is None:
+                self.docs_db = [
+                    self._get_embedding() for _ in range(self.num_docs_in_db)
+                ]
+            self.currently_suggested_docs = [
+                self.docs_db[doc_idx].astype(np.float32)
+                for doc_idx in np.random.choice(
+                    self.num_docs_in_db,
+                    size=(self.num_docs_to_select_from,),
+                    replace=False,
+                )
+            ]
+        # Pick from an infinite pool of docs.
+        else:
+            self.currently_suggested_docs = [
+                self._get_embedding() for _ in range(self.num_docs_to_select_from)
+            ]
+
+        doc = {str(i): d for i, d in enumerate(self.currently_suggested_docs)}
+
+        if not response:
+            response = self.observation_space["response"].sample()
+
+        return {
+            "user": self.current_user.astype(np.float32),
+            "doc": doc,
+            "response": response,
+        }
+
+
+if __name__ == "__main__":
+    """Test RecommSys env with random actions for baseline performance."""
+    env = ParametricRecSys(
+        num_docs_in_db=100,
+        num_users_in_db=1,
+    )
+    obs, info = env.reset()
+    num_episodes = 0
+    episode_rewards = []
+    episode_reward = 0.0
+
+    while num_episodes < 100:
+        action = env.action_space.sample()
+        obs, reward, done, truncated, _ = env.step(action)
+
+        episode_reward += reward
+        if done:
+            print(f"episode reward = {episode_reward}")
+            env.reset()
+            num_episodes += 1
+            episode_rewards.append(episode_reward)
+            episode_reward = 0.0
+
+    print(f"Avg reward={np.mean(episode_rewards)}")

rllib/examples/envs/classes/multi_agent/double_row_corridor_env.py

@@ -0,0 +1,146 @@
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import AgentID
+
+
+class DoubleRowCorridorEnv(MultiAgentEnv):
+    """A MultiAgentEnv with a single, global observation space for all agents.
+
+    There are two agents in this grid-world-style environment, `agent_0` and `agent_1`.
+    The grid has two-rows and multiple columns and agents must, each
+    separately, reach their individual goal position to receive a final reward of +10:
+
+    +---------------+
+    |0              |
+    |              1|
+    +---------------+
+    Legend:
+    0 = agent_0 + goal state for agent_1
+    1 = agent_1 + goal state for agent_0
+
+    You can change the length of the grid through providing the "length" key in the
+    `config` dict passed to the env's constructor.
+
+    The action space for both agents is Discrete(4), which encodes to moving up, down,
+    left, or right in the grid.
+
+    If the two agents collide, meaning they end up in the exact same field after both
+    taking their actions at any timestep, an additional reward of +5 is given to both
+    agents. Thus, optimal policies aim at seeking the respective other agent first, and
+    only then proceeding to their agent's goal position.
+
+    Each agent in the env has an observation space of a 2-tuple containing its own
+    x/y-position, where x is the row index, being either 0 (1st row) or 1 (2nd row),
+    and y is the column index (starting from 0).
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        config = config or {}
+
+        self.length = config.get("length", 15)
+        self.terminateds = {}
+        self.collided = False
+
+        # Provide information about agents and possible agents.
+        self.agents = self.possible_agents = ["agent_0", "agent_1"]
+        self.terminateds = {}
+
+        # Observations: x/y, where the first number is the row index, the second number
+        # is the column index, for both agents.
+        # For example: [0.0, 2.0] means the agent is in row 0 and column 2.
+        self._obs_spaces = gym.spaces.Box(
+            0.0, self.length - 1, shape=(2,), dtype=np.int32
+        )
+        self._act_spaces = gym.spaces.Discrete(4)
+
+    @override(MultiAgentEnv)
+    def reset(self, *, seed=None, options=None):
+        self.agent_pos = {
+            "agent_0": [0, 0],
+            "agent_1": [1, self.length - 1],
+        }
+        self.goals = {
+            "agent_0": [0, self.length - 1],
+            "agent_1": [1, 0],
+        }
+        self.terminateds = {agent_id: False for agent_id in self.agent_pos}
+        self.collided = False
+
+        return self._get_obs(), {}
+
+    @override(MultiAgentEnv)
+    def step(self, action_dict):
+        rewards = {
+            agent_id: -0.1
+            for agent_id in self.agent_pos
+            if not self.terminateds[agent_id]
+        }
+
+        for agent_id, action in action_dict.items():
+            row, col = self.agent_pos[agent_id]
+
+            # up
+            if action == 0 and row > 0:
+                row -= 1
+            # down
+            elif action == 1 and row < 1:
+                row += 1
+            # left
+            elif action == 2 and col > 0:
+                col -= 1
+            # right
+            elif action == 3 and col < self.length - 1:
+                col += 1
+
+            # Update positions.
+            self.agent_pos[agent_id] = [row, col]
+
+        obs = self._get_obs()
+
+        # Check for collision (only if both agents are still active).
+        if (
+            not any(self.terminateds.values())
+            and self.agent_pos["agent_0"] == self.agent_pos["agent_1"]
+        ):
+            if not self.collided:
+                rewards["agent_0"] += 5
+                rewards["agent_1"] += 5
+                self.collided = True
+
+        # Check goals.
+        for agent_id in self.agent_pos:
+            if (
+                self.agent_pos[agent_id] == self.goals[agent_id]
+                and not self.terminateds[agent_id]
+            ):
+                rewards[agent_id] += 10
+                self.terminateds[agent_id] = True
+
+        terminateds = {
+            agent_id: self.terminateds[agent_id] for agent_id in self.agent_pos
+        }
+        terminateds["__all__"] = all(self.terminateds.values())
+
+        return obs, rewards, terminateds, {}, {}
+
+    @override(MultiAgentEnv)
+    def get_observation_space(self, agent_id: AgentID) -> gym.Space:
+        return self._obs_spaces
+
+    @override(MultiAgentEnv)
+    def get_action_space(self, agent_id: AgentID) -> gym.Space:
+        return self._act_spaces
+
+    def _get_obs(self):
+        obs = {}
+        pos = self.agent_pos
+        for agent_id in self.agent_pos:
+            if self.terminateds[agent_id]:
+                continue
+            obs[agent_id] = np.array(pos[agent_id], dtype=np.int32)
+        return obs

rllib/examples/envs/classes/multi_agent/footsies/README.md

@@ -0,0 +1,10 @@
+# Footsies Environment
+
+This environment implementation is based on the [FootsiesGym project](https://github.com/chasemcd/FootsiesGym),
+specifically the version as of **July 28, 2025**.
+
+## Notes
+
+All examples in the RLlib documentation that use the Footsies environment are self-contained.
+This means that you do not need to install anything from the FootsiesGym repository or other places.
+Examples handle binary automatically (downloading, extracting, starting, stopping, etc.).

rllib/examples/envs/classes/multi_agent/footsies/encoder.py

@@ -0,0 +1,226 @@
+import collections
+import copy
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto import (
+    footsies_service_pb2 as footsies_pb2,
+)
+
+
+class FootsiesEncoder:
+    """Encoder class to generate observations from the game state"""
+
+    def __init__(self, observation_delay: int):
+        self._encoding_history = {
+            agent_id: collections.deque(maxlen=int(observation_delay))
+            for agent_id in ["p1", "p2"]
+        }
+        self.observation_delay = observation_delay
+        self._last_common_state: Optional[np.ndarray] = None
+        self._action_id_values = list(constants.FOOTSIES_ACTION_IDS.values())
+
+    @staticmethod
+    def encode_common_state(game_state: footsies_pb2.GameState) -> np.ndarray:
+        p1_state, p2_state = game_state.player1, game_state.player2
+
+        dist_x = np.abs(p1_state.player_position_x - p2_state.player_position_x) / 8.0
+
+        return np.array(
+            [
+                dist_x,
+            ],
+            dtype=np.float32,
+        )
+
+    @staticmethod
+    def _encode_input_buffer(
+        input_buffer: list[int], last_n: Optional[int] = None
+    ) -> np.ndarray:
+        """Encodes the input buffer into a one-hot vector.
+
+        :param input_buffer: The input buffer to encode
+        :type input_buffer: list[int]
+        :return: The encoded one-hot vector
+        :rtype: np.ndarray
+        """
+
+        if last_n is not None:
+            input_buffer = input_buffer[last_n:]
+
+        ib_encoding = []
+        for action_id in input_buffer:
+            arr = [0] * (len(constants.ACTION_TO_BITS) + 1)
+            arr[action_id] = 1
+            ib_encoding.extend(arr)
+
+        input_buffer_vector = np.asarray(ib_encoding, dtype=np.float32)
+
+        return input_buffer_vector
+
+    def encode(
+        self,
+        game_state: footsies_pb2.GameState,
+    ) -> dict[str, Any]:
+        """Encodes the game state into observations for all agents.
+
+        :param game_state: The game state to encode
+        :type game_state: footsies_pb2.GameState
+        :return: The encoded observations for all agents.
+        :rtype: dict[str, Any]
+        """
+        common_state = self.encode_common_state(game_state)
+        p1_encoding = self.encode_player_state(game_state.player1)
+        p2_encoding = self.encode_player_state(game_state.player2)
+
+        observation_delay = min(
+            self.observation_delay, len(self._encoding_history["p1"])
+        )
+
+        if observation_delay > 0:
+            p1_delayed_encoding = self._encoding_history["p1"][-observation_delay]
+            p2_delayed_encoding = self._encoding_history["p2"][-observation_delay]
+        else:
+            p1_delayed_encoding = copy.deepcopy(p1_encoding)
+            p2_delayed_encoding = copy.deepcopy(p2_encoding)
+
+        self._encoding_history["p1"].append(p1_encoding)
+        self._encoding_history["p2"].append(p2_encoding)
+        self._last_common_state = common_state
+
+        # Create features dictionary
+        features = {}
+        current_index = 0
+
+        # Common state
+        features["common_state"] = {
+            "start": current_index,
+            "length": len(common_state),
+        }
+        current_index += len(common_state)
+
+        # Concatenate the observations for the undelayed encoding
+        p1_encoding = np.hstack(list(p1_encoding.values()), dtype=np.float32)
+        p2_encoding = np.hstack(list(p2_encoding.values()), dtype=np.float32)
+
+        # Concatenate the observations for the delayed encoding
+        p1_delayed_encoding = np.hstack(
+            list(p1_delayed_encoding.values()), dtype=np.float32
+        )
+        p2_delayed_encoding = np.hstack(
+            list(p2_delayed_encoding.values()), dtype=np.float32
+        )
+
+        p1_centric_observation = np.hstack(
+            [common_state, p1_encoding, p2_delayed_encoding]
+        )
+
+        p2_centric_observation = np.hstack(
+            [common_state, p2_encoding, p1_delayed_encoding]
+        )
+
+        return {"p1": p1_centric_observation, "p2": p2_centric_observation}
+
+    def encode_player_state(
+        self,
+        player_state: footsies_pb2.PlayerState,
+    ) -> dict[str, Union[int, float, list, np.ndarray]]:
+        """Encodes the player state into observations.
+
+        :param player_state: The player state to encode
+        :type player_state: footsies_pb2.PlayerState
+        :return: The encoded observations for the player
+        :rtype: dict[str, Any]
+        """
+        feature_dict = {
+            "player_position_x": player_state.player_position_x
+            / constants.FeatureDictNormalizers.PLAYER_POSITION_X,
+            "velocity_x": player_state.velocity_x
+            / constants.FeatureDictNormalizers.VELOCITY_X,
+            "is_dead": int(player_state.is_dead),
+            "vital_health": player_state.vital_health,
+            "guard_health": one_hot_encoder(player_state.guard_health, [0, 1, 2, 3]),
+            "current_action_id": self._encode_action_id(player_state.current_action_id),
+            "current_action_frame": player_state.current_action_frame
+            / constants.FeatureDictNormalizers.CURRENT_ACTION_FRAME,
+            "current_action_frame_count": player_state.current_action_frame_count
+            / constants.FeatureDictNormalizers.CURRENT_ACTION_FRAME_COUNT,
+            "current_action_remaining_frames": (
+                player_state.current_action_frame_count
+                - player_state.current_action_frame
+            )
+            / constants.FeatureDictNormalizers.CURRENT_ACTION_REMAINING_FRAMES,
+            "is_action_end": int(player_state.is_action_end),
+            "is_always_cancelable": int(player_state.is_always_cancelable),
+            "current_action_hit_count": player_state.current_action_hit_count,
+            "current_hit_stun_frame": player_state.current_hit_stun_frame
+            / constants.FeatureDictNormalizers.CURRENT_HIT_STUN_FRAME,
+            "is_in_hit_stun": int(player_state.is_in_hit_stun),
+            "sprite_shake_position": player_state.sprite_shake_position,
+            "max_sprite_shake_frame": player_state.max_sprite_shake_frame
+            / constants.FeatureDictNormalizers.MAX_SPRITE_SHAKE_FRAME,
+            "is_face_right": int(player_state.is_face_right),
+            "current_frame_advantage": player_state.current_frame_advantage
+            / constants.FeatureDictNormalizers.CURRENT_FRAME_ADVANTAGE,
+            # The below features leak some information about the opponent!
+            "would_next_forward_input_dash": int(
+                player_state.would_next_forward_input_dash
+            ),
+            "would_next_backward_input_dash": int(
+                player_state.would_next_backward_input_dash
+            ),
+            "special_attack_progress": min(player_state.special_attack_progress, 1.0),
+        }
+
+        return feature_dict
+
+    def get_last_encoding(self) -> Optional[dict[str, np.ndarray]]:
+        if self._last_common_state is None:
+            return None
+
+        return {
+            "common_state": self._last_common_state.reshape(-1),
+            "p1": np.hstack(
+                list(self._encoding_history["p1"][-1].values()),
+                dtype=np.float32,
+            ),
+            "p2": np.hstack(
+                list(self._encoding_history["p2"][-1].values()),
+                dtype=np.float32,
+            ),
+        }
+
+    def reset(self):
+        self._encoding_history = {
+            agent_id: collections.deque(maxlen=int(self.observation_delay))
+            for agent_id in ["p1", "p2"]
+        }
+
+    def _encode_action_id(self, action_id: int) -> np.ndarray:
+        """Encodes the action id into a one-hot vector.
+
+        :param action_id: The action id to encode
+        :type action_id: int
+        :return: The encoded one-hot vector
+        :rtype: np.ndarray
+        """
+
+        action_vector = np.zeros(len(self._action_id_values), dtype=np.float32)
+
+        # Get the index of the action id in constants.ActionID
+        action_index = self._action_id_values.index(action_id)
+        action_vector[action_index] = 1
+
+        assert action_vector.max() == 1 and action_vector.min() == 0
+
+        return action_vector
+
+
+def one_hot_encoder(
+    value: Union[int, float, str], collection: list[Union[int, float, str]]
+) -> np.ndarray:
+    vector = np.zeros(len(collection), dtype=np.float32)
+    vector[collection.index(value)] = 1
+    return vector

rllib/examples/envs/classes/multi_agent/footsies/fixed_rlmodules.py

@@ -0,0 +1,40 @@
+import tree  # pip install dm_tree
+
+from ray.rllib.core.rl_module import RLModule
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants
+from ray.rllib.policy import sample_batch
+from ray.rllib.utils.spaces.space_utils import batch as batch_func
+
+
+class FixedRLModule(RLModule):
+    def _forward_inference(self, batch, **kwargs):
+        return self._fixed_forward(batch, **kwargs)
+
+    def _forward_exploration(self, batch, **kwargs):
+        return self._fixed_forward(batch, **kwargs)
+
+    def _forward_train(self, *args, **kwargs):
+        raise NotImplementedError(
+            f"RLlib: {self.__class__.__name__} should not be trained. "
+            f"It is a fixed RLModule, returning a fixed action for all observations."
+        )
+
+    def _fixed_forward(self, batch, **kwargs):
+        """Implements a fixed that always returns the same action."""
+        raise NotImplementedError(
+            "FixedRLModule: This method should be overridden by subclasses to implement a specific action."
+        )
+
+
+class NoopFixedRLModule(FixedRLModule):
+    def _fixed_forward(self, batch, **kwargs):
+        obs_batch_size = len(tree.flatten(batch[sample_batch.SampleBatch.OBS])[0])
+        actions = batch_func([constants.EnvActions.NONE for _ in range(obs_batch_size)])
+        return {sample_batch.SampleBatch.ACTIONS: actions}
+
+
+class BackFixedRLModule(FixedRLModule):
+    def _fixed_forward(self, batch, **kwargs):
+        obs_batch_size = len(tree.flatten(batch[sample_batch.SampleBatch.OBS])[0])
+        actions = batch_func([constants.EnvActions.BACK for _ in range(obs_batch_size)])
+        return {sample_batch.SampleBatch.ACTIONS: actions}

rllib/examples/envs/classes/multi_agent/footsies/footsies_env.py

@@ -0,0 +1,284 @@
+import logging
+from typing import Any, Optional
+
+import numpy as np
+from gymnasium import spaces
+from pettingzoo.utils.env import (
+    ActionType,
+    AgentID,
+    ObsType,
+)
+
+from ray.rllib.env import EnvContext
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+from ray.rllib.examples.envs.classes.multi_agent.footsies.encoder import FootsiesEncoder
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game.footsies_binary import (
+    FootsiesBinary,
+)
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game.footsies_game import (
+    FootsiesGame,
+)
+
+import psutil
+
+logger = logging.getLogger("ray.rllib")
+
+
+class FootsiesEnv(MultiAgentEnv):
+    metadata = {"render.modes": ["human"]}
+    SPECIAL_CHARGE_FRAMES = 60
+    GUARD_BREAK_REWARD = 0.3
+
+    observation_space = spaces.Dict(
+        {
+            agent: spaces.Box(
+                low=-np.inf,
+                high=np.inf,
+                shape=(constants.OBSERVATION_SPACE_SIZE,),
+            )
+            for agent in ["p1", "p2"]
+        }
+    )
+
+    action_space = spaces.Dict(
+        {
+            agent: spaces.Discrete(
+                len(
+                    [
+                        constants.EnvActions.NONE,
+                        constants.EnvActions.BACK,
+                        constants.EnvActions.FORWARD,
+                        constants.EnvActions.ATTACK,
+                        constants.EnvActions.BACK_ATTACK,
+                        constants.EnvActions.FORWARD_ATTACK,
+                        # This is a special input that holds down
+                        # attack for 60 frames. It's just too long of a sequence
+                        # to easily learn by holding ATTACK for so long.
+                        constants.EnvActions.SPECIAL_CHARGE,
+                    ]
+                )
+            )
+            for agent in ["p1", "p2"]
+        }
+    )
+
+    def __init__(self, config: EnvContext, port: int):
+        super().__init__()
+
+        if config is None:
+            config = {}
+        self.config = config
+        self.port = port
+        self.footsies_process_pid = (
+            None  # Store PID of the running footsies process (we assume one per env)
+        )
+        self.agents: list[AgentID] = ["p1", "p2"]
+        self.possible_agents: list[AgentID] = self.agents.copy()
+        self._agent_ids: set[AgentID] = set(self.agents)
+
+        self.t: int = 0
+        self.max_t: int = config.get("max_t", 1000)
+        self.frame_skip = config.get("frame_skip", 4)
+        observation_delay = config.get("observation_delay", 16)
+
+        assert (
+            observation_delay % self.frame_skip == 0
+        ), "observation_delay must be divisible by frame_skip"
+
+        self.encoder = FootsiesEncoder(
+            observation_delay=observation_delay // self.frame_skip
+        )
+
+        # start the game server before initializing the communication between the
+        # game server and the Python harness via gRPC
+        self._prepare_and_start_game_server()
+        self.game = FootsiesGame(
+            host=config["host"],
+            port=self.port,
+        )
+
+        self.last_game_state = None
+        self.special_charge_queue = {
+            "p1": -1,
+            "p2": -1,
+        }
+
+    @staticmethod
+    def _convert_to_charge_action(action: int) -> int:
+        if action == constants.EnvActions.BACK:
+            return constants.EnvActions.BACK_ATTACK
+        elif action == constants.EnvActions.FORWARD:
+            return constants.EnvActions.FORWARD_ATTACK
+        else:
+            return constants.EnvActions.ATTACK
+
+    def close(self):
+        """Terminate Footsies game server process.
+
+        Run to ensure no game servers are left running.
+        """
+        timeout = 2
+        try:
+            logger.info(
+                f"RLlib {self.__class__.__name__}: Terminating Footsies "
+                f"game server process with PID: {self.footsies_process_pid}..."
+            )
+            p = psutil.Process(self.footsies_process_pid)
+            p.terminate()
+            p.wait(timeout=timeout)
+        except psutil.NoSuchProcess:
+            logger.info(
+                f"RLlib {self.__class__.__name__}: Process with PID {self.footsies_process_pid} not found, "
+                f"it might have been already terminated."
+            )
+        except psutil.TimeoutExpired:
+            logger.warning(
+                f"RLlib {self.__class__.__name__}: Process with PID {self.footsies_process_pid} did not terminate "
+                f"within {timeout} seconds. "
+                f"Sending SIGKILL signal instead.",
+            )
+            p.kill()
+            p.wait(timeout=timeout)
+
+    def get_infos(self):
+        return {agent: {} for agent in self.agents}
+
+    def get_obs(self, game_state):
+        return self.encoder.encode(game_state)
+
+    def reset(
+        self,
+        *,
+        seed: Optional[int] = None,
+        options: Optional[dict] = None,
+    ) -> tuple[dict[AgentID, ObsType], dict[AgentID, Any]]:
+        """Resets the environment to the starting state
+        and returns the initial observations for all agents.
+
+        :return: Tuple of observations and infos for each agent.
+        :rtype: tuple[dict[AgentID, ObsType], dict[AgentID, Any]]
+        """
+        self.t = 0
+        self.game.reset_game()
+        self.game.start_game()
+
+        self.encoder.reset()
+        self.last_game_state = self.game.get_state()
+
+        observations = self.get_obs(self.last_game_state)
+
+        return observations, {agent: {} for agent in self.agents}
+
+    def step(
+        self, actions: dict[AgentID, ActionType]
+    ) -> tuple[
+        dict[AgentID, ObsType],
+        dict[AgentID, float],
+        dict[AgentID, bool],
+        dict[AgentID, bool],
+        dict[AgentID, dict[str, Any]],
+    ]:
+        """Step the environment with the provided actions for all agents.
+
+        :param actions: Dictionary mapping agent ids to their actions for this step.
+        :type actions: dict[AgentID, ActionType]
+        :return: Tuple of observations, rewards, terminates, truncateds and infos for all agents.
+        :rtype: tuple[ dict[AgentID, ObsType], dict[AgentID, float], dict[AgentID, bool], dict[AgentID, bool], dict[AgentID, dict[str, Any]], ]
+        """
+        self.t += 1
+
+        for agent_id in self.agents:
+            empty_queue = self.special_charge_queue[agent_id] < 0
+            action_is_special_charge = (
+                actions[agent_id] == constants.EnvActions.SPECIAL_CHARGE
+            )
+
+            # Refill the charge queue only if we're not already in a special charge.
+            if action_is_special_charge and empty_queue:
+                self.special_charge_queue[
+                    agent_id
+                ] = self._build_charged_special_queue()
+
+            if self.special_charge_queue[agent_id] >= 0:
+                self.special_charge_queue[agent_id] -= 1
+                actions[agent_id] = self._convert_to_charge_action(actions[agent_id])
+
+        p1_action = self.game.action_to_bits(actions["p1"], is_player_1=True)
+        p2_action = self.game.action_to_bits(actions["p2"], is_player_1=False)
+
+        game_state = self.game.step_n_frames(
+            p1_action=p1_action, p2_action=p2_action, n_frames=self.frame_skip
+        )
+        observations = self.get_obs(game_state)
+
+        terminated = game_state.player1.is_dead or game_state.player2.is_dead
+
+        # Zero-sum game: 1 if other player is dead, -1 if you're dead:
+        rewards = {
+            "p1": int(game_state.player2.is_dead) - int(game_state.player1.is_dead),
+            "p2": int(game_state.player1.is_dead) - int(game_state.player2.is_dead),
+        }
+
+        if self.config.get("reward_guard_break", False):
+            p1_prev_guard_health = self.last_game_state.player1.guard_health
+            p2_prev_guard_health = self.last_game_state.player2.guard_health
+            p1_guard_health = game_state.player1.guard_health
+            p2_guard_health = game_state.player2.guard_health
+
+            if p2_guard_health < p2_prev_guard_health:
+                rewards["p1"] += self.GUARD_BREAK_REWARD
+                rewards["p2"] -= self.GUARD_BREAK_REWARD
+            if p1_guard_health < p1_prev_guard_health:
+                rewards["p2"] += self.GUARD_BREAK_REWARD
+                rewards["p1"] -= self.GUARD_BREAK_REWARD
+
+        terminateds = {
+            "p1": terminated,
+            "p2": terminated,
+            "__all__": terminated,
+        }
+
+        truncated = self.t >= self.max_t
+        truncateds = {
+            "p1": truncated,
+            "p2": truncated,
+            "__all__": truncated,
+        }
+
+        self.last_game_state = game_state
+
+        return observations, rewards, terminateds, truncateds, self.get_infos()
+
+    def _build_charged_special_queue(self):
+        assert self.SPECIAL_CHARGE_FRAMES % self.frame_skip == 0
+        steps_to_apply_attack = int(self.SPECIAL_CHARGE_FRAMES // self.frame_skip)
+        return steps_to_apply_attack
+
+    def _prepare_and_start_game_server(self):
+        fb = FootsiesBinary(config=self.config, port=self.port)
+        self.footsies_process_pid = fb.start_game_server()
+
+
+def env_creator(env_config: EnvContext) -> FootsiesEnv:
+    """Creates the Footsies environment
+
+    Ensure that each game server runs on a unique port. Training and evaluation env runners have separate port ranges.
+
+    Helper function to create the FootsiesEnv with a unique port based on the worker index and vector index.
+    It's usually passed to the `register_env()`, like this: register_env(name="FootsiesEnv", env_creator=env_creator).
+    """
+    if env_config.get("env-for-evaluation", False):
+        port = (
+            env_config["eval_start_port"]
+            - 1  # "-1" to start with eval_start_port as the first port (eval worker index starts at 1)
+            + int(env_config.worker_index) * env_config.get("num_envs_per_worker", 1)
+            + env_config.get("vector_index", 0)
+        )
+    else:
+        port = (
+            env_config["train_start_port"]
+            + int(env_config.worker_index) * env_config.get("num_envs_per_worker", 1)
+            + env_config.get("vector_index", 0)
+        )
+    return FootsiesEnv(config=env_config, port=port)

rllib/examples/envs/classes/multi_agent/footsies/game/constants.py

@@ -0,0 +1,151 @@
+from dataclasses import dataclass
+
+OBSERVATION_SPACE_SIZE: int = 81
+
+
+@dataclass
+class EnvActions:
+    NONE = 0
+    BACK = 1
+    FORWARD = 2
+    ATTACK = 3
+    BACK_ATTACK = 4
+    FORWARD_ATTACK = 5
+    SPECIAL_CHARGE = 6
+
+
+@dataclass
+class GameActions:
+    NONE = 0
+    LEFT = 1
+    RIGHT = 2
+    ATTACK = 3
+    LEFT_ATTACK = 4
+    RIGHT_ATTACK = 5
+
+
+@dataclass
+class ActionBits:
+    NONE: int = 0
+    LEFT: int = 1 << 0
+    RIGHT: int = 1 << 1
+    ATTACK: int = 1 << 2
+    LEFT_ATTACK: int = LEFT | ATTACK
+    RIGHT_ATTACK: int = RIGHT | ATTACK
+
+
+@dataclass
+class ActionID:
+    STAND = 0
+    FORWARD = 1
+    BACKWARD = 2
+    DASH_FORWARD = 10
+    DASH_BACKWARD = 11
+    N_ATTACK = 100
+    B_ATTACK = 105
+    N_SPECIAL = 110
+    B_SPECIAL = 115
+    DAMAGE = 200
+    GUARD_M = 301
+    GUARD_STAND = 305
+    GUARD_CROUCH = 306
+    GUARD_BREAK = 310
+    GUARD_PROXIMITY = 350
+    DEAD = 500
+    WIN = 510
+
+
+@dataclass
+class FeatureDictNormalizers:
+    PLAYER_POSITION_X = 4.0
+    VELOCITY_X = 5.0
+    CURRENT_ACTION_FRAME = 25
+    CURRENT_ACTION_FRAME_COUNT = 25
+    CURRENT_ACTION_REMAINING_FRAMES = 25
+    CURRENT_HIT_STUN_FRAME = 10
+    MAX_SPRITE_SHAKE_FRAME = 10
+    CURRENT_FRAME_ADVANTAGE = 10
+
+
+ACTION_TO_BITS = {
+    GameActions.NONE: ActionBits.NONE,
+    GameActions.LEFT: ActionBits.LEFT,
+    GameActions.RIGHT: ActionBits.RIGHT,
+    GameActions.ATTACK: ActionBits.ATTACK,
+    GameActions.LEFT_ATTACK: ActionBits.LEFT_ATTACK,
+    GameActions.RIGHT_ATTACK: ActionBits.RIGHT_ATTACK,
+}
+
+FOOTSIES_ACTION_IDS = {
+    "STAND": ActionID.STAND,
+    "FORWARD": ActionID.FORWARD,
+    "BACKWARD": ActionID.BACKWARD,
+    "DASH_FORWARD": ActionID.DASH_FORWARD,
+    "DASH_BACKWARD": ActionID.DASH_BACKWARD,
+    "N_ATTACK": ActionID.N_ATTACK,
+    "B_ATTACK": ActionID.B_ATTACK,
+    "N_SPECIAL": ActionID.N_SPECIAL,
+    "B_SPECIAL": ActionID.B_SPECIAL,
+    "DAMAGE": ActionID.DAMAGE,
+    "GUARD_M": ActionID.GUARD_M,
+    "GUARD_STAND": ActionID.GUARD_STAND,
+    "GUARD_CROUCH": ActionID.GUARD_CROUCH,
+    "GUARD_BREAK": ActionID.GUARD_BREAK,
+    "GUARD_PROXIMITY": ActionID.GUARD_PROXIMITY,
+    "DEAD": ActionID.DEAD,
+    "WIN": ActionID.WIN,
+}
+
+# backup file location (uploaded July 29th, 2025):
+# https://ray-example-data.s3.us-west-2.amazonaws.com/rllib/env-footsies/feature_indices.json
+# Dictionary mapping feature names to their index ranges within a flat observation vector.
+# Each key is a feature name, and its value is a dictionary with keys:
+#   "start": the starting index in the observation array.
+#   "length": it's length in bytes
+feature_indices = {
+    "common_state": {"start": 0, "length": 1},
+    "frame_count": {"start": 1, "length": 1},
+    "player_position_x": {"start": 2, "length": 1},
+    "velocity_x": {"start": 3, "length": 1},
+    "is_dead": {"start": 4, "length": 1},
+    "vital_health": {"start": 5, "length": 1},
+    "guard_health": {"start": 6, "length": 4},
+    "current_action_id": {"start": 10, "length": 17},
+    "current_action_frame": {"start": 27, "length": 1},
+    "current_action_frame_count": {"start": 28, "length": 1},
+    "current_action_remaining_frames": {"start": 29, "length": 1},
+    "is_action_end": {"start": 30, "length": 1},
+    "is_always_cancelable": {"start": 31, "length": 1},
+    "current_action_hit_count": {"start": 32, "length": 1},
+    "current_hit_stun_frame": {"start": 33, "length": 1},
+    "is_in_hit_stun": {"start": 34, "length": 1},
+    "sprite_shake_position": {"start": 35, "length": 1},
+    "max_sprite_shake_frame": {"start": 36, "length": 1},
+    "is_face_right": {"start": 37, "length": 1},
+    "current_frame_advantage": {"start": 38, "length": 1},
+    "would_next_forward_input_dash": {"start": 39, "length": 1},
+    "would_next_backward_input_dash": {"start": 40, "length": 1},
+    "special_attack_progress": {"start": 41, "length": 1},
+    "opponent_frame_count": {"start": 42, "length": 1},
+    "opponent_player_position_x": {"start": 43, "length": 1},
+    "opponent_velocity_x": {"start": 44, "length": 1},
+    "opponent_is_dead": {"start": 45, "length": 1},
+    "opponent_vital_health": {"start": 46, "length": 1},
+    "opponent_guard_health": {"start": 47, "length": 4},
+    "opponent_current_action_id": {"start": 51, "length": 17},
+    "opponent_current_action_frame": {"start": 68, "length": 1},
+    "opponent_current_action_frame_count": {"start": 69, "length": 1},
+    "opponent_current_action_remaining_frames": {"start": 70, "length": 1},
+    "opponent_is_action_end": {"start": 71, "length": 1},
+    "opponent_is_always_cancelable": {"start": 72, "length": 1},
+    "opponent_current_action_hit_count": {"start": 73, "length": 1},
+    "opponent_current_hit_stun_frame": {"start": 74, "length": 1},
+    "opponent_is_in_hit_stun": {"start": 75, "length": 1},
+    "opponent_sprite_shake_position": {"start": 76, "length": 1},
+    "opponent_max_sprite_shake_frame": {"start": 77, "length": 1},
+    "opponent_is_face_right": {"start": 78, "length": 1},
+    "opponent_current_frame_advantage": {"start": 79, "length": 1},
+    "opponent_would_next_forward_input_dash": {"start": 80, "length": 1},
+    "opponent_would_next_backward_input_dash": {"start": 81, "length": 1},
+    "opponent_special_attack_progress": {"start": 82, "length": 1},
+}

rllib/examples/envs/classes/multi_agent/footsies/game/footsies_binary.py

@@ -0,0 +1,227 @@
+import logging
+import os
+import stat
+import subprocess
+import time
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+
+import grpc
+import requests
+from filelock import FileLock
+
+from ray.rllib.env import EnvContext
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto import (
+    footsies_service_pb2 as footsies_pb2,
+    footsies_service_pb2_grpc as footsies_pb2_grpc,
+)
+from ray.util import log_once
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BinaryUrls:
+    # Uploaded 07.28.2025
+    S3_ROOT = "https://ray-example-data.s3.us-west-2.amazonaws.com/rllib/env-footsies/binaries/"
+
+    # Zip file names
+    ZIP_LINUX_SERVER = "footsies_linux_server_021725.zip"
+    ZIP_LINUX_WINDOWED = "footsies_linux_windowed_021725.zip"
+    ZIP_MAC_HEADLESS = "footsies_mac_headless_5709b6d.zip"
+    ZIP_MAC_WINDOWED = "footsies_mac_windowed_5709b6d.zip"
+
+    # Full URLs
+    URL_LINUX_SERVER_BINARIES = S3_ROOT + ZIP_LINUX_SERVER
+    URL_LINUX_WINDOWED_BINARIES = S3_ROOT + ZIP_LINUX_WINDOWED
+    URL_MAC_HEADLESS_BINARIES = S3_ROOT + ZIP_MAC_HEADLESS
+    URL_MAC_WINDOWED_BINARIES = S3_ROOT + ZIP_MAC_WINDOWED
+
+
+class FootsiesBinary:
+    def __init__(self, config: EnvContext, port: int):
+        self._urls = BinaryUrls()
+        self.config = config
+        self.port = port
+        self.binary_to_download = config["binary_to_download"]
+
+        if self.binary_to_download == "linux_server":
+            self.url = self._urls.URL_LINUX_SERVER_BINARIES
+        elif self.binary_to_download == "linux_windowed":
+            self.url = self._urls.URL_LINUX_WINDOWED_BINARIES
+        elif self.binary_to_download == "mac_headless":
+            self.url = self._urls.URL_MAC_HEADLESS_BINARIES
+        elif self.binary_to_download == "mac_windowed":
+            self.url = self._urls.URL_MAC_WINDOWED_BINARIES
+        else:
+            raise ValueError(f"Invalid target binary: {self.binary_to_download}")
+
+        self.full_download_dir = Path(config["binary_download_dir"]).resolve()
+        self.full_download_path = (
+            self.full_download_dir / str.split(self.url, sep="/")[-1]
+        )
+        self.full_extract_dir = Path(config["binary_extract_dir"]).resolve()
+        self.renamed_path = self.full_extract_dir / "footsies_binaries"
+
+    @staticmethod
+    def _add_executable_permission(binary_path: Path) -> None:
+        binary_path.chmod(binary_path.stat().st_mode | stat.S_IXUSR)
+
+    def start_game_server(self) -> int:
+        """Downloads, unzips, and starts the Footsies game server binary.
+
+        Returns footsies process PID.
+        """
+        self._download_game_binary()
+        self._unzip_game_binary()
+
+        if self.binary_to_download == "mac_windowed":
+            game_binary_path = (
+                Path(self.renamed_path) / "Contents" / "MacOS" / "FOOTSIES"
+            )
+        elif self.binary_to_download == "mac_headless":
+            game_binary_path = Path(self.renamed_path) / "FOOTSIES"
+        else:
+            game_binary_path = Path(self.renamed_path) / "footsies.x86_64"
+
+        if os.access(game_binary_path, os.X_OK):
+            logger.info(
+                f"Game binary has an 'executable' permission: {game_binary_path}"
+            )
+        else:
+            self._add_executable_permission(game_binary_path)
+        logger.info(f"Game binary path: {game_binary_path}")
+
+        # The underlying game can be quite spammy. So when we are not debugging it, we can suppress the output.
+        log_output = self.config.get("log_unity_output")
+        if log_output is None:
+            if log_once("log_unity_output_not_set"):
+                logger.warning(
+                    "`log_unity_output` not set in environment config, not logging output by default"
+                )
+            log_output = False
+
+        if not log_output:
+            stdout_dest = stderr_dest = subprocess.DEVNULL
+        else:
+            stdout_dest = stderr_dest = None  # Use parent's stdout/stderr
+
+        if (
+            self.binary_to_download == "linux_server"
+            or self.binary_to_download == "linux_windowed"
+        ):
+            process = subprocess.Popen(
+                [game_binary_path, "--port", str(self.port)],
+                stdout=stdout_dest,
+                stderr=stderr_dest,
+            )
+        else:
+            process = subprocess.Popen(
+                [
+                    "arch",
+                    "-x86_64",
+                    game_binary_path,
+                    "--port",
+                    str(self.port),
+                ],
+                stdout=stdout_dest,
+                stderr=stderr_dest,
+            )
+
+        # check if the game server is running correctly
+        timeout = 2
+        channel = grpc.insecure_channel(f"localhost:{self.port}")
+        stub = footsies_pb2_grpc.FootsiesGameServiceStub(channel)
+
+        # step 1: try to start the game
+        while True:
+            try:
+                stub.StartGame(footsies_pb2.Empty())
+                logger.info("Game ready!")
+                break
+            except grpc.RpcError as e:
+                code = e.code()
+                if code in (
+                    grpc.StatusCode.UNAVAILABLE,
+                    grpc.StatusCode.DEADLINE_EXCEEDED,
+                ):
+                    logger.info(f"RLlib {self.__class__.__name__}: Game not ready...")
+                    time.sleep(timeout)
+                    continue
+                raise
+
+        # step 2: check if the game is ready
+        ready = False
+        while not ready:
+            try:
+                ready = stub.IsReady(footsies_pb2.Empty()).value
+                if not ready:
+                    logger.info(f"RLlib {self.__class__.__name__}: Game not ready...")
+                    time.sleep(timeout)
+                    continue
+                else:
+                    logger.info("Game ready!")
+                    break
+            except grpc.RpcError as e:
+                if e.code() in (
+                    grpc.StatusCode.UNAVAILABLE,
+                    grpc.StatusCode.DEADLINE_EXCEEDED,
+                ):
+                    time.sleep(timeout)
+                    logger.info(f"RLlib {self.__class__.__name__}: Game not ready...")
+                    continue
+                raise
+
+        channel.close()
+        return process.pid
+
+    def _download_game_binary(self):
+        # As multiple actors might try to download all at the same time.
+        # The file lock should force only one actor to download
+        chunk_size = 1024 * 1024  # 1MB
+
+        lock_path = self.full_download_path.parent / ".footsies-download.lock"
+        with FileLock(lock_path, timeout=300):
+            if self.full_download_path.exists():
+                logger.info(
+                    f"Game binary already exists at {self.full_download_path}, skipping download."
+                )
+
+            else:
+                try:
+                    with requests.get(self.url, stream=True) as response:
+                        response.raise_for_status()
+                        self.full_download_dir.mkdir(parents=True, exist_ok=True)
+                        with open(self.full_download_path, "wb") as f:
+                            for chunk in response.iter_content(chunk_size=chunk_size):
+                                if chunk:
+                                    f.write(chunk)
+                    logger.info(
+                        f"Downloaded game binary to {self.full_download_path}\n"
+                        f"Binary size: {self.full_download_path.stat().st_size / 1024 / 1024:.1f} MB\n"
+                    )
+                except requests.exceptions.RequestException as e:
+                    logger.error(f"Failed to download binary from {self.url}: {e}")
+
+    def _unzip_game_binary(self):
+        # As multiple actors might try to unzip or rename the paths at the same time.
+        # The file lock should force this function to be sequential
+        lock_path = self.full_download_path.parent / ".footsies-unzip.lock"
+        with FileLock(lock_path, timeout=300):
+            if self.renamed_path.exists():
+                logger.info(
+                    f"Game binary already extracted at {self.renamed_path}, skipping extraction."
+                )
+            else:
+                self.full_extract_dir.mkdir(parents=True, exist_ok=True)
+                with zipfile.ZipFile(self.full_download_path, mode="r") as zip_ref:
+                    zip_ref.extractall(self.full_extract_dir)
+
+                if self.binary_to_download == "mac_windowed":
+                    self.full_download_path.with_suffix(".app").rename(
+                        self.renamed_path
+                    )
+                else:
+                    self.full_download_path.with_suffix("").rename(self.renamed_path)
+                logger.info(f"Extracted game binary to {self.renamed_path}")

rllib/examples/envs/classes/multi_agent/footsies/game/footsies_game.py

@@ -0,0 +1,121 @@
+import logging
+import time
+
+import grpc
+import numpy as np
+
+import ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto.footsies_service_pb2 as footsies_pb2
+import ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto.footsies_service_pb2_grpc as footsies_pb2_grpc
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants
+
+logger = logging.getLogger(__name__)
+
+
+class FootsiesGame:
+    """Handles gRPC communication with game the server.
+
+    This class establishes communication between the
+    game server and the Python harness via gRPC. It provides methods
+    to start the game, reset it, get the current state, and step the
+    game by a certain number of frames.
+    """
+
+    def __init__(self, host: str, port: int):
+        self.host = host
+        self.port = port
+        self.stub = self._initialize_stub()
+
+    @staticmethod
+    def action_to_bits(action: int, is_player_1: bool) -> int:
+        """Converts an action to its corresponding bit representation."""
+
+        if isinstance(action, np.ndarray):
+            action = action.item()
+
+        if is_player_1:
+            if action == constants.EnvActions.BACK:
+                action = constants.GameActions.LEFT
+            elif action == constants.EnvActions.FORWARD:
+                action = constants.GameActions.RIGHT
+            elif action == constants.EnvActions.BACK_ATTACK:
+                action = constants.GameActions.LEFT_ATTACK
+            elif action == constants.EnvActions.FORWARD_ATTACK:
+                action = constants.GameActions.RIGHT_ATTACK
+        else:
+            if action == constants.EnvActions.BACK:
+                action = constants.GameActions.RIGHT
+            elif action == constants.EnvActions.FORWARD:
+                action = constants.GameActions.LEFT
+            elif action == constants.EnvActions.BACK_ATTACK:
+                action = constants.GameActions.RIGHT_ATTACK
+            elif action == constants.EnvActions.FORWARD_ATTACK:
+                action = constants.GameActions.LEFT_ATTACK
+
+        return constants.ACTION_TO_BITS[action]
+
+    def get_encoded_state(self) -> footsies_pb2.EncodedGameState:
+        """Gets the current encoded game state by calling the GetEncodedState RPC."""
+        try:
+            return self.stub.GetEncodedState(footsies_pb2.Empty())
+        except Exception as e:
+            logger.error(f"Error calling GetEncodedState with exception: {e}")
+            raise e
+
+    def get_state(self) -> footsies_pb2.GameState:
+        """Gets the current game state by calling the GetState RPC."""
+        try:
+            return self.stub.GetState(footsies_pb2.Empty())
+        except Exception as e:
+            logger.error(f"Error calling GetState with exception: {e}")
+            raise e
+
+    def is_ready(self) -> bool:
+        """Checks if the game is ready by calling the IsReady RPC."""
+        try:
+            return self.stub.IsReady(footsies_pb2.Empty()).value
+        except Exception as e:
+            logger.error(f"Error calling IsReady with exception: {e}")
+            raise e
+
+    def reset_game(self) -> None:
+        """Resets the game by calling the ResetGame RPC."""
+        try:
+            self.stub.ResetGame(footsies_pb2.Empty())
+        except Exception as e:
+            logger.error(f"Error calling ResetGame with exception: {e}")
+            raise e
+
+    def start_game(self) -> None:
+        """Starts the game by calling the StartGame RPC."""
+        try:
+            self.stub.StartGame(footsies_pb2.Empty())
+
+            while not self.is_ready():
+                logger.info("Game not ready...")
+                time.sleep(0.5)
+            logger.info("StartGame called successfully")
+
+        except Exception as e:
+            logger.error(f"Error calling StartGame with exception: {e}")
+            raise e
+
+    def step_n_frames(
+        self, p1_action: int, p2_action: int, n_frames: int
+    ) -> footsies_pb2.GameState:
+        """Steps the game by n_frames with the given player actions. The provided actions will be repeated for all n_frames."""
+        try:
+            step_input = footsies_pb2.StepInput(
+                p1_action=p1_action, p2_action=p2_action, nFrames=n_frames
+            )
+            return self.stub.StepNFrames(step_input)
+        except Exception as e:
+            logger.error(f"Error calling StepNFrames with exception: {e}")
+            raise e
+
+    def _initialize_stub(self) -> footsies_pb2_grpc.FootsiesGameServiceStub:
+        try:
+            channel = grpc.insecure_channel(f"{self.host}:{self.port}")
+            return footsies_pb2_grpc.FootsiesGameServiceStub(channel)
+        except grpc.RpcError as e:
+            logger.error(f"Error connecting to gRPC stub with exception: {e}")
+            raise e

rllib/examples/envs/classes/multi_agent/footsies/game/proto/footsies_service.proto

@@ -0,0 +1,63 @@
+syntax = "proto3";
+
+
+service FootsiesGameService {
+  rpc StartGame(Empty) returns (Empty) {}
+  rpc ResetGame(Empty) returns (Empty) {}
+  rpc StepNFrames(StepInput) returns (GameState) {}
+  rpc GetState(Empty) returns (GameState) {}
+  rpc GetEncodedState(Empty) returns (EncodedGameState) {}
+  rpc IsReady(Empty) returns (BoolValue) {}
+}
+
+
+message StepInput {
+  int64 p1_action = 1;
+  int64 p2_action = 2;
+  int64 nFrames = 3;
+}
+
+message PlayerState {
+  float player_position_x = 1;
+  bool is_dead = 2;
+  int64 vital_health = 3;
+  int64 guard_health = 4;
+  int64 current_action_id = 5;
+  int64 current_action_frame = 6;
+  int64 current_action_frame_count = 7;
+  bool is_action_end = 8;
+  bool is_always_cancelable = 9;
+  int64 current_action_hit_count = 10;
+  int64 current_hit_stun_frame = 11;
+  bool is_in_hit_stun = 12;
+  int64 sprite_shake_position = 13;
+  int64 max_sprite_shake_frame = 14;
+  float velocity_x = 15;
+  bool is_face_right = 16;
+  repeated int64 input_buffer = 17;
+  int64 current_frame_advantage = 18;
+  bool would_next_forward_input_dash = 19;
+  bool would_next_backward_input_dash = 20;
+  float special_attack_progress = 21;
+}
+
+message GameState {
+  PlayerState player1 = 1;
+  PlayerState player2 = 2;
+  int64 round_state = 3;
+  int64 frame_count = 4;
+}
+
+message EncodedGameState {
+  repeated float player1_encoding = 1;
+  repeated float player2_encoding = 2;
+}
+
+message BoolValue {
+  bool value = 1;
+}
+
+
+
+
+message Empty {}

rllib/examples/envs/classes/multi_agent/footsies/game/proto/footsies_service_pb2_grpc.py

@@ -0,0 +1,307 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+import ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto.footsies_service_pb2 as footsies__service__pb2
+
+
+# import footsies_service_pb2 as footsies__service__pb2
+
+
+class FootsiesGameServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.StartGame = channel.unary_unary(
+            "/FootsiesGameService/StartGame",
+            request_serializer=footsies__service__pb2.Empty.SerializeToString,
+            response_deserializer=footsies__service__pb2.Empty.FromString,
+        )
+        self.ResetGame = channel.unary_unary(
+            "/FootsiesGameService/ResetGame",
+            request_serializer=footsies__service__pb2.Empty.SerializeToString,
+            response_deserializer=footsies__service__pb2.Empty.FromString,
+        )
+        self.StepNFrames = channel.unary_unary(
+            "/FootsiesGameService/StepNFrames",
+            request_serializer=footsies__service__pb2.StepInput.SerializeToString,
+            response_deserializer=footsies__service__pb2.GameState.FromString,
+        )
+        self.GetState = channel.unary_unary(
+            "/FootsiesGameService/GetState",
+            request_serializer=footsies__service__pb2.Empty.SerializeToString,
+            response_deserializer=footsies__service__pb2.GameState.FromString,
+        )
+        self.GetEncodedState = channel.unary_unary(
+            "/FootsiesGameService/GetEncodedState",
+            request_serializer=footsies__service__pb2.Empty.SerializeToString,
+            response_deserializer=footsies__service__pb2.EncodedGameState.FromString,
+        )
+        self.IsReady = channel.unary_unary(
+            "/FootsiesGameService/IsReady",
+            request_serializer=footsies__service__pb2.Empty.SerializeToString,
+            response_deserializer=footsies__service__pb2.BoolValue.FromString,
+        )
+
+
+class FootsiesGameServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def StartGame(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def ResetGame(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def StepNFrames(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def GetState(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def GetEncodedState(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def IsReady(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+
+def add_FootsiesGameServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+        "StartGame": grpc.unary_unary_rpc_method_handler(
+            servicer.StartGame,
+            request_deserializer=footsies__service__pb2.Empty.FromString,
+            response_serializer=footsies__service__pb2.Empty.SerializeToString,
+        ),
+        "ResetGame": grpc.unary_unary_rpc_method_handler(
+            servicer.ResetGame,
+            request_deserializer=footsies__service__pb2.Empty.FromString,
+            response_serializer=footsies__service__pb2.Empty.SerializeToString,
+        ),
+        "StepNFrames": grpc.unary_unary_rpc_method_handler(
+            servicer.StepNFrames,
+            request_deserializer=footsies__service__pb2.StepInput.FromString,
+            response_serializer=footsies__service__pb2.GameState.SerializeToString,
+        ),
+        "GetState": grpc.unary_unary_rpc_method_handler(
+            servicer.GetState,
+            request_deserializer=footsies__service__pb2.Empty.FromString,
+            response_serializer=footsies__service__pb2.GameState.SerializeToString,
+        ),
+        "GetEncodedState": grpc.unary_unary_rpc_method_handler(
+            servicer.GetEncodedState,
+            request_deserializer=footsies__service__pb2.Empty.FromString,
+            response_serializer=footsies__service__pb2.EncodedGameState.SerializeToString,
+        ),
+        "IsReady": grpc.unary_unary_rpc_method_handler(
+            servicer.IsReady,
+            request_deserializer=footsies__service__pb2.Empty.FromString,
+            response_serializer=footsies__service__pb2.BoolValue.SerializeToString,
+        ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+        "FootsiesGameService", rpc_method_handlers
+    )
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+# This class is part of an EXPERIMENTAL API.
+class FootsiesGameService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def StartGame(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/StartGame",
+            footsies__service__pb2.Empty.SerializeToString,
+            footsies__service__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def ResetGame(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/ResetGame",
+            footsies__service__pb2.Empty.SerializeToString,
+            footsies__service__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def StepNFrames(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/StepNFrames",
+            footsies__service__pb2.StepInput.SerializeToString,
+            footsies__service__pb2.GameState.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def GetState(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/GetState",
+            footsies__service__pb2.Empty.SerializeToString,
+            footsies__service__pb2.GameState.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def GetEncodedState(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/GetEncodedState",
+            footsies__service__pb2.Empty.SerializeToString,
+            footsies__service__pb2.EncodedGameState.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def IsReady(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/FootsiesGameService/IsReady",
+            footsies__service__pb2.Empty.SerializeToString,
+            footsies__service__pb2.BoolValue.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )

rllib/examples/envs/classes/multi_agent/footsies/test/footsies_suppress_unity_logs.py

@@ -0,0 +1,135 @@
+import os
+import platform
+import tempfile
+import time
+import unittest
+from contextlib import contextmanager
+from pathlib import Path
+
+from ray.rllib.env import EnvContext
+from ray.rllib.examples.envs.classes.multi_agent.footsies.footsies_env import (
+    env_creator,
+)
+
+# Detect platform and choose appropriate binary
+if platform.system() == "Darwin":
+    binary_to_download = "mac_headless"
+elif platform.system() == "Linux":
+    binary_to_download = "linux_server"
+else:
+    raise RuntimeError(f"Unsupported platform: {platform.system()}")
+
+FOOTSIES_ENV_BASE_CONFIG = {
+    "max_t": 1000,
+    "frame_skip": 4,
+    "observation_delay": 16,
+    "host": "localhost",
+    "binary_download_dir": "/tmp/ray/binaries/footsies",
+    "binary_extract_dir": "/tmp/ray/binaries/footsies",
+    "binary_to_download": binary_to_download,
+}
+
+_port_counter = 45001
+
+
+def _create_env(config_overrides):
+    global _port_counter
+
+    config = {
+        **FOOTSIES_ENV_BASE_CONFIG,
+        "train_start_port": _port_counter,
+        "eval_start_port": _port_counter + 1,
+        **config_overrides,
+    }
+    _port_counter += 2
+
+    return env_creator(EnvContext(config, worker_index=0))
+
+
+@contextmanager
+def capture_stdout_stderr():
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tmp:
+        log_path = tmp.name
+
+    # Use file descriptors 1 and 2 directly (stdout and stderr)
+    # This bypasses pytest's wrapping of sys.stdout/stderr
+    stdout_fd = 1
+    stderr_fd = 2
+
+    # Save original file descriptors
+    saved_stdout = os.dup(stdout_fd)
+    saved_stderr = os.dup(stderr_fd)
+
+    log_file = None
+    try:
+        # Open log file with line buffering
+        log_file = open(log_path, "w", buffering=1)
+        os.dup2(log_file.fileno(), stdout_fd)
+        os.dup2(log_file.fileno(), stderr_fd)
+
+        yield log_path
+
+    finally:
+        # Flush the log file
+        if log_file:
+            log_file.flush()
+
+        # Restore original file descriptors
+        os.dup2(saved_stdout, stdout_fd)
+        os.dup2(saved_stderr, stderr_fd)
+
+        # Close file and saved descriptors
+        if log_file:
+            log_file.close()
+        os.close(saved_stdout)
+        os.close(saved_stderr)
+
+
+class TestFootsies(unittest.TestCase):
+    def test_default_supress_output_mode(self):
+        with capture_stdout_stderr() as log_path:
+            env = _create_env({})
+            time.sleep(2)  # Give Unity time to write output
+            env.close()
+            # Give a bit more time for any buffered output to be written
+            time.sleep(0.5)
+
+        # Read the captured output
+        with open(log_path, "r") as f:
+            captured_output = f.read()
+
+        assert (
+            "`log_unity_output` not set in environment config, not logging output by default"
+            in captured_output
+        )
+        assert "[UnityMemory]" not in captured_output
+
+        # Clean up
+        if Path(log_path).exists():
+            os.unlink(log_path)
+
+    def test_enable_output_mode(self):
+        with capture_stdout_stderr() as log_path:
+            env = _create_env({"log_unity_output": True})
+            time.sleep(2)  # Give Unity time to write output
+            env.close()
+            # Give a bit more time for any buffered output to be written
+            time.sleep(0.5)
+
+        # Read the captured output
+        with open(log_path, "r") as f:
+            captured_output = f.read()
+
+        assert "[UnityMemory]" in captured_output
+
+        # Clean up
+        if Path(log_path).exists():
+            os.unlink(log_path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/examples/envs/classes/multi_agent/footsies/utils.py

@@ -0,0 +1,383 @@
+import collections
+import logging
+import platform
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.callbacks import RLlibCallback
+from ray.rllib.core.rl_module import RLModuleSpec
+from ray.rllib.env.env_runner import EnvRunner
+from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
+from ray.rllib.examples.envs.classes.multi_agent.footsies.game.constants import (
+    FOOTSIES_ACTION_IDS,
+)
+from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS
+from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
+from ray.rllib.utils.typing import EpisodeType
+
+logger = logging.getLogger("ray.rllib")
+
+
+@dataclass
+class Matchup:
+    p1: str
+    p2: str
+    prob: float
+
+
+class Matchmaker:
+    def __init__(self, matchups: list[Matchup]):
+        self.matchups = matchups
+        self.probs = [matchup.prob for matchup in matchups]
+        self.current_matchups = collections.defaultdict(dict)
+
+    def agent_to_module_mapping_fn(
+        self, agent_id: str, episode: EpisodeType, **kwargs
+    ) -> str:
+        """Mapping function that retrieves policy_id from the sampled matchup"""
+        id_ = episode.id_
+        if self.current_matchups.get(id_) is None:
+            # step 1: sample a matchup according to the specified probabilities
+            sampled_matchup = np.random.choice(a=self.matchups, p=self.probs)
+
+            # step 2: Randomize who is player 1 and player 2
+            policies = [sampled_matchup.p1, sampled_matchup.p2]
+            p1, p2 = np.random.choice(policies, size=2, replace=False)
+
+            # step 3: Set as the current matchup for the episode in question (id_)
+            self.current_matchups[id_]["p1"] = p1
+            self.current_matchups[id_]["p2"] = p2
+
+        policy_id = self.current_matchups[id_].pop(agent_id)
+
+        # remove (an empty dict) for the current episode with id_
+        if not self.current_matchups[id_]:
+            del self.current_matchups[id_]
+
+        return policy_id
+
+
+class MetricsLoggerCallback(RLlibCallback):
+    def __init__(self, main_policy: str) -> None:
+        """Log experiment metrics
+
+        Logs metrics after each episode step and at the end of each (train or eval) episode.
+        Metrics logged at the end of each episode will be later used by MixManagerCallback
+        to decide whether to add a new opponent to the mix.
+        """
+        super().__init__()
+        self.main_policy = main_policy
+        self.action_id_to_str = {
+            action_id: action_str
+            for action_str, action_id in FOOTSIES_ACTION_IDS.items()
+        }
+
+    def on_episode_step(
+        self,
+        *,
+        episode: MultiAgentEpisode,
+        env_runner: Optional[EnvRunner] = None,
+        metrics_logger: Optional[MetricsLogger] = None,
+        env: Optional[gym.Env] = None,
+        env_index: int,
+        **kwargs,
+    ) -> None:
+        """Log action usage frequency
+
+        Log actions performed by both players at each step of the (training or evaluation) episode.
+        """
+        stage = "eval" if env_runner.config.in_evaluation else "train"
+
+        # get the ModuleID for each agent
+        p1_module = episode.module_for("p1")
+        p2_module = episode.module_for("p2")
+
+        # get action string for each agent
+        p1_action_id = env.envs[
+            env_index
+        ].unwrapped.last_game_state.player1.current_action_id
+        p2_action_id = env.envs[
+            env_index
+        ].unwrapped.last_game_state.player2.current_action_id
+        p1_action_str = self.action_id_to_str[p1_action_id]
+        p2_action_str = self.action_id_to_str[p2_action_id]
+
+        metrics_logger.log_value(
+            key=f"footsies/{stage}/actions/{p1_module}/{p1_action_str}",
+            value=1,
+            reduce="sum",
+            window=100,
+        )
+        metrics_logger.log_value(
+            key=f"footsies/{stage}/actions/{p2_module}/{p2_action_str}",
+            value=1,
+            reduce="sum",
+            window=100,
+        )
+
+    def on_episode_end(
+        self,
+        *,
+        episode: MultiAgentEpisode,
+        env_runner: Optional[EnvRunner] = None,
+        metrics_logger: Optional[MetricsLogger] = None,
+        env: Optional[gym.Env] = None,
+        env_index: int,
+        **kwargs,
+    ) -> None:
+        """Log win rates
+
+        Log win rates of the main policy against its opponent at the end of the (training or evaluation) episode.
+        """
+        stage = "eval" if env_runner.config.in_evaluation else "train"
+
+        # check status of "p1" and "p2"
+        last_game_state = env.envs[env_index].unwrapped.last_game_state
+        p1_dead = last_game_state.player1.is_dead
+        p2_dead = last_game_state.player2.is_dead
+
+        # get the ModuleID for each agent
+        p1_module = episode.module_for("p1")
+        p2_module = episode.module_for("p2")
+
+        if self.main_policy == p1_module:
+            opponent_id = p2_module
+            main_policy_win = p2_dead
+        elif self.main_policy == p2_module:
+            opponent_id = p1_module
+            main_policy_win = p1_dead
+        else:
+            logger.info(
+                f"RLlib {self.__class__.__name__}: Main policy: '{self.main_policy}' not found in this episode. "
+                f"Policies in this episode are: '{p1_module}' and '{p2_module}'. "
+                f"Check your multi_agent 'policy_mapping_fn'. "
+                f"Metrics logging for this episode will be skipped."
+            )
+            return
+
+        if p1_dead and p2_dead:
+            metrics_logger.log_value(
+                key=f"footsies/{stage}/both_dead/{self.main_policy}/vs_{opponent_id}",
+                value=1,
+                reduce="mean",
+                window=100,
+            )
+        elif not p1_dead and not p2_dead:
+            metrics_logger.log_value(
+                key=f"footsies/{stage}/both_alive/{self.main_policy}/vs_{opponent_id}",
+                value=1,
+                reduce="mean",
+                window=100,
+            )
+        else:
+            # log the win rate against the opponent with an 'opponent_id'
+            metrics_logger.log_value(
+                key=f"footsies/{stage}/win_rates/{self.main_policy}/vs_{opponent_id}",
+                value=int(main_policy_win),
+                reduce="mean",
+                window=100,
+            )
+
+            # log the win rate, without specifying the opponent
+            # this metric collected from the eval env runner
+            # will be used to decide whether to add
+            # a new opponent at the current level.
+            metrics_logger.log_value(
+                key=f"footsies/{stage}/win_rates/{self.main_policy}/vs_any",
+                value=int(main_policy_win),
+                reduce="mean",
+                window=100,
+            )
+
+
+class MixManagerCallback(RLlibCallback):
+    def __init__(
+        self,
+        win_rate_threshold: float,
+        main_policy: str,
+        target_mix_size: int,
+        starting_modules=list[str],  # default is ["lstm", "noop"]
+        fixed_modules_progression_sequence=tuple[str],  # default is ("noop", "back")
+    ) -> None:
+        """Track win rates and manage mix of opponents"""
+        super().__init__()
+        self.win_rate_threshold = win_rate_threshold
+        self.main_policy = main_policy
+        self.target_mix_size = target_mix_size
+        self.fixed_modules_progression_sequence = tuple(
+            fixed_modules_progression_sequence
+        )  # Order of RL modules to be added to the mix
+        self.modules_in_mix = list(
+            starting_modules
+        )  # RLModules that are currently in the mix
+        self._trained_policy_idx = (
+            0  # We will use this to create new opponents of the main policy
+        )
+
+    def on_evaluate_end(
+        self,
+        *,
+        algorithm: Algorithm,
+        metrics_logger: Optional[MetricsLogger] = None,
+        evaluation_metrics: dict,
+        **kwargs,
+    ) -> None:
+        """Check win rates and add new opponent if necessary.
+
+        Check the win rate of the main policy against its current opponent.
+        If the win rate exceeds the specified threshold, add a new opponent to the mix, by modifying:
+        1. update the policy_mapping_fn for (training and evaluation) env runners
+        2. if the new policy is a trained one (not a fixed RL module), modify Algorithm's state (initialize the state of the newly added RLModule by using the main policy)
+        """
+        _main_module = algorithm.get_module(self.main_policy)
+        new_module_id = None
+        new_module_spec = None
+
+        win_rate = evaluation_metrics[ENV_RUNNER_RESULTS][
+            f"footsies/eval/win_rates/{self.main_policy}/vs_any"
+        ]
+
+        if win_rate > self.win_rate_threshold:
+            logger.info(
+                f"RLlib {self.__class__.__name__}: Win rate for main policy '{self.main_policy}' "
+                f"exceeded threshold ({win_rate} > {self.win_rate_threshold})."
+                f" Adding new RL Module to the mix..."
+            )
+
+            # check if fixed RL module should be added to the mix,
+            # and if so, create new_module_id and new_module_spec for it
+            for module_id in self.fixed_modules_progression_sequence:
+                if module_id not in self.modules_in_mix:
+                    new_module_id = module_id
+                    break
+
+            # in case that all fixed RL Modules are already in the mix (together with the main policy),
+            # we will add a new RL Module by taking main policy and adding an instance of it to the mix
+            if new_module_id is None:
+                new_module_id = f"{self.main_policy}_v{self._trained_policy_idx}"
+                new_module_spec = RLModuleSpec.from_module(_main_module)
+                self._trained_policy_idx += 1
+
+            # create new policy mapping function, to ensure that the main policy plays against newly added policy
+            new_mapping_fn = Matchmaker(
+                [
+                    Matchup(
+                        p1=self.main_policy,
+                        p2=new_module_id,
+                        prob=1.0,
+                    )
+                ]
+            ).agent_to_module_mapping_fn
+
+            # STEP 1: Add the new module first (if it's a trained module)
+            if new_module_id not in self.fixed_modules_progression_sequence:
+                # Add module to Learners and EnvRunners (but don't update mapping yet)
+                algorithm.add_module(
+                    module_id=new_module_id,
+                    module_spec=new_module_spec,
+                    new_agent_to_module_mapping_fn=None,  # Don't update mapping yet!
+                )
+
+                # Initialize the new module with main policy's weights
+                algorithm.set_state(
+                    {
+                        "learner_group": {
+                            "learner": {
+                                "rl_module": {
+                                    new_module_id: _main_module.get_state(),
+                                }
+                            }
+                        },
+                    }
+                )
+
+                # STEP 2: CRITICAL - Update aggregator actors with the new module
+                # Aggregators run the learner connector pipeline which needs all modules.
+                if (
+                    hasattr(algorithm, "_aggregator_actor_manager")
+                    and algorithm._aggregator_actor_manager
+                ):
+                    logger.info(
+                        f"RLlib {self.__class__.__name__}: Updating aggregator actors "
+                        f"with new module '{new_module_id}'..."
+                    )
+
+                    # Add the new module to each aggregator actor's MultiRLModule
+                    algorithm._aggregator_actor_manager.foreach_actor(
+                        func=lambda actor, mid=new_module_id, spec=new_module_spec: (
+                            actor._module.add_module(
+                                module_id=mid,
+                                module=spec.build(),
+                            )
+                        )
+                    )
+
+                    # Sync weights from learner to aggregator actors
+                    weights = algorithm.learner_group.get_weights(
+                        module_ids=[new_module_id]
+                    )
+                    algorithm._aggregator_actor_manager.foreach_actor(
+                        func=lambda actor, w=weights: actor._module.set_state(w)
+                    )
+
+                    logger.info(
+                        f"RLlib {self.__class__.__name__}: Aggregator actors updated successfully."
+                    )
+
+            # STEP 3: NOW update the policy mapping function on all EnvRunners
+            # At this point, the module exists everywhere (Learners, EnvRunners, Aggregators)
+            algorithm.env_runner_group.foreach_env_runner(
+                lambda er: er.config.multi_agent(policy_mapping_fn=new_mapping_fn),
+                local_env_runner=True,
+            )
+
+            algorithm.eval_env_runner_group.foreach_env_runner(
+                lambda er: er.config.multi_agent(policy_mapping_fn=new_mapping_fn),
+                local_env_runner=True,
+            )
+            # Update algorithm's config to maintain consistency
+            algorithm.config._is_frozen = False
+            algorithm.config.multi_agent(policy_mapping_fn=new_mapping_fn)
+            algorithm.config.freeze()
+
+            # Update the current mix list
+            self.modules_in_mix.append(new_module_id)
+
+        else:
+            logger.info(
+                f"RLlib {self.__class__.__name__}: Win rate for main policy '{self.main_policy}' "
+                f"did not exceed threshold ({win_rate} <= {self.win_rate_threshold})."
+            )
+
+    def on_train_result(
+        self,
+        *,
+        algorithm: Algorithm,
+        metrics_logger: Optional[MetricsLogger] = None,
+        result: Dict,
+        **kwargs,
+    ) -> None:
+        """Report the current mix size at the end of training iteration.
+
+        That will tell Ray Tune, whether to stop training (once the 'target_mix_size' has been reached).
+        """
+        result["mix_size"] = len(self.modules_in_mix)
+
+
+def platform_for_binary_to_download(render: bool) -> str:
+    if platform.system() == "Darwin":
+        if render:
+            return "mac_windowed"
+        else:
+            return "mac_headless"
+    elif platform.system() == "Linux":
+        if render:
+            return "linux_windowed"
+        else:
+            return "linux_server"
+    else:
+        raise RuntimeError(f"Unsupported platform: {platform.system()}")

rllib/examples/envs/classes/multi_agent/guess_the_number_game.py

@@ -0,0 +1,89 @@
+import gymnasium as gym
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class GuessTheNumberGame(MultiAgentEnv):
+    """
+    We have two players, 0 and 1. Agent 0 has to pick a number between 0, MAX-1
+    at reset. Agent 1 has to guess the number by asking N questions of whether
+    of the form of "a <number> is higher|lower|equal to the picked number. The
+    action space is MultiDiscrete [3, MAX]. For the first index 0 means lower,
+    1 means higher and 2 means equal. The environment answers with yes (1) or
+    no (0) on the reward function. Every time step that agent 1 wastes agent 0
+    gets a reward of 1. After N steps the game is terminated. If agent 1
+    guesses the number correctly, it gets a reward of 100 points, otherwise it
+    gets a reward of 0. On the other hand if agent 0 wins they win 100 points.
+    The optimal policy controlling agent 1 should converge to a binary search
+    strategy.
+    """
+
+    MAX_NUMBER = 3
+    MAX_STEPS = 20
+
+    def __init__(self, config=None):
+        super().__init__()
+        self._agent_ids = {0, 1}
+
+        self.max_number = config.get("max_number", self.MAX_NUMBER)
+        self.max_steps = config.get("max_steps", self.MAX_STEPS)
+
+        self._number = None
+        self.observation_space = gym.spaces.Discrete(2)
+        self.action_space = gym.spaces.MultiDiscrete([3, self.max_number])
+
+    def reset(self, *, seed=None, options=None):
+        self._step = 0
+        self._number = None
+        # agent 0 has to pick a number. So the returned obs does not matter.
+        return {0: 0}, {}
+
+    def step(self, action_dict):
+        # get agent 0's action
+        agent_0_action = action_dict.get(0)
+
+        if agent_0_action is not None:
+            # ignore the first part of the action and look at the number
+            self._number = agent_0_action[1]
+            # next obs should tell agent 1 to start guessing.
+            # the returned reward and dones should be on agent 0 who picked a
+            # number.
+            return (
+                {1: 0},
+                {0: 0},
+                {0: False, "__all__": False},
+                {0: False, "__all__": False},
+                {},
+            )
+
+        if self._number is None:
+            raise ValueError(
+                "No number is selected by agent 0. Have you restarted "
+                "the environment?"
+            )
+
+        # get agent 1's action
+        direction, number = action_dict.get(1)
+        info = {}
+        # always the same, we don't need agent 0 to act ever again, agent 1 should keep
+        # guessing.
+        obs = {1: 0}
+        guessed_correctly = False
+        terminated = {1: False, "__all__": False}
+        truncated = {1: False, "__all__": False}
+        # everytime agent 1 does not guess correctly agent 0 gets a reward of 1.
+        if direction == 0:  # lower
+            reward = {1: int(number > self._number), 0: 1}
+        elif direction == 1:  # higher
+            reward = {1: int(number < self._number), 0: 1}
+        else:  # equal
+            guessed_correctly = number == self._number
+            reward = {1: guessed_correctly * 100, 0: guessed_correctly * -100}
+            terminated = {1: guessed_correctly, "__all__": guessed_correctly}
+
+        self._step += 1
+        if self._step >= self.max_steps:  # max number of steps episode is over
+            truncated["__all__"] = True
+            if not guessed_correctly:
+                reward[0] = 100  # agent 0 wins
+        return obs, reward, terminated, truncated, info

rllib/examples/envs/classes/multi_agent/pettingzoo_chess.py

@@ -0,0 +1,236 @@
+import copy
+from typing import Any, Dict
+
+import chess as ch
+import numpy as np
+from pettingzoo import AECEnv
+from pettingzoo.classic.chess.chess import raw_env as chess_v5
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class MultiAgentChess(MultiAgentEnv):
+    """An interface to the PettingZoo MARL environment library.
+    See: https://github.com/Farama-Foundation/PettingZoo
+    Inherits from MultiAgentEnv and exposes a given AEC
+    (actor-environment-cycle) game from the PettingZoo project via the
+    MultiAgentEnv public API.
+    Note that the wrapper has some important limitations:
+    1. All agents have the same action_spaces and observation_spaces.
+       Note: If, within your aec game, agents do not have homogeneous action /
+       observation spaces, apply SuperSuit wrappers
+       to apply padding functionality: https://github.com/Farama-Foundation/
+       SuperSuit#built-in-multi-agent-only-functions
+    2. Environments are positive sum games (-> Agents are expected to cooperate
+       to maximize reward). This isn't a hard restriction, it just that
+       standard algorithms aren't expected to work well in highly competitive
+       games.
+
+    .. testcode::
+        :skipif: True
+
+        from pettingzoo.butterfly import prison_v3
+        from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv
+        env = PettingZooEnv(prison_v3.env())
+        obs = env.reset()
+        print(obs)
+        # only returns the observation for the agent which should be stepping
+
+    .. testoutput::
+
+        {
+            'prisoner_0': array([[[0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                ...,
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0]]], dtype=uint8)
+        }
+
+    .. testcode::
+        :skipif: True
+
+        obs, rewards, dones, infos = env.step({
+                        "prisoner_0": 1
+                    })
+        # only returns the observation, reward, info, etc, for
+        # the agent who's turn is next.
+        print(obs)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': array([[[0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                ...,
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0]]], dtype=uint8)
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(rewards)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': 0
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(dones)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': False, '__all__': False
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(infos)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': {'map_tuple': (1, 0)}
+        }
+    """
+
+    def __init__(
+        self,
+        config: Dict[Any, Any] = None,
+        env: AECEnv = None,
+    ):
+        super().__init__()
+        if env is None:
+            self.env = chess_v5()
+        else:
+            self.env = env
+        self.env.reset()
+
+        self.config = config
+        if self.config is None:
+            self.config = {}
+        try:
+            self.config["random_start"] = self.config["random_start"]
+        except KeyError:
+            self.config["random_start"] = 4
+
+        # If these important attributes are not set, try to infer them.
+        if not self.agents:
+            self.agents = list(self._agent_ids)
+        if not self.possible_agents:
+            self.possible_agents = self.agents.copy()
+
+        # Get first observation space, assuming all agents have equal space
+        self.observation_space = self.env.observation_space(self.env.agents[0])
+
+        # Get first action space, assuming all agents have equal space
+        self.action_space = self.env.action_space(self.env.agents[0])
+
+        assert all(
+            self.env.observation_space(agent) == self.observation_space
+            for agent in self.env.agents
+        ), (
+            "Observation spaces for all agents must be identical. Perhaps "
+            "SuperSuit's pad_observations wrapper can help (useage: "
+            "`supersuit.aec_wrappers.pad_observations(env)`"
+        )
+
+        assert all(
+            self.env.action_space(agent) == self.action_space
+            for agent in self.env.agents
+        ), (
+            "Action spaces for all agents must be identical. Perhaps "
+            "SuperSuit's pad_action_space wrapper can help (usage: "
+            "`supersuit.aec_wrappers.pad_action_space(env)`"
+        )
+        self._agent_ids = set(self.env.agents)
+
+    def random_start(self, random_moves):
+        self.env.board = ch.Board()
+        for i in range(random_moves):
+            self.env.board.push(np.random.choice(list(self.env.board.legal_moves)))
+        return self.env.board
+
+    def observe(self):
+        return {
+            self.env.agent_selection: self.env.observe(self.env.agent_selection),
+            "state": self.get_state(),
+        }
+
+    def reset(self, *args, **kwargs):
+        self.env.reset()
+        if self.config["random_start"] > 0:
+            self.random_start(self.config["random_start"])
+        return (
+            {self.env.agent_selection: self.env.observe(self.env.agent_selection)},
+            {self.env.agent_selection: {}},
+        )
+
+    def step(self, action):
+        try:
+            self.env.step(action[self.env.agent_selection])
+        except (KeyError, IndexError):
+            self.env.step(action)
+        except AssertionError:
+            # Illegal action
+            print(action)
+            raise AssertionError("Illegal action")
+
+        obs_d = {}
+        rew_d = {}
+        done_d = {}
+        truncated_d = {}
+        info_d = {}
+        while self.env.agents:
+            obs, rew, done, trunc, info = self.env.last()
+            a = self.env.agent_selection
+            obs_d[a] = obs
+            rew_d[a] = rew
+            done_d[a] = done
+            truncated_d[a] = trunc
+            info_d[a] = info
+            if self.env.terminations[self.env.agent_selection]:
+                self.env.step(None)
+                done_d["__all__"] = True
+                truncated_d["__all__"] = True
+            else:
+                done_d["__all__"] = False
+                truncated_d["__all__"] = False
+                break
+
+        return obs_d, rew_d, done_d, truncated_d, info_d
+
+    def close(self):
+        self.env.close()
+
+    def seed(self, seed=None):
+        self.env.seed(seed)
+
+    def render(self, mode="human"):
+        return self.env.render(mode)
+
+    @property
+    def agent_selection(self):
+        return self.env.agent_selection
+
+    @property
+    def get_sub_environments(self):
+        return self.env.unwrapped
+
+    def get_state(self):
+        state = copy.deepcopy(self.env)
+        return state
+
+    def set_state(self, state):
+        self.env = copy.deepcopy(state)
+        return self.env.observe(self.env.agent_selection)

rllib/examples/envs/classes/multi_agent/pettingzoo_connect4.py

@@ -0,0 +1,220 @@
+import copy
+from typing import Any, Dict
+
+from pettingzoo import AECEnv
+from pettingzoo.classic.connect_four_v3 import raw_env as connect_four_v3
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class MultiAgentConnect4(MultiAgentEnv):
+    """An interface to the PettingZoo MARL environment library.
+    See: https://github.com/Farama-Foundation/PettingZoo
+    Inherits from MultiAgentEnv and exposes a given AEC
+    (actor-environment-cycle) game from the PettingZoo project via the
+    MultiAgentEnv public API.
+    Note that the wrapper has some important limitations:
+    1. All agents have the same action_spaces and observation_spaces.
+       Note: If, within your aec game, agents do not have homogeneous action /
+       observation spaces, apply SuperSuit wrappers
+       to apply padding functionality: https://github.com/Farama-Foundation/
+       SuperSuit#built-in-multi-agent-only-functions
+    2. Environments are positive sum games (-> Agents are expected to cooperate
+       to maximize reward). This isn't a hard restriction, it just that
+       standard algorithms aren't expected to work well in highly competitive
+       games.
+
+    .. testcode::
+        :skipif: True
+
+        from pettingzoo.butterfly import prison_v3
+        from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv
+        env = PettingZooEnv(prison_v3.env())
+        obs = env.reset()
+        print(obs)
+
+    .. testoutput::
+
+        # only returns the observation for the agent which should be stepping
+        {
+            'prisoner_0': array([[[0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                ...,
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0]]], dtype=uint8)
+        }
+
+    .. testcode::
+        :skipif: True
+
+        obs, rewards, dones, infos = env.step({
+                        "prisoner_0": 1
+                    })
+        # only returns the observation, reward, info, etc, for
+        # the agent who's turn is next.
+        print(obs)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': array([[[0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                ...,
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0]]], dtype=uint8)
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(rewards)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': 0
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(dones)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': False, '__all__': False
+        }
+
+    .. testcode::
+        :skipif: True
+
+        print(infos)
+
+    .. testoutput::
+
+        {
+            'prisoner_1': {'map_tuple': (1, 0)}
+        }
+    """
+
+    def __init__(
+        self,
+        config: Dict[Any, Any] = None,
+        env: AECEnv = None,
+    ):
+        super().__init__()
+        if env is None:
+            self.env = connect_four_v3()
+        else:
+            self.env = env
+        self.env.reset()
+
+        # If these important attributes are not set, try to infer them.
+        if not self.agents:
+            self.agents = list(self._agent_ids)
+        if not self.possible_agents:
+            self.possible_agents = self.agents.copy()
+
+        self.config = config
+        # Get first observation space, assuming all agents have equal space
+        self.observation_space = self.env.observation_space(self.env.agents[0])
+
+        # Get first action space, assuming all agents have equal space
+        self.action_space = self.env.action_space(self.env.agents[0])
+
+        assert all(
+            self.env.observation_space(agent) == self.observation_space
+            for agent in self.env.agents
+        ), (
+            "Observation spaces for all agents must be identical. Perhaps "
+            "SuperSuit's pad_observations wrapper can help (useage: "
+            "`supersuit.aec_wrappers.pad_observations(env)`"
+        )
+
+        assert all(
+            self.env.action_space(agent) == self.action_space
+            for agent in self.env.agents
+        ), (
+            "Action spaces for all agents must be identical. Perhaps "
+            "SuperSuit's pad_action_space wrapper can help (usage: "
+            "`supersuit.aec_wrappers.pad_action_space(env)`"
+        )
+
+        self._agent_ids = set(self.env.agents)
+
+    def observe(self):
+        return {
+            self.env.agent_selection: self.env.observe(self.env.agent_selection),
+            "state": self.get_state(),
+        }
+
+    def reset(self, *args, **kwargs):
+        self.env.reset()
+        return (
+            {self.env.agent_selection: self.env.observe(self.env.agent_selection)},
+            {self.env.agent_selection: {}},
+        )
+
+    def step(self, action):
+        try:
+            self.env.step(action[self.env.agent_selection])
+        except (KeyError, IndexError):
+            self.env.step(action)
+        except AssertionError:
+            # Illegal action
+            print(action)
+            raise AssertionError("Illegal action")
+
+        obs_d = {}
+        rew_d = {}
+        done_d = {}
+        trunc_d = {}
+        info_d = {}
+        while self.env.agents:
+            obs, rew, done, trunc, info = self.env.last()
+            a = self.env.agent_selection
+            obs_d[a] = obs
+            rew_d[a] = rew
+            done_d[a] = done
+            trunc_d[a] = trunc
+            info_d[a] = info
+            if self.env.terminations[self.env.agent_selection]:
+                self.env.step(None)
+                done_d["__all__"] = True
+                trunc_d["__all__"] = True
+            else:
+                done_d["__all__"] = False
+                trunc_d["__all__"] = False
+                break
+
+        return obs_d, rew_d, done_d, trunc_d, info_d
+
+    def close(self):
+        self.env.close()
+
+    def seed(self, seed=None):
+        self.env.seed(seed)
+
+    def render(self, mode="human"):
+        return self.env.render(mode)
+
+    @property
+    def agent_selection(self):
+        return self.env.agent_selection
+
+    @property
+    def get_sub_environments(self):
+        return self.env.unwrapped
+
+    def get_state(self):
+        state = copy.deepcopy(self.env)
+        return state
+
+    def set_state(self, state):
+        self.env = copy.deepcopy(state)
+        return self.env.observe(self.env.agent_selection)

rllib/examples/envs/classes/multi_agent/rock_paper_scissors.py

@@ -0,0 +1,125 @@
+# __sphinx_doc_1_begin__
+import gymnasium as gym
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class RockPaperScissors(MultiAgentEnv):
+    """Two-player environment for the famous rock paper scissors game.
+
+    # __sphinx_doc_1_end__
+    Optionally, the "Sheldon Cooper extension" can be activated by passing
+    `sheldon_cooper_mode=True` into the constructor, in which case two more moves
+    are allowed: Spock and Lizard. Spock is poisoned by Lizard, disproven by Paper, but
+    crushes Rock and smashes Scissors. Lizard poisons Spock and eats Paper, but is
+    decapitated by Scissors and crushed by Rock.
+
+    # __sphinx_doc_2_begin__
+    Both players always move simultaneously over a course of 10 timesteps in total.
+    The winner of each timestep receives reward of +1, the losing player -1.0.
+
+    The observation of each player is the last opponent action.
+    """
+
+    ROCK = 0
+    PAPER = 1
+    SCISSORS = 2
+    LIZARD = 3
+    SPOCK = 4
+
+    WIN_MATRIX = {
+        (ROCK, ROCK): (0, 0),
+        (ROCK, PAPER): (-1, 1),
+        (ROCK, SCISSORS): (1, -1),
+        (PAPER, ROCK): (1, -1),
+        (PAPER, PAPER): (0, 0),
+        (PAPER, SCISSORS): (-1, 1),
+        (SCISSORS, ROCK): (-1, 1),
+        (SCISSORS, PAPER): (1, -1),
+        (SCISSORS, SCISSORS): (0, 0),
+    }
+    # __sphinx_doc_2_end__
+
+    WIN_MATRIX.update(
+        {
+            # Sheldon Cooper mode:
+            (LIZARD, LIZARD): (0, 0),
+            (LIZARD, SPOCK): (1, -1),  # Lizard poisons Spock
+            (LIZARD, ROCK): (-1, 1),  # Rock crushes lizard
+            (LIZARD, PAPER): (1, -1),  # Lizard eats paper
+            (LIZARD, SCISSORS): (-1, 1),  # Scissors decapitate lizard
+            (ROCK, LIZARD): (1, -1),  # Rock crushes lizard
+            (PAPER, LIZARD): (-1, 1),  # Lizard eats paper
+            (SCISSORS, LIZARD): (1, -1),  # Scissors decapitate lizard
+            (SPOCK, SPOCK): (0, 0),
+            (SPOCK, LIZARD): (-1, 1),  # Lizard poisons Spock
+            (SPOCK, ROCK): (1, -1),  # Spock vaporizes rock
+            (SPOCK, PAPER): (-1, 1),  # Paper disproves Spock
+            (SPOCK, SCISSORS): (1, -1),  # Spock smashes scissors
+            (ROCK, SPOCK): (-1, 1),  # Spock vaporizes rock
+            (PAPER, SPOCK): (1, -1),  # Paper disproves Spock
+            (SCISSORS, SPOCK): (-1, 1),  # Spock smashes scissors
+        }
+    )
+
+    # __sphinx_doc_3_begin__
+    def __init__(self, config=None):
+        super().__init__()
+
+        self.agents = self.possible_agents = ["player1", "player2"]
+
+        # The observations are always the last taken actions. Hence observation- and
+        # action spaces are identical.
+        self.observation_spaces = self.action_spaces = {
+            "player1": gym.spaces.Discrete(3),
+            "player2": gym.spaces.Discrete(3),
+        }
+        self.last_move = None
+        self.num_moves = 0
+        # __sphinx_doc_3_end__
+
+        self.sheldon_cooper_mode = False
+        if config.get("sheldon_cooper_mode"):
+            self.sheldon_cooper_mode = True
+            self.action_spaces = self.observation_spaces = {
+                "player1": gym.spaces.Discrete(5),
+                "player2": gym.spaces.Discrete(5),
+            }
+
+    # __sphinx_doc_4_begin__
+    def reset(self, *, seed=None, options=None):
+        self.num_moves = 0
+
+        # The first observation should not matter (none of the agents has moved yet).
+        # Set them to 0.
+        return {
+            "player1": 0,
+            "player2": 0,
+        }, {}  # <- empty infos dict
+
+    # __sphinx_doc_4_end__
+
+    # __sphinx_doc_5_begin__
+    def step(self, action_dict):
+        self.num_moves += 1
+
+        move1 = action_dict["player1"]
+        move2 = action_dict["player2"]
+
+        # Set the next observations (simply use the other player's action).
+        # Note that because we are publishing both players in the observations dict,
+        # we expect both players to act in the next `step()` (simultaneous stepping).
+        observations = {"player1": move2, "player2": move1}
+
+        # Compute rewards for each player based on the win-matrix.
+        r1, r2 = self.WIN_MATRIX[move1, move2]
+        rewards = {"player1": r1, "player2": r2}
+
+        # Terminate the entire episode (for all agents) once 10 moves have been made.
+        terminateds = {"__all__": self.num_moves >= 10}
+
+        # Leave truncateds and infos empty.
+        return observations, rewards, terminateds, {}, {}
+
+
+# __sphinx_doc_5_end__

rllib/examples/envs/classes/multi_agent/tic_tac_toe.py

@@ -0,0 +1,183 @@
+# __sphinx_doc_1_begin__
+import random
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class TicTacToe(MultiAgentEnv):
+    """A two-player game in which any player tries to complete one row in a 3x3 field.
+
+    The observation space is Box(-1.0, 1.0, (9,)), where each index represents a distinct
+    field on a 3x3 board. From the current player's perspective: 1.0 means we occupy the
+    field, -1.0 means the opponent owns the field, and 0.0 means the field is empty:
+    ----------
+    | 0| 1| 2|
+    ----------
+    | 3| 4| 5|
+    ----------
+    | 6| 7| 8|
+    ----------
+
+    The action space is Discrete(9). Actions landing on an already occupied field
+    result in a -1.0 penalty for the player taking the invalid action.
+
+    Once a player completes a row, they receive +1.0 reward, the losing player receives
+    -1.0 reward. A draw results in 0.0 reward for both players.
+    """
+
+    # __sphinx_doc_1_end__
+
+    # Winning line indices: rows, columns, and diagonals.
+    WIN_LINES = [
+        [0, 1, 2],  # rows
+        [3, 4, 5],
+        [6, 7, 8],
+        [0, 3, 6],  # cols
+        [1, 4, 7],
+        [2, 5, 8],
+        [0, 4, 8],  # diagonals
+        [2, 4, 6],
+    ]
+
+    # __sphinx_doc_2_begin__
+    def __init__(self, config=None):
+        super().__init__()
+
+        # Define the agents in the game.
+        self.agents = self.possible_agents = ["player1", "player2"]
+
+        # Each agent observes a 9D tensor, representing the 3x3 fields of the board.
+        # From the current player's perspective: 1 means our piece, -1 means opponent's
+        # piece, 0 means empty. The board is flipped after each turn.
+        self.observation_spaces = {
+            "player1": gym.spaces.Box(-1.0, 1.0, (9,), np.float32),
+            "player2": gym.spaces.Box(-1.0, 1.0, (9,), np.float32),
+        }
+        # Each player has 9 actions, encoding the 9 fields each player can place a piece
+        # on during their turn.
+        self.action_spaces = {
+            "player1": gym.spaces.Discrete(9),
+            "player2": gym.spaces.Discrete(9),
+        }
+        self.max_timesteps = 30
+
+        self.board = None
+        self.current_player = None
+        self.timestep = 0
+
+    # __sphinx_doc_2_end__
+
+    # __sphinx_doc_3_begin__
+    def reset(self, *, seed=None, options=None):
+        self.board = [0] * 9
+
+        # Pick a random player to start the game and reset the current timesteps.
+        self.current_player = random.choice(self.agents)
+        self.timestep = 0
+
+        # Return observations dict (only with the starting player, which is the one
+        # we expect to act next).
+        return {self.current_player: np.array(self.board, np.float32)}, {}
+
+    # __sphinx_doc_3_end__
+
+    # __sphinx_doc_4_begin__
+    def step(self, action_dict):
+        action = action_dict[self.current_player]
+        opponent = "player2" if self.current_player == "player1" else "player1"
+        self.timestep += 1
+
+        # Invalid move: penalize and return without changing board.
+        if self.board[action] != 0:
+            # The time limit is reached
+            if self.timestep >= self.max_timesteps:
+                board_arr = np.array(self.board, np.float32)
+                return (
+                    {self.current_player: board_arr, opponent: board_arr * -1},
+                    {self.current_player: -0.5, opponent: 0.0},
+                    {"__all__": False},
+                    {"__all__": True},
+                    {},
+                )
+            else:
+                reward = {self.current_player: -0.5}
+                self.board = [-x for x in self.board]
+                self.current_player = opponent
+                return (
+                    {opponent: np.array(self.board, np.float32)},
+                    reward,
+                    {"__all__": False},
+                    {"__all__": False},
+                    {},
+                )
+
+        # Place the piece on the board.
+        self.board[action] = 1
+
+        # Check for win.
+        if any(all(self.board[i] == 1 for i in line) for line in self.WIN_LINES):
+            board_arr = np.array(self.board, np.float32)
+            return (
+                {self.current_player: board_arr, opponent: board_arr * -1},
+                {self.current_player: 1.0, opponent: -1.0},
+                {"__all__": True},
+                {"__all__": False},
+                {},
+            )
+
+        # Check for draw (board full, no winner).
+        if 0 not in self.board:
+            board_arr = np.array(self.board, np.float32)
+            return (
+                {self.current_player: board_arr, opponent: board_arr * -1},
+                {self.current_player: 0.0, opponent: 0.0},
+                {"__all__": True},
+                {"__all__": False},
+                {},
+            )
+
+        # Check for truncation.
+        if self.timestep >= self.max_timesteps:
+            board_arr = np.array(self.board, np.float32)
+            return (
+                {self.current_player: board_arr, opponent: board_arr * -1},
+                {self.current_player: 0.0, opponent: 0.0},
+                {"__all__": False},
+                {"__all__": True},
+                {},
+            )
+
+        # Continue game: flip board and switch player.
+        reward = {self.current_player: 0.0}
+        self.board = [-x for x in self.board]
+        self.current_player = opponent
+
+        return (
+            {opponent: np.array(self.board, np.float32)},
+            reward,
+            {"__all__": False},
+            {"__all__": False},
+            {},
+        )
+
+    # __sphinx_doc_4_end__
+
+    def render(self) -> str:
+        """Render the current board state as an ASCII grid.
+
+        Returns:
+            A string representation of the board where:
+            - 'X' represents the current player's pieces
+            - 'O' represents opponent player's pieces
+            - ' ' represents empty fields
+        """
+        symbols = {0: " ", 1: "X", -1: "O"}
+        rows = []
+        for i in range(3):
+            row_cells = [symbols[self.board[i * 3 + j]] for j in range(3)]
+            rows.append(" " + " | ".join(row_cells) + " ")
+        separator = "-----------"
+        return "\n" + f"\n{separator}\n".join(rows) + "\n"

rllib/examples/envs/classes/multi_agent/two_step_game.py

@@ -0,0 +1,125 @@
+import numpy as np
+from gymnasium.spaces import Dict, Discrete, MultiDiscrete, Tuple
+
+from ray.rllib.env.multi_agent_env import ENV_STATE, MultiAgentEnv
+
+
+class TwoStepGame(MultiAgentEnv):
+    action_space = Discrete(2)
+
+    def __init__(self, env_config):
+        super().__init__()
+        self.action_space = Discrete(2)
+        self.state = None
+        self.agent_1 = 0
+        self.agent_2 = 1
+        # MADDPG emits action logits instead of actual discrete actions
+        self.actions_are_logits = env_config.get("actions_are_logits", False)
+        self.one_hot_state_encoding = env_config.get("one_hot_state_encoding", False)
+        self.with_state = env_config.get("separate_state_space", False)
+        self._agent_ids = {0, 1}
+        if not self.one_hot_state_encoding:
+            self.observation_space = Discrete(6)
+            self.with_state = False
+        else:
+            # Each agent gets the full state (one-hot encoding of which of the
+            # three states are active) as input with the receiving agent's
+            # ID (1 or 2) concatenated onto the end.
+            if self.with_state:
+                self.observation_space = Dict(
+                    {
+                        "obs": MultiDiscrete([2, 2, 2, 3]),
+                        ENV_STATE: MultiDiscrete([2, 2, 2]),
+                    }
+                )
+            else:
+                self.observation_space = MultiDiscrete([2, 2, 2, 3])
+
+    def reset(self, *, seed=None, options=None):
+        if seed is not None:
+            np.random.seed(seed)
+        self.state = np.array([1, 0, 0])
+        return self._obs(), {}
+
+    def step(self, action_dict):
+        if self.actions_are_logits:
+            action_dict = {
+                k: np.random.choice([0, 1], p=v) for k, v in action_dict.items()
+            }
+
+        state_index = np.flatnonzero(self.state)
+        if state_index == 0:
+            action = action_dict[self.agent_1]
+            assert action in [0, 1], action
+            if action == 0:
+                self.state = np.array([0, 1, 0])
+            else:
+                self.state = np.array([0, 0, 1])
+            global_rew = 0
+            terminated = False
+        elif state_index == 1:
+            global_rew = 7
+            terminated = True
+        else:
+            if action_dict[self.agent_1] == 0 and action_dict[self.agent_2] == 0:
+                global_rew = 0
+            elif action_dict[self.agent_1] == 1 and action_dict[self.agent_2] == 1:
+                global_rew = 8
+            else:
+                global_rew = 1
+            terminated = True
+
+        rewards = {self.agent_1: global_rew / 2.0, self.agent_2: global_rew / 2.0}
+        obs = self._obs()
+        terminateds = {"__all__": terminated}
+        truncateds = {"__all__": False}
+        infos = {
+            self.agent_1: {"done": terminateds["__all__"]},
+            self.agent_2: {"done": terminateds["__all__"]},
+        }
+        return obs, rewards, terminateds, truncateds, infos
+
+    def _obs(self):
+        if self.with_state:
+            return {
+                self.agent_1: {"obs": self.agent_1_obs(), ENV_STATE: self.state},
+                self.agent_2: {"obs": self.agent_2_obs(), ENV_STATE: self.state},
+            }
+        else:
+            return {self.agent_1: self.agent_1_obs(), self.agent_2: self.agent_2_obs()}
+
+    def agent_1_obs(self):
+        if self.one_hot_state_encoding:
+            return np.concatenate([self.state, [1]])
+        else:
+            return np.flatnonzero(self.state)[0]
+
+    def agent_2_obs(self):
+        if self.one_hot_state_encoding:
+            return np.concatenate([self.state, [2]])
+        else:
+            return np.flatnonzero(self.state)[0] + 3
+
+
+class TwoStepGameWithGroupedAgents(MultiAgentEnv):
+    def __init__(self, env_config):
+        self._agent_ids = {"agents"}
+
+        super().__init__()
+        env = TwoStepGame(env_config)
+        tuple_obs_space = Tuple([env.observation_space, env.observation_space])
+        tuple_act_space = Tuple([env.action_space, env.action_space])
+
+        self.env = env.with_agent_groups(
+            groups={"agents": [0, 1]},
+            obs_space=tuple_obs_space,
+            act_space=tuple_act_space,
+        )
+        self.observation_space = Dict({"agents": self.env.observation_space})
+        self.action_space = Dict({"agents": self.env.action_space})
+
+    def reset(self, *, seed=None, options=None):
+        return self.env.reset(seed=seed, options=options)
+
+    def step(self, actions):
+        return self.env.step(actions)

rllib/examples/envs/classes/nested_space_repeat_after_me_env.py

@@ -0,0 +1,50 @@
+import gymnasium as gym
+import numpy as np
+import tree  # pip install dm_tree
+from gymnasium.spaces import Box, Dict, Discrete, Tuple
+
+from ray.rllib.utils.spaces.space_utils import flatten_space
+
+
+class NestedSpaceRepeatAfterMeEnv(gym.Env):
+    """Env for which policy has to repeat the (possibly complex) observation.
+
+    The action space and observation spaces are always the same and may be
+    arbitrarily nested Dict/Tuple Spaces.
+    Rewards are given for exactly matching Discrete sub-actions and for being
+    as close as possible for Box sub-actions.
+    """
+
+    def __init__(self, config=None):
+        config = config or {}
+        self.observation_space = config.get(
+            "space", Tuple([Discrete(2), Dict({"a": Box(-1.0, 1.0, (2,))})])
+        )
+        self.action_space = self.observation_space
+        self.flattened_action_space = flatten_space(self.action_space)
+        self.episode_len = config.get("episode_len", 100)
+
+    def reset(self, *, seed=None, options=None):
+        self.steps = 0
+        return self._next_obs(), {}
+
+    def step(self, action):
+        self.steps += 1
+        action = tree.flatten(action)
+        reward = 0.0
+        for a, o, space in zip(
+            action, self.current_obs_flattened, self.flattened_action_space
+        ):
+            # Box: -abs(diff).
+            if isinstance(space, gym.spaces.Box):
+                reward -= np.sum(np.abs(a - o))
+            # Discrete: +1.0 if exact match.
+            if isinstance(space, gym.spaces.Discrete):
+                reward += 1.0 if a == o else 0.0
+        truncated = self.steps >= self.episode_len
+        return self._next_obs(), reward, False, truncated, {}
+
+    def _next_obs(self):
+        self.current_obs = self.observation_space.sample()
+        self.current_obs_flattened = tree.flatten(self.current_obs)
+        return self.current_obs

rllib/examples/envs/classes/parametric_actions_cartpole.py

@@ -0,0 +1,145 @@
+import random
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Dict, Discrete
+
+
+class ParametricActionsCartPole(gym.Env):
+    """Parametric action version of CartPole.
+
+    In this env there are only ever two valid actions, but we pretend there are
+    actually up to `max_avail_actions` actions that can be taken, and the two
+    valid actions are randomly hidden among this set.
+
+    At each step, we emit a dict of:
+        - the actual cart observation
+        - a mask of valid actions (e.g., [0, 0, 1, 0, 0, 1] for 6 max avail)
+        - the list of action embeddings (w/ zeroes for invalid actions) (e.g.,
+            [[0, 0],
+             [0, 0],
+             [-0.2322, -0.2569],
+             [0, 0],
+             [0, 0],
+             [0.7878, 1.2297]] for max_avail_actions=6)
+
+    In a real environment, the actions embeddings would be larger than two
+    units of course, and also there would be a variable number of valid actions
+    per step instead of always [LEFT, RIGHT].
+    """
+
+    def __init__(self, max_avail_actions):
+        # Use simple random 2-unit action embeddings for [LEFT, RIGHT]
+        self.left_action_embed = np.random.randn(2)
+        self.right_action_embed = np.random.randn(2)
+        self.action_space = Discrete(max_avail_actions)
+        self.wrapped = gym.make("CartPole-v1")
+        self.observation_space = Dict(
+            {
+                "action_mask": Box(0, 1, shape=(max_avail_actions,), dtype=np.int8),
+                "avail_actions": Box(-10, 10, shape=(max_avail_actions, 2)),
+                "cart": self.wrapped.observation_space,
+            }
+        )
+
+    def update_avail_actions(self):
+        self.action_assignments = np.array(
+            [[0.0, 0.0]] * self.action_space.n, dtype=np.float32
+        )
+        self.action_mask = np.array([0.0] * self.action_space.n, dtype=np.int8)
+        self.left_idx, self.right_idx = random.sample(range(self.action_space.n), 2)
+        self.action_assignments[self.left_idx] = self.left_action_embed
+        self.action_assignments[self.right_idx] = self.right_action_embed
+        self.action_mask[self.left_idx] = 1
+        self.action_mask[self.right_idx] = 1
+
+    def reset(self, *, seed=None, options=None):
+        self.update_avail_actions()
+        obs, infos = self.wrapped.reset()
+        return {
+            "action_mask": self.action_mask,
+            "avail_actions": self.action_assignments,
+            "cart": obs,
+        }, infos
+
+    def step(self, action):
+        if action == self.left_idx:
+            actual_action = 0
+        elif action == self.right_idx:
+            actual_action = 1
+        else:
+            raise ValueError(
+                "Chosen action was not one of the non-zero action embeddings",
+                action,
+                self.action_assignments,
+                self.action_mask,
+                self.left_idx,
+                self.right_idx,
+            )
+        orig_obs, rew, done, truncated, info = self.wrapped.step(actual_action)
+        self.update_avail_actions()
+        self.action_mask = self.action_mask.astype(np.int8)
+        obs = {
+            "action_mask": self.action_mask,
+            "avail_actions": self.action_assignments,
+            "cart": orig_obs,
+        }
+        return obs, rew, done, truncated, info
+
+
+class ParametricActionsCartPoleNoEmbeddings(gym.Env):
+    """Same as the above ParametricActionsCartPole.
+
+    However, action embeddings are not published inside observations,
+    but will be learnt by the model.
+
+    At each step, we emit a dict of:
+        - the actual cart observation
+        - a mask of valid actions (e.g., [0, 0, 1, 0, 0, 1] for 6 max avail)
+        - action embeddings (w/ "dummy embedding" for invalid actions) are
+          outsourced in the model and will be learned.
+    """
+
+    def __init__(self, max_avail_actions):
+        # Randomly set which two actions are valid and available.
+        self.left_idx, self.right_idx = random.sample(range(max_avail_actions), 2)
+        self.valid_avail_actions_mask = np.array(
+            [0.0] * max_avail_actions, dtype=np.int8
+        )
+        self.valid_avail_actions_mask[self.left_idx] = 1
+        self.valid_avail_actions_mask[self.right_idx] = 1
+        self.action_space = Discrete(max_avail_actions)
+        self.wrapped = gym.make("CartPole-v1")
+        self.observation_space = Dict(
+            {
+                "valid_avail_actions_mask": Box(0, 1, shape=(max_avail_actions,)),
+                "cart": self.wrapped.observation_space,
+            }
+        )
+
+    def reset(self, *, seed=None, options=None):
+        obs, infos = self.wrapped.reset()
+        return {
+            "valid_avail_actions_mask": self.valid_avail_actions_mask,
+            "cart": obs,
+        }, infos
+
+    def step(self, action):
+        if action == self.left_idx:
+            actual_action = 0
+        elif action == self.right_idx:
+            actual_action = 1
+        else:
+            raise ValueError(
+                "Chosen action was not one of the non-zero action embeddings",
+                action,
+                self.valid_avail_actions_mask,
+                self.left_idx,
+                self.right_idx,
+            )
+        orig_obs, rew, done, truncated, info = self.wrapped.step(actual_action)
+        obs = {
+            "valid_avail_actions_mask": self.valid_avail_actions_mask,
+            "cart": orig_obs,
+        }
+        return obs, rew, done, truncated, info

rllib/examples/envs/classes/random_env.py

@@ -0,0 +1,126 @@
+import copy
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Discrete, Tuple
+
+from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent
+
+
+class RandomEnv(gym.Env):
+    """A randomly acting environment.
+
+    Can be instantiated with arbitrary action-, observation-, and reward
+    spaces. Observations and rewards are generated by simply sampling from the
+    observation/reward spaces. The probability of a `terminated=True` after each
+    action can be configured, as well as the max episode length.
+    """
+
+    def __init__(self, config=None):
+        config = config or {}
+
+        # Action space.
+        self.action_space = config.get("action_space", Discrete(2))
+        # Observation space from which to sample.
+        self.observation_space = config.get("observation_space", Discrete(2))
+        # Reward space from which to sample.
+        self.reward_space = config.get(
+            "reward_space",
+            gym.spaces.Box(low=-1.0, high=1.0, shape=(), dtype=np.float32),
+        )
+        self.static_samples = config.get("static_samples", False)
+        if self.static_samples:
+            self.observation_sample = self.observation_space.sample()
+            self.reward_sample = self.reward_space.sample()
+
+        # Chance that an episode ends at any step.
+        # Note that a max episode length can be specified via
+        # `max_episode_len`.
+        self.p_terminated = config.get("p_terminated")
+        if self.p_terminated is None:
+            self.p_terminated = config.get("p_done", 0.1)
+        # A max episode length. Even if the `p_terminated` sampling does not lead
+        # to a terminus, the episode will end after at most this many
+        # timesteps.
+        # Set to 0 or None for using no limit on the episode length.
+        self.max_episode_len = config.get("max_episode_len", None)
+        # Whether to check action bounds.
+        self.check_action_bounds = config.get("check_action_bounds", False)
+        # Steps taken so far (after last reset).
+        self.steps = 0
+
+    def reset(self, *, seed=None, options=None):
+        self.steps = 0
+        if not self.static_samples:
+            return self.observation_space.sample(), {}
+        else:
+            return copy.deepcopy(self.observation_sample), {}
+
+    def step(self, action):
+        if self.check_action_bounds and not self.action_space.contains(action):
+            raise ValueError(
+                "Illegal action for {}: {}".format(self.action_space, action)
+            )
+        if isinstance(self.action_space, Tuple) and len(action) != len(
+            self.action_space.spaces
+        ):
+            raise ValueError(
+                "Illegal action for {}: {}".format(self.action_space, action)
+            )
+
+        self.steps += 1
+        terminated = False
+        truncated = False
+        # We are `truncated` as per our max-episode-len.
+        if self.max_episode_len and self.steps >= self.max_episode_len:
+            truncated = True
+        # Max episode length not reached yet -> Sample `terminated` via `p_terminated`.
+        elif self.p_terminated > 0.0:
+            terminated = bool(
+                np.random.choice(
+                    [True, False], p=[self.p_terminated, 1.0 - self.p_terminated]
+                )
+            )
+
+        if not self.static_samples:
+            return (
+                self.observation_space.sample(),
+                self.reward_space.sample(),
+                terminated,
+                truncated,
+                {},
+            )
+        else:
+            return (
+                copy.deepcopy(self.observation_sample),
+                copy.deepcopy(self.reward_sample),
+                terminated,
+                truncated,
+                {},
+            )
+
+
+# Multi-agent version of the RandomEnv.
+RandomMultiAgentEnv = make_multi_agent(lambda c: RandomEnv(c))
+
+
+# Large observation space "pre-compiled" random env (for testing).
+class RandomLargeObsSpaceEnv(RandomEnv):
+    def __init__(self, config=None):
+        config = config or {}
+        config.update({"observation_space": gym.spaces.Box(-1.0, 1.0, (5000,))})
+        super().__init__(config=config)
+
+
+# Large observation space + cont. actions "pre-compiled" random env
+# (for testing).
+class RandomLargeObsSpaceEnvContActions(RandomEnv):
+    def __init__(self, config=None):
+        config = config or {}
+        config.update(
+            {
+                "observation_space": gym.spaces.Box(-1.0, 1.0, (5000,)),
+                "action_space": gym.spaces.Box(-1.0, 1.0, (5,)),
+            }
+        )
+        super().__init__(config=config)

rllib/examples/envs/classes/recommender_system_envs_with_recsim.py

@@ -0,0 +1,108 @@
+"""Examples for RecSim envs ready to be used by RLlib Algorithms.
+
+RecSim is a configurable recommender systems simulation platform.
+Source: https://github.com/google-research/recsim
+"""
+
+from recsim import choice_model
+from recsim.environments import (
+    interest_evolution as iev,
+    interest_exploration as iex,
+    long_term_satisfaction as lts,
+)
+
+from ray.rllib.env.wrappers.recsim import make_recsim_env
+from ray.tune import register_env
+
+# Some built-in RecSim envs to test with.
+# ---------------------------------------
+
+# Long-term satisfaction env: User has to pick from items that are either
+# a) unhealthy, but taste good, or b) healthy, but have bad taste.
+# Best strategy is to pick a mix of both to ensure long-term
+# engagement.
+
+
+def lts_user_model_creator(env_ctx):
+    return lts.LTSUserModel(
+        env_ctx["slate_size"],
+        user_state_ctor=lts.LTSUserState,
+        response_model_ctor=lts.LTSResponse,
+    )
+
+
+def lts_document_sampler_creator(env_ctx):
+    return lts.LTSDocumentSampler()
+
+
+LongTermSatisfactionRecSimEnv = make_recsim_env(
+    recsim_user_model_creator=lts_user_model_creator,
+    recsim_document_sampler_creator=lts_document_sampler_creator,
+    reward_aggregator=lts.clicked_engagement_reward,
+)
+
+
+# Interest exploration env: Models the problem of active exploration
+# of user interests. It is meant to illustrate popularity bias in
+# recommender systems, where myopic maximization of engagement leads
+# to bias towards documents that have wider appeal,
+# whereas niche user interests remain unexplored.
+def iex_user_model_creator(env_ctx):
+    return iex.IEUserModel(
+        env_ctx["slate_size"],
+        user_state_ctor=iex.IEUserState,
+        response_model_ctor=iex.IEResponse,
+        seed=env_ctx["seed"],
+    )
+
+
+def iex_document_sampler_creator(env_ctx):
+    return iex.IETopicDocumentSampler(seed=env_ctx["seed"])
+
+
+InterestExplorationRecSimEnv = make_recsim_env(
+    recsim_user_model_creator=iex_user_model_creator,
+    recsim_document_sampler_creator=iex_document_sampler_creator,
+    reward_aggregator=iex.total_clicks_reward,
+)
+
+
+# Interest evolution env: See https://github.com/google-research/recsim
+# for more information.
+def iev_user_model_creator(env_ctx):
+    return iev.IEvUserModel(
+        env_ctx["slate_size"],
+        choice_model_ctor=choice_model.MultinomialProportionalChoiceModel,
+        response_model_ctor=iev.IEvResponse,
+        user_state_ctor=iev.IEvUserState,
+        seed=env_ctx["seed"],
+    )
+
+
+# Extend IEvVideo to fix a bug caused by None cluster_ids.
+class SingleClusterIEvVideo(iev.IEvVideo):
+    def __init__(self, doc_id, features, video_length=None, quality=None):
+        super(SingleClusterIEvVideo, self).__init__(
+            doc_id=doc_id,
+            features=features,
+            cluster_id=0,  # single cluster.
+            video_length=video_length,
+            quality=quality,
+        )
+
+
+def iev_document_sampler_creator(env_ctx):
+    return iev.UtilityModelVideoSampler(doc_ctor=iev.IEvVideo, seed=env_ctx["seed"])
+
+
+InterestEvolutionRecSimEnv = make_recsim_env(
+    recsim_user_model_creator=iev_user_model_creator,
+    recsim_document_sampler_creator=iev_document_sampler_creator,
+    reward_aggregator=iev.clicked_watchtime_reward,
+)
+
+
+# Backward compatibility.
+register_env(
+    name="RecSim-v1", env_creator=lambda env_ctx: InterestEvolutionRecSimEnv(env_ctx)
+)

rllib/examples/envs/classes/repeat_after_me_env.py

@@ -0,0 +1,47 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+
+class RepeatAfterMeEnv(gym.Env):
+    """Env in which the observation at timestep minus n must be repeated."""
+
+    def __init__(self, config=None):
+        config = config or {}
+        if config.get("continuous"):
+            self.observation_space = Box(-1.0, 1.0, (2,))
+        else:
+            self.observation_space = Discrete(2)
+
+        self.action_space = self.observation_space
+        # Note: Set `repeat_delay` to 0 for simply repeating the seen
+        # observation (no delay).
+        self.delay = config.get("repeat_delay", 1)
+        self.episode_len = config.get("episode_len", 100)
+        self.history = []
+
+    def reset(self, *, seed=None, options=None):
+        self.history = [0] * self.delay
+        return self._next_obs(), {}
+
+    def step(self, action):
+        obs = self.history[-(1 + self.delay)]
+
+        reward = 0.0
+        # Box: -abs(diff).
+        if isinstance(self.action_space, Box):
+            reward = -np.sum(np.abs(action - obs))
+        # Discrete: +1.0 if exact match, -1.0 otherwise.
+        if isinstance(self.action_space, Discrete):
+            reward = 1.0 if action == obs else -1.0
+
+        done = truncated = len(self.history) > self.episode_len
+        return self._next_obs(), reward, done, truncated, {}
+
+    def _next_obs(self):
+        if isinstance(self.observation_space, Box):
+            token = np.random.random(size=(2,))
+        else:
+            token = np.random.choice([0, 1])
+        self.history.append(token)
+        return token

rllib/examples/envs/classes/repeat_initial_obs_env.py

@@ -0,0 +1,33 @@
+import random
+
+import gymnasium as gym
+from gymnasium.spaces import Discrete
+
+
+class RepeatInitialObsEnv(gym.Env):
+    """Env in which the initial observation has to be repeated all the time.
+
+    Runs for n steps.
+    r=1 if action correct, -1 otherwise (max. R=100).
+    """
+
+    def __init__(self, episode_len=100):
+        self.observation_space = Discrete(2)
+        self.action_space = Discrete(2)
+        self.token = None
+        self.episode_len = episode_len
+        self.num_steps = 0
+
+    def reset(self, *, seed=None, options=None):
+        self.token = random.choice([0, 1])
+        self.num_steps = 0
+        return self.token, {}
+
+    def step(self, action):
+        if action == self.token:
+            reward = 1
+        else:
+            reward = -1
+        self.num_steps += 1
+        done = truncated = self.num_steps >= self.episode_len
+        return 0, reward, done, truncated, {}

rllib/examples/envs/classes/simple_corridor.py

@@ -0,0 +1,46 @@
+import logging
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+logger = logging.getLogger("ray.rllib")
+
+
+class SimpleCorridor(gym.Env):
+    """Example of a custom env in which you have to walk down a corridor.
+
+    You can configure the length of the corridor via the env config."""
+
+    def __init__(self, config=None):
+        config = config or {}
+
+        self.action_space = Discrete(2)
+        self.observation_space = Box(0.0, 999.0, shape=(1,), dtype=np.float32)
+
+        self.set_corridor_length(config.get("corridor_length", 10))
+
+        self._cur_pos = 0
+
+    def set_corridor_length(self, length):
+        self.end_pos = length
+        logger.info(f"Set corridor length to {self.end_pos}")
+        assert self.end_pos <= 999, "The maximum `corridor_length` allowed is 999!"
+
+    def reset(self, *, seed=None, options=None):
+        self._cur_pos = 0.0
+        return self._get_obs(), {}
+
+    def step(self, action):
+        assert action in [0, 1], action
+        if action == 0 and self._cur_pos > 0:
+            self._cur_pos -= 1.0
+        elif action == 1:
+            self._cur_pos += 1.0
+        terminated = self._cur_pos >= self.end_pos
+        truncated = False
+        reward = 1.0 if terminated else -0.01
+        return self._get_obs(), reward, terminated, truncated, {}
+
+    def _get_obs(self):
+        return np.array([self._cur_pos], np.float32)

rllib/examples/envs/classes/simple_rpg.py

@@ -0,0 +1,49 @@
+import gymnasium as gym
+from gymnasium.spaces import Box, Dict, Discrete
+
+from ray.rllib.utils.spaces.repeated import Repeated
+
+# Constraints on the Repeated space.
+MAX_PLAYERS = 4
+MAX_ITEMS = 7
+MAX_EFFECTS = 2
+
+
+class SimpleRPG(gym.Env):
+    """Example of a custom env with a complex, structured observation.
+
+    The observation is a list of players, each of which is a Dict of
+    attributes, and may further hold a list of items (categorical space).
+
+    Note that the env doesn't train, it's just a dummy example to show how to
+    use spaces.Repeated in a custom model (see CustomRPGModel below).
+    """
+
+    def __init__(self, config):
+        self.cur_pos = 0
+        self.action_space = Discrete(4)
+
+        # Represents an item.
+        self.item_space = Discrete(5)
+
+        # Represents an effect on the player.
+        self.effect_space = Box(9000, 9999, shape=(4,))
+
+        # Represents a player.
+        self.player_space = Dict(
+            {
+                "location": Box(-100, 100, shape=(2,)),
+                "status": Box(-1, 1, shape=(10,)),
+                "items": Repeated(self.item_space, max_len=MAX_ITEMS),
+                "effects": Repeated(self.effect_space, max_len=MAX_EFFECTS),
+            }
+        )
+
+        # Observation is a list of players.
+        self.observation_space = Repeated(self.player_space, max_len=MAX_PLAYERS)
+
+    def reset(self, *, seed=None, options=None):
+        return self.observation_space.sample(), {}
+
+    def step(self, action):
+        return self.observation_space.sample(), 1, True, False, {}

rllib/examples/envs/classes/six_room_env.py

@@ -0,0 +1,314 @@
+import gymnasium as gym
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+# Map representation: Always six rooms (as the name suggests) with doors in between.
+MAPS = {
+    "small": [
+        "WWWWWWWWWWWWW",
+        "W   W   W   W",
+        "W       W   W",
+        "W   W       W",
+        "W WWWW WWWW W",
+        "W   W   W   W",
+        "W   W       W",
+        "W       W  GW",
+        "WWWWWWWWWWWWW",
+    ],
+    "medium": [
+        "WWWWWWWWWWWWWWWWWWW",
+        "W     W     W     W",
+        "W           W     W",
+        "W     W           W",
+        "W WWWWWWW WWWWWWW W",
+        "W     W     W     W",
+        "W     W           W",
+        "W           W    GW",
+        "WWWWWWWWWWWWWWWWWWW",
+    ],
+    "large": [
+        "WWWWWWWWWWWWWWWWWWWWWWWWW",
+        "W       W       W       W",
+        "W       W       W       W",
+        "W               W       W",
+        "W       W               W",
+        "W       W       W       W",
+        "WW WWWWWWWWW WWWWWWWWWW W",
+        "W       W       W       W",
+        "W       W               W",
+        "W       W       W       W",
+        "W               W       W",
+        "W       W       W      GW",
+        "WWWWWWWWWWWWWWWWWWWWWWWWW",
+    ],
+}
+
+
+class SixRoomEnv(gym.Env):
+    """A grid-world with six rooms (arranged as 2x3), which are connected by doors.
+
+    The agent starts in the upper left room and has to reach a designated goal state
+    in one of the rooms using primitive actions up, left, down, and right.
+
+    The agent receives a small penalty of -0.01 on each step and a reward of +10.0 when
+    reaching the goal state.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        # User can provide a custom map or a recognized map name (small, medium, large).
+        self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"]))
+        self.time_limit = config.get("time_limit", 50)
+
+        # Define observation space: Discrete, index fields.
+        self.observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0]))
+        # Primitive actions: up, down, left, right.
+        self.action_space = gym.spaces.Discrete(4)
+
+        # Initialize environment state.
+        self.reset()
+
+    def reset(self, *, seed=None, options=None):
+        self._agent_pos = (1, 1)
+        self._ts = 0
+        # Return high-level observation.
+        return self._agent_discrete_pos, {}
+
+    def step(self, action):
+        next_pos = _get_next_pos(action, self._agent_pos)
+
+        self._ts += 1
+
+        # Check if the move ends up in a wall. If so -> Ignore the move and stay
+        # where we are right now.
+        if self.map[next_pos[0]][next_pos[1]] != "W":
+            self._agent_pos = next_pos
+
+        # Check if the agent has reached the global goal state.
+        if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G":
+            return self._agent_discrete_pos, 10.0, True, False, {}
+
+        # Small step penalty.
+        return self._agent_discrete_pos, -0.01, False, self._ts >= self.time_limit, {}
+
+    @property
+    def _agent_discrete_pos(self):
+        x = self._agent_pos[0]
+        y = self._agent_pos[1]
+        # discrete position = row idx * columns + col idx
+        return x * len(self.map[0]) + y
+
+
+class HierarchicalSixRoomEnv(MultiAgentEnv):
+    def __init__(self, config=None):
+        super().__init__()
+
+        # User can provide a custom map or a recognized map name (small, medium, large).
+        self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"]))
+        self.max_steps_low_level = config.get("max_steps_low_level", 15)
+        self.time_limit = config.get("time_limit", 50)
+        self.num_low_level_agents = config.get("num_low_level_agents", 3)
+
+        self.agents = self.possible_agents = ["high_level_agent"] + [
+            f"low_level_agent_{i}" for i in range(self.num_low_level_agents)
+        ]
+
+        # Define basic observation space: Discrete, index fields.
+        observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0]))
+        # Low level agents always see where they are right now and what the target
+        # state should be.
+        low_level_observation_space = gym.spaces.Tuple(
+            (observation_space, observation_space)
+        )
+        # Primitive actions: up, down, left, right.
+        low_level_action_space = gym.spaces.Discrete(4)
+
+        self.observation_spaces = {"high_level_agent": observation_space}
+        self.observation_spaces.update(
+            {
+                f"low_level_agent_{i}": low_level_observation_space
+                for i in range(self.num_low_level_agents)
+            }
+        )
+        self.action_spaces = {
+            "high_level_agent": gym.spaces.Tuple(
+                (
+                    # The new target observation.
+                    observation_space,
+                    # Low-level policy that should get us to the new target observation.
+                    gym.spaces.Discrete(self.num_low_level_agents),
+                )
+            )
+        }
+        self.action_spaces.update(
+            {
+                f"low_level_agent_{i}": low_level_action_space
+                for i in range(self.num_low_level_agents)
+            }
+        )
+
+        # Initialize environment state.
+        self.reset()
+
+    def reset(self, *, seed=None, options=None):
+        self._agent_pos = (1, 1)
+        self._low_level_steps = 0
+        self._high_level_action = None
+        # Number of times the low-level agent reached the given target (by the high
+        # level agent).
+        self._num_targets_reached = 0
+
+        self._ts = 0
+
+        # Return high-level observation.
+        return {
+            "high_level_agent": self._agent_discrete_pos,
+        }, {}
+
+    def step(self, action_dict):
+        self._ts += 1
+
+        terminateds = {"__all__": self._ts >= self.time_limit}
+        truncateds = {"__all__": False}
+
+        # High-level agent acted: Set next goal and next low-level policy to use.
+        # Note that the agent does not move in this case and stays at its current
+        # location.
+        if "high_level_agent" in action_dict:
+            self._high_level_action = action_dict["high_level_agent"]
+            low_level_agent = f"low_level_agent_{self._high_level_action[1]}"
+            self._low_level_steps = 0
+            # Return next low-level observation for the now-active agent.
+            # We want this agent to act next.
+            return (
+                {
+                    low_level_agent: (
+                        self._agent_discrete_pos,  # current
+                        self._high_level_action[0],  # target
+                    )
+                },
+                # Penalty for a target state that's close to the current state.
+                {
+                    "high_level_agent": (
+                        self.eucl_dist(
+                            self._agent_discrete_pos,
+                            self._high_level_action[0],
+                            self.map,
+                        )
+                        / (len(self.map) ** 2 + len(self.map[0]) ** 2) ** 0.5
+                    )
+                    - 1.0,
+                },
+                terminateds,
+                truncateds,
+                {},
+            )
+        # Low-level agent made a move (primitive action).
+        else:
+            assert len(action_dict) == 1
+
+            # Increment low-level step counter.
+            self._low_level_steps += 1
+
+            target_discrete_pos, low_level_agent = self._high_level_action
+            low_level_agent = f"low_level_agent_{low_level_agent}"
+            next_pos = _get_next_pos(action_dict[low_level_agent], self._agent_pos)
+
+            # Check if the move ends up in a wall. If so -> Ignore the move and stay
+            # where we are right now.
+            if self.map[next_pos[0]][next_pos[1]] != "W":
+                self._agent_pos = next_pos
+
+            # Check if the agent has reached the global goal state.
+            if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G":
+                rewards = {
+                    "high_level_agent": 10.0,
+                    # +1.0 if the goal position was also the target position for the
+                    # low level agent.
+                    low_level_agent: float(
+                        self._agent_discrete_pos == target_discrete_pos
+                    ),
+                }
+                terminateds["__all__"] = True
+                return (
+                    {"high_level_agent": self._agent_discrete_pos},
+                    rewards,
+                    terminateds,
+                    truncateds,
+                    {},
+                )
+
+            # Low-level agent has reached its target location (given by the high-level):
+            # - Hand back control to high-level agent.
+            # - Reward low level agent and high-level agent with small rewards.
+            elif self._agent_discrete_pos == target_discrete_pos:
+                self._num_targets_reached += 1
+                rewards = {
+                    "high_level_agent": 1.0,
+                    low_level_agent: 1.0,
+                }
+                return (
+                    {"high_level_agent": self._agent_discrete_pos},
+                    rewards,
+                    terminateds,
+                    truncateds,
+                    {},
+                )
+
+            # Low-level agent has not reached anything.
+            else:
+                # Small step penalty for low-level agent.
+                rewards = {low_level_agent: -0.01}
+                # Reached time budget -> Hand back control to high level agent.
+                if self._low_level_steps >= self.max_steps_low_level:
+                    rewards["high_level_agent"] = -0.01
+                    return (
+                        {"high_level_agent": self._agent_discrete_pos},
+                        rewards,
+                        terminateds,
+                        truncateds,
+                        {},
+                    )
+                else:
+                    return (
+                        {
+                            low_level_agent: (
+                                self._agent_discrete_pos,  # current
+                                target_discrete_pos,  # target
+                            ),
+                        },
+                        rewards,
+                        terminateds,
+                        truncateds,
+                        {},
+                    )
+
+    @property
+    def _agent_discrete_pos(self):
+        x = self._agent_pos[0]
+        y = self._agent_pos[1]
+        # discrete position = row idx * columns + col idx
+        return x * len(self.map[0]) + y
+
+    @staticmethod
+    def eucl_dist(pos1, pos2, map):
+        x1, y1 = pos1 % len(map[0]), pos1 // len(map)
+        x2, y2 = pos2 % len(map[0]), pos2 // len(map)
+        return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5
+
+
+def _get_next_pos(action, pos):
+    x, y = pos
+    # Up.
+    if action == 0:
+        return x - 1, y
+    # Down.
+    elif action == 1:
+        return x + 1, y
+    # Left.
+    elif action == 2:
+        return x, y - 1
+    # Right.
+    else:
+        return x, y + 1

rllib/examples/envs/classes/stateless_cartpole.py

@@ -0,0 +1,38 @@
+import numpy as np
+from gymnasium.envs.classic_control import CartPoleEnv
+from gymnasium.spaces import Box
+
+
+class StatelessCartPole(CartPoleEnv):
+    """Partially observable variant of the CartPole gym environment.
+
+    https://github.com/openai/gym/blob/master/gym/envs/classic_control/
+    cartpole.py
+
+    We delete the x- and angular velocity components of the state, so that it
+    can only be solved by a memory enhanced model (policy).
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        # Fix our observation-space (remove 2 velocity components).
+        high = np.array(
+            [
+                self.x_threshold * 2,
+                self.theta_threshold_radians * 2,
+            ],
+            dtype=np.float32,
+        )
+
+        self.observation_space = Box(low=-high, high=high, dtype=np.float32)
+
+    def step(self, action):
+        next_obs, reward, done, truncated, info = super().step(action)
+        # next_obs is [x-pos, x-veloc, angle, angle-veloc]
+        return np.array([next_obs[0], next_obs[2]]), reward, done, truncated, info
+
+    def reset(self, *, seed=None, options=None):
+        init_obs, init_info = super().reset(seed=seed, options=options)
+        # init_obs is [x-pos, x-veloc, angle, angle-veloc]
+        return np.array([init_obs[0], init_obs[2]]), init_info

rllib/examples/envs/classes/stateless_pendulum.py

@@ -0,0 +1,34 @@
+import numpy as np
+from gymnasium.envs.classic_control import PendulumEnv
+from gymnasium.spaces import Box
+
+
+class StatelessPendulum(PendulumEnv):
+    """Partially observable variant of the Pendulum gym environment.
+
+    https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/
+    classic_control/pendulum.py
+
+    We delete the angular velocity component of the state, so that it
+    can only be solved by a memory enhanced model (policy).
+    """
+
+    def __init__(self, config=None):
+        config = config or {}
+        g = config.get("g", 10.0)
+
+        super().__init__(g=g)
+
+        # Fix our observation-space (remove angular velocity component).
+        high = np.array([1.0, 1.0], dtype=np.float32)
+        self.observation_space = Box(low=-high, high=high, dtype=np.float32)
+
+    def step(self, action):
+        next_obs, reward, done, truncated, info = super().step(action)
+        # next_obs is [cos(theta), sin(theta), theta-dot (angular velocity)]
+        return next_obs[:-1], reward, done, truncated, info
+
+    def reset(self, *, seed=None, options=None):
+        init_obs, init_info = super().reset(seed=seed, options=options)
+        # init_obs is [cos(theta), sin(theta), theta-dot (angular velocity)]
+        return init_obs[:-1], init_info

rllib/examples/envs/classes/ten_step_error_env.py

@@ -0,0 +1,49 @@
+import logging
+
+import gymnasium as gym
+
+logger = logging.getLogger(__name__)
+
+
+class TenStepErrorEnv(gym.Env):
+    """An environment that lets you sample 1 episode and raises an error during the next one.
+
+    The expectation to the env runner is that it will sample one episode and recreate the env
+    to sample the second one.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.step_count = 0
+        self.last_eps_errored = False
+        self.observation_space = gym.spaces.Box(low=0, high=1, shape=(1,))
+        self.action_space = gym.spaces.Box(low=0, high=1, shape=(1,))
+
+    def reset(self, seed=None, options=None):
+        self.step_count = 0
+        return self.observation_space.sample(), {
+            "last_eps_errored": self.last_eps_errored
+        }
+
+    def step(self, action):
+        self.step_count += 1
+        if self.step_count == 10:
+            if not self.last_eps_errored:
+                self.last_eps_errored = True
+                return (
+                    self.observation_space.sample(),
+                    0.0,
+                    True,
+                    False,
+                    {"last_eps_errored": False},
+                )
+            else:
+                raise Exception("Test error")
+
+        return (
+            self.observation_space.sample(),
+            0.0,
+            False,
+            False,
+            {"last_eps_errored": self.last_eps_errored},
+        )

rllib/examples/envs/classes/transformed_action_space_env.py

@@ -0,0 +1,62 @@
+from typing import Type
+
+import gymnasium as gym
+
+
+class ActionTransform(gym.ActionWrapper):
+    def __init__(self, env, low, high):
+        super().__init__(env)
+        self._low = low
+        self._high = high
+        self.action_space = type(env.action_space)(
+            self._low, self._high, env.action_space.shape, env.action_space.dtype
+        )
+
+    def action(self, action):
+        return (action - self._low) / (self._high - self._low) * (
+            self.env.action_space.high - self.env.action_space.low
+        ) + self.env.action_space.low
+
+
+def transform_action_space(env_name_or_creator) -> Type[gym.Env]:
+    """Wrapper for gym.Envs to have their action space transformed.
+
+    Args:
+        env_name_or_creator (Union[str, Callable[]]: String specifier or
+            env_maker function.
+
+    Returns:
+        New transformed_action_space_env function that returns an environment
+        wrapped by the ActionTransform wrapper. The constructor takes a
+        config dict with `_low` and `_high` keys specifying the new action
+        range (default -1.0 to 1.0). The reset of the config dict will be
+        passed on to the underlying/wrapped env's constructor.
+
+    .. testcode::
+        :skipif: True
+
+        # By gym string:
+        pendulum_300_to_500_cls = transform_action_space("Pendulum-v1")
+        # Create a transformed pendulum env.
+        pendulum_300_to_500 = pendulum_300_to_500_cls({"_low": -15.0})
+        pendulum_300_to_500.action_space
+
+    .. testoutput::
+
+        gym.spaces.Box(-15.0, 1.0, (1, ), "float32")
+    """
+
+    def transformed_action_space_env(config):
+        if isinstance(env_name_or_creator, str):
+            inner_env = gym.make(env_name_or_creator)
+        else:
+            inner_env = env_name_or_creator(config)
+        _low = config.pop("low", -1.0)
+        _high = config.pop("high", 1.0)
+        env = ActionTransform(inner_env, _low, _high)
+        return env
+
+    return transformed_action_space_env
+
+
+TransformedActionPendulum = transform_action_space("Pendulum-v1")

rllib/examples/envs/classes/utils/cartpole_observations.proto

@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+message CartPoleObservation {
+  double x_pos = 1;
+  double x_veloc = 2;
+  double angle_pos = 3;
+  double angle_veloc = 4;
+}

rllib/examples/envs/classes/utils/cartpole_observations_proto.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: cartpole_observations.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import (
+    descriptor as _descriptor,
+    descriptor_pool as _descriptor_pool,
+    symbol_database as _symbol_database,
+)
+from google.protobuf.internal import builder as _builder
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n\x1b\x63\x61rtpole_observations.proto"]\n\x13\x43\x61rtPoleObservation\x12\r\n\x05x_pos\x18\x01 \x01(\x01\x12\x0f\n\x07x_veloc\x18\x02 \x01(\x01\x12\x11\n\tangle_pos\x18\x03 \x01(\x01\x12\x13\n\x0b\x61ngle_veloc\x18\x04 \x01(\x01\x62\x06proto3'  # noqa
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "cartpole_observations_proto", _globals
+)
+if not _descriptor._USE_C_DESCRIPTORS:
+    DESCRIPTOR._loaded_options = None
+    _globals["_CARTPOLEOBSERVATION"]._serialized_start = 31
+    _globals["_CARTPOLEOBSERVATION"]._serialized_end = 124
+# @@protoc_insertion_point(module_scope)

rllib/examples/envs/classes/utils/dummy_external_client.py

@@ -0,0 +1,126 @@
+import pickle
+import socket
+import time
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.core import (
+    COMPONENT_RL_MODULE,
+    Columns,
+)
+from ray.rllib.env.external.rllink import (
+    RLlink,
+    get_rllink_message,
+    send_rllink_message,
+)
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import softmax
+
+torch, _ = try_import_torch()
+
+
+def _dummy_external_client(port: int = 5556):
+    """A dummy client that runs CartPole and acts as a testing external env."""
+
+    def _set_state(msg_body, rl_module):
+        rl_module.set_state(msg_body[COMPONENT_RL_MODULE])
+        # return msg_body[WEIGHTS_SEQ_NO]
+
+    # Connect to server.
+    while True:
+        try:
+            print(f"Trying to connect to localhost:{port} ...")
+            sock_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock_.connect(("localhost", port))
+            break
+        except ConnectionRefusedError:
+            time.sleep(5)
+
+    # Send ping-pong.
+    send_rllink_message(sock_, {"type": RLlink.PING.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.PONG
+
+    # Request config.
+    send_rllink_message(sock_, {"type": RLlink.GET_CONFIG.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.SET_CONFIG
+
+    config = pickle.loads(msg_body["config"])
+    # Create the RLModule.
+    rl_module = config.get_rl_module_spec().build()
+
+    # Request state/weights.
+    send_rllink_message(sock_, {"type": RLlink.GET_STATE.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.SET_STATE
+    _set_state(msg_body["state"], rl_module)
+
+    env_steps_per_sample = config.get_rollout_fragment_length()
+
+    # Start actual env loop.
+    env = gym.make("CartPole-v1")
+    obs, _ = env.reset()
+    episode = SingleAgentEpisode(observations=[obs])
+    episodes = [episode]
+
+    while True:
+        # Perform action inference using the RLModule.
+        logits = rl_module.forward_exploration(
+            batch={
+                Columns.OBS: torch.tensor(np.array([obs], np.float32)),
+            }
+        )[Columns.ACTION_DIST_INPUTS][
+            0
+        ].numpy()  # [0]=batch size 1
+
+        # Stochastic sample.
+        action_probs = softmax(logits)
+        action = int(np.random.choice(list(range(env.action_space.n)), p=action_probs))
+        logp = float(np.log(action_probs[action]))
+
+        # Perform the env step.
+        obs, reward, terminated, truncated, _ = env.step(action)
+
+        # Collect step data.
+        episode.add_env_step(
+            action=action,
+            reward=reward,
+            observation=obs,
+            terminated=terminated,
+            truncated=truncated,
+            extra_model_outputs={
+                Columns.ACTION_DIST_INPUTS: logits,
+                Columns.ACTION_LOGP: logp,
+            },
+        )
+
+        # We collected enough samples -> Send them to server.
+        if sum(map(len, episodes)) == env_steps_per_sample:
+            # Send the data to the server.
+            send_rllink_message(
+                sock_,
+                {
+                    "type": RLlink.EPISODES_AND_GET_STATE.name,
+                    "episodes": [e.get_state() for e in episodes],
+                    "timesteps": env_steps_per_sample,
+                },
+            )
+            # We are forced to sample on-policy. Have to wait for a response
+            # with the state (weights) in it.
+            msg_type, msg_body = get_rllink_message(sock_)
+            assert msg_type == RLlink.SET_STATE
+            _set_state(msg_body["state"], rl_module)
+
+            episodes = []
+            if not episode.is_done:
+                episode = episode.cut()
+                episodes.append(episode)
+
+        # If episode is done, reset env and create a new episode.
+        if episode.is_done:
+            obs, _ = env.reset()
+            episode = SingleAgentEpisode(observations=[obs])
+            episodes.append(episode)

rllib/examples/envs/classes/windy_maze_env.py

@@ -0,0 +1,160 @@
+import logging
+import random
+
+import gymnasium as gym
+from gymnasium.spaces import Box, Discrete, Tuple
+
+from ray.rllib.env import MultiAgentEnv
+
+logger = logging.getLogger(__name__)
+
+# Agent has to traverse the maze from the starting position S -> F
+# Observation space [x_pos, y_pos, wind_direction]
+# Action space: stay still OR move in current wind direction
+MAP_DATA = """
+#########
+#S      #
+####### #
+      # #
+      # #
+####### #
+#F      #
+#########"""
+
+
+class WindyMazeEnv(gym.Env):
+    def __init__(self, env_config):
+        self.map = [m for m in MAP_DATA.split("\n") if m]
+        self.x_dim = len(self.map)
+        self.y_dim = len(self.map[0])
+        logger.info("Loaded map {} {}".format(self.x_dim, self.y_dim))
+        for x in range(self.x_dim):
+            for y in range(self.y_dim):
+                if self.map[x][y] == "S":
+                    self.start_pos = (x, y)
+                elif self.map[x][y] == "F":
+                    self.end_pos = (x, y)
+        logger.info("Start pos {} end pos {}".format(self.start_pos, self.end_pos))
+        self.observation_space = Tuple(
+            [
+                Box(0, 100, shape=(2,)),  # (x, y)
+                Discrete(4),  # wind direction (N, E, S, W)
+            ]
+        )
+        self.action_space = Discrete(2)  # whether to move or not
+
+    def reset(self, *, seed=None, options=None):
+        self.wind_direction = random.choice([0, 1, 2, 3])
+        self.pos = self.start_pos
+        self.num_steps = 0
+        return [[self.pos[0], self.pos[1]], self.wind_direction], {}
+
+    def step(self, action):
+        if action == 1:
+            self.pos = self._get_new_pos(self.pos, self.wind_direction)
+        self.num_steps += 1
+        self.wind_direction = random.choice([0, 1, 2, 3])
+        at_goal = self.pos == self.end_pos
+        truncated = self.num_steps >= 200
+        done = at_goal or truncated
+        return (
+            [[self.pos[0], self.pos[1]], self.wind_direction],
+            100 * int(at_goal),
+            done,
+            truncated,
+            {},
+        )
+
+    def _get_new_pos(self, pos, direction):
+        if direction == 0:
+            new_pos = (pos[0] - 1, pos[1])
+        elif direction == 1:
+            new_pos = (pos[0], pos[1] + 1)
+        elif direction == 2:
+            new_pos = (pos[0] + 1, pos[1])
+        elif direction == 3:
+            new_pos = (pos[0], pos[1] - 1)
+        if (
+            new_pos[0] >= 0
+            and new_pos[0] < self.x_dim
+            and new_pos[1] >= 0
+            and new_pos[1] < self.y_dim
+            and self.map[new_pos[0]][new_pos[1]] != "#"
+        ):
+            return new_pos
+        else:
+            return pos  # did not move
+
+
+class HierarchicalWindyMazeEnv(MultiAgentEnv):
+    def __init__(self, env_config):
+        super().__init__()
+        self.flat_env = WindyMazeEnv(env_config)
+
+    def reset(self, *, seed=None, options=None):
+        self.cur_obs, infos = self.flat_env.reset()
+        self.current_goal = None
+        self.steps_remaining_at_level = None
+        self.num_high_level_steps = 0
+        # current low level agent id. This must be unique for each high level
+        # step since agent ids cannot be reused.
+        self.low_level_agent_id = "low_level_{}".format(self.num_high_level_steps)
+        return {
+            "high_level_agent": self.cur_obs,
+        }, {"high_level_agent": infos}
+
+    def step(self, action_dict):
+        assert len(action_dict) == 1, action_dict
+        if "high_level_agent" in action_dict:
+            return self._high_level_step(action_dict["high_level_agent"])
+        else:
+            return self._low_level_step(list(action_dict.values())[0])
+
+    def _high_level_step(self, action):
+        logger.debug("High level agent sets goal")
+        self.current_goal = action
+        self.steps_remaining_at_level = 25
+        self.num_high_level_steps += 1
+        self.low_level_agent_id = "low_level_{}".format(self.num_high_level_steps)
+        obs = {self.low_level_agent_id: [self.cur_obs, self.current_goal]}
+        rew = {self.low_level_agent_id: 0}
+        done = truncated = {"__all__": False}
+        return obs, rew, done, truncated, {}
+
+    def _low_level_step(self, action):
+        logger.debug("Low level agent step {}".format(action))
+        self.steps_remaining_at_level -= 1
+        cur_pos = tuple(self.cur_obs[0])
+        goal_pos = self.flat_env._get_new_pos(cur_pos, self.current_goal)
+
+        # Step in the actual env
+        f_obs, f_rew, f_terminated, f_truncated, info = self.flat_env.step(action)
+        new_pos = tuple(f_obs[0])
+        self.cur_obs = f_obs
+
+        # Calculate low-level agent observation and reward
+        obs = {self.low_level_agent_id: [f_obs, self.current_goal]}
+        if new_pos != cur_pos:
+            if new_pos == goal_pos:
+                rew = {self.low_level_agent_id: 1}
+            else:
+                rew = {self.low_level_agent_id: -1}
+        else:
+            rew = {self.low_level_agent_id: 0}
+
+        # Handle env termination & transitions back to higher level.
+        terminated = {"__all__": False}
+        truncated = {"__all__": False}
+        if f_terminated or f_truncated:
+            terminated["__all__"] = f_terminated
+            truncated["__all__"] = f_truncated
+            logger.debug("high level final reward {}".format(f_rew))
+            rew["high_level_agent"] = f_rew
+            obs["high_level_agent"] = f_obs
+        elif self.steps_remaining_at_level == 0:
+            terminated[self.low_level_agent_id] = True
+            truncated[self.low_level_agent_id] = False
+            rew["high_level_agent"] = 0
+            obs["high_level_agent"] = f_obs
+
+        return obs, rew, terminated, truncated, {self.low_level_agent_id: info}