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/docs/api/mujoco.rst b/docs/api/mujoco.rst
index deb2ade0..ab195d58 100644
--- a/docs/api/mujoco.rst
+++ b/docs/api/mujoco.rst
@@ -92,6 +92,21 @@ InvertedPendulum-v2/v4
 - Action space: ``(1)``, with range ``[-3, 3]``.
 
 
+Pusher-v2/v4
+------------
+
+`gym Pusher-v2 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/pusher.py>`_
+
+`gym Pusher-v4 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/pusher_v4.py>`_
+
+- Observation space: ``(23)``, first 7 elements for ``qpos[:7]``, next 7
+  elements for ``qvel[:7]``, next 3 elements for ``tips_arm``, next 3
+  elements for ``object``, next 3 elements for ``goal``;
+- Action space: ``(7)``, with range ``[-2, 2]``.
+
+
 Reacher-v2/v4
 -------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 428410fc..31d68112 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -82,6 +82,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "pusher",
+    hdrs = ["pusher.h"],
+    deps = [
+        ":mujoco_env",
+        "//envpool/core:async_envpool",
+    ],
+)
+
 cc_library(
     name = "reacher",
     hdrs = ["reacher.h"],
@@ -110,6 +119,7 @@ pybind_extension(
         ":humanoid",
         ":humanoid_standup",
         ":inverted_pendulum",
+        ":pusher",
         ":reacher",
         ":swimmer",
         "//envpool/core:py_envpool",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index e6b944e7..a3efbb1e 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -28,6 +28,8 @@
   _HumanoidStandupEnvSpec,
   _InvertedPendulumEnvPool,
   _InvertedPendulumEnvSpec,
+  _PusherEnvPool,
+  _PusherEnvSpec,
   _ReacherEnvPool,
   _ReacherEnvSpec,
   _SwimmerEnvPool,
@@ -60,6 +62,10 @@
   InvertedPendulumGymEnvPool,
 ) = py_env(_InvertedPendulumEnvSpec, _InvertedPendulumEnvPool)
 
+PusherEnvSpec, PusherDMEnvPool, PusherGymEnvPool = py_env(
+  _PusherEnvSpec, _PusherEnvPool
+)
+
 ReacherEnvSpec, ReacherDMEnvPool, ReacherGymEnvPool = py_env(
   _ReacherEnvSpec, _ReacherEnvPool
 )
@@ -87,6 +93,9 @@
   "InvertedPendulumEnvSpec",
   "InvertedPendulumDMEnvPool",
   "InvertedPendulumGymEnvPool",
+  "PusherEnvSpec",
+  "PusherDMEnvPool",
+  "PusherGymEnvPool",
   "ReacherEnvSpec",
   "ReacherDMEnvPool",
   "ReacherGymEnvPool",
diff --git a/envpool/mujoco/mujoco_envpool.cc b/envpool/mujoco/mujoco_envpool.cc
index 456dbecc..e73700f8 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/pusher.h"
 #include "envpool/mujoco/reacher.h"
 #include "envpool/mujoco/swimmer.h"
 
@@ -40,6 +41,9 @@ typedef PyEnvPool<mujoco::HumanoidStandupEnvPool> HumanoidStandupEnvPool;
 typedef PyEnvSpec<mujoco::InvertedPendulumEnvSpec> InvertedPendulumEnvSpec;
 typedef PyEnvPool<mujoco::InvertedPendulumEnvPool> InvertedPendulumEnvPool;
 
+typedef PyEnvSpec<mujoco::PusherEnvSpec> PusherEnvSpec;
+typedef PyEnvPool<mujoco::PusherEnvPool> PusherEnvPool;
+
 typedef PyEnvSpec<mujoco::ReacherEnvSpec> ReacherEnvSpec;
 typedef PyEnvPool<mujoco::ReacherEnvPool> ReacherEnvPool;
 
@@ -53,6 +57,7 @@ PYBIND11_MODULE(mujoco_envpool, m) {
   REGISTER(m, HumanoidEnvSpec, HumanoidEnvPool)
   REGISTER(m, HumanoidStandupEnvSpec, HumanoidStandupEnvPool)
   REGISTER(m, InvertedPendulumEnvSpec, InvertedPendulumEnvPool)
+  REGISTER(m, PusherEnvSpec, PusherEnvPool)
   REGISTER(m, ReacherEnvSpec, ReacherEnvPool)
   REGISTER(m, SwimmerEnvSpec, SwimmerEnvPool)
 }
diff --git a/envpool/mujoco/pusher.h b/envpool/mujoco/pusher.h
new file mode 100644
index 00000000..c121848d
--- /dev/null
+++ b/envpool/mujoco/pusher.h
@@ -0,0 +1,175 @@
+/*
+ * 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_PUSHER_H_
+#define ENVPOOL_MUJOCO_PUSHER_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 PusherEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict(
+        "max_episode_steps"_.bind(100), "reward_threshold"_.bind(0.0),
+        "frame_skip"_.bind(5), "post_constraint"_.bind(true),
+        "ctrl_cost_weight"_.bind(0.1), "dist_cost_weight"_.bind(1.0),
+        "near_cost_weight"_.bind(0.5), "reset_qvel_scale"_.bind(0.005),
+        "cylinder_x_min"_.bind(-0.3), "cylinder_x_max"_.bind(0.0),
+        "cylinder_y_min"_.bind(-0.2), "cylinder_y_max"_.bind(0.2),
+        "cylinder_dist_min"_.bind(0.17));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    mjtNum inf = std::numeric_limits<mjtNum>::infinity();
+    return MakeDict("obs"_.bind(Spec<mjtNum>({23}, {-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>({11})),
+                    "info:qvel0"_.bind(Spec<mjtNum>({11})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<mjtNum>({-1, 7}, {-2.0f, 2.0f})));
+  }
+};
+
+typedef class EnvSpec<PusherEnvFns> PusherEnvSpec;
+
+class PusherEnv : public Env<PusherEnvSpec>, public MujocoEnv {
+ protected:
+  int max_episode_steps_, elapsed_step_;
+  mjtNum ctrl_cost_weight_, dist_cost_weight_, near_cost_weight_;
+  mjtNum cylinder_dist_min_;
+  std::unique_ptr<mjtNum> qpos0_, qvel0_;  // for align check
+  std::uniform_real_distribution<> dist_qpos_x_, dist_qpos_y_, dist_qvel_;
+  bool done_;
+
+ public:
+  PusherEnv(const Spec& spec, int env_id)
+      : Env<PusherEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_] + "/mujoco/assets/pusher.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"_]),
+        dist_cost_weight_(spec.config["dist_cost_weight"_]),
+        near_cost_weight_(spec.config["near_cost_weight"_]),
+        cylinder_dist_min_(spec.config["cylinder_dist_min"_]),
+        qpos0_(new mjtNum[model_->nq]),
+        qvel0_(new mjtNum[model_->nv]),
+        dist_qpos_x_(spec.config["cylinder_x_min"_],
+                     spec.config["cylinder_x_max"_]),
+        dist_qpos_y_(spec.config["cylinder_y_min"_],
+                     spec.config["cylinder_y_max"_]),
+        dist_qvel_(-spec.config["reset_qvel_scale"_],
+                   spec.config["reset_qvel_scale"_]),
+        done_(true) {}
+
+  void MujocoResetModel() {
+    for (int i = 0; i < model_->nq - 4; ++i) {
+      data_->qpos[i] = qpos0_.get()[i] = init_qpos_[i];
+    }
+    while (1) {
+      mjtNum x = dist_qpos_x_(gen_), y = dist_qpos_y_(gen_);
+      if (std::sqrt(x * x + y * y) > cylinder_dist_min_) {
+        data_->qpos[model_->nq - 4] = qpos0_.get()[model_->nq - 4] = x;
+        data_->qpos[model_->nq - 3] = qpos0_.get()[model_->nq - 3] = y;
+        data_->qpos[model_->nq - 2] = qpos0_.get()[model_->nq - 2] = 0.0;
+        data_->qpos[model_->nq - 1] = qpos0_.get()[model_->nq - 1] = 0.0;
+        break;
+      }
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      data_->qvel[i] = qvel0_.get()[i] =
+          i < model_->nv - 4 ? 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());
+    mjtNum near_cost = GetDist(30, 33);
+    mjtNum dist_cost = GetDist(33, 36);
+    MujocoStep(act);
+
+    // ctrl_cost
+    mjtNum ctrl_cost = 0.0;
+    for (int i = 0; i < model_->nu; ++i) {
+      ctrl_cost += act[i] * act[i];
+    }
+
+    // reward and done
+    float reward = -ctrl_cost * ctrl_cost_weight_ -
+                   dist_cost * dist_cost_weight_ -
+                   near_cost * near_cost_weight_;
+    done_ = (++elapsed_step_ >= max_episode_steps_);
+    WriteObs(reward, ctrl_cost, dist_cost);
+  }
+
+ private:
+  mjtNum GetDist(int off0, int off1) {
+    mjtNum x = data_->xpos[off0 + 0] - data_->xpos[off1 + 0];
+    mjtNum y = data_->xpos[off0 + 1] - data_->xpos[off1 + 1];
+    mjtNum z = data_->xpos[off0 + 2] - data_->xpos[off1 + 2];
+    return std::sqrt(x * x + y * y + z * z);
+  }
+
+  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());
+    for (int i = 0; i < 7; ++i) {
+      *(obs++) = data_->qpos[i];
+    }
+    for (int i = 0; i < 7; ++i) {
+      *(obs++) = data_->qvel[i];
+    }
+    for (int i = 30; i < 3 * model_->nbody; ++i) {
+      *(obs++) = data_->xpos[i];
+    }
+    // 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<PusherEnv> PusherEnvPool;
+
+}  // namespace mujoco
+
+#endif  // ENVPOOL_MUJOCO_PUSHER_H_
diff --git a/envpool/mujoco/reacher.h b/envpool/mujoco/reacher.h
index 2044e0d3..2b236800 100644
--- a/envpool/mujoco/reacher.h
+++ b/envpool/mujoco/reacher.h
@@ -34,8 +34,9 @@ class ReacherEnvFns {
     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));
+        "ctrl_cost_weight"_.bind(1.0), "dist_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) {
@@ -58,7 +59,8 @@ 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_;
+  mjtNum ctrl_cost_weight_, dist_cost_weight_;
+  mjtNum 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_;
@@ -71,6 +73,7 @@ class ReacherEnv : public Env<ReacherEnvSpec>, public MujocoEnv {
         max_episode_steps_(spec.config["max_episode_steps"_]),
         elapsed_step_(max_episode_steps_ + 1),
         ctrl_cost_weight_(spec.config["ctrl_cost_weight"_]),
+        dist_cost_weight_(spec.config["dist_cost_weight"_]),
         reset_goal_scale_(spec.config["reset_goal_scale"_]),
         qpos0_(new mjtNum[model_->nq]),
         qvel0_(new mjtNum[model_->nv]),
@@ -117,6 +120,7 @@ class ReacherEnv : public Env<ReacherEnvSpec>, public MujocoEnv {
 
     // dist_cost
     mjtNum dist_cost =
+        dist_cost_weight_ *
         std::sqrt(dist_x_ * dist_x_ + dist_y_ * dist_y_ + dist_z_ * dist_z_);
     // ctrl_cost
     mjtNum ctrl_cost = 0.0;
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index b095c723..5ae29172 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -32,6 +32,8 @@
   ("HumanoidStandup", "v4", True),
   ("InvertedPendulum", "v2", False),
   ("InvertedPendulum", "v4", True),
+  ("Pusher", "v2", False),
+  ("Pusher", "v4", True),
   ("Reacher", "v2", False),
   ("Reacher", "v4", True),
   ("Swimmer", "v3", False),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index ec4907de..9d77e0a5 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",
+        "Pusher-v2",
+        "Pusher-v4",
         "Reacher-v2",
         "Reacher-v4",
         "Swimmer-v3",
diff --git a/envpool/mujoco/mujoco_test.py b/envpool/mujoco/mujoco_test.py
index 000cde45..f33f5b39 100644
--- a/envpool/mujoco/mujoco_test.py
+++ b/envpool/mujoco/mujoco_test.py
@@ -34,6 +34,8 @@
   HumanoidStandupGymEnvPool,
   InvertedPendulumEnvSpec,
   InvertedPendulumGymEnvPool,
+  PusherEnvSpec,
+  PusherGymEnvPool,
   ReacherEnvSpec,
   ReacherGymEnvPool,
   SwimmerEnvSpec,
@@ -171,6 +173,13 @@ def test_inverted_pendulum(self) -> None:
       InvertedPendulumEnvSpec, InvertedPendulumGymEnvPool
     )
 
+  def test_pusher(self) -> None:
+    env0 = mjc_mwe.PusherEnv()
+    env1 = PusherGymEnvPool(PusherEnvSpec(PusherEnvSpec.gen_config()))
+    self.run_space_check(env0, env1)
+    self.run_align_check(env0, env1, no_time_limit=True)
+    self.run_deterministic_check(PusherEnvSpec, PusherGymEnvPool)
+
   def test_reacher(self) -> None:
     env0 = mjc_mwe.ReacherEnv()
     env1 = ReacherGymEnvPool(ReacherEnvSpec(ReacherEnvSpec.gen_config()))