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/toy_text.rst b/docs/api/toy_text.rst
index 15dbe082..f454788a 100644
--- a/docs/api/toy_text.rst
+++ b/docs/api/toy_text.rst
@@ -32,3 +32,17 @@ The difference between ``FrozenLake-v1`` and ``FrozenLake8x8-v1`` is that the
 former has ``size`` 4, ``max_episode_steps`` 100 with 0.7 reward threshold,
 while the latter has ``size`` 8, ``max_episode_steps`` 200 with 0.85 reward
 threshold.
+
+
+Taxi-v3
+-------
+
+`gym Taxi-v3 source code
+<https://github.com/openai/gym/blob/master/gym/envs/toy_text/taxi.py>`_
+
+There are four designated locations in the grid world indicated by Red, Green,
+Yellow, and Blue. When the episode starts, the taxi starts off at a random
+square and the passenger is at a random location. The taxi drives to the
+passenger's location, picks up the passenger, drives to the passenger's
+destination (another one of the four specified locations), and then drops off
+the passenger. Once the passenger is dropped off, the episode ends.
diff --git a/envpool/toy_text/BUILD b/envpool/toy_text/BUILD
index d9fe35c7..a77d5917 100644
--- a/envpool/toy_text/BUILD
+++ b/envpool/toy_text/BUILD
@@ -19,6 +19,14 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "taxi",
+    hdrs = ["taxi.h"],
+    deps = [
+        "//envpool/core:async_envpool",
+    ],
+)
+
 pybind_extension(
     name = "toy_text_envpool",
     srcs = [
@@ -27,6 +35,7 @@ pybind_extension(
     deps = [
         ":catch",
         ":frozen_lake",
+        ":taxi",
         "//envpool/core:py_envpool",
     ],
 )
diff --git a/envpool/toy_text/__init__.py b/envpool/toy_text/__init__.py
index c53535aa..cd5cbcee 100644
--- a/envpool/toy_text/__init__.py
+++ b/envpool/toy_text/__init__.py
@@ -20,6 +20,8 @@
   _CatchEnvSpec,
   _FrozenLakeEnvPool,
   _FrozenLakeEnvSpec,
+  _TaxiEnvPool,
+  _TaxiEnvSpec,
 )
 
 CatchEnvSpec, CatchDMEnvPool, CatchGymEnvPool = py_env(
@@ -30,6 +32,8 @@
   _FrozenLakeEnvSpec, _FrozenLakeEnvPool
 )
 
+TaxiEnvSpec, TaxiDMEnvPool, TaxiGymEnvPool = py_env(_TaxiEnvSpec, _TaxiEnvPool)
+
 __all__ = [
   "CatchEnvSpec",
   "CatchDMEnvPool",
@@ -37,4 +41,7 @@
   "FrozenLakeEnvSpec",
   "FrozenLakeDMEnvPool",
   "FrozenLakeGymEnvPool",
+  "TaxiEnvSpec",
+  "TaxiDMEnvPool",
+  "TaxiGymEnvPool",
 ]
diff --git a/envpool/toy_text/registration.py b/envpool/toy_text/registration.py
index 40f056f4..7cc3624c 100644
--- a/envpool/toy_text/registration.py
+++ b/envpool/toy_text/registration.py
@@ -46,3 +46,13 @@
   max_episode_steps=200,
   reward_threshold=0.85,
 )
