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 feb3c2ae..25c5b2ee 100644
--- a/docs/api/toy_text.rst
+++ b/docs/api/toy_text.rst
@@ -64,3 +64,20 @@ The end of the chain, however, presents a large reward, and by moving
 At each action, there is a small probability that the agent 'slips' and the
 opposite transition is instead taken.
 The observed state is the current state in the chain (0 to n-1).
+
+
+CliffWalking-v0
+---------------
+
+`gym CliffWalking-v0 source code
+<https://github.com/openai/gym/blob/master/gym/envs/toy_text/cliffwalking.py>`_
+
+The board is a 4x12 matrix, with (using NumPy matrix indexing):
+
+- [3, 0] as the start at bottom-left
+- [3, 11] as the goal at bottom-right
+- [3, 1..10] as the cliff at bottom-center
+
+Each time step incurs -1 reward, and stepping into the cliff incurs -100
+reward and a reset to the start. An episode terminates when the agent reaches
+the goal.
diff --git a/envpool/toy_text/BUILD b/envpool/toy_text/BUILD
index 83b5b365..27c6e925 100644
--- a/envpool/toy_text/BUILD
+++ b/envpool/toy_text/BUILD
@@ -35,6 +35,14 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "cliffwalking",
+    hdrs = ["cliffwalking.h"],
+    deps = [
+        "//envpool/core:async_envpool",
+    ],
+)
+
 pybind_extension(
     name = "toy_text_envpool",
     srcs = [
@@ -42,6 +50,7 @@ pybind_extension(
     ],
     deps = [
         ":catch",
+        ":cliffwalking",
         ":frozen_lake",
         ":nchain",
         ":taxi",
diff --git a/envpool/toy_text/__init__.py b/envpool/toy_text/__init__.py
index 83b039a2..8347e599 100644
--- a/envpool/toy_text/__init__.py
+++ b/envpool/toy_text/__init__.py
@@ -18,6 +18,8 @@
 from .toy_text_envpool import (
   _CatchEnvPool,
   _CatchEnvSpec,
+  _CliffWalkingEnvPool,
+  _CliffWalkingEnvSpec,
   _FrozenLakeEnvPool,
   _FrozenLakeEnvSpec,
   _NChainEnvPool,
@@ -40,6 +42,10 @@
   _NChainEnvSpec, _NChainEnvPool
 )
 
+CliffWalkingEnvSpec, CliffWalkingDMEnvPool, CliffWalkingGymEnvPool = py_env(
+  _CliffWalkingEnvSpec, _CliffWalkingEnvPool
+)
+
 __all__ = [
   "CatchEnvSpec",
   "CatchDMEnvPool",
@@ -53,4 +59,7 @@
   "NChainEnvSpec",
   "NChainDMEnvPool",
   "NChainGymEnvPool",
+  "CliffWalkingEnvSpec",
+  "CliffWalkingDMEnvPool",
+  "CliffWalkingGymEnvPool",
 ]
diff --git a/envpool/toy_text/cliffwalking.h b/envpool/toy_text/cliffwalking.h
new file mode 100644
index 00000000..bc3a0b15
--- /dev/null
+++ b/envpool/toy_text/cliffwalking.h
@@ -0,0 +1,103 @@
+/*
+ * 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/cliffwalking.py
+
+#ifndef ENVPOOL_TOY_TEXT_CLIFFWALKING_H_
+#define ENVPOOL_TOY_TEXT_CLIFFWALKING_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 CliffWalkingEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() { return MakeDict(); }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs"_.bind(Spec<int>({-1}, {0, 47})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<int>({-1}, {0, 3})));
+  }
+};
+
+typedef class EnvSpec<CliffWalkingEnvFns> CliffWalkingEnvSpec;
+
+class CliffWalkingEnv : public Env<CliffWalkingEnvSpec> {
+ protected:
+  int x_, y_;
+  bool done_;
+
+ public:
+  CliffWalkingEnv(const Spec& spec, int env_id)
+      : Env<CliffWalkingEnvSpec>(spec, env_id), done_(true) {}
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    x_ = 3;
+    y_ = 0;
+    done_ = false;
+    State state = Allocate();
+    WriteObs(state, 0.0f);
+  }
+
+  void Step(const Action& action) override {
+    int act = action["action"_];
+    float reward = -1.0f;
+    if (act == 0) {
+      --x_;
+    } else if (act == 1) {
+      ++y_;
+    } else if (act == 2) {
+      ++x_;
+    } else {
+      --y_;
+    }
+    x_ = std::min(3, std::max(0, x_));
+    y_ = std::min(11, std::max(0, y_));
+    if (x_ == 3 && y_ > 0 && y_ < 11) {
+      reward = -100.0f;
+      x_ = 3;
+      y_ = 0;
+    }
+    if (x_ == 3 && y_ == 11) {
+      done_ = true;
+    }
+    State state = Allocate();
+    WriteObs(state, reward);
+  }
+
+ private:
+  void WriteObs(State& state, float reward) {  // NOLINT
+    state["obs"_] = x_ * 12 + y_;
+    state["reward"_] = reward;
+  }
+};
+
+typedef AsyncEnvPool<CliffWalkingEnv> CliffWalkingEnvPool;
+
+}  // namespace toy_text
+
+#endif  // ENVPOOL_TOY_TEXT_CLIFFWALKING_H_
diff --git a/envpool/toy_text/registration.py b/envpool/toy_text/registration.py
index 012037b1..776f7ea2 100644
--- a/envpool/toy_text/registration.py
+++ b/envpool/toy_text/registration.py
@@ -65,3 +65,11 @@
   gym_cls="NChainGymEnvPool",
   max_episode_steps=1000,
 )
+
+register(
+  task_id="CliffWalking-v0",
+  import_path="envpool.toy_text",
+  spec_cls="CliffWalkingEnvSpec",
+  dm_cls="CliffWalkingDMEnvPool",
+  gym_cls="CliffWalkingGymEnvPool",
+)
diff --git a/envpool/toy_text/toy_text.cc b/envpool/toy_text/toy_text.cc
index 203331d9..735da694 100644
--- a/envpool/toy_text/toy_text.cc
+++ b/envpool/toy_text/toy_text.cc
@@ -14,6 +14,7 @@
 
 #include "envpool/core/py_envpool.h"
 #include "envpool/toy_text/catch.h"
+#include "envpool/toy_text/cliffwalking.h"
 #include "envpool/toy_text/frozen_lake.h"
 #include "envpool/toy_text/nchain.h"
 #include "envpool/toy_text/taxi.h"
@@ -30,9 +31,13 @@ typedef PyEnvPool<toy_text::TaxiEnvPool> TaxiEnvPool;
 typedef PyEnvSpec<toy_text::NChainEnvSpec> NChainEnvSpec;
 typedef PyEnvPool<toy_text::NChainEnvPool> NChainEnvPool;
 
+typedef PyEnvSpec<toy_text::CliffWalkingEnvSpec> CliffWalkingEnvSpec;
+typedef PyEnvPool<toy_text::CliffWalkingEnvPool> CliffWalkingEnvPool;
+
 PYBIND11_MODULE(toy_text_envpool, m) {
   REGISTER(m, CatchEnvSpec, CatchEnvPool)
   REGISTER(m, FrozenLakeEnvSpec, FrozenLakeEnvPool)
   REGISTER(m, TaxiEnvSpec, TaxiEnvPool)
   REGISTER(m, NChainEnvSpec, NChainEnvPool)
+  REGISTER(m, CliffWalkingEnvSpec, CliffWalkingEnvPool)
 }


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 896214ce..c60a70b5 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -49,7 +49,8 @@ def test_make_classic_and_toytext(self) -> None:
       "MountainCarContinuous-v0", "Acrobot-v1"
     ]
     toytext = [
-      "Catch-v0", "FrozenLake-v1", "FrozenLake8x8-v1", "Taxi-v3", "NChain-v0"
+      "Catch-v0", "FrozenLake-v1", "FrozenLake8x8-v1", "Taxi-v3", "NChain-v0",
+      "CliffWalking-v0"
     ]
     for task_id in classic + toytext:
       envpool.make_spec(task_id)
diff --git a/envpool/toy_text/toy_text_test.py b/envpool/toy_text/toy_text_test.py
index f8c7b5d6..fef5e517 100644
--- a/envpool/toy_text/toy_text_test.py
+++ b/envpool/toy_text/toy_text_test.py
@@ -23,6 +23,8 @@
   CatchDMEnvPool,
   CatchEnvSpec,
   CatchGymEnvPool,
+  CliffWalkingEnvSpec,
+  CliffWalkingGymEnvPool,
   FrozenLakeEnvSpec,
   FrozenLakeGymEnvPool,
   NChainEnvSpec,
@@ -208,6 +210,25 @@ def test_nchain(self) -> None:
       reward += rew
     assert abs(np.mean(reward) - 1310) < 30 and abs(np.std(reward) - 78) < 15
 
+  def test_cliffwalking(self) -> None:
+    spec = CliffWalkingEnvSpec(CliffWalkingEnvSpec.gen_config())
+    env = CliffWalkingGymEnvPool(spec)
+    assert isinstance(env.observation_space, gym.spaces.Discrete)
+    assert env.observation_space.n == 48
+    assert isinstance(env.action_space, gym.spaces.Discrete)
+    assert env.action_space.n == 4
+    ref = gym.make("CliffWalking-v0")
+    for i in range(12):
+      action = [0] * 4 + [1] * i + [2] * 4
+      assert env.reset()[0] == ref.reset()
+      while len(action) > 0:
+        a = action.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]
+        if ref_done:
+          break
+
 
 if __name__ == "__main__":
   absltest.main()