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

Issue:
[Feature Request] Refactor mujoco_env.h with dm composer callback API
## Motivation

https://github.com/sail-sg/envpool/issues/59#issuecomment-1110289093

We are going to refactor the `MujocoEnv` class with deepmind composer API. It is more general with various kinds of robotics research.

## Resources

- https://arxiv.org/pdf/2006.12983.pdf page 16
- composer task: https://github.com/deepmind/dm_control/blob/main/dm_control/composer/task.py
- composer environment: https://github.com/deepmind/dm_control/blob/main/dm_control/composer/environment.py


Discussion / hints:

Fix patch:
diff --git a/docs/api/dm_control.rst b/docs/api/dm_control.rst
index 81f31b0a..a0a56576 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -20,6 +20,19 @@ converted into ``DomainNameTaskName-v1`` in envpool, e.g.,
   envpool.make_dm("BallInCupCatch-v1", num_envs=1)
 
 
+CheetahRun-v1
+-------------
+
+`dm_control suite cheetah source code
+<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/cheetah.py>`_
+
+- Observation spec: a namedtuple with two keys: ``position (8)`` and
+  ``velocity (9)``;
+- Action spec: ``(6)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_steps``: 1000;
+
+
 HopperStand-v1, HopperHop-v1
 ----------------------------
 
@@ -31,4 +44,3 @@ HopperStand-v1, HopperHop-v1
 - Action spec: ``(4)``, with range ``[-1, 1]``;
 - ``frame_skip``: 4;
 - ``max_episode_steps``: 1000;
