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 dc056dc2..e0b1cf4d 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -66,6 +66,19 @@ HopperStand-v1, HopperHop-v1
 - ``max_episode_steps``: 1000;
 
 
+PendulumSwingup-v1
+------------------
+
+`dm_control suite pendulum source code
+<https://github.com/deepmind/dm_control/blob/main/dm_control/suite/pendulum.py>`_
+
+- Observation spec: a namedtuple with three keys: ``orientations (2)``,
+  ``velocity (1)``;
+- 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 9a50d854..73058f92 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -67,6 +67,7 @@ cc_library(
         "dmc/finger.h",
         "dmc/hopper.h",
         "dmc/mujoco_env.h",
+        "dmc/pendulum.h",
         "dmc/reacher.h",
         "dmc/utils.h",
         "dmc/walker.h",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index 896f13dd..807c9cce 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -23,6 +23,8 @@
   _DmcFingerEnvSpec,
   _DmcHopperEnvPool,
   _DmcHopperEnvSpec,
+  _DmcPendulumEnvPool,
+  _DmcPendulumEnvSpec,
   _DmcReacherEnvPool,
   _DmcReacherEnvSpec,
   _DmcWalkerEnvPool,
@@ -59,14 +61,15 @@
 DmcCheetahEnvSpec, DmcCheetahDMEnvPool, DmcCheetahGymEnvPool = py_env(
   _DmcCheetahEnvSpec, _DmcCheetahEnvPool
 )
-
 DmcFingerEnvSpec, DmcFingerDMEnvPool, DmcFingerGymEnvPool = py_env(
   _DmcFingerEnvSpec, _DmcFingerEnvPool
 )
-
 DmcHopperEnvSpec, DmcHopperDMEnvPool, DmcHopperGymEnvPool = py_env(
   _DmcHopperEnvSpec, _DmcHopperEnvPool
 )
+DmcPendulumEnvSpec, DmcPendulumDMEnvPool, DmcPendulumGymEnvPool = py_env(
+  _DmcPendulumEnvSpec, _DmcPendulumEnvPool
+)
 DmcReacherEnvSpec, DmcReacherDMEnvPool, DmcReacherGymEnvPool = py_env(
   _DmcReacherEnvSpec, _DmcReacherEnvPool
 )
@@ -126,6 +129,9 @@
   "DmcHopperEnvSpec",
   "DmcHopperDMEnvPool",
   "DmcHopperGymEnvPool",
+  "DmcPendulumEnvSpec",
+  "DmcPendulumDMEnvPool",
+  "DmcPendulumGymEnvPool",
   "DmcReacherEnvSpec",
   "DmcReacherDMEnvPool",
   "DmcReacherGymEnvPool",
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index c86fca67..4da1ff75 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -16,6 +16,7 @@
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/finger.h"
 #include "envpool/mujoco/dmc/hopper.h"
+#include "envpool/mujoco/dmc/pendulum.h"
 #include "envpool/mujoco/dmc/reacher.h"
 #include "envpool/mujoco/dmc/walker.h"
 
@@ -28,6 +29,9 @@ using DmcFingerEnvPool = PyEnvPool<mujoco_dmc::FingerEnvPool>;
 using DmcHopperEnvSpec = PyEnvSpec<mujoco_dmc::HopperEnvSpec>;
 using DmcHopperEnvPool = PyEnvPool<mujoco_dmc::HopperEnvPool>;
 
+using DmcPendulumEnvSpec = PyEnvSpec<mujoco_dmc::PendulumEnvSpec>;
+using DmcPendulumEnvPool = PyEnvPool<mujoco_dmc::PendulumEnvPool>;
+
 using DmcReacherEnvSpec = PyEnvSpec<mujoco_dmc::ReacherEnvSpec>;
 using DmcReacherEnvPool = PyEnvPool<mujoco_dmc::ReacherEnvPool>;
 
@@ -38,6 +42,7 @@ PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
   REGISTER(m, DmcHopperEnvSpec, DmcHopperEnvPool)
+  REGISTER(m, DmcPendulumEnvSpec, DmcPendulumEnvPool)
   REGISTER(m, DmcReacherEnvSpec, DmcReacherEnvPool)
   REGISTER(m, DmcWalkerEnvSpec, DmcWalkerEnvPool)
 }