+
+register(
+  task_id="Taxi-v3",
+  import_path="envpool.toy_text",
+  spec_cls="TaxiEnvSpec",
+  dm_cls="TaxiDMEnvPool",
+  gym_cls="TaxiGymEnvPool",
+  max_episode_steps=200,
+  reward_threshold=8.0,
+)
diff --git a/envpool/toy_text/taxi.h b/envpool/toy_text/taxi.h
new file mode 100644
index 00000000..29e16d56
--- /dev/null
+++ b/envpool/toy_text/taxi.h
@@ -0,0 +1,136 @@
+/*
+ * 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.
+ */
+// https://github.com/openai/gym/blob/master/gym/envs/toy_text/taxi.py
+
+#ifndef ENVPOOL_TOY_TEXT_TAXI_H_
+#define ENVPOOL_TOY_TEXT_TAXI_H_
+
+#include <algorithm>
+#include <cmath>
+#include <random>
+#include <string>
+#include <vector>
+
+#include "envpool/core/async_envpool.h"
+#include "envpool/core/env.h"
+
+namespace toy_text {
+
+class TaxiEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.bind(200),
+                    "reward_threshold"_.bind(8.0f));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs"_.bind(Spec<int>({-1}, {0, 499})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<int>({-1}, {0, 5})));
+  }
+};
+
+typedef class EnvSpec<TaxiEnvFns> TaxiEnvSpec;
+
+class TaxiEnv : public Env<TaxiEnvSpec> {
+ protected:
+  int x_, y_, s_, t_, max_episode_steps_, elapsed_step_;
+  std::uniform_int_distribution<> dist_car_, dist_loc_;
+  bool done_;
+  std::vector<std::vector<int>> loc_;
+  std::vector<std::string> map_, loc_map_;
+
+ public:
+  TaxiEnv(const Spec& spec, int env_id)
+      : Env<TaxiEnvSpec>(spec, env_id),
+        max_episode_steps_(spec.config["max_episode_steps"_]),
+        dist_car_(0, 3),
+        dist_loc_(0, 4),
+        done_(true),
+        loc_({{0, 0}, {0, 4}, {4, 0}, {4, 3}}),
+        map_({"|:|::|", "|:|::|", "|::::|", "||:|:|", "||:|:|"}),
+        loc_map_({"0   1", "     ", "     ", "     ", "2  3 "}) {}
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    x_ = dist_loc_(gen_);
+    y_ = dist_loc_(gen_);
+    s_ = dist_car_(gen_);
+    t_ = dist_car_(gen_);
+    done_ = false;
+    elapsed_step_ = 0;
+    State state = Allocate();
+    WriteObs(state, 0.0f);
+  }
+
+  void Step(const Action& action) override {
+    done_ = (++elapsed_step_ >= max_episode_steps_);
+    int act = action["action"_];
+    float reward = -1.0f;
+    if (act == 0) {
+      if (x_ < 4) {
+        ++x_;
+      }
+    } else if (act == 1) {
+      if (x_ > 0) {
+        --x_;
+      }
+    } else if (act == 2) {
+      if (map_[x_][y_ + 1] == ':') {
+        ++y_;
+      }
+    } else if (act == 3) {
+      if (map_[x_][y_] == ':') {
+        --y_;
+      }
+    } else if (act == 4) {
+      // pick up
+      if (s_ < 4 && x_ == loc_[s_][0] && y_ == loc_[s_][1]) {
+        s_ = 4;
+      } else {
+        reward = -10.0f;
+      }
+    } else {
+      // drop off
+      if (s_ == 4 && x_ == loc_[t_][0] && y_ == loc_[t_][1]) {
+        s_ = t_;
+        done_ = true;
+        reward = 20.0f;
+      } else if (s_ == 4 && loc_map_[x_][y_] != ' ') {
+        s_ = loc_map_[x_][y_] - '0';
+      } else {
+        reward = -10.0f;
+      }
+    }
+    State state = Allocate();
+    WriteObs(state, reward);
+  }
+
+ private:
+  void WriteObs(State& state, float reward) {  // NOLINT
+    state["obs"_] = ((x_ * 5 + y_) * 5 + s_) * 4 + t_;
+    state["reward"_] = reward;
+  }
+};
+
+typedef AsyncEnvPool<TaxiEnv> TaxiEnvPool;
+
+}  // namespace toy_text
+
+#endif  // ENVPOOL_TOY_TEXT_TAXI_H_
diff --git a/envpool/toy_text/toy_text.cc b/envpool/toy_text/toy_text.cc
index 50d1b741..6f7b8e1d 100644
--- a/envpool/toy_text/toy_text.cc
+++ b/envpool/toy_text/toy_text.cc
@@ -15,6 +15,7 @@
 #include "envpool/core/py_envpool.h"
 #include "envpool/toy_text/catch.h"
 #include "envpool/toy_text/frozen_lake.h"
+#include "envpool/toy_text/taxi.h"
 
 typedef PyEnvSpec<toy_text::CatchEnvSpec> CatchEnvSpec;
 typedef PyEnvPool<toy_text::CatchEnvPool> CatchEnvPool;
@@ -22,7 +23,11 @@ typedef PyEnvPool<toy_text::CatchEnvPool> CatchEnvPool;
 typedef PyEnvSpec<toy_text::FrozenLakeEnvSpec> FrozenLakeEnvSpec;
 typedef PyEnvPool<toy_text::FrozenLakeEnvPool> FrozenLakeEnvPool;
 
