Contact forces silently understated for multi-geom body segments (FlyBody)

[sibocw]

Summary

Simulation.get_bodysegment_contact_forces reports contact forces for only
one geom per body segment — the last one registered. Any segment built from
more than one collision geom (the FlyBody model has ~15 of them: c_thorax,
c_head, the abdomen/wing/antenna segments, …) has the contacts on all of its
other geoms silently dropped, so the reported force is understated or zero.
There is no error — the number just comes out wrong.

NeuroMechFly / the musculoskeletal model are unaffected because each of their
segments maps to exactly one geom.

Root cause

In #279 fly.bodyseg_to_mjcfgeom became dict[BodySegment, list[geom]] (a
segment can own several geoms). The contact-force consumer was not updated to
match — it still assumes one geom per segment.

Simulation._map_internal_geom_ids (src/flygym/simulation.py:516-532) writes
a scalar id inside the per-geom loop, so only the last geom survives:

for bodyseg, mjcf_geom_elements in fly.bodyseg_to_mjcfgeom.items():
    for mjcf_geom_element in mjcf_geom_elements:
        internal_geom_id = mj.mj_name2id(self.mj_model, mj.mjtObj.mjOBJ_GEOM,
                                         mjcf_geom_element.name)
        internal_geomids_by_bodyseg_by_fly[fly_name][bodyseg] = internal_geom_id  # last-wins

get_bodysegment_contact_forces (src/flygym/simulation.py:283-326) then builds
a one-id-per-segment lookup and matches contacts against it via np.isin, so
contacts on the dropped geoms never match:

geom_ids_by_segment = self._internal_geomid_by_bodyseg_by_fly[fly_name]
requested_geom_to_output = {geom_ids_by_segment[seg]: i
                            for i, seg in enumerate(requested_segments)}

Why this is a design issue (not a one-line patch)

The "one geom id per segment" assumption is baked into the
_internal_geomid_by_bodyseg_by_fly map and its consumer, and it implies a
semantic decision: a multi-geom segment's contact force should be the sum
over its geoms. So the fix is to (a) make the map hold all geom ids per
segment and (b) update the consumer (and any other reader of that map) to
aggregate. Worth doing deliberately with a regression test rather than slipping
into the release PR.

Suggested fix

1. Store a list of geom ids per segment in _map_internal_geom_ids:

   internal_geomids_by_bodyseg_by_fly[fly_name][bodyseg] = [
       mj.mj_name2id(self.mj_model, mj.mjtObj.mjOBJ_GEOM, g.name)
       for g in mjcf_geom_elements
   ]

2. In get_bodysegment_contact_forces, build the reverse map from every
   geom of each segment to that segment's output row, so a contact on any of a
   segment's geoms accumulates into the same row:

   requested_geom_to_output = {
       geom_id: i
       for i, seg in enumerate(requested_segments)
       for geom_id in geom_ids_by_segment[seg]
   }

   The rest of the method (the np.isin filter and the per-contact
   accumulation at lines 322-326) then works unchanged and naturally sums
   contributions from all geoms of a segment.

3. Audit other readers of _internal_geomid_by_bodyseg_by_fly for the same
   scalar assumption.

Test

Add a regression test on a FlyBody (or a synthetic multi-geom segment) that
places contacts on a non-last geom of a multi-geom segment and asserts the
reported force is non-zero / equals the analytic sum. The current code returns
~0 for that case.

---

Found during review of #279.

[omnilink-tech]

Two things that may be useful here: the fix appears to have already landed in an
open PR, and the regression test this issue asked for is not in it.

The code change is in #305

src/flygym/simulation.py in #305 makes exactly the change described above --
_map_internal_geom_ids stores a list per segment, and requested_geom_to_output
maps every geom of a segment to the same output row so their forces sum:

requested_geom_to_output = {
    gid: i
    for i, seg in enumerate(requested_segments)
    for gid in geom_ids_by_segment[seg]
}

That reads as a correct fix for the root cause named here. But #305's test
changes are confined to tests/core/test_compose.py and
tests/core/test_rendering.py -- neither exercises
get_bodysegment_contact_forces. So the behaviour this issue is about would
be fixed without a test, which is the outcome the issue text specifically wanted
to avoid ("worth doing deliberately with a regression test rather than slipping
into the release PR").

Since #294 has no comments, I suspect the two just have not been linked.

The magnitude, measured

A minimal check of what MuJoCo actually reports for a multi-geom body -- one
body, mass 0.02 kg, two box geoms, resting on a plane, mj_contactForce
summed per geom:

contacts reported          : 8
|force| on part_a          : 0.098100 N
|force| on part_b          : 0.098100 N
SUM over both geoms        : 0.196200 N
analytic m*g               : 0.196200 N
error of the sum           : 4.59e-10 N
if only ONE geom is counted: 0.098100 N  -> understated by 50.0%

So contacts are reported per geom, the sum recovers m*g to ~5e-10 N, and
last-wins costs exactly the fraction of the segment's contact area carried by
the dropped geoms -- 50% for a symmetric two-geom segment, and 0% only when the
contact happens to fall on the geom that won. That last part is why the failure
is intermittent-looking rather than obviously broken.

The test

Analytic rather than golden, so it cannot drift with the model:

def test_multigeom_segment_contact_force_sums_over_all_geoms():
    """A body segment built from several collision geoms must report the SUM of
    the contact forces on all of them.

    MuJoCo reports contacts per geom, so a consumer keeping one geom id per
    segment silently understates the force -- by 50% for a symmetric two-geom
    segment, and by 0% when the contact happens to land on the geom that won,
    which is what makes the bug look intermittent.
    """
    mass, g = 0.02, 9.81
    model = mujoco.MjModel.from_xml_string(TWO_GEOM_SEGMENT_XML)   # see below
    data = mujoco.MjData(model)
    for _ in range(3000):
        mujoco.mj_step(model, data)

    total = np.zeros(3)
    buf = np.zeros(6)
    seg_geoms = {mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, n)
                 for n in ("part_a", "part_b")}
    for i in range(data.ncon):
        contact = data.contact[i]
        if seg_geoms & {contact.geom1, contact.geom2}:
            mujoco.mj_contactForce(model, data, i, buf)
            total += buf[0] * contact.frame[0:3]

    assert np.linalg.norm(total) == pytest.approx(mass * g, rel=1e-6)

with

TWO_GEOM_SEGMENT_XML = """
<mujoco>
  <option gravity="0 0 -9.81"/>
  <worldbody>
    <geom name="floor" type="plane" size="5 5 .1"/>
    <body name="segment" pos="0 0 .05">
      <freejoint/>
      <inertial pos="0 0 0" mass="0.02" diaginertia="1e-6 1e-6 1e-6"/>
      <geom name="part_a" type="box" size=".05 .05 .05" pos="-.05 0 0"/>
      <geom name="part_b" type="box" size=".05 .05 .05" pos=".05 0 0"/>
    </body>
  </worldbody>
</mujoco>"""

The directed case the issue asks for -- a contact on a non-last geom -- falls
out of the same fixture by shifting the body so only part_a touches; the
last-wins code returns ~0 there while the summed version returns m*g.

Ideally this would go through Simulation.get_bodysegment_contact_forces on a
real FlyBody rather than raw MuJoCo, so it guards your API rather than
MuJoCo's. I wrote it at the MuJoCo level because I do not have a flygym
environment here and did not want to guess at the fixture conventions -- happy
to rework it against Simulation if you tell me which fixture to build on, or
to open it against #305's branch so it lands with the fix rather than competing
with it.