Skip to content

Add sharding support and fix the world axis #99

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ratheron merged 2 commits into from
Aug 20, 2026
Merged

Add sharding support and fix the world axis #99

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
04a1204
Add sharding support and fix the world axis
amacati Aug 18, 2026
cacc8a7
Rework sharding to work with core ndims
amacati Aug 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions crazyflow/control/mellinger/control.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 CORE_NDIM_KEY, leaf_replace

if TYPE_CHECKING:
from jax import Device
Expand Down Expand Up @@ -311,19 +311,19 @@ def force_torque2rotor_vel(

@dataclass
class MellingerStateData:
cmd: Array # (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 # (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 # (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 # (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]
Expand All @@ -344,20 +344,20 @@ def create(

@dataclass
class MellingerAttitudeData:
cmd: Array # (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 # (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 # (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 # (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 # (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]
Expand All @@ -384,14 +384,14 @@ def create(

@dataclass
class MellingerForceTorqueData:
cmd: Array # (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 # (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 # (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."""
Expand Down
26 changes: 13 additions & 13 deletions crazyflow/dynamics/first_principles/dynamics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 CORE_NDIM_KEY, to_xp

if TYPE_CHECKING:
from jax import Device
Expand Down Expand Up @@ -298,27 +298,27 @@ def symbolic_dynamics(

@dataclass
class Params:
mass: Array # (N, M, 1)
mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Mass of the drone."""
L: Array # (N, M, 1)
L: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
Comment thread
amacati marked this conversation as resolved.
"""Arm length of the drone."""
prop_inertia: Array # (N, M, 1)
prop_inertia: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Inertia of the propeller."""
gravity_vec: Array # (N, M, 3)
gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Gravity vector of the drone."""
J: Array # (N, M, 3, 3)
J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3)
"""Inertia matrix of the drone."""
J_inv: Array # (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 # (N, M, 1)
rpm2thrust: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Force constant of the drone."""
rpm2torque: Array # (N, M, 1)
rpm2torque: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Torque constant of the drone."""
mixing_matrix: Array # (N, M, 3, 4)
mixing_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 4)
"""Mixing matrix of the drone."""
drag_matrix: Array # (N, M, 3, 3)
drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 3)
"""Drag matrix of the drone."""
rotor_dyn_coef: Array # (N, M, 4)
rotor_dyn_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (4,)
"""Rotor speed dynamics time constant of the drone."""

@staticmethod
Expand Down
22 changes: 11 additions & 11 deletions crazyflow/dynamics/so_rpy/dynamics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 CORE_NDIM_KEY, to_xp

if TYPE_CHECKING:
from jax import Device
Expand Down Expand Up @@ -316,31 +316,31 @@ def symbolic_dynamics_euler(

@dataclass
class Params:
mass: Array # (N, M, 1)
mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Mass of the drone."""

gravity_vec: Array # (N, M, 3)
gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Gravity vector of the drone."""

J: Array # (N, M, 3, 3)
J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3)
"""Inertia matrix of the drone."""

J_inv: Array # (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 # (N, M, 1)
acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
Comment thread
amacati marked this conversation as resolved.
"""Coefficient for the acceleration."""

cmd_f_coef: Array # (N, M, 1)
cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Coefficient for the collective thrust."""

rpy_coef: Array # (N, M, 1)
rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Coefficient for the roll pitch yaw dynamics."""

rpy_rates_coef: Array # (N, M, 1)
rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Coefficient for the roll pitch yaw rates dynamics."""

cmd_rpy_coef: Array # (N, M, 1)
cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Coefficient for the roll pitch yaw command dynamics."""

@staticmethod
Expand Down
24 changes: 12 additions & 12 deletions crazyflow/dynamics/so_rpy_rotor/dynamics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 CORE_NDIM_KEY, to_xp

if TYPE_CHECKING:
from jax import Device
Expand Down Expand Up @@ -375,25 +375,25 @@ def symbolic_dynamics_euler(

@dataclass
class Params:
mass: Array # (N, M, 1)
mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Mass of the drone."""
gravity_vec: Array # (N, M, 3)
gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Gravity vector of the drone."""
J: Array # (N, M, 3, 3)
J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3)
"""Inertia matrix of the drone."""
J_inv: Array # (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 # (N, M, 1)
thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
Comment thread
amacati marked this conversation as resolved.
"""Rotor coefficient of the drone."""
acc_coef: Array # (N, M, 1)
acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Acceleration coefficient of the drone."""
cmd_f_coef: Array # (N, M, 1)
cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Collective thrust coefficient of the drone."""
rpy_coef: Array # (N, M, 1)
rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw coefficient of the drone."""
rpy_rates_coef: Array # (N, M, 1)
rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw rates coefficient of the drone."""
cmd_rpy_coef: Array # (N, M, 1)
cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw command coefficient of the drone."""

@staticmethod
Expand Down
26 changes: 13 additions & 13 deletions crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py
View file
Open in desktop
Comment thread
ratheron marked this conversation as resolved.
Original file line number Diff line number Diff line change
Expand Up @@ -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 CORE_NDIM_KEY, to_xp

if TYPE_CHECKING:
from jax import Device
Expand Down Expand Up @@ -409,27 +409,27 @@ def symbolic_dynamics_euler(

@dataclass
class Params:
mass: Array # (N, M, 1)
mass: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Mass of the drone."""
gravity_vec: Array # (N, M, 3)
gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Gravity vector of the drone."""
J: Array # (N, M, 3, 3)
J: Array = field(metadata={CORE_NDIM_KEY: 2}) # (N, M, 3, 3)
"""Inertia matrix of the drone."""
J_inv: Array # (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 # (N, M, 1)
thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
Comment thread
amacati marked this conversation as resolved.
"""Rotor coefficient of the drone."""
acc_coef: Array # (N, M, 1)
acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Acceleration coefficient of the drone."""
cmd_f_coef: Array # (N, M, 1)
cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0}) # ()
"""Collective thrust coefficient of the drone."""
rpy_coef: Array # (N, M, 1)
rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw coefficient of the drone."""
rpy_rates_coef: Array # (N, M, 1)
rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw rates coefficient of the drone."""
cmd_rpy_coef: Array # (N, M, 1)
cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1}) # (3,)
"""Roll pitch yaw command coefficient of the drone."""
drag_matrix: Array # (N, M, 3, 3)
drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2}) # (3, 3)
"""Linear drag coefficient matrix of the drone."""

@staticmethod
Expand Down
Loading
Loading