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/env/dm_control.rst b/docs/env/dm_control.rst
index 95ea5428..86be6273 100644
--- a/docs/env/dm_control.rst
+++ b/docs/env/dm_control.rst
@@ -46,6 +46,20 @@ BallInCupCatch-v1
 - ``max_episode_steps``: 1000;
 
 
+CartpoleBalance-v1, CartpoleBalanceSparse-v1, CarpoletSwingup-v1, CartpoleSwingupSparse-v1, CartpoleTwoPoles-v1, CartPoleThreePoles-v1
+--------------------------------------------------------------------------------------------------------------------------------------
+
+`dm_control suite cartpole source code
+<https://github.com/deepmind/dm_control/blob/1.0.2/dm_control/suite/cartpole.py>`_
+
+- Observation spec: a namedtuple with two keys: ``position (5 for two_poles, 7
+  for three_poles, 3 for others)``, ``velocity (3 for two_poles, 4 for
+  three_poles, 2 for others)``;
+- Action spec: ``(1)``, with range ``[-1, 1]``;
+- ``frame_skip``: 1;
+- ``max_episode_stes``: 1000;
+
+
 CheetahRun-v1
 -------------
 
diff --git a/envpool/mujoco/BUILD b/envpool/mujoco/BUILD
index f17aeaac..48523957 100644
--- a/envpool/mujoco/BUILD
+++ b/envpool/mujoco/BUILD
@@ -69,6 +69,7 @@ cc_library(
     hdrs = [
         "dmc/acrobot.h",
         "dmc/ball_in_cup.h",
+        "dmc/cartpole.h",
         "dmc/cheetah.h",
         "dmc/finger.h",
         "dmc/fish.h",
diff --git a/envpool/mujoco/dmc/__init__.py b/envpool/mujoco/dmc/__init__.py
index 9d604aef..2971026d 100644
--- a/envpool/mujoco/dmc/__init__.py
+++ b/envpool/mujoco/dmc/__init__.py
@@ -18,6 +18,8 @@
   _DmcAcrobotEnvSpec,
   _DmcBallInCupEnvPool,
   _DmcBallInCupEnvSpec,
+  _DmcCartpoleEnvPool,
+  _DmcCartpoleEnvSpec,
   _DmcCheetahEnvPool,
   _DmcCheetahEnvSpec,
   _DmcFingerEnvPool,
@@ -47,6 +49,9 @@
 DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool, DmcBallInCupGymEnvPool = py_env(
   _DmcBallInCupEnvSpec, _DmcBallInCupEnvPool
 )
+DmcCartpoleEnvSpec, DmcCartpoleDMEnvPool, DmcCartpoleGymEnvPool = py_env(
+  _DmcCartpoleEnvSpec, _DmcCartpoleEnvPool
+)
 DmcCheetahEnvSpec, DmcCheetahDMEnvPool, DmcCheetahGymEnvPool = py_env(
   _DmcCheetahEnvSpec, _DmcCheetahEnvPool
 )
@@ -87,6 +92,9 @@
   "DmcBallInCupEnvSpec",
   "DmcBallInCupDMEnvPool",
   "DmcBallInCupGymEnvPool",
+  "DmcCartpoleEnvSpec",
+  "DmcCartpoleDMEnvPool",
+  "DmcCartpoleGymEnvPool",
   "DmcCheetahEnvSpec",
   "DmcCheetahDMEnvPool",
   "DmcCheetahGymEnvPool",
diff --git a/envpool/mujoco/dmc/cartpole.h b/envpool/mujoco/dmc/cartpole.h
new file mode 100644
index 00000000..18acb29a
--- /dev/null
+++ b/envpool/mujoco/dmc/cartpole.h
@@ -0,0 +1,239 @@
+/*
+ * 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/cartpole.py
+
+#ifndef ENVPOOL_MUJOCO_DMC_CARTPOLE_H_
+#define ENVPOOL_MUJOCO_DMC_CARTPOLE_H_
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <random>
+#include <regex>
+#include <set>
+#include <string>
+#include <vector>
+
+#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 GetCartpoleXML(const std::string& base_path,
+                           const std::string& task_name) {
+  auto content = GetFileContent(base_path, "cartpole.xml");
+  if (task_name == "two_poles") {
+    return XMLAddPoles(content, 2);
+  }
+  if (task_name == "three_poles") {
+    return XMLAddPoles(content, 3);
+  }
+  return content;
+}
+
+class CartpoleEnvFns {
+ public:
+  static decltype(auto) DefaultConfig() {
+    return MakeDict("max_episode_steps"_.Bind(1000), "frame_skip"_.Bind(1),
+                    "task_name"_.Bind(std::string("balance")));
+  }
+  template <typename Config>
+  static decltype(auto) StateSpec(const Config& conf) {
+    const std::string task_name = conf["task_name"_];
+    int n_poles;
+    if (task_name == "two_poles") {
+      n_poles = 2;
+    } else if (task_name == "three_poles") {
+      n_poles = 3;
+    } else if (task_name == "swingup" || task_name == "swingup_sparse" ||
+               task_name == "balance" || task_name == "balance_sparse") {
+      n_poles = 1;
+    } else {
+      throw std::runtime_error("Unknown task_name " + task_name +
+                               " for dmc cartpole.");
+    }
+    return MakeDict("obs:position"_.Bind(Spec<mjtNum>({1 + 2 * n_poles})),
+                    "obs:velocity"_.Bind(Spec<mjtNum>({1 + n_poles})),
+#ifdef ENVPOOL_TEST
+                    "info:qpos0"_.Bind(Spec<mjtNum>({1 + n_poles})),
+                    "info:qvel0"_.Bind(Spec<mjtNum>({1 + n_poles})),
+#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 CartpoleEnvSpec = EnvSpec<CartpoleEnvFns>;
+
+class CartpoleEnv : public Env<CartpoleEnvSpec>, public MujocoEnv {
+ protected:
+  int id_slider_, id_hinge1_;
+  bool is_sparse_, is_swingup_;
+
+#ifdef ENVPOOL_TEST
+  std::unique_ptr<mjtNum> qvel0_;
+#endif
+
+ public:
+  CartpoleEnv(const Spec& spec, int env_id)
+      : Env<CartpoleEnvSpec>(spec, env_id),
+        MujocoEnv(spec.config["base_path"_],
+                  GetCartpoleXML(spec.config["base_path"_],
+                                 spec.config["task_name"_]),
+                  spec.config["frame_skip"_],
+                  spec.config["max_episode_steps"_]),
+        id_slider_(GetQposId(model_, "slider")),
+        id_hinge1_(GetQposId(model_, "hinge_1")),
+        is_sparse_(spec.config["task_name"_] == "balance_sparse" ||
+                   spec.config["task_name"_] == "swingup_sparse"),
+        is_swingup_(spec.config["task_name"_] == "swingup" ||
+                    spec.config["task_name"_] == "swingup_sparse" ||
+                    spec.config["task_name"_] == "two_poles" ||
+                    spec.config["task_name"_] == "three_poles") {
+#ifdef ENVPOOL_TEST
+    qvel0_.reset(new mjtNum[model_->nv]);
+#endif
+  }
+
+  void TaskInitializeEpisode() override {
+    if (is_swingup_) {
+      data_->qpos[id_slider_] = RandNormal(0, 0.01)(gen_);
+      data_->qpos[id_hinge1_] = RandNormal(M_PI, 0.01)(gen_);
+      for (int i = 2; i < model_->nq; ++i) {
+        data_->qpos[i] = RandNormal(0, 0.01)(gen_);
+      }
+    } else {
+      data_->qpos[id_slider_] = RandUniform(-0.1, 0.1)(gen_);
+      for (int i = 1; i < model_->nq; ++i) {
+        data_->qpos[i] = RandUniform(-0.034, 0.034)(gen_);
+      }
+    }
+    for (int i = 0; i < model_->nv; ++i) {
+      data_->qvel[i] = RandNormal(0, 0.01)(gen_);
+    }
+#ifdef ENVPOOL_TEST
+    std::memcpy(qpos0_.get(), data_->qpos, sizeof(mjtNum) * model_->nq);
+    std::memcpy(qvel0_.get(), data_->qvel, sizeof(mjtNum) * model_->nv);
+#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 {
+    const auto& pole_angle_cosine = PoleAngleCosine();
+    if (is_sparse_) {
+      auto cart_in_bounds = RewardTolerance(CartPosition(), -0.25, 0.25);
+      mjtNum angle_in_bounds = 1.0;
+      for (auto x : pole_angle_cosine) {
+        angle_in_bounds *= RewardTolerance(x, 0.995, 1);
+      }
+      return static_cast<float>(cart_in_bounds * angle_in_bounds);
+    }
+    mjtNum upright = 0.0;
+    for (auto x : pole_angle_cosine) {
+      upright += (x + 1) / 2;
+    }
+    upright /= pole_angle_cosine.size();
+    auto centered = (1 + RewardTolerance(CartPosition(), 0.0, 0.0, 2)) / 2;
+    auto small_control = RewardTolerance(data_->ctrl[0], 0.0, 0.0, 1.0, 0.0,
+                                         SigmoidType::kQuadratic);
+    small_control = (small_control + 4) / 5;
+    auto angular_vel = AngularVel();
+    for (auto& x : angular_vel) {
+      x = RewardTolerance(x, 0.0, 0.0, 5.0);
+    }
+    mjtNum small_velocity = angular_vel[0];
+    for (auto x : angular_vel) {
+      small_velocity = std::min(small_velocity, x);
+    }
+    small_velocity = (small_velocity + 1) / 2;
+    return static_cast<float>(upright * small_control * small_velocity *
+                              centered);
+  }
+  bool TaskShouldTerminateEpisode() override { return false; }
+
+ private:
+  void WriteState() {
+    State state = Allocate();
+    state["reward"_] = reward_;
+    state["discount"_] = discount_;
+    // obs
+    const auto& position = BoundedPosition();
+    state["obs:position"_].Assign(position.data(), position.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:qvel0"_].Assign(qvel0_.get(), model_->nv);
+#endif
+  }
+
+  mjtNum CartPosition() {
+    // return self.named.data.qpos['slider'][0]
+    return data_->qpos[id_slider_];
+  }
+  std::vector<mjtNum> AngularVel() {
+    // return self.data.qvel[1:]
+    std::vector<mjtNum> result;
+    for (int i = 1; i < model_->nv; ++i) {
+      result.emplace_back(data_->qvel[i]);
+    }
+    return result;
+  }
+
+  std::vector<mjtNum> PoleAngleCosine() {
+    // return self.named.data.xmat[2:, 'zz']
+    std::vector<mjtNum> result;
+    for (int i = 2; i < model_->nbody; ++i) {
+      result.emplace_back(data_->xmat[i * 9 + 8]);
+    }
+    return result;
+  }
+
+  std::vector<mjtNum> BoundedPosition() {
+    // return np.hstack((self.cart_position(),
+    //                   self.named.data.xmat[2:, ['zz', 'xz']].ravel()))
+    std::vector<mjtNum> result = {CartPosition()};
+    for (int i = 2; i < model_->nbody; ++i) {
+      result.emplace_back(data_->xmat[i * 9 + 8]);
+      result.emplace_back(data_->xmat[i * 9 + 2]);
+    }
+    return result;
+  }
+};
+
+using CartpoleEnvPool = AsyncEnvPool<CartpoleEnv>;
+
+}  // namespace mujoco_dmc
+
+#endif  // ENVPOOL_MUJOCO_DMC_CARTPOLE_H_
diff --git a/envpool/mujoco/dmc/mujoco_envpool.cc b/envpool/mujoco/dmc/mujoco_envpool.cc
index 41b2b5e6..ab3923ac 100644
--- a/envpool/mujoco/dmc/mujoco_envpool.cc
+++ b/envpool/mujoco/dmc/mujoco_envpool.cc
@@ -15,6 +15,7 @@
 #include "envpool/core/py_envpool.h"
 #include "envpool/mujoco/dmc/acrobot.h"
 #include "envpool/mujoco/dmc/ball_in_cup.h"
+#include "envpool/mujoco/dmc/cartpole.h"
 #include "envpool/mujoco/dmc/cheetah.h"
 #include "envpool/mujoco/dmc/finger.h"
 #include "envpool/mujoco/dmc/fish.h"
@@ -32,6 +33,9 @@ using DmcAcrobotEnvPool = PyEnvPool<mujoco_dmc::AcrobotEnvPool>;
 using DmcBallInCupEnvSpec = PyEnvSpec<mujoco_dmc::BallInCupEnvSpec>;
 using DmcBallInCupEnvPool = PyEnvPool<mujoco_dmc::BallInCupEnvPool>;
 
+using DmcCartpoleEnvSpec = PyEnvSpec<mujoco_dmc::CartpoleEnvSpec>;
+using DmcCartpoleEnvPool = PyEnvPool<mujoco_dmc::CartpoleEnvPool>;
+
 using DmcCheetahEnvSpec = PyEnvSpec<mujoco_dmc::CheetahEnvSpec>;
 using DmcCheetahEnvPool = PyEnvPool<mujoco_dmc::CheetahEnvPool>;
 
@@ -65,6 +69,7 @@ using DmcWalkerEnvPool = PyEnvPool<mujoco_dmc::WalkerEnvPool>;
 PYBIND11_MODULE(mujoco_dmc_envpool, m) {
   REGISTER(m, DmcAcrobotEnvSpec, DmcAcrobotEnvPool)
   REGISTER(m, DmcBallInCupEnvSpec, DmcBallInCupEnvPool)
+  REGISTER(m, DmcCartpoleEnvSpec, DmcCartpoleEnvPool)
   REGISTER(m, DmcCheetahEnvSpec, DmcCheetahEnvPool)
   REGISTER(m, DmcFingerEnvSpec, DmcFingerEnvPool)
   REGISTER(m, DmcFishEnvSpec, DmcFishEnvPool)
diff --git a/envpool/mujoco/dmc/registration.py b/envpool/mujoco/dmc/registration.py
index 5b225999..93ac6ecb 100644
--- a/envpool/mujoco/dmc/registration.py
+++ b/envpool/mujoco/dmc/registration.py
@@ -26,6 +26,12 @@
   ("acrobot", "swingup"),
   ("acrobot", "swingup_sparse"),
   ("ball_in_cup", "catch"),
+  ("cartpole", "balance"),
+  ("cartpole", "balance_sparse"),
+  ("cartpole", "swingup"),
+  ("cartpole", "swingup_sparse"),
+  ("cartpole", "three_poles"),
+  ("cartpole", "two_poles"),
   ("cheetah", "run"),
   ("finger", "spin"),
   ("finger", "turn_easy"),
@@ -34,10 +40,10 @@
   ("fish", "upright"),
   ("hopper", "hop"),
   ("hopper", "stand"),
-  ("humanoid", "stand"),
-  ("humanoid", "walk"),
   ("humanoid", "run"),
   ("humanoid", "run_pure_state"),
+  ("humanoid", "stand"),
+  ("humanoid", "walk"),
   ("manipulator", "bring_ball"),
   ("manipulator", "bring_peg"),
   ("manipulator", "insert_ball"),
diff --git a/envpool/mujoco/dmc/utils.cc b/envpool/mujoco/dmc/utils.cc
index ce569b12..383fb23c 100644
--- a/envpool/mujoco/dmc/utils.cc
+++ b/envpool/mujoco/dmc/utils.cc
@@ -55,6 +55,42 @@ std::string XMLRemoveByBodyName(const std::string& content,
   return writer.result;
 }
 
+std::string XMLAddPoles(const std::string& content, int n_poles) {
+  pugi::xml_document doc;
+  doc.load_string(content.c_str());
+
+  pugi::xml_node body = doc.select_node("//worldbody/body/body").node();
+  for (int i = 2; i <= n_poles; ++i) {
+    pugi::xml_node new_pole = body.append_child("body");
+    new_pole.append_attribute("childclass") = "pole";
+    new_pole.append_attribute("name") = ("pole_" + std::to_string(i)).c_str();
+    new_pole.append_attribute("pos") = "0 0 1";
+    pugi::xml_node joint = new_pole.append_child("joint");
+    joint.append_attribute("name") = ("hinge_" + std::to_string(i)).c_str();
+    pugi::xml_node geom = new_pole.append_child("geom");
+    geom.append_attribute("name") = ("pole_" + std::to_string(i)).c_str();
+    body = new_pole;
+  }
+
+  pugi::xml_node floor = doc.select_node("//worldbody/geom").node();
+  floor.attribute("pos").set_value(
+      ("0 0 " + std::to_string(1 - n_poles - 0.05)).c_str());
+  pugi::xpath_node_set cameras = doc.select_nodes("//worldbody/camera");
+  for (const pugi::xpath_node& c : cameras) {
+    std::string name = c.node().attribute("name").value();
+    if (name == "fixed") {
+      c.node().attribute("pos").set_value(
+          ("0 " + std::to_string(-1 - 2 * n_poles) + " 1").c_str());
+    } else if (name == "lookatcart") {
+      c.node().attribute("pos").set_value(
+          ("0 " + std::to_string(-2 * n_poles) + " 2").c_str());
+    }
+  }
+  XMLStringWriter writer;
+  doc.print(writer);
+  return writer.result;
+}
+
 int GetQposId(mjModel* model, const std::string& name) {
   return model->jnt_qposadr[mj_name2id(model, mjOBJ_JOINT, name.c_str())];
 }
diff --git a/envpool/mujoco/dmc/utils.h b/envpool/mujoco/dmc/utils.h
index 6ee64737..d2329bee 100644
--- a/envpool/mujoco/dmc/utils.h
+++ b/envpool/mujoco/dmc/utils.h
@@ -35,6 +35,7 @@ std::string GetFileContent(const std::string& base_path,
                            const std::string& asset_name);
 std::string XMLRemoveByBodyName(const std::string& content,
                                 const std::vector<std::string>& body_names);
+std::string XMLAddPoles(const std::string& content, int n_poles);
 
 // the following id is not 1 on 1 mapping
 int GetQposId(mjModel* model, const std::string& name);


Test patch:
diff --git a/envpool/make_test.py b/envpool/make_test.py
index f0ab68f7..5d9bb83c 100644
--- a/envpool/make_test.py
+++ b/envpool/make_test.py
@@ -130,6 +130,12 @@ def test_make_mujoco_dmc(self) -> None:
         "AcrobotSwingup-v1",
         "AcrobotSwingupSparse-v1",
         "BallInCupCatch-v1",
+        "CartpoleBalance-v1",
+        "CartpoleBalanceSparse-v1",
+        "CartpoleSwingup-v1",
+        "CartpoleSwingupSparse-v1",
+        "CartpoleThreePoles-v1",
+        "CartpoleTwoPoles-v1",
         "CheetahRun-v1",
         "FingerSpin-v1",
         "FingerTurnEasy-v1",
@@ -138,10 +144,10 @@ def test_make_mujoco_dmc(self) -> None:
         "FishUpright-v1",
         "HopperHop-v1",
         "HopperStand-v1",
-        "HumanoidStand-v1",
-        "HumanoidWalk-v1",
         "HumanoidRun-v1",
         "HumanoidRunPureState-v1",
+        "HumanoidStand-v1",
+        "HumanoidWalk-v1",
         "ManipulatorBringBall-v1",
         "ManipulatorBringPeg-v1",
         "ManipulatorInsertBall-v1",
diff --git a/envpool/mujoco/dmc/mujoco_dmc_align_test.py b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
index 486e8c17..a47726ac 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_align_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_align_test.py
@@ -26,6 +26,8 @@
   DmcAcrobotEnvSpec,
   DmcBallInCupDMEnvPool,
   DmcBallInCupEnvSpec,
+  DmcCartpoleDMEnvPool,
+  DmcCartpoleEnvSpec,
   DmcCheetahDMEnvPool,
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
@@ -74,6 +76,8 @@ def reset_state(
         for _ in range(200):
           env.physics.step()
         env.physics.data.time = 0
+      elif domain == "cartpole":
+        env.physics.data.qvel = ts.observation.qvel0[0]
       elif domain == "reacher":
         target = ts.observation.target[0]
         env.physics.named.model.geom_pos["target", "x"] = target[0]
@@ -193,6 +197,14 @@ def test_ball_in_cup(self) -> None:
       "ball_in_cup", ["catch"], DmcBallInCupEnvSpec, DmcBallInCupDMEnvPool
     )
 
+  def test_cartpole(self) -> None:
+    self.run_align_check_entry(
+      "cartpole", [
+        "balance", "balance_sparse", "swingup", "swingup_sparse", "two_poles",
+        "three_poles"
+      ], DmcCartpoleEnvSpec, DmcCartpoleDMEnvPool
+    )
+
   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 383bc3f6..02a84fb8 100644
--- a/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
+++ b/envpool/mujoco/dmc/mujoco_dmc_deterministic_test.py
@@ -24,6 +24,8 @@
   DmcAcrobotEnvSpec,
   DmcBallInCupDMEnvPool,
   DmcBallInCupEnvSpec,
+  DmcCartpoleDMEnvPool,
+  DmcCartpoleEnvSpec,
   DmcCheetahDMEnvPool,
   DmcCheetahEnvSpec,
   DmcFingerDMEnvPool,
@@ -108,6 +110,14 @@ def test_ball_in_cup(self) -> None:
         blacklist=["velocity"],
       )
 
+  def test_cartpole(self) -> None:
+    obs_keys = ["position", "velocity"]
+    for task in [
+      "balance", "balance_sparse", "swingup", "swingup_sparse", "two_poles",
+      "three_poles"
+    ]:
+      self.check(DmcCartpoleEnvSpec, DmcCartpoleDMEnvPool, task, obs_keys)
+
   def test_cheetah(self) -> None:
     obs_keys = ["position", "velocity"]
     for task in ["run"]: