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 9ae8e293..d9aa57a8 100644
--- a/docs/api/mujoco.rst
+++ b/docs/api/mujoco.rst
@@ -22,3 +22,14 @@ Ant-v4
   elements for ``qvel``, other elements for clipped ``cfrc_ext`` (contact
   force);
 - Action space: ``(8)``, with range ``[-1, 1]``.
+
+
+HalfCheetah-v4
+--------------
+
+`gym HalfCheetah-v4 source code
+<https://github.com/openai/gym/blob/master/gym/envs/mujoco/half_cheetah_v4.py>`_
+
+- Observation space: ``(17)``, first 8 elements for ``qpos[1:]``, next 9
+  elements for ``qvel``;
+- Action space: ``(6)``, with range ``[-1, 1]``.
diff --git a/envpool/__init__.py b/envpool/__init__.py
index 60fb4949..e1545253 100644
--- a/envpool/__init__.py
+++ b/envpool/__init__.py
@@ -23,7 +23,7 @@
   register,
 )
 
-__version__ = "0.5.0"
+__version__ = "0.5.1"
 __all__ = [
   "register",
   "make",
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index ceaf2cb7..8c58eabf 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -37,11 +37,21 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "half_cheetah",
+    hdrs = ["half_cheetah.h"],
+    deps = [
+        ":mujoco_env",
+        "//envpool/core:async_envpool",
+    ],
+)
+
 pybind_extension(
     name = "mujoco_envpool",
     srcs = ["mujoco_envpool.cc"],
     deps = [
         ":ant",
+        ":half_cheetah",
         "//envpool/core:py_envpool",
     ],
 )
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 17e5044b..9cdf7214 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -15,12 +15,23 @@
 
 from envpool.python.api import py_env
 
-from .mujoco_envpool import _AntEnvPool, _AntEnvSpec
+from .mujoco_envpool import (
+  _AntEnvPool,
+  _AntEnvSpec,
+  _HalfCheetahEnvPool,
+  _HalfCheetahEnvSpec,
+)
 
 AntEnvSpec, AntDMEnvPool, AntGymEnvPool = py_env(_AntEnvSpec, _AntEnvPool)
+HalfCheetahEnvSpec, HalfCheetahDMEnvPool, HalfCheetahGymEnvPool = py_env(
+  _HalfCheetahEnvSpec, _HalfCheetahEnvPool
+)
 
 __all__ = [
   "AntEnvSpec",
   "AntDMEnvPool",
   "AntGymEnvPool",
+  "HalfCheetahEnvSpec",
+  "HalfCheetahDMEnvPool",
+  "HalfCheetahGymEnvPool",
 ]
