feat: multi-ROM support for vectorized environments

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:

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<fs::path> expand_rom_paths(const std::vector<fs::path>& rom_paths, int num_envs) {
+    if (num_envs <= 0) return rom_paths;
+    std::vector<fs::path> 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<typename T>
 nb::ndarray<nb::numpy, T> make_numpy_array(T* data, std::vector<std::size_t> shape) {
@@ -126,7 +144,7 @@ nb::tuple wrap_step_result(EnvVectorizer& vec, BatchResult&& result) {
 void init_vector_module(nb::module_& m) {
     nb::class_<EnvVectorizer>(m, "ALEVectorInterface")
         .def("__init__", [](EnvVectorizer* t,
-                const fs::path& rom_path,
+                const std::vector<fs::path>& 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)
 

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,

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<fs::path>& 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<int>(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<PreprocessedEnv>(
-            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<ActionQueue>(num_envs_ * 2);

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<fs::path>& 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<ActionVect>& 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<std::unique_ptr<PreprocessedEnv>> envs_;
-    ActionVect action_set_;
+    std::vector<ActionVect> action_sets_;
 
     // Worker threads
     std::vector<std::thread> workers_;