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

Issue:
[Feature Request] Mujoco integration
https://github.com/openai/gym/tree/master/gym/envs/mujoco

Env List:

- [x] Ant-v4 (#74)
- [x] HalfCheetah-v4 (#75)
- [x] Hopper-v4 (#76)
- [x] Humanoid-v4 (#77)
- [x] HumanoidStandup-v4 (#78)
- [x] InvertedDoublePendulum-v4 (@Benjamin-eecs, #83)
- [x] InvertedPendulum-v4 (#79)
- [x] Pusher-v4 (#82)
- [x] Reacher-v4 (#81)
- [x] Swimmer-v4 (#80)
- [x] Walker2d-v4 (@Benjamin-eecs, #86)
- [x] add other options to align with gym (#93)

Road Map:

- [x] Get comfortable with current codebase, go through https://envpool.readthedocs.io/en/latest/pages/env.html and add a toy environment by yourself locally;
- [x] Download Mujoco and run on your local machine [1] [5], try with different env settings and see the actual behavior;
- [x] Go through their code [1] [2] (I think it's better to go through both openai and deepmind versions, but only use deepmind's solution as reference), understand their ctype APIs and what we can use to bind with EnvPool APIs [3];
- [x] Integrate only one game and let it work;
- [x] Add some unit tests (good to submit the first PR here);
- [x] Integrate other environments (submit another PR) and related tests.

Resources:
1. https://github.com/openai/mujoco-py 
2. https://github.com/deepmind/dm_control/tree/master/dm_control/mujoco 
3. https://github.com/deepmind/mujoco/blob/main/doc/programming.rst
4. It is quite similar with Atari games which we have already integrated: https://github.com/mgbellemare/Arcade-Learning-Environment
5. First install gym and mujoco, then run with
```python
import gym
env = gym.make("Ant-v3")
env.reset()
for _ in range(10):
  env.step(env.action_space.sample())
  env.render()
```
6. https://github.com/ikostrikov/gym_dmc/blob/master/compare.py a checker


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 2de28df7..d1888bf7 100644
--- a/README.md
+++ b/README.md
@@ -16,11 +16,11 @@
 **EnvPool** is a C++-based batched environment pool with pybind11 and thread pool. It has high performance (\~1M raw FPS in DGX on Atari games) and compatible APIs (supports both gym and dm\_env, both sync and async, both single and multi player environment). Currently it supports:
 
 - [x] Atari games
+- [x] Mujoco
 - [x] Classic RL envs: CartPole, MountainCar, Pendulum, Acrobot
 - [x] Toy text RL envs: Catch, FrozenLake, Taxi, NChain, CliffWalking, Blackjack
 - [x] ViZDoom single player
 - [ ] Box2D
-- [ ] Mujoco
 - [ ] Procgen
 - [ ] Minigrid
 
diff --git a/docs/api/mujoco.rst b/docs/api/mujoco.rst
index 3fe870da..deb2ade0 100644
--- a/docs/api/mujoco.rst
+++ b/docs/api/mujoco.rst
@@ -12,8 +12,8 @@ You can set ``post_constraint`` to ``False`` to disable the bug fix with
 environments' standard approach.
 
 
-Ant-v3/4
---------
+Ant-v3/v4
+---------
 
 `gym Ant-v3 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/ant_v3.py>`_
@@ -27,8 +27,8 @@ Ant-v3/4
 - Action space: ``(8)``, with range ``[-1, 1]``.
 
 
-HalfCheetah-v3/4
-----------------
+HalfCheetah-v3/v4
+-----------------
 
 `gym HalfCheetah-v3 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/half_cheetah_v3.py>`_
@@ -41,8 +41,8 @@ HalfCheetah-v3/4
 - Action space: ``(6)``, with range ``[-1, 1]``.
 
 
-Hopper-v3/4
------------
+Hopper-v3/v4
+------------
 
 `gym Hopper-v3 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/hopper_v3.py>`_
@@ -55,8 +55,8 @@ Hopper-v3/4
 - Action space: ``(3)``, with range ``[-1, 1]``.
 
 
-Humanoid-v3/4, HumanoidStandup-v2/4
------------------------------------
+Humanoid-v3/v4, HumanoidStandup-v2/v4
+-------------------------------------
 
 `gym Humanoid-v3 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v3.py>`_
@@ -78,8 +78,8 @@ Humanoid-v3/4, HumanoidStandup-v2/4
 - Action space: ``(17)``, with range ``[-0.4, 0.4]``.
 
 
-InvertedPendulum-v2/4
----------------------
+InvertedPendulum-v2/v4
+----------------------
 
 `gym InvertedPendulum-v2 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/inverted_pendulum.py>`_
@@ -92,8 +92,24 @@ InvertedPendulum-v2/4
 - Action space: ``(1)``, with range ``[-3, 3]``.
 
 
-Swimmer-v3/4
-------------
+Reacher-v2/v4
+-------------
+
+`gym Reacher-v2 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/reacher.py>`_
+
+`gym Reacher-v4 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/reacher_v4.py>`_
+
+- Observation space: ``(11)``, first 2 elements for ``cos(qpos[:2])``, next 2
+  elements for ``sin(qpos[:2])``, next 2 elements for ``qpos[2:]``, next 2
+  elements for ``qvel[:2]``, next 3 elements for ``dist``, a.k.a.
+  ``fingertip - target``;
+- Action space: ``(2)``, with range ``[-1, 1]``.
+
+
+Swimmer-v3/v4
+-------------
 
 `gym Swimmer-v3 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/swimmer_v3.py>`_
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index 016d9ad0..00067b54 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -48,3 +48,4 @@ RGB
 py
 Mujoco
 tran
+Reacher
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 8dbc9c73..428410fc 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -82,6 +82,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "reacher",
+    hdrs = ["reacher.h"],
+    deps = [
+        ":mujoco_env",
+        "//envpool/core:async_envpool",
+    ],
+)
+
 cc_library(
     name = "swimmer",
     hdrs = ["swimmer.h"],
@@ -101,6 +110,7 @@ pybind_extension(
         ":humanoid",
         ":humanoid_standup",
         ":inverted_pendulum",
+        ":reacher",
         ":swimmer",
         "//envpool/core:py_envpool",
     ],
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index db183a4a..e6b944e7 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -28,6 +28,8 @@
   _HumanoidStandupEnvSpec,
   _InvertedPendulumEnvPool,
   _InvertedPendulumEnvSpec,
+  _ReacherEnvPool,
+  _ReacherEnvSpec,
   _SwimmerEnvPool,
   _SwimmerEnvSpec,
 )
@@ -58,6 +60,10 @@
   InvertedPendulumGymEnvPool,
 ) = py_env(_InvertedPendulumEnvSpec, _InvertedPendulumEnvPool)
 
+ReacherEnvSpec, ReacherDMEnvPool, ReacherGymEnvPool = py_env(
+  _ReacherEnvSpec, _ReacherEnvPool
+)
+
 SwimmerEnvSpec, SwimmerDMEnvPool, SwimmerGymEnvPool = py_env(
   _SwimmerEnvSpec, _SwimmerEnvPool
 )
@@ -81,6 +87,9 @@
   "InvertedPendulumEnvSpec",
   "InvertedPendulumDMEnvPool",
   "InvertedPendulumGymEnvPool",
+  "ReacherEnvSpec",
+  "ReacherDMEnvPool",
+  "ReacherGymEnvPool",
   "SwimmerEnvSpec",
   "SwimmerDMEnvPool",
   "SwimmerGymEnvPool",
diff --git a/envpool/mujoco/mujoco_envpool.cc b/envpool/mujoco/mujoco_envpool.cc
index c364e980..456dbecc 100644
--- a/envpool/mujoco/mujoco_envpool.cc
+++ b/envpool/mujoco/mujoco_envpool.cc
@@ -19,6 +19,7 @@
 #include "envpool/mujoco/humanoid.h"
 #include "envpool/mujoco/humanoid_standup.h"
 #include "envpool/mujoco/inverted_pendulum.h"
+#include "envpool/mujoco/reacher.h"
 #include "envpool/mujoco/swimmer.h"
 
 typedef PyEnvSpec<mujoco::AntEnvSpec> AntEnvSpec;
@@ -39,6 +40,9 @@ typedef PyEnvPool<mujoco::HumanoidStandupEnvPool> HumanoidStandupEnvPool;
 typedef PyEnvSpec<mujoco::InvertedPendulumEnvSpec> InvertedPendulumEnvSpec;
 typedef PyEnvPool<mujoco::InvertedPendulumEnvPool> InvertedPendulumEnvPool;
 
+typedef PyEnvSpec<mujoco::ReacherEnvSpec> ReacherEnvSpec;
+typedef PyEnvPool<mujoco::ReacherEnvPool> ReacherEnvPool;
+
 typedef PyEnvSpec<mujoco::SwimmerEnvSpec> SwimmerEnvSpec;
 typedef PyEnvPool<mujoco::SwimmerEnvPool> SwimmerEnvPool;
 
@@ -49,5 +53,6 @@ PYBIND11_MODULE(mujoco_envpool, m) {
   REGISTER(m, HumanoidEnvSpec, HumanoidEnvPool)
   REGISTER(m, HumanoidStandupEnvSpec, HumanoidStandupEnvPool)
   REGISTER(m, InvertedPendulumEnvSpec, InvertedPendulumEnvPool)
+  REGISTER(m, ReacherEnvSpec, ReacherEnvPool)
   REGISTER(m, SwimmerEnvSpec, SwimmerEnvPool)
 }
diff --git a/envpool/mujoco/reacher.h b/envpool/mujoco/reacher.h
new file mode 100644
index 00000000..2044e0d3
--- /dev/null
+++ b/envpool/mujoco/reacher.h
@@ -0,0 +1,171 @@
+/*
+ * Copyright 2022 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.
+ */
+
+#ifndef ENVPOOL_MUJOCO_REACHER_H_
+#define ENVPOOL_MUJOCO_REACHER_H_
+
+#include <algorithm>
+#include <limits>
+#include <memory>
+#include <string>
+
+#include "envpool/core/async_envpool.h"
+#include "envpool/core/env.h"
+#include "envpool/mujoco/mujoco_env.h"
+
+namespace mujoco {
+
+class ReacherEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict(
+        "max_episode_steps"_.bind(50), "reward_threshold"_.bind(-3.75),
+        "frame_skip"_.bind(2), "post_constraint"_.bind(true),
+        "ctrl_cost_weight"_.bind(1.0), "reset_qpos_scale"_.bind(0.1),
+        "reset_qvel_scale"_.bind(0.005), "reset_goal_scale"_.bind(0.2));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    mjtNum inf = std::numeric_limits<mjtNum>::infinity();
+    return MakeDict("obs"_.bind(Spec<mjtNum>({11}, {-inf, inf})),
+                    "info:reward_dist"_.bind(Spec<mjtNum>({-1})),
+                    "info:reward_ctrl"_.bind(Spec<mjtNum>({-1})),
+                    // TODO(jiayi): remove these two lines for speed
+                    "info:qpos0"_.bind(Spec<mjtNum>({4})),
+                    "info:qvel0"_.bind(Spec<mjtNum>({4})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<mjtNum>({-1, 2}, {-1.0f, 1.0f})));
+  }
+};
+
+typedef class EnvSpec<ReacherEnvFns> ReacherEnvSpec;
+
+class ReacherEnv : public Env<ReacherEnvSpec>, public MujocoEnv {
+ protected:
+  int max_episode_steps_, elapsed_step_;
+  mjtNum ctrl_cost_weight_, reset_goal_scale_, dist_x_, dist_y_, dist_z_;
+  std::unique_ptr<mjtNum> qpos0_, qvel0_;  // for align check
+  std::uniform_real_distribution<> dist_qpos_, dist_qvel_, dist_goal_;
+  bool done_;
+
+ public:
+  ReacherEnv(const Spec& spec, int env_id)
+      : Env<ReacherEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_] + "/mujoco/assets/reacher.xml",
+                  spec.config["frame_skip"_], spec.config["post_constraint"_]),
+        max_episode_steps_(spec.config["max_episode_steps"_]),
+        elapsed_step_(max_episode_steps_ + 1),
+        ctrl_cost_weight_(spec.config["ctrl_cost_weight"_]),
+        reset_goal_scale_(spec.config["reset_goal_scale"_]),
+        qpos0_(new mjtNum[model_->nq]),
+        qvel0_(new mjtNum[model_->nv]),
+        dist_qpos_(-spec.config["reset_qpos_scale"_],
+                   spec.config["reset_qpos_scale"_]),
+        dist_qvel_(-spec.config["reset_qvel_scale"_],
+                   spec.config["reset_qvel_scale"_]),
+        dist_goal_(-spec.config["reset_goal_scale"_],
+                   spec.config["reset_goal_scale"_]),
+        done_(true) {}
+
+  void MujocoResetModel() {
+    for (int i = 0; i < model_->nq - 2; ++i) {
+      data_->qpos[i] = qpos0_.get()[i] = init_qpos_[i] + dist_qpos_(gen_);
+    }
+    while (1) {
+      mjtNum x = dist_goal_(gen_), y = dist_goal_(gen_);
+      if (std::sqrt(x * x + y * y) < reset_goal_scale_) {
+        data_->qpos[model_->nq - 2] = qpos0_.get()[model_->nq - 2] = x;
+        data_->qpos[model_->nq - 1] = qpos0_.get()[model_->nq - 1] = y;
+        break;
+      }
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      data_->qvel[i] = qvel0_.get()[i] =
+          i < model_->nv - 2 ? init_qvel_[i] + dist_qvel_(gen_) : 0.0;
+    }
+  }
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    done_ = false;
+    elapsed_step_ = 0;
+    MujocoReset();
+    WriteObs(0.0f, 0, 0);
+  }
+
+  void Step(const Action& action) override {
+    // step
+    mjtNum* act = static_cast<mjtNum*>(action["action"_].data());
+    GetDist();
+    MujocoStep(act);
+
+    // dist_cost
+    mjtNum dist_cost =
+        std::sqrt(dist_x_ * dist_x_ + dist_y_ * dist_y_ + dist_z_ * dist_z_);
+    // ctrl_cost
+    mjtNum ctrl_cost = 0.0;
+    for (int i = 0; i < model_->nu; ++i) {
+      ctrl_cost += ctrl_cost_weight_ * act[i] * act[i];
+    }
+
+    // reward and done
+    float reward = -dist_cost - ctrl_cost;
+    done_ = (++elapsed_step_ >= max_episode_steps_);
+    WriteObs(reward, ctrl_cost, dist_cost);
+  }
+
+ private:
+  void GetDist() {
+    dist_x_ = data_->xpos[9] - data_->xpos[12];
+    dist_y_ = data_->xpos[10] - data_->xpos[13];
+    dist_z_ = data_->xpos[11] - data_->xpos[14];
+  }
+
+  void WriteObs(float reward, mjtNum ctrl_cost, mjtNum dist_cost) {
+    State state = Allocate();
+    state["reward"_] = reward;
+    // obs
+    mjtNum* obs = static_cast<mjtNum*>(state["obs"_].data());
+    *(obs++) = std::cos(data_->qpos[0]);
+    *(obs++) = std::cos(data_->qpos[1]);
+    *(obs++) = std::sin(data_->qpos[0]);
+    *(obs++) = std::sin(data_->qpos[1]);
+    for (int i = 2; i < model_->nq; ++i) {
+      *(obs++) = data_->qpos[i];
+    }
+    for (int i = 0; i < 2; ++i) {
+      *(obs++) = data_->qvel[i];
+    }
+    GetDist();
+    *(obs++) = dist_x_;
+    *(obs++) = dist_y_;
+    *(obs++) = dist_z_;
+    // info
+    state["info:reward_dist"_] = -dist_cost;
+    state["info:reward_ctrl"_] = -ctrl_cost;
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:qvel0"_].Assign(qvel0_.get(), model_->nv);
+  }
+};
+
+typedef AsyncEnvPool<ReacherEnv> ReacherEnvPool;
+
+}  // namespace mujoco
+
+#endif  // ENVPOOL_MUJOCO_REACHER_H_
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 5b1a9de6..b095c723 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -32,6 +32,8 @@
   ("HumanoidStandup", "v4", True),
   ("InvertedPendulum", "v2", False),
   ("InvertedPendulum", "v4", True),
+  ("Reacher", "v2", False),
+  ("Reacher", "v4", True),
   ("Swimmer", "v3", False),
   ("Swimmer", "v4", True),
 ]


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 9bc4f93a..ec4907de 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -104,6 +104,8 @@ def test_make_mujoco(self) -> None:
         "HumanoidStandup-v4",
         "InvertedPendulum-v2",
         "InvertedPendulum-v4",
+        "Reacher-v2",
+        "Reacher-v4",
         "Swimmer-v3",
         "Swimmer-v4",
       ]
diff --git a/envpool/mujoco/mujoco_test.py b/envpool/mujoco/mujoco_test.py
index 253475f4..000cde45 100644
--- a/envpool/mujoco/mujoco_test.py
+++ b/envpool/mujoco/mujoco_test.py
@@ -34,6 +34,8 @@
   HumanoidStandupGymEnvPool,
   InvertedPendulumEnvSpec,
   InvertedPendulumGymEnvPool,
+  ReacherEnvSpec,
+  ReacherGymEnvPool,
   SwimmerEnvSpec,
   SwimmerGymEnvPool,
 )
@@ -99,6 +101,8 @@ def run_align_check(
       a = env0.action_space.sample()
       obs, _, _, info = env1.step(np.array([a]), np.array([0]))
       self.reset_state(env0, info["qpos0"][0], info["qvel0"][0])
+      logging.info(f'reset qpos {info["qpos0"][0]}')
+      logging.info(f'reset qvel {info["qvel0"][0]}')
       d1 = np.array([False])
       cnt = 0
       while not d1[0]:
@@ -167,6 +171,13 @@ def test_inverted_pendulum(self) -> None:
       InvertedPendulumEnvSpec, InvertedPendulumGymEnvPool
     )
 
+  def test_reacher(self) -> None:
+    env0 = mjc_mwe.ReacherEnv()
+    env1 = ReacherGymEnvPool(ReacherEnvSpec(ReacherEnvSpec.gen_config()))
+    self.run_space_check(env0, env1)
+    self.run_align_check(env0, env1, no_time_limit=True)
+    self.run_deterministic_check(ReacherEnvSpec, ReacherGymEnvPool)
+
   def test_swimmer(self) -> None:
     env0 = mjc_mwe.SwimmerEnv()
     env1 = SwimmerGymEnvPool(SwimmerEnvSpec(SwimmerEnvSpec.gen_config()))