diff --git a/crazyflow/sim/sharding.py b/crazyflow/sim/sharding.py index f6d32200..fdd8bf1c 100644 --- a/crazyflow/sim/sharding.py +++ b/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) diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 35103c61..e9d9f91a 100644 --- a/crazyflow/sim/sim.py +++ b/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: diff --git a/crazyflow/utils.py b/crazyflow/utils.py index e0f97f7a..50fef31c 100644 --- a/crazyflow/utils.py +++ b/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 diff --git a/tests/unit/test_sharding.py b/tests/unit/test_sharding.py index 04e7bf6b..eda3ad90 100644 --- a/tests/unit/test_sharding.py +++ b/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()