Add body/simple attribute to control simple body optimization.

PiperOrigin-RevId: 942164766
Change-Id: I1e83bb99b6a1955917eb9b2e9bed0dc7b8e6a161

doc/XMLreference.rst

@@ -2152,6 +2152,31 @@ defined. Its body name is automatically defined as "world".
 
    See :ref:`implementation notes<siSleep>` for more details.
 
+.. _body-simple:
+
+:at:`simple`: :at-val:`[false, auto], "auto"`
+   Controls the *simple body* optimization. When a body qualifies as "simple", its inertial matrix block in the mass
+   matrix is diagonal, representing independent translational and rotational degrees of freedom. The optimization
+   omits storing of the zero-valued off-diagonal entries, reducing memory footprint and computation.
+
+   A body qualifies for this optimization if it satisfies all of the following:
+
+   - **Inertial frame alignment**: The body's inertial frame coincides with its body frame.
+   - **Kinematic root**: The body's parent is either the world body or a static body.
+   - **Leaf body**: The body is a leaf node in the kinematic tree (it has no child bodies).
+   - **Origin-centered joints**: All joints belonging to this body must reside at the body's origin.
+   - **Aligned joint axes**: Any hinge or slide joint axes must be aligned with the local coordinate axes, and at most
+     one joint with rotational degrees of freedom (hinge or ball) is permitted.
+   - **No inertia-bearing tendons**: The body must not contain sites or geoms used as wrap objects by any tendon that
+     has non-zero :ref:`armature<tendon-spatial-armature>`.
+
+   Setting this attribute to :at-val:`false` disables the optimization for this body. This is necessary for domain
+   randomization workflows where model parameters (such as joint/inertial offsets or angles) are perturbed dynamically
+   during simulation and updated via :ref:`mj_setConst`. Because a body compiled with the simple optimization active
+   cannot dynamically lose its simple state at runtime (which would require reallocation of sparse matrix structures),
+   any runtime parameter change that violates the simple conditions will trigger a validation error unless
+   ``simple="false"`` was explicitly declared in the XML.
+
 .. _body-user:
 
 :at:`user`: :at-val:`real(nbody_user), "0 0 ..."`

doc/XMLschema.rst

@@ -713,6 +713,9 @@
          .. grid-item::
             :ref:`sleep<body-sleep>`
 
+         .. grid-item::
+            :ref:`simple<body-simple>`
+
          .. grid-item::
             :ref:`user<body-user>`
 

doc/changelog.rst

@@ -12,6 +12,8 @@ General
 - :ref:`mj_encode` now supports encoding of MJB and TXT files.
 - :ref:`mj_setConst` now recomputes the ``mjModel.{body,geom,site}_sameframe`` flags, to account for changes in
   body/geom/site frames after compilation.
+- Added :ref:`body/simple<body-simple>` attribute ("false"/"auto") to disable the *simple body* mass matrix
+  optimization. This is useful for domain randomization, where model parameters may change post-compilation.
 - The :el:`attach` element now supports self-attachment (attaching elements of the current model to itself) by omitting
   the :at:`model` attribute. It also supports attaching a frame via the new :at:`frame` attribute, which is mutually
   exclusive with :at:`body`.

doc/includes/references.h

@@ -1661,6 +1661,7 @@ typedef struct mjsBody_ {          // body specification
   mjtByte mocap;                   // is this a mocap body
   double gravcomp;                 // gravity compensation
   mjtSleepPolicy sleep;            // sleep policy
+  mjtByte simple;                  // simple body optimization (0: false, 1: auto)
   mjDoubleVec* userdata;           // user data
   mjtByte explicitinertial;        // whether to save the body with explicit inertial clause
   mjsPlugin plugin;                // passive force plugin

include/mujoco/mjspec.h

@@ -279,6 +279,7 @@ typedef struct mjsBody_ {          // body specification
   mjtByte mocap;                   // is this a mocap body
   double gravcomp;                 // gravity compensation
   mjtSleepPolicy sleep;            // sleep policy
+  mjtByte simple;                  // simple body optimization (0: false, 1: auto)
   mjDoubleVec* userdata;           // user data
   mjtByte explicitinertial;        // whether to save the body with explicit inertial clause
   mjsPlugin plugin;                // passive force plugin

include/mujoco/mjspecmacro.h

@@ -114,6 +114,7 @@
     X   ( mjtByte,         mocap,            1 ) \
     X   ( double,          gravcomp,         1 ) \
     X   ( mjtSleepPolicy,  sleep,            1 ) \
+    X   ( mjtByte,         simple,           1 ) \
     X   ( mjDoubleVec*,    userdata,         1 ) \
     X   ( mjtByte,         explicitinertial, 1 ) \
     X   ( mjsPlugin,       plugin,           1 ) \

python/mujoco/introspect/structs.py

@@ -7391,6 +7391,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='mjtSleepPolicy'),
                  doc='sleep policy',
              ),
