From bf74b13fb20ab8ac3578de8fbce58d8af6794004 Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Sat, 30 May 2026 23:05:48 -0600 Subject: [PATCH 1/6] feat: multi-ROM support for vectorized environments - EnvVectorizer constructor takes vector rom_paths directly; action_sets_ stores per-environment action sets via action_sets() - ALEVectorInterface binding takes rom_paths + num_envs multiplier via expand_rom_paths helper; parse_autoreset_mode accepts canonical PascalCase - AtariVectorEnv accepts games: list[str]|str with optional num_envs multiplier; action space is always MultiDiscrete over all environments --- .../python/ale_vector_python_interface.cpp | 40 +++++---- src/ale/python/vector_env.py | 83 +++++++++++-------- src/ale/vector/env_vectorizer.cpp | 15 ++-- src/ale/vector/env_vectorizer.hpp | 7 +- tests/python/test_atari_vector_env.py | 2 +- 5 files changed, 87 insertions(+), 60 deletions(-) diff --git a/src/ale/python/ale_vector_python_interface.cpp b/src/ale/python/ale_vector_python_interface.cpp index 18ee4beae..b9b246b81 100644 --- a/src/ale/python/ale_vector_python_interface.cpp +++ b/src/ale/python/ale_vector_python_interface.cpp @@ -18,6 +18,24 @@ using ale::vector::Action; namespace { +// Repeat each ROM num_envs times in order that it was sent. +// E.g., ["pong", "breakout"], num_envs=2 -> ["pong", "pong", "breakout", "breakout"] +std::vector expand_rom_paths(const std::vector& rom_paths, int num_envs) { + if (num_envs <= 0) return rom_paths; + std::vector paths; + paths.reserve(rom_paths.size() * num_envs); + for (const auto& p : rom_paths) + for (int i = 0; i < num_envs; ++i) + paths.push_back(p); + return paths; +} + +AutoresetMode parse_autoreset_mode(const std::string& s) { + if (s == "NextStep") return AutoresetMode::NextStep; + if (s == "SameStep") return AutoresetMode::SameStep; + throw std::invalid_argument("Invalid autoreset_mode: " + s); +} + /// Helper to create numpy array from raw pointer with capsule ownership template nb::ndarray make_numpy_array(T* data, std::vector shape) { @@ -126,7 +144,7 @@ nb::tuple wrap_step_result(EnvVectorizer& vec, BatchResult&& result) { void init_vector_module(nb::module_& m) { nb::class_(m, "ALEVectorInterface") .def("__init__", [](EnvVectorizer* t, - const fs::path& rom_path, + const std::vector& rom_paths, int num_envs, int frame_skip, int stack_num, @@ -147,25 +165,17 @@ void init_vector_module(nb::module_& m) { int thread_affinity_offset, const std::string& autoreset_mode_str ) { - AutoresetMode autoreset_mode; - if (autoreset_mode_str == "NextStep") { - autoreset_mode = AutoresetMode::NextStep; - } else if (autoreset_mode_str == "SameStep") { - autoreset_mode = AutoresetMode::SameStep; - } else { - throw std::invalid_argument("Invalid autoreset_mode: " + autoreset_mode_str); - } - new (t) EnvVectorizer( - rom_path, num_envs, batch_size, num_threads, thread_affinity_offset, - autoreset_mode, img_height, img_width, stack_num, grayscale, + expand_rom_paths(rom_paths, num_envs), batch_size, num_threads, thread_affinity_offset, + parse_autoreset_mode(autoreset_mode_str), + img_height, img_width, stack_num, grayscale, frame_skip, maxpool, noop_max, use_fire_reset, episodic_life, life_loss_info, reward_clipping, max_episode_steps, repeat_action_probability, full_action_space ); }, - nb::arg("rom_path"), - nb::arg("num_envs"), + nb::arg("rom_paths"), + nb::arg("num_envs") = 0, nb::arg("frame_skip") = 4, nb::arg("stack_num") = 4, nb::arg("img_height") = 84, @@ -223,7 +233,7 @@ void init_vector_module(nb::module_& m) { return wrap_step_result(self, std::move(result)); }) - .def("get_action_set", &EnvVectorizer::action_set) + .def("get_action_sets", &EnvVectorizer::action_sets) .def("get_num_envs", &EnvVectorizer::num_envs) diff --git a/src/ale/python/vector_env.py b/src/ale/python/vector_env.py index 91a8f85a0..1f4bbdc38 100644 --- a/src/ale/python/vector_env.py +++ b/src/ale/python/vector_env.py @@ -10,7 +10,7 @@ from ale_py import roms from ale_py.env import AtariEnv from gymnasium.core import ObsType -from gymnasium.spaces import Box, Discrete +from gymnasium.spaces import Box, Discrete, MultiDiscrete from gymnasium.vector import AutoresetMode, VectorEnv @@ -19,9 +19,10 @@ class AtariVectorEnv(VectorEnv): def __init__( self, - game: str, - num_envs: int, + games: list[str] | str | None = None, + num_envs: int | None = None, *, + game: str | None = None, batch_size: int = 0, num_threads: int = 0, thread_affinity_offset: int = -1, @@ -47,8 +48,9 @@ def __init__( """Constructor for vector environment. Args: - game: ROM name - num_envs: Number of environments + games: List of ROM names + game: Single ROM name (used by gymnasium registration) + num_envs: Repeat each ROM this many times batch_size: If to provide a batch of environments (in async mode) num_threads: The number of threads to use for parallel environments thread_affinity_offset: The CPU core offset for thread affinity (-1 means no affinity, default: -1) @@ -70,14 +72,20 @@ def __init__( reward_clipping: If to clip rewards between -1 and 1 use_fire_reset: If to take fire action on reset if available """ - rom_path = roms.get_rom_path(game) - assert ( - rom_path is not None - ), f'{game} is not a ROM name, it should be snake_case not camel-case, i.e., "ms_pacman" not "MsPacman"' + if game is not None: + games = [game] + if isinstance(games, str): + games = [games] + rom_paths = [] + for game in games: + rom_path = roms.get_rom_path(game) + assert ( + rom_path is not None + ), f'{game} is not a ROM name, it should be snake_case not camel-case, i.e., "ms_pacman" not "MsPacman"' + rom_paths.extend([rom_path] * (num_envs or 1)) self.ale = ale_py.ALEVectorInterface( - rom_path=rom_path, - num_envs=num_envs, + rom_paths=rom_paths, frame_skip=frameskip, stack_num=stack_num, img_height=img_height, @@ -102,17 +110,11 @@ def __init__( ), ) - self.continuous = continuous - self.continuous_action_threshold = continuous_action_threshold - self.grayscale = grayscale - self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) - for h in (-1, 0, 1): - for v in (-1, 0, 1): - for f in (0, 1): - action = AtariEnv.map_action_idx(h, v, bool(f)).value - self.map_action_idx[h + 1, v + 1, f] = action + self.num_envs = len(rom_paths) + self.batch_size = self.num_envs if batch_size == 0 else batch_size # Set up the observation space based on grayscale or RGB format + self.grayscale = grayscale obs_shape = (stack_num, img_height, img_width) if not grayscale: obs_shape += (3,) @@ -120,6 +122,23 @@ def __init__( shape=obs_shape, low=0, high=255, dtype=np.uint8 ) + self.autoreset_mode = AutoresetMode(autoreset_mode) + self.metadata["autoreset_mode"] = self.autoreset_mode + + self.observation_space = gymnasium.vector.utils.batch_space( + self.single_observation_space, self.batch_size + ) + + # Set up the action space to use continuous or discrete actions + self.continuous = continuous + self.continuous_action_threshold = continuous_action_threshold + self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) + for h in (-1, 0, 1): + for v in (-1, 0, 1): + for f in (0, 1): + action = AtariEnv.map_action_idx(h, v, bool(f)).value + self.map_action_idx[h + 1, v + 1, f] = action + if self.continuous: # Actions are radius, theta, and fire, where first two are the parameters of polar coordinates. self.single_action_space = Box( @@ -128,20 +147,18 @@ def __init__( dtype=np.float32, shape=(3,), ) + self.action_space = gymnasium.vector.utils.batch_space( + self.single_action_space, self.batch_size + ) else: - self.single_action_space = Discrete(len(self.ale.get_action_set())) - - self.batch_size = num_envs if batch_size == 0 else batch_size - self.num_envs = num_envs - self.autoreset_mode = AutoresetMode(autoreset_mode) - self.metadata["autoreset_mode"] = self.autoreset_mode - - self.observation_space = gymnasium.vector.utils.batch_space( - self.single_observation_space, self.batch_size - ) - self.action_space = gymnasium.vector.utils.batch_space( - self.single_action_space, self.batch_size - ) + self.single_action_space = ( + None if self.full_action_space is None else Discrete(16) + ) + # Note: we expose the action space for all games instead of filtering up to the batch size, + # the user will need the full list to determine the action size for each environment. + self.action_space = MultiDiscrete( + np.array([len(s) for s in self.ale.get_action_sets()], dtype=np.int64) + ) self.is_xla_registered = False diff --git a/src/ale/vector/env_vectorizer.cpp b/src/ale/vector/env_vectorizer.cpp index 35a705167..210c0be00 100644 --- a/src/ale/vector/env_vectorizer.cpp +++ b/src/ale/vector/env_vectorizer.cpp @@ -14,8 +14,7 @@ namespace ale::vector { EnvVectorizer::EnvVectorizer( - const fs::path& rom_path, - int num_envs, + const std::vector& rom_paths, int batch_size, int num_threads, int thread_affinity_offset, @@ -34,28 +33,30 @@ EnvVectorizer::EnvVectorizer( int max_episode_steps, float repeat_action_probability, bool full_action_space -) : num_envs_(num_envs), - batch_size_(batch_size > 0 ? batch_size : num_envs), +) : num_envs_(static_cast(rom_paths.size())), + batch_size_(batch_size > 0 ? batch_size : num_envs_), img_height_(img_height), img_width_(img_width), stack_num_(stack_num), grayscale_(grayscale), autoreset_mode_(autoreset_mode), - last_recv_env_ids_(batch_size_ > 0 ? batch_size_ : num_envs) + last_recv_env_ids_(batch_size_ > 0 ? batch_size_ : num_envs_) { // Create environments envs_.reserve(num_envs_); + action_sets_.reserve(num_envs_); + for (int i = 0; i < num_envs_; ++i) { envs_.push_back(std::make_unique( - i, rom_path, img_height, img_width, frame_skip, maxpool, + i, rom_paths[i], img_height, img_width, frame_skip, maxpool, grayscale, stack_num, noop_max, use_fire_reset, episodic_life, life_loss_info, reward_clipping, max_episode_steps, repeat_action_probability, full_action_space, -1 )); + action_sets_.push_back(envs_.back()->action_set()); } stacked_obs_size_ = envs_[0]->stacked_obs_size(); - action_set_ = envs_[0]->action_set(); // Create action queue (capacity = 2x num_envs for safety) action_queue_ = std::make_unique(num_envs_ * 2); diff --git a/src/ale/vector/env_vectorizer.hpp b/src/ale/vector/env_vectorizer.hpp index 0c6720547..1ccbe4a5c 100644 --- a/src/ale/vector/env_vectorizer.hpp +++ b/src/ale/vector/env_vectorizer.hpp @@ -22,8 +22,7 @@ namespace ale::vector { class EnvVectorizer { public: EnvVectorizer( - const fs::path& rom_path, - int num_envs, + const std::vector& rom_paths, int batch_size = 0, int num_threads = 0, int thread_affinity_offset = -1, @@ -70,7 +69,7 @@ class EnvVectorizer { int num_envs() const { return num_envs_; } int batch_size() const { return batch_size_; } std::size_t stacked_obs_size() const { return stacked_obs_size_; } - const ActionVect& action_set() const { return action_set_; } + const std::vector& action_sets() const { return action_sets_; } AutoresetMode autoreset_mode() const { return autoreset_mode_; } bool is_grayscale() const { return grayscale_; } @@ -98,7 +97,7 @@ class EnvVectorizer { // Environments std::vector> envs_; - ActionVect action_set_; + std::vector action_sets_; // Worker threads std::vector workers_; diff --git a/tests/python/test_atari_vector_env.py b/tests/python/test_atari_vector_env.py index c0e28cf4f..ba05dcd6e 100644 --- a/tests/python/test_atari_vector_env.py +++ b/tests/python/test_atari_vector_env.py @@ -333,7 +333,7 @@ def test_batch_size_async( assert sync_envs.single_action_space == async_envs.single_action_space assert sync_envs.single_observation_space == async_envs.single_observation_space - assert sync_envs.action_space != async_envs.action_space + assert sync_envs.action_space == async_envs.action_space assert sync_envs.observation_space != async_envs.observation_space sync_envs.action_space.seed(action_seed) From 947d2e40cefe4fc34babdea3344fac9f0eeb330b Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Mon, 1 Jun 2026 12:58:18 -0600 Subject: [PATCH 2/6] updated method signatures, fixed missing class parameter full_action_space --- src/ale/python/__init__.pyi | 21 ++++++++++----------- src/ale/python/vector_env.py | 1 + 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/ale/python/__init__.pyi b/src/ale/python/__init__.pyi index 9b425f457..1bc2d1d3c 100644 --- a/src/ale/python/__init__.pyi +++ b/src/ale/python/__init__.pyi @@ -1,7 +1,7 @@ from __future__ import annotations import os -from typing import Any, Dict, List, Optional, Tuple, TypeAlias, overload +from typing import Any, Dict, List, Optional, Tuple, Union, overload import gymnasium as gym import numpy as np @@ -180,7 +180,7 @@ class ALEInterface: class ALEVectorInterface: def __init__( self, - rom_path: os.PathLike, + rom_paths: List[os.PathLike], num_envs: int, frame_skip: int, stack_num: int, @@ -204,22 +204,21 @@ class ALEVectorInterface: def reset( self, reset_indices: np.ndarray, reset_seeds: np.ndarray ) -> Tuple[np.ndarray, Dict[str, Any]]: ... - def send(self, action_idx: np.ndarray, paddle_strength: np.ndarray): ... + def send(self, action_ids: np.ndarray, paddle_strengths: np.ndarray) -> None: ... def recv( self, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]: ... - def get_action_set(self) -> List[Action]: ... + def get_action_sets(self) -> List[List[Action]]: ... def get_num_envs(self) -> int: ... - def get_observation_shape(self) -> Tuple[int, int, int]: ... + def get_observation_shape( + self, + ) -> Union[Tuple[int, int, int], Tuple[int, int, int, int]]: ... def handle(self) -> np.ndarray: ... try: - from ale_py.env import AtariEnvStepMetadata - from ale_py.vector_env import AtariVectorEnv - - AtariEnv: TypeAlias = AtariEnv - AtariEnvStepMetadata: TypeAlias = AtariEnvStepMetadata - AtariVectorEnv: TypeAlias = AtariVectorEnv + from ale_py.env import AtariEnv as AtariEnv + from ale_py.env import AtariEnvStepMetadata as AtariEnvStepMetadata + from ale_py.vector_env import AtariVectorEnv as AtariVectorEnv __all__ += ["AtariEnv", "AtariEnvStepMetadata", "AtariVectorEnv"] except ImportError: diff --git a/src/ale/python/vector_env.py b/src/ale/python/vector_env.py index 1f4bbdc38..e077af62b 100644 --- a/src/ale/python/vector_env.py +++ b/src/ale/python/vector_env.py @@ -130,6 +130,7 @@ def __init__( ) # Set up the action space to use continuous or discrete actions + self.full_action_space = full_action_space self.continuous = continuous self.continuous_action_threshold = continuous_action_threshold self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) From 297c528a199d8694a8283349838e5e0f0b3e0fe3 Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Mon, 1 Jun 2026 14:22:45 -0600 Subject: [PATCH 3/6] refactor action and observation setup in AtariVectorEnv --- src/ale/python/vector_env.py | 72 ++++++++++++++++++------------------ 1 file changed, 37 insertions(+), 35 deletions(-) diff --git a/src/ale/python/vector_env.py b/src/ale/python/vector_env.py index e077af62b..a677e0c33 100644 --- a/src/ale/python/vector_env.py +++ b/src/ale/python/vector_env.py @@ -113,55 +113,57 @@ def __init__( self.num_envs = len(rom_paths) self.batch_size = self.num_envs if batch_size == 0 else batch_size - # Set up the observation space based on grayscale or RGB format - self.grayscale = grayscale - obs_shape = (stack_num, img_height, img_width) - if not grayscale: - obs_shape += (3,) - self.single_observation_space = Box( - shape=obs_shape, low=0, high=255, dtype=np.uint8 - ) - self.autoreset_mode = AutoresetMode(autoreset_mode) self.metadata["autoreset_mode"] = self.autoreset_mode - self.observation_space = gymnasium.vector.utils.batch_space( - self.single_observation_space, self.batch_size + # Set up the observation space + self.grayscale = grayscale + self.single_observation_space, self.observation_space = self._setup_obs( + stack_num, img_height, img_width ) - # Set up the action space to use continuous or discrete actions + # Set up the action space self.full_action_space = full_action_space self.continuous = continuous self.continuous_action_threshold = continuous_action_threshold + self.single_action_space, self.action_space = ( + self._setup_continuous_action() + if self.continuous + else self._setup_discrete_action() + ) + self.is_xla_registered = False + + def _setup_obs(self, stack_num: int, img_height: int, img_width: int) -> tuple: + obs_shape = (stack_num, img_height, img_width) + if not self.grayscale: + obs_shape += (3,) + single = Box(shape=obs_shape, low=0, high=255, dtype=np.uint8) + return single, gymnasium.vector.utils.batch_space(single, self.batch_size) + + def _setup_continuous_action(self) -> tuple: self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) for h in (-1, 0, 1): for v in (-1, 0, 1): for f in (0, 1): action = AtariEnv.map_action_idx(h, v, bool(f)).value self.map_action_idx[h + 1, v + 1, f] = action - - if self.continuous: - # Actions are radius, theta, and fire, where first two are the parameters of polar coordinates. - self.single_action_space = Box( - low=np.array([0.0, -np.pi, 0.0]).astype(np.float32), - high=np.array([1.0, np.pi, 1.0]).astype(np.float32), - dtype=np.float32, - shape=(3,), - ) - self.action_space = gymnasium.vector.utils.batch_space( - self.single_action_space, self.batch_size - ) - else: - self.single_action_space = ( - None if self.full_action_space is None else Discrete(16) - ) - # Note: we expose the action space for all games instead of filtering up to the batch size, - # the user will need the full list to determine the action size for each environment. - self.action_space = MultiDiscrete( - np.array([len(s) for s in self.ale.get_action_sets()], dtype=np.int64) - ) - - self.is_xla_registered = False + # Actions are radius, theta, and fire, where first two are the parameters of polar coordinates. + single = Box( + low=np.array([0.0, -np.pi, 0.0]).astype(np.float32), + high=np.array([1.0, np.pi, 1.0]).astype(np.float32), + dtype=np.float32, + shape=(3,), + ) + return single, gymnasium.vector.utils.batch_space(single, self.batch_size) + + def _setup_discrete_action(self) -> tuple: + single = None if self.full_action_space is None else Discrete(16) + # Note: we expose the action space for all games instead of filtering up to the batch size, + # the user will need the full list to determine the action size for each environment. + action_space = MultiDiscrete( + np.array([len(s) for s in self.ale.get_action_sets()], dtype=np.int64) + ) + return single, action_space def reset( self, From e67843ef4d0333cb359bb57d888ed9a9faa4da94 Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Mon, 1 Jun 2026 16:17:55 -0600 Subject: [PATCH 4/6] test multirom, reset, step, send, recv, and batch size --- tests/python/test_atari_vector_env.py | 103 +++++++++++++++++++++++++- 1 file changed, 102 insertions(+), 1 deletion(-) diff --git a/tests/python/test_atari_vector_env.py b/tests/python/test_atari_vector_env.py index ba05dcd6e..96ae1685c 100644 --- a/tests/python/test_atari_vector_env.py +++ b/tests/python/test_atari_vector_env.py @@ -98,7 +98,6 @@ def assert_gym_ale_rollout_equivalence( # "env_id", [env_id for env_id in gym.registry if "ALE/" in env_id] # ) class TestVectorEnv: - disable_vector_args = dict( noop_max=0, use_fire_reset=False, @@ -665,3 +664,105 @@ def test_same_step_autoreset_mode( gym_envs.close() ale_envs.close() + + +class TestMultiRomVectorEnv: + """Tests for multi-ROM support in AtariVectorEnv.""" + + def test_multi_rom_num_envs(self): + from ale_py.vector_env import AtariVectorEnv + from gymnasium.spaces import MultiDiscrete + + env = AtariVectorEnv(games=["pong", "breakout"], num_envs=4) + + # 2 ROMs each with 4 duplicates = 8 total environments + assert env.num_envs == 8 + assert env.batch_size == 8 + + # Pong has 6 actions, Breakout has 4 actions, the games are expanded in order of the list + pong_actions = 6 + breakout_actions = 4 + assert np.all(env.action_space.nvec[:4] == pong_actions) + assert np.all(env.action_space.nvec[4:] == breakout_actions) + + # Observations + assert env.observation_space.shape == (8, 4, 84, 84) + assert env.single_observation_space.shape == (4, 84, 84) + + # Action + assert isinstance(env.action_space, MultiDiscrete) + assert env.action_space.shape == (8,) + + def test_multi_rom_reset_and_step(self): + from ale_py.vector_env import AtariVectorEnv + from gymnasium.spaces import MultiDiscrete + + env = AtariVectorEnv(games=["pong", "breakout"], num_envs=2) + assert isinstance(env.action_space, MultiDiscrete) + + # reset + obs, info = env.reset(seed=0) + assert obs.shape == (4, 4, 84, 84) + assert obs.dtype == np.uint8 + assert all(len(v) == 4 for v in info.values()) + + # sampling actions is now a vector, with different action spaces + for _ in range(30): + actions = env.action_space.sample() + assert np.all((actions >= 0) & (actions < env.action_space.nvec)) + + # step multiple roms + obs, rewards, terminations, truncations, info = env.step(actions) + assert obs.shape == (4, 4, 84, 84) + assert rewards.shape == (4,) + assert terminations.shape == (4,) + assert truncations.shape == (4,) + assert all(len(v) == 4 for v in info.values()) + + # send and recv + actions = env.action_space.sample() + env.send(actions) + obs, rewards, terminations, truncations, info = env.recv() + assert obs.shape == (4, 4, 84, 84) + assert rewards.shape == (4,) + assert terminations.shape == (4,) + assert truncations.shape == (4,) + assert all(len(v) == 4 for v in info.values()) + + def test_backward_compatibility_single_game(self): + from ale_py.vector_env import AtariVectorEnv + + str_game = AtariVectorEnv(game="breakout", num_envs=4) + str_games = AtariVectorEnv(games="breakout", num_envs=4) + multi_rom = AtariVectorEnv(games=["breakout"], num_envs=4) + for env in (str_game, str_games, multi_rom): + assert env.num_envs == str_game.num_envs + assert env.action_space == str_game.action_space + assert env.observation_space == str_game.observation_space + + def test_multi_rom_async_batch_size(self): + from ale_py.vector_env import AtariVectorEnv + + env = AtariVectorEnv(games=["pong", "breakout"], num_envs=4, batch_size=4) + + assert env.num_envs == 8 + assert env.batch_size == 4 + + obs, info = env.reset(seed=0) + assert obs.shape == (4, 4, 84, 84) + + # action_space covers all 8 envs; mask to the batch that just returned + actions = env.action_space.sample() + assert len(actions) == 8 + + env_ids = info.pop("env_id") + masked_actions = actions[env_ids] + + # step: send to first batch, recv second batch + _, _, _, _, info = env.step(masked_actions) + + # send and recv separately for the second batch + env_ids = info.pop("env_id") + masked_actions = actions[env_ids] + env.send(masked_actions) + obs, rewards, terminations, truncations, info = env.recv() From cb313128962366648a45427bdd960cac7d5512fc Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Thu, 4 Jun 2026 20:58:06 -0600 Subject: [PATCH 5/6] single action space and tests for documenting behavior --- src/ale/python/vector_env.py | 10 ++++++---- tests/python/test_atari_vector_env.py | 28 +++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/ale/python/vector_env.py b/src/ale/python/vector_env.py index a677e0c33..95bb55d22 100644 --- a/src/ale/python/vector_env.py +++ b/src/ale/python/vector_env.py @@ -157,12 +157,14 @@ def _setup_continuous_action(self) -> tuple: return single, gymnasium.vector.utils.batch_space(single, self.batch_size) def _setup_discrete_action(self) -> tuple: - single = None if self.full_action_space is None else Discrete(16) # Note: we expose the action space for all games instead of filtering up to the batch size, # the user will need the full list to determine the action size for each environment. - action_space = MultiDiscrete( - np.array([len(s) for s in self.ale.get_action_sets()], dtype=np.int64) - ) + sizes = [len(s) for s in self.ale.get_action_sets()] + action_space = MultiDiscrete(np.array(sizes, dtype=np.int64)) + # When all envs share the same action-set size, single_action_space is that + # Discrete space (e.g. Discrete(4) for Breakout, Discrete(18) for the full set). + # For different ROMs there is no single canonical space, we use None. + single = Discrete(sizes[0]) if len(set(sizes)) == 1 else None return single, action_space def reset( diff --git a/tests/python/test_atari_vector_env.py b/tests/python/test_atari_vector_env.py index 96ae1685c..6fba32a34 100644 --- a/tests/python/test_atari_vector_env.py +++ b/tests/python/test_atari_vector_env.py @@ -4,6 +4,7 @@ import gymnasium as gym import numpy as np import pytest +from ale_py.vector_env import AtariVectorEnv from gymnasium.utils.env_checker import data_equivalence gym.register_envs(ale_py) @@ -665,10 +666,37 @@ def test_same_step_autoreset_mode( gym_envs.close() ale_envs.close() + @pytest.mark.parametrize("game,n_minimal", [("breakout", 4), ("pong", 6)]) + def test_single_action_space(game, n_minimal): + """single_action_space reflects the real action-set size, not a constant.""" + env = AtariVectorEnv(game=game, num_envs=3) + assert env.single_action_space == gym.spaces.Discrete(n_minimal) + assert np.all(env.action_space.nvec == n_minimal) + # For homogeneous envs the batched action_space is the single one repeated. + assert env.action_space == gym.vector.utils.batch_space( + env.single_action_space, env.num_envs + ) + env.close() + + # Full action space is the 18 PLAYER_A_* actions (not 16). + env = AtariVectorEnv(game=game, num_envs=3, full_action_space=True) + assert env.single_action_space == gym.spaces.Discrete(18) + assert np.all(env.action_space.nvec == 18) + env.close() + class TestMultiRomVectorEnv: """Tests for multi-ROM support in AtariVectorEnv.""" + def test_single_action_space_none_for_different_roms(self): + """Different ROMs (different action-set sizes) have no single space.""" + # num_envs repeats each ROM, so this is [pong, pong, breakout, breakout]. + env = AtariVectorEnv(games=["pong", "breakout"], num_envs=2) + assert env.single_action_space is None + # pong has 6 actions, breakout has 4, in ROM order. + assert np.all(env.action_space.nvec == [6, 6, 4, 4]) + env.close() + def test_multi_rom_num_envs(self): from ale_py.vector_env import AtariVectorEnv from gymnasium.spaces import MultiDiscrete From 57b266ee4c0898459f3cf8bef66014458e9225ba Mon Sep 17 00:00:00 2001 From: Adrian Orenstein Date: Fri, 10 Jul 2026 15:38:19 +0000 Subject: [PATCH 6/6] fixed tests --- src/ale/python/vector_env.py | 12 ++++++------ tests/python/test_atari_vector_env.py | 10 +++++----- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/ale/python/vector_env.py b/src/ale/python/vector_env.py index 95bb55d22..cffad2dfc 100644 --- a/src/ale/python/vector_env.py +++ b/src/ale/python/vector_env.py @@ -126,6 +126,12 @@ def __init__( self.full_action_space = full_action_space self.continuous = continuous self.continuous_action_threshold = continuous_action_threshold + self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) + for h in (-1, 0, 1): + for v in (-1, 0, 1): + for f in (0, 1): + action = AtariEnv.map_action_idx(h, v, bool(f)).value + self.map_action_idx[h + 1, v + 1, f] = action self.single_action_space, self.action_space = ( self._setup_continuous_action() if self.continuous @@ -141,12 +147,6 @@ def _setup_obs(self, stack_num: int, img_height: int, img_width: int) -> tuple: return single, gymnasium.vector.utils.batch_space(single, self.batch_size) def _setup_continuous_action(self) -> tuple: - self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32) - for h in (-1, 0, 1): - for v in (-1, 0, 1): - for f in (0, 1): - action = AtariEnv.map_action_idx(h, v, bool(f)).value - self.map_action_idx[h + 1, v + 1, f] = action # Actions are radius, theta, and fire, where first two are the parameters of polar coordinates. single = Box( low=np.array([0.0, -np.pi, 0.0]).astype(np.float32), diff --git a/tests/python/test_atari_vector_env.py b/tests/python/test_atari_vector_env.py index 6fba32a34..123c6945e 100644 --- a/tests/python/test_atari_vector_env.py +++ b/tests/python/test_atari_vector_env.py @@ -666,8 +666,12 @@ def test_same_step_autoreset_mode( gym_envs.close() ale_envs.close() + +class TestMultiRomVectorEnv: + """Tests for multi-ROM support in AtariVectorEnv.""" + @pytest.mark.parametrize("game,n_minimal", [("breakout", 4), ("pong", 6)]) - def test_single_action_space(game, n_minimal): + def test_single_action_space(self, game, n_minimal): """single_action_space reflects the real action-set size, not a constant.""" env = AtariVectorEnv(game=game, num_envs=3) assert env.single_action_space == gym.spaces.Discrete(n_minimal) @@ -684,10 +688,6 @@ def test_single_action_space(game, n_minimal): assert np.all(env.action_space.nvec == 18) env.close() - -class TestMultiRomVectorEnv: - """Tests for multi-ROM support in AtariVectorEnv.""" - def test_single_action_space_none_for_different_roms(self): """Different ROMs (different action-set sizes) have no single space.""" # num_envs repeats each ROM, so this is [pong, pong, breakout, breakout].