Allow contact sensor subtree1/subtree2 to be any body.

PiperOrigin-RevId: 799632590
Change-Id: If9cd6c5a3a5f83d89ae024a8f64232ab1f1d603a

Author: yuvaltassa

doc/XMLreference.rst

@@ -7502,9 +7502,7 @@ Extraction
 .. _sensor-contact-subtree2:
 
 :at:`subtree1`, :at:`subtree2`: :at-val:`string, optional`
-   Name of a body whose subtree is participating in a contact. See **matching** :ref:`above <sensor-contact>`. Note
-   currently only "entire" subtrees are supported, in the sense that the specified body must be a direct child of the
-   world. General subtrees could be added in the future.
+   Name of a body whose subtree is participating in a contact. See **matching** :ref:`above <sensor-contact>`.
 
 .. _sensor-contact-site:
 

doc/changelog.rst

@@ -10,6 +10,8 @@ General
 - Constraint island discovery and construction, previously an experimental feature, is now :ref:`documented<soIsland>`
   and promoted to default; disable it with :ref:`option/flag/island <option-flag-island>`. We expect islanding to be
   a strict improvement over the monolithic constraint solver, please let us know if you experience any issues.
+- :ref:`Contact sensor<sensor-contact>` :at-val:`subtree1/subtree2` specification is now available for any body, not
+  just direct children of the world.
 
 .. admonition:: Breaking API changes
    :class: attention

python/mujoco/specs_test.py

@@ -1428,22 +1428,6 @@ def test_bad_contact_sensor(self):
                 refname='cam',
             ),
         ),
-        dict(
-            expected_error='subtree1 must be a child of the world',
-            sensor_params=dict(
-                type=mujoco.mjtSensor.mjSENS_CONTACT,
-                objtype=mujoco.mjtObj.mjOBJ_XBODY,
-                objname='non_root',
-            ),
-        ),
-        dict(
-            expected_error='subtree2 must be a child of the world',
-            sensor_params=dict(
-                type=mujoco.mjtSensor.mjSENS_CONTACT,
-                reftype=mujoco.mjtObj.mjOBJ_XBODY,
-                refname='non_root',
-            ),
-        ),
     ]
 
     for params in test_cases:

src/engine/engine_sensor.c

@@ -247,7 +247,14 @@ static int checkMatch(const mjModel* m, int body, int geom, mjtObj type, int id)
   if (type == mjOBJ_SITE)    return 1;  // already passed site filter test
   if (type == mjOBJ_GEOM)    return id == geom;
   if (type == mjOBJ_BODY)    return id == body;
-  if (type == mjOBJ_XBODY)   return body >= 0 && m->body_rootid[id] == m->body_rootid[body];
+  if (type == mjOBJ_XBODY) {
+    // traverse up the tree from body, return true if we land on id
+    while (body > id) {
+      body = m->body_parentid[body];
+    }
+    return body == id;
+  }
+
   return 0;
 }
 

src/user/user_objects.cc

@@ -7152,11 +7152,6 @@ void mjCSensor::Compile(void) {
           throw mjCError(this, "first matching criterion: if set, must be (x)body, geom or site");
         }
 
-        // check that subtree1 is a full tree
-        if (objtype == mjOBJ_XBODY && static_cast<mjCBody*>(obj)->GetParent()->id != 0) {
-          throw mjCError(this, "subtree1 must be a child of the world");
-        }
-
         // check second matching criterion
         if (reftype != mjOBJ_BODY &&
             reftype != mjOBJ_XBODY &&
@@ -7165,11 +7160,6 @@ void mjCSensor::Compile(void) {
           throw mjCError(this, "second matching criterion: if set, must be (x)body or geom");
         }
 
-        // check that subtree2 is a full tree
-        if (reftype == mjOBJ_XBODY && static_cast<mjCBody*>(ref)->GetParent()->id != 0) {
-          throw mjCError(this, "subtree2 must be a child of the world");
-        }
-
         // check for dataspec correctness
         int dataspec = intprm[0];
         if (dataspec <= 0) {

test/engine/engine_sensor_test.cc

@@ -708,8 +708,6 @@ TEST_F(SensorTest, BadContact) {
        "at most one of (geom1, body1, subtree1, site) can be specified"},
       {"geom2='sphere1' body2='body'",
        "at most one of (geom2, body2, subtree2) can be specified"},
-      {"subtree1='non_root'",
-       "must be a child of the world"}
   };
 
   for (const auto& test : test_cases) {
@@ -870,6 +868,31 @@ TEST_F(SensorTest, ContactSubtree) {
   mj_deleteModel(model);
 }
 
+TEST_F(SensorTest, ContactSubtreePartial) {
+  const string xml_path =
+      GetTestDataFilePath("engine/testdata/sensor/contact_subtree_partial.xml");
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+
+  mjData* data = mj_makeData(model);
+
+  while (data->time < 0.6) {
+    mj_step(model, data);
+  }
+
+  EXPECT_EQ(GetSensor(model, data, "all")[0],     4);
+  EXPECT_EQ(GetSensor(model, data, "world")[0],   4);
+  EXPECT_EQ(GetSensor(model, data, "thigh")[0],   4);
+  EXPECT_EQ(GetSensor(model, data, "shin")[0],    2);
+  EXPECT_EQ(GetSensor(model, data, "foot")[0],    1);
+  EXPECT_EQ(GetSensor(model, data, "foot_w")[0],  0);
+  EXPECT_EQ(GetSensor(model, data, "foot_w2")[0], 1);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 TEST_F(SensorTest, ContactNet) {
   const string xml_path =
       GetTestDataFilePath("engine/testdata/sensor/contact_net.xml");

test/engine/testdata/sensor/contact_subtree_partial.xml

@@ -0,0 +1,28 @@
+<mujoco model="contact subtree">
+  <worldbody>
+    <light pos="0 0 5"/>
+    <geom name="floor" type="plane" size="2 2 .01"/>
+    <body name="thigh" pos="-1 0 .1">
+      <joint type="slide"/>
+      <geom type="capsule" size=".1" fromto="0 0 0 .5 0 0"/>
+      <body name="shin" pos=".7 0 0">
+        <joint axis="0 1 0"/>
+        <geom type="capsule" size=".1" fromto="0 0 0 0 0 .5"/>
+        <body name="foot" pos="0 0 .7">
+          <joint axis="0 1 0"/>
+          <geom type="capsule" size=".1" fromto="0 0 0 -.7 0 0"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <sensor>
+    <contact name="all"/>
+    <contact name="world" subtree1="world"/>
+    <contact name="thigh" subtree1="thigh"/>
+    <contact name="shin" subtree1="shin"/>
+    <contact name="foot" subtree1="foot"/>
+    <contact name="foot_w" subtree1="foot" body2="world"/>
+    <contact name="foot_w2" subtree1="foot" subtree2="world"/>
+  </sensor>
+</mujoco>