+             StructFieldDecl(
+                 name='simple',
+                 type=ValueType(name='mjtByte'),
+                 doc='simple body optimization (0: false, 1: auto)',
+             ),
              StructFieldDecl(
                  name='userdata',
                  type=PointerType(

src/engine/engine_setconst.c

@@ -1324,6 +1324,14 @@ void mj_setConst(mjModel* m, mjData* d) {
   // recompute sameframe flags from current model geometry
   setSameframe(m);
 
+  // error if simple body lost sameframe (user must set simple="false")
+  for (int i = 1; i < m->nbody; i++) {
+    if (m->body_simple[i] > 0 && m->body_sameframe[i] != mjSAMEFRAME_BODY) {
+      mjERROR("body %d is compiled as simple but sameframe no longer holds, "
+              "use body/simple='false'", i);
+    }
+  }
+
   // set fixed quantities
   setFixed(m, d);
 

src/user/user_init.c

@@ -82,6 +82,9 @@ void mjs_defaultBody(mjsBody* body) {
   body->ipos[0] = mjNAN;
   body->iquat[0] = 1;
   body->fullinertia[0] = mjNAN;
+
+  // simple optimization: auto
+  body->simple = 1;
 }
 
 

src/user/user_model.cc

@@ -1002,6 +1002,11 @@ void mjCModel::ComputeSparseSizes() {
                            bodies_[parentid]->parent          &&
                            bodies_[parentid]->parent->id == 0 &&
                            bodies_[parentid]->dofnum == 0)));
+
+    // user override: disable simple optimization
+    if (!pb->simple) {
+      body_simple_pre[i] = 0;
+    }
   }
 
   // a parent body is never simple (unless world)
@@ -2816,6 +2821,11 @@ void mjCModel::CopyTree(mjModel* m) {
                           (m->body_parentid[parentid] == 0 &&
                            m->body_dofnum[parentid] == 0)));
 
+    // user override: disable simple optimization
+    if (!pb->simple) {
+      m->body_simple[i] = 0;
+    }
+
     // a parent body is never simple (unless world)
     if (m->body_parentid[i] > 0) {
       m->body_simple[m->body_parentid[i]] = 0;

src/xml/xml_base.h

@@ -54,6 +54,7 @@ extern const mjMap enable_map[];
 extern const mjMap bool_map[];
 extern const mjMap fluid_map[];
 extern const mjMap TFAuto_map[];
+extern const mjMap FAuto_map[];
 extern const mjMap joint_map[];
 extern const mjMap bodysleep_map[];
 extern const mjMap geom_map[];

src/xml/xml_native_reader.cc

@@ -335,9 +335,8 @@ std::vector<const char*> MJCF[nMJCF] = {
         {">"},
         {"model", "*", "name", "file", "content_type"},
     {">"},
-
-    {"body", "R", "name", "childclass", "pos", "quat", "mocap",
-        "axisangle", "xyaxes", "zaxis", "euler", "gravcomp", "sleep", "user"},
+    {"body", "R", "name", "childclass", "pos", "quat", "mocap", "axisangle",
+        "xyaxes", "zaxis", "euler", "gravcomp", "sleep", "simple", "user"},
     {"<"},
         {"inertial", "?", "pos", "quat", "mass", "diaginertia",
             "axisangle", "xyaxes", "zaxis", "euler", "fullinertia"},
@@ -644,6 +643,13 @@ const mjMap TFAuto_map[3] = {
 };
 
 
+// FAuto type
+const mjMap FAuto_map[2] = {
+  {"false",   0},
+  {"auto",    1}
+};
+
+
 // body sleep type
 const int bodysleep_sz = 4;
 const mjMap bodysleep_map[bodysleep_sz] = {
@@ -3926,6 +3932,9 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
       if (MapValue(elem, "sleep", &n, bodysleep_map, bodysleep_sz)) {
         child->sleep = (mjtSleepPolicy) n;
       }
+      if (MapValue(elem, "simple", &n, FAuto_map, 2)) {
+        child->simple = (mjtByte) n;
+      }
 
       // read userdata
       std::vector<double> userdata;

src/xml/xml_native_writer.cc

@@ -1726,6 +1726,9 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi
       WriteAttrKey(elem, "sleep", bodysleep_map, bodysleep_sz, body->sleep);
     }
 
+    // simple optimization
+    WriteAttrKey(elem, "simple", FAuto_map, 2, body->simple, 1);
+
     // userdata
     WriteVector(elem, "user", body->get_userdata());
 

test/engine/engine_setconst_test.cc

@@ -329,10 +329,12 @@ TEST_F(SetConstTest, TendonTreeId) {
   int t_static_id = mj_name2id(model.get(), mjOBJ_TENDON, "T_static");
   int t_tree1_id = mj_name2id(model.get(), mjOBJ_TENDON, "T_tree1");
   int t_intertree12_id = mj_name2id(model.get(), mjOBJ_TENDON, "T_intertree12");
-  int t_intertree123_id = mj_name2id(model.get(), mjOBJ_TENDON, "T_intertree123");
-
-  int b1_1_treeid = model->body_treeid[mj_name2id(model.get(), mjOBJ_BODY, "B1_1")];
-  int b2_1_treeid = model->body_treeid[mj_name2id(model.get(), mjOBJ_BODY, "B2_1")];
+  int t_intertree123_id =
+      mj_name2id(model.get(), mjOBJ_TENDON, "T_intertree123");
+  int b1_1_treeid =
+      model->body_treeid[mj_name2id(model.get(), mjOBJ_BODY, "B1_1")];
+  int b2_1_treeid =
+      model->body_treeid[mj_name2id(model.get(), mjOBJ_BODY, "B2_1")];
 
   // Tendon 1: Not associated with any tree
   EXPECT_EQ(model->tendon_treenum[t_static_id], 0);
@@ -461,7 +463,7 @@ TEST_F(SetConstTest, BodySameframeRecomputed) {
   constexpr char xml[] = R"(
   <mujoco>
     <worldbody>
-      <body name="B1">
+      <body name="B1" simple="false">
         <joint type="slide"/>
         <geom size=".1"/>
       </body>
@@ -570,7 +572,7 @@ TEST_F(SetConstTest, SameframeKinematicsCorrect) {
   constexpr char xml[] = R"(
   <mujoco>
     <worldbody>
-      <body name="B1" pos="1 0 0">
+      <body name="B1" pos="1 0 0" simple="false">
         <joint type="slide" axis="1 0 0"/>
         <geom name="G1" size=".1"/>
         <site name="S1"/>
@@ -604,5 +606,34 @@ TEST_F(SetConstTest, SameframeKinematicsCorrect) {
   EXPECT_NEAR(d->geom_xpos[3*g+2], 0.3, MjTol(1e-10, 1e-6));
 }
 
+TEST_F(SetConstTest, SimpleBodyLostSameframeError) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="B1">
+        <joint type="slide"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(m.get(), NotNull()) << error;
+  MjDataPtr d(mj_makeData(m.get()));
+
+  int b = mj_name2id(m.get(), mjOBJ_BODY, "B1");
+
+  // confirm body is compiled as simple
+  EXPECT_GT(m->body_simple[b], 0);
+
+  // perturb body_ipos, breaking sameframe; calling mj_setConst should fail
+  m->body_ipos[3*b+0] = 1.0;
+
+  std::string err = MjuErrorMessageFrom(mj_setConst)(m.get(), d.get());
+  EXPECT_THAT(err, HasSubstr("body 1 is compiled as simple but "
+                             "sameframe no longer holds"));
+}
+
 }  // namespace
 }  // namespace mujoco

test/mjspecmacro_test.cc

@@ -140,7 +140,7 @@ CHECK_STRUCT(mjsDefault, MJSDEFAULT_FIELDS)
 #undef XIMPL
 
 TEST(MjspecmacroTest, CompileTimeChecks) {
-  // Verifies that the test binary compiled successfully and all static_asserts passed.
+  // Verifies that the test compiled successfully and all static_asserts passed.
   EXPECT_TRUE(true);
 }
 

test/user/user_objects_test.cc

@@ -1287,6 +1287,55 @@ TEST_F(MjCJointTest, AlignFree) {
   mj_deleteSpec(s);
 }
 
+TEST_F(MjCJointTest, BodySimpleFalse) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="B1">
+        <joint type="slide"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  constexpr char xml_nosimple[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="B1" simple="false">
+        <joint type="slide"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+
+  // compile default model: body should be simple
+  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(m.get(), NotNull()) << error;
+  int b = mj_name2id(m.get(), mjOBJ_BODY, "B1");
+  EXPECT_GT(m->body_simple[b], 0);
+  EXPECT_GT(m->dof_simplenum[0], 0);
+
+  // compile with simple="false": body should not be simple
+  MjModelPtr m_ns = LoadModelFromString(xml_nosimple, error, sizeof(error));
+  ASSERT_THAT(m_ns.get(), NotNull()) << error;
+  EXPECT_EQ(m_ns->body_simple[b], 0);
+  EXPECT_EQ(m_ns->dof_simplenum[0], 0);
+
+  // forward kinematics should produce identical results
+  MjDataPtr d = MakeData(m);
+  MjDataPtr d_ns = MakeData(m_ns);
+  d->qpos[0] = d_ns->qpos[0] = 0.5;
+  mj_forward(m.get(), d.get());
+  mj_forward(m_ns.get(), d_ns.get());
+  EXPECT_THAT(d_ns->xpos[3*b+0], MjNear(d->xpos[3*b+0], 1e-10, 1e-6));
+  EXPECT_THAT(d_ns->xpos[3*b+1], MjNear(d->xpos[3*b+1], 1e-10, 1e-6));
+  EXPECT_THAT(d_ns->xpos[3*b+2], MjNear(d->xpos[3*b+2], 1e-10, 1e-6));
+}
+
 // ------------- test height fields --------------------------------------------
 
 using MjCHFieldTest = MujocoTest;

test/xml/xml_native_writer_test.cc

@@ -1590,5 +1590,41 @@ TEST_F(XMLWriterTest, WritesActuatorDelayHistory) {
   EXPECT_THAT(saved_xml, HasSubstr("interp=\"cubic\""));
 }
 
+TEST_F(XMLWriterTest, BodySimpleRoundtrip) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="B1" simple="false">
+        <joint type="ball"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  MjModelPtr model = LoadModelFromString(xml);
+  ASSERT_THAT(model.get(), NotNull());
+  std::string saved_xml = SaveAndReadXml(model.get());
+  EXPECT_THAT(saved_xml, HasSubstr("simple=\"false\""));
+
+  // auto should not be written
+  static constexpr char xml_auto[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="B1">
+        <joint type="ball"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  MjModelPtr model_auto = LoadModelFromString(xml_auto);
+  ASSERT_THAT(model_auto.get(), NotNull());
+  std::string saved_auto = SaveAndReadXml(model_auto.get());
+  EXPECT_THAT(saved_auto, Not(HasSubstr("simple=")));
+
+  // nC should increase when simple is disabled
+  EXPECT_GT(model->nC, model_auto->nC);
+}
+
 }  // namespace
 }  // namespace mujoco

wasm/codegen/generated/bindings.cc

@@ -5593,6 +5593,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("plugin", &MjsBody::plugin, reference())
     .property("pos", &MjsBody::pos)
     .property("quat", &MjsBody::quat)
+    .property("simple", &MjsBody::simple, &MjsBody::set_simple, reference())
     .property("sleep", &MjsBody::sleep, &MjsBody::set_sleep, reference())
     .property("userdata", &MjsBody::userdata, reference());
   emscripten::class_<MjsCamera>("MjsCamera")

wasm/codegen/generated/bindings.h

@@ -6087,6 +6087,12 @@ struct MjsBody {
   void set_sleep(mjtSleepPolicy value) {
     ptr_->sleep = value;
   }
+  mjtByte simple() const {
+    return ptr_->simple;
+  }
+  void set_simple(mjtByte value) {
+    ptr_->simple = value;
+  }
   mjDoubleVec &userdata() const {
     return *(ptr_->userdata);
   }