From c2d7847cfc536b3a3670feb7f14a66759c04232b Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Fri, 23 Oct 2026 15:35:02 +0200 Subject: [PATCH 1/3] Rework contact APIs. Closes #124 --- .gitignore | 1 + crazyflow/sim/sim.py | 77 +++++++++++++++++++++++++++---------- docs/user-guide/mujoco.md | 5 ++- tests/unit/test_sharding.py | 2 +- tests/unit/test_sim.py | 51 +++++++++++++++++++----- 5 files changed, 104 insertions(+), 32 deletions(-) diff --git a/.gitignore b/.gitignore index 226ae7989..19c713457 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,4 @@ site/ # Temporary files .DS_Store Thumbs.db +.claude diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index ecd8c1c95..425018ca0 100644 --- a/crazyflow/sim/sim.py +++ b/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,49 @@ 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 @staticmethod def _reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData: @@ -697,15 +733,16 @@ def increment_steps(data: SimData) -> SimData: return data.replace(core=data.core.replace(steps=data.core.steps + 1)) -@jax.jit -def contacts(geom_start: int, geom_count: int, data: Data) -> Array: - """Filter contacts from MuJoCo data.""" +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 - 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) + 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 diff --git a/docs/user-guide/mujoco.md b/docs/user-guide/mujoco.md index b5b9e45a6..83aa554cb 100644 --- a/docs/user-guide/mujoco.md +++ b/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`. diff --git a/tests/unit/test_sharding.py b/tests/unit/test_sharding.py index eda3ad904..fdb7a34bd 100644 --- a/tests/unit/test_sharding.py +++ b/tests/unit/test_sharding.py @@ -114,5 +114,5 @@ def test_sharded_contacts(platform: str): sim.data = sim.data.replace( states=sim.data.states.replace(pos=pos), core=sim.data.core.replace(mjx_synced=False) ) - assert jnp.array_equal(sim.contacts("drone:0").any(axis=-1), grounded), "wrong worlds collide" + assert jnp.array_equal(sim.contacts()[:, 0], grounded), "wrong worlds collide" sim.close() diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index 5fe1dc9ed..e0b1a570c 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -503,12 +503,11 @@ def test_thrust_clip(dynamics: Dynamics, integrator: Integrator): @pytest.mark.unit @pytest.mark.parametrize("dynamics", Dynamics) def test_contacts(dynamics: Dynamics): - sim = Sim(dynamics=dynamics, control=Control.attitude, freq=500, device="cpu") + 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() - # The contact buffer must contain the drone-floor contact - assert contacts.shape[-1] > 0, "Contact buffer should not be empty" + 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() @@ -531,8 +530,41 @@ def test_contacts_between_drones(box_collision: bool): 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.any(sim.contacts("drone:0")), "Overlapping drones should be in contact" - assert jnp.any(sim.contacts("drone:1")), "Overlapping drones should be in contact" + 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() @@ -572,7 +604,9 @@ def test_contacts_with_obstacle(mocap: bool): 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.any(sim.contacts("drone:0")), "Drone should collide with the obstacle" + 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() @@ -798,7 +832,6 @@ def test_capped_contacts_identify_colliding_drones(): 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() - assert jnp.any(sim.contacts("drone:3")), "Overlapping drones should be in contact" - assert jnp.any(sim.contacts("drone:4")), "Overlapping drones should be in contact" - assert not jnp.any(sim.contacts("drone:5")), "Distant drones should not be in contact" + 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() From 886b4fc720ca45868cc8cf39f2817ad000faa9c3 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Fri, 23 Oct 2026 18:09:02 +0200 Subject: [PATCH 2/3] Add contact bodies API --- crazyflow/sim/sim.py | 20 ++++ docs/examples/contacts.md | 6 + examples/contacts/collision_pairs.py | 49 ++++++++ tests/unit/test_contacts.py | 171 +++++++++++++++++++++++++++ tests/unit/test_sim.py | 144 +--------------------- 5 files changed, 247 insertions(+), 143 deletions(-) create mode 100644 examples/contacts/collision_pairs.py create mode 100644 tests/unit/test_contacts.py diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 425018ca0..5c3fa34ee 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -638,6 +638,19 @@ def build_contact_fn( 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: raise NotInitializedError("_reset call before building the simulation pipeline.") @@ -745,6 +758,13 @@ def contacts(data: Data, *, geom2drone: Array, geom2body: Array, ignored: Array) return jax.vmap(any_per_drone)(touching, drone_ids) +@jax.jit +def contact_bodies(data: Data, geom2body: Array) -> tuple[Array, Array, Array]: + """Map both geoms of each contact to their bodies.""" + contact = data._impl.contact + return geom2body[contact.geom1], geom2body[contact.geom2], contact.dist < 0 + + @jax.jit def sync_sim2mjx(data: SimData, mjx_data: Data, mjx_model: Model) -> tuple[SimData, Data]: """Synchronize the simulation data with the MuJoCo mjx_model.""" diff --git a/docs/examples/contacts.md b/docs/examples/contacts.md index 7780b152e..b7e59f8c1 100644 --- a/docs/examples/contacts.md +++ b/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" diff --git a/examples/contacts/collision_pairs.py b/examples/contacts/collision_pairs.py new file mode 100644 index 000000000..1bc124f34 --- /dev/null +++ b/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 = """ + + + + + + + +""" + + +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() diff --git a/tests/unit/test_contacts.py b/tests/unit/test_contacts.py new file mode 100644 index 000000000..2b8711ceb --- /dev/null +++ b/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 = """ + + + + + + + + """ + 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() diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index e0b1a570c..481fb5106 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -19,7 +19,7 @@ from crazyflow.sim import Dynamics, Sim from crazyflow.sim.data import ControlData, SimData from crazyflow.sim.integration import Integrator -from crazyflow.sim.sim import rotor_vel_limits, sync_sim2mjx, use_box_collision +from crazyflow.sim.sim import rotor_vel_limits, sync_sim2mjx from crazyflow.sim.visualize import change_material if TYPE_CHECKING: @@ -500,116 +500,6 @@ def test_thrust_clip(dynamics: Dynamics, integrator: Integrator): sim.close() -@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. - - Drones are mocap bodies that are all welded to the world. MuJoCo/MJX filter 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() - - -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 = """ - - - - - - - - """ - 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_spec_copy_during_attach(): """sim.spec must copy bodies on attach so the same source body can be attached repeatedly.""" @@ -803,35 +693,3 @@ def test_full_reset_restores_shared_arrays(): assert jnp.array_equal(sim.data.params.gravity_vec, default_gravity) # The random key is the only thing that does not reset assert jnp.array_equal(jax.random.key_data(sim.data.core.rng_key), jax.random.key_data(rng_key)) - - -@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() From f3e9ad0ca646a224cd006f00fbb1519fcf5f813f Mon Sep 17 00:00:00 2001 From: Martin Schuck <57562633+amacati@users.noreply.github.com> Date: Sun, 25 Oct 2026 17:02:46 +0200 Subject: [PATCH 3/3] Rename variable for clarity --- crazyflow/sim/sim.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 5c3fa34ee..10bf93e14 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -614,8 +614,8 @@ def build_contact_fn( 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) + no_drone_id = self.n_drones # One past the last drone, gets ignored later + body2drone = np.full(mj_model.nbody, no_drone_id) 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)