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 2b79f461..78d7c47c 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -46,6 +46,19 @@ HopperStand-v1, HopperHop-v1
 - ``max_episode_steps``: 1000;
 
 
+ReacherEasy-v1, ReacherHard-v1
+------------------------------
+
+`dm_control suite reacher source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/reacher.py>`_
+
+- Observation spec: a namedtuple with three keys: ``position (2)``,
+  ``to_target (2)`` and ``velocity (2)``;
+- Action spec: ``(2)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_steps``: 1000;
+
+
 WalkerRun-v1, WalkerStand-v1, WalkerWalk-v1
 -------------------------------------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 2f3a1771..6d3fc718 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -66,6 +66,7 @@ cc_library(
         "dmc/cheetah.h",
         "dmc/hopper.h",
         "dmc/mujoco_env.h",
+        "dmc/reacher.h",
         "dmc/utils.h",
         "dmc/walker.h",
     ],
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 0fdb61ab..f0fcc877 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -21,6 +21,8 @@
   _DmcCheetahEnvSpec,
   _DmcHopperEnvPool,
   _DmcHopperEnvSpec,
+  _DmcReacherEnvPool,
+  _DmcReacherEnvSpec,
   _DmcWalkerEnvPool,
   _DmcWalkerEnvSpec,
 )
@@ -58,6 +60,9 @@
 DmcHopperEnvSpec, DmcHopperDMEnvPool, DmcHopperGymEnvPool = py_env(
   _DmcHopperEnvSpec, _DmcHopperEnvPool
 )
+DmcReacherEnvSpec, DmcReacherDMEnvPool, DmcReacherGymEnvPool = py_env(
+  _DmcReacherEnvSpec, _DmcReacherEnvPool
+)
 DmcWalkerEnvSpec, DmcWalkerDMEnvPool, DmcWalkerGymEnvPool = py_env(
   _DmcWalkerEnvSpec, _DmcWalkerEnvPool
 )
@@ -114,6 +119,9 @@
   "DmcHopperEnvSpec",
   "DmcHopperDMEnvPool",
   "DmcHopperGymEnvPool",
+  "DmcReacherEnvSpec",
+  "DmcReacherDMEnvPool",
+  "DmcReacherGymEnvPool",
   "DmcWalkerEnvSpec",
   "DmcWalkerDMEnvPool",
   "DmcWalkerGymEnvPool",
diff --git a/envpool/mujoco/dmc/cheetah.h b/envpool/mujoco/dmc/cheetah.h
index 88fb28bd..4019a8ff 100644
--- a/envpool/mujoco/dmc/cheetah.h
+++ b/envpool/mujoco/dmc/cheetah.h
@@ -13,6 +13,7 @@
  * 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/cheetah.py
 
 #ifndef ENVPOOL_MUJOCO_DMC_CHEETAH_H_
 #define ENVPOOL_MUJOCO_DMC_CHEETAH_H_