+typedef PyEnvSpec<toy_text::TaxiEnvSpec> TaxiEnvSpec;
+typedef PyEnvPool<toy_text::TaxiEnvPool> TaxiEnvPool;
+
 PYBIND11_MODULE(toy_text_envpool, m) {
   REGISTER(m, CatchEnvSpec, CatchEnvPool)
   REGISTER(m, FrozenLakeEnvSpec, FrozenLakeEnvPool)
+  REGISTER(m, TaxiEnvSpec, TaxiEnvPool)
 }


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 98d8a4a1..67581831 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -48,7 +48,7 @@ def test_make_classic_and_toytext(self) -> None:
       "CartPole-v0", "CartPole-v1", "Pendulum-v0", "MountainCar-v0",
       "MountainCarContinuous-v0", "Acrobot-v1"
     ]
-    toytext = ["Catch-v0"]
+    toytext = ["Catch-v0", "FrozenLake-v1", "FrozenLake8x8-v1", "Taxi-v3"]
     for task_id in classic + toytext:
       envpool.make_spec(task_id)
       env_gym = envpool.make_gym(task_id)
diff --git a/envpool/toy_text/toy_text_test.py b/envpool/toy_text/toy_text_test.py
index c7713664..5f9386ec 100644
--- a/envpool/toy_text/toy_text_test.py
+++ b/envpool/toy_text/toy_text_test.py
@@ -25,6 +25,8 @@
   CatchGymEnvPool,
   FrozenLakeEnvSpec,
   FrozenLakeGymEnvPool,
+  TaxiEnvSpec,
+  TaxiGymEnvPool,
 )
 
 
@@ -128,6 +130,65 @@ def test_frozen_lake(self) -> None:
           elapsed_step = 0
           last_obs = env.reset()
 
+  def test_taxi(self) -> None:
+    spec = TaxiEnvSpec(TaxiEnvSpec.gen_config(num_envs=1))
+    env = TaxiGymEnvPool(spec)
+    assert isinstance(env.observation_space, gym.spaces.Discrete)
+    assert env.observation_space.n == 500
+    assert isinstance(env.action_space, gym.spaces.Discrete)
+    assert env.action_space.n == 6
+    ref = gym.make("Taxi-v3")
+    for _ in range(10):
+      # random agent
+      ref.reset()
+      ref.unwrapped.s = env.reset()[0]
+      done = [False]
+      while not done[0]:
+        act = np.random.randint(6, size=(1,))
+        obs, rew, done, info = env.step(act)
+        ref_obs, ref_rew, ref_done, ref_info = ref.step(int(act[0]))
+        assert ref_obs == obs[0] and ref_rew == rew[0] and ref_done == done[0]
+    locs = ref.unwrapped.locs
+    left_point = [(4, 4), (0, 3), (4, 2), (0, 1)]
+    right_point = [(0, 0), (4, 1), (0, 2), (4, 3)]
+    left = {0: 1, 1: -1, 2: 1, 3: -1, 4: 1}
+    right = {0: -1, 1: 1, 2: -1, 3: 1, 4: -1}
+    for _ in range(10):
+      # 10IQ agent
+      ref.reset()
+      obs = ref.unwrapped.s = env.reset()[0]
+      x, y, s, t = ref.unwrapped.decode(obs)
+      actions = []
+      for i, g in [[0, s], [1, t]]:
+        # to (4, 0)
+        while (x, y) != (4, 0):
+          if (x, y) in left_point:
+            y -= 1
+            actions.append(3)
+          else:
+            x += left[y]
+            actions.append(int(left[y] == -1))
+        if (x, y) == locs[g]:
+          actions.append(4 + i)
+        else:
+          # to (0, 4)
+          while (x, y) != (0, 4):
+            if (x, y) in right_point:
+              y += 1
+              actions.append(2)
+            else:
+              x += right[y]
+              actions.append(int(right[y] == -1))
+            if (x, y) == locs[g]:
+              actions.append(4 + i)
+              break
+      while len(actions):
+        a = actions.pop(0)
+        ref_obs, ref_rew, ref_done, ref_info = ref.step(a)
+        obs, rew, done, info = env.step(np.array([a], int))
+        assert ref_obs == obs[0] and ref_rew == rew[0] and ref_done == done[0]
+      assert ref_rew == 20 and ref_done
+
 
 if __name__ == "__main__":
   absltest.main()