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/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..cffad2dfc 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,48 +110,62 @@ def __init__( ), ) + self.num_envs = len(rom_paths) + self.batch_size = self.num_envs if batch_size == 0 else batch_size + + self.autoreset_mode = AutoresetMode(autoreset_mode) + self.metadata["autoreset_mode"] = self.autoreset_mode + + # 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 + self.full_action_space = full_action_space 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.single_action_space, self.action_space = ( + self._setup_continuous_action() + if self.continuous + else self._setup_discrete_action() + ) + self.is_xla_registered = False - # Set up the observation space based on grayscale or RGB format + def _setup_obs(self, stack_num: int, img_height: int, img_width: int) -> tuple: obs_shape = (stack_num, img_height, img_width) - if not grayscale: + if not self.grayscale: obs_shape += (3,) - self.single_observation_space = Box( - shape=obs_shape, low=0, high=255, dtype=np.uint8 - ) - - 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,), - ) - 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 + 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: + # 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,), ) - - self.is_xla_registered = False + return single, gymnasium.vector.utils.batch_space(single, self.batch_size) + + def _setup_discrete_action(self) -> tuple: + # 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. + 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( self, 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..123c6945e 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) @@ -98,7 +99,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, @@ -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) @@ -665,3 +665,132 @@ 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(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) + 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() + + 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 + + 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()