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 9e8d759b..6218412c 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -92,6 +92,19 @@ PendulumSwingup-v1
 - ``max_episode_stes``: 1000;
 
 
+PointMassEasy-v1, PointMassHard-v1
+----------------------------------
+
+`dm_control suite point mass source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/point_mass.py>`_
+
+- Observation spec: a namedtuple with three keys: ``position (2)``,
+  ``velocity (2)``;
+- Action spec: ``(1)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_stes``: 1000;
+
+
 ReacherEasy-v1, ReacherHard-v1
 ------------------------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 635d0783..a6d9d260 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -69,6 +69,7 @@ cc_library(
         "dmc/hopper.h",
         "dmc/mujoco_env.h",
         "dmc/pendulum.h",
+        "dmc/point_mass.h",
         "dmc/reacher.h",
         "dmc/utils.h",
         "dmc/walker.h",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 3f82fa13..e417ca3f 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -27,6 +27,8 @@
   _DmcHopperEnvSpec,
   _DmcPendulumEnvPool,
   _DmcPendulumEnvSpec,
+  _DmcPointMassEnvPool,
+  _DmcPointMassEnvSpec,
   _DmcReacherEnvPool,
   _DmcReacherEnvSpec,
   _DmcWalkerEnvPool,
@@ -75,6 +77,9 @@
 DmcPendulumEnvSpec, DmcPendulumDMEnvPool, DmcPendulumGymEnvPool = py_env(
   _DmcPendulumEnvSpec, _DmcPendulumEnvPool
 )
+DmcPointMassEnvSpec, DmcPointMassDMEnvPool, DmcPointMassGymEnvPool = py_env(
+  _DmcPointMassEnvSpec, _DmcPointMassEnvPool
+)
 DmcReacherEnvSpec, DmcReacherDMEnvPool, DmcReacherGymEnvPool = py_env(
   _DmcReacherEnvSpec, _DmcReacherEnvPool
 )
@@ -140,6 +145,9 @@
   "DmcPendulumEnvSpec",
   "DmcPendulumDMEnvPool",
   "DmcPendulumGymEnvPool",
+  "DmcPointMassEnvSpec",
+  "DmcPointMassDMEnvPool",
+  "DmcPointMassGymEnvPool",
   "DmcReacherEnvSpec",
   "DmcReacherDMEnvPool",
   "DmcReacherGymEnvPool",
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 1d1cff5a..eb82ffa9 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -18,6 +18,7 @@
 #include "envpool/mujoco/dmc/finger.h"
 #include "envpool/mujoco/dmc/hopper.h"
 #include "envpool/mujoco/dmc/pendulum.h"
+#include "envpool/mujoco/dmc/point_mass.h"
 #include "envpool/mujoco/dmc/reacher.h"
 #include "envpool/mujoco/dmc/walker.h"
 
@@ -36,6 +37,9 @@ using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 using DmcPendulumEnvSpec = PyEnvSpec<mujoco_dmc::PendulumEnvSpec>;
 using DmcPendulumEnvPool = PyEnvPool<mujoco_dmc::PendulumEnvPool>;
 
+using DmcPointMassEnvSpec = PyEnvSpec<mujoco_dmc::PointMassEnvSpec>;
+using DmcPointMassEnvPool = PyEnvPool<mujoco_dmc::PointMassEnvPool>;
+
 using DmcReacherEnvSpec = PyEnvSpec<mujoco_dmc::ReacherEnvSpec>;
 using DmcReacherEnvPool = PyEnvPool<mujoco_dmc::ReacherEnvPool>;
 
@@ -48,6 +52,7 @@ PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
   REGISTER(m, DmcPendulumEnvSpec, DmcPendulumEnvPool)
+  REGISTER(m, DmcPointMassEnvSpec, DmcPointMassEnvPool)
   REGISTER(m, DmcReacherEnvSpec, DmcReacherEnvPool)
   REGISTER(m, DmcWalkerEnvSpec, DmcWalkerEnvPool)
 }
