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 6218412c..1a2dd3b6 100644
--- a/docs/api/dm_control.rst
+++ b/docs/api/dm_control.rst
@@ -20,6 +20,19 @@ converted into ``DomainNameTaskName-v1`` in envpool, e.g.,
   envpool.make_dm("BallInCupCatch-v1", num_envs=1)
 
 
+AcrobotSwingup-v1, AcrobotSwingupSparse-v1
+------------------------------------------
+
+`dm_control suite acrobot source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/acrobot.py>`_
+
+- Observation spec: a namedtuple with two keys: ``orientations (4)``,
+  ``velocity (2)``;
+- Action spec: ``(1)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_stes``: 1000;
+
+
 BallInCupCatch-v1
 -----------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index a6d9d260..81711a60 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -63,6 +63,7 @@ pybind_extension(
 cc_library(
     name = "mujoco_dmc_env",
     hdrs = [
+        "dmc/acrobot.h",
         "dmc/ball_in_cup.h",
         "dmc/cheetah.h",
         "dmc/finger.h",
diff --git a/envpool/mujoco/__init__.py b/envpool/mujoco/__init__.py
index e417ca3f..dff158b5 100644
--- a/envpool/mujoco/__init__.py
+++ b/envpool/mujoco/__init__.py
@@ -17,6 +17,8 @@
 
 # dmc mujoco
 from .mujoco_dmc_envpool import (
+  _DmcAcrobotEnvPool,
+  _DmcAcrobotEnvSpec,
   _DmcBallInCupEnvPool,
   _DmcBallInCupEnvSpec,
   _DmcCheetahEnvPool,
@@ -62,6 +64,9 @@
 )
 
 # dmc
+DmcAcrobotEnvSpec, DmcAcrobotDMEnvPool, DmcAcrobotGymEnvPool = py_env(
+  _DmcAcrobotEnvSpec, _DmcAcrobotEnvPool
+)
 DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool, DmcBallInCupGymEnvPool = py_env(
   _DmcBallInCupEnvSpec, _DmcBallInCupEnvPool
 )
@@ -133,6 +138,9 @@
 )
 
 __all__ = [
+  "DmcAcrobotEnvSpec",
+  "DmcAcrobotDMEnvPool",
+  "DmcAcrobotGymEnvPool",
   "DmcBallInCupEnvSpec",
   "DmcBallInCupDMEnvPool",
   "DmcBallInCupGymEnvPool",
diff --git a/envpool/mujoco/dmc/acrobot.h b/envpool/mujoco/dmc/acrobot.h
new file mode 100644
index 00000000..fe361a6c
--- /dev/null
+++ b/envpool/mujoco/dmc/acrobot.h
@@ -0,0 +1,171 @@
+/*
+ * 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/acrobot.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_ACROBOT_H_
+#define ENVPOOL_MUJOCO_DMC_ACROBOT_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 GetAcrobotXML(const std::string& base_path,
+                          const std::string& task_name) {
+  return GetFileContent(base_path, "acrobot.xml");
+}
+
+class AcrobotEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
+                    "task_name"_.Bind(std::string("swingup")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    return MakeDict("obs:orientations"_.Bind(Spec<mjtNum>({4})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({2})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({2})),
+#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 AcrobotEnvSpec = EnvSpec<AcrobotEnvFns>;
+
+class AcrobotEnv : public Env<AcrobotEnvSpec>, public MujocoEnv {
+ protected:
+  int id_upper_arm_, id_lower_arm_, id_target_, id_tip_, id_shoulder_,
+      id_elbow_;
+  std::uniform_real_distribution<> dist_uniform_;
+  bool is_sparse_;
+
+ public:
+  AcrobotEnv(const Spec& spec, int env_id)
+      : Env<AcrobotEnvSpec>(spec, env_id),
+        MujocoEnv(
+            spec.config["base_path"_],
+            GetAcrobotXML(spec.config["base_path"_], spec.config["task_name"_]),
+            spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
+        id_upper_arm_(mj_name2id(model_, mjOBJ_XBODY, "upper_arm")),
+        id_lower_arm_(mj_name2id(model_, mjOBJ_XBODY, "lower_arm")),
+        id_target_(mj_name2id(model_, mjOBJ_SITE, "target")),
+        id_tip_(mj_name2id(model_, mjOBJ_SITE, "tip")),
+        id_shoulder_(mj_name2id(model_, mjOBJ_JOINT, "shoulder")),
+        id_elbow_(mj_name2id(model_, mjOBJ_JOINT, "elbow")),
+        dist_uniform_(-M_PI, M_PI),
+        is_sparse_(spec.config["task_name"_] == "swingup_sparse") {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name != "swingup" && task_name != "swingup_sparse") {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc acrobot.");
+    }
+  }
+
+  void TaskInitializeEpisode() override {
+    data_->qpos[id_shoulder_] = dist_uniform_(gen_);
+    data_->qpos[id_elbow_] = 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 {
+    mjtNum target_radius = model_->site_size[id_target_];
+    return static_cast<float>(RewardTolerance(ToTarget(), 0.0, target_radius,
+                                              is_sparse_ ? 0.0 : 1.0));
+  }
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    const auto& orientations = Orientations();
+    state["obs:orientations"_].Assign(orientations.begin(),
+                                      orientations.size());
+    state["obs:velocity"_].Assign(data_->qvel, model_->nv);
+    // info for check alignment
+#ifdef ENVPOOL_TEST
+    state["info:qpos0"_].Assign(qpos0_.get(), model_->nq);
+#endif
+  }
+
+  std::array<mjtNum, 2> Horizontal() {
+    // return self.named.data.xmat[['upper_arm', 'lower_arm'], 'xz']
+    return {data_->xmat[id_upper_arm_ * 9 + 2],
+            data_->xmat[id_lower_arm_ * 9 + 2]};
+  }
+  std::array<mjtNum, 2> Vertical() {
+    // return self.named.data.xmat[['upper_arm', 'lower_arm'], 'zz']
+    return {data_->xmat[id_upper_arm_ * 9 + 8],
+            data_->xmat[id_lower_arm_ * 9 + 8]};
+  }
+  mjtNum ToTarget() {
+    // tip_to_target = (self.named.data.site_xpos['target'] -
+    //                  self.named.data.site_xpos['tip'])
+    // return np.linalg.norm(tip_to_target)
+    std::array<mjtNum, 3> tip_to_target = {
+        data_->site_xpos[id_target_ * 3] - data_->site_xpos[id_tip_ * 3],
+        data_->site_xpos[id_target_ * 3 + 1] -
+            data_->site_xpos[id_tip_ * 3 + 1],
+        data_->site_xpos[id_target_ * 3 + 2] -
+            data_->site_xpos[id_tip_ * 3 + 2]};
+    return std::sqrt(tip_to_target[0] * tip_to_target[0] +
+                     tip_to_target[1] * tip_to_target[1] +
+                     tip_to_target[2] * tip_to_target[2]);
+  }
+  std::array<mjtNum, 4> Orientations() {
+    const auto& horizontal = Horizontal();
+    const auto& vertical = Vertical();
+    return {horizontal[0], horizontal[1], vertical[0], vertical[1]};
+  }
+};
+
+using AcrobotEnvPool = AsyncEnvPool<AcrobotEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_ACROBOT_H_
diff --git a/envpool/mujoco/dmc/ball_in_cup.h b/envpool/mujoco/dmc/ball_in_cup.h
index 6bc764f1..6c5292b3 100644
--- a/envpool/mujoco/dmc/ball_in_cup.h
+++ b/envpool/mujoco/dmc/ball_in_cup.h
@@ -77,7 +77,13 @@ class BallInCupEnv : public Env<BallInCupEnvSpec>, public MujocoEnv {
         id_ball_x_(mj_name2id(model_, mjOBJ_JOINT, "ball_x")),
         id_ball_z_(mj_name2id(model_, mjOBJ_JOINT, "ball_z")),
         dist_ball_x_(-0.2, 0.2),
-        dist_ball_z_(0.2, 0.5) {}
+        dist_ball_z_(0.2, 0.5) {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name != "catch") {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc ball_in_cup.");
+    }
+  }
 
   void TaskInitializeEpisode() override {
     while (true) {
diff --git a/envpool/mujoco/dmc/cheetah.h b/envpool/mujoco/dmc/cheetah.h
index af285536..dec4b61e 100644
--- a/envpool/mujoco/dmc/cheetah.h
+++ b/envpool/mujoco/dmc/cheetah.h
@@ -73,7 +73,13 @@ class CheetahEnv : public Env<CheetahEnvSpec>, public MujocoEnv {
             GetCheetahXML(spec.config["base_path"_], spec.config["task_name"_]),
             spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
         id_torso_subtreelinvel_(GetSensorId(model_, "torso_subtreelinvel")),
-        dist_uniform_(0, 1) {}
+        dist_uniform_(0, 1) {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name != "run") {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc cheetah.");
+    }
+  }
 
   void TaskInitializeEpisode() override {
     assert(model_->njnt == model_->nq);
diff --git a/envpool/mujoco/dmc/finger.h b/envpool/mujoco/dmc/finger.h
index 11890a05..df2f11e8 100644
--- a/envpool/mujoco/dmc/finger.h
+++ b/envpool/mujoco/dmc/finger.h
@@ -112,7 +112,8 @@ class FingerEnv : public Env<FingerEnvSpec>, public MujocoEnv {
     } else if (task_name == "turn_hard") {
       target_radius_ = kHardTargetSize;
     } else if (task_name != "spin") {
-      throw std::runtime_error("Unknown task_name for dmc finger.");
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc finger.");
     }
   }
 
diff --git a/envpool/mujoco/dmc/hopper.h b/envpool/mujoco/dmc/hopper.h
index e6ffa572..82d65347 100644
--- a/envpool/mujoco/dmc/hopper.h
+++ b/envpool/mujoco/dmc/hopper.h
@@ -78,13 +78,14 @@ class HopperEnv : public Env<HopperEnvSpec>, public MujocoEnv {
         id_torso_subtreelinvel_(GetSensorId(model_, "torso_subtreelinvel")),
         id_touch_toe_(GetSensorId(model_, "touch_toe")),
         id_touch_heel_(GetSensorId(model_, "touch_heel")) {
-    std::string task_name = spec.config["task_name"_];
+    const std::string& task_name = spec.config["task_name"_];
     if (task_name == "stand") {
       hopping_ = false;
     } else if (task_name == "hop") {
       hopping_ = true;
     } else {
-      throw std::runtime_error("Unknown task_name for dmc hopper.");
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc hopper.");
     }
   }
 
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index eb82ffa9..a326e58d 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/acrobot.h"
 #include "envpool/mujoco/dmc/ball_in_cup.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/finger.h"
@@ -22,6 +23,9 @@
 #include "envpool/mujoco/dmc/reacher.h"
 #include "envpool/mujoco/dmc/walker.h"
 
+using DmcAcrobotEnvSpec = PyEnvSpec<mujoco_dmc::AcrobotEnvSpec>;
+using DmcAcrobotEnvPool = PyEnvPool<mujoco_dmc::AcrobotEnvPool>;
+
 using DmcBallInCupEnvSpec = PyEnvSpec<mujoco_dmc::BallInCupEnvSpec>;
 using DmcBallInCupEnvPool = PyEnvPool<mujoco_dmc::BallInCupEnvPool>;
 
@@ -47,6 +51,7 @@ using DmcWalkerEnvSpec = PyEnvSpec<mujoco_dmc::WalkerEnvSpec>;
 using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
 
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
+  REGISTER(m, DmcAcrobotEnvSpec, DmcAcrobotEnvPool)
   REGISTER(m, DmcBallInCupEnvSpec, DmcBallInCupEnvPool)
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
diff --git a/envpool/mujoco/dmc/pendulum.h b/envpool/mujoco/dmc/pendulum.h
index 45d130cf..6751e742 100644
--- a/envpool/mujoco/dmc/pendulum.h
+++ b/envpool/mujoco/dmc/pendulum.h
@@ -76,7 +76,13 @@ class PendulumEnv : public Env<PendulumEnvSpec>, public MujocoEnv {
                   spec.config["max_episode_steps"_]),
         id_hinge_(mj_name2id(model_, mjOBJ_JOINT, "hinge")),
         id_pole_(mj_name2id(model_, mjOBJ_XBODY, "pole")),
-        dist_uniform_(-M_PI, M_PI) {}
+        dist_uniform_(-M_PI, M_PI) {
+    const std::string& task_name = spec.config["task_name"_];
+    if (task_name != "swingup") {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc pendulum.");
+    }
+  }
 
   void TaskInitializeEpisode() override {
     data_->qpos[0] = dist_uniform_(gen_);
diff --git a/envpool/mujoco/dmc/point_mass.h b/envpool/mujoco/dmc/point_mass.h
index a834b1ab..a47c467d 100644
--- a/envpool/mujoco/dmc/point_mass.h
+++ b/envpool/mujoco/dmc/point_mass.h
@@ -80,13 +80,14 @@ class PointMassEnv : public Env<PointMassEnvSpec>, public MujocoEnv {
         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"_];
+    const 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.");
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc point_mass.");
     }
 #ifdef ENVPOOL_TEST
     wrap_prm_.reset(new mjtNum[model_->nwrap]);
diff --git a/envpool/mujoco/dmc/reacher.h b/envpool/mujoco/dmc/reacher.h
index 95366344..7d215d2b 100644
--- a/envpool/mujoco/dmc/reacher.h
+++ b/envpool/mujoco/dmc/reacher.h
@@ -84,13 +84,14 @@ class ReacherEnv : public Env<ReacherEnvSpec>, public MujocoEnv {
         id_finger_(mj_name2id(model_, mjOBJ_GEOM, "finger")),
         dist_angle_(0, 2 * M_PI),
         dist_radius_(0.05, 0.2) {
-    std::string task_name = spec.config["task_name"_];
+    const 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.");
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc reacher.");
     }
   }
 
diff --git a/envpool/mujoco/dmc/walker.h b/envpool/mujoco/dmc/walker.h
index 9e3dd0da..50b479ff 100644
--- a/envpool/mujoco/dmc/walker.h
+++ b/envpool/mujoco/dmc/walker.h
@@ -79,7 +79,7 @@ class WalkerEnv : public Env<WalkerEnvSpec>, public MujocoEnv {
             spec.config["frame_skip"_], spec.config["max_episode_steps"_]),
         id_torso_(mj_name2id(model_, mjOBJ_XBODY, "torso")),
         id_torso_subtreelinvel_(GetSensorId(model_, "torso_subtreelinvel")) {
-    std::string task_name = spec.config["task_name"_];
+    const std::string& task_name = spec.config["task_name"_];
     if (task_name == "stand") {
       move_speed_ = 0;
     } else if (task_name == "walk") {
@@ -87,7 +87,8 @@ 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 walker.");
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc walker.");
     }
   }
 
diff --git a/envpool/mujoco/registration.py b/envpool/mujoco/registration.py
index 2aa0926e..812c2af3 100644
--- a/envpool/mujoco/registration.py
+++ b/envpool/mujoco/registration.py
@@ -21,6 +21,8 @@
 
 # from suite.BENCHMARKING
 dmc_mujoco_envs = [
+  ("acrobot", "swingup"),
+  ("acrobot", "swingup_sparse"),
   ("ball_in_cup", "catch"),
   ("cheetah", "run"),
   ("finger", "spin"),


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index c5c9553e..adcbb3ff 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -121,6 +121,8 @@ def test_make_mujoco_gym(self) -> None:
   def test_make_mujoco_dmc(self) -> None:
     self.check_step(
       [
+        "AcrobotSwingup-v1",
+        "AcrobotSwingupSparse-v1",
         "BallInCupCatch-v1",
         "CheetahRun-v1",
         "FingerSpin-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 036f8adf..3d856fc6 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 (
+  DmcAcrobotDMEnvPool,
+  DmcAcrobotEnvSpec,
   DmcBallInCupDMEnvPool,
   DmcBallInCupEnvSpec,
   DmcCheetahDMEnvPool,
@@ -127,6 +129,12 @@ def run_align_check_entry(
       self.run_space_check(env0, env1)
       self.run_align_check(env0, env1, domain, task)
 
+  def test_acrobot(self) -> None:
+    self.run_align_check_entry(
+      "acrobot", ["swingup", "swingup_sparse"], DmcAcrobotEnvSpec,
+      DmcAcrobotDMEnvPool
+    )
+
   def test_ball_in_cup(self) -> None:
     self.run_align_check_entry(
       "ball_in_cup", ["catch"], DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool
diff --git a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
index fccc3895..5fba8d79 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -20,6 +20,8 @@
 from absl.testing import absltest
 
 from envpool.mujoco import (
+  DmcAcrobotDMEnvPool,
+  DmcAcrobotEnvSpec,
   DmcBallInCupDMEnvPool,
   DmcBallInCupEnvSpec,
   DmcCheetahDMEnvPool,
@@ -83,6 +85,11 @@ def check(
         if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
           self.assertFalse(np.allclose(o0, o2), (k, o0, o2))
 
+  def test_acrobot(self) -> None:
+    obs_keys = ["orientations", "velocity"]
+    for task in ["swingup", "swingup_sparse"]:
+      self.check(DmcAcrobotEnvSpec, DmcAcrobotDMEnvPool, task, obs_keys)
+
   def test_ball_in_cup(self) -> None:
     obs_keys = ["position", "velocity"]
     for task in ["catch"]: