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 2902cbf1..95ea5428 100644
--- a/docs/env/dm_control.rst
+++ b/docs/env/dm_control.rst
@@ -79,6 +79,24 @@ FingerSpin-v1, FingerTurnEasy-v1, FingerTurnHard-v1
     values are meaningless in ``FingerSpin-v1``.
 
 
+FishSwim-v1, FishUpright-v1
+---------------------------
+
+`dm_control suite fish source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/fish.py>`_
+
+- Observation spec: a namedtuple with four keys: ``joint_angles (7)``,
+  ``upright ()``, ``target (3)``, ``velocity (13)``;
+- Action spec: ``(5)``, with range ``[-1, 1]``;
+- ``frame_skip``: 10;
+- ``max_episode_steps``: 1000;
+
+.. note ::
+
+    The observation key ``target`` is only available in ``FishSwim-v1`` task.
+    The value is meaningless in ``FishUpright-v1``.
+
+
 HopperStand-v1, HopperHop-v1
 ----------------------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 3be423cd..f17aeaac 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -71,6 +71,7 @@ cc_library(
         "dmc/ball_in_cup.h",
         "dmc/cheetah.h",
         "dmc/finger.h",
+        "dmc/fish.h",
         "dmc/hopper.h",
         "dmc/humanoid.h",
         "dmc/manipulator.h",
diff --git a/envpool/mujoco/dmc/__init__.py b/envpool/mujoco/dmc/__init__.py
index f86a3fea..9d604aef 100644
--- a/envpool/mujoco/dmc/__init__.py
+++ b/envpool/mujoco/dmc/__init__.py
@@ -22,6 +22,8 @@
   _DmcCheetahEnvSpec,
   _DmcFingerEnvPool,
   _DmcFingerEnvSpec,
+  _DmcFishEnvPool,
+  _DmcFishEnvSpec,
   _DmcHopperEnvPool,
   _DmcHopperEnvSpec,
   _DmcHumanoidEnvPool,
@@ -51,6 +53,9 @@
 DmcFingerEnvSpec, DmcFingerDMEnvPool, DmcFingerGymEnvPool = py_env(
   _DmcFingerEnvSpec, _DmcFingerEnvPool
 )
+DmcFishEnvSpec, DmcFishDMEnvPool, DmcFishGymEnvPool = py_env(
+  _DmcFishEnvSpec, _DmcFishEnvPool
+)
 DmcHopperEnvSpec, DmcHopperDMEnvPool, DmcHopperGymEnvPool = py_env(
   _DmcHopperEnvSpec, _DmcHopperEnvPool
 )
@@ -85,6 +90,12 @@
   "DmcCheetahEnvSpec",
   "DmcCheetahDMEnvPool",
   "DmcCheetahGymEnvPool",
+  "DmcFingerEnvSpec",
+  "DmcFingerDMEnvPool",
+  "DmcFingerGymEnvPool",
+  "DmcFishEnvSpec",
+  "DmcFishDMEnvPool",
+  "DmcFishGymEnvPool",
   "DmcHopperEnvSpec",
   "DmcHopperDMEnvPool",
   "DmcHopperGymEnvPool",
diff --git a/envpool/mujoco/dmc/finger.h b/envpool/mujoco/dmc/finger.h
index d250ffef..545fcd4d 100644
--- a/envpool/mujoco/dmc/finger.h
+++ b/envpool/mujoco/dmc/finger.h
@@ -92,7 +92,7 @@ class FingerEnv : public Env<FingerEnvSpec>, public MujocoEnv {
             spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
         id_site_target_(mj_name2id(model_, mjOBJ_SITE, "target")),
         id_site_tip_(mj_name2id(model_, mjOBJ_SITE, "tip")),
-        id_hinge_(model_->jnt_dofadr[mj_name2id(model_, mjOBJ_JOINT, "hinge")]),
+        id_hinge_(GetQvelId(model_, "hinge")),
         id_cap1_(mj_name2id(model_, mjOBJ_GEOM, "cap1")),
         id_proximal_(GetSensorId(model_, "proximal")),
         id_distal_(GetSensorId(model_, "distal")),
diff --git a/envpool/mujoco/dmc/fish.h b/envpool/mujoco/dmc/fish.h
new file mode 100644
index 00000000..886d9c73
--- /dev/null
+++ b/envpool/mujoco/dmc/fish.h
@@ -0,0 +1,248 @@
+/*
+ * 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/fish.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_FISH_H_
+#define ENVPOOL_MUJOCO_DMC_FISH_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 GetFishXML(const std::string& base_path,
+                       const std::string& task_name_) {
+  return GetFileContent(base_path, "fish.xml");
+}
+
+class FishEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(10),
+                    "task_name"_.Bind(std::string("upright")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:joint_angles"_.Bind(Spec<mjtNum>({7})),
+                    "obs:upright"_.Bind(Spec<mjtNum>({})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({13})),
+                    "obs:target"_.Bind(Spec<mjtNum>({3})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({14})),
+                    "info:target0"_.Bind(Spec<mjtNum>({3})),
+#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 FishEnvSpec = EnvSpec<FishEnvFns>;
+
+class FishEnv : public Env<FishEnvSpec>, public MujocoEnv {
+  const std::array<std::string, 7> kJoints = {
+      "tail1",          "tail_twist",   "tail2",        "finright_roll",
+      "finright_pitch", "finleft_roll", "finleft_pitch"};
+
+ protected:
+  int id_mouth_, id_qpos_root_, id_torso_, id_target_;
+  std::array<int, 7> id_qpos_joint_, id_qvel_joint_;
+  std::normal_distribution<> dist_normal_;
+  std::uniform_real_distribution<> dist_uniform_;
+  bool is_swim_;
+#ifdef ENVPOOL_TEST
+  std::array<mjtNum, 3> target0_;
+#endif
+
+ public:
+  FishEnv(const Spec& spec, int env_id)
+      : Env<FishEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_],
+            GetFishXML(spec.config["base_path"_], spec.config["task_name"_]),
+            spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
+        id_mouth_(mj_name2id(model_, mjOBJ_GEOM, "mouth")),
+        id_qpos_root_(GetQposId(model_, "root")),
+        id_torso_(mj_name2id(model_, mjOBJ_XBODY, "torso")),
+        id_target_(mj_name2id(model_, mjOBJ_GEOM, "target")),
+        dist_normal_(0, 1),
+        dist_uniform_(0, 1),
+        is_swim_(spec.config["task_name"_] == "swim") {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name != "upright" && task_name != "swim") {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc fish.");
+    }
+    for (std::size_t i = 0; i < kJoints.size(); ++i) {
+      id_qpos_joint_[i] = GetQposId(model_, kJoints[i]);
+      id_qvel_joint_[i] = GetQvelId(model_, kJoints[i]);
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    // quat = self.random.randn(4)
+    // physics.named.data.qpos['root'][3:7] = quat / np.linalg.norm(quat)
+    std::array<mjtNum, 4> quat = {dist_normal_(gen_), dist_normal_(gen_),
+                                  dist_normal_(gen_), dist_normal_(gen_)};
+    mjtNum quat_norm = std::sqrt(quat[0] * quat[0] + quat[1] * quat[1] +
+                                 quat[2] * quat[2] + quat[3] * quat[3]);
+    for (int i = 0; i < 4; ++i) {
+      data_->qpos[id_qpos_root_ + 3 + i] = quat[i] / quat_norm;
+    }
+    // for joint in _JOINTS:
+    //   physics.named.data.qpos[joint] = self.random.uniform(-.2, .2)
+    for (int id : id_qpos_joint_) {
+      data_->qpos[id] = dist_uniform_(gen_) * 0.4 - 0.2;
+    }
+    if (is_swim_) {
+      // Randomize target position.
+      // physics.named.model.geom_pos['target', 'x'] = uniform(-.4, .4)
+      // physics.named.model.geom_pos['target', 'y'] = uniform(-.4, .4)
+      // physics.named.model.geom_pos['target', 'z'] = uniform(.1, .3)
+      mjtNum target_x = dist_uniform_(gen_) * 0.8 - 0.4;
+      mjtNum target_y = dist_uniform_(gen_) * 0.8 - 0.4;
+      mjtNum target_z = dist_uniform_(gen_) * 0.2 + 0.1;
+      model_->geom_pos[id_target_ * 3 + 0] = target_x;
+      model_->geom_pos[id_target_ * 3 + 1] = target_y;
+      model_->geom_pos[id_target_ * 3 + 2] = target_z;
+    } else {
+      // Hide the target. It's irrelevant for this task.
+      // physics.named.model.geom_rgba['target', 3] = 0
+      model_->geom_rgba[id_target_ * 4 + 3] = 0;
+    }
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+    target0_[0] = model_->geom_pos[id_target_ * 3 + 0];
+    target0_[1] = model_->geom_pos[id_target_ * 3 + 1];
+    target0_[2] = model_->geom_pos[id_target_ * 3 + 2];
+#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 {
+    if (!is_swim_) {
+      return static_cast<float>(RewardTolerance(Upright(), 1.0, 1.0, 1.0));
+    }
+    mjtNum radii =
+        model_->geom_size[id_mouth_ * 3] + model_->geom_size[id_target_ * 3];
+    const auto& target = MouthToTarget();
+    auto target_norm = std::sqrt(target[0] * target[0] + target[1] * target[1] +
+                                 target[2] * target[2]);
+    auto in_target = RewardTolerance(target_norm, 0.0, radii, 2 * radii);
+    auto is_upright = 0.5 * (Upright() + 1);
+    return static_cast<float>((7 * in_target + is_upright) / 8);
+  }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    const auto& joint_angles = JointAngles();
+    state["obs:joint_angles"_].Assign(joint_angles.begin(),
+                                      joint_angles.size());
+    state["obs:upright"_] = Upright();
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+    if (is_swim_) {
+      const auto& target = MouthToTarget();
+      state["obs:target"_].Assign(target.begin(), target.size());
+    }
+    // info
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:target0"_].Assign(target0_.begin(), target0_.size());
+#endif
+  }
+
+  mjtNum Upright() {
+    // return self.named.data.xmat['torso', 'zz']
+    return data_->xmat[id_torso_ * 9 + 8];
+  }
+
+  std::array<mjtNum, 6> TorsoVelocity() {
+    // return self.data.sensordata
+    return {data_->sensordata[0], data_->sensordata[1], data_->sensordata[2],
+            data_->sensordata[3], data_->sensordata[4], data_->sensordata[5]};
+  }
+
+  std::array<mjtNum, 7> JointVelocities() {
+    // return self.named.data.qvel[_JOINTS]
+    std::array<mjtNum, 7> result;
+    for (std::size_t i = 0; i < id_qvel_joint_.size(); ++i) {
+      result[i] = data_->qvel[id_qvel_joint_[i]];
+    }
+    return result;
+  }
+
+  std::array<mjtNum, 7> JointAngles() {
+    // return self.named.data.qpos[_JOINTS]
+    std::array<mjtNum, 7> result;
+    for (std::size_t i = 0; i < id_qpos_joint_.size(); ++i) {
+      result[i] = data_->qpos[id_qpos_joint_[i]];
+    }
+    return result;
+  }
+
+  std::array<mjtNum, 3> MouthToTarget() {
+    // data.geom_xpos['target'] - data.geom_xpos['mouth']
+    std::array<mjtNum, 3> mouth_to_target_global;
+    for (int i = 0; i < 3; i++) {
+      mouth_to_target_global[i] = (data_->geom_xpos[id_target_ * 3 + i] -
+                                   data_->geom_xpos[id_mouth_ * 3 + i]);
+    }
+    // mouth_to_target_global.dot(data.geom_xmat['mouth'].reshape(3, 3))
+    std::array<mjtNum, 3> mouth_to_target;
+    for (int i = 0; i < 3; i++) {
+      mouth_to_target[i] =
+          mouth_to_target_global[0] * data_->geom_xmat[id_mouth_ * 9 + i + 0] +
+          mouth_to_target_global[1] * data_->geom_xmat[id_mouth_ * 9 + i + 3] +
+          mouth_to_target_global[2] * data_->geom_xmat[id_mouth_ * 9 + i + 6];
+    }
+    return mouth_to_target;
+  }
+};
+
+using FishEnvPool = AsyncEnvPool<FishEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_FISH_H_
diff --git a/envpool/mujoco/dmc/manipulator.h b/envpool/mujoco/dmc/manipulator.h
index 703c6b46..c8cd487b 100644
--- a/envpool/mujoco/dmc/manipulator.h
+++ b/envpool/mujoco/dmc/manipulator.h
@@ -103,10 +103,10 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
   // object_angle, qvel_objx
 
   // ids
-  std::array<int, 8> id_arm_joints_, id_arm_qpos_;
+  std::array<int, 8> id_arm_joints_, id_arm_qpos_, id_arm_qvel_;
   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, 3> id_qpos_object_joints_, id_qvel_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_;
@@ -136,7 +136,7 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
             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_object_x_(GetQvelId(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")),
@@ -152,6 +152,7 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
       id_arm_joints_[i] =
           mj_name2id(model_, mjOBJ_JOINT, kArmJoints[i].c_str());
       id_arm_qpos_[i] = GetQposId(model_, kArmJoints[i]);
+      id_arm_qvel_[i] = GetQvelId(model_, kArmJoints[i]);
     }
     std::array<std::string, 3> object_joints;
     if (use_peg_) {
@@ -160,7 +161,8 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
       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]);
+      id_qpos_object_joints_[i] = GetQposId(model_, object_joints[i]);
+      id_qvel_object_joints_[i] = GetQvelId(model_, object_joints[i]);
     }
     for (std::size_t i = 0; i < kTouchSensors.size(); ++i) {
       id_touch_sensors_[i] = GetSensorId(model_, kTouchSensors[i]);
@@ -250,9 +252,9 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
         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;
+      data_->qpos[id_qpos_object_joints_[0]] = random_info_[4] = object_x;
+      data_->qpos[id_qpos_object_joints_[1]] = random_info_[5] = object_z;
+      data_->qpos[id_qpos_object_joints_[2]] = random_info_[6] = object_angle;
       // Check for collisions.
       PhysicsAfterReset();
       penetrating = data_->ncon > 0;
@@ -294,16 +296,16 @@ class ManipulatorEnv : public Env<ManipulatorEnvSpec>, public MujocoEnv {
   // 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]];
+    for (std::size_t i = 0; i < id_arm_qvel_.size(); i++) {
+      joint[i] = data_->qvel[id_arm_qvel_[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]];
+    for (std::size_t i = 0; i < id_qvel_object_joints_.size(); i++) {
+      joint[i] = data_->qvel[id_qvel_object_joints_[i]];
     }
     return joint;
   }
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index e07a4440..41b2b5e6 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -17,6 +17,7 @@
 #include "envpool/mujoco/dmc/ball_in_cup.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/finger.h"
+#include "envpool/mujoco/dmc/fish.h"
 #include "envpool/mujoco/dmc/hopper.h"
 #include "envpool/mujoco/dmc/humanoid.h"
 #include "envpool/mujoco/dmc/manipulator.h"
@@ -37,6 +38,9 @@ using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
 using DmcFingerEnvSpec = PyEnvSpec<mujoco_dmc::FingerEnvSpec>;
 using DmcFingerEnvPool = PyEnvPool<mujoco_dmc::FingerEnvPool>;
 
+using DmcFishEnvSpec = PyEnvSpec<mujoco_dmc::FishEnvSpec>;
+using DmcFishEnvPool = PyEnvPool<mujoco_dmc::FishEnvPool>;
+
 using DmcHopperEnvSpec = PyEnvSpec<mujoco_dmc::HopperEnvSpec>;
 using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 
@@ -63,6 +67,7 @@ PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcBallInCupEnvSpec, DmcBallInCupEnvPool)
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
+  REGISTER(m, DmcFishEnvSpec, DmcFishEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
   REGISTER(m, DmcHumanoidEnvSpec, DmcHumanoidEnvPool)
   REGISTER(m, DmcManipulatorEnvSpec, DmcManipulatorEnvPool)
diff --git a/envpool/mujoco/dmc/pendulum.h b/envpool/mujoco/dmc/pendulum.h
index b97f9478..422ab9ed 100644
--- a/envpool/mujoco/dmc/pendulum.h
+++ b/envpool/mujoco/dmc/pendulum.h
@@ -74,7 +74,7 @@ class PendulumEnv : public Env<PendulumEnvSpec>, public MujocoEnv {
                                  spec.config["task_name"_]),
                   spec.config["frame_skip"_],
                   spec.config["max_episode_steps"_]),
-        id_hinge_(GetQposId(model_, "hinge")),
+        id_hinge_(GetQvelId(model_, "hinge")),
         id_pole_(mj_name2id(model_, mjOBJ_XBODY, "pole")),
         dist_uniform_(-M_PI, M_PI) {
     const std::string& task_name = spec.config["task_name"_];
diff --git a/envpool/mujoco/dmc/registration.py b/envpool/mujoco/dmc/registration.py
index 5a875217..5b225999 100644
--- a/envpool/mujoco/dmc/registration.py
+++ b/envpool/mujoco/dmc/registration.py
@@ -30,6 +30,8 @@
   ("finger", "spin"),
   ("finger", "turn_easy"),
   ("finger", "turn_hard"),
+  ("fish", "swim"),
+  ("fish", "upright"),
   ("hopper", "hop"),
   ("hopper", "stand"),
   ("humanoid", "stand"),
diff --git a/envpool/mujoco/dmc/utils.cc b/envpool/mujoco/dmc/utils.cc
index bf6caa2b..ce569b12 100644
--- a/envpool/mujoco/dmc/utils.cc
+++ b/envpool/mujoco/dmc/utils.cc
@@ -59,6 +59,10 @@ int GetQposId(mjModel* model, const std::string& name) {
   return model->jnt_qposadr[mj_name2id(model, mjOBJ_JOINT, name.c_str())];
 }
 
+int GetQvelId(mjModel* model, const std::string& name) {
+  return model->jnt_dofadr[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())];
 }
diff --git a/envpool/mujoco/dmc/utils.h b/envpool/mujoco/dmc/utils.h
index 91a0bfa0..045529da 100644
--- a/envpool/mujoco/dmc/utils.h
+++ b/envpool/mujoco/dmc/utils.h
@@ -33,6 +33,7 @@ std::string XMLRemoveByBodyName(const std::string& content,
 
 // the following id is not 1 on 1 mapping
 int GetQposId(mjModel* model, const std::string& name);
+int GetQvelId(mjModel* model, const std::string& name);
 int GetSensorId(mjModel* model, const std::string& name);
 
 // rewards


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 091ff28d..f0ab68f7 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -134,6 +134,8 @@ def test_make_mujoco_dmc(self) -> None:
         "FingerSpin-v1",
         "FingerTurnEasy-v1",
         "FingerTurnHard-v1",
+        "FishSwim-v1",
+        "FishUpright-v1",
         "HopperHop-v1",
         "HopperStand-v1",
         "HumanoidStand-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 33cb19e8..486e8c17 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -30,6 +30,8 @@
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
   DmcFingerEnvSpec,
+  DmcFishDMEnvPool,
+  DmcFishEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
   DmcHumanoidDMEnvPool,
@@ -76,6 +78,11 @@ def reset_state(
         target = ts.observation.target[0]
         env.physics.named.model.geom_pos["target", "x"] = target[0]
         env.physics.named.model.geom_pos["target", "y"] = target[1]
+      elif domain == "fish" and task == "swim":
+        target = ts.observation.target0[0]
+        env.physics.named.model.geom_pos["target", "x"] = target[0]
+        env.physics.named.model.geom_pos["target", "y"] = target[1]
+        env.physics.named.model.geom_pos["target", "z"] = target[2]
       elif domain in ["finger", "ball_in_cup", "humanoid"]:
         if domain == "finger" and task in ["turn_easy", "turn_hard"]:
           target_angle = ts.observation.target[0][0]
@@ -197,6 +204,11 @@ def test_finger(self) -> None:
       DmcFingerDMEnvPool
     )
 
+  def test_fish(self) -> None:
+    self.run_align_check_entry(
+      "fish", ["swim", "upright"], DmcFishEnvSpec, DmcFishDMEnvPool
+    )
+
   def test_hopper(self) -> None:
     self.run_align_check_entry(
       "hopper", ["hop", "stand"], DmcHopperEnvSpec, DmcHopperDMEnvPool
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index e551c4b5..383bc3f6 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -28,6 +28,8 @@
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
   DmcFingerEnvSpec,
+  DmcFishDMEnvPool,
+  DmcFishEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
   DmcHumanoidDMEnvPool,
@@ -67,7 +69,7 @@ def check(
       spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=1, task_name=task))
     )
     act_spec = env0.action_spec()
-    for _ in range(3000):
+    for t in range(3000):
       action = np.array(
         [
           np.random.uniform(
@@ -87,7 +89,7 @@ def check(
         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), (k, o0, o2))
+          self.assertFalse(np.allclose(o0, o2), (t, k, o0, o2))
 
   def test_acrobot(self) -> None:
     obs_keys = ["orientations", "velocity"]
@@ -119,6 +121,17 @@ def test_finger(self) -> None:
     for task in ["turn_easy", "turn_hard"]:
       self.check(DmcFingerEnvSpec, DmcFingerDMEnvPool, task, obs_keys)
 
+  def test_fish(self) -> None:
+    obs_keys = ["joint_angles", "upright", "velocity"]
+    for task in ["swim", "upright"]:
+      self.check(
+        DmcFishEnvSpec,
+        DmcFishDMEnvPool,
+        task,
+        obs_keys + (["target"] if task == "swim" else []),
+        blacklist=["joint_angles"],
+      )
+
   def test_hopper(self) -> None:
     obs_keys = ["position", "velocity", "touch"]
     for task in ["stand", "hop"]: