diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index e9d9f91..9bd722d 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -130,6 +130,7 @@ def __init__( self.n_worlds = n_worlds self.n_drones = n_drones self.freq = freq + self.max_geom_pairs = -1 if n_drones < 64 else 2 * n_drones self.max_visual_geom = 1000 # Initialize MuJoCo world and data @@ -287,6 +288,7 @@ def build_mjx_spec(self) -> mujoco.MjSpec: assert self._xml_path.exists(), f"Model file {self._xml_path} does not exist" spec = mujoco.MjSpec.from_file(str(self._xml_path)) spec.option.timestep = 1 / self.freq + spec.add_numeric(name="max_geom_pairs", data=[self.max_geom_pairs]) spec.copy_during_attach = True drone_spec = mujoco.MjSpec.from_file(str(self.drone_path)) frame = spec.worldbody.add_frame(name="world") @@ -502,6 +504,7 @@ def build_mjx(self): if self.viewer is not None: self.viewer.close() 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) def init_data( @@ -573,11 +576,22 @@ def controllable(self) -> Array: def contacts(self, body: str | None = None) -> Array: """Get contact information from the simulation. + Note: + ``sim.max_geom_pairs`` limits the maximum detectable collision contacts per collision + group. This is relevant for swarms, where the full pairwise collision buffer grows + quadratically. By default, we allocate 2*n_drones contact pairs if the swarm size + exceeds 64. That gives us enough capacity to detect all drone-drone contacts. However, + if the swarm collapses e.g. into a single position, this will no longer be correct. If + you need to truly detect all contacts, set ``sim.max_geom_pairs`` to -1 and rebuild the + simulation. + Args: body: Optional body name to filter contacts for. If None, returns flags for all bodies. Returns: - An boolean array of shape (n_worlds,) that is True if any contact is present. + 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``. """ if body is None: return self.mjx_data._impl.contact.dist < 0 diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index ea6b69a..027a43f 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -767,3 +767,36 @@ 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()