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 e0b1cf4d..9e8d759b 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -20,11 +20,24 @@ converted into ``DomainNameTaskName-v1`` in envpool, e.g.,
   envpool.make_dm("BallInCupCatch-v1", num_envs=1)
 
 
+BallInCupCatch-v1
+-----------------
+
+`dm_control suite ball-in-cup source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/ball_in_cup.py>`_
+
+- Observation spec: a namedtuple with two keys: ``position (4)`` and
+  ``velocity (4)``;
+- Action spec: ``(2)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_steps``: 1000;
+
+
 CheetahRun-v1
 -------------
 
 `dm_control suite cheetah source code
-<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/cheetah.py>`_
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/cheetah.py>`_
 
 - Observation spec: a namedtuple with two keys: ``position (8)`` and
   ``velocity (9)``;
@@ -37,7 +50,7 @@ FingerSpin-v1, FingerTurnEasy-v1, FingerTurnHard-v1
 ---------------------------------------------------
 
 `dm_control suite finger source code
-<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/finger.py>`_
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/finger.py>`_
 
 - Observation spec: a namedtuple with five keys: ``position (4)``,
   ``velocity (3)``, ``touch (2)``, ``target_position (2)``,
@@ -57,7 +70,7 @@ HopperStand-v1, HopperHop-v1
 ----------------------------
 
 `dm_control suite hopper source code
-<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/hopper.py>`_
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/hopper.py>`_
 
 - Observation spec: a namedtuple with three keys: ``position (6)``,
   ``velocity (7)``, ``touch (2)``;
@@ -70,7 +83,7 @@ PendulumSwingup-v1
 ------------------
 
 `dm_control suite pendulum source code
