From 78a4d2f3894fcb1f962621c5216076439708f620 Mon Sep 17 00:00:00 2001
From: Tom Power <tjpower@google.com>
Date: Fri, 14 Feb 2025 11:25:31 -0800
Subject: [PATCH] Fix mass/inertia checking for moving bodies with nested
 children.

Previously, we only checked immediate children of moving bodies.

PiperOrigin-RevId: 726999705
Change-Id: I133a557daf86223b8751c6d46e25809781e168a1
---
 src/user/user_model.cc       | 72 +++++++++++++++++++++-------------
 src/user/user_model.h        |  7 ++++
 test/user/user_model_test.cc | 76 ++++++++++++++++++++++++++++++++++++
 3 files changed, 128 insertions(+), 27 deletions(-)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index fc21d29f..02c8a018 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -4319,33 +4319,11 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     }
   }
 
-  // check body mass and inertia
-  for (int i=1; i<bodies_.size(); i++) {
-    mjCBody* b = bodies_[i];
-
-    // find moving body with small mass or inertia
-    if (!b->joints.empty() &&
-        (b->mass<mjMINVAL ||
-          b->inertia[0]<mjMINVAL ||
-          b->inertia[1]<mjMINVAL ||
-          b->inertia[2]<mjMINVAL)) {
-      // does it have static children with mass and inertia
-      bool ok = false;
-      for (size_t j=0; j<b->bodies.size(); j++) {
-        if (b->bodies[j]->joints.empty() &&
-            b->bodies[j]->mass>=mjMINVAL &&
-            b->bodies[j]->inertia[0]>=mjMINVAL &&
-            b->bodies[j]->inertia[1]>=mjMINVAL &&
-            b->bodies[j]->inertia[2]>=mjMINVAL) {
-          ok = true;
-          break;
-        }
-      }
-
-      // error
-      if (!ok) {
-        throw mjCError(b, "mass and inertia of moving bodies must be larger than mjMINVAL");
-      }
+  // check mass and inertia of moving bodies
+  if (bodies_.size() > 1) {
+    // we ignore the first body as it is the world body
+    if (!CheckBodiesMassInertia(std::vector<mjCBody*>(bodies_.begin()+1, bodies_.end()))) {
+      throw mjCError(0, "mass and inertia of moving bodies must be larger than mjMINVAL");
     }
   }
 
@@ -4497,6 +4475,46 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
 
 
 
+bool mjCModel::CheckBodiesMassInertia(std::vector<mjCBody*> bodies) {
+  // check mass and inertia of moving bodies
+  for (int i=0; i<bodies.size(); i++) {
+    if (!bodies[i]->joints.empty()) {
+      if (!CheckBodyMassInertia(bodies[i])) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+
+
+bool mjCModel::CheckBodyMassInertia(mjCBody* body) {
+  // check if body has valid mass and inertia
+  if (body->mass>=mjMINVAL &&
+      body->inertia[0]>=mjMINVAL &&
+      body->inertia[1]>=mjMINVAL &&
+      body->inertia[2]>=mjMINVAL) {
+      return true;
+    }
+
+  // body is valid if we find a single static child with valid mass and inertia
+  for (int i=0; i<body->Bodies().size(); i++) {
+    // if we find a child with a joint, time to move on to the next moving body
+    if (!body->Bodies()[i]->joints.empty()) {
+      continue;
+    }
+    if (CheckBodyMassInertia(body->Bodies()[i])) {
+      return true;
+    }
+  }
+
+  // we did not find a child with valid mass and inertia
+  return false;
+}
+
+
+
 //------------------------------- DECOMPILER -------------------------------------------------------
 
 // get numeric data back from mjModel
diff --git a/src/user/user_model.h b/src/user/user_model.h
index 119654a3..7a99f6d8 100644
--- a/src/user/user_model.h
+++ b/src/user/user_model.h
@@ -414,6 +414,13 @@ class mjCModel : public mjCModel_, private mjSpec {
   // compute qpos0
   void ComputeReference();
 
+  // return true if all bodies have valid mass and inertia
+  bool CheckBodiesMassInertia(std::vector<mjCBody*> bodies);
+
+  // return true if body has valid mass and inertia
+  bool CheckBodyMassInertia(mjCBody* body);
+
+
   mjListKeyMap ids;   // map from object names to ids
   mjCError errInfo;   // last error info
   std::vector<mjKeyInfo> key_pending_;  // attached keyframes
diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc
index 2a17bc00..9f85669e 100644
--- a/test/user/user_model_test.cc
+++ b/test/user/user_model_test.cc
@@ -140,6 +140,82 @@ TEST_F(UserCModelTest, ActuatorSparsity) {
   mj_deleteModel(m);
 }
 
+TEST_F(UserCModelTest, NestedZeroMassBodiesOK) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <freejoint/>
+        <body>
+          <body>
+            <body>
+              <geom size="1"/>
+            </body>
+          </body>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  mj_deleteModel(model);
+}
+
+TEST_F(UserCModelTest, NestedZeroMassBodiesWithJointOK) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <freejoint/>
+        <body>
+          <body>
+            <body>
+              <joint/>
+              <geom size="1"/>
+            </body>
+            <body>
+              <geom size="1"/>
+            </body>
+          </body>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  mj_deleteModel(model);
+}
+
+TEST_F(UserCModelTest, NestedZeroMassBodiesFail) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <geom size="1"/>
+        <body>
+          <freejoint/>
+          <body>
+            <body>
+            </body>
+          </body>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, IsNull());
+  EXPECT_THAT(
+      error,
+      HasSubstr(
+          "mass and inertia of moving bodies must be larger than mjMINVAL"));
+  mj_deleteModel(model);
+}
 
 // ------------- test automatic inference of nuser_xxx -------------------------