From 04a1204458987ab42da7ec3eec9f9efff0114b4d Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Tue, 18 Aug 2026 13:46:30 +0200 Subject: [PATCH 1/2] Add sharding support and fix the world axis --- crazyflow/control/mellinger/control.py | 26 +++--- .../dynamics/first_principles/dynamics.py | 26 +++--- crazyflow/dynamics/so_rpy/dynamics.py | 22 ++--- crazyflow/dynamics/so_rpy_rotor/dynamics.py | 24 ++--- .../dynamics/so_rpy_rotor_drag/dynamics.py | 26 +++--- crazyflow/sim/data.py | 46 +++++----- crazyflow/sim/functional.py | 11 ++- crazyflow/sim/sharding.py | 73 +++++++++++++++ crazyflow/sim/sim.py | 28 +++++- crazyflow/utils.py | 71 +++++++++++---- docs/api/index.md | 1 + docs/examples/index.md | 11 +++ docs/gen_ref_pages.py | 1 + docs/user-guide/index.md | 1 + docs/user-guide/oo-api.md | 2 + docs/user-guide/sharding.md | 55 ++++++++++++ docs/user-guide/world-axis.md | 71 +++++++++++++++ examples/jax/sharding.py | 39 ++++++++ properdocs.yml | 2 + tests/conftest.py | 2 + tests/unit/test_plugins.py | 56 +++++++----- tests/unit/test_sharding.py | 89 +++++++++++++++++++ tests/unit/test_sim.py | 25 +++++- tests/unit/test_utils.py | 87 +++++++++++++++++- 24 files changed, 666 insertions(+), 129 deletions(-) create mode 100644 crazyflow/sim/sharding.py create mode 100644 docs/user-guide/sharding.md create mode 100644 docs/user-guide/world-axis.md create mode 100644 examples/jax/sharding.py create mode 100644 tests/unit/test_sharding.py diff --git a/crazyflow/control/mellinger/control.py b/crazyflow/control/mellinger/control.py index b122970f..85ef2608 100644 --- a/crazyflow/control/mellinger/control.py +++ b/crazyflow/control/mellinger/control.py @@ -21,7 +21,7 @@ from crazyflow.control.core import controllable, load_params from crazyflow.control.transform import force2pwm, motor_force2rotor_vel, pwm2force -from crazyflow.utils import leaf_replace +from crazyflow.utils import WORLD_INDEXED_KEY, leaf_replace if TYPE_CHECKING: from jax import Device @@ -311,19 +311,19 @@ def force_torque2rotor_vel( @dataclass class MellingerStateData: - cmd: Array # (N, M, 13) + cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 13) """Full state control command for the drone. A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]. We currently do not use the acceleration and angle rate components. This is subject to change. """ - staged_cmd: Array # (N, M, 13) + staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 13) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array # (N, 1) + steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) """Last simulation steps that the state control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the state control command.""" - pos_err_i: Array # (N, M, 3) + pos_err_i: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Integral errors of the state control command.""" # Parameters for the state controller params: dict[str, Array] @@ -344,20 +344,20 @@ def create( @dataclass class MellingerAttitudeData: - cmd: Array # (N, M, 4) + cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Full attitude control command for the drone. A command consists of [roll, pitch, yaw, collective thrust]. """ - staged_cmd: Array # (N, M, 4) + staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array # (N, 1) + steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) """Last simulation steps that the attitude control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the attitude control command.""" - r_int_error: Array # (N, M, 3) + r_int_error: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Integral errors of the attitude control command.""" - last_ang_vel: Array # (N, M, 3) + last_ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Last angular velocity of the drone.""" # Parameters for the attitude controller params: dict[str, Array] @@ -384,14 +384,14 @@ def create( @dataclass class MellingerForceTorqueData: - cmd: Array # (N, M, 4) + cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Force-torque command for the drone. A command consists of [fz, tx, ty, tz]. """ - staged_cmd: Array # (N, M, 4) + staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array # (N, 1) + steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) """Last simulation steps that the force and torque control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the force and torque control command.""" diff --git a/crazyflow/dynamics/first_principles/dynamics.py b/crazyflow/dynamics/first_principles/dynamics.py index 70ee0fdc..05f84866 100644 --- a/crazyflow/dynamics/first_principles/dynamics.py +++ b/crazyflow/dynamics/first_principles/dynamics.py @@ -22,13 +22,13 @@ import jax.numpy as jnp from array_api_compat import array_namespace from array_api_compat import device as xp_device -from flax.struct import dataclass +from flax.struct import dataclass, field from scipy.spatial.transform import Rotation as R import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import to_xp +from crazyflow.utils import WORLD_INDEXED_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -298,27 +298,27 @@ def symbolic_dynamics( @dataclass class Params: - mass: Array # (N, M, 1) + mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) """Mass of the drone.""" - L: Array # (N, M, 1) + L: Array # () """Arm length of the drone.""" - prop_inertia: Array # (N, M, 1) + prop_inertia: Array # () """Inertia of the propeller.""" - gravity_vec: Array # (N, M, 3) + gravity_vec: Array # (3,) """Gravity vector of the drone.""" - J: Array # (N, M, 3, 3) + J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array # (N, M, 3, 3) + J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - rpm2thrust: Array # (N, M, 1) + rpm2thrust: Array # (3,) """Force constant of the drone.""" - rpm2torque: Array # (N, M, 1) + rpm2torque: Array # (3,) """Torque constant of the drone.""" - mixing_matrix: Array # (N, M, 3, 4) + mixing_matrix: Array # (3, 4) """Mixing matrix of the drone.""" - drag_matrix: Array # (N, M, 3, 3) + drag_matrix: Array # (3, 3) """Drag matrix of the drone.""" - rotor_dyn_coef: Array # (N, M, 4) + rotor_dyn_coef: Array # (4,) """Rotor speed dynamics time constant of the drone.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy/dynamics.py b/crazyflow/dynamics/so_rpy/dynamics.py index a5bffb15..1a828620 100644 --- a/crazyflow/dynamics/so_rpy/dynamics.py +++ b/crazyflow/dynamics/so_rpy/dynamics.py @@ -21,13 +21,13 @@ import jax.numpy as jnp from array_api_compat import array_namespace from array_api_compat import device as xp_device -from flax.struct import dataclass +from flax.struct import dataclass, field from scipy.spatial.transform import Rotation as R import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import to_xp +from crazyflow.utils import WORLD_INDEXED_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -316,31 +316,31 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array # (N, M, 1) + mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (N, M, 3) + gravity_vec: Array # (3,) """Gravity vector of the drone.""" - J: Array # (N, M, 3, 3) + J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array # (N, M, 3, 3) + J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - acc_coef: Array # (N, M, 1) + acc_coef: Array # () """Coefficient for the acceleration.""" - cmd_f_coef: Array # (N, M, 1) + cmd_f_coef: Array # () """Coefficient for the collective thrust.""" - rpy_coef: Array # (N, M, 1) + rpy_coef: Array # (3,) """Coefficient for the roll pitch yaw dynamics.""" - rpy_rates_coef: Array # (N, M, 1) + rpy_rates_coef: Array # (3,) """Coefficient for the roll pitch yaw rates dynamics.""" - cmd_rpy_coef: Array # (N, M, 1) + cmd_rpy_coef: Array # (3,) """Coefficient for the roll pitch yaw command dynamics.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy_rotor/dynamics.py b/crazyflow/dynamics/so_rpy_rotor/dynamics.py index eef3723f..6f0157fb 100644 --- a/crazyflow/dynamics/so_rpy_rotor/dynamics.py +++ b/crazyflow/dynamics/so_rpy_rotor/dynamics.py @@ -23,13 +23,13 @@ import jax.numpy as jnp from array_api_compat import array_namespace from array_api_compat import device as xp_device -from flax.struct import dataclass +from flax.struct import dataclass, field from scipy.spatial.transform import Rotation as R import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import to_xp +from crazyflow.utils import WORLD_INDEXED_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -375,25 +375,25 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array # (N, M, 1) + mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (N, M, 3) + gravity_vec: Array # (3,) """Gravity vector of the drone.""" - J: Array # (N, M, 3, 3) + J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array # (N, M, 3, 3) + J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - thrust_time_coef: Array # (N, M, 1) + thrust_time_coef: Array # () """Rotor coefficient of the drone.""" - acc_coef: Array # (N, M, 1) + acc_coef: Array # () """Acceleration coefficient of the drone.""" - cmd_f_coef: Array # (N, M, 1) + cmd_f_coef: Array # () """Collective thrust coefficient of the drone.""" - rpy_coef: Array # (N, M, 1) + rpy_coef: Array # (3,) """Roll pitch yaw coefficient of the drone.""" - rpy_rates_coef: Array # (N, M, 1) + rpy_rates_coef: Array # (3,) """Roll pitch yaw rates coefficient of the drone.""" - cmd_rpy_coef: Array # (N, M, 1) + cmd_rpy_coef: Array # (3,) """Roll pitch yaw command coefficient of the drone.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py index c72fdc61..0127c46d 100644 --- a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py +++ b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py @@ -25,13 +25,13 @@ import jax.numpy as jnp from array_api_compat import array_namespace from array_api_compat import device as xp_device -from flax.struct import dataclass +from flax.struct import dataclass, field from scipy.spatial.transform import Rotation as R import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import to_xp +from crazyflow.utils import WORLD_INDEXED_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -409,27 +409,27 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array # (N, M, 1) + mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (N, M, 3) + gravity_vec: Array # (3,) """Gravity vector of the drone.""" - J: Array # (N, M, 3, 3) + J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array # (N, M, 3, 3) + J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - thrust_time_coef: Array # (N, M, 1) + thrust_time_coef: Array # () """Rotor coefficient of the drone.""" - acc_coef: Array # (N, M, 1) + acc_coef: Array # () """Acceleration coefficient of the drone.""" - cmd_f_coef: Array # (N, M, 1) + cmd_f_coef: Array # () """Collective thrust coefficient of the drone.""" - rpy_coef: Array # (N, M, 1) + rpy_coef: Array # (3,) """Roll pitch yaw coefficient of the drone.""" - rpy_rates_coef: Array # (N, M, 1) + rpy_rates_coef: Array # (3,) """Roll pitch yaw rates coefficient of the drone.""" - cmd_rpy_coef: Array # (N, M, 1) + cmd_rpy_coef: Array # (3,) """Roll pitch yaw command coefficient of the drone.""" - drag_matrix: Array # (N, M, 3, 3) + drag_matrix: Array # (3, 3) """Linear drag coefficient matrix of the drone.""" @staticmethod diff --git a/crazyflow/sim/data.py b/crazyflow/sim/data.py index 6872853d..edf8ce57 100644 --- a/crazyflow/sim/data.py +++ b/crazyflow/sim/data.py @@ -1,6 +1,7 @@ from __future__ import annotations import typing +from typing import Any import jax import jax.numpy as jnp @@ -18,23 +19,24 @@ from crazyflow.dynamics.so_rpy import Params as SoRpyParams from crazyflow.dynamics.so_rpy_rotor import Params as SoRpyRotorParams from crazyflow.dynamics.so_rpy_rotor_drag import Params as SoRpyRotorDragParams +from crazyflow.utils import WORLD_INDEXED_KEY @dataclass class SimState: - pos: Array # (N, M, 3) + pos: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Position of the drone's center of mass.""" - quat: Array # (N, M, 4) + quat: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Quaternion of the drone's orientation.""" - vel: Array # (N, M, 3) + vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Velocity of the drone's center of mass in the world frame.""" - ang_vel: Array # (N, M, 3) + ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Angular velocity of the drone in the body frame.""" - force: Array # (N, M, 3) # CoM force + force: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) CoM force """Force applied to the drone's center of mass in the world frame.""" - torque: Array # (N, M, 3) # CoM torque + torque: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) CoM torque """Torque applied to the drone's center of mass in the world frame.""" - rotor_vel: Array # (N, M, 4) # Motor forces along body frame z axis + rotor_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) in RPM """Motor forces along body frame z axis.""" @staticmethod @@ -57,15 +59,15 @@ def create(n_worlds: int, n_drones: int, device: Device) -> SimState: @dataclass class SimStateDeriv: - vel: Array # (N, M, 3) + vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Derivative of the position of the drone's center of mass.""" - ang_vel: Array # (N, M, 3) + ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Derivative of the quaternion of the drone's orientation as angular velocity.""" - acc: Array # (N, M, 3) + acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Derivative of the velocity of the drone's center of mass.""" - ang_acc: Array # (N, M, 3) + ang_acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) """Derivative of the angular velocity of the drone's center of mass.""" - rotor_acc: Array # (N, M, 4) + rotor_acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Derivative of the rotor velocity.""" @staticmethod @@ -108,7 +110,7 @@ class SimControls: """Attitude control data.""" force_torque: ControlData | None """Force and torque control data.""" - rotor_vel: Array # (N, M, 4) + rotor_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) """Desired motor speed.""" @staticmethod @@ -205,9 +207,7 @@ def create( class SimCore: freq: int = field(pytree_node=False) """Frequency of the simulation.""" - device: Device = field(pytree_node=False) - """Device of the simulation.""" - steps: Array # (N, 1) + steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) """Simulation steps taken since the last reset.""" n_worlds: int = field(pytree_node=False) """Number of worlds in the simulation.""" @@ -215,9 +215,9 @@ class SimCore: """Number of drones in the simulation.""" drone_mocap_ids: Array # (M,) """MuJoCo mocap IDs of the drone bodies.""" - rng_key: Array # (N, 1) + rng_key: Array # () """Random number generator key for the simulation.""" - mjx_synced: Array # (1,) + mjx_synced: Array # () """Whether the simulation data is synchronized with the MuJoCo mjx_data.""" @staticmethod @@ -236,7 +236,6 @@ def create( rng_key = jax.device_put(rng_key, device) return SimCore( freq=freq, - device=device, steps=steps, n_worlds=n_worlds, n_drones=n_drones, @@ -258,5 +257,10 @@ class SimData: """Drone parameters.""" core: SimCore """Core parameters of the simulation.""" - plugins: dict[str, Array] = field(default_factory=dict) - """Arbitrary data for plugins to store state in the simulation.""" + plugins: dict[str, Any] = field(default_factory=dict) + """Arbitrary data for plugins to store state in the simulation. + + Data that is batched over worlds has to declare this using the WORLD_INDEXED_KEY metadata, so + that resets can mask it and sharding can distribute it. This requires data to be stored in a + struct. + """ diff --git a/crazyflow/sim/functional.py b/crazyflow/sim/functional.py index 354af0b7..e69a769a 100644 --- a/crazyflow/sim/functional.py +++ b/crazyflow/sim/functional.py @@ -6,7 +6,6 @@ from crazyflow.control import Control from crazyflow.control.core import controllable as _controllable -from crazyflow.utils import to_device if TYPE_CHECKING: from jax import Array @@ -18,7 +17,7 @@ def state_control(data: SimData, controls: Array) -> SimData: """State control function.""" assert data.controls.mode == Control.state, f"control type {data.controls.mode} not enabled" assert controls.shape == (data.core.n_worlds, data.core.n_drones, 13), "controls shape mismatch" - controls = to_device(controls, data.core.device) + controls = jnp.asarray(controls) data = data.replace( controls=data.controls.replace(state=data.controls.state.replace(staged_cmd=controls)) ) @@ -36,7 +35,7 @@ def attitude_control(data: SimData, controls: Array) -> SimData: """ assert data.controls.mode == Control.attitude, f"control type {data.controls.mode} not enabled" assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch" - controls = to_device(controls, data.core.device) + controls = jnp.asarray(controls) data = data.replace( controls=data.controls.replace(attitude=data.controls.attitude.replace(staged_cmd=controls)) ) @@ -49,7 +48,7 @@ def force_torque_control(data: SimData, controls: Array) -> SimData: f"control type {data.controls.mode} not enabled" ) assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch" - controls = to_device(controls, data.core.device) + controls = jnp.asarray(controls) data = data.replace( controls=data.controls.replace( force_torque=data.controls.force_torque.replace(staged_cmd=controls) @@ -65,7 +64,7 @@ def rotor_vel_control(data: SimData, controls: Array) -> SimData: """ assert data.controls.mode == Control.rotor_vel, f"control type {data.controls.mode} not enabled" assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch" - controls = to_device(controls, data.core.device) + controls = jnp.asarray(controls) return data.replace(controls=data.controls.replace(rotor_vel=controls)) @@ -81,7 +80,7 @@ def controllable(data: SimData) -> Array: control_steps = controls.force_torque.steps control_freq = controls.force_torque.freq case Control.rotor_vel: - return jnp.ones((data.core.n_worlds, 1), dtype=bool, device=data.core.device) + return jnp.ones_like(data.core.steps, dtype=bool) case _: raise NotImplementedError(f"Control mode {data.controls.mode} not implemented") return _controllable(data.core.steps, data.core.freq, control_steps, control_freq) diff --git a/crazyflow/sim/sharding.py b/crazyflow/sim/sharding.py new file mode 100644 index 00000000..2312a7c7 --- /dev/null +++ b/crazyflow/sim/sharding.py @@ -0,0 +1,73 @@ +"""Distribution of crazyflow over multiple devices. + +Sharding arrays along the world axis allows efficient parallelization of crazyflow across multiple +devices. The data itself specifies which arrays carry the world axis. See +[world_mask][crazyflow.utils.world_mask]. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import jax +from jax.sharding import AxisType, NamedSharding, PartitionSpec + +from crazyflow.utils import world_mask + +if TYPE_CHECKING: + from collections.abc import Sequence + + from jax import Device + from jax.sharding import Mesh + + from crazyflow.sim.data import SimData + +WORLD_AXIS = "worlds" +"""Name of the mesh axis that the worlds are distributed over.""" + + +def world_mesh(devices: Sequence[Device]) -> Mesh: + """Create a mesh that distributes the worlds over the devices. + + The mesh uses automatic axis types. Explicit axis types cause problems in SciPy. + + Args: + devices: Devices to distribute the worlds over. + + Returns: + A one-dimensional mesh over the world axis. + """ + axes, names = (len(devices),), (WORLD_AXIS,) + return jax.make_mesh(axes, names, axis_types=(AxisType.Auto,), devices=devices) + + +def placement(data: SimData, mesh: Mesh) -> SimData: + """Build the placement that distributes the worlds of the simulation data over a mesh. + + Note: + Pass the placement to `jax.device_put`, or modify it first to place the data by hand. + + Args: + data: Simulation data to place. + mesh: Mesh to distribute the worlds over, as built by + [world_mesh][crazyflow.sim.sharding.world_mesh]. + + Returns: + A pytree of shardings matching the simulation data. + """ + world = NamedSharding(mesh, PartitionSpec(WORLD_AXIS)) + replicated = NamedSharding(mesh, PartitionSpec()) + return jax.tree.map(lambda indexed: world if indexed else replicated, world_mask(data)) + + +def shard(data: SimData, mesh: Mesh) -> SimData: + """Distribute the worlds of the simulation data over a mesh. + + Args: + data: Simulation data to place. + mesh: Mesh to distribute the worlds over. + + Returns: + The placed simulation data. + """ + return jax.device_put(data, placement(data, mesh)) diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index a4abc884..fe955a8b 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -31,9 +31,11 @@ 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.utils import grid_2d, pytree_replace +from crazyflow.sim.sharding import placement +from crazyflow.utils import grid_2d, pytree_replace, world_mask if TYPE_CHECKING: + from jax.sharding import Mesh from mujoco.mjx import Data, Model from numpy.typing import NDArray @@ -430,6 +432,20 @@ 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. + + Args: + mesh: Mesh to distribute the worlds over, as built by + [world_mesh][crazyflow.sim.sharding.world_mesh]. + + Returns: + The placed simulation data. + """ + 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)) + return self.data + def build_default_data(self) -> SimData: """Initialize the default data for the simulation. @@ -601,8 +617,14 @@ def select_integrate_fn( def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData: - """Reset the simulation data to the default data for the worlds specified by the mask.""" - return pytree_replace(data, default_data, mask) # Does not overwrite rng_key + """Reset the simulation data to the default data for the worlds specified by the mask. + + Without a mask, the full data is restored. The mask selects along the world axis, so it only + restores per-world arrays. The rng key is never restored. + """ + if mask is None: + return default_data.replace(core=default_data.core.replace(rng_key=data.core.rng_key)) + return pytree_replace(data, default_data, world_mask(data), mask) def increment_steps(data: SimData) -> SimData: diff --git a/crazyflow/utils.py b/crazyflow/utils.py index 5136444d..c4110ac3 100644 --- a/crazyflow/utils.py +++ b/crazyflow/utils.py @@ -3,6 +3,7 @@ import inspect import os from collections.abc import Mapping +from dataclasses import fields, is_dataclass from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar @@ -15,6 +16,9 @@ if TYPE_CHECKING: from types import ModuleType +WORLD_INDEXED_KEY = "world_indexed" +"""Metadata key for dataclass fields that have a world axis.""" + def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array: """Generate a 2D grid of points.""" @@ -32,21 +36,64 @@ def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array: R = TypeVar("R") -def pytree_replace(tree: T, new_tree: T, mask: Array | None = None) -> T: - """Overwrite elements of a PyTree with values from another PyTree filtered by a mask. +def world_mask(tree: T) -> T: + """Flag the arrays of a PyTree that are indexed by world. + + Structs declare fields with a world axis with the WORLD_INDEXED_KEY metadata. Arrays without + fields or the metadata tag are shared. Structs must be flax dataclasses. + + Args: + tree: PyTree of arrays and structs to classify. + + Returns: + A PyTree of booleans matching the input. + """ + return _flag(tree) + + +def _flag(node: Any) -> Any: + """Flag the leaves of a node, recursing through the declared structure of the data.""" + # Case 1: a dataclass with potential WORLD_INDEXED_KEY fields. We filter out the dynamic fields + # and check which of them declare a world axis + if _declares_fields(node): + dynamic = [f for f in fields(node) if f.metadata.get("pytree_node", True)] + tags = {f.name: f.metadata.get(WORLD_INDEXED_KEY, False) for f in dynamic} + values = {f: getattr(node, f) for f in tags} + # Structs with fields must not be tagged. Tagging is array-level only + if tagged_structs := [f for f, tag in tags.items() if tag and _declares_fields(values[f])]: + raise ValueError(f"{type(node).__name__} tags structs: {tagged_structs}") + # Tagged fields hold no structs, so all of their arrays are indexed by world + indexed = {f: jax.tree.map(lambda _: True, values[f]) for f in tags if tags[f]} + return node.replace(**indexed, **{f: _flag(values[f]) for f in tags if not tags[f]}) + # Case 2: any other node. Its arrays are shared. Nested structs are handed back to us via + # is_leaf so that we can further recurse into them + return jax.tree.map( + lambda x: _flag(x) if _declares_fields(x) else False, node, is_leaf=_declares_fields + ) + + +def _declares_fields(node: Any) -> bool: + """Check whether a node is a dataclass that jax traverses.""" + return is_dataclass(node) and not jax.tree_util.all_leaves([node]) + + +def pytree_replace(tree: T, new_tree: T, batched: T, mask: Array | None = None) -> T: + """Overwrite batched leaves of a PyTree with values from another PyTree filtered by a mask. - The mask indicates which elements of the leaf arrays to overwrite with new values, and which - ones to leave unchanged. + Args: + tree: PyTree to overwrite. + new_tree: PyTree to take the new values from. + batched: PyTree of booleans flagging the leaves that carry the batch axis. + mask: Boolean array matching the leading axis of the batched leaves. """ - def _replace(x: Array, y: Array) -> Array: - """Replace broadcastable leaves in tree.map.""" - # Do not replace leaves that are not broadcastable - if x.ndim < 1 or (mask is not None and mask.shape[0] != x.shape[0]): + def _replace(x: Array, y: Array, batched: bool) -> Array: + """Replace batched leaves in tree.map.""" + if not batched: return x return jnp.where(broadcast_mask(mask, x.shape), y, x) - return jax.tree.map(_replace, tree, new_tree) + return jax.tree.map(_replace, tree, new_tree, batched) def leaf_replace(tree: T, mask: Array | None = None, **kwargs: dict[str, Array]) -> T: @@ -71,12 +118,6 @@ def broadcast_mask(mask: Array | None, shape: tuple[int, ...]) -> Array: return mask.reshape(*mask.shape, *[1] * (len(shape) - mask.ndim)) -@partial(jax.jit, static_argnames="device") -def to_device(data: Array, device: str) -> Array: - """Turn an array into a jax array on the specified device.""" - return jnp.array(data, device=device) - - def enable_cache( cache_path: Path | None = None, min_entry_size_bytes: int = -1, diff --git a/docs/api/index.md b/docs/api/index.md index 6463c252..8c602fab 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -10,6 +10,7 @@ This section is auto-generated from the Crazyflow source code using [mkdocstring | `crazyflow.sim.pipeline` | `OrderedDict`-based pipeline helpers (`append_fn`, `insert_fn_before`, `replace_fn`, etc.) | | `crazyflow.sim.data` | `SimData`, `SimState`, `SimControls`, `SimParams`, `SimCore` pytrees | | `crazyflow.sim.functional` | Pure functional control API for use inside `jax.jit` | +| `crazyflow.sim.sharding` | Placement of the simulation data on multiple devices | | `crazyflow.dynamics` | `Dynamics` enum and dynamics implementations | | `crazyflow.sim.integration` | `Integrator` enum, Euler, RK4, and symplectic Euler | | `crazyflow.sim.sensors` | Raycast depth (`sensors.depth`) and gaussian splat RGB (`sensors.splat`) sensors | diff --git a/docs/examples/index.md b/docs/examples/index.md index 76c876cb..ab5322e2 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -51,6 +51,17 @@ Because the simulator is built entirely from JAX operations, `jax.grad` can diff --- +## Sharding across devices + +We can also distribute the simulation across devices. The host backend is asked for four logical devices. We then compare this against running on a single device. See [Sharding](../user-guide/sharding.md) for the details. + + +```{ .python notest } +--8<-- "examples/jax/sharding.py" +``` + +--- + ## Domain randomization Randomizing mass and inertia through the reset pipeline. An optional mask limits randomization to selected worlds. diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py index 52d418ac..88e40cbf 100644 --- a/docs/gen_ref_pages.py +++ b/docs/gen_ref_pages.py @@ -53,6 +53,7 @@ * [sim.sensors](crazyflow/sim/sensors/index.md) * [sim.sensors.depth](crazyflow/sim/sensors/depth.md) * [sim.sensors.splat](crazyflow/sim/sensors/splat.md) + * [sim.sharding](crazyflow/sim/sharding.md) * [sim.sim](crazyflow/sim/sim.md) * [sim.splat](crazyflow/sim/splat.md) * [sim.visualize](crazyflow/sim/visualize.md) diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index 21663c91..99fbe308 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -8,6 +8,7 @@ In-depth documentation for every part of the simulator. - [Dynamics](dynamics/index.md) — first-principles vs. fitted dynamics, when to use each - [Control Modes](control/index.md) — state, attitude, force/torque, and rotor velocity control - [Pipelines](pipelines.md) — composable step and reset pipelines, randomization, and disturbances +- [The world axis](world-axis.md) — which arrays are batched over worlds, and what resets and sharding do with them - [Visualization](visualization.md) — rendering modes, cameras, raycasting, and materials - [MuJoCo Integration](mujoco.md) — MJCF scene construction, adding objects, and sync internals - [Gymnasium Environments](gymnasium-envs.md) — vectorized environments for RL training diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index b8389bc3..c13fb36e 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -138,6 +138,8 @@ Passing more steps to a single `step(n_steps)` call is more efficient than multi `sim.reset()` reinitialises all worlds to their default state. Pass a boolean mask of shape `(n_worlds,)` to reset only selected worlds: `True` resets that world, `False` leaves it unchanged. This is useful in RL training loops where episodes end at different times. +A full reset restores everything except the rng key. A mask selects along the world axis, so it only restores per-world arrays, and leaves parameters shared by all worlds untouched. See [The world axis](world-axis.md) for more details. + ```python import numpy as np from crazyflow.sim import Sim diff --git a/docs/user-guide/sharding.md b/docs/user-guide/sharding.md new file mode 100644 index 00000000..b595b0e3 --- /dev/null +++ b/docs/user-guide/sharding.md @@ -0,0 +1,55 @@ +# Sharding + +Crazyflow can be sharded across devices for even more performance. We partition the simulation data along the world axis and replicate shared data. `sim.shard` places `sim.data` and `sim.default_data` on multiple devices, and subsequent runs will use all devices to compute the simulation step. + +!!! note + Workloads that do not make use of sharding are unaffected by replication etc. + +```python +import jax + +from crazyflow.sim import Sim +from crazyflow.sim.sharding import world_mesh + +devices = jax.devices() +sim = Sim(n_worlds=4 * len(devices)) +sim.shard(world_mesh(devices)) +sim.step() + +assert sim.data.states.pos.sharding.spec == jax.sharding.PartitionSpec("worlds") +assert sim.data.params.gravity_vec.sharding.spec == jax.sharding.PartitionSpec() +``` + +The number of worlds has to be divisible by the number of devices. Arrays declare that they are batched over worlds on their fields, see [The world axis](world-axis.md). Data shared across worlds, such as the gravity vector above, is replicated on every device. + +!!! warning + Sharding relies on JAX's automatic sharding mode, which `world_mesh` requests for you. `jax.make_mesh` defaults to explicit sharding, under which the Mellinger controller fails inside SciPy's rotation backend. + +## Functional API + +`shard` places a `SimData` without going through the `Sim` object: + +```{ .python continuation } +from crazyflow.sim.sharding import shard + +mesh = world_mesh(devices) +data, default_data = shard(sim.data, mesh), shard(sim.default_data, mesh) +``` + +`placement` returns the shardings that `shard` applies, as a pytree mirroring the simulation data. Modify it to place the data by hand: + +```{ .python continuation } +from crazyflow.sim.sharding import placement + +data = jax.device_put(sim.data, placement(sim.data, mesh)) +``` + +## Plugin data + +Bare arrays in `sim.data.plugins` are replicated, since a dict has no fields to declare the world axis. Store per-world plugin state in a struct that declares it, as shown in [The world axis](world-axis.md#your-own-state). + +## Next steps + +- [The world axis](world-axis.md) — declaring which arrays are batched over worlds +- [Functional API](functional-api.md) — composing the pure step, reset and control functions +- [Batching & domain randomization](dynamics/batching.md) — varying the drone parameters per world diff --git a/docs/user-guide/world-axis.md b/docs/user-guide/world-axis.md new file mode 100644 index 00000000..87cbce3b --- /dev/null +++ b/docs/user-guide/world-axis.md @@ -0,0 +1,71 @@ +# The world axis + +The simulation is batched over worlds, so most arrays in `sim.data` carry a leading `n_worlds` axis. Resetting selected worlds and distributing worlds across devices both need to know which arrays those are. Here we explain how that is declared and what depends on it. + +## Why it is declared, not inferred + +A world axis is an ordinary axis, so no shape tells you whether it indexes worlds. With three worlds, `params.drag_matrix` of shape `(3, 3)` and `params.gravity_vec` of shape `(3,)` both look world-batched while neither is. + +Each field therefore declares if it has a leading world axis with the `WORLD_INDEXED_KEY` metadata: + +```python +from crazyflow.sim import Sim +from crazyflow.utils import world_mask + +sim = Sim(n_worlds=3) +mask = world_mask(sim.data) + +assert mask.states.pos # (n_worlds, n_drones, 3), indexed by world +assert mask.params.mass # (n_worlds, n_drones, 1), one mass per drone +assert not mask.params.gravity_vec # (3,), shared by all worlds +assert not mask.params.drag_matrix # (3, 3), shared by all worlds +``` + +`world_mask` returns a pytree of booleans matching the data, with one flag per array. Arrays that no field declares are shared by all worlds. + +## Resets + +`sim.reset()` restores the whole simulation, both the world-indexed and the shared arrays, with the rng key as the only exception. + +`sim.reset(mask)` is the constrained form. It selects worlds along the world axis, so it can only restore arrays that have one. A shared array cannot be restored for some worlds and not others, so a masked reset leaves it untouched: + +```{ .python continuation } +import jax.numpy as jnp + +gravity = jnp.array([1.0, 2.0, 3.0]) +sim.data = sim.data.replace(params=sim.data.params.replace(gravity_vec=gravity)) +sim.reset(jnp.array([True, False, False])) + +assert jnp.array_equal(sim.data.params.gravity_vec, gravity) # Shared, so left alone +assert jnp.all(sim.data.states.pos[0] == 0.0) # World 0 was reset +``` + +To keep a change to a shared parameter across masked resets, write it into `sim.default_data` as well by calling `sim.build_default_data()`. See [Stepping and resetting](oo-api.md#stepping-and-resetting) for the reset API. + +## Sharding + +Distributing the simulation partitions the world axis of the declared arrays and replicates the rest, using the same declarations. See [Sharding](sharding.md) for the mesh and placement API. + +## Your own state + +Anything you add to `sim.data.plugins` follows the same rule. A `dict` has no fields, so bare arrays are shared by default, which means they are neither reset nor distributed. If you need per-world states, you must declare it with a struct: + +```python +import flax.struct +from jax import Array + +from crazyflow.utils import WORLD_INDEXED_KEY + + +@flax.struct.dataclass +class PluginState: + buffer: Array = flax.struct.field(metadata={WORLD_INDEXED_KEY: True}) + lookup: Array = None +``` + +If you store this under `sim.data.plugins`, `buffer` is restored for the masked worlds and partitioned when the simulation is sharded, while `lookup` is only restored by a full reset and replicated across devices. + +## Next steps + +- [Sharding](sharding.md) — distributing the worlds over several devices +- [Pipelines](pipelines.md) — customising what runs on a step and on a reset diff --git a/examples/jax/sharding.py b/examples/jax/sharding.py new file mode 100644 index 00000000..9630fe40 --- /dev/null +++ b/examples/jax/sharding.py @@ -0,0 +1,39 @@ +"""Example on how to shard the simulation across devices.""" + +import os + +os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4" + +import jax +import numpy as np + +from crazyflow.sim import Sim +from crazyflow.sim.sharding import world_mesh + + +def main(): + devices = jax.devices("cpu") + sim = Sim(n_worlds=4 * len(devices), device="cpu") # Worlds must be divisible by n_devices + sim.step(10) + single_device_pos = np.asarray(sim.data.states.pos) + + sim.reset() + sim.shard(world_mesh(devices)) + sim.step(10) + + pos, gravity = sim.data.states.pos, sim.data.params.gravity_vec + print(f"Distributing {sim.n_worlds} worlds over {len(devices)} devices") + # Arrays that carry a world axis are partitioned. Shared params are replicated on every device + print(" Placement") + print(f" {'position':<{10}} {str(pos.shape):<{10}} {pos.sharding.spec}") + print(f" {'gravity':<{10}} {str(gravity.shape):<{10}} {gravity.sharding.spec}") + print(" Shards of states.pos") + for shard in pos.addressable_shards: + print(f" {shard.device} {shard.data.shape[0]} worlds") + + assert np.allclose(np.asarray(pos), single_device_pos, atol=1e-6) + sim.close() + + +if __name__ == "__main__": + main() diff --git a/properdocs.yml b/properdocs.yml index 04f557a6..3463a604 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -70,6 +70,8 @@ nav: - Batching: user-guide/control/batching.md - JIT compilation: user-guide/control/jit.md - Pipelines: user-guide/pipelines.md + - The world axis: user-guide/world-axis.md + - Sharding: user-guide/sharding.md - Visualization: user-guide/visualization.md - Gaussian Splat Rendering: user-guide/splats.md - MuJoCo Integration: user-guide/mujoco.md diff --git a/tests/conftest.py b/tests/conftest.py index e0611592..7c88d620 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,8 @@ os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" os.environ["SCIPY_ARRAY_API"] = "1" +# We need multiple devices for sharding tests +os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2" import jax import pytest diff --git a/tests/unit/test_plugins.py b/tests/unit/test_plugins.py index f31950aa..d728aa67 100644 --- a/tests/unit/test_plugins.py +++ b/tests/unit/test_plugins.py @@ -1,34 +1,53 @@ """Unit tests for the simulation plugin system. Plugins are callables of the form ``fn(data: SimData) -> SimData`` that are inserted into -``sim.step_pipeline``. They can store arbitrary state in ``sim.data.plugins`` (a dict of JAX -arrays). Tests use a simple per-world step counter as a canonical plugin. +``sim.step_pipeline``. They can store arbitrary state in ``sim.data.plugins``. State that is batched +over worlds has to declare its world axis so that resets can mask it. Tests use a simple per-world +step counter as a canonical plugin. """ from __future__ import annotations from typing import TYPE_CHECKING +import flax import jax.numpy as jnp import pytest +from flax.struct import field from crazyflow.sim import Sim from crazyflow.sim.pipeline import append_fn +from crazyflow.utils import WORLD_INDEXED_KEY if TYPE_CHECKING: + from jax import Array + from crazyflow.sim.data import SimData +@flax.struct.dataclass +class Counter: + """Per-world plugin state.""" + + count: Array = field(metadata={WORLD_INDEXED_KEY: True}) + + @staticmethod + def create(n_worlds: int) -> Counter: + """Create a zeroed counter for all worlds.""" + return Counter(jnp.zeros((n_worlds, 1), jnp.int32)) + + def counter_plugin(data: SimData) -> SimData: """Increment plugins["counter"] by 1 on every simulation step.""" - return data.replace(plugins=data.plugins | {"counter": data.plugins["counter"] + 1}) + counter = data.plugins["counter"] + return data.replace(plugins=data.plugins | {"counter": Counter(counter.count + 1)}) def accumulater_plugin(data: SimData) -> SimData: """Accumulate plugins["counter"] on every simulation step.""" + accumulated, counter = data.plugins["accumulated"], data.plugins["counter"] return data.replace( - plugins=data.plugins - | {"accumulated": data.plugins["accumulated"] + data.plugins["counter"]} + plugins=data.plugins | {"accumulated": Counter(accumulated.count + counter.count)} ) @@ -56,12 +75,12 @@ def test_builds(): def test_plugin_data_changes(n_worlds: int): """Plugin data changes across simulation steps.""" sim = Sim(n_worlds=n_worlds) - sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)}) + sim.data = sim.data.replace(plugins={"counter": Counter.create(n_worlds)}) append_fn(sim.step_pipeline, counter_plugin) sim.build_step_fn() n_steps = 7 sim.step(n_steps) - assert jnp.all(sim.data.plugins["counter"] == n_steps), ( + assert jnp.all(sim.data.plugins["counter"].count == n_steps), ( f"Expected counter={n_steps}, got {sim.data.plugins['counter']}" ) sim.close() @@ -71,14 +90,14 @@ def test_plugin_data_changes(n_worlds: int): def test_plugin_resets(): """sim.reset() restores the counter to its default value (0).""" sim = Sim() - sim.data = sim.data.replace(plugins={"counter": jnp.zeros((1, 1), dtype=jnp.int32)}) + sim.data = sim.data.replace(plugins={"counter": Counter.create(1)}) append_fn(sim.step_pipeline, counter_plugin) sim.build_default_data() sim.build_step_fn() sim.step(10) - assert jnp.all(sim.data.plugins["counter"] == 10), "Precondition: counter should be 10" + assert jnp.all(sim.data.plugins["counter"].count == 10), "Counter should be 10" sim.reset() - assert jnp.all(sim.data.plugins["counter"] == 0), "Counter should be 0 after full reset" + assert jnp.all(sim.data.plugins["counter"].count == 0), "Counter should be 0 after full reset" sim.close() @@ -87,17 +106,17 @@ def test_plugin_masked_reset(): """Masked reset only resets the plugin data for selected worlds.""" n_worlds = 2 sim = Sim(n_worlds=n_worlds) - sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)}) + sim.data = sim.data.replace(plugins={"counter": Counter.create(n_worlds)}) append_fn(sim.step_pipeline, counter_plugin) sim.build_default_data() sim.build_step_fn() sim.step(10) - assert jnp.all(sim.data.plugins["counter"] == 10), "Precondition: both counters should be 10" + assert jnp.all(sim.data.plugins["counter"].count == 10), "Both counters should be 10" mask = jnp.array([True, False]) # reset world 0 only sim.reset(mask) - assert jnp.all(sim.data.plugins["counter"][0] == 0), "World 0 counter must be reset to 0" - assert jnp.all(sim.data.plugins["counter"][1] == 10), "World 1 counter must remain at 10" + assert jnp.all(sim.data.plugins["counter"].count[0] == 0), "World 0 counter must be reset to 0" + assert jnp.all(sim.data.plugins["counter"].count[1] == 10), "World 1 counter must remain at 10" sim.close() @@ -106,10 +125,7 @@ def test_chained_plugins(): """Two chained plugins produce different results depending on their order in the pipeline.""" sim = Sim() sim.data = sim.data.replace( - plugins={ - "counter": jnp.zeros((1, 1), dtype=jnp.int32), - "accumulated": jnp.zeros((1, 1), dtype=jnp.int32), - } + plugins={"counter": Counter.create(1), "accumulated": Counter.create(1)} ) append_fn(sim.step_pipeline, counter_plugin) append_fn(sim.step_pipeline, accumulater_plugin) @@ -117,10 +133,10 @@ def test_chained_plugins(): sim.build_step_fn() n_steps = 3 sim.step(n_steps) - assert int(sim.data.plugins["counter"][0, 0]) == n_steps, ( + assert int(sim.data.plugins["counter"].count[0, 0]) == n_steps, ( f"Expected counter={n_steps}, got {sim.data.plugins['counter']}" ) - assert int(sim.data.plugins["accumulated"][0, 0]) == sum(range(1, n_steps + 1)), ( + assert int(sim.data.plugins["accumulated"].count[0, 0]) == sum(range(1, n_steps + 1)), ( f"Expected accumulated={sum(range(1, n_steps + 1))}, got {sim.data.plugins['accumulated']}" ) sim.close() diff --git a/tests/unit/test_sharding.py b/tests/unit/test_sharding.py new file mode 100644 index 00000000..55488c81 --- /dev/null +++ b/tests/unit/test_sharding.py @@ -0,0 +1,89 @@ +"""Unit tests for distributing the simulation across devices.""" + +from __future__ import annotations + +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from conftest import available_backends +from jax.sharding import PartitionSpec + +from crazyflow.control import Control +from crazyflow.sim import Sim +from crazyflow.sim.functional import attitude_control +from crazyflow.sim.sharding import placement, shard, world_mesh + +# Sharding tests require at least 2 devices +multi_device = [ + pytest.param( + platform, + marks=pytest.mark.skipif( + len(jax.devices(platform)) < 2, reason=f"requires at least 2 {platform} devices" + ), + ) + for platform in available_backends() +] + + +def assert_sharding(x: Any, sharding: Any): + assert x.sharding.is_equivalent_to(sharding, x.ndim), f"{x.sharding} != {sharding}" + + +@pytest.mark.unit +@pytest.mark.parametrize("platform", multi_device) +def test_placement(platform: str): + devices = jax.devices(platform) + sim = Sim(n_worlds=2 * len(devices), device=platform) + sim.shard(world_mesh(devices)) + assert sim.data.states.pos.sharding.spec == PartitionSpec("worlds") + assert sim.data.params.mass.sharding.spec == PartitionSpec("worlds") + assert sim.data.params.gravity_vec.sharding.spec == PartitionSpec() + assert sim.data.params.rotor_dyn_coef.sharding.spec == PartitionSpec() + + +@pytest.mark.unit +@pytest.mark.parametrize("platform", multi_device) +def test_sharded_oop_api(platform: str): + devices = jax.devices(platform) + sim = Sim(n_worlds=2 * len(devices), device=platform, control=Control.attitude) + sim.shard(world_mesh(devices)) + expected = placement(sim.data, world_mesh(jax.devices(platform))) + jax.tree.map(assert_sharding, sim.default_data, expected) + sim.attitude_control(jnp.zeros((sim.n_worlds, 1, 4))) + sim.step() + jax.tree.map(assert_sharding, sim.data, expected) + sim.reset() + jax.tree.map(assert_sharding, sim.data, expected) + + +@pytest.mark.unit +@pytest.mark.parametrize("platform", multi_device) +def test_sharded_functional_api(platform: str): + devices = jax.devices(platform) + sim = Sim(n_worlds=2 * len(devices), control=Control.attitude, device=platform) + step_fn, reset_fn = sim.build_step_fn(), sim.build_reset_fn() + mesh = world_mesh(devices) + data, default_data = shard(sim.data, mesh), shard(sim.default_data, mesh) + expected = placement(sim.data, mesh) + data = attitude_control(data, jnp.zeros((sim.n_worlds, 1, 4))) + data = step_fn(data) + jax.tree.map(assert_sharding, data, expected) + data = reset_fn(data, default_data) + jax.tree.map(assert_sharding, data, expected) + + +@pytest.mark.unit +@pytest.mark.parametrize("platform", multi_device) +def test_sharded_step_values(platform: str): + # Sharding must not change the simulation results + devices = jax.devices(platform) + sim = Sim(n_worlds=2 * len(devices), device=platform) + sim.step(10) + pos = np.asarray(sim.data.states.pos) + sim.reset() + sim.shard(world_mesh(devices)) + sim.step(10) + assert np.allclose(np.asarray(sim.data.states.pos), pos, atol=1e-6) diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index 7d91a382..3a6dde58 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -530,7 +530,7 @@ def assert_committed(obj0: Array | Any, path: str = "data"): elif isinstance(obj0, (list, tuple)): # Handle sequences for i, item0 in enumerate(obj0): assert_committed(item0, f"{path}[{i}]") - elif isinstance(obj0, type(sim.data.core.device)): # Device objects + elif isinstance(obj0, type(sim.device)): # Device objects pass # Devices themselves don't have committed attribute elif isinstance(obj0, dict): for key, value0 in obj0.items(): @@ -640,3 +640,26 @@ def test_fused_model(device: str, drone: str): sim.reset() sim.step(1) sim.close() + + +@pytest.mark.unit +def test_partial_reset_keeps_shared_arrays(): + # A partial reset with a world mask must leave shared arrays intact + sim = Sim(n_worlds=3) + gravity = jnp.array([1.0, 2.0, 3.0]) + sim.data = sim.data.replace(params=sim.data.params.replace(gravity_vec=gravity)) + sim.reset(jnp.array([True, False, False])) + assert jnp.array_equal(sim.data.params.gravity_vec, gravity) + + +@pytest.mark.unit +def test_full_reset_restores_shared_arrays(): + sim = Sim(n_worlds=3) + default_gravity, rng_key = sim.default_data.params.gravity_vec, sim.data.core.rng_key + sim.data = sim.data.replace( + params=sim.data.params.replace(gravity_vec=jnp.array([1.0, 2.0, 3.0])) + ) + sim.reset() + 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)) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index bb83d018..afbe1075 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1,7 +1,20 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import flax import jax +import jax.numpy as jnp import pytest +from flax.struct import field + +from crazyflow.control import Control +from crazyflow.sim import Sim +from crazyflow.utils import WORLD_INDEXED_KEY, enable_cache, world_mask -from crazyflow.utils import enable_cache +if TYPE_CHECKING: + from jax import Array @pytest.mark.unit @@ -47,3 +60,75 @@ def test_enable_cache(enable_xla: bool): jax.config.update("jax_persistent_cache_min_compile_time_secs", orig_min_time) if orig_xla is not None: jax.config.update("jax_persistent_cache_enable_xla_caches", orig_xla) + + +@pytest.mark.unit +def test_world_mask(): + # Check that the mask correctly handles cases that previously failed under shape-based detection + mask = world_mask(Sim(n_worlds=3, control=Control.attitude).data) + assert mask.states.pos + assert mask.states_deriv.acc + assert mask.core.steps + assert mask.params.mass + assert mask.controls.attitude.cmd + assert not mask.params.gravity_vec + assert not mask.params.drag_matrix + assert not mask.core.drone_mocap_ids + assert not mask.controls.attitude.params["kR"] + + +@pytest.mark.unit +def test_world_mask_covers_batched_arrays(): + # Make a shape-based check that we catch all world-batched arrays using a unique world length + sim = Sim(n_worlds=7, control=Control.attitude) + paths, leaves = jax.tree.flatten_with_path(sim.data)[0], jax.tree.leaves(world_mask(sim.data)) + for (path, x), batched in zip(paths, leaves): + world_axis = x.ndim >= 2 and x.shape[0] == 7 + assert batched == world_axis, f"{jax.tree_util.keystr(path)}: {x.shape=} is {batched=}" + + +@pytest.mark.unit +def test_world_mask_user_dataclass(): + @flax.struct.dataclass + class PluginState: + buffer: Array = field(metadata={WORLD_INDEXED_KEY: True}) + constant: Array = None + + state = PluginState(buffer=jnp.zeros((3, 2)), constant=jnp.zeros(3)) + mask = world_mask({"state": state, "array": jnp.zeros((3, 2))}) + assert mask["state"].buffer + assert not mask["state"].constant + assert not mask["array"] # Arrays that no struct declares are shared + + +@pytest.mark.unit +def test_world_mask_matches_structure(): + # The mask has to match the tree leaf for leaf, whatever container holds the arrays + @flax.struct.dataclass + class PluginState: # With world data + buffer: Array = field(metadata={WORLD_INDEXED_KEY: True}) + + @dataclass + class Opaque: # State that is treated as leaf + a: int + + array = jnp.zeros((3, 2)) + tree = { + "containers": (array, [array], {"x": array}), + "nested": (PluginState(buffer=array),), + "opaque": Opaque(1), + } + mask = world_mask(tree) + assert jax.tree.structure(mask) == jax.tree.structure(tree) + assert mask["nested"][0].buffer # A declaration survives inside a container + + +@pytest.mark.unit +def test_world_mask_rejects_tagged_struct(): + # The world axis is declared per array. Tags on structs are not allowed. + @flax.struct.dataclass + class TaggedStruct: + nested: Any = field(metadata={WORLD_INDEXED_KEY: True}) + + with pytest.raises(ValueError, match="tags structs"): + world_mask(TaggedStruct(nested=TaggedStruct(nested=None))) From cacc8a7eae372a2780c5e8b8ee6230fc816bd092 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Tue, 18 Aug 2026 21:55:28 +0200 Subject: [PATCH 2/2] Rework sharding to work with core ndims --- crazyflow/control/mellinger/control.py | 26 +++++------ .../dynamics/first_principles/dynamics.py | 24 +++++----- crazyflow/dynamics/so_rpy/dynamics.py | 20 ++++----- crazyflow/dynamics/so_rpy_rotor/dynamics.py | 22 +++++----- .../dynamics/so_rpy_rotor_drag/dynamics.py | 24 +++++----- crazyflow/sim/data.py | 44 +++++++++---------- crazyflow/utils.py | 42 +++++++++--------- docs/user-guide/sharding.md | 4 +- docs/user-guide/world-axis.md | 34 ++++++++++---- tests/unit/test_plugins.py | 7 ++- tests/unit/test_utils.py | 24 +++++++--- 11 files changed, 149 insertions(+), 122 deletions(-) diff --git a/crazyflow/control/mellinger/control.py b/crazyflow/control/mellinger/control.py index 85ef2608..9f0ef661 100644 --- a/crazyflow/control/mellinger/control.py +++ b/crazyflow/control/mellinger/control.py @@ -21,7 +21,7 @@ from crazyflow.control.core import controllable, load_params from crazyflow.control.transform import force2pwm, motor_force2rotor_vel, pwm2force -from crazyflow.utils import WORLD_INDEXED_KEY, leaf_replace +from crazyflow.utils import CORE_NDIM_KEY, leaf_replace if TYPE_CHECKING: from jax import Device @@ -311,19 +311,19 @@ def force_torque2rotor_vel( @dataclass class MellingerStateData: - cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 13) + cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 13) """Full state control command for the drone. A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]. We currently do not use the acceleration and angle rate components. This is subject to change. """ - staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 13) + staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 13) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) + steps: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, 1) """Last simulation steps that the state control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the state control command.""" - pos_err_i: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + pos_err_i: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Integral errors of the state control command.""" # Parameters for the state controller params: dict[str, Array] @@ -344,20 +344,20 @@ def create( @dataclass class MellingerAttitudeData: - cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Full attitude control command for the drone. A command consists of [roll, pitch, yaw, collective thrust]. """ - staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) + steps: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, 1) """Last simulation steps that the attitude control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the attitude control command.""" - r_int_error: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + r_int_error: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Integral errors of the attitude control command.""" - last_ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + last_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Last angular velocity of the drone.""" # Parameters for the attitude controller params: dict[str, Array] @@ -384,14 +384,14 @@ def create( @dataclass class MellingerForceTorqueData: - cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Force-torque command for the drone. A command consists of [fz, tx, ty, tz]. """ - staged_cmd: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Staging buffer to store the most recent command until the next controller tick.""" - steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) + steps: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, 1) """Last simulation steps that the force and torque control command was applied.""" freq: int = field(pytree_node=False) """Frequency of the force and torque control command.""" diff --git a/crazyflow/dynamics/first_principles/dynamics.py b/crazyflow/dynamics/first_principles/dynamics.py index 05f84866..1295ff10 100644 --- a/crazyflow/dynamics/first_principles/dynamics.py +++ b/crazyflow/dynamics/first_principles/dynamics.py @@ -28,7 +28,7 @@ import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import WORLD_INDEXED_KEY, to_xp +from crazyflow.utils import CORE_NDIM_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -298,27 +298,27 @@ def symbolic_dynamics( @dataclass class Params: - mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) + mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1) """Mass of the drone.""" - L: Array # () + L: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Arm length of the drone.""" - prop_inertia: Array # () + prop_inertia: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Inertia of the propeller.""" - gravity_vec: Array # (3,) + gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Gravity vector of the drone.""" - J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J_inv: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - rpm2thrust: Array # (3,) + rpm2thrust: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Force constant of the drone.""" - rpm2torque: Array # (3,) + rpm2torque: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Torque constant of the drone.""" - mixing_matrix: Array # (3, 4) + mixing_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 4) """Mixing matrix of the drone.""" - drag_matrix: Array # (3, 3) + drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 3) """Drag matrix of the drone.""" - rotor_dyn_coef: Array # (4,) + rotor_dyn_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (4,) """Rotor speed dynamics time constant of the drone.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy/dynamics.py b/crazyflow/dynamics/so_rpy/dynamics.py index 1a828620..4af83291 100644 --- a/crazyflow/dynamics/so_rpy/dynamics.py +++ b/crazyflow/dynamics/so_rpy/dynamics.py @@ -27,7 +27,7 @@ import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import WORLD_INDEXED_KEY, to_xp +from crazyflow.utils import CORE_NDIM_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -316,31 +316,31 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) + mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (3,) + gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Gravity vector of the drone.""" - J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J_inv: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - acc_coef: Array # () + acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Coefficient for the acceleration.""" - cmd_f_coef: Array # () + cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Coefficient for the collective thrust.""" - rpy_coef: Array # (3,) + rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Coefficient for the roll pitch yaw dynamics.""" - rpy_rates_coef: Array # (3,) + rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Coefficient for the roll pitch yaw rates dynamics.""" - cmd_rpy_coef: Array # (3,) + cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Coefficient for the roll pitch yaw command dynamics.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy_rotor/dynamics.py b/crazyflow/dynamics/so_rpy_rotor/dynamics.py index 6f0157fb..ae929e5f 100644 --- a/crazyflow/dynamics/so_rpy_rotor/dynamics.py +++ b/crazyflow/dynamics/so_rpy_rotor/dynamics.py @@ -29,7 +29,7 @@ import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import WORLD_INDEXED_KEY, to_xp +from crazyflow.utils import CORE_NDIM_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -375,25 +375,25 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) + mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (3,) + gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Gravity vector of the drone.""" - J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J_inv: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - thrust_time_coef: Array # () + thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Rotor coefficient of the drone.""" - acc_coef: Array # () + acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Acceleration coefficient of the drone.""" - cmd_f_coef: Array # () + cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Collective thrust coefficient of the drone.""" - rpy_coef: Array # (3,) + rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw coefficient of the drone.""" - rpy_rates_coef: Array # (3,) + rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw rates coefficient of the drone.""" - cmd_rpy_coef: Array # (3,) + cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw command coefficient of the drone.""" @staticmethod diff --git a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py index 0127c46d..c39a27c0 100644 --- a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py +++ b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py @@ -31,7 +31,7 @@ import crazyflow.dynamics.symbols as symbols from crazyflow.dynamics.core import load_params, supports from crazyflow.dynamics.utils import rotation -from crazyflow.utils import WORLD_INDEXED_KEY, to_xp +from crazyflow.utils import CORE_NDIM_KEY, to_xp if TYPE_CHECKING: from jax import Device @@ -409,27 +409,27 @@ def symbolic_dynamics_euler( @dataclass class Params: - mass: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 1) + mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1) """Mass of the drone.""" - gravity_vec: Array # (3,) + gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Gravity vector of the drone.""" - J: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inertia matrix of the drone.""" - J_inv: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3, 3) + J_inv: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3) """Inverse of the inertia matrix of the drone.""" - thrust_time_coef: Array # () + thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Rotor coefficient of the drone.""" - acc_coef: Array # () + acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Acceleration coefficient of the drone.""" - cmd_f_coef: Array # () + cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Collective thrust coefficient of the drone.""" - rpy_coef: Array # (3,) + rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw coefficient of the drone.""" - rpy_rates_coef: Array # (3,) + rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw rates coefficient of the drone.""" - cmd_rpy_coef: Array # (3,) + cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,) """Roll pitch yaw command coefficient of the drone.""" - drag_matrix: Array # (3, 3) + drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 3) """Linear drag coefficient matrix of the drone.""" @staticmethod diff --git a/crazyflow/sim/data.py b/crazyflow/sim/data.py index edf8ce57..42c313c1 100644 --- a/crazyflow/sim/data.py +++ b/crazyflow/sim/data.py @@ -19,24 +19,24 @@ from crazyflow.dynamics.so_rpy import Params as SoRpyParams from crazyflow.dynamics.so_rpy_rotor import Params as SoRpyRotorParams from crazyflow.dynamics.so_rpy_rotor_drag import Params as SoRpyRotorDragParams -from crazyflow.utils import WORLD_INDEXED_KEY +from crazyflow.utils import CORE_NDIM_KEY @dataclass class SimState: - pos: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + pos: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Position of the drone's center of mass.""" - quat: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + quat: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Quaternion of the drone's orientation.""" - vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Velocity of the drone's center of mass in the world frame.""" - ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Angular velocity of the drone in the body frame.""" - force: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) CoM force + force: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) CoM force """Force applied to the drone's center of mass in the world frame.""" - torque: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) CoM torque + torque: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) CoM torque """Torque applied to the drone's center of mass in the world frame.""" - rotor_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) in RPM + rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) in RPM """Motor forces along body frame z axis.""" @staticmethod @@ -59,15 +59,15 @@ def create(n_worlds: int, n_drones: int, device: Device) -> SimState: @dataclass class SimStateDeriv: - vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Derivative of the position of the drone's center of mass.""" - ang_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Derivative of the quaternion of the drone's orientation as angular velocity.""" - acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + acc: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Derivative of the velocity of the drone's center of mass.""" - ang_acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 3) + ang_acc: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) """Derivative of the angular velocity of the drone's center of mass.""" - rotor_acc: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + rotor_acc: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Derivative of the rotor velocity.""" @staticmethod @@ -92,7 +92,7 @@ class ControlData(typing.Protocol): """Control command for the drone.""" staged_cmd: Array # (N, M, X) """Staged control command for the drone.""" - steps: Array # (N, 1) + steps: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, 1) """Last simulation steps that the state control command was applied.""" freq: int """Frequency of the state control command.""" @@ -110,7 +110,7 @@ class SimControls: """Attitude control data.""" force_torque: ControlData | None """Force and torque control data.""" - rotor_vel: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, M, 4) + rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) """Desired motor speed.""" @staticmethod @@ -207,17 +207,17 @@ def create( class SimCore: freq: int = field(pytree_node=False) """Frequency of the simulation.""" - steps: Array = field(metadata={WORLD_INDEXED_KEY: True}) # (N, 1) + steps: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, 1) """Simulation steps taken since the last reset.""" n_worlds: int = field(pytree_node=False) """Number of worlds in the simulation.""" n_drones: int = field(pytree_node=False) """Number of drones in the simulation.""" - drone_mocap_ids: Array # (M,) + drone_mocap_ids: Array = field(metadata={CORE_NDIM_KEY: 1}) # (M,) """MuJoCo mocap IDs of the drone bodies.""" - rng_key: Array # () + rng_key: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Random number generator key for the simulation.""" - mjx_synced: Array # () + mjx_synced: Array = field(metadata={CORE_NDIM_KEY: 0}) # () """Whether the simulation data is synchronized with the MuJoCo mjx_data.""" @staticmethod @@ -260,7 +260,7 @@ class SimData: plugins: dict[str, Any] = field(default_factory=dict) """Arbitrary data for plugins to store state in the simulation. - Data that is batched over worlds has to declare this using the WORLD_INDEXED_KEY metadata, so - that resets can mask it and sharding can distribute it. This requires data to be stored in a - struct. + Data that is batched over worlds has to declare its dimensions without batch axes using the + CORE_NDIM_KEY metadata, so that resets can mask it and sharding can distribute it. Dicts have no + fields to declare on, so batched data has to be stored in a struct. """ diff --git a/crazyflow/utils.py b/crazyflow/utils.py index c4110ac3..e0f97f7a 100644 --- a/crazyflow/utils.py +++ b/crazyflow/utils.py @@ -16,8 +16,8 @@ if TYPE_CHECKING: from types import ModuleType -WORLD_INDEXED_KEY = "world_indexed" -"""Metadata key for dataclass fields that have a world axis.""" +CORE_NDIM_KEY = "core_ndim" +"""Metadata key for the number of dimensions a dataclass field has without any batch axes.""" def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array: @@ -39,8 +39,9 @@ def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array: def world_mask(tree: T) -> T: """Flag the arrays of a PyTree that are indexed by world. - Structs declare fields with a world axis with the WORLD_INDEXED_KEY metadata. Arrays without - fields or the metadata tag are shared. Structs must be flax dataclasses. + Fields declare their number of dimensions without batch axes with the CORE_NDIM_KEY metadata. + An array carries a world axis when its ndim exceeds that. Arrays without the metadata are + treated as not batched. Args: tree: PyTree of arrays and structs to classify. @@ -53,23 +54,22 @@ def world_mask(tree: T) -> T: def _flag(node: Any) -> Any: """Flag the leaves of a node, recursing through the declared structure of the data.""" - # Case 1: a dataclass with potential WORLD_INDEXED_KEY fields. We filter out the dynamic fields - # and check which of them declare a world axis - if _declares_fields(node): - dynamic = [f for f in fields(node) if f.metadata.get("pytree_node", True)] - tags = {f.name: f.metadata.get(WORLD_INDEXED_KEY, False) for f in dynamic} - values = {f: getattr(node, f) for f in tags} - # Structs with fields must not be tagged. Tagging is array-level only - if tagged_structs := [f for f, tag in tags.items() if tag and _declares_fields(values[f])]: - raise ValueError(f"{type(node).__name__} tags structs: {tagged_structs}") - # Tagged fields hold no structs, so all of their arrays are indexed by world - indexed = {f: jax.tree.map(lambda _: True, values[f]) for f in tags if tags[f]} - return node.replace(**indexed, **{f: _flag(values[f]) for f in tags if not tags[f]}) - # Case 2: any other node. Its arrays are shared. Nested structs are handed back to us via - # is_leaf so that we can further recurse into them - return jax.tree.map( - lambda x: _flag(x) if _declares_fields(x) else False, node, is_leaf=_declares_fields - ) + if not _declares_fields(node): + # Any other node. Nested structs are handed back to us via is_leaf so that we can recurse + return jax.tree.map( + lambda x: _flag(x) if _declares_fields(x) else False, node, is_leaf=_declares_fields + ) + flags = {} + for f in (f for f in fields(node) if f.metadata.get("pytree_node", True)): + value, core_ndim = getattr(node, f.name), f.metadata.get(CORE_NDIM_KEY) + if core_ndim is None: # Structs and containers classify their own contents + flags[f.name] = _flag(value) + elif _declares_fields(value): + name = f"{type(node).__name__}.{f.name}" + raise ValueError(f"{name} holds a struct, which declares its own fields") + else: + flags[f.name] = jax.tree.map(lambda x, core=core_ndim: x.ndim > core, value) + return node.replace(**flags) def _declares_fields(node: Any) -> bool: diff --git a/docs/user-guide/sharding.md b/docs/user-guide/sharding.md index b595b0e3..f5873612 100644 --- a/docs/user-guide/sharding.md +++ b/docs/user-guide/sharding.md @@ -20,7 +20,7 @@ assert sim.data.states.pos.sharding.spec == jax.sharding.PartitionSpec("worlds") assert sim.data.params.gravity_vec.sharding.spec == jax.sharding.PartitionSpec() ``` -The number of worlds has to be divisible by the number of devices. Arrays declare that they are batched over worlds on their fields, see [The world axis](world-axis.md). Data shared across worlds, such as the gravity vector above, is replicated on every device. +The number of worlds has to be divisible by the number of devices. Whether an array is batched over worlds is determined from its shape and the core ndim its field declares, see [The world axis](world-axis.md). Data shared across worlds, such as the gravity vector above, is replicated on every device. !!! warning Sharding relies on JAX's automatic sharding mode, which `world_mesh` requests for you. `jax.make_mesh` defaults to explicit sharding, under which the Mellinger controller fails inside SciPy's rotation backend. @@ -46,7 +46,7 @@ data = jax.device_put(sim.data, placement(sim.data, mesh)) ## Plugin data -Bare arrays in `sim.data.plugins` are replicated, since a dict has no fields to declare the world axis. Store per-world plugin state in a struct that declares it, as shown in [The world axis](world-axis.md#your-own-state). +Bare arrays in `sim.data.plugins` are replicated, since a dict has no fields to declare a core ndim on. Store per-world plugin state in a struct that declares one, as shown in [The world axis](world-axis.md#your-own-state). ## Next steps diff --git a/docs/user-guide/world-axis.md b/docs/user-guide/world-axis.md index 87cbce3b..e97f97cf 100644 --- a/docs/user-guide/world-axis.md +++ b/docs/user-guide/world-axis.md @@ -2,11 +2,11 @@ The simulation is batched over worlds, so most arrays in `sim.data` carry a leading `n_worlds` axis. Resetting selected worlds and distributing worlds across devices both need to know which arrays those are. Here we explain how that is declared and what depends on it. -## Why it is declared, not inferred +## How it is declared -A world axis is an ordinary axis, so no shape tells you whether it indexes worlds. With three worlds, `params.drag_matrix` of shape `(3, 3)` and `params.gravity_vec` of shape `(3,)` both look world-batched while neither is. +A world axis is an ordinary axis, so no shape on its own tells you whether it indexes worlds. With three worlds, `params.drag_matrix` of shape `(3, 3)` and `params.gravity_vec` of shape `(3,)` both look world-batched while neither is. -Each field therefore declares if it has a leading world axis with the `WORLD_INDEXED_KEY` metadata: +Each field therefore declares its number of dimensions without batch axes with the `CORE_NDIM_KEY` metadata. `states.pos` is a vector, so its core ndim is 1, and `params.drag_matrix` is a matrix, so its core ndim is 2. An array carries a world axis exactly when its own `ndim` exceeds the core ndim of its field: ```python from crazyflow.sim import Sim @@ -21,7 +21,7 @@ assert not mask.params.gravity_vec # (3,), shared by all worlds assert not mask.params.drag_matrix # (3, 3), shared by all worlds ``` -`world_mask` returns a pytree of booleans matching the data, with one flag per array. Arrays that no field declares are shared by all worlds. +`world_mask` returns a pytree of booleans matching the data, with one flag per array. Arrays whose field declares no core ndim are shared by all worlds. ## Resets @@ -44,26 +44,42 @@ To keep a change to a shared parameter across masked resets, write it into `sim. ## Sharding -Distributing the simulation partitions the world axis of the declared arrays and replicates the rest, using the same declarations. See [Sharding](sharding.md) for the mesh and placement API. +Distributing the simulation partitions the world axis of the arrays that have one and replicates the rest, using the same declarations. See [Sharding](sharding.md) for the mesh and placement API. + +## Randomizing a shared parameter + +A parameter that is shared by default gains a world axis the moment you randomize it per world, and its `ndim` then exceeds the declared core ndim: + +```{ .python continuation } +n = sim.n_worlds +scale = jnp.arange(1, n + 1)[:, None, None, None] +drag = jnp.broadcast_to(sim.data.params.drag_matrix, (n, 1, 3, 3)) * scale +sim.data = sim.data.replace(params=sim.data.params.replace(drag_matrix=drag)) +sim.build_default_data() # Keep the randomized shape across full resets + +assert world_mask(sim.data).params.drag_matrix +``` + +From there it behaves like any other world-indexed array. A masked reset restores it for the masked worlds, and sharding partitions it instead of replicating it. Setting it back to a `(3, 3)` array makes it shared again. ## Your own state -Anything you add to `sim.data.plugins` follows the same rule. A `dict` has no fields, so bare arrays are shared by default, which means they are neither reset nor distributed. If you need per-world states, you must declare it with a struct: +Anything you add to `sim.data.plugins` follows the same rule. A `dict` has no fields, so bare arrays are shared, which means a masked reset leaves them alone and sharding replicates them. If you need per-world state, declare it on a struct: ```python import flax.struct from jax import Array -from crazyflow.utils import WORLD_INDEXED_KEY +from crazyflow.utils import CORE_NDIM_KEY @flax.struct.dataclass class PluginState: - buffer: Array = flax.struct.field(metadata={WORLD_INDEXED_KEY: True}) + buffer: Array = flax.struct.field(metadata={CORE_NDIM_KEY: 1}) lookup: Array = None ``` -If you store this under `sim.data.plugins`, `buffer` is restored for the masked worlds and partitioned when the simulation is sharded, while `lookup` is only restored by a full reset and replicated across devices. +If you store this under `sim.data.plugins`, a `buffer` of shape `(n_worlds, 3)` exceeds its core ndim of 1, so it is restored for the masked worlds and partitioned when the simulation is sharded. A `lookup` of any shape declares no core ndim, so it is only restored by a full reset and replicated across devices. ## Next steps diff --git a/tests/unit/test_plugins.py b/tests/unit/test_plugins.py index d728aa67..1a27d76b 100644 --- a/tests/unit/test_plugins.py +++ b/tests/unit/test_plugins.py @@ -2,8 +2,7 @@ Plugins are callables of the form ``fn(data: SimData) -> SimData`` that are inserted into ``sim.step_pipeline``. They can store arbitrary state in ``sim.data.plugins``. State that is batched -over worlds has to declare its world axis so that resets can mask it. Tests use a simple per-world -step counter as a canonical plugin. +over worlds has to declare its dimensions without batch axes so that resets can mask it. """ from __future__ import annotations @@ -17,7 +16,7 @@ from crazyflow.sim import Sim from crazyflow.sim.pipeline import append_fn -from crazyflow.utils import WORLD_INDEXED_KEY +from crazyflow.utils import CORE_NDIM_KEY if TYPE_CHECKING: from jax import Array @@ -29,7 +28,7 @@ class Counter: """Per-world plugin state.""" - count: Array = field(metadata={WORLD_INDEXED_KEY: True}) + count: Array = field(metadata={CORE_NDIM_KEY: 1}) @staticmethod def create(n_worlds: int) -> Counter: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index afbe1075..3ae49075 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -11,7 +11,7 @@ from crazyflow.control import Control from crazyflow.sim import Sim -from crazyflow.utils import WORLD_INDEXED_KEY, enable_cache, world_mask +from crazyflow.utils import CORE_NDIM_KEY, enable_cache, world_mask if TYPE_CHECKING: from jax import Array @@ -91,7 +91,7 @@ def test_world_mask_covers_batched_arrays(): def test_world_mask_user_dataclass(): @flax.struct.dataclass class PluginState: - buffer: Array = field(metadata={WORLD_INDEXED_KEY: True}) + buffer: Array = field(metadata={CORE_NDIM_KEY: 1}) constant: Array = None state = PluginState(buffer=jnp.zeros((3, 2)), constant=jnp.zeros(3)) @@ -106,7 +106,7 @@ def test_world_mask_matches_structure(): # The mask has to match the tree leaf for leaf, whatever container holds the arrays @flax.struct.dataclass class PluginState: # With world data - buffer: Array = field(metadata={WORLD_INDEXED_KEY: True}) + buffer: Array = field(metadata={CORE_NDIM_KEY: 1}) @dataclass class Opaque: # State that is treated as leaf @@ -125,10 +125,22 @@ class Opaque: # State that is treated as leaf @pytest.mark.unit def test_world_mask_rejects_tagged_struct(): - # The world axis is declared per array. Tags on structs are not allowed. + # A struct declares its own fields, so a core ndim on a field holding a struct is an error. @flax.struct.dataclass class TaggedStruct: - nested: Any = field(metadata={WORLD_INDEXED_KEY: True}) + nested: Any = field(metadata={CORE_NDIM_KEY: 1}) - with pytest.raises(ValueError, match="tags structs"): + with pytest.raises(ValueError, match="declares its own fields"): world_mask(TaggedStruct(nested=TaggedStruct(nested=None))) + + +@pytest.mark.unit +def test_world_mask_infers_randomized_parameter(): + # A parameter that is shared by default declares itself once it gains a world axis + sim = Sim(n_worlds=3) + assert not world_mask(sim.data).params.drag_matrix # (3, 3), shared + drag = jnp.broadcast_to(sim.data.params.drag_matrix, (3, 1, 3, 3)) + sim.data = sim.data.replace(params=sim.data.params.replace(drag_matrix=drag)) + mask = world_mask(sim.data) + assert mask.params.drag_matrix # (3, 1, 3, 3), now world-indexed + assert not mask.params.gravity_vec # Siblings are unaffected