diff --git a/envpool/mujoco/half_cheetah.h b/envpool/mujoco/half_cheetah.h
new file mode 100644
index 00000000..f84ef4f7
--- /dev/null
+++ b/envpool/mujoco/half_cheetah.h
@@ -0,0 +1,151 @@
+/*
+ * 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_HALF_CHEETAH_H_
+#define ENVPOOL_MUJOCO_HALF_CHEETAH_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 HalfCheetahEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict(
+        "max_episode_steps"_.bind(1000), "reward_threshold"_.bind(4800.0),
+        "frame_skip"_.bind(5), "post_constraint"_.bind(true),
+        "forward_reward_weight"_.bind(1.0), "ctrl_cost_weight"_.bind(0.1),
+        "reset_noise_scale"_.bind(0.1));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    mjtNum inf = std::numeric_limits<mjtNum>::infinity();
+    return MakeDict("obs"_.bind(Spec<mjtNum>({17}, {-inf, inf})),
+                    "info:reward_run"_.bind(Spec<mjtNum>({-1})),
+                    "info:reward_ctrl"_.bind(Spec<mjtNum>({-1})),
+                    "info:x_position"_.bind(Spec<mjtNum>({-1})),
+                    "info:x_velocity"_.bind(Spec<mjtNum>({-1})),
+                    // TODO(jiayi): remove these two lines for speed
+                    "info:qpos0"_.bind(Spec<mjtNum>({9})),
+                    "info:qvel0"_.bind(Spec<mjtNum>({9})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.bind(Spec<mjtNum>({-1, 6}, {-1.0f, 1.0f})));
+  }
+};
+
+typedef class EnvSpec<HalfCheetahEnvFns> HalfCheetahEnvSpec;
+
+class HalfCheetahEnv : public Env<HalfCheetahEnvSpec>, public MujocoEnv {
+ protected:
+  int max_episode_steps_, elapsed_step_;
+  mjtNum forward_reward_weight_, ctrl_cost_weight_;
+  std::unique_ptr<mjtNum> qpos0_, qvel0_;  // for align check
+  std::uniform_real_distribution<> dist_qpos_;
+  std::normal_distribution<> dist_qvel_;
+  bool done_;
+
+ public:
+  HalfCheetahEnv(const Spec& spec, int env_id)
+      : Env<HalfCheetahEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_] + "/mujoco/assets/half_cheetah.xml",
+                  spec.config["frame_skip"_], spec.config["post_constraint"_]),
+        max_episode_steps_(spec.config["max_episode_steps"_]),
+        elapsed_step_(max_episode_steps_ + 1),
+        forward_reward_weight_(spec.config["forward_reward_weight"_]),
+        ctrl_cost_weight_(spec.config["ctrl_cost_weight"_]),
+        qpos0_(new mjtNum[model_->nq]),
+        qvel0_(new mjtNum[model_->nv]),
+        dist_qpos_(-spec.config["reset_noise_scale"_],
+                   spec.config["reset_noise_scale"_]),
+        dist_qvel_(0, 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_qpos_(gen_);
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      data_->qvel[i] = qvel0_.get()[i] = init_qvel_[i] + dist_qvel_(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());
+    mjtNum x_before = data_->qpos[0];
+    MujocoStep(act);
+    mjtNum x_after = data_->qpos[0];
+
+    // 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 dt = frame_skip_ * model_->opt.timestep;
+    mjtNum xv = (x_after - x_before) / dt;
+    // reward and done
+    float reward = forward_reward_weight_ * xv - ctrl_cost;
+    done_ = (++elapsed_step_ >= max_episode_steps_);
+    WriteObs(reward, xv, ctrl_cost, x_after);
+  }
+
+ private:
+  void WriteObs(float reward, mjtNum xv, mjtNum ctrl_cost,
+                mjtNum x_after) {  // NOLINT
+    State state = Allocate();
+    state["reward"_] = reward;
+    // obs
+    mjtNum* obs = static_cast<mjtNum*>(state["obs"_].data());
+    for (int i = 1; i < model_->nq; ++i) {
+      *(obs++) = data_->qpos[i];
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      *(obs++) = data_->qvel[i];
+    }
+    // info
+    state["info:reward_run"_] = forward_reward_weight_ * xv;
+    state["info:reward_ctrl"_] = -ctrl_cost;
+    state["info:x_position"_] = x_after;
+    state["info:x_velocity"_] = xv;
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:qvel0"_].Assign(qvel0_.get(), model_->nv);
+  }
+};
+
+typedef AsyncEnvPool<HalfCheetahEnv> HalfCheetahEnvPool;
+
+}  // namespace mujoco
+
+#endif  // ENVPOOL_MUJOCO_HALF_CHEETAH_H_
diff --git a/envpool/mujoco/mujoco_envpool.cc b/envpool/mujoco/mujoco_envpool.cc
index 99507bcd..7bff230a 100644
--- a/envpool/mujoco/mujoco_envpool.cc
+++ b/envpool/mujoco/mujoco_envpool.cc
@@ -14,8 +14,15 @@
 
 #include "envpool/core/py_envpool.h"
 #include "envpool/mujoco/ant.h"
+#include "envpool/mujoco/half_cheetah.h"
 
 typedef PyEnvSpec<mujoco::AntEnvSpec> AntEnvSpec;
 typedef PyEnvPool<mujoco::AntEnvPool> AntEnvPool;
 
-PYBIND11_MODULE(mujoco_envpool, m) { REGISTER(m, AntEnvSpec, AntEnvPool) }
+typedef PyEnvSpec<mujoco::HalfCheetahEnvSpec> HalfCheetahEnvSpec;
+typedef PyEnvPool<mujoco::HalfCheetahEnvPool> HalfCheetahEnvPool;
+
+PYBIND11_MODULE(mujoco_envpool, m) {
+  REGISTER(m, AntEnvSpec, AntEnvPool)
+  REGISTER(m, HalfCheetahEnvSpec, HalfCheetahEnvPool)
+}
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index b2fbdccd..70cc537b 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -29,3 +29,14 @@
   reward_threshold=6000.0,
   base_path=base_path,
 )
+
+register(
+  task_id="HalfCheetah-v4",
+  import_path="envpool.mujoco",
+  spec_cls="HalfCheetahEnvSpec",
+  dm_cls="HalfCheetahDMEnvPool",
+  gym_cls="HalfCheetahGymEnvPool",
+  max_episode_steps=1000,
+  reward_threshold=4800.0,
+  base_path=base_path,
+)
diff --git a/setup.cfg b/setup.cfg
index fd76f91b..6e5a7eef 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = envpool
-version = 0.5.0
+version = 0.5.1
 author = "EnvPool Contributors"
 author_email = "sail@sea.com"
 description = "C++-based high-performance parallel environment execution engine (vectorized env) for general RL environments."


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 269a72de..10089752 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -72,7 +72,7 @@ def test_make_classic_and_toytext(self) -> None:
       env_gym.reset()
 
   def test_make_mujoco(self) -> None:
-    mujoco = ["Ant-v4"]
+    mujoco = ["Ant-v4", "HalfCheetah-v4"]
     for task_id in mujoco:
       envpool.make_spec(task_id)
       env_gym = envpool.make_gym(task_id)
diff --git a/envpool/mujoco/mujoco_test.py b/envpool/mujoco/mujoco_test.py
index 1d5367c5..bc1f39bf 100644
--- a/envpool/mujoco/mujoco_test.py
+++ b/envpool/mujoco/mujoco_test.py
@@ -21,7 +21,12 @@
 from absl import logging
 from absl.testing import absltest
 
-from envpool.mujoco import AntEnvSpec, AntGymEnvPool
+from envpool.mujoco import (
+  AntEnvSpec,
+  AntGymEnvPool,
+  HalfCheetahEnvSpec,
+  HalfCheetahGymEnvPool,
+)
 
 
 class _MujocoEnvPoolTest(absltest.TestCase):
@@ -74,7 +79,10 @@ def reset_state(
     env.set_state(qpos, qvel)
     env._mujoco_bindings.mj_forward(env.model, env.data)
 
-  def run_align_check(self, env0: gym.Env, env1: Any) -> None:
+  def run_align_check(
+    self, env0: gym.Env, env1: Any, no_time_limit: bool = False
+  ) -> None:
+    logging.info(f"align check for {env1.__class__.__name__}")
     for i in range(5):
       env0.action_space.seed(i)
       env0.reset()
@@ -86,12 +94,13 @@ def run_align_check(self, env0: gym.Env, env1: Any) -> None:
       while not d1[0]:
         cnt += 1
         a = env0.action_space.sample()
-        logging.info(f"{cnt} {a}")
+        # 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(r0, r1[0], atol=1e-4)
-        np.testing.assert_allclose(d0, d1[0])
+        if not no_time_limit:
+          np.testing.assert_allclose(d0, d1[0])
         for k in i0:
           if k in i1:
             np.testing.assert_allclose(i0[k], i1[k][0], atol=1e-4)
@@ -103,6 +112,15 @@ def test_ant(self) -> None:
     self.run_align_check(env0, env1)
     self.run_deterministic_check(AntEnvSpec, AntGymEnvPool)
 
+  def test_half_cheetah(self) -> None:
+    env0 = mjc_mwe.HalfCheetahEnv()
+    env1 = HalfCheetahGymEnvPool(
+      HalfCheetahEnvSpec(HalfCheetahEnvSpec.gen_config())
+    )
+    self.run_space_check(env0, env1)
+    self.run_align_check(env0, env1, no_time_limit=True)
+    self.run_deterministic_check(HalfCheetahEnvSpec, HalfCheetahGymEnvPool)
+
 
 if __name__ == "__main__":
   absltest.main()