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 a0a56576..2b79f461 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -44,3 +44,16 @@ HopperStand-v1, HopperHop-v1
 - Action spec: ``(4)``, with range ``[-1, 1]``;
 - ``frame_skip``: 4;
 - ``max_episode_steps``: 1000;
+
+
+WalkerRun-v1, WalkerStand-v1, WalkerWalk-v1
+-------------------------------------------
+
+`dm_control suite walker source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/walker.py>`_
+
+- Observation spec: a namedtuple with three keys: ``orientations (14)``,
+  ``height ()`` and ``velocity (9)``;
+- Action spec: ``(6)``, with range ``[-1, 1]``;
+- ``frame_skip``: 10;
+- ``max_episode_steps``: 1000;
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 30d238a7..2f3a1771 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -67,6 +67,7 @@ cc_library(
         "dmc/hopper.h",
         "dmc/mujoco_env.h",
         "dmc/utils.h",
+        "dmc/walker.h",
     ],
     data = [":gen_mujoco_dmc_xml"],
     deps = [
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index f62d74af..0fdb61ab 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -21,6 +21,8 @@
   _DmcCheetahEnvSpec,
   _DmcHopperEnvPool,
   _DmcHopperEnvSpec,
+  _DmcWalkerEnvPool,
+  _DmcWalkerEnvSpec,
 )
 
 # gym mujoco
@@ -56,6 +58,9 @@
 DmcHopperEnvSpec, DmcHopperDMEnvPool, DmcHopperGymEnvPool = py_env(
   _DmcHopperEnvSpec, _DmcHopperEnvPool
 )
+DmcWalkerEnvSpec, DmcWalkerDMEnvPool, DmcWalkerGymEnvPool = py_env(
+  _DmcWalkerEnvSpec, _DmcWalkerEnvPool
+)
 
 # gym
 GymAntEnvSpec, GymAntDMEnvPool, GymAntGymEnvPool = py_env(
@@ -109,6 +114,9 @@
   "DmcHopperEnvSpec",
   "DmcHopperDMEnvPool",
   "DmcHopperGymEnvPool",
+  "DmcWalkerEnvSpec",
+  "DmcWalkerDMEnvPool",
+  "DmcWalkerGymEnvPool",
   "GymAntEnvSpec",
   "GymAntDMEnvPool",
   "GymAntGymEnvPool",
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 25e3b7d5..769e8932 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -15,6 +15,7 @@
 #include "envpool/core/py_envpool.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/hopper.h"
+#include "envpool/mujoco/dmc/walker.h"
 
 using DmcHopperEnvSpec = PyEnvSpec<mujoco_dmc::HopperEnvSpec>;
 using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
@@ -22,7 +23,11 @@ using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
 using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
 
+using DmcWalkerEnvSpec = PyEnvSpec<mujoco_dmc::WalkerEnvSpec>;
+using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
+
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
+  REGISTER(m, DmcWalkerEnvSpec, DmcWalkerEnvPool)
 }
