Repository: sail-sg/envpool (language: python)

Issue:
[Feature Request] Support for full action space in ALE
## Motivation

Currently, there is no way to use the full action space of ALE of 18 actions. (https://envpool.readthedocs.io/en/latest/env/atari.html)

## Solution

Maybe we can add an option such as `full_action_space=False` similar to how ale-py does it? (https://brosa.ca/blog/ale-release-v0.7)



Discussion / hints:

Fix patch:
diff --git a/docs/env/atari.rst b/docs/env/atari.rst
index 88754a77..f3b988ff 100644
--- a/docs/env/atari.rst
+++ b/docs/env/atari.rst
@@ -52,6 +52,8 @@ Options
   image resize, default to ``True``.
 * ``use_fire_reset (bool)``: whether to use ``fire-reset`` wrapper, default to
   ``True``.
+* ``full_action_space (bool)``: whether to use full action space of ALE of 18
+  actions, default to ``False``.
 
 
 Observation Space
diff --git a/envpool/__init__.py b/envpool/__init__.py
index 7466c5a6..a6fe9578 100644
--- a/envpool/__init__.py
+++ b/envpool/__init__.py
@@ -24,7 +24,7 @@
   register,
 )
 
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 __all__ = [
   "register",
   "make",
diff --git a/envpool/atari/atari_env.h b/envpool/atari/atari_env.h
index bc601fa5..48cb83ce 100644
--- a/envpool/atari/atari_env.h
+++ b/envpool/atari/atari_env.h
@@ -54,7 +54,7 @@ class AtariEnvFns {
         "zero_discount_on_life_loss"_.Bind(false), "episodic_life"_.Bind(false),
         "reward_clip"_.Bind(false), "use_fire_reset"_.Bind(true),
         "img_height"_.Bind(84), "img_width"_.Bind(84),
-        "task"_.Bind(std::string("pong")),
+        "task"_.Bind(std::string("pong")), "full_action_space"_.Bind(false),
         "repeat_action_probability"_.Bind(0.0f),
         "use_inter_area_resize"_.Bind(true), "gray_scale"_.Bind(true));
   }
@@ -72,7 +72,9 @@ class AtariEnvFns {
   static decltype(auto) ActionSpec(const Config& conf) {
     ale::ALEInterface env;
     env.loadROM(GetRomPath(conf["base_path"_], conf["task"_]));
-    int action_size = env.getMinimalActionSet().size();
+    int action_size = conf["full_action_space"_]
+                          ? env.getLegalActionSet().size()
+                          : env.getMinimalActionSet().size();
     return MakeDict("action"_.Bind(Spec<int>({-1}, {0, action_size - 1})));
   }
 };
@@ -124,7 +126,11 @@ class AtariEnv : public Env<AtariEnvSpec> {
                    spec.config["repeat_action_probability"_]);
     env_->setInt("random_seed", seed_);
     env_->loadROM(rom_path_);
-    action_set_ = env_->getMinimalActionSet();
+    if (spec.config["full_action_space"_]) {
+      action_set_ = env_->getLegalActionSet();
+    } else {
+      action_set_ = env_->getMinimalActionSet();
+    }
     if (spec.config["use_fire_reset"_]) {
       // https://github.com/sail-sg/envpool/issues/221
       for (auto a : action_set_) {
diff --git a/setup.cfg b/setup.cfg
index 2e27e7c3..215129fa 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = envpool
-version = 0.7.0
+version = 0.7.1
 author = "EnvPool Contributors"
 author_email = "sail@sea.com"
 description = "C++-based high-performance parallel environment execution engine (vectorized env) for general RL environments."


Test patch:
diff --git a/envpool/atari/atari_envpool_test.py b/envpool/atari/atari_envpool_test.py
index 86e19d52..a687060e 100644
--- a/envpool/atari/atari_envpool_test.py
+++ b/envpool/atari/atari_envpool_test.py
@@ -61,6 +61,12 @@ def test_raw_envpool(self) -> None:
     fps = total * batch / duration * 4
     logging.info(f"Raw envpool FPS = {fps:.6f}")
 
+  def test_full_action_space(self) -> None:
+    env = make_gym("Pong-v5", full_action_space=True)
+    self.assertEqual(env.action_space.n, 18)
+    env = make_gym("Breakout-v5", full_action_space=True)
+    self.assertEqual(env.action_space.n, 18)
+
   def test_align(self) -> None:
     """Make sure gym's envpool and dm_env's envpool generate the same data."""
     num_envs = 4