Add body/simple attribute to control simple body optimization.

PiperOrigin-RevId: 942164766
Change-Id: I1e83bb99b6a1955917eb9b2e9bed0dc7b8e6a161

Author: yuvaltassa

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 @@
                  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;