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..10bf93e14 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,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_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)
+ 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
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/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/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_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..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,82 +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(dynamics=dynamics, control=Control.attitude, 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 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.any(sim.contacts("drone:0")), "Overlapping drones should be in contact"
- assert jnp.any(sim.contacts("drone:1")), "Overlapping drones should be in contact"
- 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.any(sim.contacts("drone:0")), "Drone should collide with the obstacle"
- 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."""
@@ -769,36 +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()
- 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"
- sim.close()