Rework contact APIs

.gitignore

@@ -29,3 +29,4 @@ site/
 # Temporary files
 .DS_Store
 Thumbs.db
+.claude

crazyflow/sim/sim.py

@@ -43,6 +43,8 @@
 from crazyflow.utils import grid_2d, pytree_replace, world_mask
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from jax.sharding import Mesh
     from mujoco.mjx import Data, Model
     from numpy.typing import NDArray
@@ -142,6 +144,7 @@ def __init__(
         self.drone_path = Path(__file__).parents[1] / f"drones/{drone}.xml"
         self.spec = self.build_mjx_spec()
         self.mj_model, self.mj_data, self.mjx_model, self.mjx_data = self.build_mjx_model(self.spec)
+        self._contacts = self.build_contact_fn()
         self.viewer: MujocoRenderer | None = None
 
         freqs = (state_freq, attitude_freq, body_rate_freq, force_torque_freq)
@@ -509,6 +512,7 @@ def build_mjx(self):
             self.viewer = None
         self.spec.numeric("max_geom_pairs").data = [self.max_geom_pairs]
         self.mj_model, self.mj_data, self.mjx_model, self.mjx_data = self.build_mjx_model(self.spec)
+        self._contacts = self.build_contact_fn()
 
     def init_data(
         self,
@@ -575,8 +579,11 @@ def controllable(self) -> Array:
         return F.controllable(self.data)
 
     @requires_mujoco_sync
-    def contacts(self, body: str | None = None) -> Array:
-        """Get contact information from the simulation.
+    def contacts(self) -> Array:
+        """Flag each drone that is in contact with any other body.
+
+        Use [build_contact_fn][crazyflow.sim.Sim.build_contact_fn] to ignore bodies or to check
+        contacts on your own MJX data.
 
         Note:
             ``sim.max_geom_pairs`` limits the maximum detectable collision contacts per collision
@@ -587,20 +594,62 @@ def contacts(self, body: str | None = None) -> Array:
             you need to truly detect all contacts, set ``sim.max_geom_pairs`` to -1 and rebuild the
             simulation.
 
+        Returns:
+            A boolean array of shape (n_worlds, n_drones).
+        """
+        return self._contacts(self.mjx_data)
+
+    def build_contact_fn(
+        self, ignore: Sequence[str | int | tuple[str | int, str | int]] = ()
+    ) -> Callable[[Data], Array]:
+        """Build a function that flags each drone in contact with any other body.
+
         Args:
-            body: Optional body name to filter contacts for. If None, returns flags for all bodies.
+            ignore: Bodies, given by name or id, whose contacts are not counted. A single body is
+                ignored for all drones, a pair of bodies only for that pair.
 
         Returns:
-            A boolean array of shape (n_worlds, n_contacts), one flag per slot in the contact
-            buffer. Which geoms a slot holds is given by the matching entries of
-            ``sim.mjx_data._impl.contact.geom1`` and ``geom2``.
+            A contact function from MJX data to a boolean contact array (n_worlds, n_drones).
         """
-        if body is None:
-            return self.mjx_data._impl.contact.dist < 0
-        body_id = self.mj_model.body(body).id
-        geom_start = self.mj_model.body_geomadr[body_id]
-        geom_count = self.mj_model.body_geomnum[body_id]
-        return contacts(geom_start, geom_count, self.mjx_data)
+        mj_model = self.mj_model
+        base = "drone_fused" if self.fused_mjx_model else "drone"
+        drone_ids = [mj_model.body(f"{base}:{i}").id for i in range(self.n_drones)]
+        no_drone = self.n_drones  # One past the last drone, gets ignored later
+        body2drone = np.full(mj_model.nbody, no_drone)
+        body2drone[drone_ids] = np.arange(self.n_drones)
+        # We add one row that absorbes all ignored contacts for non-drone bodies
+        ignored = np.zeros((self.n_drones + 1, mj_model.nbody), dtype=bool)
+        for entry in ignore:
+            if isinstance(entry, str | int):  # Single bodies are paired with all drones
+                b = mj_model.body_rootid[mj_model.body(entry).id]
+                ignored[:, b] = ignored[body2drone[b], drone_ids] = True
+            else:
+                a, b = mj_model.body_rootid[[mj_model.body(body).id for body in entry]]
+                ignored[body2drone[a], b] = ignored[body2drone[b], a] = True
+        geom2body = mj_model.body_rootid[mj_model.geom_bodyid]
+        contact_fn = jax.jit(
+            partial(
+                contacts,
+                geom2drone=jnp.asarray(body2drone[geom2body]),
+                geom2body=jnp.asarray(geom2body),
+                ignored=jnp.asarray(ignored[:-1]),
+            )
+        )
+        self._contacts = contact_fn
+        return contact_fn
+
+    @requires_mujoco_sync
+    def contact_bodies(self) -> tuple[Array, Array, Array]:
+        """Return the bodies on both sides of each contact.
+
+        Bodies are MuJoCo body ids, e.g. from ``sim.mj_model.body(name).id``. Child bodies count as
+        their root body.
+
+        Returns:
+            The two body id arrays and the contact flag, each of shape (n_worlds, n_contacts).
+        """
+        geom2body = self.mj_model.body_rootid[self.mj_model.geom_bodyid]
+        return contact_bodies(self.mjx_data, geom2body)
 
     @staticmethod
     def _reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
@@ -697,15 +746,23 @@ def increment_steps(data: SimData) -> SimData:
     return data.replace(core=data.core.replace(steps=data.core.steps + 1))
 
 
+def contacts(data: Data, *, geom2drone: Array, geom2body: Array, ignored: Array) -> Array:
+    """Detect contacts for each drone in each world with all pairs that are not ignored."""
+    contact = data._impl.contact
+    own_geom = jnp.concat((contact.geom1, contact.geom2), axis=-1)  # Each contact from both sides
+    other_geom = jnp.concat((contact.geom2, contact.geom1), axis=-1)
+    drone_ids = geom2drone[own_geom]
+    n_drones = ignored.shape[0]
+    touching = jnp.tile(contact.dist < 0, 2) & ~ignored[drone_ids, geom2body[other_geom]]
+    any_per_drone = partial(jax.ops.segment_max, num_segments=n_drones)
+    return jax.vmap(any_per_drone)(touching, drone_ids)
+
+
 @jax.jit
-def contacts(geom_start: int, geom_count: int, data: Data) -> Array:
-    """Filter contacts from MuJoCo data."""
+def contact_bodies(data: Data, geom2body: Array) -> tuple[Array, Array, Array]:
+    """Map both geoms of each contact to their bodies."""
     contact = data._impl.contact
-    geom1_valid = contact.geom1 >= geom_start
-    geom1_valid &= contact.geom1 < geom_start + geom_count
-    geom2_valid = contact.geom2 >= geom_start
-    geom2_valid &= contact.geom2 < geom_start + geom_count
-    return (contact.dist < 0) & (geom1_valid | geom2_valid)
+    return geom2body[contact.geom1], geom2body[contact.geom2], contact.dist < 0
 
 
 @jax.jit

docs/examples/contacts.md

@@ -22,3 +22,9 @@ Querying contacts between drones and the scene, and choosing the collision geome
 ## Crash detection
 
 --8<-- example "examples/contacts/crash.py"
+
+---
+
+## Collision pairs
+
+--8<-- example "examples/contacts/collision_pairs.py"

docs/user-guide/mujoco.md

@@ -150,16 +150,17 @@ sim = Sim(n_worlds=1, n_drones=1)
 sim.reset()
 
 _mjx_data = sim.mjx_data  # captured in closure
+contact_fn = sim.build_contact_fn()
 
 
 def check_contacts(sim_data: SimData, obstacle_mocap_pos: Array) -> Array:
     # Update obstacle positions and sync inside JIT
     mjx_data = _mjx_data.replace(mocap_pos=obstacle_mocap_pos)
     _, mjx_data = sync_sim2mjx(sim_data, mjx_data, sim.mjx_model)
-    return mjx_data._impl.contact.dist < 0
+    return contact_fn(mjx_data)
 
 
-in_contact = check_contacts(sim.data, sim.mjx_data.mocap_pos)
+in_contact = check_contacts(sim.data, sim.mjx_data.mocap_pos)  # shape (n_worlds, n_drones)
 ```
 
 `_mjx_data` is fused into the closure and compiled as a constant. Only `sim_data` and the obstacle positions cross the JIT boundary at runtime — a much smaller pytree than passing the full `mjx_data`.

examples/contacts/collision_pairs.py

@@ -0,0 +1,49 @@
+"""Find out which drone collides with which body."""
+
+import jax.numpy as jnp
+import mujoco
+import numpy as np
+
+from crazyflow.sim import Sim
+
+OBSTACLE_XML = """
+<mujoco>
+  <worldbody>
+    <body name="obstacle">
+      <geom type="box" size="0.1 0.1 0.1" rgba="0.8 0.2 0.2 1"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def main():
+    """Place drones into obstacles, the floor and each other, and print each collision."""
+    sim = Sim(n_drones=5)
+    obstacle_spec = mujoco.MjSpec.from_string(OBSTACLE_XML)
+    frame = sim.spec.worldbody.add_frame()
+    for i, pos in enumerate([[1.0, 0.0, 0.5], [2.0, 0.0, 0.5]]):
+        frame.attach_body(obstacle_spec.body("obstacle"), "", f":{i}").pos = pos
+    sim.build_mjx()
+    sim.reset()
+
+    pos = jnp.array(
+        [[0.85, 0.0, 0.5], [2.15, 0.0, 0.5], [1.5, 0.5, 0.5], [1.5, 0.5, 0.5], [1.5, -0.5, 0.0]]
+    )
+    sim.data = sim.data.replace(
+        states=sim.data.states.replace(pos=pos[None]), core=sim.data.core.replace(mjx_synced=False)
+    )
+
+    body1, body2, touching = sim.contact_bodies()
+    for i in np.flatnonzero(touching[0]):
+        name1 = sim.mj_model.body(int(body1[0, i])).name
+        name2 = sim.mj_model.body(int(body2[0, i])).name
+        print(f"{name1} collides with {name2}")
+
+    for _ in range(120):
+        sim.render()
+    sim.close()
+
+
+if __name__ == "__main__":
+    main()

tests/unit/test_contacts.py

@@ -0,0 +1,171 @@
+"""Unit tests for contact detection."""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import mujoco
+import numpy as np
+import pytest
+
+from crazyflow.control import Control
+from crazyflow.sim import Dynamics, Sim
+from crazyflow.sim.sim import use_box_collision
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("dynamics", Dynamics)
+def test_contacts(dynamics: Dynamics):
+    sim = Sim(n_worlds=2, n_drones=2, dynamics=dynamics, freq=500, device="cpu")
+    sim.reset()
+    sim.step(10)  # Make sure the drone is on the ground
+    contacts = sim.contacts()
+    assert contacts.shape == (2, 2), "Contacts should be flagged per world and drone"
+    assert jnp.all(contacts), "Drone should be in contact with the floor"
+    sim.close()
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("box_collision", [False, True])
+def test_contacts_between_drones(box_collision: bool):
+    """Two overlapping drones must register a contact."""
+    # Regression test: Drones are mocap bodies that are all welded to the world. MuJoCo/MJX filters
+    # out collisions between geoms welded into the same kinematic tree, so without explicit contact
+    # pairs no drone-drone contact is ever generated.
+    sim = Sim(n_drones=2, control=Control.attitude, freq=500, device="cpu")
+    if box_collision:
+        use_box_collision(sim, True)
+    sim.reset()
+    states = sim.data.states
+    # Place both drones at the same position above the floor
+    pos = states.pos.at[:, 1, :].set(states.pos[:, 0, :]).at[..., 2].set(1.0)
+    sim.data = sim.data.replace(states=states.replace(pos=pos))
+    sim.data = sim.data.replace(core=sim.data.core.replace(mjx_synced=False))
+    assert jnp.all(sim.contacts()), "Overlapping drones should be in contact"
+    sim.close()
+
+
+@pytest.mark.unit
+def test_contacts_ignore():
+    sim = Sim(n_worlds=2, n_drones=3)
+    sim.reset()
+    pos = sim.data.states.pos.at[..., 2].set(1.0)
+    pos = pos.at[0, 1].set(pos[0, 0])  # World 0: drones 0 and 1 overlap
+    pos = pos.at[1, 2, 2].set(-0.05)  # World 1: drone 2 is in the floor
+    sim.data = sim.data.replace(states=sim.data.states.replace(pos=pos))
+    sim.data = sim.data.replace(core=sim.data.core.replace(mjx_synced=False))
+    expected = jnp.array([[True, True, False], [False, False, True]])
+    assert jnp.array_equal(sim.contacts(), expected)
+
+    # Drone pairs are correctly ignored
+    sim.build_contact_fn(ignore=[("drone:1", "drone:0")])
+    expected = jnp.array([[False, False, False], [False, False, True]])
+    assert jnp.array_equal(sim.contacts(), expected), "Pairs apply in both directions"
+
+    # Single pairs are correctly ignored. Using the body ID also works
+    sim.build_contact_fn(ignore=[(sim.mj_model.body("drone:2").id, "world")])
+    expected = jnp.array([[True, True, False], [False, False, False]])
+    assert jnp.array_equal(sim.contacts(), expected), "Floor contact not ignored"
+
+    # Single bodies are ignored by all drones
+    sim.build_contact_fn(ignore=["world"])
+    expected = jnp.array([[True, True, False], [False, False, False]])
+    assert jnp.array_equal(sim.contacts(), expected), "Floor contact not ignored"
+
+    # Single drones are also ignored by all other drones
+    sim.build_contact_fn(ignore=["drone:0"])
+    expected = jnp.array([[False, False, False], [False, False, True]])
+    assert jnp.array_equal(sim.contacts(), expected), "Drone 0 not ignored"
+    sim.close()
+
+
+@pytest.mark.unit
+def test_contact_bodies():
+    sim = Sim(n_worlds=2, n_drones=3)
+    sim.reset()
+    pos = sim.data.states.pos.at[..., 2].set(1.0)
+    pos = pos.at[0, 1].set(pos[0, 0])  # World 0: drones 0 and 1 overlap
+    pos = pos.at[1, 2, 2].set(-0.05)  # World 1: drone 2 is in the floor
+    sim.data = sim.data.replace(states=sim.data.states.replace(pos=pos))
+    sim.data = sim.data.replace(core=sim.data.core.replace(mjx_synced=False))
+    body1, body2, touching = sim.contact_bodies()
+    assert body1.shape == body2.shape == touching.shape
+    assert touching.sum(axis=-1).tolist() == [1, 1], "Each collision should appear once"
+    drone_0, drone_1, drone_2 = (sim.mj_model.body(f"drone:{i}").id for i in range(3))
+    floor = sim.mj_model.body("world").id
+    collisions = [{int(body1[w][touching[w]][0]), int(body2[w][touching[w]][0])} for w in range(2)]
+    assert collisions == [{drone_0, drone_1}, {floor, drone_2}]
+    sim.close()
+
+
+def _attach_obstacle(sim: Sim, pos: list[float], mocap: bool):
+    """Attach the minimal obstacle to the sim at ``pos``, mirroring downstream track loading."""
+    OBSTACLE_XML = """
+    <mujoco>
+    <worldbody>
+        <body name="obstacle">
+        <geom name="obstacle" type="box" size="0.1 0.1 0.1"/>
+        </body>
+    </worldbody>
+    </mujoco>
+    """
+    obstacle_spec = mujoco.MjSpec.from_string(OBSTACLE_XML)
+    frame = sim.spec.worldbody.add_frame()
+    obstacle = frame.attach_body(obstacle_spec.body("obstacle"), "", ":0")
+    obstacle.pos = pos
+    obstacle.mocap = mocap  # Mocap bodies (like the gates) are still welded to the world
+    sim.build_mjx()
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("mocap", [True, False])
+def test_contacts_with_obstacle(mocap: bool):
+    """A drone placed inside an obstacle must register a contact.
+
+    Covers both a mocap obstacle (attached like the gates) and a static one. Both are welded to the
+    world just like the drones, so the collision is only detected if explicit contact pairs are
+    generated.
+    """
+    obstacle_pos = [0.0, 0.0, 1.0]
+    sim = Sim(control=Control.attitude, freq=500, device="cpu")
+    _attach_obstacle(sim, obstacle_pos, mocap=mocap)
+    sim.reset()
+    states = sim.data.states
+    pos = states.pos.at[..., :].set(jnp.array(obstacle_pos))
+    sim.data = sim.data.replace(states=states.replace(pos=pos))
+    sim.data = sim.data.replace(core=sim.data.core.replace(mjx_synced=False))
+    assert jnp.all(sim.contacts()), "Drone should collide with the obstacle"
+    sim.build_contact_fn(ignore=["obstacle:0"])
+    assert not jnp.any(sim.contacts()), "Obstacle contact should be ignored"
+    sim.close()
+
+
+@pytest.mark.unit
+def test_max_geom_pairs_caps_contact_buffer():
+    """Small swarms check all geom pairs, large swarms cap the buffer to stay linear."""
+    sim = Sim(n_drones=32)
+    assert sim.mjx_data._impl.contact.dist.shape[-1] == 32 * 33 // 2
+    sim.close()
+    n_drones = 64
+    sim = Sim(n_drones=n_drones)
+    assert sim.mjx_data._impl.contact.dist.shape[-1] == n_drones + 2 * n_drones  # Drones + floor
+    sim.max_geom_pairs = -1  # applies on the next build
+    sim.build_mjx()
+    assert sim.mjx_data._impl.contact.dist.shape[-1] == n_drones * (n_drones + 1) // 2
+    sim.close()
+
+
+@pytest.mark.unit
+def test_capped_contacts_identify_colliding_drones():
+    """Test that capping the contact buffer still correctly identifies drone collisions by name."""
+    n_drones = 16
+    sim = Sim(n_drones=n_drones)
+    sim.max_geom_pairs = 4
+    sim.build_mjx()
+    sim.reset()
+    pos = np.stack([[i * 2.0, 0.0, 1.0] for i in range(n_drones)])[None]
+    pos[0, 4] = pos[0, 3]  # overlap drones 3 and 4, leave the rest far apart
+    sim.data = sim.data.replace(states=sim.data.states.replace(pos=jnp.array(pos)))
+    sim.step()
+    expected = (jnp.arange(n_drones) == 3) | (jnp.arange(n_drones) == 4)
+    assert jnp.array_equal(sim.contacts()[0], expected), "Only drones 3 and 4 should collide"
+    sim.close()