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 583b273b..2902cbf1 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -92,6 +92,20 @@ HopperStand-v1, HopperHop-v1
 - ``max_episode_steps``: 1000;
 
 
+ManipulatorBringBall-v1, ManipulatorBringPeg-v1, ManipulatorInsertBall-v1, ManipulatorInsertPeg-v1
+--------------------------------------------------------------------------------------------------
+
+`dm_control suite manipulator source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/manipulator.py>`_
+
+- Observation spec: a namedtuple with three keys: ``arm_pos (8,2)``,
+  ``arm_vel (8)``, ``touch (5)``, ``hand_pos (4)``, ``object_pos (4)``,
+  ``object_vel (3)``, ``target_pos (4)``;
+- Action spec: ``(5)``, with range ``[-1, 1]``;
+- ``frame_skip``: 10;
+- ``max_episode_steps``: 1000;
+
+
 HumanoidStand-v1, HumanoidWalk-v1, HumanoidRun-v1, HumanoidRunPureState-v1
 --------------------------------------------------------------------------
 
@@ -105,14 +119,13 @@ HumanoidStand-v1, HumanoidWalk-v1, HumanoidRun-v1, HumanoidRunPureState-v1
 - ``frame_skip``: 5;
 - ``max_episode_steps``: 1000;
 
-
 .. note ::
 
     The observation keys ``joint_angles``, ``head_height``, ``extremities``,
-    ``torso_vertical`` and ``com_velocity`` are only
-    available in ``HumanoidStand-v1``, ``HumanoidWalk-v1`` and ``HumanoidRun-v1``.
-    The observation keys ``position`` are only
-    available in ``HumanoidRunPureState-v1`` tasks.
+    ``torso_vertical`` and ``com_velocity`` are only available in
+    ``HumanoidStand-v1``, ``HumanoidWalk-v1`` and ``HumanoidRun-v1``.
+    The observation keys ``position`` are only available in
+    ``HumanoidRunPureState-v1`` tasks.
 
 
 PendulumSwingup-v1
diff --git a/docs/index.rst b/docs/index.rst
index 8fc8fe89..e0028405 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -7,9 +7,9 @@ Welcome to EnvPool!
 ===================
 
 **EnvPool** is a C++-based batched environment pool with pybind11 and thread
-pool. It has high performance (~1M raw FPS in DGX on Atari games) and
-compatible APIs (supports both gym and dm_env, both sync and async, both single
-and multi player environment).
+pool. It has high performance (~1M raw FPS on Atari games / ~3M FPS with Mujoco
+physics engine in DGX-A100) and compatible APIs (supports both gym and dm_env,
+both sync and async, both single and multi player environment).
 
 Here are EnvPool's several highlights:
 
@@ -18,8 +18,9 @@ Here are EnvPool's several highlights:
 - Support both synchronous execution and asynchronous execution;
 - Support both single player and multi-player environment;
 - Easy C++ developer API to add new envs;
-- **1 Million** Atari frames per second simulation with 256 CPU cores,
-  **~13x** throughput of Python subprocess-based vector env;
+- **1 Million** Atari frames / **3 Million** Mujoco steps per second
+  simulation with 256 CPU cores, **~20x** throughput of Python
+  subprocess-based vector env;
 - **~3x** throughput of Python subprocess-based vector env on low resource
   setup like 12 CPU cores;
 - Comparing with the existing GPU-based solution
@@ -27,8 +28,10 @@ Here are EnvPool's several highlights:
   `Isaac-gym <https://developer.nvidia.com/isaac-gym>`_), EnvPool is a
   **general** solution for various kinds of speeding-up RL environment
   parallelization;
-- Compatible with some existing RL libraries, e.g.,
-  `Tianshou <https://github.com/thu-ml/tianshou>`_.
+  `Stable-Baselines3 <https://github.com/DLR-RM/stable-baselines3>`_,
+  `Tianshou <https://github.com/thu-ml/tianshou>`_,
+  or `CleanRL <https://github.com/vwxyzjn/cleanrl>`_
+  (`Solving Pong in 5 mins <https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/#solving-pong-in-5-minutes-with-ppo--envpool>`_).
 
 Installation
 ------------
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 7cca54a9..3be423cd 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -62,6 +62,10 @@ pybind_extension(
 
 cc_library(
     name = "mujoco_dmc_env",
+    srcs = [
+        "dmc/mujoco_env.cc",
+        "dmc/utils.cc",
+    ],
     hdrs = [
         "dmc/acrobot.h",
         "dmc/ball_in_cup.h",
@@ -69,6 +73,7 @@ cc_library(
         "dmc/finger.h",
         "dmc/hopper.h",
         "dmc/humanoid.h",
+        "dmc/manipulator.h",
         "dmc/mujoco_env.h",
         "dmc/pendulum.h",
         "dmc/point_mass.h",
@@ -80,6 +85,7 @@ cc_library(
     deps = [
         "//envpool/core:async_envpool",
         "@mujoco//:mujoco_lib",
+        "@pugixml",
     ],
 )
 
diff --git a/envpool/mujoco/dmc/__init__.py b/envpool/mujoco/dmc/__init__.py
index f3e0f60b..f86a3fea 100644
--- a/envpool/mujoco/dmc/__init__.py
+++ b/envpool/mujoco/dmc/__init__.py
@@ -26,6 +26,8 @@
   _DmcHopperEnvSpec,
   _DmcHumanoidEnvPool,
   _DmcHumanoidEnvSpec,
+  _DmcManipulatorEnvPool,
+  _DmcManipulatorEnvSpec,
   _DmcPendulumEnvPool,
   _DmcPendulumEnvSpec,
   _DmcPointMassEnvPool,
@@ -55,6 +57,11 @@
 DmcHumanoidEnvSpec, DmcHumanoidDMEnvPool, DmcHumanoidGymEnvPool = py_env(
   _DmcHumanoidEnvSpec, _DmcHumanoidEnvPool
 )
+(
+  DmcManipulatorEnvSpec,
+  DmcManipulatorDMEnvPool,
+  DmcManipulatorGymEnvPool,
+) = py_env(_DmcManipulatorEnvSpec, _DmcManipulatorEnvPool)
 DmcPendulumEnvSpec, DmcPendulumDMEnvPool, DmcPendulumGymEnvPool = py_env(
   _DmcPendulumEnvSpec, _DmcPendulumEnvPool
 )
@@ -84,6 +91,9 @@
   "DmcHumanoidEnvSpec",
   "DmcHumanoidDMEnvPool",
   "DmcHumanoidGymEnvPool",
+  "DmcManipulatorEnvSpec",
+  "DmcManipulatorDMEnvPool",
+  "DmcManipulatorGymEnvPool",
   "DmcPendulumEnvSpec",
   "DmcPendulumDMEnvPool",
   "DmcPendulumGymEnvPool",
diff --git a/envpool/mujoco/dmc/cheetah.h b/envpool/mujoco/dmc/cheetah.h
index dec4b61e..d4efd128 100644
--- a/envpool/mujoco/dmc/cheetah.h
+++ b/envpool/mujoco/dmc/cheetah.h
@@ -82,14 +82,14 @@ class CheetahEnv : public Env<CheetahEnvSpec>, public MujocoEnv {
   }
 
   void TaskInitializeEpisode() override {
-    assert(model_->njnt == model_->nq);
-    for (int i = 0; i < model_->njnt; i++) {
-      bool is_limited = model_->jnt_limited[i] == 1;
+    for (int id_joint = 0; id_joint < model_->njnt; ++id_joint) {
+      bool is_limited = model_->jnt_limited[id_joint] == 1;
       if (is_limited) {
-        mjtNum range_min = model_->jnt_range[i * 2 + 0];
-        mjtNum range_max = model_->jnt_range[i * 2 + 1];
+        mjtNum range_min = model_->jnt_range[id_joint * 2 + 0];
+        mjtNum range_max = model_->jnt_range[id_joint * 2 + 1];
         mjtNum range = range_max - range_min;
-        data_->qpos[i] = dist_uniform_(gen_) * range + range_min;
+        data_->qpos[model_->jnt_qposadr[id_joint]] =
+            dist_uniform_(gen_) * range + range_min;
       }
     }
 #ifdef ENVPOOL_TEST
diff --git a/envpool/mujoco/dmc/finger.h b/envpool/mujoco/dmc/finger.h
index 7476ac79..d250ffef 100644
--- a/envpool/mujoco/dmc/finger.h
+++ b/envpool/mujoco/dmc/finger.h
@@ -33,7 +33,7 @@
 namespace mujoco_dmc {
 
 std::string GetFingerXML(const std::string& base_path,
-                         const std::string& task_name_) {
+                         const std::string& task_name) {
   return GetFileContent(base_path, "finger.xml");
 }
 
diff --git a/envpool/mujoco/dmc/humanoid.h b/envpool/mujoco/dmc/humanoid.h
index 4b8cdf54..6b2e1658 100644
--- a/envpool/mujoco/dmc/humanoid.h
+++ b/envpool/mujoco/dmc/humanoid.h
@@ -33,7 +33,7 @@
 namespace mujoco_dmc {
 
 std::string GetHumanoidXML(const std::string& base_path,
-                           const std::string& task_name_) {
+                           const std::string& task_name) {
   return GetFileContent(base_path, "humanoid.xml");
 }
 
diff --git a/envpool/mujoco/dmc/manipulator.h b/envpool/mujoco/dmc/manipulator.h
new file mode 100644
index 00000000..703c6b46
--- /dev/null
+++ b/envpool/mujoco/dmc/manipulator.h
@@ -0,0 +1,389 @@
+/*
+ * 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/manipulator.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_MANIPULATOR_H_
+#define ENVPOOL_MUJOCO_DMC_MANIPULATOR_H_
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <random>
+#include <regex>
+#include <set>
+#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 GetManipulatorXML(const std::string& base_path,
+                              const std::string& task_name) {
+  auto content = GetFileContent(base_path, "manipulator.xml");
+  if (task_name == "bring_ball") {
+    return XMLRemoveByBodyName(content, {"slot", "target_peg", "cup", "peg"});
+  }
+  if (task_name == "bring_peg") {
+    return XMLRemoveByBodyName(content, {"slot", "target_ball", "cup", "ball"});
+  }
+  if (task_name == "insert_ball") {
+    return XMLRemoveByBodyName(content, {"slot", "target_peg", "peg"});
+  }
+  if (task_name == "insert_peg") {
+    return XMLRemoveByBodyName(content, {"target_ball", "cup", "ball"});
+  }
+  throw std::runtime_error("Unknown task_name " + task_name +
+                           " for dmc manipulator.");
+}
+
+class ManipulatorEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(10),
+                    "task_name"_.Bind(std::string("bring_ball")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:arm_pos"_.Bind(Spec<mjtNum>({8, 2})),
+                    "obs:arm_vel"_.Bind(Spec<mjtNum>({8})),
+                    "obs:touch"_.Bind(Spec<mjtNum>({5})),
+                    "obs:hand_pos"_.Bind(Spec<mjtNum>({4})),
+                    "obs:object_pos"_.Bind(Spec<mjtNum>({4})),
+                    "obs:object_vel"_.Bind(Spec<mjtNum>({3})),
+                    "obs:target_pos"_.Bind(Spec<mjtNum>({4})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({11})),
+                    "info:random_info"_.Bind(Spec<mjtNum>({8})),
+#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, 5}, {-1.0, 1.0})));
+  }
+};
+
+using ManipulatorEnvSpec = EnvSpec<ManipulatorEnvFns>;
+
+class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
+ protected:
+  const mjtNum kClose = 0.01;
+  const mjtNum kPInHand = 0.1;
+  const mjtNum kPInTarget = 0.1;
+  const std::array<std::string, 8> kArmJoints = {
+      "arm_root", "arm_shoulder", "arm_elbow", "arm_wrist",
+      "finger",   "fingertip",    "thumb",     "thumbtip"};
+  const std::array<std::string, 6> kAllProps = {"ball", "target_ball", "cup",
+                                                "peg",  "target_peg",  "slot"};
+  const std::array<std::string, 5> kTouchSensors = {
+      "palm_touch", "finger_touch", "thumb_touch", "fingertip_touch",
+      "thumbtip_touch"};
+
+  bool use_peg_, insert_;
+  std::uniform_real_distribution<> dist_uniform_;
+  std::array<mjtNum, 8> random_info_;
+  // target_x, target_z, target_angle, init_type, object_x, object_z,
+  // object_angle, qvel_objx
+
+  // ids
+  std::array<int, 8> id_arm_joints_, id_arm_qpos_;
+  int id_finger_, id_thumb_, id_body_receptacle_, id_body_target_;
+  int id_xbody_hand_, id_xbody_object_, id_xbody_target_, id_object_x_;
+  std::array<int, 3> id_object_joints_;
+  std::array<int, 5> id_touch_sensors_;
+  int id_site_peg_grasp_, id_site_grasp_, id_site_peg_pinch_, id_site_pinch_;
+  int id_site_peg_, id_site_target_peg_, id_site_target_peg_tip_;
+  int id_site_peg_tip_, id_site_ball_, id_site_target_ball_;
+
+ public:
+  ManipulatorEnv(const Spec& spec, int env_id)
+      : Env<ManipulatorEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetManipulatorXML(spec.config["base_path"_],
+                                    spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+        use_peg_(spec.config["task_name"_] == "bring_peg" ||
+                 spec.config["task_name"_] == "insert_peg"),
+        insert_(spec.config["task_name"_] == "insert_peg" ||
+                spec.config["task_name"_] == "insert_ball"),
+        dist_uniform_(0, 1),
+        id_finger_(GetQposId(model_, "finger")),
+        id_thumb_(GetQposId(model_, "thumb")),
+        id_body_receptacle_(
+            mj_name2id(model_, mjOBJ_BODY, use_peg_ ? "slot" : "cup")),
+        id_body_target_(mj_name2id(model_, mjOBJ_BODY,
+                                   use_peg_ ? "target_peg" : "target_ball")),
+        id_xbody_hand_(mj_name2id(model_, mjOBJ_XBODY, "hand")),
+        id_xbody_object_(
+            mj_name2id(model_, mjOBJ_XBODY, use_peg_ ? "peg" : "ball")),
+        id_xbody_target_(mj_name2id(model_, mjOBJ_XBODY,
+                                    use_peg_ ? "target_peg" : "target_ball")),
+        id_object_x_(GetQposId(model_, use_peg_ ? "peg_x" : "ball_x")),
+        id_site_peg_grasp_(mj_name2id(model_, mjOBJ_SITE, "peg_grasp")),
+        id_site_grasp_(mj_name2id(model_, mjOBJ_SITE, "grasp")),
+        id_site_peg_pinch_(mj_name2id(model_, mjOBJ_SITE, "peg_pinch")),
+        id_site_pinch_(mj_name2id(model_, mjOBJ_SITE, "pinch")),
+        id_site_peg_(mj_name2id(model_, mjOBJ_SITE, "peg")),
+        id_site_target_peg_(mj_name2id(model_, mjOBJ_SITE, "target_peg")),
+        id_site_target_peg_tip_(
+            mj_name2id(model_, mjOBJ_SITE, "target_peg_tip")),
+        id_site_peg_tip_(mj_name2id(model_, mjOBJ_SITE, "peg_tip")),
+        id_site_ball_(mj_name2id(model_, mjOBJ_SITE, "ball")),
+        id_site_target_ball_(mj_name2id(model_, mjOBJ_SITE, "target_ball")) {
+    for (std::size_t i = 0; i < kArmJoints.size(); ++i) {
+      id_arm_joints_[i] =
+          mj_name2id(model_, mjOBJ_JOINT, kArmJoints[i].c_str());
+      id_arm_qpos_[i] = GetQposId(model_, kArmJoints[i]);
+    }
+    std::array<std::string, 3> object_joints;
+    if (use_peg_) {
+      object_joints = {"peg_x", "peg_z", "peg_y"};
+    } else {
+      object_joints = {"ball_x", "ball_z", "ball_y"};
+    }
+    for (std::size_t i = 0; i < object_joints.size(); ++i) {
+      id_object_joints_[i] = GetQposId(model_, object_joints[i]);
+    }
+    for (std::size_t i = 0; i < kTouchSensors.size(); ++i) {
+      id_touch_sensors_[i] = GetSensorId(model_, kTouchSensors[i]);
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    bool penetrating = true;
+    while (penetrating) {
+      for (std::size_t i = 0; i < kArmJoints.size(); ++i) {
+        int id_joint = id_arm_joints_[i];
+        bool is_limited = model_->jnt_limited[id_joint] == 1 ? true : false;
+        mjtNum lower = is_limited ? model_->jnt_range[id_joint * 2 + 0] : -M_PI;
+        mjtNum upper = is_limited ? model_->jnt_range[id_joint * 2 + 1] : M_PI;
+        data_->qpos[id_arm_qpos_[i]] =
+            dist_uniform_(gen_) * (upper - lower) + lower;
+      }
+      data_->qpos[id_finger_] = data_->qpos[id_thumb_];
+#ifdef ENVPOOL_TEST
+      std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#endif
+      mjtNum target_x = random_info_[0] = dist_uniform_(gen_) * 0.8 - 0.4;
+      mjtNum target_z = dist_uniform_(gen_) * 0.3 + 0.1;
+      mjtNum target_angle;
+      if (insert_) {
+        target_angle = dist_uniform_(gen_) * (M_PI * 2 / 3) - M_PI / 3;
+        // model.body_pos[self._receptacle, ['x', 'z']]
+        model_->body_pos[id_body_receptacle_ * 3 + 0] = target_x;
+        model_->body_pos[id_body_receptacle_ * 3 + 2] = target_z;
+        // model.body_quat[self._receptacle, ['qw', 'qy']]
+        model_->body_quat[id_body_receptacle_ * 4 + 0] =
+            std::cos(target_angle / 2);
+        model_->body_quat[id_body_receptacle_ * 4 + 2] =
+            std::sin(target_angle / 2);
+      } else {
+        target_angle = dist_uniform_(gen_) * 2 * M_PI - M_PI;
+      }
+      random_info_[0] = target_x;
+      random_info_[1] = target_z;
+      random_info_[2] = target_angle;
+
+      // model.body_pos[self._target, ['x', 'z']] = target_x, target_z
+      // model.body_quat[self._target, ['qw', 'qy']] = [
+      //     np.cos(target_angle/2), np.sin(target_angle/2)]
+      model_->body_pos[id_body_target_ * 3 + 0] = target_x;
+      model_->body_pos[id_body_target_ * 3 + 2] = target_z;
+      model_->body_quat[id_body_target_ * 4 + 0] = std::cos(target_angle / 2);
+      model_->body_quat[id_body_target_ * 4 + 2] = std::sin(target_angle / 2);
+
+      mjtNum choice = dist_uniform_(gen_);
+      mjtNum object_x;
+      mjtNum object_z;
+      mjtNum object_angle;
+      if (choice <= kPInTarget) {
+        // in_target
+        random_info_[3] = 1;
+        object_x = target_x;
+        object_z = target_z;
+        object_angle = target_angle;
+      } else if (choice <= kPInTarget + kPInHand) {
+        // in_hand
+        random_info_[3] = 2;
+        // physics.after_reset()
+        // object_x = data.site_xpos['grasp', 'x']
+        // object_z = data.site_xpos['grasp', 'z']
+        // grasp_direction = data.site_xmat['grasp', ['xx', 'zx']]
+        // object_angle = np.pi-np.arctan2(grasp_direction[1],
+        // grasp_direction[0])
+        PhysicsAfterReset();
+        object_x = data_->site_xpos[id_site_grasp_ * 3 + 0];
+        object_z = data_->site_xpos[id_site_grasp_ * 3 + 2];
+        std::array<mjtNum, 2> grasp_direction = {
+            data_->site_xmat[id_site_grasp_ * 9 + 0],
+            data_->site_xmat[id_site_grasp_ * 9 + 6]};
+        object_angle =
+            M_PI - std::atan2(grasp_direction[1], grasp_direction[0]);
+      } else {
+        // uniform
+        random_info_[3] = 3;
+        // object_x = uniform(-.5, .5)
+        // object_z = uniform(0, .7)
+        // object_angle = uniform(0, 2*np.pi)
+        // data.qvel[self._object + '_x'] = uniform(-5, 5)
+        object_x = dist_uniform_(gen_) * 1 - 0.5;
+        object_z = dist_uniform_(gen_) * 0.7;
+        object_angle = dist_uniform_(gen_) * 2 * M_PI;
+        data_->qvel[id_object_x_] = random_info_[7] =
+            dist_uniform_(gen_) * 10 - 5;
+      }
+      data_->qpos[id_object_joints_[0]] = random_info_[4] = object_x;
+      data_->qpos[id_object_joints_[1]] = random_info_[5] = object_z;
+      data_->qpos[id_object_joints_[2]] = random_info_[6] = object_angle;
+      // Check for collisions.
+      PhysicsAfterReset();
+      penetrating = data_->ncon > 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 {
+    if (use_peg_) {
+      return static_cast<float>(PegReward());
+    }
+    return static_cast<float>(BallReward());
+  }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  std::array<mjtNum, 16> BoundedJointPos() {
+    std::array<mjtNum, 16> bound;
+    for (std::size_t i = 0; i < id_arm_qpos_.size(); i++) {
+      bound[i * 2 + 0] = std::sin(data_->qpos[id_arm_qpos_[i]]);
+      bound[i * 2 + 1] = std::cos(data_->qpos[id_arm_qpos_[i]]);
+    }
+    return bound;
+  }
+
+  // bug :(
+  std::array<mjtNum, 8> JointVelArm() {
+    std::array<mjtNum, 8> joint;
+    for (std::size_t i = 0; i < id_arm_qpos_.size(); i++) {
+      joint[i] = data_->qvel[id_arm_qpos_[i]];
+    }
+    return joint;
+  }
+
+  std::array<mjtNum, 3> JointVelObj() {
+    std::array<mjtNum, 3> joint;
+    for (std::size_t i = 0; i < id_object_joints_.size(); i++) {
+      joint[i] = data_->qvel[id_object_joints_[i]];
+    }
+    return joint;
+  }
+
+  std::array<mjtNum, 5> Touch() {
+    std::array<mjtNum, 5> touch;
+    for (std::size_t i = 0; i < id_touch_sensors_.size(); i++) {
+      touch[i] = std::log1p(data_->sensordata[id_touch_sensors_[i]]);
+    }
+    return touch;
+  }
+
+  std::array<mjtNum, 4> Body2dPose(int id) {
+    // self.named.data.xpos[body_names, ['x', 'z']]
+    // self.named.data.xquat[body_names, ['qw', 'qy']]
+    return {data_->xpos[id * 3 + 0], data_->xpos[id * 3 + 2],
+            data_->xquat[id * 4 + 0], data_->xquat[id * 4 + 2]};
+  }
+
+  mjtNum SiteDistance(int id_site1, int id_site2) {
+    std::array<mjtNum, 3> diff = {
+        data_->site_xpos[id_site1 * 3 + 0] - data_->site_xpos[id_site2 * 3 + 0],
+        data_->site_xpos[id_site1 * 3 + 1] - data_->site_xpos[id_site2 * 3 + 1],
+        data_->site_xpos[id_site1 * 3 + 2] -
+            data_->site_xpos[id_site2 * 3 + 2]};
+    return std::sqrt(diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]);
+  }
+
+  [[nodiscard]] mjtNum IsClose(mjtNum distance) const {
+    // return rewards.tolerance(distance, (0, _CLOSE), _CLOSE * 2)
+    return RewardTolerance(distance, 0, kClose, kClose * 2);
+  }
+
+  mjtNum PegReward() {
+    auto grasping =
+        (IsClose(SiteDistance(id_site_peg_grasp_, id_site_grasp_)) +
+         IsClose(SiteDistance(id_site_peg_pinch_, id_site_pinch_))) /
+        2;
+    auto bringing =
+        (IsClose(SiteDistance(id_site_peg_, id_site_target_peg_)) +
+         IsClose(SiteDistance(id_site_target_peg_tip_, id_site_peg_tip_))) /
+        2;
+    return std::max(bringing, grasping / 3);
+  }
+  mjtNum BallReward() {
+    // return self._is_close(physics.site_distance('ball', 'target_ball'))
+    return IsClose(SiteDistance(id_site_ball_, id_site_target_ball_));
+  }
+
+  void WriteState() {
+    const auto& bounded_joint_pos = BoundedJointPos();
+    const auto& joint_vel_arm = JointVelArm();
+    const auto& touch = Touch();
+    const auto& hand_pos = Body2dPose(id_xbody_hand_);
+    const auto& object_pos = Body2dPose(id_xbody_object_);
+    const auto& joint_vel_obj = JointVelObj();
+    const auto& target_pos = Body2dPose(id_xbody_target_);
+
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    state["obs:arm_pos"_].Assign(bounded_joint_pos.begin(),
+                                 bounded_joint_pos.size());
+    state["obs:arm_vel"_].Assign(joint_vel_arm.begin(), joint_vel_arm.size());
+    state["obs:touch"_].Assign(touch.begin(), touch.size());
+    state["obs:hand_pos"_].Assign(hand_pos.begin(), hand_pos.size());
+    state["obs:object_pos"_].Assign(object_pos.begin(), object_pos.size());
+    state["obs:object_vel"_].Assign(joint_vel_obj.begin(),
+                                    joint_vel_obj.size());
+    state["obs:target_pos"_].Assign(target_pos.begin(), target_pos.size());
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:random_info"_].Assign(random_info_.begin(), 8);
+#endif
+  }
+};
+
+using ManipulatorEnvPool = AsyncEnvPool<ManipulatorEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_MANIPULATOR_H_
diff --git a/envpool/mujoco/dmc/mujoco_env.cc b/envpool/mujoco/dmc/mujoco_env.cc
new file mode 100644
index 00000000..9bbdd72d
--- /dev/null
+++ b/envpool/mujoco/dmc/mujoco_env.cc
@@ -0,0 +1,209 @@
+// 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.
+
+#include "envpool/mujoco/dmc/mujoco_env.h"
+
+#include <cassert>
+#include <cmath>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+
+namespace mujoco_dmc {
+
+MujocoEnv::MujocoEnv(const std::string& base_path, const std::string& raw_xml,
+                     int n_sub_steps, int max_episode_steps)
+    : n_sub_steps_(n_sub_steps),
+      max_episode_steps_(max_episode_steps),
+      elapsed_step_(max_episode_steps + 1),
+      done_(true),
+      dist_uniform_(0, 1),
+      dist_normal_(0, 1) {
+  // initialize vfs from common assets and raw xml
+  // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/wrapper/core.py#L158
+  // https://github.com/deepmind/mujoco/blob/main/python/mujoco/structs.cc
+  // MjModelWrapper::LoadXML
+  std::unique_ptr<mjVFS, void (*)(mjVFS*)> vfs(new mjVFS, [](mjVFS* vfs) {
+    mj_deleteVFS(vfs);
+    delete vfs;
+  });
+  mj_defaultVFS(vfs.get());
+  // save raw_xml into vfs
+  std::string model_filename("model_.xml");
+  mj_makeEmptyFileVFS(vfs.get(), model_filename.c_str(), raw_xml.size());
+  std::memcpy(vfs->filedata[vfs->nfile - 1], raw_xml.c_str(), raw_xml.size());
+  // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/common/__init__.py#L28
+  std::vector<std::string> common_assets_name(
+      {"./common/materials.xml", "./common/skybox.xml", "./common/visual.xml"});
+  for (const auto& asset_name : common_assets_name) {
+    std::string content = GetFileContent(base_path, asset_name);
+    mj_makeEmptyFileVFS(vfs.get(), asset_name.c_str(), content.size());
+    std::memcpy(vfs->filedata[vfs->nfile - 1], content.c_str(), content.size());
+  }
+  // create model and data
+  model_ = mj_loadXML(model_filename.c_str(), vfs.get(), error_.begin(), 1000);
+  data_ = mj_makeData(model_);
+#ifdef ENVPOOL_TEST
+  qpos0_.reset(new mjtNum[model_->nq]);
+#endif
+}
+
+MujocoEnv::~MujocoEnv() {
+  mj_deleteModel(model_);
+  mj_deleteData(data_);
+}
+
+// rl control Environment
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/rl/control.py#L77
+void MujocoEnv::ControlReset() {
+  elapsed_step_ = 0;
+  discount_ = 1.0;
+  done_ = false;
+  TaskInitializeEpisodeMjcf();
+  // attention: no keyframe_id
+  PhysicsReset();  // first mj_forward
+  TaskInitializeEpisode();
+  PhysicsAfterReset();  // second mj_forward
+}
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/rl/control.py#L94
+void MujocoEnv::ControlStep(const mjtNum* action) {
+  TaskBeforeStep(action);
+  PhysicsStep(n_sub_steps_, action);
+  TaskAfterStep();
+  reward_ = TaskGetReward();
+  if (++elapsed_step_ >= max_episode_steps_) {
+    discount_ = 1.0;
+    done_ = true;
+  } else if (TaskShouldTerminateEpisode()) {
+    // copy from composer task
+    // this is different because we don't use None in C++
+    discount_ = 0.0;
+    done_ = true;
+  } else {
+    discount_ = TaskGetDiscount();
+    done_ = false;
+  }
+}
+
+// Task
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/base.py#L73
+void MujocoEnv::TaskBeforeStep(const mjtNum* action) {
+  PhysicsSetControl(action);
+}
+float MujocoEnv::TaskGetReward() {
+  throw std::runtime_error("GetReward not implemented");
+}
+float MujocoEnv::TaskGetDiscount() { return 1.0; }
+bool MujocoEnv::TaskShouldTerminateEpisode() { return false; }
+
+// Physics
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L263
+void MujocoEnv::PhysicsReset(int keyframe_id) {
+  if (keyframe_id < 0) {
+    mj_resetData(model_, data_);
+  } else {
+    // actually no one steps to this line
+    assert(keyframe_id < model_->nkey);
+    mj_resetDataKeyframe(model_, data_, keyframe_id);
+  }
+
+  // PhysicsAfterReset may be overwritten?
+  int old_flags = model_->opt.disableflags;
+  model_->opt.disableflags |= mjDSBL_ACTUATION;
+  PhysicsForward();
+  model_->opt.disableflags = old_flags;
+}
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L286
+void MujocoEnv::PhysicsAfterReset() {
+  int old_flags = model_->opt.disableflags;
+  model_->opt.disableflags |= mjDSBL_ACTUATION;
+  PhysicsForward();
+  model_->opt.disableflags = old_flags;
+}
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L138
+void MujocoEnv::PhysicsSetControl(const mjtNum* control) {
+  std::memcpy(data_->ctrl, control, sizeof(mjtNum) * model_->nu);
+}
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L292
+void MujocoEnv::PhysicsForward() { mj_forward(model_, data_); }
+
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L146
+void MujocoEnv::PhysicsStep(int nstep, const mjtNum* action) {
+  TaskBeforeSubStep(action);
+  if (model_->opt.integrator != mjINT_RK4) {
+    mj_step2(model_, data_);
+  } else {
+    mj_step(model_, data_);
+  }
+  TaskAftersSubStep();
+  if (nstep > 1) {
+    for (int i = 0; i < nstep - 1; ++i) {
+      TaskBeforeSubStep(action);
+      mj_step(model_, data_);
+      TaskAftersSubStep();
+    }
+  }
+  mj_step1(model_, data_);
+}
+
+// randomizer
+// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/utils/randomizers.py#L35
+void MujocoEnv::RandomizeLimitedAndRotationalJoints(std::mt19937* gen) {
+  for (int joint_id = 0; joint_id < model_->njnt; ++joint_id) {
+    int joint_type = model_->jnt_type[joint_id];
+    mjtByte is_limited = model_->jnt_limited[joint_id];
+    mjtNum range_min = model_->jnt_range[joint_id * 2 + 0];
+    mjtNum range_max = model_->jnt_range[joint_id * 2 + 1];
+    mjtNum range = range_max - range_min;
+    int qpos_offset = model_->jnt_qposadr[joint_id];
+    if (is_limited != 0) {
+      if (joint_type == mjJNT_HINGE || joint_type == mjJNT_SLIDE) {
+        data_->qpos[qpos_offset] = dist_uniform_(*gen) * range + range_min;
+      } else if (joint_type == mjJNT_BALL) {
+        // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/utils/randomizers.py#L23
+        std::array<mjtNum, 3> axis = {dist_normal_(*gen), dist_normal_(*gen),
+                                      dist_normal_(*gen)};
+        auto norm = std::sqrt(axis[0] * axis[0] + axis[1] * axis[1] +
+                              axis[2] * axis[2]);
+        axis = {axis[0] / norm, axis[1] / norm, axis[2] / norm};
+        auto angle = dist_uniform_(*gen) * range_max;
+        mju_axisAngle2Quat(data_->qpos + qpos_offset, axis.begin(), angle);
+      }
+    } else if (joint_type == mjJNT_HINGE) {
+      data_->qpos[qpos_offset] = dist_uniform_(*gen) * M_PI * 2 - M_PI;
+    } else if (joint_type == mjJNT_BALL || joint_type == mjJNT_FREE) {
+      std::array<mjtNum, 4> quat;
+      if (joint_type == mjJNT_BALL) {
+        quat = {dist_normal_(*gen), dist_normal_(*gen), dist_normal_(*gen),
+                dist_normal_(*gen)};
+      } else {
+        quat = {dist_uniform_(*gen), dist_uniform_(*gen), dist_uniform_(*gen),
+                dist_uniform_(*gen)};
+      }
+      auto norm = std::sqrt(quat[0] * quat[0] + quat[1] * quat[1] +
+                            quat[2] * quat[2] + quat[3] * quat[3]);
+      int extra_offset = joint_type == mjJNT_BALL ? 0 : 3;
+      data_->qpos[qpos_offset + extra_offset + 0] = quat[0] / norm;
+      data_->qpos[qpos_offset + extra_offset + 1] = quat[1] / norm;
+      data_->qpos[qpos_offset + extra_offset + 2] = quat[2] / norm;
+      data_->qpos[qpos_offset + extra_offset + 3] = quat[3] / norm;
+    }
+  }
+}
+
+}  // namespace mujoco_dmc
diff --git a/envpool/mujoco/dmc/mujoco_env.h b/envpool/mujoco/dmc/mujoco_env.h
index e1cb0e5f..a8a17cb5 100644
--- a/envpool/mujoco/dmc/mujoco_env.h
+++ b/envpool/mujoco/dmc/mujoco_env.h
@@ -20,33 +20,13 @@
 #include <mjxmacro.h>
 #include <mujoco.h>
 
-#include <cmath>
-#include <fstream>
 #include <memory>
 #include <random>
-#include <sstream>
 #include <string>
-#include <vector>
 
-namespace mujoco_dmc {
-
-std::string GetFileContent(const std::string& base_path,
-                           const std::string& asset_name) {
-  // hardcode path here :(
-  std::string filename = base_path + "/mujoco/assets_dmc/" + asset_name;
-  std::ifstream ifs(filename);
-  std::stringstream ss;
-  ss << ifs.rdbuf();
-  return ss.str();
-}
+#include "envpool/mujoco/dmc/utils.h"
 
-int GetQposId(mjModel* model, const std::string& name) {
-  return model->jnt_qposadr[mj_name2id(model, mjOBJ_JOINT, name.c_str())];
-}
-
-int GetSensorId(mjModel* model, const std::string& name) {
-  return model->sensor_adr[mj_name2id(model, mjOBJ_SENSOR, name.c_str())];
-}
+namespace mujoco_dmc {
 
 /*
  * This class combines with dmc Task and Physics API.
@@ -77,196 +57,47 @@ class MujocoEnv {
 
  public:
   MujocoEnv(const std::string& base_path, const std::string& raw_xml,
-            int n_sub_steps, int max_episode_steps)
-      : n_sub_steps_(n_sub_steps),
-        max_episode_steps_(max_episode_steps),
-        elapsed_step_(max_episode_steps + 1),
-        done_(true),
-        dist_uniform_(0, 1),
-        dist_normal_(0, 1) {
-    // initialize vfs from common assets and raw xml
-    // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/wrapper/core.py#L158
-    // https://github.com/deepmind/mujoco/blob/main/python/mujoco/structs.cc
-    // MjModelWrapper::LoadXML
-    std::unique_ptr<mjVFS, void (*)(mjVFS*)> vfs(new mjVFS, [](mjVFS* vfs) {
-      mj_deleteVFS(vfs);
-      delete vfs;
-    });
-    mj_defaultVFS(vfs.get());
-    // save raw_xml into vfs
-    std::string model_filename("model_.xml");
-    mj_makeEmptyFileVFS(vfs.get(), model_filename.c_str(), raw_xml.size());
-    std::memcpy(vfs->filedata[vfs->nfile - 1], raw_xml.c_str(), raw_xml.size());
-    // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/common/__init__.py#L28
-    std::vector<std::string> common_assets_name({"./common/materials.xml",
-                                                 "./common/skybox.xml",
-                                                 "./common/visual.xml"});
-    for (const auto& asset_name : common_assets_name) {
-      std::string content = GetFileContent(base_path, asset_name);
-      mj_makeEmptyFileVFS(vfs.get(), asset_name.c_str(), content.size());
-      std::memcpy(vfs->filedata[vfs->nfile - 1], content.c_str(),
-                  content.size());
-    }
-    // create model and data
-    model_ =
-        mj_loadXML(model_filename.c_str(), vfs.get(), error_.begin(), 1000);
-    data_ = mj_makeData(model_);
-#ifdef ENVPOOL_TEST
-    qpos0_.reset(new mjtNum[model_->nq]);
-#endif
-  }
-
-  ~MujocoEnv() {
-    mj_deleteModel(model_);
-    mj_deleteData(data_);
-  }
+            int n_sub_steps, int max_episode_steps);
+  ~MujocoEnv();
 
   // rl control Environment
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/rl/control.py#L77
-  void ControlReset() {
-    elapsed_step_ = 0;
-    discount_ = 1.0;
-    done_ = false;
-    TaskInitializeEpisodeMjcf();
-    // attention: no keyframe_id
-    PhysicsReset();  // first mj_forward
-    TaskInitializeEpisode();
-    PhysicsAfterReset();  // second mj_forward
-  }
+  void ControlReset();
 
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/rl/control.py#L94
-  void ControlStep(const mjtNum* action) {
-    TaskBeforeStep(action);
-    PhysicsStep(n_sub_steps_, action);
-    TaskAfterStep();
-    reward_ = TaskGetReward();
-    if (++elapsed_step_ >= max_episode_steps_) {
-      discount_ = 1.0;
-      done_ = true;
-    } else if (TaskShouldTerminateEpisode()) {
-      // copy from composer task
-      // this is different because we don't use None in C++
-      discount_ = 0.0;
-      done_ = true;
-    } else {
-      discount_ = TaskGetDiscount();
-      done_ = false;
-    }
-  }
+  void ControlStep(const mjtNum* action);
 
   // Task
   virtual void TaskInitializeEpisodeMjcf() {}
   virtual void TaskInitializeEpisode() {}
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/base.py#L73
-  virtual void TaskBeforeStep(const mjtNum* action) {
-    PhysicsSetControl(action);
-  }
+  virtual void TaskBeforeStep(const mjtNum* action);
   virtual void TaskBeforeSubStep(const mjtNum* action) {}
   virtual void TaskAfterStep() {}
   virtual void TaskAftersSubStep() {}
-  virtual float TaskGetReward() {
-    throw std::runtime_error("GetReward not implemented");
-  }
-  virtual float TaskGetDiscount() { return 1.0; }
-  virtual bool TaskShouldTerminateEpisode() { return false; }
+  virtual float TaskGetReward();
+  virtual float TaskGetDiscount();
+  virtual bool TaskShouldTerminateEpisode();
 
   // Physics
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L263
-  void PhysicsReset(int keyframe_id = -1) {
-    if (keyframe_id < 0) {
-      mj_resetData(model_, data_);
-    } else {
-      // actually no one steps to this line
-      assert(keyframe_id < model_->nkey);
-      mj_resetDataKeyframe(model_, data_, keyframe_id);
-    }
-
-    // PhysicsAfterReset may be overwritten?
-    int old_flags = model_->opt.disableflags;
-    model_->opt.disableflags |= mjDSBL_ACTUATION;
-    PhysicsForward();
-    model_->opt.disableflags = old_flags;
-  }
+  void PhysicsReset(int keyframe_id = -1);
 
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L286
-  void PhysicsAfterReset() {
-    int old_flags = model_->opt.disableflags;
-    model_->opt.disableflags |= mjDSBL_ACTUATION;
-    PhysicsForward();
-    model_->opt.disableflags = old_flags;
-  }
+  void PhysicsAfterReset();
 
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L138
-  void PhysicsSetControl(const mjtNum* control) {
-    std::memcpy(data_->ctrl, control, sizeof(mjtNum) * model_->nu);
-  }
+  void PhysicsSetControl(const mjtNum* control);
 
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L292
-  void PhysicsForward() { mj_forward(model_, data_); }
+  void PhysicsForward();
 
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/mujoco/engine.py#L146
-  void PhysicsStep(int nstep, const mjtNum* action) {
-    TaskBeforeSubStep(action);
-    if (model_->opt.integrator != mjINT_RK4) {
-      mj_step2(model_, data_);
-    } else {
-      mj_step(model_, data_);
-    }
-    TaskAftersSubStep();
-    if (nstep > 1) {
-      for (int i = 0; i < nstep - 1; ++i) {
-        TaskBeforeSubStep(action);
-        mj_step(model_, data_);
-        TaskAftersSubStep();
-      }
-    }
-    mj_step1(model_, data_);
-  }
+  void PhysicsStep(int nstep, const mjtNum* action);
 
   // randomizer
   // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/utils/randomizers.py#L35
-  void RandomizeLimitedAndRotationalJoints(std::mt19937* gen) {
-    for (int joint_id = 0; joint_id < model_->njnt; ++joint_id) {
-      auto joint_type = model_->jnt_type[joint_id];
-      auto is_limited = model_->jnt_limited[joint_id];
-      auto range_min = model_->jnt_range[joint_id * 2 + 0];
-      auto range_max = model_->jnt_range[joint_id * 2 + 1];
-      auto range = range_max - range_min;
-      auto qpos_offset = model_->jnt_qposadr[joint_id];
-      if (is_limited != 0) {
-        if (joint_type == mjJNT_HINGE || joint_type == mjJNT_SLIDE) {
-          data_->qpos[qpos_offset] = dist_uniform_(*gen) * range + range_min;
-        } else if (joint_type == mjJNT_BALL) {
-          // https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/utils/randomizers.py#L23
-          std::array<mjtNum, 3> axis = {dist_normal_(*gen), dist_normal_(*gen),
-                                        dist_normal_(*gen)};
-          auto norm = std::sqrt(axis[0] * axis[0] + axis[1] * axis[1] +
-                                axis[2] * axis[2]);
-          axis = {axis[0] / norm, axis[1] / norm, axis[2] / norm};
-          auto angle = dist_uniform_(*gen) * range_max;
-          mju_axisAngle2Quat(data_->qpos + qpos_offset, axis.begin(), angle);
-        }
-      } else if (joint_type == mjJNT_HINGE) {
-        data_->qpos[qpos_offset] = dist_uniform_(*gen) * M_PI * 2 - M_PI;
-      } else if (joint_type == mjJNT_BALL || joint_type == mjJNT_FREE) {
-        std::array<mjtNum, 4> quat;
-        if (joint_type == mjJNT_BALL) {
-          quat = {dist_normal_(*gen), dist_normal_(*gen), dist_normal_(*gen),
-                  dist_normal_(*gen)};
-        } else {
-          quat = {dist_uniform_(*gen), dist_uniform_(*gen), dist_uniform_(*gen),
-                  dist_uniform_(*gen)};
-        }
-        auto norm = std::sqrt(quat[0] * quat[0] + quat[1] * quat[1] +
-                              quat[2] * quat[2] + quat[3] * quat[3]);
-        int extra_offset = joint_type == mjJNT_BALL ? 0 : 3;
-        data_->qpos[qpos_offset + extra_offset + 0] = quat[0] / norm;
-        data_->qpos[qpos_offset + extra_offset + 1] = quat[1] / norm;
-        data_->qpos[qpos_offset + extra_offset + 2] = quat[2] / norm;
-        data_->qpos[qpos_offset + extra_offset + 3] = quat[3] / norm;
-      }
-    }
-  }
+  void RandomizeLimitedAndRotationalJoints(std::mt19937* gen);
 };
 
 }  // namespace mujoco_dmc
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 16352c2a..e07a4440 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -19,6 +19,7 @@
 #include "envpool/mujoco/dmc/finger.h"
 #include "envpool/mujoco/dmc/hopper.h"
 #include "envpool/mujoco/dmc/humanoid.h"
+#include "envpool/mujoco/dmc/manipulator.h"
 #include "envpool/mujoco/dmc/pendulum.h"
 #include "envpool/mujoco/dmc/point_mass.h"
 #include "envpool/mujoco/dmc/reacher.h"
@@ -42,6 +43,9 @@ using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 using DmcHumanoidEnvSpec = PyEnvSpec<mujoco_dmc::HumanoidEnvSpec>;
 using DmcHumanoidEnvPool = PyEnvPool<mujoco_dmc::HumanoidEnvPool>;
 
+using DmcManipulatorEnvSpec = PyEnvSpec<mujoco_dmc::ManipulatorEnvSpec>;
+using DmcManipulatorEnvPool = PyEnvPool<mujoco_dmc::ManipulatorEnvPool>;
+
 using DmcPendulumEnvSpec = PyEnvSpec<mujoco_dmc::PendulumEnvSpec>;
 using DmcPendulumEnvPool = PyEnvPool<mujoco_dmc::PendulumEnvPool>;
 
@@ -61,6 +65,7 @@ PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
   REGISTER(m, DmcHumanoidEnvSpec, DmcHumanoidEnvPool)
+  REGISTER(m, DmcManipulatorEnvSpec, DmcManipulatorEnvPool)
   REGISTER(m, DmcPendulumEnvSpec, DmcPendulumEnvPool)
   REGISTER(m, DmcPointMassEnvSpec, DmcPointMassEnvPool)
   REGISTER(m, DmcReacherEnvSpec, DmcReacherEnvPool)
diff --git a/envpool/mujoco/dmc/registration.py b/envpool/mujoco/dmc/registration.py
index 15606d23..5a875217 100644
--- a/envpool/mujoco/dmc/registration.py
+++ b/envpool/mujoco/dmc/registration.py
@@ -36,6 +36,10 @@
   ("humanoid", "walk"),
   ("humanoid", "run"),
   ("humanoid", "run_pure_state"),
+  ("manipulator", "bring_ball"),
+  ("manipulator", "bring_peg"),
+  ("manipulator", "insert_ball"),
+  ("manipulator", "insert_peg"),
   ("pendulum", "swingup"),
   ("point_mass", "easy"),
   ("point_mass", "hard"),
diff --git a/envpool/mujoco/dmc/utils.cc b/envpool/mujoco/dmc/utils.cc
new file mode 100644
index 00000000..bf6caa2b
--- /dev/null
+++ b/envpool/mujoco/dmc/utils.cc
@@ -0,0 +1,136 @@
+// 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.
+
+#include "envpool/mujoco/dmc/utils.h"
+
+#include <cmath>
+#include <fstream>
+#include <sstream>
+
+#include "pugixml.hpp"
+
+namespace mujoco_dmc {
+
+std::string GetFileContent(const std::string& base_path,
+                           const std::string& asset_name) {
+  // hardcode path here :(
+  std::string filename = base_path + "/mujoco/assets_dmc/" + asset_name;
+  std::ifstream ifs(filename);
+  std::stringstream ss;
+  ss << ifs.rdbuf();
+  return ss.str();
+}
+
+class XMLStringWriter : public pugi::xml_writer {
+ public:
+  std::string result;
+  void write(const void* data, size_t size) override {
+    result.append(static_cast<const char*>(data), size);
+  }
+};
+
+std::string XMLRemoveByBodyName(const std::string& content,
+                                const std::vector<std::string>& body_names) {
+  pugi::xml_document doc;
+  doc.load_string(content.c_str());
+  for (const auto& name : body_names) {
+    std::string xpath = "//body[@name='" + name + "']";
+    pugi::xml_node node = doc.select_node(xpath.c_str()).node();
+    auto parent = node.parent();
+    parent.remove_child(node);
+  }
+  XMLStringWriter writer;
+  doc.print(writer);
+  return writer.result;
+}
+
+int GetQposId(mjModel* model, const std::string& name) {
+  return model->jnt_qposadr[mj_name2id(model, mjOBJ_JOINT, name.c_str())];
+}
+
+int GetSensorId(mjModel* model, const std::string& name) {
+  return model->sensor_adr[mj_name2id(model, mjOBJ_SENSOR, name.c_str())];
+}
+
+// rewards
+double RewardTolerance(double x, double bound_min, double bound_max,
+                       double margin, double value_at_margin,
+                       SigmoidType sigmoid_type) {
+  if (bound_min <= x && x <= bound_max) {
+    return 1.0;
+  }
+  if (margin <= 0.0) {
+    return 0.0;
+  }
+  x = (x < bound_min ? bound_min - x : x - bound_max) / margin;
+  if (sigmoid_type == SigmoidType::kGaussian) {
+    // scale = np.sqrt(-2 * np.log(value_at_1))
+    // return np.exp(-0.5 * (x*scale)**2)
+    double scaled_x = std::sqrt(-2 * std::log(value_at_margin)) * x;
+    return std::exp(-0.5 * scaled_x * scaled_x);
+  }
+  if (sigmoid_type == SigmoidType::kHyperbolic) {
+    // scale = np.arccosh(1/value_at_1)
+    // return 1 / np.cosh(x*scale)
+    double scaled_x = std::acosh(1 / value_at_margin) * x;
+    return 1 / std::cosh(scaled_x);
+  }
+  if (sigmoid_type == SigmoidType::kLongTail) {
+    // scale = np.sqrt(1/value_at_1 - 1)
+    // return 1 / ((x*scale)**2 + 1)
+    double scaled_x = std::sqrt(1 / value_at_margin - 1) * x;
+    return 1 / (scaled_x * scaled_x + 1);
+  }
+  if (sigmoid_type == SigmoidType::kReciprocal) {
+    // scale = 1/value_at_1 - 1
+    // return 1 / (abs(x)*scale + 1)
+    double scale = 1 / value_at_margin - 1;
+    return 1 / (std::abs(x) * scale + 1);
+  }
+  if (sigmoid_type == SigmoidType::kCosine) {
+    // scale = np.arccos(2*value_at_1 - 1) / np.pi
+    // scaled_x = x*scale
+    // with warnings.catch_warnings():
+    //   warnings.filterwarnings(
+    //       action='ignore', message='invalid value encountered in cos')
+    //   cos_pi_scaled_x = np.cos(np.pi*scaled_x)
+    // return np.where(abs(scaled_x) < 1, (1 + cos_pi_scaled_x)/2, 0.0)
+    const double pi = std::acos(-1);
+    double scaled_x = std::acos(2 * value_at_margin - 1) / pi * x;
+    return std::abs(scaled_x) < 1 ? (1 + std::cos(pi * scaled_x)) / 2 : 0.0;
+  }
+  if (sigmoid_type == SigmoidType::kLinear) {
+    // scale = 1-value_at_1
+    // scaled_x = x*scale
+    // return np.where(abs(scaled_x) < 1, 1 - scaled_x, 0.0)
+    double scaled_x = (1 - value_at_margin) * x;
+    return std::abs(scaled_x) < 1 ? 1 - scaled_x : 0.0;
+  }
+  if (sigmoid_type == SigmoidType::kQuadratic) {
+    // scale = np.sqrt(1-value_at_1)
+    // scaled_x = x*scale
+    // return np.where(abs(scaled_x) < 1, 1 - scaled_x**2, 0.0)
+    double scaled_x = std::sqrt(1 - value_at_margin) * x;
+    return std::abs(scaled_x) < 1 ? 1 - scaled_x * scaled_x : 0.0;
+  }
+  if (sigmoid_type == SigmoidType::kTanhSquared) {
+    // scale = np.arctanh(np.sqrt(1-value_at_1))
+    // return 1 - np.tanh(x*scale)**2
+    double scaled_x = std::atanh(std::sqrt(1 - value_at_margin)) * x;
+    return 1 - std::tanh(scaled_x) * std::tanh(scaled_x);
+  }
+  throw std::runtime_error("Unknown sigmoid type for RewardTolerance.");
+}
+
+}  // namespace mujoco_dmc
diff --git a/envpool/mujoco/dmc/utils.h b/envpool/mujoco/dmc/utils.h
index 39824d20..91a0bfa0 100644
--- a/envpool/mujoco/dmc/utils.h
+++ b/envpool/mujoco/dmc/utils.h
@@ -17,10 +17,25 @@
 #ifndef ENVPOOL_MUJOCO_DMC_UTILS_H_
 #define ENVPOOL_MUJOCO_DMC_UTILS_H_
 
-#include <cmath>
+#include <mjxmacro.h>
+#include <mujoco.h>
+
+#include <string>
+#include <vector>
 
 namespace mujoco_dmc {
 
+// xml related
+std::string GetFileContent(const std::string& base_path,
+                           const std::string& asset_name);
+std::string XMLRemoveByBodyName(const std::string& content,
+                                const std::vector<std::string>& body_names);
+
+// the following id is not 1 on 1 mapping
+int GetQposId(mjModel* model, const std::string& name);
+int GetSensorId(mjModel* model, const std::string& name);
+
+// rewards
 enum class SigmoidType {
   kGaussian,
   kHyperbolic,
@@ -34,72 +49,7 @@ enum class SigmoidType {
 
 double RewardTolerance(double x, double bound_min = 0.0, double bound_max = 0.0,
                        double margin = 0.0, double value_at_margin = 0.1,
-                       SigmoidType sigmoid_type = SigmoidType::kGaussian) {
-  if (bound_min <= x && x <= bound_max) {
-    return 1.0;
-  }
-  if (margin <= 0.0) {
-    return 0.0;
-  }
-  x = (x < bound_min ? bound_min - x : x - bound_max) / margin;
-  if (sigmoid_type == SigmoidType::kGaussian) {
-    // scale = np.sqrt(-2 * np.log(value_at_1))
-    // return np.exp(-0.5 * (x*scale)**2)
-    double scaled_x = std::sqrt(-2 * std::log(value_at_margin)) * x;
-    return std::exp(-0.5 * scaled_x * scaled_x);
-  }
-  if (sigmoid_type == SigmoidType::kHyperbolic) {
-    // scale = np.arccosh(1/value_at_1)
-    // return 1 / np.cosh(x*scale)
-    double scaled_x = std::acosh(1 / value_at_margin) * x;
-    return 1 / std::cosh(scaled_x);
-  }
-  if (sigmoid_type == SigmoidType::kLongTail) {
-    // scale = np.sqrt(1/value_at_1 - 1)
-    // return 1 / ((x*scale)**2 + 1)
-    double scaled_x = std::sqrt(1 / value_at_margin - 1) * x;
-    return 1 / (scaled_x * scaled_x + 1);
-  }
-  if (sigmoid_type == SigmoidType::kReciprocal) {
-    // scale = 1/value_at_1 - 1
-    // return 1 / (abs(x)*scale + 1)
-    double scale = 1 / value_at_margin - 1;
-    return 1 / (std::abs(x) * scale + 1);
-  }
-  if (sigmoid_type == SigmoidType::kCosine) {
-    // scale = np.arccos(2*value_at_1 - 1) / np.pi
-    // scaled_x = x*scale
-    // with warnings.catch_warnings():
-    //   warnings.filterwarnings(
-    //       action='ignore', message='invalid value encountered in cos')
-    //   cos_pi_scaled_x = np.cos(np.pi*scaled_x)
-    // return np.where(abs(scaled_x) < 1, (1 + cos_pi_scaled_x)/2, 0.0)
-    const double pi = std::acos(-1);
-    double scaled_x = std::acos(2 * value_at_margin - 1) / pi * x;
-    return std::abs(scaled_x) < 1 ? (1 + std::cos(pi * scaled_x)) / 2 : 0.0;
-  }
-  if (sigmoid_type == SigmoidType::kLinear) {
-    // scale = 1-value_at_1
-    // scaled_x = x*scale
-    // return np.where(abs(scaled_x) < 1, 1 - scaled_x, 0.0)
-    double scaled_x = (1 - value_at_margin) * x;
-    return std::abs(scaled_x) < 1 ? 1 - scaled_x : 0.0;
-  }
-  if (sigmoid_type == SigmoidType::kQuadratic) {
-    // scale = np.sqrt(1-value_at_1)
-    // scaled_x = x*scale
-    // return np.where(abs(scaled_x) < 1, 1 - scaled_x**2, 0.0)
-    double scaled_x = std::sqrt(1 - value_at_margin) * x;
-    return std::abs(scaled_x) < 1 ? 1 - scaled_x * scaled_x : 0.0;
-  }
-  if (sigmoid_type == SigmoidType::kTanhSquared) {
-    // scale = np.arctanh(np.sqrt(1-value_at_1))
-    // return 1 - np.tanh(x*scale)**2
-    double scaled_x = std::atanh(std::sqrt(1 - value_at_margin)) * x;
-    return 1 - std::tanh(scaled_x) * std::tanh(scaled_x);
-  }
-  throw std::runtime_error("Unknown sigmoid type for RewardTolerance.");
-}
+                       SigmoidType sigmoid_type = SigmoidType::kGaussian);
 
 }  // namespace mujoco_dmc
 
diff --git a/envpool/workspace0.bzl b/envpool/workspace0.bzl
index 988fd29a..63189320 100644
--- a/envpool/workspace0.bzl
+++ b/envpool/workspace0.bzl
@@ -168,6 +168,18 @@ def workspace():
         build_file = "//third_party/opencv:opencv.BUILD",
     )
 
+    maybe(
+        http_archive,
+        name = "pugixml",
+        urls = [
+            "https://github.com/zeux/pugixml/releases/download/v1.12.1/pugixml-1.12.1.tar.gz",
+            "https://ml.cs.tsinghua.edu.cn/~jiayi/envpool/zeux/pugixml/pugixml-1.12.1.tar.gz",
+        ],
+        sha256 = "dcf671a919cc4051210f08ffd3edf9e4247f79ad583c61577a13ee93af33afc7",
+        strip_prefix = "pugixml-1.12/src",
+        build_file = "//third_party/pugixml:pugixml.BUILD",
+    )
+
     maybe(
         http_archive,
         name = "ale",
diff --git a/third_party/pugixml/BUILD b/third_party/pugixml/BUILD
new file mode 100644
index 00000000..e69de29b
diff --git a/third_party/pugixml/pugixml.BUILD b/third_party/pugixml/pugixml.BUILD
new file mode 100644
index 00000000..d13608ba
--- /dev/null
+++ b/third_party/pugixml/pugixml.BUILD
@@ -0,0 +1,10 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "pugixml",
+    srcs = ["pugixml.cpp"],
+    hdrs = [
+        "pugiconfig.hpp",
+        "pugixml.hpp",
+    ],
+)


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 5d983825..091ff28d 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -140,6 +140,10 @@ def test_make_mujoco_dmc(self) -> None:
         "HumanoidWalk-v1",
         "HumanoidRun-v1",
         "HumanoidRunPureState-v1",
+        "ManipulatorBringBall-v1",
+        "ManipulatorBringPeg-v1",
+        "ManipulatorInsertBall-v1",
+        "ManipulatorInsertPeg-v1",
         "PendulumSwingup-v1",
         "PointMassEasy-v1",
         "PointMassHard-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 2583d801..33cb19e8 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -34,6 +34,8 @@
   DmcHopperEnvSpec,
   DmcHumanoidDMEnvPool,
   DmcHumanoidEnvSpec,
+  DmcManipulatorDMEnvPool,
+  DmcManipulatorEnvSpec,
   DmcPendulumDMEnvPool,
   DmcPendulumEnvSpec,
   DmcPointMassDMEnvPool,
@@ -85,7 +87,49 @@ def reset_state(
                                            ["x", "z"]] = target_x, target_z
         env.physics.after_reset()
       elif domain == "point_mass":
-        env.physics.model.wrap_prm = ts.observation.wrap_prm
+        env.physics.model.wrap_prm = ts.observation.wrap_prm[0]
+      elif domain == "manipulator":
+        # config
+        use_peg = "peg" in task
+        insert = "insert" in task
+        obj = "peg" if use_peg else "ball"
+        target = "target_" + obj
+        obj_joints = [obj + "_x", obj + "_z", obj + "_y"]
+        receptacle = "slot" if use_peg else "cup"
+        (
+          target_x, target_z, target_angle, init_type, object_x, object_z,
+          object_angle, qvel_objx
+        ) = ts.observation.random_info[0]
+        logging.info(ts.observation.random_info[0])
+        p = env.physics.named
+        # assign random value from envpool
+        if insert:
+          p.model.body_pos[receptacle, ["x", "z"]] = target_x, target_z
+          p.model.body_quat[receptacle, ["qw", "qy"]] = (
+            np.cos(target_angle / 2), np.sin(target_angle / 2)
+          )
+        p.model.body_pos[target, ["x", "z"]] = target_x, target_z
+        p.model.body_quat[target, ["qw", "qy"]] = \
+          (np.cos(target_angle / 2), np.sin(target_angle / 2))
+        if np.isclose(init_type, 1):
+          np.testing.assert_allclose(
+            [object_x, object_z, object_angle],
+            [target_x, target_z, target_angle]
+          )
+        elif np.isclose(init_type, 2):
+          env.physics.after_reset()
+          np.testing.assert_allclose(
+            [object_x, object_z], p.data.site_xpos["grasp", ["x", "z"]]
+          )
+          grasp_direction = p.data.site_xmat["grasp", ["xx", "zx"]]
+          np.testing.assert_allclose(
+            object_angle,
+            np.pi - np.arctan2(grasp_direction[1], grasp_direction[0])
+          )
+        else:
+          p.data.qvel[obj + "_x"] = qvel_objx
+        p.data.qpos[obj_joints] = object_x, object_z, object_angle
+        env.physics.after_reset()
 
   def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
     return np.random.uniform(
@@ -164,6 +208,12 @@ def test_humanoid(self) -> None:
       DmcHumanoidEnvSpec, DmcHumanoidDMEnvPool
     )
 
+  def test_manipulator(self) -> None:
+    self.run_align_check_entry(
+      "manipulator", ["bring_ball", "bring_peg", "insert_ball", "insert_peg"],
+      DmcManipulatorEnvSpec, DmcManipulatorDMEnvPool
+    )
+
   def test_pendulum(self) -> None:
     self.run_align_check_entry(
       "pendulum", ["swingup"], DmcPendulumEnvSpec, DmcPendulumDMEnvPool
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index 73680dd8..e551c4b5 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -32,6 +32,8 @@
   DmcHopperEnvSpec,
   DmcHumanoidDMEnvPool,
   DmcHumanoidEnvSpec,
+  DmcManipulatorDMEnvPool,
+  DmcManipulatorEnvSpec,
   DmcPendulumDMEnvPool,
   DmcPendulumEnvSpec,
   DmcPointMassDMEnvPool,
@@ -133,6 +135,16 @@ def test_humanoid(self) -> None:
     for task in ["run_pure_state"]:
       self.check(DmcHumanoidEnvSpec, DmcHumanoidDMEnvPool, task, obs_keys)
 
+  def test_manipulator(self) -> None:
+    obs_keys = [
+      "arm_pos", "arm_vel", "touch", "hand_pos", "object_pos", "object_vel",
+      "target_pos"
+    ]
+    for task in ["bring_ball", "bring_peg", "insert_ball", "insert_peg"]:
+      self.check(
+        DmcManipulatorEnvSpec, DmcManipulatorDMEnvPool, task, obs_keys
+      )
+
   def test_pendulum(self) -> None:
     obs_keys = ["orientation", "velocity"]
     for task in ["swingup"]: