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 82889723..cbe3d7db 100644
--- a/docs/api/mujoco.rst
+++ b/docs/api/mujoco.rst
@@ -46,12 +46,15 @@ Hopper-v4
 - Action space: ``(3)``, with range ``[-1, 1]``.
 
 
-Humanoid-v4
------------
+Humanoid-v4, HumanoidStandup-v4
+-------------------------------
 
 `gym Humanoid-v4 source code
 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v4.py>`_
 
+`gym HumanoidStandup-v4 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoidstandup_v4.py>`_
+
 - Observation space: ``(376)``, first 22 elements for ``qpos[2:]``, next 23
   elements for ``qvel``, next 140 elements for ``cinert`` (com-based body
   inertia and mass), next 84 elements for ``cvel`` (com-based velocity [3D
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 7eb1d4b9..67396ce8 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -64,6 +64,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "humanoid_standup",
+    hdrs = ["humanoid_standup.h"],
+    deps = [
+        ":mujoco_env",
+        "//envpool/core:async_envpool",
+    ],
+)
+
 pybind_extension(
     name = "mujoco_envpool",
     srcs = ["mujoco_envpool.cc"],
@@ -72,6 +81,7 @@ pybind_extension(
         ":half_cheetah",
         ":hopper",
         ":humanoid",
+        ":humanoid_standup",
         "//envpool/core:py_envpool",
     ],
 )
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 404ef4f3..46723799 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -24,6 +24,8 @@
   _HopperEnvSpec,
   _HumanoidEnvPool,
   _HumanoidEnvSpec,
+  _HumanoidStandupEnvPool,
+  _HumanoidStandupEnvSpec,
 )
 
 AntEnvSpec, AntDMEnvPool, AntGymEnvPool = py_env(_AntEnvSpec, _AntEnvPool)
@@ -40,6 +42,9 @@
   _HumanoidEnvSpec, _HumanoidEnvPool
 )
 
+HumanoidStandupEnvSpec, HumanoidStandupDMEnvPool, HumanoidStandupGymEnvPool \
+  = py_env(_HumanoidStandupEnvSpec, _HumanoidStandupEnvPool)
+
 __all__ = [
   "AntEnvSpec",
   "AntDMEnvPool",
@@ -53,4 +58,7 @@
   "HumanoidEnvSpec",
   "HumanoidDMEnvPool",
   "HumanoidGymEnvPool",
+  "HumanoidStandupEnvSpec",
+  "HumanoidStandupDMEnvPool",
+  "HumanoidStandupGymEnvPool",
 ]
