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 f454788a..feb3c2ae 100644
--- a/docs/api/toy_text.rst
+++ b/docs/api/toy_text.rst
@@ -46,3 +46,21 @@ 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.
+
+
+NChain-v0
+---------
+
+`gym NChain-v0 source code
+<https://github.com/openai/gym/blob/v0.20.0/gym/envs/toy_text/nchain.py>`_
+
+This game presents moves along a linear chain of states, with two actions:
+
+0. forward, which moves along the chain but returns no reward
+1. backward, which returns to the beginning and has a small reward
+
+The end of the chain, however, presents a large reward, and by moving
+'forward' at the end of the chain this large reward can be repeated.
+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).
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index b0e2618e..8110c217 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -38,3 +38,4 @@ Acrobot
 frictionless
 effector
 walkable
+NChain
diff --git a/envpool/toy_text/BUILD b/envpool/toy_text/BUILD
index a77d5917..83b5b365 100644
--- a/envpool/toy_text/BUILD
+++ b/envpool/toy_text/BUILD
@@ -27,6 +27,14 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "nchain",
+    hdrs = ["nchain.h"],
+    deps = [
+        "//envpool/core:async_envpool",
+    ],
+)
+
 pybind_extension(
     name = "toy_text_envpool",
     srcs = [
@@ -35,6 +43,7 @@ pybind_extension(
     deps = [
         ":catch",
         ":frozen_lake",
+        ":nchain",
         ":taxi",
         "//envpool/core:py_envpool",
     ],
diff --git a/envpool/toy_text/__init__.py b/envpool/toy_text/__init__.py
index cd5cbcee..83b039a2 100644
--- a/envpool/toy_text/__init__.py
+++ b/envpool/toy_text/__init__.py
@@ -20,6 +20,8 @@
   _CatchEnvSpec,
   _FrozenLakeEnvPool,
   _FrozenLakeEnvSpec,
+  _NChainEnvPool,
+  _NChainEnvSpec,
   _TaxiEnvPool,
   _TaxiEnvSpec,
 )
@@ -34,6 +36,10 @@
 
 TaxiEnvSpec, TaxiDMEnvPool, TaxiGymEnvPool = py_env(_TaxiEnvSpec, _TaxiEnvPool)
 
+NChainEnvSpec, NChainDMEnvPool, NChainGymEnvPool = py_env(
+  _NChainEnvSpec, _NChainEnvPool
+)
+
 __all__ = [
   "CatchEnvSpec",
   "CatchDMEnvPool",
@@ -44,4 +50,7 @@
   "TaxiEnvSpec",
   "TaxiDMEnvPool",
   "TaxiGymEnvPool",
+  "NChainEnvSpec",
+  "NChainDMEnvPool",
+  "NChainGymEnvPool",
 ]
diff --git a/envpool/toy_text/nchain.h b/envpool/toy_text/nchain.h
new file mode 100644
index 00000000..c1839520
--- /dev/null
+++ b/envpool/toy_text/nchain.h
@@ -0,0 +1,102 @@
+/*
+ * 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/v0.20.0/gym/envs/toy_text/nchain.py
+
+#ifndef ENVPOOL_TOY_TEXT_NCHAIN_H_
+#define ENVPOOL_TOY_TEXT_NCHAIN_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 NChainEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.bind(1000));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs"_.bind(Spec<int>({-1}, {0, 4})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<int>({-1}, {0, 1})));
+  }
+};
+
+typedef class EnvSpec<NChainEnvFns> NChainEnvSpec;
+
+class NChainEnv : public Env<NChainEnvSpec> {
+ protected:
+  int s_, max_episode_steps_, elapsed_step_;
+  std::uniform_real_distribution<> dist_;
+  bool done_;
+
+ public:
+  NChainEnv(const Spec& spec, int env_id)
+      : Env<NChainEnvSpec>(spec, env_id),
+        max_episode_steps_(spec.config["max_episode_steps"_]),
+        dist_(0, 1),
+        done_(true) {}
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    s_ = 0;
+    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"_];
+    if (dist_(gen_) < 0.2) {
+      act = 1 - act;
+    }
+    float reward = 0.0f;
+    if (act) {
+      reward = 2.0f;
+      s_ = 0;
+    } else if (s_ < 4) {
+      ++s_;
+    } else {
+      reward = 10.0f;
+    }
+    State state = Allocate();
+    WriteObs(state, reward);
+  }
+
+ private:
+  void WriteObs(State& state, float reward) {  // NOLINT
+    state["obs"_] = s_;
+    state["reward"_] = reward;
+  }
+};
+
+typedef AsyncEnvPool<NChainEnv> NChainEnvPool;
+
+}  // namespace toy_text
+
+#endif  // ENVPOOL_TOY_TEXT_NCHAIN_H_
diff --git a/envpool/toy_text/registration.py b/envpool/toy_text/registration.py
index 7cc3624c..012037b1 100644
--- a/envpool/toy_text/registration.py
+++ b/envpool/toy_text/registration.py
@@ -56,3 +56,12 @@
   max_episode_steps=200,
   reward_threshold=8.0,
 )
+
+register(
+  task_id="NChain-v0",
+  import_path="envpool.toy_text",
+  spec_cls="NChainEnvSpec",
+  dm_cls="NChainDMEnvPool",
+  gym_cls="NChainGymEnvPool",
+  max_episode_steps=1000,
+)
diff --git a/envpool/toy_text/toy_text.cc b/envpool/toy_text/toy_text.cc
index 6f7b8e1d..203331d9 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/nchain.h"
 #include "envpool/toy_text/taxi.h"
 
 typedef PyEnvSpec<toy_text::CatchEnvSpec> CatchEnvSpec;
@@ -26,8 +27,12 @@ typedef PyEnvPool<toy_text::FrozenLakeEnvPool> FrozenLakeEnvPool;
 typedef PyEnvSpec<toy_text::TaxiEnvSpec> TaxiEnvSpec;
 typedef PyEnvPool<toy_text::TaxiEnvPool> TaxiEnvPool;
 
+typedef PyEnvSpec<toy_text::NChainEnvSpec> NChainEnvSpec;
+typedef PyEnvPool<toy_text::NChainEnvPool> NChainEnvPool;
+
 PYBIND11_MODULE(toy_text_envpool, m) {
   REGISTER(m, CatchEnvSpec, CatchEnvPool)
   REGISTER(m, FrozenLakeEnvSpec, FrozenLakeEnvPool)
   REGISTER(m, TaxiEnvSpec, TaxiEnvPool)
+  REGISTER(m, NChainEnvSpec, NChainEnvPool)
 }


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 67581831..896214ce 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -48,7 +48,9 @@ def test_make_classic_and_toytext(self) -> None:
       "CartPole-v0", "CartPole-v1", "Pendulum-v0", "MountainCar-v0",
       "MountainCarContinuous-v0", "Acrobot-v1"
     ]
-    toytext = ["Catch-v0", "FrozenLake-v1", "FrozenLake8x8-v1", "Taxi-v3"]
+    toytext = [
+      "Catch-v0", "FrozenLake-v1", "FrozenLake8x8-v1", "Taxi-v3", "NChain-v0"
+    ]
     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 5f9386ec..f8c7b5d6 100644
--- a/envpool/toy_text/toy_text_test.py
+++ b/envpool/toy_text/toy_text_test.py
@@ -25,6 +25,8 @@
   CatchGymEnvPool,
   FrozenLakeEnvSpec,
   FrozenLakeGymEnvPool,
+  NChainEnvSpec,
+  NChainGymEnvPool,
   TaxiEnvSpec,
   TaxiGymEnvPool,
 )
@@ -189,6 +191,23 @@ def test_taxi(self) -> None:
         assert ref_obs == obs[0] and ref_rew == rew[0] and ref_done == done[0]
       assert ref_rew == 20 and ref_done
 
+  def test_nchain(self) -> None:
+    num_envs = 100
+    spec = NChainEnvSpec(NChainEnvSpec.gen_config(num_envs=num_envs))
+    env = NChainGymEnvPool(spec)
+    assert isinstance(env.observation_space, gym.spaces.Discrete)
+    assert env.observation_space.n == 5
+    assert isinstance(env.action_space, gym.spaces.Discrete)
+    assert env.action_space.n == 2
+    done = [False]
+    env.reset()
+    reward = 0
+    while not done[0]:
+      actions = np.random.randint(2, size=(num_envs,))
+      obs, rew, done, info = env.step(actions)
+      reward += rew
+    assert abs(np.mean(reward) - 1310) < 30 and abs(np.std(reward) - 78) < 15
+
 
 if __name__ == "__main__":
   absltest.main()