-<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/pendulum.py>`_
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/pendulum.py>`_
 
 - Observation spec: a namedtuple with three keys: ``orientations (2)``,
   ``velocity (1)``;
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index 73058f92..635d0783 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -63,6 +63,7 @@ pybind_extension(
 cc_library(
     name = "mujoco_dmc_env",
     hdrs = [
+        "dmc/ball_in_cup.h",
         "dmc/cheetah.h",
         "dmc/finger.h",
         "dmc/hopper.h",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 807c9cce..3f82fa13 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -17,6 +17,8 @@
 
 # dmc mujoco
 from .mujoco_dmc_envpool import (
+  _DmcBallInCupEnvPool,
+  _DmcBallInCupEnvSpec,
   _DmcCheetahEnvPool,
   _DmcCheetahEnvSpec,
   _DmcFingerEnvPool,
@@ -58,6 +60,9 @@
 )
 
 # dmc
+DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool, DmcBallInCupGymEnvPool = py_env(
+  _DmcBallInCupEnvSpec, _DmcBallInCupEnvPool
+)
 DmcCheetahEnvSpec, DmcCheetahDMEnvPool, DmcCheetahGymEnvPool = py_env(
   _DmcCheetahEnvSpec, _DmcCheetahEnvPool
 )
@@ -123,6 +128,9 @@
 )
 
 __all__ = [
+  "DmcBallInCupEnvSpec",
+  "DmcBallInCupDMEnvPool",
+  "DmcBallInCupGymEnvPool",
   "DmcCheetahEnvSpec",
   "DmcCheetahDMEnvPool",
   "DmcCheetahGymEnvPool",
diff --git a/envpool/mujoco/dmc/ball_in_cup.h b/envpool/mujoco/dmc/ball_in_cup.h
new file mode 100644
index 00000000..0fff5e2f
--- /dev/null
+++ b/envpool/mujoco/dmc/ball_in_cup.h
@@ -0,0 +1,150 @@
+/*
+ * 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/ball_in_cup.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_BALL_IN_CUP_H_
+#define ENVPOOL_MUJOCO_DMC_BALL_IN_CUP_H_
+
+#include <algorithm>
+#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 GetBallInCupXML(const std::string& base_path,
+                            const std::string& task_name) {
+  return GetFileContent(base_path, "ball_in_cup.xml");
+}
+
+class BallInCupEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(10),
+                    "task_name"_.Bind(std::string("catch")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:position"_.Bind(Spec<mjtNum>({4})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({4})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.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 BallInCupEnvSpec = EnvSpec<BallInCupEnvFns>;
+
+class BallInCupEnv : public Env<BallInCupEnvSpec>, public MujocoEnv {
+ protected:
+  std::uniform_real_distribution<> dist_uniform_;
+
+ public:
+  BallInCupEnv(const Spec& spec, int env_id)
+      : Env<BallInCupEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetBallInCupXML(spec.config["base_path"_],
+                                  spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+        dist_uniform_(0, 1) {}
+
+  void TaskInitializeEpisode() override {
+    while (true) {
+      // Assign a random ball position.
+      data_->qpos[2] = dist_uniform_(gen_) * 0.4 - 0.2;  // ball_x
+      data_->qpos[3] = dist_uniform_(gen_) * 0.3 + 0.2;  // ball_z
+#ifdef ENVPOOL_TEST
+      std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#endif
+      PhysicsAfterReset();
+      if (data_->ncon <= 0) {
+        break;
+      }
+    }
+  }
+
+  bool IsDone() override { return done_; }
+
+  void Reset() override {
+    ControlReset();
+    WriteState();
+  }
+
+  void Step(const Action& action) override {
+    mjtNum* act = static_cast<mjtNum*>(action["action"_].Data());
+    ControlStep(act);
+    WriteState();
+  }
+
+  float TaskGetReward() override { return static_cast<float>(InTarget()); }
+
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  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);
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  std::array<mjtNum, 2> BallToTarget() {
+    // target = self.named.data.site_xpos['target', ['x', 'z']]
+    // ball = self.named.data.xpos['ball', ['x', 'z']]
+    // return target - ball
+    std::array<mjtNum, 2> target{data_->site_xpos[1 * 3],
+                                 data_->site_xpos[1 * 3 + 2]};
+    std::array<mjtNum, 2> ball{data_->site_xpos[2 * 3],
+                               data_->site_xpos[2 * 3 + 2]};
+    return {target[0] - ball[0], target[1] - ball[1]};
+  }
+
+  bool InTarget() {
+    // ball_to_target = abs(self.ball_to_target())
+    // target_size = self.named.model.site_size['target', [0, 2]]
+    // ball_size = self.named.model.geom_size['ball', 0]
+    // return float(all(ball_to_target < target_size - ball_size))
+    const auto& ball_to_target = BallToTarget();
+    std::array<mjtNum, 2> target_size{model_->site_size[1 * 3],
+                                      model_->site_size[1 * 3 + 2]};
+    auto ball_size = model_->geom_size[2 * 3];
+    return std::abs(ball_to_target[0]) < target_size[0] - ball_size &&
+           std::abs(ball_to_target[1]) < target_size[1] - ball_size;
+  }
+};
+
+using BallInCupEnvPool = AsyncEnvPool<BallInCupEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_BALL_IN_CUP_H_
diff --git a/envpool/mujoco/dmc/finger.h b/envpool/mujoco/dmc/finger.h
index 942b9204..f1098c65 100644
--- a/envpool/mujoco/dmc/finger.h
+++ b/envpool/mujoco/dmc/finger.h
@@ -100,7 +100,6 @@ class FingerEnv : public Env<FingerEnvSpec>, public MujocoEnv {
       model_->site_rgba[3 * 0 + 3] = 0;
       model_->site_rgba[3 * 3 + 3] = 0;
       model_->dof_damping[2] = 0.03;
-      SetRandomJointAngles();
     } else {
       mjtNum target_angle = dist_uniform_(gen_);
       mjtNum hinge_x = data_->xanchor[2 * 3 + 0];
@@ -116,11 +115,8 @@ class FingerEnv : public Env<FingerEnvSpec>, public MujocoEnv {
       target_[0] = target_x;
       target_[1] = target_z;
 #endif
-      SetRandomJointAngles();
     }
-#ifdef ENVPOOL_TEST
-    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
-#endif
+    SetRandomJointAngles();
   }
 
   bool IsDone() override { return done_; }
@@ -176,6 +172,9 @@ class FingerEnv : public Env<FingerEnvSpec>, public MujocoEnv {
     int i = 0;
     for (int i = 0; i < max_attempts; i++) {
       RandomizeLimitedAndRotationalJoints(&gen_);
+#ifdef ENVPOOL_TEST
+      std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#endif
       PhysicsAfterReset();
       if (data_->ncon == 0) {
         break;
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 4da1ff75..1d1cff5a 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "envpool/core/py_envpool.h"
+#include "envpool/mujoco/dmc/ball_in_cup.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/finger.h"
 #include "envpool/mujoco/dmc/hopper.h"
@@ -20,6 +21,9 @@
 #include "envpool/mujoco/dmc/reacher.h"
 #include "envpool/mujoco/dmc/walker.h"
 
+using DmcBallInCupEnvSpec = PyEnvSpec<mujoco_dmc::BallInCupEnvSpec>;
+using DmcBallInCupEnvPool = PyEnvPool<mujoco_dmc::BallInCupEnvPool>;
+
 using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
 using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
 
@@ -39,6 +43,7 @@ using DmcWalkerEnvSpec = PyEnvSpec<mujoco_dmc::WalkerEnvSpec>;
 using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
 
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
+  REGISTER(m, DmcBallInCupEnvSpec, DmcBallInCupEnvPool)
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index f1cf3944..204d744f 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -21,6 +21,7 @@
 
 # from suite.BENCHMARKING
 dmc_mujoco_envs = [
+  ("ball_in_cup", "catch"),
   ("cheetah", "run"),
   ("finger", "spin"),
   ("finger", "turn_easy"),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 29958668..dc4502ef 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -120,6 +120,7 @@ def test_make_mujoco_gym(self) -> None:
   def test_make_mujoco_dmc(self) -> None:
     self.check_step(
       [
+        "BallInCupCatch-v1",
         "CheetahRun-v1",
         "FingerSpin-v1",
         "FingerTurnEasy-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 0489be46..476bb840 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -22,6 +22,8 @@
 from dm_control import suite
 
 from envpool.mujoco import (
+  DmcBallInCupDMEnvPool,
+  DmcBallInCupEnvSpec,
   DmcCheetahDMEnvPool,
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
@@ -66,9 +68,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 == "finger" and task in ["turn_easy", "turn_hard"]:
-        obs = ts.observation
-        env.physics.named.model.site_pos["target", ["x", "z"]] = obs.target[0]
+      elif domain in ["finger", "ball_in_cup"]:
+        if domain == "finger" and task in ["turn_easy", "turn_hard"]:
+          target = ts.observation.target[0]
+          env.physics.named.model.site_pos["target", ["x", "z"]] = target
+        env.physics.after_reset()
 
   def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
     return np.random.uniform(
@@ -82,7 +86,7 @@ def run_align_check(
   ) -> None:
     logging.info(f"align check for {domain} {task}")
     obs_spec, action_spec = env0.observation_spec(), env0.action_spec()
-    for i in range(5):
+    for i in range(3):
       np.random.seed(i)
       env0.reset()
       a = self.sample_action(action_spec)
@@ -114,6 +118,11 @@ def run_align_check_entry(
       self.run_space_check(env0, env1)
       self.run_align_check(env0, env1, domain, task)
 
+  def test_ball_in_cup(self) -> None:
+    self.run_align_check_entry(
+      "ball_in_cup", ["catch"], DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool
+    )
+
   def test_cheetah(self) -> None:
     self.run_align_check_entry(
       "cheetah", ["run"], DmcCheetahEnvSpec, DmcCheetahDMEnvPool
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index 5dbdea83..945f6ec7 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -13,13 +13,15 @@
 # limitations under the License.
 """Unit tests for Mujoco dm_control deterministic check."""
 
-from typing import Any, List
+from typing import Any, List, Optional
 
 import dm_env
 import numpy as np
 from absl.testing import absltest
 
 from envpool.mujoco import (
+  DmcBallInCupDMEnvPool,
+  DmcBallInCupEnvSpec,
   DmcCheetahDMEnvPool,
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
@@ -43,8 +45,10 @@ def check(
     envpool_cls: Any,
     task: str,
     obs_keys: List[str],
+    blacklist: Optional[List[str]] = None,
     num_envs: int = 4,
   ) -> None:
+    np.random.seed(0)
     env0 = envpool_cls(
       spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=0, task_name=task))
     )
@@ -72,8 +76,22 @@ def check(
         o1 = getattr(obs1, k)
         o2 = getattr(obs2, k)
         np.testing.assert_allclose(o0, o1)
+        if blacklist and k in blacklist:
+          continue
         if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
-          self.assertFalse(np.allclose(o0, o2), (o0, o2))
+          self.assertFalse(np.allclose(o0, o2), (k, o0, o2))
+
+  def test_ball_in_cup(self) -> None:
+    obs_keys = ["position", "velocity"]
+    for task in ["catch"]:
+      self.check(
+        DmcBallInCupEnvSpec,
+        DmcBallInCupDMEnvPool,
+        task,
+        obs_keys,
+        # https://github.com/sail-sg/envpool/pull/124#issuecomment-1127860698
+        blacklist=["velocity"],
+      )
 
   def test_cheetah(self) -> None:
     obs_keys = ["position", "velocity"]