Fix locomotion packaging and platform validation

Author: Trinkle23897

docs/env/dm_control_locomotion.rst

@@ -191,7 +191,7 @@ unchanged. The four CMU egocentric cameras allow at most five color levels per
 channel and a total absolute error of 20 across the entire 64-by-64 frame;
 the rodent maze camera allows one level in one channel. Public renders allow
 one level in at most three channels for go-to-target and escape-bowl, or one
-channel for tracking. Other images and all native dynamics/reward replays
+channel for tracking and heterogeneous forage. Other images and all native dynamics/reward replays
 remain bitwise. These limits are checked per frame, not averaged over a rollout.
 
 Oracle reward checks retain only small derived-math residuals: tracking's

envpool/mujoco/locomotion/cmu.cc

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+#include <array>
 #include <memory>
 #include <set>
 #include <sstream>
@@ -31,14 +32,20 @@ std::vector<double> Read(pugi::xml_node node, const char* attribute) {
   std::istringstream stream(node.attribute(attribute).value());
   std::vector<double> values;
   double value;
-  while (stream >> value) values.push_back(value);
+  while (stream >> value) {
+    values.push_back(value);
+  }
   return values;
 }
 
 void Scale(pugi::xml_node node, const char* attribute, double factor) {
   auto values = Read(node, attribute);
-  if (values.empty()) return;
-  for (double& value : values) value *= factor;
+  if (values.empty()) {
+    return;
+  }
+  for (double& value : values) {
+    value *= factor;
+  }
   Set(node, attribute, Numbers(values.data(), values.size()));
 }
 
@@ -65,22 +72,26 @@ void Rescale(pugi::xml_node model) {
   mjVFS vfs;
   mj_defaultVFS(&vfs);
   mj_addBufferVFS(&vfs, "cmu.xml", contents.data(), contents.size());
-  char error[1024]{};
+  std::array<char, 1024> error{};
   std::unique_ptr<mjModel, decltype(&mj_deleteModel)> physics(
-      mj_loadXML("cmu.xml", &vfs, error, sizeof(error)), mj_deleteModel);
+      mj_loadXML("cmu.xml", &vfs, error.data(), sizeof(error)), mj_deleteModel);
   mj_deleteVFS(&vfs);
-  if (!physics) throw std::runtime_error(error);
+  if (!physics) {
+    throw std::runtime_error(error.data());
+  }
   const int root = mj_name2id(physics.get(), mjOBJ_BODY, "root");
   const double mass_factor = 70 / physics->body_subtreemass[root];
-  for (auto entry : Find(model, "body", "root").select_nodes(".//inertial"))
+  for (auto entry : Find(model, "body", "root").select_nodes(".//inertial")) {
     Scale(entry.node(), "mass", mass_factor);
+  }
   for (auto entry : Find(model, "body", "root").select_nodes(".//geom")) {
     auto geom = entry.node();
-    if (geom.attribute("mass"))
+    if (geom.attribute("mass") != nullptr) {
       Scale(geom, "mass", mass_factor);
-    else
+    } else {
       Set(geom, "density",
           geom.attribute("density").as_double(1000) * mass_factor);
+    }
   }
 }
 
@@ -124,8 +135,9 @@ void Scene::CmuVisuals(pugi::xml_node model, Walker walker, int player,
       auto visual = geom.parent().append_child("geom");
       Set(visual, "name", name + "_visual");
       for (const char* attr : {"type", "quat", "pos", "size"}) {
-        if (geom.attribute(attr))
+        if (geom.attribute(attr) != nullptr) {
           Set(visual, attr, geom.attribute(attr).value());
+        }
       }
       Scale(visual, "size", name == hand ? 1.3 : 1.5);
       Scale(visual, "pos", 1.5);
@@ -153,33 +165,38 @@ void Scene::CmuVisuals(pugi::xml_node model, Walker walker, int player,
   for (const char* name :
        {"lhipjoint", "rhipjoint", "lfemur", "lowerback", "upperback",
         "rclavicle", "lclavicle", "thorax", "lhumerus", "root_geom",
-        "lowerneck", "rhumerus", "rfemur"})
+        "lowerneck", "rhumerus", "rfemur"}) {
     Set(Find(model, "geom", name), "rgba", {0, 0, 0, 0});
-  const double neck = .066 - .0452401, arm = .20 - .138421,
-               leg = .384 - .202473;
-  const char* bodies[]{"lowerneck", "lhumerus", "rhumerus", "lfemur", "rfemur"};
-  const char* names[]{"halfneck", "lelbow", "relbow", "lknee", "rknee"};
-  const double sizes[5][2]{{.05, .02279225 - neck},
-                           {.035, .1245789 - arm},
-                           {.035, .1245789 - arm},
-                           {.055, .1822257 - leg},
-                           {.055, .1822257 - leg}};
-  const double positions[5][3]{{-.00165071, .0452401 + neck, .00534359},
-                               {0, -.138421 - arm, 0},
-                               {0, -.138421 - arm, 0},
-                               {-5.0684e-8, -.202473 - leg, 0},
-                               {-5.0684e-8, -.202473 - leg, 0}};
-  const double quats[5][4]{{.66437, .746906, .027253, 0},
-                           {.612372, -.612372, .353553, .353553},
-                           {.612372, -.612372, -.353553, -.353553},
-                           {.696364, -.696364, -.122788, -.122788},
-                           {.696364, -.696364, .122788, .122788}};
+  }
+  const double neck = .066 - .0452401;
+  const double arm = .20 - .138421;
+  const double leg = .384 - .202473;
+  const std::array bodies{"lowerneck", "lhumerus", "rhumerus", "lfemur",
+                          "rfemur"};
+  const std::array names{"halfneck", "lelbow", "relbow", "lknee", "rknee"};
+  const std::array<std::array<double, 2>, 5> sizes{{{.05, .02279225 - neck},
+                                                    {.035, .1245789 - arm},
+                                                    {.035, .1245789 - arm},
+                                                    {.055, .1822257 - leg},
+                                                    {.055, .1822257 - leg}}};
+  const std::array<std::array<double, 3>, 5> positions{
+      {{-.00165071, .0452401 + neck, .00534359},
+       {0, -.138421 - arm, 0},
+       {0, -.138421 - arm, 0},
+       {-5.0684e-8, -.202473 - leg, 0},
+       {-5.0684e-8, -.202473 - leg, 0}}};
+  const std::array<std::array<double, 4>, 5> quats{
+      {{.66437, .746906, .027253, 0},
+       {.612372, -.612372, .353553, .353553},
+       {.612372, -.612372, -.353553, -.353553},
+       {.696364, -.696364, -.122788, -.122788},
+       {.696364, -.696364, .122788, .122788}}};
   for (int i = 0; i < 5; ++i) {
     auto geom = Find(model, "body", bodies[i]).append_child("geom");
     Set(geom, "name", names[i]);
-    Set(geom, "size", Numbers(sizes[i], 2));
-    Set(geom, "pos", Numbers(positions[i], 3));
-    Set(geom, "quat", Numbers(quats[i], 4));
+    Set(geom, "size", Numbers(sizes[i].data(), 2));
+    Set(geom, "pos", Numbers(positions[i].data(), 3));
+    Set(geom, "quat", Numbers(quats[i].data(), 4));
     Set(geom, "mass", 0);
     Set(geom, "contype", 0);
     Set(geom, "conaffinity", 0);

envpool/mujoco/locomotion/locomotion_env.h

@@ -76,11 +76,14 @@ class LocomotionEnvSpec : public EnvSpec<LocomotionEnvFns> {
     Config config(values);
     const auto task = GetTaskConfig(config["task_name"_]);
     const int team = config["team_size"_];
-    if (team < 1 || team > 11)
+    if (team < 1 || team > 11) {
       throw std::invalid_argument("team_size must be between 1 and 11");
+    }
     config["max_num_players"_] = task.task == Task::kSoccer ? 2 * team : 1;
     double& limit = config["time_limit"_];
-    if (limit == -1) limit = task.time_limit;
+    if (limit == -1) {
+      limit = task.time_limit;
+    }
     if (!std::isfinite(limit) || limit <= 0 ||
         config["max_episode_steps"_] <= 0) {
       throw std::invalid_argument(
@@ -97,16 +100,16 @@ class LocomotionEnv : public Env<LocomotionEnvSpec>, public RenderableEnv {
   LocomotionEnv(const Spec& spec, int env_id)
       : Base(spec, env_id), simulation_(MakeOptions(spec, seed_)) {}
 
-  bool IsDone() override { return simulation_.done(); }
+  bool IsDone() override { return simulation_.Done(); }
 
   void Reset() override {
-    std::lock_guard<std::mutex> lock(mutex_);
+    std::scoped_lock lock(mutex_);
     simulation_.Reset();
     WriteState();
   }
 
   void Step(const Action& action) override {
-    std::lock_guard<std::mutex> lock(mutex_);
+    std::scoped_lock lock(mutex_);
     simulation_.Step(static_cast<const double*>(action["action"_].Data()));
     WriteState();
   }
@@ -117,9 +120,10 @@ class LocomotionEnv : public Env<LocomotionEnvSpec>, public RenderableEnv {
 
   void Render(int width, int height, int camera,
               unsigned char* output) override {
-    std::lock_guard<std::mutex> lock(mutex_);
-    if (simulation_.model() == nullptr)
+    std::scoped_lock lock(mutex_);
+    if (simulation_.Model() == nullptr) {
       throw std::runtime_error("reset before rendering");
+    }
     simulation_.Render(width, height, camera, output);
   }
 
@@ -128,9 +132,10 @@ class LocomotionEnv : public Env<LocomotionEnvSpec>, public RenderableEnv {
   // changes the simulator or synchronizes either side of an oracle rollout.
   template <typename Function>
   auto Inspect(Function function) {
-    std::lock_guard<std::mutex> lock(mutex_);
-    if (simulation_.model() == nullptr)
+    std::scoped_lock lock(mutex_);
+    if (simulation_.Model() == nullptr) {
       throw std::runtime_error("reset before inspecting physics");
+    }
     return function(simulation_);
   }
 #endif
@@ -154,7 +159,7 @@ class LocomotionEnv : public Env<LocomotionEnvSpec>, public RenderableEnv {
   }
 
   void WriteState() {
-    State state = Allocate(simulation_.players());
+    State state = Allocate(simulation_.Players());
     state["obs:continuous"_] = 0;
     state["obs:discrete"_] = 0;
     state["obs:pixels"_] = 0;
@@ -164,13 +169,13 @@ class LocomotionEnv : public Env<LocomotionEnvSpec>, public RenderableEnv {
               static_cast<int64_t*>(state["obs:discrete"_].Data()));
     std::copy(simulation_.pixels.begin(), simulation_.pixels.end(),
               static_cast<uint8_t*>(state["obs:pixels"_].Data()));
-    for (int i = 0; i < simulation_.players(); ++i) {
+    for (int i = 0; i < simulation_.Players(); ++i) {
       state["reward"_][i] = simulation_.rewards[i];
       state["reward64"_][i] = simulation_.rewards[i];
       state["discount"_][i] = simulation_.discount;
     }
-    state["terminated"_] = simulation_.terminated();
-    state["trunc"_] = simulation_.truncated();
+    state["terminated"_] = simulation_.Terminated();
+    state["trunc"_] = simulation_.Truncated();
   }
 
   Simulation simulation_;
@@ -196,8 +201,9 @@ class LocomotionEnvPool : public AsyncEnvPool<LocomotionEnv> {
 #ifdef ENVPOOL_TEST
   template <typename Function>
   auto Inspect(int env_id, Function function) {
-    if (env_id < 0 || env_id >= static_cast<int>(envs_.size()))
+    if (env_id < 0 || env_id >= static_cast<int>(envs_.size())) {
       throw std::out_of_range("invalid env_id");
+    }
     return envs_[env_id]->Inspect(function);
   }
 #endif
@@ -225,13 +231,15 @@ class LocomotionEnvPool : public AsyncEnvPool<LocomotionEnv> {
     for (std::size_t i = 0; i < actions[0].size; ++i) {
       const int id = env_ids[i];
       if (id < 0 || id >= static_cast<int>(envs_.size()) ||
-          !selected.insert(id).second)
+          !selected.insert(id).second) {
         throw std::invalid_argument("invalid or duplicate env_id");
+      }
     }
     for (std::size_t i = 0; i < actions[1].size; ++i) {
       const int id = player_ids[i];
-      if (!selected.count(id) || ++counts[id] > players)
+      if ((selected.count(id) == 0u) || ++counts[id] > players) {
         throw std::invalid_argument("expected one action for every player");
+      }
     }
   }
 };

envpool/mujoco/locomotion/locomotion_envpool.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <array>
 #include <cstring>
 #include <map>
 #include <string>
@@ -45,8 +46,8 @@ PYBIND11_MODULE(locomotion_envpool, m) {
         return self.Inspect(
             env_id, [include_model](const mujoco_locomotion::Simulation& sim) {
               py::dict result;
-              const mjModel* model = sim.model();
-              const mjData* data = sim.data();
+              const mjModel* model = sim.Model();
+              const mjData* data = sim.Data();
               for (const auto& [key, ptr, size] :
                    std::vector<std::tuple<const char*, const double*, int>>{
                        {"qpos", data->qpos, model->nq},
@@ -65,8 +66,8 @@ PYBIND11_MODULE(locomotion_envpool, m) {
                 std::vector<char> bytes(mj_sizeModel(model));
                 mj_saveModel(model, nullptr, bytes.data(), bytes.size());
                 result["model"] = py::bytes(bytes.data(), bytes.size());
-                result["maze"] = sim.scene().maze_entities;
-                result["variations"] = sim.scene().maze_variations;
+                result["maze"] = sim.GetScene().maze_entities;
+                result["variations"] = sim.GetScene().maze_variations;
               }
               return result;
             });

envpool/mujoco/locomotion/locomotion_test.py

@@ -54,7 +54,10 @@ def assert_pixels(
             "rodent_escape_bowl",
         }:
             peak, total = 1, 3
-        elif not egocentric and task == "cmu_humanoid_tracking":
+        elif not egocentric and task in {
+            "cmu_humanoid_tracking",
+            "cmu_humanoid_heterogeneous_forage",
+        }:
             peak = total = 1
     if total:
         np.testing.assert_equal(actual.shape, expected.shape)