diff --git a/envpool/mujoco/dmc/pendulum.h b/envpool/mujoco/dmc/pendulum.h
index 3153dd9b..45d130cf 100644
--- a/envpool/mujoco/dmc/pendulum.h
+++ b/envpool/mujoco/dmc/pendulum.h
@@ -41,7 +41,7 @@ class PendulumEnvFns {
  public:
   static decltype(auto) DefaultConfig() {
     return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
-                    "task_name"_.Bind(std::string("stand")));
+                    "task_name"_.Bind(std::string("swingup")));
   }
   template <typename Config>
   static decltype(auto) StateSpec(const Config& conf) {
diff --git a/envpool/mujoco/dmc/point_mass.h b/envpool/mujoco/dmc/point_mass.h
new file mode 100644
index 00000000..609e040f
--- /dev/null
+++ b/envpool/mujoco/dmc/point_mass.h
@@ -0,0 +1,188 @@
+/*
+ * 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/point_mass.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_POINT_MASS_H_
+#define ENVPOOL_MUJOCO_DMC_POINT_MASS_H_
+
+#include <algorithm>
+#include <limits>
+#include <memory>
+#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 GetPointMassXML(const std::string& base_path,
+                            const std::string& task_name) {
+  return GetFileContent(base_path, "point_mass.xml");
+}
+
+class PointMassEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
+                    "task_name"_.Bind(std::string("easy")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:position"_.Bind(Spec<mjtNum>({2})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({2})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({2})),
+                    "info:wrap_prm"_.Bind(Spec<mjtNum>({4})),
+#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, 2}, {-1.0, 1.0})));
+  }
+};
+
+using PointMassEnvSpec = EnvSpec<PointMassEnvFns>;
+
+class PointMassEnv : public Env<PointMassEnvSpec>, public MujocoEnv {
+ protected:
+  bool randomize_gains_;
+  int id_geom_target_, id_geom_pointmass_;
+  std::normal_distribution<> dist_normal_;
+#ifdef ENVPOOL_TEST
+  std::unique_ptr<mjtNum> wrap_prm_;
+#endif
+
+ public:
+  PointMassEnv(const Spec& spec, int env_id)
+      : Env<PointMassEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetPointMassXML(spec.config["base_path"_],
+                                  spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+
+        id_geom_target_(mj_name2id(model_, mjOBJ_GEOM, "target")),
+        id_geom_pointmass_(mj_name2id(model_, mjOBJ_GEOM, "pointmass")),
+        dist_normal_(0, 1) {
+    std::string task_name = spec.config["task_name"_];
+    if (task_name == "easy") {
+      randomize_gains_ = false;
+    } else if (task_name == "hard") {
+      randomize_gains_ = true;
+    } else {
+      throw std::runtime_error("Unknown task_name for dmc point_mass.");
+    }
+    wrap_prm_.reset(new mjtNum[model_->nwrap]);
+  }
+
+  void TaskInitializeEpisode() override {
+    RandomizeLimitedAndRotationalJoints(&gen_);
+    if (randomize_gains_) {
+      const auto& dir1 = GetDir();
+      bool parallel = true;
+      std::array<mjtNum, 2> dir2;
+      while (parallel) {
+        dir2 = GetDir();
+        parallel = std::abs(dir1[0] * dir2[0] + dir1[1] * dir2[1]) > 0.9;
+      }
+      model_->wrap_prm[0] = dir1[0];
+      model_->wrap_prm[1] = dir1[1];
+      model_->wrap_prm[2] = dir2[0];
+      model_->wrap_prm[3] = dir2[1];
+    }
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+    std::memcpy(wrap_prm_.get(), model_->wrap_prm,
+                sizeof(mjtNum) * model_->nwrap);
+#endif
+  }
+
+  std::array<mjtNum, 2> GetDir() {
+    std::array<mjtNum, 2> dir = {dist_normal_(gen_), dist_normal_(gen_)};
+    mjtNum norm_of_dir = std::sqrt(dir[0] * dir[0] + dir[1] * dir[1]);
+    return {dir[0] / norm_of_dir, dir[1] / norm_of_dir};
+  }
+
+  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 {
+    mjtNum target_size = model_->geom_size[id_geom_target_ * 3 + 0];
+    mjtNum near_target =
+        RewardTolerance(MassToTargetDist(), 0, target_size, target_size);
+    mjtNum control_reward =
+        (RewardTolerance(data_->ctrl[0], 0, 0, 1, 0, SigmoidType::kQuadratic) +
+         RewardTolerance(data_->ctrl[1], 0, 0, 1, 0, SigmoidType::kQuadratic)) /
+        2;
+    mjtNum small_control = (control_reward + 4) / 5;
+    return static_cast<float>(near_target * small_control);
+  }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  mjtNum MassToTargetDist() {
+    // return np.linalg.norm(self.mass_to_target())
+    const auto& mass_to_target = MassToTarget();
+    return std::sqrt(mass_to_target[0] * mass_to_target[0] +
+                     mass_to_target[1] * mass_to_target[1] +
+                     mass_to_target[2] * mass_to_target[2]);
+  }
+
+  std::array<mjtNum, 3> MassToTarget() {
+    // return (self.named.data.geom_xpos['target'] -
+    //         self.named.data.geom_xpos['pointmass'])
+    return {data_->geom_xpos[id_geom_target_ * 3 + 0] -
+                data_->geom_xpos[id_geom_pointmass_ * 3 + 0],
+            data_->geom_xpos[id_geom_target_ * 3 + 1] -
+                data_->geom_xpos[id_geom_pointmass_ * 3 + 1],
+            data_->geom_xpos[id_geom_target_ * 3 + 2] -
+                data_->geom_xpos[id_geom_pointmass_ * 3 + 2]};
+  }
+
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    state["obs:position"_].Assign(data_->qpos, model_->nq);
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+    // info for check alignment
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+    state["info:wrap_prm"_].Assign(wrap_prm_.get(), model_->nwrap);
+#endif
+  }
+};
+
+using PointMassEnvPool = AsyncEnvPool<PointMassEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_POINT_MASS_H_
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 204d744f..2aa0926e 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -29,6 +29,8 @@
   ("hopper", "hop"),
   ("hopper", "stand"),
   ("pendulum", "swingup"),
+  ("point_mass", "easy"),
+  ("point_mass", "hard"),
   ("reacher", "easy"),
   ("reacher", "hard"),
   ("walker", "run"),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 668601e2..c5c9553e 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -129,6 +129,8 @@ def test_make_mujoco_dmc(self) -> None:
         "HopperHop-v1",
         "HopperStand-v1",
         "PendulumSwingup-v1",
+        "PointMassEasy-v1",
+        "PointMassHard-v1",
         "ReacherEasy-v1",
         "ReacherHard-v1",
         "WalkerRun-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 9bc9de5b..036f8adf 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -32,6 +32,8 @@
   DmcHopperEnvSpec,
   DmcPendulumDMEnvPool,
   DmcPendulumEnvSpec,
+  DmcPointMassDMEnvPool,
+  DmcPointMassEnvSpec,
   DmcReacherDMEnvPool,
   DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
@@ -78,6 +80,8 @@ def reset_state(
           env.physics.named.model.site_pos["target",
                                            ["x", "z"]] = target_x, target_z
         env.physics.after_reset()
+      elif domain == "point_mass":
+        env.physics.model.wrap_prm = ts.observation.wrap_prm
 
   def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
     return np.random.uniform(
@@ -111,7 +115,7 @@ def run_align_check(
         for k in obs_spec:
           np.testing.assert_allclose(o0[k], getattr(o1, k)[0])
         np.testing.assert_allclose(ts0.step_type, ts1.step_type[0])
-        np.testing.assert_allclose(ts0.reward, ts1.reward[0])
+        np.testing.assert_allclose(ts0.reward, ts1.reward[0], atol=1e-8)
         np.testing.assert_allclose(ts0.discount, ts1.discount[0])
 
   def run_align_check_entry(
@@ -149,6 +153,12 @@ def test_pendulum(self) -> None:
       "pendulum", ["swingup"], DmcPendulumEnvSpec, DmcPendulumDMEnvPool
     )
 
+  def test_point_mass(self) -> None:
+    self.run_align_check_entry(
+      "point_mass", ["easy", "hard"], DmcPointMassEnvSpec,
+      DmcPointMassDMEnvPool
+    )
+
   def test_reacher(self) -> None:
     self.run_align_check_entry(
       "reacher", ["easy", "hard"], DmcReacherEnvSpec, DmcReacherDMEnvPool
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index 945f6ec7..fccc3895 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -30,6 +30,8 @@
   DmcHopperEnvSpec,
   DmcPendulumDMEnvPool,
   DmcPendulumEnvSpec,
+  DmcPointMassDMEnvPool,
+  DmcPointMassEnvSpec,
   DmcReacherDMEnvPool,
   DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
@@ -116,6 +118,17 @@ def test_pendulum(self) -> None:
     for task in ["swingup"]:
       self.check(DmcPendulumEnvSpec, DmcPendulumDMEnvPool, task, obs_keys)
 
+  def test_point_mass(self) -> None:
+    obs_keys = ["position", "velocity"]
+    for task in ["easy", "hard"]:
+      self.check(
+        DmcPointMassEnvSpec,
+        DmcPointMassDMEnvPool,
+        task,
+        obs_keys,
+        blacklist=["velocity"] if task == "easy" else [],
+      )
+
   def test_reacher(self) -> None:
     obs_keys = ["position", "to_target", "velocity"]
     for task in ["easy", "hard"]: