Optimize large swarms (#125)

Author: amacati

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

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()