-
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 3b00ceb6..30d238a7 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -63,6 +63,7 @@ pybind_extension(
 cc_library(
     name = "mujoco_dmc_env",
     hdrs = [
+        "dmc/cheetah.h",
         "dmc/hopper.h",
         "dmc/mujoco_env.h",
         "dmc/utils.h",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index e7825a8c..f62d74af 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -16,7 +16,12 @@
 from envpool.python.api import py_env
 
 # dmc mujoco
-from .mujoco_dmc_envpool import _DmcHopperEnvPool, _DmcHopperEnvSpec
+from .mujoco_dmc_envpool import (
+  _DmcCheetahEnvPool,
+  _DmcCheetahEnvSpec,
+  _DmcHopperEnvPool,
+  _DmcHopperEnvSpec,
+)
 
 # gym mujoco
 from .mujoco_gym_envpool import (
@@ -45,6 +50,9 @@
 )
 
 # dmc
+DmcCheetahEnvSpec, DmcCheetahDMEnvPool, DmcCheetahGymEnvPool = py_env(
+  _DmcCheetahEnvSpec, _DmcCheetahEnvPool
+)
 DmcHopperEnvSpec, DmcHopperDMEnvPool, DmcHopperGymEnvPool = py_env(
   _DmcHopperEnvSpec, _DmcHopperEnvPool
 )
@@ -95,6 +103,9 @@
 )
 
 __all__ = [
+  "DmcCheetahEnvSpec",
+  "DmcCheetahDMEnvPool",
+  "DmcCheetahGymEnvPool",
   "DmcHopperEnvSpec",
   "DmcHopperDMEnvPool",
   "DmcHopperGymEnvPool",
diff --git a/envpool/mujoco/dmc/cheetah.h b/envpool/mujoco/dmc/cheetah.h
new file mode 100644
index 00000000..88fb28bd
--- /dev/null
+++ b/envpool/mujoco/dmc/cheetah.h
@@ -0,0 +1,138 @@
+/*
+ * 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_DMC_CHEETAH_H_
+#define ENVPOOL_MUJOCO_DMC_CHEETAH_H_
+
+#include <algorithm>
+#include <limits>
+#include <memory>
+#include <random>
+#include <string>
+
+#include "envpool/core/async_envpool.h"
+#include "envpool/core/env.h"
+#include "envpool/mujoco/dmc/mujoco_env.h"
+#include "envpool/mujoco/dmc/utils.h"
+
+namespace mujoco_dmc {
+
+std::string GetCheetahXML(const std::string& base_path,
+                          const std::string& task_name) {
+  return GetFileContent(base_path, "cheetah.xml");
+}
+
+class CheetahEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
+                    "task_name"_.Bind(std::string("run")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:position"_.Bind(Spec<mjtNum>({8})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({9})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({9})),
+#endif
+                    "discount"_.Bind(Spec<float>({-1}, {0.0, 1.0})));
+  }
+  template <typename Config>
+  static decltype(auto) ActionSpec(const Config& conf) {
+    return MakeDict("action"_.Bind(Spec<mjtNum>({-1, 6}, {-1.0, 1.0})));
+  }
+};
+
+using CheetahEnvSpec = EnvSpec<CheetahEnvFns>;
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/cheetah.py#L60
+class CheetahEnv : public Env<CheetahEnvSpec>, public MujocoEnv {
+ protected:
+  const mjtNum kRunSpeed = 10;
+  std::uniform_real_distribution<> dist_uniform_;
+
+ public:
+  CheetahEnv(const Spec& spec, int env_id)
+      : Env<CheetahEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_],
+            GetCheetahXML(spec.config["base_path"_], spec.config["task_name"_]),
+            spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
+        dist_uniform_(0, 1) {}
+
+  void TaskInitializeEpisode() override {
+    assert(model_->njnt == model_->nq);
+    for (int i = 0; i < model_->njnt; i++) {
+      bool is_limited = model_->jnt_limited[i] == 1;
+      if (is_limited) {
+        mjtNum range_min = model_->jnt_range[i * 2 + 0];
+        mjtNum range_max = model_->jnt_range[i * 2 + 1];
+        mjtNum range = range_max - range_min;
+        data_->qpos[i] = dist_uniform_(gen_) * range + range_min;
+      }
+    }
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#endif
+    PhysicsStep(200, nullptr);
+    data_->time = 0;
+  }
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    ControlReset();
+    WriteState();
+  }
+
+  void Step(const Action& action) override {
+    mjtNum* act = static_cast<mjtNum*>(action["action"_].Data());
+    ControlStep(act);
+    WriteState();
+  }
+
+  float TaskGetReward() override {
+    return static_cast<float>(RewardTolerance(
+        Speed(), kRunSpeed, std::numeric_limits<double>::infinity(), kRunSpeed,
+        0, SigmoidType::kLinear));
+  }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    state["obs:position"_].Assign(data_->qpos + 1, model_->nq - 1);
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  mjtNum Speed() {
+    // return self.named.data.sensordata['torso_subtreelinvel'][0]
+    return data_->sensordata[0];
+  }
+};
+
+using CheetahEnvPool = AsyncEnvPool<CheetahEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_CHEETAH_H_
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 99d6df06..25e3b7d5 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -13,11 +13,16 @@
 // limitations under the License.
 
 #include "envpool/core/py_envpool.h"
+#include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/hopper.h"
 
 using DmcHopperEnvSpec = PyEnvSpec<mujoco_dmc::HopperEnvSpec>;
 using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 
+using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
+using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
+
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
+  REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
 }
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 9126c7cd..eafaca3d 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -21,6 +21,7 @@
 
 # from suite.BENCHMARKING
 dmc_mujoco_envs = [
+  ("cheetah", "run"),
   ("hopper", "stand"),
   ("hopper", "hop"),
 ]


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index d1e0d287..b812656a 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -119,6 +119,7 @@ def test_make_mujoco_gym(self) -> None:
 
   def test_make_mujoco_dmc(self) -> None:
     self.check_step([
+      "CheetahRun-v1",
       "HopperStand-v1",
       "HopperHop-v1",
     ])
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 8df83a77..059b77e6 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -21,7 +21,12 @@
 from absl.testing import absltest
 from dm_control import suite
 
-from envpool.mujoco import DmcHopperDMEnvPool, DmcHopperEnvSpec
+from envpool.mujoco import (
+  DmcCheetahDMEnvPool,
+  DmcCheetahEnvSpec,
+  DmcHopperDMEnvPool,
+  DmcHopperEnvSpec,
+)
 
 
 class _MujocoDmcAlignTest(absltest.TestCase):
@@ -39,10 +44,16 @@ def run_space_check(self, env0: dm_env.Environment, env1: Any) -> None:
     np.testing.assert_allclose(act0.maximum, act1.maximum)
 
   @no_type_check
-  def reset_state(self, env: dm_env.Environment, ts: dm_env.TimeStep) -> None:
+  def reset_state(
+    self, env: dm_env.Environment, ts: dm_env.TimeStep, task: str
+  ) -> None:
     # manually reset
     with env.physics.reset_context():
       env.physics.data.qpos = ts.observation.qpos0[0]
+      if task == "cheetah":
+        for _ in range(200):
+          env.physics.step()
+        env.physics.data.time = 0
 
   def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
     return np.random.uniform(
@@ -51,7 +62,9 @@ def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
       size=action_spec.shape,
     )
 
-  def run_align_check(self, env0: dm_env.Environment, env1: Any) -> None:
+  def run_align_check(
+    self, env0: dm_env.Environment, env1: Any, task: str
+  ) -> None:
     logging.info(f"align check for {env1.__class__.__name__}")
     obs_spec, action_spec = env0.observation_spec(), env0.action_spec()
     for i in range(5):
@@ -59,7 +72,7 @@ def run_align_check(self, env0: dm_env.Environment, env1: Any) -> None:
       env0.reset()
       a = self.sample_action(action_spec)
       ts = env1.reset(np.array([0]))
-      self.reset_state(env0, ts)
+      self.reset_state(env0, ts, task)
       logging.info(f'reset qpos {ts.observation.qpos0[0]}')
       cnt = 0
       done = False
@@ -70,12 +83,12 @@ def run_align_check(self, env0: dm_env.Environment, env1: Any) -> None:
         ts0 = env0.step(a)
         ts1 = env1.step(np.array([a]), np.array([0]))
         done = ts0.step_type == dm_env.StepType.LAST
-        np.testing.assert_allclose(ts0.step_type, ts1.step_type[0])
-        np.testing.assert_allclose(ts0.reward, ts1.reward[0])
-        np.testing.assert_allclose(ts0.discount, ts1.discount[0])
         o0, o1 = ts0.observation, ts1.observation
         for k in obs_spec:
           np.testing.assert_allclose(o0[k], getattr(o1, k)[0])
+        np.testing.assert_allclose(ts0.step_type, ts1.step_type[0])
+        np.testing.assert_allclose(ts0.reward, ts1.reward[0])
+        np.testing.assert_allclose(ts0.discount, ts1.discount[0])
 
   def test_hopper(self) -> None:
     env0 = suite.load("hopper", "stand")
@@ -83,14 +96,22 @@ def test_hopper(self) -> None:
       DmcHopperEnvSpec(DmcHopperEnvSpec.gen_config(task_name="stand"))
     )
     self.run_space_check(env0, env1)
-    self.run_align_check(env0, env1)
+    self.run_align_check(env0, env1, "hopper")
 
     env0 = suite.load("hopper", "hop")
     env1 = DmcHopperDMEnvPool(
       DmcHopperEnvSpec(DmcHopperEnvSpec.gen_config(task_name="hop"))
     )
     self.run_space_check(env0, env1)
-    self.run_align_check(env0, env1)
+    self.run_align_check(env0, env1, "hopper")
+
+  def test_cheetah(self) -> None:
+    env0 = suite.load("cheetah", "run")
+    env1 = DmcCheetahDMEnvPool(
+      DmcCheetahEnvSpec(DmcCheetahEnvSpec.gen_config(task_name="run"))
+    )
+    self.run_space_check(env0, env1)
+    self.run_align_check(env0, env1, "cheetah")
 
 
 if __name__ == "__main__":
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index 184abeae..bd924a89 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -18,7 +18,12 @@
 import numpy as np
 from absl.testing import absltest
 
-from envpool.mujoco import DmcHopperDMEnvPool, DmcHopperEnvSpec
+from envpool.mujoco import (
+  DmcCheetahDMEnvPool,
+  DmcCheetahEnvSpec,
+  DmcHopperDMEnvPool,
+  DmcHopperEnvSpec,
+)
 
 
 class _MujocoDmcDeterministicTest(absltest.TestCase):
@@ -65,6 +70,10 @@ def test_hopper(self) -> None:
     self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, "stand", obs_keys)
     self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, "hop", obs_keys)
 
+  def test_cheetah(self) -> None:
+    obs_keys = ["position", "velocity"]
+    self.check(DmcCheetahEnvSpec, DmcCheetahDMEnvPool, "run", obs_keys)
+
 
 if __name__ == "__main__":
   absltest.main()