diff --git a/envpool/mujoco/dmc/pendulum.h b/envpool/mujoco/dmc/pendulum.h
new file mode 100644
index 00000000..d0df4c78
--- /dev/null
+++ b/envpool/mujoco/dmc/pendulum.h
@@ -0,0 +1,129 @@
+/*
+ * 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/pendulum.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_PENDULUM_H_
+#define ENVPOOL_MUJOCO_DMC_PENDULUM_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 GetPendulumXML(const std::string& base_path,
+                           const std::string& task_name) {
+  return GetFileContent(base_path, "pendulum.xml");
+}
+
+class PendulumEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
+                    "task_name"_.Bind(std::string("stand")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:orientation"_.Bind(Spec<mjtNum>({2})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({1})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({1})),
+#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, 1}, {-1.0, 1.0})));
+  }
+};
+
+using PendulumEnvSpec = EnvSpec<PendulumEnvFns>;
+
+class PendulumEnv : public Env<PendulumEnvSpec>, public MujocoEnv {
+ protected:
+  const mjtNum kCosineBound = 0.9902680687415704;
+  std::uniform_real_distribution<> dist_uniform_;
+
+ public:
+  PendulumEnv(const Spec& spec, int env_id)
+      : Env<PendulumEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetPendulumXML(spec.config["base_path"_],
+                                 spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+        dist_uniform_(-M_PI, M_PI) {}
+
+  void TaskInitializeEpisode() override {
+    data_->qpos[0] = dist_uniform_(gen_);
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+#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 {
+    return static_cast<float>(RewardTolerance(PoleVertical(), kCosineBound, 1));
+  }
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    const auto& pole_orient = PoleOrientation();
+    state["obs:orientation"_].Assign(pole_orient.begin(), pole_orient.size());
+    state["obs:velocity"_] = AngularVelocity();
+    // info for check alignment
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  mjtNum PoleVertical() { return data_->xmat[1 * 9 + 8]; }
+  mjtNum AngularVelocity() { return data_->qvel[0]; }
+  std::array<mjtNum, 2> PoleOrientation() {
+    return {data_->xmat[1 * 9 + 8], data_->xmat[1 * 9 + 2]};
+  }
+};
+
+using PendulumEnvPool = AsyncEnvPool<PendulumEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_PENDULUM_H_
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 1351f6e3..f1cf3944 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -27,6 +27,7 @@
   ("finger", "turn_hard"),
   ("hopper", "hop"),
   ("hopper", "stand"),
+  ("pendulum", "swingup"),
   ("reacher", "easy"),
   ("reacher", "hard"),
   ("walker", "run"),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index 153521ff..29958668 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -126,6 +126,7 @@ def test_make_mujoco_dmc(self) -> None:
         "FingerTurnHard-v1",
         "HopperHop-v1",
         "HopperStand-v1",
+        "PendulumSwingup-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 dc97a940..0489be46 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -28,6 +28,8 @@
   DmcFingerEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcPendulumDMEnvPool,
+  DmcPendulumEnvSpec,
   DmcReacherDMEnvPool,
   DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
@@ -128,6 +130,11 @@ def test_hopper(self) -> None:
       "hopper", ["hop", "stand"], DmcHopperEnvSpec, DmcHopperDMEnvPool
     )
 
+  def test_pendulum(self) -> None:
+    self.run_align_check_entry(
+      "pendulum", ["swingup"], DmcPendulumEnvSpec, DmcPendulumDMEnvPool
+    )
+
   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 f9dece5b..5dbdea83 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -26,6 +26,8 @@
   DmcFingerEnvSpec,
   DmcHopperDMEnvPool,
   DmcHopperEnvSpec,
+  DmcPendulumDMEnvPool,
+  DmcPendulumEnvSpec,
   DmcReacherDMEnvPool,
   DmcReacherEnvSpec,
   DmcWalkerDMEnvPool,
@@ -91,6 +93,11 @@ def test_hopper(self) -> None:
     for task in ["stand", "hop"]:
       self.check(DmcHopperEnvSpec, DmcHopperDMEnvPool, task, obs_keys)
 
+  def test_pendulum(self) -> None:
+    obs_keys = ["orientation", "velocity"]
+    for task in ["swingup"]:
+      self.check(DmcPendulumEnvSpec, DmcPendulumDMEnvPool, task, obs_keys)
+
   def test_reacher(self) -> None:
     obs_keys = ["position", "to_target", "velocity"]
     for task in ["easy", "hard"]: