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/env/dm_control.rst b/docs/env/dm_control.rst
index 10ea5e96..be16e5ff 100644
--- a/docs/env/dm_control.rst
+++ b/docs/env/dm_control.rst
@@ -160,6 +160,20 @@ HumanoidStand-v1, HumanoidWalk-v1, HumanoidRun-v1, HumanoidRunPureState-v1
     ``HumanoidRunPureState-v1`` tasks.
 
 
+HumanoidCMUStand-v1, HumanoidCMURun-v1
+--------------------------------------
+
+`dm_control suite humanoid-CMU source code
+<https://github.com/deepmind/dm_control/blob/1.0.3.post1/dm_control/suite/humanoid_CMU.py>`_
+
+- Observation spec: a namedtuple with six keys: ``joint_angles (56)``,
+  ``head_height ()``, ``extremities (12)``, ``torso_vertical (3)``,
+  ``com_velocity (3)`` and ``velocity (62)``;
+- Action spec: ``(56)``, with range ``[-1, 1]``;
+- ``frame_skip``: 10;
+- ``max_episode_steps``: 1000;
+
+
 PendulumSwingup-v1
 ------------------
 
@@ -176,7 +190,7 @@ PendulumSwingup-v1
 PointMassEasy-v1, PointMassHard-v1
 ----------------------------------
 
-`dm_control suite point mass source code
+`dm_control suite point-mass source code
 <https://github.com/deepmind/dm_control/blob/1.0.3.post1/dm_control/suite/point_mass.py>`_
 
 - Observation spec: a namedtuple with three keys: ``position (2)``,
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index cc159672..bd8fb982 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -75,6 +75,7 @@ cc_library(
         "dmc/fish.h",
         "dmc/hopper.h",
         "dmc/humanoid.h",
+        "dmc/humanoid_CMU.h",
         "dmc/manipulator.h",
         "dmc/mujoco_env.h",
         "dmc/pendulum.h",
@@ -152,9 +153,9 @@ py_test(
 )
 
 py_test(
-    name = "mujoco_dmc_deterministic_test",
+    name = "mujoco_dmc_suite_deterministic_test",
     size = "enormous",
-    srcs = ["dmc/mujoco_dmc_deterministic_test.py"],
+    srcs = ["dmc/mujoco_dmc_suite_deterministic_test.py"],
     deps = [
         ":mujoco_dmc",
         requirement("numpy"),
@@ -163,9 +164,34 @@ py_test(
 )
 
 py_test(
-    name = "mujoco_dmc_align_test",
+    name = "mujoco_dmc_suite_align_test",
     size = "enormous",
-    srcs = ["dmc/mujoco_dmc_align_test.py"],
+    srcs = ["dmc/mujoco_dmc_suite_align_test.py"],
+    deps = [
+        ":mujoco_dmc",
+        requirement("numpy"),
+        requirement("absl-py"),
+        requirement("mujoco"),
+        requirement("dm-env"),
+        requirement("dm-control"),
+    ],
+)
+
+py_test(
+    name = "mujoco_dmc_suite_ext_deterministic_test",
+    size = "enormous",
+    srcs = ["dmc/mujoco_dmc_suite_ext_deterministic_test.py"],
+    deps = [
+        ":mujoco_dmc",
+        requirement("numpy"),
+        requirement("absl-py"),
+    ],
+)
+
+py_test(
+    name = "mujoco_dmc_suite_ext_align_test",
+    size = "enormous",
+    srcs = ["dmc/mujoco_dmc_suite_ext_align_test.py"],
     deps = [
         ":mujoco_dmc",
         requirement("numpy"),
diff --git a/envpool/mujoco/dmc/__init__.py b/envpool/mujoco/dmc/__init__.py
index 8e1abc00..1754d110 100644
--- a/envpool/mujoco/dmc/__init__.py
+++ b/envpool/mujoco/dmc/__init__.py
@@ -28,6 +28,8 @@
   _DmcFishEnvSpec,
   _DmcHopperEnvPool,
   _DmcHopperEnvSpec,
+  _DmcHumanoidCMUEnvPool,
+  _DmcHumanoidCMUEnvSpec,
   _DmcHumanoidEnvPool,
   _DmcHumanoidEnvSpec,
   _DmcManipulatorEnvPool,
@@ -69,6 +71,11 @@
 DmcHumanoidEnvSpec, DmcHumanoidDMEnvPool, DmcHumanoidGymEnvPool = py_env(
   _DmcHumanoidEnvSpec, _DmcHumanoidEnvPool
 )
+(
+  DmcHumanoidCMUEnvSpec,
+  DmcHumanoidCMUDMEnvPool,
+  DmcHumanoidCMUGymEnvPool,
+) = py_env(_DmcHumanoidCMUEnvSpec, _DmcHumanoidCMUEnvPool)
 (
   DmcManipulatorEnvSpec,
   DmcManipulatorDMEnvPool,
@@ -115,6 +122,9 @@
   "DmcHumanoidEnvSpec",
   "DmcHumanoidDMEnvPool",
   "DmcHumanoidGymEnvPool",
+  "DmcHumanoidCMUEnvSpec",
+  "DmcHumanoidCMUDMEnvPool",
+  "DmcHumanoidCMUGymEnvPool",
   "DmcManipulatorEnvSpec",
   "DmcManipulatorDMEnvPool",
   "DmcManipulatorGymEnvPool",
diff --git a/envpool/mujoco/dmc/humanoid_CMU.h b/envpool/mujoco/dmc/humanoid_CMU.h
new file mode 100644
index 00000000..119f9065
--- /dev/null
+++ b/envpool/mujoco/dmc/humanoid_CMU.h
@@ -0,0 +1,298 @@
+/*
+ * 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.
+ */
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/humanoid_CMU.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_HUMANOID_CMU_H_
+#define ENVPOOL_MUJOCO_DMC_HUMANOID_CMU_H_
+
+#include <algorithm>
+#include <cmath>
+#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 GetHumanoidCMUXML(const std::string& base_path,
+                              const std::string& task_name) {
+  return GetFileContent(base_path, "humanoid_CMU.xml");
+}
+
+class HumanoidCMUEnvFns {
+ 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:joint_angles"_.Bind(Spec<mjtNum>({56})),
+                    "obs:head_height"_.Bind(Spec<mjtNum>({})),
+                    "obs:extremities"_.Bind(Spec<mjtNum>({12})),
+                    "obs:torso_vertical"_.Bind(Spec<mjtNum>({3})),
+                    "obs:com_velocity"_.Bind(Spec<mjtNum>({3})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({62})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({63})),
+#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, 56}, {-1.0, 1.0})));
+  }
+};
+
+using HumanoidCMUEnvSpec = EnvSpec<HumanoidCMUEnvFns>;
+
+class HumanoidCMUEnv : public Env<HumanoidCMUEnvSpec>, public MujocoEnv {
+ protected:
+  // Height of head above which stand reward is 1.
+  const mjtNum kStandHeight = 1.4;
+  // Horizontal speeds above which move reward is 1.
+  const mjtNum kRunSpeed = 10;
+  int id_head_;
+  int id_lhand_;
+  int id_lfoot_;
+  int id_rhand_;
+  int id_rfoot_;
+  int id_thorax_;
+  int id_thorax_subtreelinvel_;
+  mjtNum move_speed_;
+
+ public:
+  HumanoidCMUEnv(const Spec& spec, int env_id)
+      : Env<HumanoidCMUEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetHumanoidCMUXML(spec.config["base_path"_],
+                                    spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+        id_head_(mj_name2id(model_, mjOBJ_XBODY, "head")),
+        id_lhand_(mj_name2id(model_, mjOBJ_XBODY, "lhand")),
+        id_lfoot_(mj_name2id(model_, mjOBJ_XBODY, "lfoot")),
+        id_rhand_(mj_name2id(model_, mjOBJ_XBODY, "rhand")),
+        id_rfoot_(mj_name2id(model_, mjOBJ_XBODY, "rfoot")),
+        id_thorax_(mj_name2id(model_, mjOBJ_XBODY, "thorax")),
+        id_thorax_subtreelinvel_(GetSensorId(model_, "thorax_subtreelinvel")) {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name == "stand") {
+      move_speed_ = 0;
+    } else if (task_name == "run") {
+      move_speed_ = kRunSpeed;
+    } else {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc humanoid_CMU.");
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    while (true) {
+      // Find a collision-free random initial configuration.
+      // 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
+      PhysicsAfterReset();
+      if (data_->ncon <= 0) {
+        break;
+      }
+    }
+  }
+
+  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(HeadHeight(), kStandHeight,
+                                    std::numeric_limits<double>::infinity(),
+                                    kStandHeight / 4);
+    auto upright = RewardTolerance(ThoraxUpright(), 0.9,
+                                   std::numeric_limits<double>::infinity(), 1.9,
+                                   0.0, SigmoidType::kLinear);
+    auto stand_reward = standing * upright;
+    mjtNum small_control = 0.0;
+    for (int i = 0; i < model_->nu; ++i) {
+      small_control += RewardTolerance(data_->ctrl[i], 0.0, 0.0, 1.0, 0.0,
+                                       SigmoidType::kQuadratic);
+    }
+    small_control = (small_control / model_->nu + 4.0) / 5.0;
+    auto center_of_mass_velocity = CenterOfMassVelocity();
+    if (move_speed_ == 0) {
+      mjtNum dont_move =
+          (RewardTolerance(center_of_mass_velocity[0], 0.0, 0.0, 2.0) +
+           RewardTolerance(center_of_mass_velocity[1], 0.0, 0.0, 2.0)) /
+          2;
+      return static_cast<float>(small_control * stand_reward * dont_move);
+    }
+    auto com_velocity =
+        std::sqrt(center_of_mass_velocity[0] * center_of_mass_velocity[0] +
+                  center_of_mass_velocity[1] * center_of_mass_velocity[1]);
+    auto move = RewardTolerance(com_velocity, move_speed_,
+                                std::numeric_limits<double>::infinity(),
+                                move_speed_, 0.0, SigmoidType::kLinear);
+    move = (5.0 * move + 1.0) / 6.0;
+    return static_cast<float>(small_control * stand_reward * move);
+  }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    const auto& joint_angles = JointAngles();
+    const auto& extremities = Extremities();
+    const auto& com_velocity = CenterOfMassVelocity();
+    const auto& torso_vertical = TorsoVerticalOrientation();
+    state["obs:joint_angles"_].Assign(joint_angles.begin(),
+                                      joint_angles.size());
+    state["obs:head_height"_] = HeadHeight();
+    state["obs:extremities"_].Assign(extremities.begin(), extremities.size());
+    state["obs:torso_vertical"_].Assign(torso_vertical.begin(),
+                                        torso_vertical.size());
+    state["obs:com_velocity"_].Assign(com_velocity.begin(),
+                                      com_velocity.size());
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+    // info
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  mjtNum ThoraxUpright() {
+    // return self.named.data.xmat['thorax', 'zy']
+    return data_->xmat[id_thorax_ * 9 + 7];
+  }
+  mjtNum HeadHeight() {
+    // return self.named.data.xpos['head', 'z']
+    return data_->xpos[id_head_ * 3 + 2];
+  }
+  std::array<mjtNum, 3> CenterOfMassPosition() {
+    // return self.named.data.subtree_com['thorax']
+    return {
+        data_->subtree_com[id_thorax_ * 3],
+        data_->subtree_com[id_thorax_ * 3 + 1],
+        data_->subtree_com[id_thorax_ * 3 + 2],
+    };
+  }
+  std::array<mjtNum, 3> CenterOfMassVelocity() {
+    // return self.named.data.sensordata['thorax_subtreelinvel'].copy()
+    return {
+        data_->sensordata[id_thorax_subtreelinvel_ * 3],
+        data_->sensordata[id_thorax_subtreelinvel_ * 3 + 1],
+        data_->sensordata[id_thorax_subtreelinvel_ * 3 + 2],
+    };
+  }
+  std::array<mjtNum, 3> TorsoVerticalOrientation() {
+    // return self.named.data.xmat['thorax', ['zx', 'zy', 'zz']]
+    return {
+        data_->xmat[id_thorax_ * 9 + 6],
+        data_->xmat[id_thorax_ * 9 + 7],
+        data_->xmat[id_thorax_ * 9 + 8],
+    };
+  }
+  std::array<mjtNum, 56> JointAngles() {
+    // return self.data.qpos[7:].copy()
+    std::array<mjtNum, 56> joint_angles;
+    for (int i = 0; i < 56; i++) {
+      joint_angles[i] = data_->qpos[7 + i];
+    }
+    return joint_angles;
+  }
+  std::array<mjtNum, 12> Extremities() {
+    // torso_frame = self.named.data.xmat['thorax'].reshape(3, 3)
+    // torso_pos = self.named.data.xpos['thorax']
+    // positions = []
+    // for side in ('l', 'r'):
+    //   for limb in ('hand', 'foot'):
+    //     torso_to_limb = self.named.data.xpos[side + limb] - torso_pos
+    //     positions.append(torso_to_limb.dot(torso_frame))
+    // return np.hstack(positions)
+    std::array<mjtNum, 9> torso_frame;
+    for (int i = 0; i < 9; i++) {
+      torso_frame[i] = data_->xmat[id_thorax_ * 9 + i];
+    }
+    std::array<mjtNum, 3> torso_pos;
+    for (int i = 0; i < 3; i++) {
+      torso_pos[i] = data_->xpos[id_thorax_ * 3 + i];
+    }
+    // left hand
+    std::array<mjtNum, 3> torso_to_limb_lh;
+    for (int i = 0; i < 3; i++) {
+      torso_to_limb_lh[i] =
+          (data_->xpos[id_lhand_ * 3] - torso_pos[0]) * torso_frame[i] +
+          (data_->xpos[id_lhand_ * 3 + 1] - torso_pos[1]) * torso_frame[i + 3] +
+          (data_->xpos[id_lhand_ * 3 + 2] - torso_pos[2]) * torso_frame[i + 6];
+    }
+    // left foot
+    std::array<mjtNum, 3> torso_to_limb_lf;
+    for (int i = 0; i < 3; i++) {
+      torso_to_limb_lf[i] =
+          (data_->xpos[id_lfoot_ * 3] - torso_pos[0]) * torso_frame[i] +
+          (data_->xpos[id_lfoot_ * 3 + 1] - torso_pos[1]) * torso_frame[i + 3] +
+          (data_->xpos[id_lfoot_ * 3 + 2] - torso_pos[2]) * torso_frame[i + 6];
+    }
+    // right hand
+    std::array<mjtNum, 3> torso_to_limb_rh;
+    for (int i = 0; i < 3; i++) {
+      torso_to_limb_rh[i] =
+          (data_->xpos[id_rhand_ * 3] - torso_pos[0]) * torso_frame[i] +
+          (data_->xpos[id_rhand_ * 3 + 1] - torso_pos[1]) * torso_frame[i + 3] +
+          (data_->xpos[id_rhand_ * 3 + 2] - torso_pos[2]) * torso_frame[i + 6];
+    }
+    // right foot
+    std::array<mjtNum, 3> torso_to_limb_rf;
+    for (int i = 0; i < 3; i++) {
+      torso_to_limb_rf[i] =
+          (data_->xpos[id_rfoot_ * 3] - torso_pos[0]) * torso_frame[i] +
+          (data_->xpos[id_rfoot_ * 3 + 1] - torso_pos[1]) * torso_frame[i + 3] +
+          (data_->xpos[id_rfoot_ * 3 + 2] - torso_pos[2]) * torso_frame[i + 6];
+    }
+    return {
+        torso_to_limb_lh[0], torso_to_limb_lh[1], torso_to_limb_lh[2],
+        torso_to_limb_lf[0], torso_to_limb_lf[1], torso_to_limb_lf[2],
+        torso_to_limb_rh[0], torso_to_limb_rh[1], torso_to_limb_rh[2],
+        torso_to_limb_rf[0], torso_to_limb_rf[1], torso_to_limb_rf[2],
+    };
+  }
+};
+
+using HumanoidCMUEnvPool = AsyncEnvPool<HumanoidCMUEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_HUMANOID_CMU_H_
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index c9ab7e7e..34f98fe1 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -21,6 +21,7 @@
 #include "envpool/mujoco/dmc/fish.h"
 #include "envpool/mujoco/dmc/hopper.h"
 #include "envpool/mujoco/dmc/humanoid.h"
+#include "envpool/mujoco/dmc/humanoid_CMU.h"
 #include "envpool/mujoco/dmc/manipulator.h"
 #include "envpool/mujoco/dmc/pendulum.h"
 #include "envpool/mujoco/dmc/point_mass.h"
@@ -52,6 +53,9 @@ using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 using DmcHumanoidEnvSpec = PyEnvSpec<mujoco_dmc::HumanoidEnvSpec>;
 using DmcHumanoidEnvPool = PyEnvPool<mujoco_dmc::HumanoidEnvPool>;
 
+using DmcHumanoidCMUEnvSpec = PyEnvSpec<mujoco_dmc::HumanoidCMUEnvSpec>;
+using DmcHumanoidCMUEnvPool = PyEnvPool<mujoco_dmc::HumanoidCMUEnvPool>;
+
 using DmcManipulatorEnvSpec = PyEnvSpec<mujoco_dmc::ManipulatorEnvSpec>;
 using DmcManipulatorEnvPool = PyEnvPool<mujoco_dmc::ManipulatorEnvPool>;
 
@@ -79,6 +83,7 @@ PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcFishEnvSpec, DmcFishEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
   REGISTER(m, DmcHumanoidEnvSpec, DmcHumanoidEnvPool)
+  REGISTER(m, DmcHumanoidCMUEnvSpec, DmcHumanoidCMUEnvPool)
   REGISTER(m, DmcManipulatorEnvSpec, DmcManipulatorEnvPool)
   REGISTER(m, DmcPendulumEnvSpec, DmcPendulumEnvPool)
   REGISTER(m, DmcPointMassEnvSpec, DmcPointMassEnvPool)
diff --git a/envpool/mujoco/dmc/registration.py b/envpool/mujoco/dmc/registration.py
index 7d3462e8..0b4bc826 100644
--- a/envpool/mujoco/dmc/registration.py
+++ b/envpool/mujoco/dmc/registration.py
@@ -44,6 +44,8 @@
   ("humanoid", "run_pure_state"),
   ("humanoid", "stand"),
   ("humanoid", "walk"),
+  ("humanoid_CMU", "run"),
+  ("humanoid_CMU", "stand"),
   ("manipulator", "bring_ball"),
   ("manipulator", "bring_peg"),
   ("manipulator", "insert_ball"),
@@ -61,8 +63,8 @@
 ]
 
 for domain, task in dmc_mujoco_envs:
-  domain_name = "".join([g.capitalize() for g in domain.split("_")])
-  task_name = "".join([g.capitalize() for g in task.split("_")])
+  domain_name = "".join([g[:1].upper() + g[1:] for g in domain.split("_")])
+  task_name = "".join([g[:1].upper() + g[1:] for g in task.split("_")])
   register(
     task_id=f"{domain_name}{task_name}-v1",
     import_path="envpool.mujoco.dmc",


Test patch:
diff --git a/benchmark/test_dmc.py b/benchmark/test_dmc.py
index 01914bfe..c3ddbea9 100644
--- a/benchmark/test_dmc.py
+++ b/benchmark/test_dmc.py
@@ -48,7 +48,7 @@ def run_envpool(env, action, frame_skip, total_step):
   return fps
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--domain", type=str, default="cheetah")
   parser.add_argument("--task", type=str, default="run")
diff --git a/benchmark/test_gym.py b/benchmark/test_gym.py
index dc8051a7..8a2afa98 100644
--- a/benchmark/test_gym.py
+++ b/benchmark/test_gym.py
@@ -70,7 +70,7 @@ def run(env, num_envs, total_step, async_):
   print(f"FPS = {frame_skip * total_step * num_envs / (time.time() - t):.2f}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument(
     "--env", type=str, default="atari", choices=["atari", "mujoco", "box2d"]
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 7ac7f163..9225bb9d 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -151,6 +151,8 @@ def test_make_mujoco_dmc(self) -> None:
         "HumanoidRunPureState-v1",
         "HumanoidStand-v1",
         "HumanoidWalk-v1",
+        "HumanoidCMURun-v1",
+        "HumanoidCMUStand-v1",
         "ManipulatorBringBall-v1",
         "ManipulatorBringPeg-v1",
         "ManipulatorInsertBall-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_suite_align_test.py
similarity index 100%
rename from envpool/mujoco/dmc/mujoco_dmc_align_test.py
rename to envpool/mujoco/dmc/mujoco_dmc_suite_align_test.py
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_suite_deterministic_test.py
similarity index 100%
rename from envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
rename to envpool/mujoco/dmc/mujoco_dmc_suite_deterministic_test.py
diff --git a/envpool/mujoco/dmc/mujoco_dmc_suite_ext_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_suite_ext_align_test.py
new file mode 100644
index 00000000..ed078fdf
--- /dev/null
+++ b/envpool/mujoco/dmc/mujoco_dmc_suite_ext_align_test.py
@@ -0,0 +1,103 @@
+# 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.
+"""Unit tests for Mujoco dm_control suite align check."""
+
+from typing import Any, List, no_type_check
+
+import dm_env
+import numpy as np
+from absl import logging
+from absl.testing import absltest
+from dm_control import suite
+
+from envpool.mujoco.dmc import DmcHumanoidCMUDMEnvPool, DmcHumanoidCMUEnvSpec
+
+
+class _MujocoDmcSuiteExtAlignTest(absltest.TestCase):
+
+  @no_type_check
+  def run_space_check(self, env0: dm_env.Environment, env1: Any) -> None:
+    """Check observation_spec() and action_spec()."""
+    obs0, obs1 = env0.observation_spec(), env1.observation_spec()
+    for k in obs0:
+      self.assertTrue(hasattr(obs1, k))
+      np.testing.assert_allclose(obs0[k].shape, getattr(obs1, k).shape)
+    act0, act1 = env0.action_spec(), env1.action_spec()
+    np.testing.assert_allclose(act0.shape, act1.shape)
+    np.testing.assert_allclose(act0.minimum, act1.minimum)
+    np.testing.assert_allclose(act0.maximum, act1.maximum)
+
+  @no_type_check
+  def reset_state(
+    self, env: dm_env.Environment, ts: dm_env.TimeStep, domain: str, task: str
+  ) -> None:
+    # manually reset, mimic initialize_episode
+    with env.physics.reset_context():
+      env.physics.data.qpos = ts.observation.qpos0[0]
+      if domain in ["humanoid_CMU"]:
+        env.physics.after_reset()
+
+  def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
+    return np.random.uniform(
+      low=action_spec.minimum,
+      high=action_spec.maximum,
+      size=action_spec.shape,
+    )
+
+  def run_align_check(
+    self, env0: dm_env.Environment, env1: Any, domain: str, task: str
+  ) -> None:
+    logging.info(f"align check for {domain} {task}")
+    obs_spec, action_spec = env0.observation_spec(), env0.action_spec()
+    for i in range(3):
+      np.random.seed(i)
+      env0.reset()
+      a = self.sample_action(action_spec)
+      ts = env1.reset(np.array([0]))
+      self.reset_state(env0, ts, domain, task)
+      logging.info(f"reset qpos {ts.observation.qpos0[0]}")
+      cnt = 0
+      done = False
+      while not done:
+        cnt += 1
+        a = self.sample_action(action_spec)
+        # logging.info(f"{cnt} {a}")
+        ts0 = env0.step(a)
+        ts1 = env1.step(np.array([a]), np.array([0]))
+        done = ts0.step_type == dm_env.StepType.LAST
+        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], atol=1e-8)
+        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, task)
+
+  def test_humanoid_CMU(self) -> None:
+    self.run_align_check_entry(
+      "humanoid_CMU", ["stand", "run"], DmcHumanoidCMUEnvSpec,
+      DmcHumanoidCMUDMEnvPool
+    )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/envpool/mujoco/dmc/mujoco_dmc_suite_ext_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_suite_ext_deterministic_test.py
new file mode 100644
index 00000000..c410e40a
--- /dev/null
+++ b/envpool/mujoco/dmc/mujoco_dmc_suite_ext_deterministic_test.py
@@ -0,0 +1,81 @@
+# 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.
+"""Unit tests for Mujoco dm_control deterministic check."""
+
+from typing import Any, List, Optional
+
+import dm_env
+import numpy as np
+from absl.testing import absltest
+
+from envpool.mujoco.dmc import DmcHumanoidCMUDMEnvPool, DmcHumanoidCMUEnvSpec
+
+
+class _MujocoDmcSuiteExtDeterministicTest(absltest.TestCase):
+
+  def check(
+    self,
+    spec_cls: Any,
+    envpool_cls: Any,
+    task: str,
+    obs_keys: List[str],
+    blacklist: Optional[List[str]] = None,
+    num_envs: int = 2,
+  ) -> None:
+    np.random.seed(0)
+    env0 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=0, task_name=task))
+    )
+    env1 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=0, task_name=task))
+    )
+    env2 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=1, task_name=task))
+    )
+    act_spec = env0.action_spec()
+    for t in range(3000):
+      action = np.array(
+        [
+          np.random.uniform(
+            low=act_spec.minimum, high=act_spec.maximum, size=act_spec.shape
+          ) for _ in range(num_envs)
+        ]
+      )
+      ts0 = env0.step(action)
+      obs0 = ts0.observation
+      obs1 = env1.step(action).observation
+      obs2 = env2.step(action).observation
+      for k in obs_keys:
+        o0 = getattr(obs0, k)
+        o1 = getattr(obs1, k)
+        o2 = getattr(obs2, k)
+        np.testing.assert_allclose(o0, o1)
+        if blacklist and k in blacklist:
+          continue
+        if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
+          self.assertFalse(np.allclose(o0, o2), (t, k, o0, o2))
+
+  def test_humanoid_CMU(self) -> None:
+    obs_keys = [
+      "joint_angles", "head_height", "extremities", "torso_vertical",
+      "com_velocity", "velocity"
+    ]
+    for task in ["stand", "run"]:
+      self.check(
+        DmcHumanoidCMUEnvSpec, DmcHumanoidCMUDMEnvPool, task, obs_keys
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()