@@ -58,7 +59,6 @@ class CheetahEnvFns {
 
 using CheetahEnvSpec = EnvSpec<CheetahEnvFns>;
 
-// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/cheetah.py#L60
 class CheetahEnv : public Env<CheetahEnvSpec>, public MujocoEnv {
  protected:
   const mjtNum kRunSpeed = 10;
diff --git a/envpool/mujoco/dmc/hopper.h b/envpool/mujoco/dmc/hopper.h
index 17e904e1..7ba74830 100644
--- a/envpool/mujoco/dmc/hopper.h
+++ b/envpool/mujoco/dmc/hopper.h
@@ -13,6 +13,7 @@
  * 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/hopper.py
 
 #ifndef ENVPOOL_MUJOCO_DMC_HOPPER_H_
 #define ENVPOOL_MUJOCO_DMC_HOPPER_H_
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 769e8932..687526f4 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -15,13 +15,17 @@
 #include "envpool/core/py_envpool.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/hopper.h"
+#include "envpool/mujoco/dmc/reacher.h"
 #include "envpool/mujoco/dmc/walker.h"
 
+using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
+using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
+
 using DmcHopperEnvSpec = PyEnvSpec<mujoco_dmc::HopperEnvSpec>;
 using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 
-using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
-using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
+using DmcReacherEnvSpec = PyEnvSpec<mujoco_dmc::ReacherEnvSpec>;
+using DmcReacherEnvPool = PyEnvPool<mujoco_dmc::ReacherEnvPool>;
 
 using DmcWalkerEnvSpec = PyEnvSpec<mujoco_dmc::WalkerEnvSpec>;
 using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
@@ -29,5 +33,6 @@ using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
+  REGISTER(m, DmcReacherEnvSpec, DmcReacherEnvPool)
   REGISTER(m, DmcWalkerEnvSpec, DmcWalkerEnvPool)
 }
diff --git a/envpool/mujoco/dmc/reacher.h b/envpool/mujoco/dmc/reacher.h
new file mode 100644
index 00000000..0f594887
--- /dev/null
+++ b/envpool/mujoco/dmc/reacher.h
@@ -0,0 +1,158 @@
+/*
+ * 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/reacher.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_REACHER_H_
+#define ENVPOOL_MUJOCO_DMC_REACHER_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 GetReacherXML(const std::string& base_path,
+                          const std::string& task_name) {
+  return GetFileContent(base_path, "reacher.xml");
+}
+
+class ReacherEnvFns {
+ 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:to_target"_.Bind(Spec<mjtNum>({2})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({2})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({2})),
+                    "info:target"_.Bind(Spec<mjtNum>({30})),
+#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 ReacherEnvSpec = EnvSpec<ReacherEnvFns>;
+
+class ReacherEnv : public Env<ReacherEnvSpec>, public MujocoEnv {
+ protected:
+  const mjtNum kBigTarget = 0.05;
+  const mjtNum kSmallTarget = 0.015;
+  mjtNum target_size_;
+  std::uniform_real_distribution<> dist_uniform_;
+#ifdef ENVPOOL_TEST
+  std::array<mjtNum, 2> target_;
+#endif
+
+ public:
+  ReacherEnv(const Spec& spec, int env_id)
+      : Env<ReacherEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_],
+            GetReacherXML(spec.config["base_path"_], spec.config["task_name"_]),
+            spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
+        dist_uniform_(0, 1) {
+    std::string task_name = spec.config["task_name"_];
+    if (task_name == "easy") {
+      target_size_ = kBigTarget;
+    } else if (task_name == "hard") {
+      target_size_ = kSmallTarget;
+    } else {
+      throw std::runtime_error("Unknown task_name for dmc reacher.");
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    model_->geom_size[6 * 3] = target_size_;
+    RandomizeLimitedAndRotationalJoints(&gen_);
+    mjtNum angle = dist_uniform_(gen_) * 2 * M_PI;
+    mjtNum radius = dist_uniform_(gen_) * 0.15 + 0.05;
+    model_->geom_pos[6 * 3 + 0] = radius * std::sin(angle);
+    model_->geom_pos[6 * 3 + 1] = radius * std::cos(angle);
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+    target_[0] = model_->geom_pos[6 * 3 + 0];
+    target_[1] = model_->geom_pos[6 * 3 + 1];
+#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 {
+    mjtNum radii = model_->geom_size[6 * 3] + model_->geom_size[9 * 3];
+    return static_cast<float>(RewardTolerance(FingerToTargetDist(), 0, radii));
+  }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    state["obs:position"_].Assign(data_->qpos, model_->nq);
+    auto finger = FingerToTarget();
+    state["obs:to_target"_].Assign(finger.begin(), finger.size());
+    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:target"_].Assign(target_.begin(), target_.size());
+#endif
+  }
+
+  std::array<mjtNum, 2> FingerToTarget() {
+    std::array<mjtNum, 2> finger;
+    finger[0] = data_->geom_xpos[6 * 3 + 0] - data_->geom_xpos[9 * 3 + 0];
+    finger[1] = data_->geom_xpos[6 * 3 + 1] - data_->geom_xpos[9 * 3 + 1];
+    return finger;
+  }
+  mjtNum FingerToTargetDist() {
+    std::array<mjtNum, 2> finger = FingerToTarget();
+    return std::sqrt(finger[0] * finger[0] + finger[1] * finger[1]);
+  }
+};
+
+using ReacherEnvPool = AsyncEnvPool<ReacherEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_REACHER_H_
diff --git a/envpool/mujoco/dmc/walker.h b/envpool/mujoco/dmc/walker.h
index 0d1064db..fc646008 100644
--- a/envpool/mujoco/dmc/walker.h
+++ b/envpool/mujoco/dmc/walker.h
@@ -13,6 +13,7 @@
  * 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/walker.py
 
 #ifndef ENVPOOL_MUJOCO_DMC_WALKER_H_
 #define ENVPOOL_MUJOCO_DMC_WALKER_H_
@@ -28,7 +29,6 @@
 #include "envpool/mujoco/dmc/mujoco_env.h"
 #include "envpool/mujoco/dmc/utils.h"
 
-// https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/walker.py
 namespace mujoco_dmc {
 
 std::string GetWalkerXML(const std::string& base_path,
@@ -84,7 +84,7 @@ class WalkerEnv : public Env<WalkerEnvSpec>, public MujocoEnv {
     } else if (task_name == "run") {
       move_speed_ = kRunSpeed;
     } else {
-      throw std::runtime_error("Unknown task_name for dmc hopper.");
+      throw std::runtime_error("Unknown task_name for dmc walker.");
     }
   }
 
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 4a9c7f19..cd575bcd 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -24,6 +24,8 @@
   ("cheetah", "run"),
   ("hopper", "hop"),
   ("hopper", "stand"),
+  ("reacher", "easy"),
+  ("reacher", "hard"),
   ("walker", "run"),
   ("walker", "stand"),
   ("walker", "walk"),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index c110cf87..fbb76d50 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -123,6 +123,8 @@ def test_make_mujoco_dmc(self) -> None:
         "CheetahRun-v1",
         "HopperHop-v1",
         "HopperStand-v1",
+        "ReacherEasy-v1",
+        "ReacherHard-v1",
         "WalkerRun-v1",
         "WalkerStand-v1",
         "WalkerWalk-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 050c8abb..0a5ed667 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -26,6 +26,8 @@
   DmcCheetahEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcReacherDMEnvPool,
+  DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
   DmcWalkerEnvSpec,
 )
@@ -49,13 +51,17 @@ def run_space_check(self, env0: dm_env.Environment, env1: Any) -> None:
   def reset_state(
     self, env: dm_env.Environment, ts: dm_env.TimeStep, task: str
   ) -> None:
-    # manually reset
+    # manually reset, mimic initialize_episode
     with env.physics.reset_context():
       env.physics.data.qpos = ts.observation.qpos0[0]
       if task == "cheetah":
         for _ in range(200):
           env.physics.step()
         env.physics.data.time = 0
+      elif task == "reacher":
+        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]
 
   def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
     return np.random.uniform(
@@ -101,14 +107,19 @@ def run_align_check_entry(
       self.run_space_check(env0, env1)
       self.run_align_check(env0, env1, domain)
 
+  def test_cheetah(self) -> None:
+    self.run_align_check_entry(
+      "cheetah", ["run"], DmcCheetahEnvSpec, DmcCheetahDMEnvPool
+    )
+
   def test_hopper(self) -> None:
     self.run_align_check_entry(
       "hopper", ["hop", "stand"], DmcHopperEnvSpec, DmcHopperDMEnvPool
     )
 
-  def test_cheetah(self) -> None:
+  def test_reacher(self) -> None:
     self.run_align_check_entry(
-      "cheetah", ["run"], DmcCheetahEnvSpec, DmcCheetahDMEnvPool
+      "reacher", ["easy", "hard"], DmcReacherEnvSpec, DmcReacherDMEnvPool
     )
 
   def test_walker(self) -> None:
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index bb7de3e9..8449d786 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -24,6 +24,8 @@
   DmcCheetahEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcReacherDMEnvPool,
+  DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
   DmcWalkerEnvSpec,
 )
@@ -69,14 +71,20 @@ def check(
         if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
           self.assertFalse(np.allclose(o0, o2), (o0, o2))
 
+  def test_cheetah(self) -> None:
+    obs_keys = ["position", "velocity"]
+    for task in ["run"]:
+      self.check(DmcCheetahEnvSpec, DmcCheetahDMEnvPool, task, obs_keys)
+
   def test_hopper(self) -> None:
     obs_keys = ["position", "velocity", "touch"]
-    self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, "stand", obs_keys)
-    self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, "hop", obs_keys)
+    for task in ["stand", "hop"]:
+      self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, task, obs_keys)
 
-  def test_cheetah(self) -> None:
-    obs_keys = ["position", "velocity"]
-    self.check(DmcCheetahEnvSpec, DmcCheetahDMEnvPool, "run", obs_keys)
+  def test_reacher(self) -> None:
+    obs_keys = ["position", "to_target", "velocity"]
+    for task in ["easy", "hard"]:
+      self.check(DmcReacherEnvSpec, DmcReacherDMEnvPool, task, obs_keys)
 
   def test_walker(self) -> None:
     obs_keys = ["orientations", "height", "velocity"]