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

Issue:
Add toy_text
move catch-v0 to toy_text


Discussion / hints:

Fix patch:
diff --git a/docs/api/classic.rst b/docs/api/classic.rst
index 216e0902..18978a4b 100644
--- a/docs/api/classic.rst
+++ b/docs/api/classic.rst
@@ -1,19 +1,8 @@
 Classic Control
 ===============
 
-Most of the environments in classic control borrow from `Gym
-<https://github.com/openai/gym/tree/master/gym/envs/classic_control>`_ and
-`bsuite <https://github.com/deepmind/bsuite/tree/master/bsuite/environments>`_.
-
-
-Catch-v0
---------
-
-`bsuite catch source code
-<https://github.com/deepmind/bsuite/blob/master/bsuite/environments/catch.py>`_
-
-The agent must move a paddle to intercept falling balls. Falling balls only
-move downwards on the column they are in.
+All of the environments in classic control borrow from `Gym
+<https://github.com/openai/gym/tree/master/gym/envs/classic_control>`_.
 
 
 CartPole-v0/1
diff --git a/docs/api/toy_text.rst b/docs/api/toy_text.rst
new file mode 100644
index 00000000..9789535d
--- /dev/null
+++ b/docs/api/toy_text.rst
@@ -0,0 +1,16 @@
+Toy Text
+========
+
+Most of the environments in classic control borrow from `Gym
+<https://github.com/openai/gym/tree/master/gym/envs/toy_text>`_ and
+`bsuite <https://github.com/deepmind/bsuite/tree/master/bsuite/environments>`_.
+
+
+Catch-v0
+--------
+
+`bsuite catch source code
+<https://github.com/deepmind/bsuite/blob/master/bsuite/environments/catch.py>`_
+
+The agent must move a paddle to intercept falling balls. Falling balls only
+move downwards on the column they are in.
diff --git a/docs/index.rst b/docs/index.rst
index b118384d..8c29af15 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -74,6 +74,7 @@ stable version through `envpool.readthedocs.io/en/stable/
 
    api/atari
    api/classic
+   api/toy_text
 
 
 Indices and tables
diff --git a/envpool/BUILD b/envpool/BUILD
index 8caaea8c..c8d733c6 100644
--- a/envpool/BUILD
+++ b/envpool/BUILD
@@ -18,6 +18,7 @@ py_library(
     deps = [
         "//envpool/atari:atari_registration",
         "//envpool/classic_control:classic_control_registration",
+        "//envpool/toy_text:toy_text_registration",
     ],
 )
 
@@ -30,6 +31,7 @@ py_library(
         "//envpool/atari",
         "//envpool/classic_control",
         "//envpool/python",
+        "//envpool/toy_text",
     ],
 )
 
diff --git a/envpool/classic_control/BUILD b/envpool/classic_control/BUILD
index 8268f356..dadaec69 100644
--- a/envpool/classic_control/BUILD
+++ b/envpool/classic_control/BUILD
@@ -3,14 +3,6 @@ load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
 
 package(default_visibility = ["//visibility:public"])
 
-cc_library(
-    name = "catch",
-    hdrs = ["catch.h"],
-    deps = [
-        "//envpool/core:async_envpool",
-    ],
-)
-
 cc_library(
     name = "cartpole",
     hdrs = ["cartpole.h"],
@@ -59,7 +51,6 @@ pybind_extension(
     deps = [
         ":acrobot",
         ":cartpole",
-        ":catch",
         ":mountain_car",
         ":mountain_car_continuous",
         ":pendulum",
diff --git a/envpool/classic_control/__init__.py b/envpool/classic_control/__init__.py
index 7a1b8bc4..faa8bbb9 100644
--- a/envpool/classic_control/__init__.py
+++ b/envpool/classic_control/__init__.py
@@ -20,8 +20,6 @@
   _AcrobotEnvSpec,
   _CartPoleEnvPool,
   _CartPoleEnvSpec,
-  _CatchEnvPool,
-  _CatchEnvSpec,
   _MountainCarContinuousEnvPool,
   _MountainCarContinuousEnvSpec,
   _MountainCarEnvPool,
@@ -30,10 +28,6 @@
   _PendulumEnvSpec,
 )
 
-CatchEnvSpec, CatchDMEnvPool, CatchGymEnvPool = py_env(
-  _CatchEnvSpec, _CatchEnvPool
-)
-
 CartPoleEnvSpec, CartPoleDMEnvPool, CartPoleGymEnvPool = py_env(
   _CartPoleEnvSpec, _CartPoleEnvPool
 )
@@ -55,9 +49,6 @@
 )
 
 __all__ = [
-  "CatchEnvSpec",
-  "CatchDMEnvPool",
-  "CatchGymEnvPool",
   "CartPoleEnvSpec",
   "CartPoleDMEnvPool",
   "CartPoleGymEnvPool",
diff --git a/envpool/classic_control/classic_control.cc b/envpool/classic_control/classic_control.cc
index 9e80d9f9..f5154ca3 100644
--- a/envpool/classic_control/classic_control.cc
+++ b/envpool/classic_control/classic_control.cc
@@ -14,15 +14,11 @@
 
 #include "envpool/classic_control/acrobot.h"
 #include "envpool/classic_control/cartpole.h"
-#include "envpool/classic_control/catch.h"
 #include "envpool/classic_control/mountain_car.h"
 #include "envpool/classic_control/mountain_car_continuous.h"
 #include "envpool/classic_control/pendulum.h"
 #include "envpool/core/py_envpool.h"
 
-typedef PyEnvSpec<classic_control::CatchEnvSpec> CatchEnvSpec;
-typedef PyEnvPool<classic_control::CatchEnvPool> CatchEnvPool;
-
 typedef PyEnvSpec<classic_control::CartPoleEnvSpec> CartPoleEnvSpec;
 typedef PyEnvPool<classic_control::CartPoleEnvPool> CartPoleEnvPool;
 
@@ -41,7 +37,6 @@ typedef PyEnvSpec<classic_control::AcrobotEnvSpec> AcrobotEnvSpec;
 typedef PyEnvPool<classic_control::AcrobotEnvPool> AcrobotEnvPool;
 
 PYBIND11_MODULE(classic_control_envpool, m) {
-  REGISTER(m, CatchEnvSpec, CatchEnvPool)
   REGISTER(m, CartPoleEnvSpec, CartPoleEnvPool)
   REGISTER(m, PendulumEnvSpec, PendulumEnvPool)
   REGISTER(m, MountainCarEnvSpec, MountainCarEnvPool)
diff --git a/envpool/classic_control/registration.py b/envpool/classic_control/registration.py
index 73211738..18af08e7 100644
--- a/envpool/classic_control/registration.py
+++ b/envpool/classic_control/registration.py
@@ -15,16 +15,6 @@
 
 from envpool.registration import register
 
-register(
-  task_id="Catch-v0",
-  import_path="envpool.classic_control",
-  spec_cls="CatchEnvSpec",
-  dm_cls="CatchDMEnvPool",
-  gym_cls="CatchGymEnvPool",
-  height=10,
-  width=5,
-)
-
 register(
   task_id="CartPole-v0",
   import_path="envpool.classic_control",
diff --git a/envpool/entry.py b/envpool/entry.py
index 1f32a485..8776ed3b 100644
--- a/envpool/entry.py
+++ b/envpool/entry.py
@@ -15,3 +15,4 @@
 
 import envpool.atari.registration  # noqa: F401
 import envpool.classic_control.registration  # noqa: F401
+import envpool.toy_text.registration  # noqa: F401
diff --git a/envpool/toy_text/BUILD b/envpool/toy_text/BUILD
new file mode 100644
index 00000000..2a2a5613
--- /dev/null
+++ b/envpool/toy_text/BUILD
@@ -0,0 +1,50 @@
+load("@pip_requirements//:requirements.bzl", "requirement")
+load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "catch",
+    hdrs = ["catch.h"],
+    deps = [
+        "//envpool/core:async_envpool",
+    ],
+)
+
+pybind_extension(
+    name = "toy_text_envpool",
+    srcs = [
+        "toy_text.cc",
+    ],
+    deps = [
+        ":catch",
+        "//envpool/core:py_envpool",
+    ],
+)
+
+py_library(
+    name = "toy_text",
+    srcs = ["__init__.py"],
+    data = [":toy_text_envpool.so"],
+    deps = ["//envpool/python:api"],
+)
+
+py_test(
+    name = "toy_text_test",
+    srcs = ["toy_text_test.py"],
+    deps = [
+        ":toy_text",
+        requirement("absl-py"),
+        requirement("dm_env"),
+        requirement("gym"),
+        requirement("numpy"),
+    ],
+)
+
+py_library(
+    name = "toy_text_registration",
+    srcs = ["registration.py"],
+    deps = [
+        "//envpool:registration",
+    ],
+)
diff --git a/envpool/toy_text/__init__.py b/envpool/toy_text/__init__.py
new file mode 100644
index 00000000..b89c904a
--- /dev/null
+++ b/envpool/toy_text/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2021 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Toy text env in EnvPool."""
+
+from envpool.python.api import py_env
+
+from .toy_text_envpool import _CatchEnvPool, _CatchEnvSpec
+
+CatchEnvSpec, CatchDMEnvPool, CatchGymEnvPool = py_env(
+  _CatchEnvSpec, _CatchEnvPool
+)
+
+__all__ = [
+  "CatchEnvSpec",
+  "CatchDMEnvPool",
+  "CatchGymEnvPool",
+]
diff --git a/envpool/classic_control/catch.h b/envpool/toy_text/catch.h
similarity index 93%
rename from envpool/classic_control/catch.h
rename to envpool/toy_text/catch.h
index 973c13ad..a55cc112 100644
--- a/envpool/classic_control/catch.h
+++ b/envpool/toy_text/catch.h
@@ -15,8 +15,8 @@
  */
 // https://github.com/deepmind/bsuite/blob/master/bsuite/environments/catch.py
 
-#ifndef ENVPOOL_CLASSIC_CONTROL_CATCH_H_
-#define ENVPOOL_CLASSIC_CONTROL_CATCH_H_
+#ifndef ENVPOOL_TOY_TEXT_CATCH_H_
+#define ENVPOOL_TOY_TEXT_CATCH_H_
 
 #include <cmath>
 #include <random>
@@ -24,7 +24,7 @@
 #include "envpool/core/async_envpool.h"
 #include "envpool/core/env.h"
 
-namespace classic_control {
+namespace toy_text {
 
 class CatchEnvFns {
  public:
@@ -93,6 +93,6 @@ class CatchEnv : public Env<CatchEnvSpec> {
 
 typedef AsyncEnvPool<CatchEnv> CatchEnvPool;
 
-}  // namespace classic_control
+}  // namespace toy_text
 
-#endif  // ENVPOOL_CLASSIC_CONTROL_CATCH_H_
+#endif  // ENVPOOL_TOY_TEXT_CATCH_H_
diff --git a/envpool/toy_text/registration.py b/envpool/toy_text/registration.py
new file mode 100644
index 00000000..2c600f28
--- /dev/null
+++ b/envpool/toy_text/registration.py
@@ -0,0 +1,26 @@
+# Copyright 2021 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Classic control env registration."""
+
+from envpool.registration import register
+
+register(
+  task_id="Catch-v0",
+  import_path="envpool.toy_text",
+  spec_cls="CatchEnvSpec",
+  dm_cls="CatchDMEnvPool",
+  gym_cls="CatchGymEnvPool",
+  height=10,
+  width=5,
+)
diff --git a/envpool/toy_text/toy_text.cc b/envpool/toy_text/toy_text.cc
new file mode 100644
index 00000000..8c605fbb
--- /dev/null
+++ b/envpool/toy_text/toy_text.cc
@@ -0,0 +1,21 @@
+// Copyright 2021 Garena Online Private Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "envpool/core/py_envpool.h"
+#include "envpool/toy_text/catch.h"
+
+typedef PyEnvSpec<toy_text::CatchEnvSpec> CatchEnvSpec;
+typedef PyEnvPool<toy_text::CatchEnvPool> CatchEnvPool;
+
+PYBIND11_MODULE(toy_text_envpool, m) { REGISTER(m, CatchEnvSpec, CatchEnvPool) }


Test patch:
diff --git a/envpool/classic_control/classic_control_test.py b/envpool/classic_control/classic_control_test.py
index 88d4e3a1..c223e494 100644
--- a/envpool/classic_control/classic_control_test.py
+++ b/envpool/classic_control/classic_control_test.py
@@ -18,16 +18,12 @@
 import gym
 import numpy as np
 from absl.testing import absltest
-from dm_env import TimeStep
 
 from envpool.classic_control import (
   AcrobotEnvSpec,
   AcrobotGymEnvPool,
   CartPoleEnvSpec,
   CartPoleGymEnvPool,
-  CatchDMEnvPool,
-  CatchEnvSpec,
-  CatchGymEnvPool,
   MountainCarContinuousEnvSpec,
   MountainCarContinuousGymEnvPool,
   MountainCarEnvSpec,
@@ -153,62 +149,5 @@ def test_acrobot(self) -> None:
     # self.run_align_check(env0, env1, reset_fn)
 
 
-class _CatchEnvTest(absltest.TestCase):
-
-  def test_catch(self) -> None:
-    num_envs = 3
-    row, col = 10, 5
-    config = CatchEnvSpec.gen_config(num_envs=num_envs)
-    spec = CatchEnvSpec(config)
-    for env_type in ["dm", "gym"]:
-      if env_type == "dm":
-        e = CatchDMEnvPool(spec)
-      else:
-        e = CatchGymEnvPool(spec)
-      # get a successful trajectory
-      if env_type == "dm":
-        obs = e.reset().observation.obs  # type: ignore
-      else:
-        obs = e.reset()
-      assert obs.shape == (num_envs, row, col)
-      ball_pos = np.where(obs[:, 0] == 1)[1]
-      paddle_pos = np.where(obs[:, -1] == 1)[1]
-      for t in range(row - 1):
-        action = np.sign(ball_pos - paddle_pos) + 1
-        if env_type == "dm":
-          ts: TimeStep = e.step(action, np.arange(num_envs))
-          obs, rew, done = ts.observation.obs, ts.reward, ts.last()
-        else:
-          obs, rew, done, _ = e.step(action, np.arange(num_envs))
-        assert obs.shape == (num_envs, row, col)
-        paddle_pos = np.where(obs[:, -1] == 1)[1]
-        if t != row - 2:
-          assert np.all(rew == 0) and np.all(~done)
-        else:
-          assert np.all(rew == 1) and np.all(done)
-      # get a failure trajectory
-      if env_type == "dm":
-        obs = e.reset().observation.obs  # type: ignore
-      else:
-        obs = e.reset()
-      assert obs.shape == (num_envs, row, col)
-      ball_pos = np.where(obs[:, 0] == 1)[1]
-      paddle_pos = np.where(obs[:, -1] == 1)[1]
-      for t in range(row - 1):
-        action = np.sign(paddle_pos - ball_pos) + 1
-        action[action == 1] = 0
-        if env_type == "dm":
-          ts = e.step(action, np.arange(num_envs))
-          obs, rew, done = ts.observation.obs, ts.reward, ts.last()
-        else:
-          obs, rew, done, _ = e.step(action, np.arange(num_envs))
-        assert obs.shape == (num_envs, row, col)
-        paddle_pos = np.where(obs[:, -1] == 1)[1]
-        if t != row - 2:
-          assert np.all(rew == 0) and np.all(~done)
-        else:
-          assert np.all(rew == -1) and np.all(done)
-
-
 if __name__ == "__main__":
   absltest.main()
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 31e1a3b9..c7679abb 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -40,11 +40,13 @@ def test_make_atari(self) -> None:
     self.assertEqual(env_gym.action_space.n, 18)
     self.assertEqual(env_dm.action_spec().num_values, 18)
 
-  def test_make_classic(self) -> None:
-    for task_id in [
+  def test_make_classic_and_toytext(self) -> None:
+    classic = [
       "CartPole-v0", "CartPole-v1", "Pendulum-v0", "MountainCar-v0",
-      "MountainCarContinuous-v0", "Acrobot-v1", "Catch-v0"
-    ]:
+      "MountainCarContinuous-v0", "Acrobot-v1"
+    ]
+    toytext = ["Catch-v0"]
+    for task_id in classic + toytext:
       envpool.make_spec(task_id)
       env_gym = envpool.make_gym(task_id)
       env_dm = envpool.make_dm(task_id)
diff --git a/envpool/toy_text/toy_text_test.py b/envpool/toy_text/toy_text_test.py
new file mode 100644
index 00000000..45f3060a
--- /dev/null
+++ b/envpool/toy_text/toy_text_test.py
@@ -0,0 +1,81 @@
+# Copyright 2021 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for classic control environments."""
+
+import numpy as np
+from absl.testing import absltest
+from dm_env import TimeStep
+
+from envpool.toy_text import CatchDMEnvPool, CatchEnvSpec, CatchGymEnvPool
+
+
+class _ToyTextEnvTest(absltest.TestCase):
+
+  def test_catch(self) -> None:
+    num_envs = 3
+    row, col = 10, 5
+    config = CatchEnvSpec.gen_config(num_envs=num_envs)
+    spec = CatchEnvSpec(config)
+    for env_type in ["dm", "gym"]:
+      if env_type == "dm":
+        e = CatchDMEnvPool(spec)
+      else:
+        e = CatchGymEnvPool(spec)
+      # get a successful trajectory
+      if env_type == "dm":
+        obs = e.reset().observation.obs  # type: ignore
+      else:
+        obs = e.reset()
+      assert obs.shape == (num_envs, row, col)
+      ball_pos = np.where(obs[:, 0] == 1)[1]
+      paddle_pos = np.where(obs[:, -1] == 1)[1]
+      for t in range(row - 1):
+        action = np.sign(ball_pos - paddle_pos) + 1
+        if env_type == "dm":
+          ts: TimeStep = e.step(action, np.arange(num_envs))
+          obs, rew, done = ts.observation.obs, ts.reward, ts.last()
+        else:
+          obs, rew, done, _ = e.step(action, np.arange(num_envs))
+        assert obs.shape == (num_envs, row, col)
+        paddle_pos = np.where(obs[:, -1] == 1)[1]
+        if t != row - 2:
+          assert np.all(rew == 0) and np.all(~done)
+        else:
+          assert np.all(rew == 1) and np.all(done)
+      # get a failure trajectory
+      if env_type == "dm":
+        obs = e.reset().observation.obs  # type: ignore
+      else:
+        obs = e.reset()
+      assert obs.shape == (num_envs, row, col)
+      ball_pos = np.where(obs[:, 0] == 1)[1]
+      paddle_pos = np.where(obs[:, -1] == 1)[1]
+      for t in range(row - 1):
+        action = np.sign(paddle_pos - ball_pos) + 1
+        action[action == 1] = 0
+        if env_type == "dm":
+          ts = e.step(action, np.arange(num_envs))
+          obs, rew, done = ts.observation.obs, ts.reward, ts.last()
+        else:
+          obs, rew, done, _ = e.step(action, np.arange(num_envs))
+        assert obs.shape == (num_envs, row, col)
+        paddle_pos = np.where(obs[:, -1] == 1)[1]
+        if t != row - 2:
+          assert np.all(rew == 0) and np.all(~done)
+        else:
+          assert np.all(rew == -1) and np.all(done)
+
+
+if __name__ == "__main__":
+  absltest.main()