diff --git a/envpool/mujoco/dmc/walker.h b/envpool/mujoco/dmc/walker.h
new file mode 100644
index 00000000..0d1064db
--- /dev/null
+++ b/envpool/mujoco/dmc/walker.h
@@ -0,0 +1,173 @@
+/*
+ * 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_WALKER_H_
+#define ENVPOOL_MUJOCO_DMC_WALKER_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"
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/walker.py
+namespace mujoco_dmc {
+
+std::string GetWalkerXML(const std::string& base_path,
+                         const std::string& task_name) {
+  return GetFileContent(base_path, "walker.xml");
+}
+
+class WalkerEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(10),
+                    "task_name"_.Bind(std::string("stand")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:orientations"_.Bind(Spec<mjtNum>({14})),
+                    "obs:height"_.Bind(Spec<mjtNum>({})),
+                    "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 WalkerEnvSpec = EnvSpec<WalkerEnvFns>;
+
+class WalkerEnv : public Env<WalkerEnvSpec>, public MujocoEnv {
+ protected:
+  // Minimal height of torso over foot above which stand reward is 1.
+  const mjtNum kStandHeight = 1.2;
+  // Horizontal speeds(meters / second) above which move reward is 1.
+  const mjtNum kWalkSpeed = 1;
+  const mjtNum kRunSpeed = 8;
+  float move_speed_;
+
+ public:
+  WalkerEnv(const Spec& spec, int env_id)
+      : Env<WalkerEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_],
+            GetWalkerXML(spec.config["base_path"_], spec.config["task_name"_]),
+            spec.config["frame_skip"_], spec.config["max_episode_steps"_]) {
+    std::string task_name = spec.config["task_name"_];
+    if (task_name == "stand") {
+      move_speed_ = 0;
+    } else if (task_name == "walk") {
+      move_speed_ = kWalkSpeed;
+    } else if (task_name == "run") {
+      move_speed_ = kRunSpeed;
+    } else {
+      throw std::runtime_error("Unknown task_name for dmc hopper.");
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    // randomizers.randomize_limited_and_rotational_joints(physics,
+    // self.random)
+    RandomizeLimitedAndRotationalJoints(&gen_);
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#endif
+  }
+
+  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 {
+    auto standing = RewardTolerance(TorsoHeight(), kStandHeight,
+                                    std::numeric_limits<double>::infinity(),
+                                    kStandHeight / 2);
+    auto upright = (1 + TorsoUpright()) / 2;
+    auto stand_reward = (3 * standing + upright) / 4;
+    if (move_speed_ == 0) {
+      return static_cast<float>(stand_reward);
+    }
+    auto move_reward =
+        RewardTolerance(HorizontalVelocity(), move_speed_,
+                        std::numeric_limits<double>::infinity(),
+                        move_speed_ / 2, 0.5, SigmoidType::kLinear);
+    return static_cast<float>(stand_reward * (5 * move_reward + 1) / 6);
+  }
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    auto orient = Orientations();
+    state["obs:orientations"_].Assign(orient.begin(), orient.size());
+    state["obs:height"_] = TorsoHeight();
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+    // info for check alignment
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  mjtNum TorsoUpright() {
+    // return self.named.data.xmat['torso', 'zz']
+    return data_->xmat[1 * 9 + 8];
+  }
+  mjtNum TorsoHeight() {
+    // return self.named.data.xpos['torso', 'z']
+    return data_->xpos[1 * 3 + 2];
+  }
+  mjtNum HorizontalVelocity() {
+    // return self.named.data.sensordata['torso_subtreelinvel'][0]
+    return data_->sensordata[0];
+  }
+  std::array<mjtNum, 14> Orientations() {
+    // return self.named.data.xmat[1:, ['xx', 'xz']].ravel()
+    std::array<mjtNum, 14> orient;
+    for (int i = 0; i < 7; i++) {
+      orient[i * 2 + 0] = data_->xmat[(1 + i) * 9 + 0];
+      orient[i * 2 + 1] = data_->xmat[(1 + i) * 9 + 2];
+    }
+    return orient;
+  }
+};
+
+using WalkerEnvPool = AsyncEnvPool<WalkerEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_WALKER_H_
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index eafaca3d..4a9c7f19 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -22,8 +22,11 @@
 # from suite.BENCHMARKING
 dmc_mujoco_envs = [
   ("cheetah", "run"),
-  ("hopper", "stand"),
   ("hopper", "hop"),
+  ("hopper", "stand"),
+  ("walker", "run"),
+  ("walker", "stand"),
+  ("walker", "walk"),
 ]
 
 gym_mujoco_envs = [


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index b812656a..c110cf87 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -118,11 +118,16 @@ def test_make_mujoco_gym(self) -> None:
     )
 
   def test_make_mujoco_dmc(self) -> None:
-    self.check_step([
-      "CheetahRun-v1",
-      "HopperStand-v1",
-      "HopperHop-v1",
-    ])
+    self.check_step(
+      [
+        "CheetahRun-v1",
+        "HopperHop-v1",
+        "HopperStand-v1",
+        "WalkerRun-v1",
+        "WalkerStand-v1",
+        "WalkerWalk-v1",
+      ]
+    )
 
 
 if __name__ == "__main__":
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 059b77e6..050c8abb 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Unit tests for Mujoco dm_control suite align check."""
 
-from typing import Any, no_type_check
+from typing import Any, List, no_type_check
 
 import dm_env
 import numpy as np
@@ -26,6 +26,8 @@
   DmcCheetahEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcWalkerDMEnvPool,
+  DmcWalkerEnvSpec,
 )
 
 
@@ -90,28 +92,29 @@ def run_align_check(
         np.testing.assert_allclose(ts0.reward, ts1.reward[0])
         np.testing.assert_allclose(ts0.discount, ts1.discount[0])
 
+  def run_align_check_entry(
+    self, domain: str, tasks: List[str], spec_cls: Any, envpool_cls: Any
+  ) -> None:
+    for task in tasks:
+      env0 = suite.load(domain, task)
+      env1 = envpool_cls(spec_cls(spec_cls.gen_config(task_name=task)))
+      self.run_space_check(env0, env1)
+      self.run_align_check(env0, env1, domain)
+
   def test_hopper(self) -> None:
-    env0 = suite.load("hopper", "stand")
-    env1 = DmcHopperDMEnvPool(
-      DmcHopperEnvSpec(DmcHopperEnvSpec.gen_config(task_name="stand"))
+    self.run_align_check_entry(
+      "hopper", ["hop", "stand"], DmcHopperEnvSpec, DmcHopperDMEnvPool
     )
-    self.run_space_check(env0, env1)
-    self.run_align_check(env0, env1, "hopper")
 
-    env0 = suite.load("hopper", "hop")
-    env1 = DmcHopperDMEnvPool(
-      DmcHopperEnvSpec(DmcHopperEnvSpec.gen_config(task_name="hop"))
+  def test_cheetah(self) -> None:
+    self.run_align_check_entry(
+      "cheetah", ["run"], DmcCheetahEnvSpec, DmcCheetahDMEnvPool
     )
-    self.run_space_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"))
+  def test_walker(self) -> None:
+    self.run_align_check_entry(
+      "walker", ["run", "stand", "walk"], DmcWalkerEnvSpec, DmcWalkerDMEnvPool
     )
-    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 bd924a89..bb7de3e9 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -15,6 +15,7 @@
 
 from typing import Any, List
 
+import dm_env
 import numpy as np
 from absl.testing import absltest
 
@@ -23,6 +24,8 @@
   DmcCheetahEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcWalkerDMEnvPool,
+  DmcWalkerEnvSpec,
 )
 
 
@@ -54,7 +57,8 @@ def check(
           ) for _ in range(num_envs)
         ]
       )
-      obs0 = env0.step(action).observation
+      ts0 = env0.step(action)
+      obs0 = ts0.observation
       obs1 = env1.step(action).observation
       obs2 = env2.step(action).observation
       for k in obs_keys:
@@ -62,7 +66,7 @@ def check(
         o1 = getattr(obs1, k)
         o2 = getattr(obs2, k)
         np.testing.assert_allclose(o0, o1)
-        if np.abs(o0).sum() > 0:
+        if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
           self.assertFalse(np.allclose(o0, o2), (o0, o2))
 
   def test_hopper(self) -> None:
@@ -74,6 +78,11 @@ def test_cheetah(self) -> None:
     obs_keys = ["position", "velocity"]
     self.check(DmcCheetahEnvSpec, DmcCheetahDMEnvPool, "run", obs_keys)
 
+  def test_walker(self) -> None:
+    obs_keys = ["orientations", "height", "velocity"]
+    for task in ["run", "stand", "walk"]:
+      self.check(DmcWalkerEnvSpec, DmcWalkerDMEnvPool, task, obs_keys)
+
 
 if __name__ == "__main__":
   absltest.main()