diff --git a/envpool/mujoco/humanoid.h b/envpool/mujoco/humanoid.h
index 73b09478..7b58fd5d 100644
--- a/envpool/mujoco/humanoid.h
+++ b/envpool/mujoco/humanoid.h
@@ -127,8 +127,7 @@ class HumanoidEnv : public Env<HumanoidEnvSpec>, public MujocoEnv {
     // ctrl_cost
     mjtNum ctrl_cost = 0.0;
     for (int i = 0; i < model_->nu; ++i) {
-      // ctrl_cost += ctrl_cost_weight_ * act[i] * act[i];
-      ctrl_cost += ctrl_cost_weight_ * data_->ctrl[i] * data_->ctrl[i];
+      ctrl_cost += ctrl_cost_weight_ * act[i] * act[i];
     }
     // xv and yv
     mjtNum dt = frame_skip_ * model_->opt.timestep;
diff --git a/envpool/mujoco/humanoid_standup.h b/envpool/mujoco/humanoid_standup.h
new file mode 100644
index 00000000..e2a09945
--- /dev/null
+++ b/envpool/mujoco/humanoid_standup.h
@@ -0,0 +1,174 @@
+/*
+ * 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_HUMANOID_STANDUP_H_
+#define ENVPOOL_MUJOCO_HUMANOID_STANDUP_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 HumanoidStandupEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict(
+        "max_episode_steps"_.bind(1000), "frame_skip"_.bind(5),
+        "post_constraint"_.bind(true), "forward_reward_weight"_.bind(1.0),
+        "ctrl_cost_weight"_.bind(0.1), "contact_cost_weight"_.bind(5e-7),
+        "contact_cost_max"_.bind(10.0), "healthy_reward"_.bind(1.0),
+        "reset_noise_scale"_.bind(1e-2));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    mjtNum inf = std::numeric_limits<mjtNum>::infinity();
+    return MakeDict("obs"_.bind(Spec<mjtNum>({376}, {-inf, inf})),
+                    "info:reward_linup"_.bind(Spec<mjtNum>({-1})),
+                    "info:reward_quadctrl"_.bind(Spec<mjtNum>({-1})),
+                    "info:reward_alive"_.bind(Spec<mjtNum>({-1})),
+                    "info:reward_impact"_.bind(Spec<mjtNum>({-1})),
+                    // TODO(jiayi): remove these two lines for speed
+                    "info:qpos0"_.bind(Spec<mjtNum>({24})),
+                    "info:qvel0"_.bind(Spec<mjtNum>({23})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<mjtNum>({-1, 17}, {-0.4f, 0.4f})));
+  }
+};
+
+typedef class EnvSpec<HumanoidStandupEnvFns> HumanoidStandupEnvSpec;
+
+class HumanoidStandupEnv : public Env<HumanoidStandupEnvSpec>,
+                           public MujocoEnv {
+ protected:
+  int max_episode_steps_, elapsed_step_;
+  mjtNum ctrl_cost_weight_, contact_cost_weight_, contact_cost_max_;
+  mjtNum forward_reward_weight_, healthy_reward_;
+  std::unique_ptr<mjtNum> qpos0_, qvel0_;  // for align check
+  std::uniform_real_distribution<> dist_;
+  bool done_;
+
+ public:
+  HumanoidStandupEnv(const Spec& spec, int env_id)
+      : Env<HumanoidStandupEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_] + "/mujoco/assets/humanoidstandup.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"_]),
+        contact_cost_weight_(spec.config["contact_cost_weight"_]),
+        contact_cost_max_(spec.config["contact_cost_max"_]),
+        forward_reward_weight_(spec.config["forward_reward_weight"_]),
+        healthy_reward_(spec.config["healthy_reward"_]),
+        qpos0_(new mjtNum[model_->nq]),
+        qvel0_(new mjtNum[model_->nv]),
+        dist_(-spec.config["reset_noise_scale"_],
+              spec.config["reset_noise_scale"_]),
+        done_(true) {}
+
+  void MujocoResetModel() {
+    for (int i = 0; i < model_->nq; ++i) {
+      data_->qpos[i] = qpos0_.get()[i] = init_qpos_[i] + dist_(gen_);
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      data_->qvel[i] = qvel0_.get()[i] = init_qvel_[i] + dist_(gen_);
+    }
+  }
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    done_ = false;
+    elapsed_step_ = 0;
+    MujocoReset();
+    WriteObs(0.0f, 0, 0, 0);
+  }
+
+  void Step(const Action& action) override {
+    // step
+    mjtNum* act = static_cast<mjtNum*>(action["action"_].data());
+    MujocoStep(act);
+
+    // ctrl_cost
+    mjtNum ctrl_cost = 0.0;
+    for (int i = 0; i < model_->nu; ++i) {
+      ctrl_cost += ctrl_cost_weight_ * act[i] * act[i];
+    }
+    // xv
+    mjtNum xv = data_->qpos[2] / model_->opt.timestep;
+    // contact cost
+    mjtNum contact_cost = 0.0;
+    for (int i = 0; i < 6 * model_->nbody; ++i) {
+      mjtNum x = data_->cfrc_ext[i];
+      contact_cost += contact_cost_weight_ * x * x;
+    }
+    contact_cost = std::min(contact_cost, contact_cost_max_);
+
+    // reward and done
+    float reward = xv * forward_reward_weight_ + healthy_reward_ - ctrl_cost -
+                   contact_cost;
+    done_ = (++elapsed_step_ >= max_episode_steps_);
+    WriteObs(reward, xv, ctrl_cost, contact_cost);
+  }
+
+ private:
+  void WriteObs(float reward, mjtNum xv, mjtNum ctrl_cost,
+                mjtNum contact_cost) {
+    State state = Allocate();
+    state["reward"_] = reward;
+    // obs
+    mjtNum* obs = static_cast<mjtNum*>(state["obs"_].data());
+    for (int i = 2; i < model_->nq; ++i) {
+      *(obs++) = data_->qpos[i];
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      *(obs++) = data_->qvel[i];
+    }
+    for (int i = 0; i < 10 * model_->nbody; ++i) {
+      *(obs++) = data_->cinert[i];
+    }
+    for (int i = 0; i < 6 * model_->nbody; ++i) {
+      *(obs++) = data_->cvel[i];
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      *(obs++) = data_->qfrc_actuator[i];
+    }
+    for (int i = 0; i < 6 * model_->nbody; ++i) {
+      *(obs++) = data_->cfrc_ext[i];
+    }
+    // info
+    state["info:reward_linup"_] = xv * forward_reward_weight_;
+    state["info:reward_quadctrl"_] = -ctrl_cost;
+    state["info:reward_impact"_] = -contact_cost;
+    state["info:reward_alive"_] = healthy_reward_;
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:qvel0"_].Assign(qvel0_.get(), model_->nv);
+  }
+};
+
+typedef AsyncEnvPool<HumanoidStandupEnv> HumanoidStandupEnvPool;
+
+}  // namespace mujoco
+
+#endif  // ENVPOOL_MUJOCO_HUMANOID_STANDUP_H_
diff --git a/envpool/mujoco/mujoco_envpool.cc b/envpool/mujoco/mujoco_envpool.cc
index ab1d8108..1ae1deae 100644
--- a/envpool/mujoco/mujoco_envpool.cc
+++ b/envpool/mujoco/mujoco_envpool.cc
@@ -17,6 +17,7 @@
 #include "envpool/mujoco/half_cheetah.h"
 #include "envpool/mujoco/hopper.h"
 #include "envpool/mujoco/humanoid.h"
+#include "envpool/mujoco/humanoid_standup.h"
 
 typedef PyEnvSpec<mujoco::AntEnvSpec> AntEnvSpec;
 typedef PyEnvPool<mujoco::AntEnvPool> AntEnvPool;
@@ -30,9 +31,13 @@ typedef PyEnvPool<mujoco::HopperEnvPool> HopperEnvPool;
 typedef PyEnvSpec<mujoco::HumanoidEnvSpec> HumanoidEnvSpec;
 typedef PyEnvPool<mujoco::HumanoidEnvPool> HumanoidEnvPool;
 
+typedef PyEnvSpec<mujoco::HumanoidStandupEnvSpec> HumanoidStandupEnvSpec;
+typedef PyEnvPool<mujoco::HumanoidStandupEnvPool> HumanoidStandupEnvPool;
+
 PYBIND11_MODULE(mujoco_envpool, m) {
   REGISTER(m, AntEnvSpec, AntEnvPool)
   REGISTER(m, HalfCheetahEnvSpec, HalfCheetahEnvPool)
   REGISTER(m, HopperEnvSpec, HopperEnvPool)
   REGISTER(m, HumanoidEnvSpec, HumanoidEnvPool)
+  REGISTER(m, HumanoidStandupEnvSpec, HumanoidStandupEnvPool)
 }
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index ba6e56b9..84ab4825 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -61,3 +61,13 @@
   max_episode_steps=1000,
   base_path=base_path,
 )
+
+register(
+  task_id="HumanoidStandup-v4",
+  import_path="envpool.mujoco",
+  spec_cls="HumanoidStandupEnvSpec",
+  dm_cls="HumanoidStandupDMEnvPool",
+  gym_cls="HumanoidStandupGymEnvPool",
+  max_episode_steps=1000,
+  base_path=base_path,
+)


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index b13b6d21..fd184f63 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -90,12 +90,15 @@ def test_make_toytext(self) -> None:
     )
 
   def test_make_mujoco(self) -> None:
-    self.check_step([
-      "Ant-v4",
-      "HalfCheetah-v4",
-      "Hopper-v4",
-      "Humanoid-v4",
-    ])
+    self.check_step(
+      [
+        "Ant-v4",
+        "HalfCheetah-v4",
+        "Hopper-v4",
+        "Humanoid-v4",
+        "HumanoidStandup-v4",
+      ]
+    )
 
 
 if __name__ == "__main__":
diff --git a/envpool/mujoco/mujoco_test.py b/envpool/mujoco/mujoco_test.py
index b43e83f4..f16d20b2 100644
--- a/envpool/mujoco/mujoco_test.py
+++ b/envpool/mujoco/mujoco_test.py
@@ -30,6 +30,8 @@
   HopperGymEnvPool,
   HumanoidEnvSpec,
   HumanoidGymEnvPool,
+  HumanoidStandupEnvSpec,
+  HumanoidStandupGymEnvPool,
 )
 
 
@@ -101,7 +103,7 @@ def run_align_check(
         # logging.info(f"{cnt} {a}")
         o0, r0, d0, i0 = env0.step(a)
         o1, r1, d1, i1 = env1.step(np.array([a]), np.array([0]))
-        np.testing.assert_allclose(o0, o1[0], atol=1e-4)
+        np.testing.assert_allclose(o0, o1[0], atol=3e-4)
         np.testing.assert_allclose(r0, r1[0], atol=1e-4)
         if not no_time_limit:
           np.testing.assert_allclose(d0, d1[0])
@@ -139,6 +141,17 @@ def test_humanoid(self) -> None:
     self.run_align_check(env0, env1)
     self.run_deterministic_check(HumanoidEnvSpec, HumanoidGymEnvPool)
 
+  def test_humanoid_standup(self) -> None:
+    env0 = mjc_mwe.HumanoidStandupEnv()
+    env1 = HumanoidStandupGymEnvPool(
+      HumanoidStandupEnvSpec(HumanoidStandupEnvSpec.gen_config())
+    )
+    self.run_space_check(env0, env1)
+    self.run_align_check(env0, env1, no_time_limit=True)
+    self.run_deterministic_check(
+      HumanoidStandupEnvSpec, HumanoidStandupGymEnvPool
+    )
+
 
 if __name__ == "__main__":
   absltest.main()