Fix sharding for swarms, collisions, and MJX data

crazyflow/sim/sharding.py

@@ -7,9 +7,11 @@
 
 from __future__ import annotations
 
+from functools import partial
 from typing import TYPE_CHECKING
 
 import jax
+import jax.numpy as jnp
 from jax.sharding import AxisType, NamedSharding, PartitionSpec
 
 from crazyflow.utils import world_mask
@@ -19,6 +21,7 @@
 
     from jax import Array, Device
     from jax.sharding import Mesh
+    from mujoco.mjx import Data
 
     from crazyflow.sim.data import SimData
 
@@ -92,3 +95,19 @@ def build_sharded_data(
     if isinstance(rng_key, int):  # Tracing turns a seed into an array that is not a key
         rng_key = jax.random.key(rng_key)
     return jax.jit(create, out_shardings=placement(jax.eval_shape(create, rng_key), mesh))(rng_key)
+
+
+def build_sharded_mjx_data(data: Data, n_worlds: int, mesh: Mesh) -> Data:
+    """Distribute the per-world MJX data over a mesh without materialising on a single device.
+
+    Args:
+        data: MJX data of a single world.
+        n_worlds: Number of worlds to copy the data into.
+        mesh: Mesh to distribute the worlds over.
+
+    Returns:
+        The placed MJX data.
+    """
+    data = jax.device_put(data, NamedSharding(mesh, PartitionSpec()))  # Trace on the mesh devices
+    broadcast = partial(jax.tree.map, lambda x: jnp.broadcast_to(x, (n_worlds, *x.shape)))
+    return jax.jit(broadcast, out_shardings=NamedSharding(mesh, PartitionSpec(WORLD_AXIS)))(data)

crazyflow/sim/sim.py

@@ -14,6 +14,7 @@
 import numpy as np
 from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 from jax import Array, Device
+from jax.sharding import NamedSharding, PartitionSpec
 
 import crazyflow.sim.functional as F
 from crazyflow.control import Control
@@ -35,7 +36,7 @@
 from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
 from crazyflow.sim.integration import Integrator, euler, rk4, symplectic_euler
 from crazyflow.sim.pipeline import append_fn
-from crazyflow.sim.sharding import build_sharded_data, placement
+from crazyflow.sim.sharding import WORLD_AXIS, build_sharded_data, build_sharded_mjx_data, placement
 from crazyflow.utils import grid_2d, pytree_replace, world_mask
 
 if TYPE_CHECKING:
@@ -348,7 +349,12 @@ def build_mjx_model(self, spec: mujoco.MjSpec) -> tuple[Any, Any, Model, Data]:
         mj_data = mujoco.MjData(mj_model)
         mjx_model = mjx.put_model(mj_model, device=self.device)
         mjx_data = mjx.put_data(mj_model, mj_data, device=self.device)
-        mjx_data = jax.vmap(lambda _: mjx_data)(jnp.arange(self.n_worlds))
+        if self.mesh is None:
+            mjx_data = jax.vmap(lambda _: mjx_data)(jnp.arange(self.n_worlds))
+        else:
+            # mjx_model has no world axis, so we replicate it to keep it compatible with the mesh
+            mjx_model = jax.device_put(mjx_model, NamedSharding(self.mesh, PartitionSpec()))
+            mjx_data = build_sharded_mjx_data(mjx_data, self.n_worlds, self.mesh)
         return mj_model, mj_data, mjx_model, mjx_data
 
     def _unweld_drones(self, mj_model: mujoco.MjModel):
@@ -458,7 +464,7 @@ def build_data(self) -> SimData:
         return self.data
 
     def shard(self, mesh: Mesh) -> SimData:
-        """Distribute the data and default data over a mesh along the world axis.
+        """Distribute the data, default data and MJX data over a mesh along the world axis.
 
         Args:
             mesh: Mesh to distribute the worlds over, as built by
@@ -470,6 +476,11 @@ def shard(self, mesh: Mesh) -> SimData:
         self.mesh = mesh
         self.data = jax.device_put(self.data, placement(self.data, mesh))
         self.default_data = jax.device_put(self.default_data, placement(self.default_data, mesh))
+        # We also have to move the mjx_model and mjx_data to the mesh. mjx_model is replicated, data
+        # is sharded along its world axis
+        self.mjx_model = jax.device_put(self.mjx_model, NamedSharding(mesh, PartitionSpec()))
+        world = NamedSharding(mesh, PartitionSpec(WORLD_AXIS))
+        self.mjx_data = jax.device_put(self.mjx_data, world)
         return self.data
 
     def build_default_data(self) -> SimData:

crazyflow/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import math
 import os
 from collections.abc import Mapping
 from dataclasses import fields, is_dataclass
@@ -22,8 +23,9 @@
 
 def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array:
     """Generate a 2D grid of points."""
+    assert n > 0, "Number of points must be positive"
     center = jnp.zeros(2) if center is None else center
-    N = int(jnp.ceil(jnp.sqrt(n)))
+    N = math.isqrt(n - 1) + 1
     points = jnp.linspace(-0.5 * spacing * (N - 1), 0.5 * spacing * (N - 1), N)
     x, y = jnp.meshgrid(points, points)
     grid = jnp.stack((x.flatten(), y.flatten()), axis=-1) + center

tests/unit/test_sharding.py

@@ -9,7 +9,7 @@
 import numpy as np
 import pytest
 from conftest import available_backends
-from jax.sharding import PartitionSpec
+from jax.sharding import NamedSharding, PartitionSpec
 
 from crazyflow.control import Control
 from crazyflow.sim import Sim
@@ -34,14 +34,27 @@ def assert_sharding(x: Any, sharding: Any):
 
 @pytest.mark.unit
 @pytest.mark.parametrize("platform", multi_device)
-def test_placement(platform: str):
+@pytest.mark.parametrize("at_construction", [True, False])
+@pytest.mark.parametrize("n_drones", [1, 2])
+def test_placement(platform: str, at_construction: bool, n_drones: int):
     devices = jax.devices(platform)
-    sim = Sim(n_worlds=2 * len(devices), device=platform)
-    sim.shard(world_mesh(devices))
+    mesh = world_mesh(devices)
+    n_worlds = 2 * len(devices)
+    if at_construction:
+        sim = Sim(n_worlds=n_worlds, n_drones=n_drones, device=platform, mesh=mesh)
+    else:
+        sim = Sim(n_worlds=n_worlds, n_drones=n_drones, device=platform)
+        sim.shard(mesh)
     assert sim.data.states.pos.sharding.spec == PartitionSpec("worlds")
     assert sim.data.params.mass.sharding.spec == PartitionSpec()
     assert sim.data.params.gravity_vec.sharding.spec == PartitionSpec()
     assert sim.data.params.rotor_dyn_coef.sharding.spec == PartitionSpec()
+    expected = placement(sim.data, mesh)
+    jax.tree.map(assert_sharding, sim.data, expected)  # Sanity-check the rest
+    jax.tree.map(assert_sharding, sim.default_data, expected)
+    # MJX data is also per-world, so should be sharded as well
+    world = NamedSharding(mesh, PartitionSpec("worlds"))
+    jax.tree.map(lambda leaf: assert_sharding(leaf, world), sim.mjx_data)
 
 
 @pytest.mark.unit
@@ -82,8 +95,24 @@ def test_sharded_step_values(platform: str):
     devices = jax.devices(platform)
     sim = Sim(n_worlds=2 * len(devices), device=platform)
     sim.step(10)
-    pos = np.asarray(sim.data.states.pos)
+    pos = np.asarray(sim.data.states.pos)  # Copy to np for comparison across shardings
     sim.reset()
     sim.shard(world_mesh(devices))
     sim.step(10)
     assert np.allclose(np.asarray(sim.data.states.pos), pos, atol=1e-6)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", multi_device)
+def test_sharded_contacts(platform: str):
+    # Put drones in collision every other world, check that contacts can be computed and are correct
+    devices = jax.devices(platform)
+    n_worlds = 2 * len(devices)
+    grounded = jnp.arange(n_worlds) % 2 == 0
+    sim = Sim(n_worlds=n_worlds, device=platform, mesh=world_mesh(devices))
+    pos = sim.data.states.pos.at[:, 0, 2].set(jnp.where(grounded, -0.05, 1.0))
+    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"
+    sim.close()