diff --git a/crazyflow/control/__init__.py b/crazyflow/control/__init__.py index bb651c4b..b1c3e84e 100644 --- a/crazyflow/control/__init__.py +++ b/crazyflow/control/__init__.py @@ -16,11 +16,13 @@ from crazyflow.control.core import Control, load_params, parametrize from crazyflow.control.mellinger import attitude2force_torque as mellinger_attitude2force_torque +from crazyflow.control.mellinger import body_rate2force_torque as mellinger_body_rate2force_torque from crazyflow.control.mellinger import state2attitude as mellinger_state2attitude available_controller: dict[str, Callable] = { "mellinger_state2attitude": mellinger_state2attitude, "mellinger_attitude2force_torque": mellinger_attitude2force_torque, + "mellinger_body_rate2force_torque": mellinger_body_rate2force_torque, } __all__ = ["Control", "load_params", "parametrize"] diff --git a/crazyflow/control/core.py b/crazyflow/control/core.py index 7581c2d8..ccd50dd7 100644 --- a/crazyflow/control/core.py +++ b/crazyflow/control/core.py @@ -100,6 +100,12 @@ class Control(StrEnum): Note: Recommended frequency is >=100 Hz. """ + body_rate = "body_rate" + """Body rate control takes [roll_rate, pitch_rate, yaw_rate, collective thrust]. + + Note: + Recommended frequency is >=200 Hz. + """ force_torque = "force_torque" """Force and torque control takes [fc, tx, ty, tz]. diff --git a/crazyflow/control/mellinger/__init__.py b/crazyflow/control/mellinger/__init__.py index 570613bb..1c18ff4b 100644 --- a/crazyflow/control/mellinger/__init__.py +++ b/crazyflow/control/mellinger/__init__.py @@ -5,10 +5,13 @@ from crazyflow.control.mellinger.control import ( MellingerAttitudeData, + MellingerBodyRateData, MellingerForceTorqueData, MellingerStateData, attitude2force_torque, + body_rate2force_torque, control_attitude2force_torque, + control_body_rate2force_torque, control_commit_attitude, control_force_torque2rotor_vel, control_state2attitude, @@ -19,12 +22,15 @@ __all__ = [ "state2attitude", "attitude2force_torque", + "body_rate2force_torque", "force_torque2rotor_vel", "MellingerStateData", "MellingerAttitudeData", + "MellingerBodyRateData", "MellingerForceTorqueData", "control_state2attitude", "control_attitude2force_torque", + "control_body_rate2force_torque", "control_commit_attitude", "control_force_torque2rotor_vel", ] diff --git a/crazyflow/control/mellinger/control.py b/crazyflow/control/mellinger/control.py index 016dfdd8..384b6dc8 100644 --- a/crazyflow/control/mellinger/control.py +++ b/crazyflow/control/mellinger/control.py @@ -3,7 +3,8 @@ The controller is split into three pure functions that form a pipeline: ``state2attitude`` → ``attitude2force_torque`` → ``force_torque2rotor_vel``. Each stage can be used independently or chained together to produce per-motor -RPM commands from a full-state setpoint. +RPM commands from a full-state setpoint. ``body_rate2force_torque`` replaces +the second stage for body rate setpoints. Reference: D. Mellinger and V. Kumar, "Minimum snap trajectory generation and control for quadrotors", ICRA 2011. @@ -175,6 +176,8 @@ def attitude2force_torque( quat: Drone orientation as xyzw quaternion with shape (..., 4). ang_vel: Drone angular drone velocity in rad/s with shape (..., 3). cmd: Commanded attitude (roll, pitch, yaw) and total thrust [rad, rad, rad, N]. + prev_ang_vel: Angular velocity in rad/s from the previous call. If None, it is initialised + to zero. r_int_error: Angular velocity integral error (..., 3) from the previous call. If None, it is initialised to zero. ctrl_freq: Control frequency in Hz @@ -187,7 +190,6 @@ def attitude2force_torque( thrust_max: Maximum thrust in N. pwm_min: Minimum PWM value. pwm_max: Maximum PWM value. - prev_ang_vel: Previous angular velocity in rad/s. L: Distance from the center of the quadrotor to the center of the rotor in m. thrust2torque: Conversion factor (m). mixing_matrix: Mixing matrix for the motor forces with shape (4, 3). @@ -196,8 +198,183 @@ def attitude2force_torque( Desired force (1,), torques (3,) and i_error_m """ xp = array_namespace(quat) - force_des = cmd[..., 3] # Total thrust in N - rpy_des = cmd[..., :3] + ang_vel_des = xp.zeros_like(ang_vel) # Attitude control assumes a zero body rate setpoint + return _attitude2force_torque( + quat, + ang_vel, + cmd[..., :3], + ang_vel_des, + cmd[..., 3], + prev_ang_vel, + ang_vel_des, + r_int_error, + ctrl_freq, + kR=kR, + kw=kw, + ki_m=ki_m, + kd_omega=kd_omega, + int_err_max=int_err_max, + torque_pwm_max=torque_pwm_max, + thrust_max=thrust_max, + pwm_min=pwm_min, + pwm_max=pwm_max, + L=L, + thrust2torque=thrust2torque, + mixing_matrix=mixing_matrix, + ) + + +def body_rate2force_torque( + quat: Array, + ang_vel: Array, + cmd: Array, + prev_ang_vel: Array | None = None, + prev_cmd: Array | None = None, + r_int_error: Array | None = None, + ctrl_freq: int = 500, + *, + kR: Array, + kw: Array, + ki_m: Array, + kd_omega: Array, + int_err_max: Array, + torque_pwm_max: Array, + thrust_max: float, + pwm_min: float, + pwm_max: float, + L: float, + thrust2torque: float, + mixing_matrix: Array, +) -> tuple[Array, Array, Array]: + """Compute the body rate to desired force-torque part of the Mellinger controller. + + The firmware Mellinger controller has no dedicated body rate mode. A body rate setpoint enters + the angular velocity error and its derivative, while the attitude terms level the drone at its + current yaw. This function reproduces this behavior with the gains of the attitude controller. + Set ``kR`` and ``ki_m`` to zero to track body rates without the attitude terms. + + Note: + We omit the axis flip in the firmware as it has only been introduced to make the controller + compatible with the new frame of the Crazyflie 2.1. + + Args: + quat: Drone orientation as xyzw quaternion with shape (..., 4). + ang_vel: Drone angular velocity in the body frame in rad/s with shape (..., 3). + cmd: Commanded body rates (wx, wy, wz) and total thrust [rad/s, rad/s, rad/s, N]. + prev_ang_vel: Angular velocity in rad/s from the previous call. If None, it is initialised + to zero. + prev_cmd: Command from the previous call with shape (..., 4). The firmware includes the + derivative of the body rate setpoint in the derivative term. If None, the setpoint is + assumed to be constant. + r_int_error: Angular velocity integral error (..., 3) from the previous call. If None, it + is initialised to zero. + ctrl_freq: Control frequency in Hz + kR: Proportional gain for the rotation error with shape (3,). + kw: Proportional gain for the angular velocity error with shape (3,). + ki_m: Integral gain for the rotation error with shape (3,). + kd_omega: Derivative gain for the angular velocity error with shape (3,). + int_err_max: Range of the integral error with shape (3,). i_range in the firmware. + torque_pwm_max: Maximum torque in PWM. + thrust_max: Maximum thrust in N. + pwm_min: Minimum PWM value. + pwm_max: Maximum PWM value. + L: Distance from the center of the quadrotor to the center of the rotor in m. + thrust2torque: Conversion factor (m). + mixing_matrix: Mixing matrix for the motor forces with shape (4, 3). + + Returns: + Desired force (1,), torques (3,) and i_error_m + """ + xp = array_namespace(quat) + # l. 215 ff Without a position or attitude setpoint, the firmware levels the drone at the + # current yaw + yaw = R.from_quat(quat).as_euler("xyz", degrees=False)[..., 2] + rpy_des = xp.stack((xp.zeros_like(yaw), xp.zeros_like(yaw), yaw), axis=-1) + ang_vel_des = cmd[..., :3] + prev_ang_vel_des = ang_vel_des if prev_cmd is None else prev_cmd[..., :3] + return _attitude2force_torque( + quat, + ang_vel, + rpy_des, + ang_vel_des, + cmd[..., 3], + prev_ang_vel, + prev_ang_vel_des, + r_int_error, + ctrl_freq, + kR=kR, + kw=kw, + ki_m=ki_m, + kd_omega=kd_omega, + int_err_max=int_err_max, + torque_pwm_max=torque_pwm_max, + thrust_max=thrust_max, + pwm_min=pwm_min, + pwm_max=pwm_max, + L=L, + thrust2torque=thrust2torque, + mixing_matrix=mixing_matrix, + ) + + +def _attitude2force_torque( + quat: Array, + ang_vel: Array, + rpy_des: Array, + ang_vel_des: Array, + force_des: Array, + prev_ang_vel: Array | None, + prev_ang_vel_des: Array, + r_int_error: Array | None, + ctrl_freq: int, + *, + kR: Array, + kw: Array, + ki_m: Array, + kd_omega: Array, + int_err_max: Array, + torque_pwm_max: Array, + thrust_max: float, + pwm_min: float, + pwm_max: float, + L: float, + thrust2torque: float, + mixing_matrix: Array, +) -> tuple[Array, Array, Array]: + """Attitude and body rate controller of the Mellinger controller. + + This function follows the structure of the firmware implementation. The firmware setpoint + carries both an attitude and a body rate. The attitude and body rate controllers route their + commands into the respective setpoint. + + Args: + quat: Drone orientation as xyzw quaternion with shape (..., 4). + ang_vel: Drone angular velocity in the body frame in rad/s with shape (..., 3). + rpy_des: Desired attitude as roll, pitch, yaw in rad with shape (..., 3). + ang_vel_des: Desired angular velocity in the body frame in rad/s with shape (..., 3). + force_des: Desired total thrust in N with shape (...,). + prev_ang_vel: Angular velocity from the previous call. If None, it is initialised to zero. + prev_ang_vel_des: Desired angular velocity from the previous call with shape (..., 3). + r_int_error: Rotation integral error (..., 3) from the previous call. If None, it is + initialised to zero. + ctrl_freq: Control frequency in Hz + kR: Proportional gain for the rotation error with shape (3,). + kw: Proportional gain for the angular velocity error with shape (3,). + ki_m: Integral gain for the rotation error with shape (3,). + kd_omega: Derivative gain for the angular velocity error with shape (3,). + int_err_max: Range of the integral error with shape (3,). i_range in the firmware. + torque_pwm_max: Maximum torque in PWM. + thrust_max: Maximum thrust in N. + pwm_min: Minimum PWM value. + pwm_max: Maximum PWM value. + L: Distance from the center of the quadrotor to the center of the rotor in m. + thrust2torque: Conversion factor (m). + mixing_matrix: Mixing matrix for the motor forces with shape (4, 3). + + Returns: + Desired force (..., 1), torques (..., 3) and i_error_m + """ + xp = array_namespace(quat) dt = 1 / ctrl_freq # l. 220 ff [eR]. We're using the "inefficient" code path from the firmware rot = R.from_quat(quat) @@ -210,11 +387,10 @@ def attitude2force_torque( # Vee operator (SO3 to R3) eR = xp.stack((eRM[..., 2, 1], eRM[..., 0, 2], eRM[..., 1, 0]), axis=-1) # l.248 ff [ew] - # Warning: We assume zero desired angular velocity - ang_vel_des = xp.zeros_like(ang_vel) - prev_ang_vel_des = xp.zeros_like(ang_vel) + # The firmware negates the pitch components of the gyro and the rate setpoint to convert them + # to the legacy Crazyflie frame, matching the sign flip of eR.y. We omit both flips and keep all + # terms in the standard body frame, so the setpoint enters without a sign change. ew = ang_vel_des - ang_vel - # WARNING: if the setpoint is ever != 0 => change sign of ew.y! # l.259 ff [err_d_rpy] prev_ang_vel = xp.zeros_like(ang_vel) if prev_ang_vel is None else prev_ang_vel @@ -389,6 +565,46 @@ def create( ) +@dataclass +class MellingerBodyRateData: + cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) + """Body rate control command for the drone. + + A command consists of [roll_rate, pitch_rate, yaw_rate, collective thrust]. + """ + 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={CORE_NDIM_KEY: 1}) # (N, 1) + """Last simulation steps that the body rate control command was applied.""" + freq: int = field(pytree_node=False) + """Frequency of the body rate control command.""" + r_int_error: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) + """Integral errors of the body rate control command.""" + last_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3) + """Last angular velocity of the drone.""" + # Parameters for the body rate controller + params: dict[str, Array] + + @staticmethod + def create( + n_worlds: int, n_drones: int, freq: int, drone: str, device: Device + ) -> MellingerBodyRateData: + """Create a default set of body rate data for the simulation.""" + zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device) + zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device) + steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device) + params = load_params(body_rate2force_torque, drone, xp=jnp, device=device) + return MellingerBodyRateData( + cmd=zeros_4d.copy(), + staged_cmd=zeros_4d.copy(), + steps=steps, + freq=freq, + r_int_error=zeros_3d.copy(), + last_ang_vel=zeros_3d.copy(), + params=params, + ) + + @dataclass class MellingerForceTorqueData: cmd: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) @@ -469,6 +685,39 @@ def control_attitude2force_torque(data: SimData) -> SimData: ) +def control_body_rate2force_torque(data: SimData) -> SimData: + """Compute the updated controls for the body rate controller.""" + states = data.states + body_rate_ctrl: MellingerBodyRateData = data.controls.body_rate + assert body_rate_ctrl is not None, "Using body rate controller without initialized data" + mask = controllable(data.core.steps, data.core.freq, body_rate_ctrl.steps, body_rate_ctrl.freq) + prev_cmd = body_rate_ctrl.cmd + body_rate_ctrl = leaf_replace(body_rate_ctrl, mask, cmd=body_rate_ctrl.staged_cmd) + force, torque, r_int_error = body_rate2force_torque( + states.quat, + states.ang_vel, + body_rate_ctrl.cmd, + prev_ang_vel=body_rate_ctrl.last_ang_vel, + prev_cmd=prev_cmd, + r_int_error=body_rate_ctrl.r_int_error, + ctrl_freq=body_rate_ctrl.freq, + **body_rate_ctrl.params, + ) + body_rate_ctrl = leaf_replace( + body_rate_ctrl, + mask, + r_int_error=r_int_error, + last_ang_vel=states.ang_vel, + steps=data.core.steps, + ) + ft_ctrl = leaf_replace( + data.controls.force_torque, mask, staged_cmd=jnp.concat([force, torque], axis=-1) + ) + return data.replace( + controls=data.controls.replace(body_rate=body_rate_ctrl, force_torque=ft_ctrl) + ) + + def control_commit_attitude(data: SimData) -> SimData: """Commit the staged attitude command to the controller setpoint.""" attitude_ctrl: MellingerAttitudeData = data.controls.attitude diff --git a/crazyflow/control/mellinger/params.toml b/crazyflow/control/mellinger/params.toml index 9589be8d..e4270c01 100644 --- a/crazyflow/control/mellinger/params.toml +++ b/crazyflow/control/mellinger/params.toml @@ -31,6 +31,13 @@ ki_m = [0.0, 0.0, 500.0] kd_omega = [200.0, 200.0, 0.0] int_err_max = [1.0, 1.0, 1500.0] +[cf2x_L250.body_rate2force_torque] +kR = [70000.0, 70000.0, 60000.0] +kw = [20000.0, 20000.0, 12000.0] +ki_m = [0.0, 0.0, 500.0] +kd_omega = [200.0, 200.0, 0.0] +int_err_max = [1.0, 1.0, 1500.0] + [cf2x_P250] [cf2x_P250.core] mass = 0.029 # The controller is using the wrong mass by default @@ -64,6 +71,13 @@ ki_m = [0.0, 0.0, 500.0] kd_omega = [200.0, 200.0, 0.0] int_err_max = [1.0, 1.0, 1500.0] +[cf2x_P250.body_rate2force_torque] +kR = [70000.0, 70000.0, 60000.0] +kw = [20000.0, 20000.0, 12000.0] +ki_m = [0.0, 0.0, 500.0] +kd_omega = [200.0, 200.0, 0.0] +int_err_max = [1.0, 1.0, 1500.0] + [cf2x_T350] [cf2x_T350.core] mass = 0.0325 # The controller is using the wrong mass by default @@ -97,6 +111,13 @@ ki_m = [0.0, 0.0, 500.0] kd_omega = [200.0, 200.0, 0.0] int_err_max = [1.0, 1.0, 1500.0] +[cf2x_T350.body_rate2force_torque] +kR = [70000.0, 70000.0, 60000.0] +kw = [20000.0, 20000.0, 12000.0] +ki_m = [0.0, 0.0, 500.0] +kd_omega = [200.0, 200.0, 0.0] +int_err_max = [1.0, 1.0, 1500.0] + [cf21B_500] [cf21B_500.core] gravity_vec = [0.0, 0.0, -9.81] @@ -129,3 +150,10 @@ kw = [20000.0, 20000.0, 12000.0] ki_m = [0.0, 0.0, 500.0] kd_omega = [200.0, 200.0, 0.0] int_err_max = [1.0, 1.0, 1500.0] + +[cf21B_500.body_rate2force_torque] +kR = [70000.0, 70000.0, 60000.0] +kw = [20000.0, 20000.0, 12000.0] +ki_m = [0.0, 0.0, 500.0] +kd_omega = [200.0, 200.0, 0.0] +int_err_max = [1.0, 1.0, 1500.0] diff --git a/crazyflow/sim/data.py b/crazyflow/sim/data.py index fe8d96d0..b66ec24e 100644 --- a/crazyflow/sim/data.py +++ b/crazyflow/sim/data.py @@ -11,6 +11,7 @@ from crazyflow.control import Control from crazyflow.control.mellinger import ( MellingerAttitudeData, + MellingerBodyRateData, MellingerForceTorqueData, MellingerStateData, ) @@ -111,6 +112,8 @@ class SimControls: """State control data.""" attitude: ControlData | None """Attitude control data.""" + body_rate: ControlData | None + """Body rate control data.""" force_torque: ControlData | None """Force and torque control data.""" rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 4) @@ -124,6 +127,7 @@ def create( drone: str, state_freq: int | None, attitude_freq: int | None, + body_rate_freq: int | None, force_torque_freq: int | None, device: Device, ) -> SimControls: @@ -142,11 +146,12 @@ def create( mode=control, state=state, attitude=attitude, + body_rate=None, force_torque=force_torque, rotor_vel=rotor_vel, ) case Control.attitude: - attitude = attitude = MellingerAttitudeData.create( + attitude = MellingerAttitudeData.create( n_worlds, n_drones, attitude_freq, drone, device ) force_torque = MellingerForceTorqueData.create( @@ -156,6 +161,22 @@ def create( mode=control, state=None, attitude=attitude, + body_rate=None, + force_torque=force_torque, + rotor_vel=rotor_vel, + ) + case Control.body_rate: + body_rate = MellingerBodyRateData.create( + n_worlds, n_drones, body_rate_freq, drone, device + ) + force_torque = MellingerForceTorqueData.create( + n_worlds, n_drones, force_torque_freq, drone, device + ) + return SimControls( + mode=control, + state=None, + attitude=None, + body_rate=body_rate, force_torque=force_torque, rotor_vel=rotor_vel, ) @@ -167,12 +188,18 @@ def create( mode=control, state=None, attitude=None, + body_rate=None, force_torque=force_torque, rotor_vel=rotor_vel, ) case Control.rotor_vel: return SimControls( - mode=control, state=None, attitude=None, force_torque=None, rotor_vel=rotor_vel + mode=control, + state=None, + attitude=None, + body_rate=None, + force_torque=None, + rotor_vel=rotor_vel, ) case _: raise ValueError(f"Control mode {control} not implemented") diff --git a/crazyflow/sim/functional.py b/crazyflow/sim/functional.py index e69a769a..7a592f2b 100644 --- a/crazyflow/sim/functional.py +++ b/crazyflow/sim/functional.py @@ -42,6 +42,18 @@ def attitude_control(data: SimData, controls: Array) -> SimData: return data +def body_rate_control(data: SimData, controls: Array) -> SimData: + """Body rate control function.""" + assert data.controls.mode == Control.body_rate, f"control type {data.controls.mode} not enabled" + assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch" + controls = jnp.asarray(controls) + return data.replace( + controls=data.controls.replace( + body_rate=data.controls.body_rate.replace(staged_cmd=controls) + ) + ) + + def force_torque_control(data: SimData, controls: Array) -> SimData: """Force-torque control function.""" assert data.controls.mode == Control.force_torque, ( @@ -76,6 +88,8 @@ def controllable(data: SimData) -> Array: control_steps, control_freq = controls.state.steps, controls.state.freq case Control.attitude: control_steps, control_freq = controls.attitude.steps, controls.attitude.freq + case Control.body_rate: + control_steps, control_freq = controls.body_rate.steps, controls.body_rate.freq case Control.force_torque: control_steps = controls.force_torque.steps control_freq = controls.force_torque.freq diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 70705ccd..055f8a54 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -19,6 +19,7 @@ from crazyflow.control import Control from crazyflow.control.mellinger import ( control_attitude2force_torque, + control_body_rate2force_torque, control_commit_attitude, control_force_torque2rotor_vel, control_state2attitude, @@ -82,6 +83,7 @@ def __init__( freq: int = 500, state_freq: int = 100, attitude_freq: int = 500, + body_rate_freq: int = 500, force_torque_freq: int = 500, device: str = "cpu", xml_path: Path | None = None, @@ -100,6 +102,7 @@ def __init__( freq: Dynamics step frequency in Hz. state_freq: Frequency in Hz at which the state controller runs. attitude_freq: Frequency in Hz at which the attitude controller runs. + body_rate_freq: Frequency in Hz at which the body rate controller runs. force_torque_freq: Frequency in Hz at which the force/torque controller runs. device: Device to place the simulation data on (e.g. ``"cpu"`` or ``"gpu"``). xml_path: Path to a custom scene XML. Defaults to ``crazyflow/scene.xml``. @@ -111,7 +114,7 @@ def __init__( assert Dynamics(dynamics) in Dynamics, f"Dynamics mode {dynamics} not implemented" assert Control(control) in Control, f"Control mode {control} not implemented" if dynamics != Dynamics.first_principles: - if control in (Control.force_torque, Control.rotor_vel): + if control in (Control.body_rate, Control.force_torque, Control.rotor_vel): raise ConfigError(f"Control mode {control} requires first principles dynamics") if freq > 10_000 and not jax.config.jax_enable_x64: raise ConfigError("High frequency simulations require double precision mode") @@ -133,7 +136,9 @@ def __init__( self.mj_model, self.mj_data, self.mjx_model, self.mjx_data = self.build_mjx_model(self.spec) self.viewer: MujocoRenderer | None = None - self.data = self.init_data(state_freq, attitude_freq, force_torque_freq, rng_key) + self.data = self.init_data( + state_freq, attitude_freq, body_rate_freq, force_torque_freq, rng_key + ) self.default_data: SimData = self.build_default_data() # Build the simulation pipeline and overwrite the default _step implementation with it @@ -185,6 +190,10 @@ def attitude_control(self, controls: Array): """Set the desired attitude for all drones in all worlds.""" self.data = F.attitude_control(self.data, controls) + def body_rate_control(self, controls: Array): + """Set the desired body rates and collective thrust for all drones in all worlds.""" + self.data = F.body_rate_control(self.data, controls) + def force_torque_control(self, controls: Array): """Set the desired force and torque for all drones in all worlds.""" self.data = F.force_torque_control(self.data, controls) @@ -433,9 +442,10 @@ def build_data(self) -> SimData: """ state_freq = 0 if (s := self.data.controls.state) is None else s.freq attitude_freq = 0 if (a := self.data.controls.attitude) is None else a.freq + body_rate_freq = 0 if (br := self.data.controls.body_rate) is None else br.freq force_torque_freq = 0 if (ft := self.data.controls.force_torque) is None else ft.freq self.data = self.init_data( - state_freq, attitude_freq, force_torque_freq, self.data.core.rng_key + state_freq, attitude_freq, body_rate_freq, force_torque_freq, self.data.core.rng_key ) return self.data @@ -475,7 +485,12 @@ def build_mjx(self): self.mj_model, self.mj_data, self.mjx_model, self.mjx_data = self.build_mjx_model(self.spec) def init_data( - self, state_freq: int, attitude_freq: int, force_torque_freq: int, rng_key: Array + self, + state_freq: int, + attitude_freq: int, + body_rate_freq: int, + force_torque_freq: int, + rng_key: Array, ) -> SimData: """Initialize the simulation data.""" drone_name = "drone_fused" if self.fused_mjx_model else "drone" @@ -493,6 +508,7 @@ def init_data( self.drone, state_freq, attitude_freq, + body_rate_freq, force_torque_freq, self.device, ), @@ -515,6 +531,8 @@ def control_freq(self) -> int: return self.data.controls.state.freq if self.control == Control.attitude: return self.data.controls.attitude.freq + if self.control == Control.body_rate: + return self.data.controls.body_rate.freq if self.control == Control.force_torque: return self.data.controls.force_torque.freq raise NotImplementedError(f"Control mode {self.control} not implemented") @@ -567,6 +585,7 @@ def build_control_fns( """ state = ("state_controller", control_state2attitude) attitude = ("attitude_controller", control_attitude2force_torque) + body_rate = ("body_rate_controller", control_body_rate2force_torque) force_torque = ("force_torque_controller", control_force_torque2rotor_vel) commit_attitude = ("commit_attitude", control_commit_attitude) match control: @@ -581,6 +600,8 @@ def build_control_fns( stages = (commit_attitude,) else: raise NotImplementedError(f"Control mode {control} not implemented for {dynamics}") + case Control.body_rate: + stages = (body_rate, force_torque) case Control.force_torque: stages = (force_torque,) case Control.rotor_vel: diff --git a/docs/examples/index.md b/docs/examples/index.md index 97a133cb..20036a77 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -30,6 +30,17 @@ Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses --- +## Body rate control + +Commanding body-frame angular rates and collective thrust. The firmware controller has no dedicated body rate mode and levels the drone with its attitude terms, so the example sets the `kR` and `ki_m` gains of the body rate controller to zero. + + +```{ .python notest } +--8<-- "examples/control/body_rate.py" +``` + +--- + ## Sampling-based MPC A sampling-based model predictive controller tracks a Lissajous curve while avoiding a grid of obstacles. It rolls out thousands of candidate control sequences in parallel using identified dynamics, then applies the first action from a cost-weighted update of the best samples. The controller automatically uses a GPU when one is available and lowers the sample count on CPU. diff --git a/docs/user-guide/control/controllers.md b/docs/user-guide/control/controllers.md index 6c04330d..2dd8891b 100644 --- a/docs/user-guide/control/controllers.md +++ b/docs/user-guide/control/controllers.md @@ -16,11 +16,13 @@ The Mellinger controller [[1]](#references) is split into three stages that form | 2 | [`attitude2force_torque`](mellinger.md#attitude-to-force-torque) | Attitude + RPYT command | Collective force, body torques + angular velocity integral error | | 3 | [`force_torque2rotor_vel`](mellinger.md#force-torque-to-rotor-velocities) | Force + torques | 4 motor speeds [RPM] | +[`body_rate2force_torque`](mellinger.md#body-rate-to-force-torque) replaces stage 2 when the command is a body rate setpoint instead of an attitude. It runs the same controller with the rate setpoint in the angular velocity error and a level attitude setpoint. + ## Available controllers | Module | Controller | Stages | |---|---|---| -| `crazyflow.control.mellinger` | Mellinger | `state2attitude`, `attitude2force_torque`, `force_torque2rotor_vel` | +| `crazyflow.control.mellinger` | Mellinger | `state2attitude`, `attitude2force_torque`, `body_rate2force_torque`, `force_torque2rotor_vel` | ## References diff --git a/docs/user-guide/control/index.md b/docs/user-guide/control/index.md index 6da257b5..5fafdbd8 100644 --- a/docs/user-guide/control/index.md +++ b/docs/user-guide/control/index.md @@ -1,22 +1,22 @@ # Control Modes -Crazyflow provides four levels of control abstraction, from high-level position setpoints down to direct motor commands. Each level is a separate control mode selected at construction time. +Crazyflow provides multiple control modes, from high-level position setpoints down to direct motor commands. Each mode is selected at construction time. ## Control hierarchy -Commands flow down a hierarchy. A state command is converted to an attitude command by the Mellinger controller; an attitude command is converted to force/torque by the geometric controller; force/torque is converted to rotor velocities by the mixer. +Commands flow down a hierarchy. A state command is converted to an attitude command by the Mellinger controller; an attitude command is converted to force/torque by the geometric controller; force/torque is converted to rotor velocities by the mixer. Body rate control feeds the geometric controller with a rate setpoint instead of an attitude, so it enters the hierarchy at the same level as attitude control. ``` State (13D) └─ Mellinger controller - └─ Attitude (4D: roll, pitch, yaw, thrust) + └─ Attitude (4D: roll, pitch, yaw, thrust) | Body rates (4D: ωx, ωy, ωz, thrust) └─ Geometric controller └─ Force/torque (4D: Fc, Tx, Ty, Tz) └─ Mixer └─ Rotor velocities (4D: ω₁…ω₄) ``` -When you select `Control.state`, the full chain runs on every control tick. When you select `Control.attitude`, only the lower two stages run. +When you select `Control.state`, the full chain runs on every control tick. When you select `Control.attitude` or `Control.body_rate`, only the lower two stages run. ## State control @@ -94,6 +94,45 @@ sim.attitude_control(cmd) sim.step(sim.freq // sim.control_freq) ``` +## Body rate control + +Commands body-frame angular rates and a collective thrust. The Mellinger controller tracks the rates with the same gains as in attitude control. As in the firmware, its attitude terms level the drone at the current yaw. Set the `kR` and `ki_m` parameters of the body rate controller to zero to track body rates without the levelling terms, see the [body rate example](../../examples/index.md#body-rate-control). Requires `Dynamics.first_principles`. + +```python +from crazyflow.sim import Sim, Dynamics +from crazyflow.control import Control + +sim = Sim(control=Control.body_rate, dynamics=Dynamics.first_principles, body_rate_freq=500) +sim.reset() +``` + +Command shape: `(n_worlds, n_drones, 4)` + +| Index | Variable | Units | +|---|---|---| +| 0 | Roll rate \(\omega_x\) | rad/s | +| 1 | Pitch rate \(\omega_y\) | rad/s | +| 2 | Yaw rate \(\omega_z\) | rad/s | +| 3 | Collective thrust | N | + +Zero rates and hover thrust hold the current attitude: + +```python +import numpy as np +from crazyflow.sim import Sim, Dynamics +from crazyflow.control import Control + +sim = Sim(control=Control.body_rate, dynamics=Dynamics.first_principles) +sim.reset() + +mass = float(sim.data.params.mass[0]) +cmd = np.zeros((1, 1, 4), dtype=np.float32) +cmd[0, 0, 3] = mass * 9.81 + +sim.body_rate_control(cmd) +sim.step(sim.freq // sim.control_freq) +``` + ## Force-torque control Direct force and torque input. Requires `Dynamics.first_principles`. @@ -157,6 +196,7 @@ Each control mode has its own update rate. The dynamics tick (`freq`) is always |---|---|---| | `state` | `state_freq` | 100 Hz | | `attitude` | `attitude_freq` | 500 Hz | +| `body_rate` | `body_rate_freq` | 500 Hz | | `force_torque` | `force_torque_freq` | 500 Hz | | `rotor_vel` | — | every dynamics step | @@ -167,7 +207,7 @@ The simulator applies a new command only when the control tick fires. Between ti The control modes above are how the simulator drives the onboard controllers. Those controllers also live in `crazyflow.control` as a self-contained library of pure functions, usable on their own for control design, learning-based policies, or as a reference implementation, independent of `Sim`. The following guides cover that standalone API: - [Controllers](controllers.md): the controller interface and the Mellinger pipeline -- [Mellinger controller](mellinger.md): the three stages, their inputs and outputs +- [Mellinger controller](mellinger.md): the three stages and the body rate variant, their inputs and outputs - [Parametrization](parametrize.md): binding a controller to a drone configuration - [Integral errors](integral-errors.md): carrying controller state across calls - [Batching](batching.md): evaluating many drones at once diff --git a/docs/user-guide/control/integral-errors.md b/docs/user-guide/control/integral-errors.md index 2f6223e0..6280cdde 100644 --- a/docs/user-guide/control/integral-errors.md +++ b/docs/user-guide/control/integral-errors.md @@ -54,7 +54,7 @@ for _ in range(10): ## Both stages have integral errors -`state2attitude` tracks position error via `pos_err_i`. `attitude2force_torque` tracks angular velocity error via `r_int_error`. Manage them independently: +`state2attitude` tracks position error via `pos_err_i`. `attitude2force_torque` and `body_rate2force_torque` track angular velocity error via `r_int_error`. Manage them independently: ```python import numpy as np diff --git a/docs/user-guide/control/mellinger.md b/docs/user-guide/control/mellinger.md index 07136e5b..a9386370 100644 --- a/docs/user-guide/control/mellinger.md +++ b/docs/user-guide/control/mellinger.md @@ -1,6 +1,6 @@ # Mellinger controller -The Mellinger controller converts a full-state setpoint into individual motor speeds through three chained pure functions. The implementation closely follows the Crazyflie firmware to minimise sim-to-real gap. +The Mellinger controller converts a full-state setpoint into individual motor speeds through three chained pure functions. The implementation closely follows the Crazyflie firmware to minimise sim-to-real gap. A fourth function, `body_rate2force_torque`, replaces the second stage for body rate setpoints. ## State representation @@ -91,6 +91,47 @@ force.shape # (1,) torque.shape # (3,) ``` +## Stage 2b: Body rates to force/torque {#body-rate-to-force-torque} + +`body_rate2force_torque` replaces stage 2 when the command is a body rate setpoint. The firmware has no dedicated body rate mode. Instead, a rate setpoint enters the angular velocity error and its derivative, while the attitude terms level the drone at its current yaw. Our implementation reproduces this behaviour with the same gains as `attitude2force_torque`. To track body rates without the levelling terms, set `kR` and `ki_m` to zero. + +**Inputs:** + +| Argument | Shape | Description | +|---|---|---| +| `quat` | `(..., 4)` | Current attitude, xyzw | +| `ang_vel` | `(..., 3)` | Current angular velocity in body frame [rad/s] | +| `cmd` | `(..., 4)` | Body rate command: `[roll_rate, pitch_rate, yaw_rate, thrust_N]` | +| `prev_ang_vel` | `(..., 3)` or `None` | Angular velocity from the previous call. `None` initialises to zero | +| `prev_cmd` | `(..., 4)` or `None` | Command from the previous call, used for the setpoint derivative. `None` assumes a constant setpoint | +| `r_int_error` | `(..., 3)` or `None` | Angular velocity integral error from the previous call. `None` initialises to zero | +| `ctrl_freq` | `int` | Control frequency in Hz (default 500) | + +**Outputs:** + +| Return | Shape | Description | +|---|---|---| +| `force` | `(..., 1)` | Collective thrust [N] | +| `torque` | `(..., 3)` | Body-frame torques [N·m] | +| `r_int_error` | `(..., 3)` | Angular velocity integral error. Pass back as `r_int_error` on the next call | + +```python +import numpy as np +from crazyflow.control import load_params +from crazyflow.control.mellinger import body_rate2force_torque + +params = load_params(body_rate2force_torque, "cf2x_L250") +params["kR"], params["ki_m"] = np.zeros(3), np.zeros(3) # pure body rate tracking + +quat = np.array([0.0, 0.0, 0.0, 1.0]) +ang_vel = np.zeros(3) +cmd = np.array([0.5, 0.0, 0.0, 0.3]) # 0.5 rad/s roll rate, 0.3 N thrust + +force, torque, r_int_err = body_rate2force_torque(quat, ang_vel, cmd, **params) +force.shape # (1,) +torque.shape # (3,) +``` + ## Stage 3: Force/torque to rotor velocities {#force-torque-to-rotor-velocities} `force_torque2rotor_vel` converts collective thrust and body-frame torques into individual motor speeds, accounting for the motor mixing matrix. diff --git a/docs/user-guide/dynamics/index.md b/docs/user-guide/dynamics/index.md index f142965a..5f98feff 100644 --- a/docs/user-guide/dynamics/index.md +++ b/docs/user-guide/dynamics/index.md @@ -27,7 +27,7 @@ The first-principles dynamics derives forces and torques analytically from motor from crazyflow.sim import Sim, Dynamics from crazyflow.control import Control -# Force-torque and rotor_vel control modes require first_principles +# Body rate, force-torque and rotor_vel control modes require first_principles sim = Sim(dynamics=Dynamics.first_principles, control=Control.rotor_vel) sim.reset() ``` @@ -69,15 +69,15 @@ The `so_rpy_rotor_drag` variant includes translational drag, which captures the ## Control mode compatibility -| Dynamics | `Control.state` | `Control.attitude` | `Control.force_torque` | `Control.rotor_vel` | -|---|---|---|---|---| -| `first_principles` | ✓ | ✓ | ✓ | ✓ | -| `so_rpy` | ✓ | ✓ | ✗ | ✗ | -| `so_rpy_rotor` | ✓ | ✓ | ✗ | ✗ | -| `so_rpy_rotor_drag` | ✓ | ✓ | ✗ | ✗ | +| Dynamics | `Control.state` | `Control.attitude` | `Control.body_rate` | `Control.force_torque` | `Control.rotor_vel` | +|---|---|---|---|---|---| +| `first_principles` | ✓ | ✓ | ✓ | ✓ | ✓ | +| `so_rpy` | ✓ | ✓ | ✗ | ✗ | ✗ | +| `so_rpy_rotor` | ✓ | ✓ | ✗ | ✗ | ✗ | +| `so_rpy_rotor_drag` | ✓ | ✓ | ✗ | ✗ | ✗ | !!! warning - Using `Control.force_torque` or `Control.rotor_vel` with a fitted dynamics raises `ConfigError` at construction time. + Using `Control.body_rate`, `Control.force_torque` or `Control.rotor_vel` with a fitted dynamics raises `ConfigError` at construction time. ## Using the dynamics standalone diff --git a/docs/user-guide/functional-api.md b/docs/user-guide/functional-api.md index a491649e..52114f1d 100644 --- a/docs/user-guide/functional-api.md +++ b/docs/user-guide/functional-api.md @@ -58,7 +58,7 @@ From this point, `data` is a plain JAX pytree and `step` and `reset` are compile ## Purely functional controller functions -`crazyflow.sim.functional` mirrors all four `Sim` control methods as pure functions: +`crazyflow.sim.functional` mirrors all `Sim` control methods as pure functions: ```python import crazyflow.sim.functional as F @@ -68,6 +68,7 @@ import crazyflow.sim.functional as F |---|---| | `F.state_control(data, controls)` | Stage a state command | | `F.attitude_control(data, controls)` | Stage an attitude command | +| `F.body_rate_control(data, controls)` | Stage a body rate command | | `F.force_torque_control(data, controls)` | Stage a force/torque command | | `F.rotor_vel_control(data, controls)` | Stage rotor velocity commands | | `F.controllable(data)` | Boolean mask — which worlds may update their controller this step | diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index 99fbe308..a9ce0692 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -6,7 +6,7 @@ In-depth documentation for every part of the simulator. - [Object-Oriented API](oo-api.md) — `Sim` class, control methods, rendering, and reset - [Functional API](functional-api.md) — purely functional interface for JAX transformations - [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 +- [Control Modes](control/index.md) — state, attitude, body rate, 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 diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index a22574d2..277ae22c 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -77,7 +77,7 @@ import numpy as np from crazyflow.sim import Sim, Dynamics from crazyflow.control import Control -sim = Sim(n_worlds=1, n_drones=1, control=Control.attitude, dynamics=Dynamics.so_rpy) +sim = Sim(n_worlds=1, n_drones=1, control=Control.attitude) sim.reset() # [roll, pitch, yaw, collective_thrust_N] @@ -88,6 +88,26 @@ sim.attitude_control(cmd) sim.step(sim.freq // sim.control_freq) ``` +### Body rate control + +Commands body-frame angular rates (rad/s) and a collective thrust (N). The Mellinger controller tracks the rates and, as in the firmware, levels the drone with its attitude terms. Set `kR` and `ki_m` to zero for pure rate tracking, see [Control Modes](control/index.md#body-rate-control). Requires `Dynamics.first_principles`. + +```python +import numpy as np +from crazyflow.sim import Sim, Dynamics +from crazyflow.control import Control + +sim = Sim(n_worlds=1, n_drones=1, control=Control.body_rate) +sim.reset() + +# [body_rate_x, body_rate_y, body_rate_z, collective_thrust_N] +cmd = np.zeros((1, 1, 4), dtype=np.float32) +cmd[0, 0, 3] = float(sim.data.params.mass[0]) * 9.81 # hover thrust + +sim.body_rate_control(cmd) +sim.step(sim.freq // sim.control_freq) +``` + ### Force-torque control Direct force and torque input, useful for testing dynamics or custom controllers. Requires `Dynamics.first_principles`. diff --git a/examples/control/attitude.py b/examples/control/attitude.py index 66c7fc37..3b4ed69b 100644 --- a/examples/control/attitude.py +++ b/examples/control/attitude.py @@ -1,58 +1,20 @@ -import os +from functools import partial import numpy as np -os.environ["SCIPY_ARRAY_API"] = "1" - -from scipy.spatial.transform import Rotation as R - -from crazyflow.control import Control +from crazyflow.control import Control, parametrize +from crazyflow.control.mellinger import state2attitude from crazyflow.sim import Sim -kp = np.array([0.4, 0.4, 1.25]) -ki = np.array([0.05, 0.05, 0.05]) -kd = np.array([0.2, 0.2, 0.4]) -g = 9.81 - - -def control( - t: float, obs: dict[str, np.ndarray], pos_start: np.ndarray, drone_mass: float -) -> np.ndarray: - des_pos = np.zeros(3) - des_pos[..., :2] = pos_start[:2] + np.array([np.cos(t) - 1, np.sin(t)]) - des_pos[..., 2] = 0.2 * t - des_vel = np.zeros_like(des_pos) - des_yaw = t - - # Calculate the deviations from the desired trajectory - pos_error = des_pos - np.array(obs["pos"]) - vel_error = des_vel - np.array(obs["vel"]) - - # Compute target thrust - target_thrust = np.zeros(3) - target_thrust += kp * pos_error - target_thrust += kd * vel_error - target_thrust[2] += drone_mass * g - - # Update z_axis to the current orientation of the drone - z_axis = R.from_quat(obs["quat"]).as_matrix()[:, 2] - - # update current thrust - thrust_desired = target_thrust.dot(z_axis) - - # update z_axis_desired - z_axis_desired = target_thrust / np.linalg.norm(target_thrust) - x_c_des = np.array([np.cos(des_yaw), np.sin(des_yaw), 0.0]) - y_axis_desired = np.cross(z_axis_desired, x_c_des) - y_axis_desired /= np.linalg.norm(y_axis_desired) - x_axis_desired = np.cross(y_axis_desired, z_axis_desired) - - R_desired = np.vstack([x_axis_desired, y_axis_desired, z_axis_desired]).T - euler_desired = R.from_matrix(R_desired).as_euler("xyz", degrees=False) - - action = np.concatenate([euler_desired, [thrust_desired]], dtype=np.float32) - return action +def control(t: float, pos_start: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Compute the attitude command to track a circle with a slow climb.""" + cmd = np.zeros(13) + cmd[:3] = pos_start + np.array([np.cos(t) - 1, np.sin(t), 0.2 * t]) + cmd[3:6] = np.array([-np.sin(t), np.cos(t), 0.2]) + cmd[6:9] = np.array([-np.cos(t), -np.sin(t), 0.0]) + cmd[9] = t # Yaw + return cmd def main(): @@ -61,15 +23,17 @@ def main(): duration = 6.5 fps = 60 + # We use the Mellinger position controller to generate attitude commands. This could be any + # controller that outputs [roll, pitch, yaw, thrust], e.g. a learned policy. + position_ctrl = partial(parametrize(state2attitude, sim.drone), ctrl_freq=sim.control_freq) + pos_err_i = np.zeros(3) cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) # [roll, pitch, yaw, thrust] - pos_start = sim.data.states.pos + pos_start = np.asarray(sim.data.states.pos[0, 0]) for i in range(int(duration * sim.control_freq)): - obs = { - "pos": sim.data.states.pos[0, 0], - "vel": sim.data.states.vel[0, 0], - "quat": sim.data.states.quat[0, 0], - } - cmd[0, 0, :] = control(i / sim.control_freq, obs, pos_start[0, 0], sim.data.params.mass[0]) + pos, quat = np.asarray(sim.data.states.pos[0, 0]), np.asarray(sim.data.states.quat[0, 0]) + vel = np.asarray(sim.data.states.vel[0, 0]) + ref = control(i / sim.control_freq, pos_start) + cmd[0, 0, :], pos_err_i = position_ctrl(pos, quat, vel, ref, pos_err_i) sim.attitude_control(cmd) sim.step(sim.freq // sim.control_freq) if ((i * fps) % sim.control_freq) < fps: diff --git a/examples/control/body_rate.py b/examples/control/body_rate.py new file mode 100644 index 00000000..f979c099 --- /dev/null +++ b/examples/control/body_rate.py @@ -0,0 +1,67 @@ +import os + +os.environ["SCIPY_ARRAY_API"] = "1" + +from functools import partial + +import jax.numpy as jnp +import numpy as np +from scipy.spatial.transform import Rotation as R + +from crazyflow.control import Control, parametrize +from crazyflow.control.mellinger import state2attitude +from crazyflow.sim import Sim + +kp_att = 8.0 # Proportional gain from the attitude error to body rates + + +def trajectory(t: float, pos_start: np.ndarray) -> np.ndarray: + """Compute the full state command of a circle with a slow climb.""" + cmd = np.zeros(13) + cmd[:3] = pos_start + np.array([np.cos(t) - 1, np.sin(t), 0.2 * t]) + cmd[3:6] = np.array([-np.sin(t), np.cos(t), 0.2]) + cmd[6:9] = np.array([-np.cos(t), -np.sin(t), 0.0]) + cmd[9] = t # Yaw + return cmd + + +def control(quat: np.ndarray, rpyt: np.ndarray) -> np.ndarray: + """Convert an attitude command into body rates with a proportional attitude loop.""" + rot_err = (R.from_quat(quat).inv() * R.from_euler("xyz", rpyt[:3])).as_rotvec() + return np.concatenate([kp_att * rot_err, rpyt[3:]]) + + +def main(): + sim = Sim(control=Control.body_rate, body_rate_freq=250) + # The firmware has no dedicated body rate mode. Its attitude terms level the drone at the + # current yaw and would counteract the commanded rates. Disable them to track body rates. + body_rate = sim.data.controls.body_rate + params = body_rate.params | {"kR": jnp.zeros(3), "ki_m": jnp.zeros(3)} + controls = sim.data.controls.replace(body_rate=body_rate.replace(params=params)) + sim.data = sim.data.replace(controls=controls) + sim.build_default_data() + sim.reset() + duration = 6.5 + fps = 60 + + # We use the Mellinger position controller to generate attitude commands, which we then convert + # to body rates. This could be any controller that outputs [w_x, w_y, w_z, thrust]. + position_ctrl = partial(parametrize(state2attitude, sim.drone), ctrl_freq=sim.control_freq) + pos_err_i = np.zeros(3) + cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) # [roll_rate, pitch_rate, yaw_rate, thrust] + pos_start = np.asarray(sim.data.states.pos[0, 0]) + for i in range(int(duration * sim.control_freq)): + pos, quat = np.asarray(sim.data.states.pos[0, 0]), np.asarray(sim.data.states.quat[0, 0]) + vel = np.asarray(sim.data.states.vel[0, 0]) + ref = trajectory(i / sim.control_freq, pos_start) + rpyt, pos_err_i = position_ctrl(pos, quat, vel, ref, pos_err_i) + cmd[0, 0, :] = control(quat, rpyt) + sim.body_rate_control(cmd) + sim.step(sim.freq // sim.control_freq) + if ((i * fps) % sim.control_freq) < fps: + sim.render() + sim.close() + + +if __name__ == "__main__": + main() diff --git a/pixi.lock b/pixi.lock index f2a42e28..cf14b19a 100644 --- a/pixi.lock +++ b/pixi.lock @@ -139,13 +139,12 @@ environments: - pypi: https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl @@ -183,6 +182,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fa/de/ceae2adf7034e07e9910299fe412e1819c4f0dd520700a888bcb03625448/pandas-3.0.5-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl default: channels: - url: https://conda.anaconda.org/conda-forge/ @@ -279,13 +279,12 @@ environments: - pypi: https://files.pythonhosted.org/packages/55/15/5709e379c1191b1d354e74a04cb01d61abe29ce4bba29833229525882566/jax_cuda12_pjrt-0.11.1-py3-none-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/86/16/1a8fd2b19544b84575cf84ef7aa3ad4c173b756d5f087c91f85d1b295777/array_api_compat-1.15.0-py3-none-any.whl @@ -317,6 +316,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl osx-arm64: - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.3-pyhc364b38_0.conda @@ -386,9 +386,10 @@ environments: - pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/84/6a/441eb053b078954f7fea284dfb288701884d0a1404d39babb858e1649023/ml_dtypes-0.6.0-cp312-cp312-macosx_10_13_universal2.whl @@ -409,7 +410,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/eb/a3/b69efbf4143b5b9859b977770bbbabcc2796b702fa69dc40271e45cd5a56/prometheus_client-0.26.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fa/3c/74050e09a490c0cfed6016430226710d29fe4e165e12145368c7de4c6426/tensorstore-0.1.85-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl dist: channels: - url: https://conda.anaconda.org/conda-forge/ @@ -513,14 +513,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5b/67/314f6151bad77a93d751978a344033e1fc890822f05f0416079338e34231/nh3-0.3.6-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl @@ -562,6 +561,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl osx-arm64: - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda @@ -648,9 +648,10 @@ environments: - pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl @@ -679,7 +680,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/f3/ab/a7653bce9a3b204be6a6931767a9e23595807bb84790ce6685e4d7e5bd08/nh3-0.3.6-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl - pypi: https://files.pythonhosted.org/packages/fa/3c/74050e09a490c0cfed6016430226710d29fe4e165e12145368c7de4c6426/tensorstore-0.1.85-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl docs: channels: @@ -803,14 +803,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/55/15/5709e379c1191b1d354e74a04cb01d61abe29ce4bba29833229525882566/jax_cuda12_pjrt-0.11.1-py3-none-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5d/5b/4c1902e8bdd5c4db63284e9d101dece4038d4025d6d88850ffe0a1578980/mkdocstrings-1.0.6-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/86/16/1a8fd2b19544b84575cf84ef7aa3ad4c173b756d5f087c91f85d1b295777/array_api_compat-1.15.0-py3-none-any.whl @@ -846,6 +845,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl osx-arm64: - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.18.0-pyhcf101f3_1.conda - conda: https://conda.anaconda.org/conda-forge/noarch/backrefs-8.0-pyhcf101f3_0.conda @@ -952,10 +952,11 @@ environments: - pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5d/5b/4c1902e8bdd5c4db63284e9d101dece4038d4025d6d88850ffe0a1578980/mkdocstrings-1.0.6-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/84/6a/441eb053b078954f7fea284dfb288701884d0a1404d39babb858e1649023/ml_dtypes-0.6.0-cp312-cp312-macosx_10_13_universal2.whl @@ -980,7 +981,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/f6/b6/f65ac785d4ac90dcf7c831ac6256f5dd4a19780f4e1575b2c0d6eeebe319/griffelib-2.2.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fa/3c/74050e09a490c0cfed6016430226710d29fe4e165e12145368c7de4c6426/tensorstore-0.1.85-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl gpu: channels: - url: https://conda.anaconda.org/conda-forge/ @@ -1066,13 +1066,12 @@ environments: - pypi: https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl @@ -1108,6 +1107,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl gpu-tests: channels: - url: https://conda.anaconda.org/conda-forge/ @@ -1293,13 +1293,12 @@ environments: - pypi: https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl @@ -1335,6 +1334,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl release: channels: - url: https://conda.anaconda.org/conda-forge/ @@ -1438,14 +1438,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5b/67/314f6151bad77a93d751978a344033e1fc890822f05f0416079338e34231/nh3-0.3.6-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl @@ -1487,6 +1486,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl osx-arm64: - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda @@ -1573,9 +1573,10 @@ environments: - pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl @@ -1604,7 +1605,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/f3/ab/a7653bce9a3b204be6a6931767a9e23595807bb84790ce6685e4d7e5bd08/nh3-0.3.6-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl - pypi: https://files.pythonhosted.org/packages/fa/3c/74050e09a490c0cfed6016430226710d29fe4e165e12145368c7de4c6426/tensorstore-0.1.85-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl tests: channels: @@ -1802,13 +1802,12 @@ environments: - pypi: https://files.pythonhosted.org/packages/55/15/5709e379c1191b1d354e74a04cb01d61abe29ce4bba29833229525882566/jax_cuda12_pjrt-0.11.1-py3-none-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/73/b1/37b6fefb2f2fe5f11ae305d0a428e1cdb4f5868112c42ab29c1ce1629f06/viser-1.1.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/78/91/3635cdb13318cb0a328abaa69e2b91251caad39d6779aa308098f341f6cb/simplejson-4.1.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/85/2c/e42cbf7c1c3ac7905fd4a4c6578267dff4e541b38d28c08357fdae02a176/flax-0.12.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/86/16/1a8fd2b19544b84575cf84ef7aa3ad4c173b756d5f087c91f85d1b295777/array_api_compat-1.15.0-py3-none-any.whl @@ -1840,6 +1839,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl osx-arm64: - conda: https://conda.anaconda.org/conda-forge/noarch/array-api-strict-2.6.1-pyhc364b38_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda @@ -1954,9 +1954,10 @@ environments: - pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/84/6a/441eb053b078954f7fea284dfb288701884d0a1404d39babb858e1649023/ml_dtypes-0.6.0-cp312-cp312-macosx_10_13_universal2.whl @@ -1977,7 +1978,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/eb/a3/b69efbf4143b5b9859b977770bbbabcc2796b702fa69dc40271e45cd5a56/prometheus_client-0.26.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fa/3c/74050e09a490c0cfed6016430226710d29fe4e165e12145368c7de4c6426/tensorstore-0.1.85-cp312-cp312-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl packages: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda build_number: 20 @@ -6021,7 +6021,7 @@ packages: - numpy>=2.0.0 - scipy>=1.17.0 - jax>=0.7.0,!=0.10.2 - - mujoco>=3.3.0 + - mujoco>=3.3.0,<3.11 - mujoco-mjx>=3.3.0 - gymnasium[mujoco]>=1.2.0 - imageio @@ -6737,6 +6737,22 @@ packages: version: 0.3.6 sha256: 905f877dc66dd7aea4a76e54bcb26acb5ff8216f720c0017ccf63e0e6035698e requires_python: '>=3.8' +- pypi: https://files.pythonhosted.org/packages/5c/75/a5195ee4239baf62a5f2e6a40914c55ad78b0b429500ec67117bd41cded7/mujoco_mjx-3.10.0-py3-none-any.whl + name: mujoco-mjx + version: 3.10.0 + sha256: ab7e92ababdbdb1b8202d27e428b266eda103466d8c30a234c34aa9c34a6dfa7 + requires_dist: + - absl-py + - etils[epath] + - jax + - jaxlib + - mujoco>=3.10.0.dev0 + - scipy + - trimesh + - warp-lang==1.13.0 ; extra == 'warp' + - isort ; extra == 'dev' + - pyink ; extra == 'dev' + requires_python: '>=3.10' - pypi: https://files.pythonhosted.org/packages/5d/5b/4c1902e8bdd5c4db63284e9d101dece4038d4025d6d88850ffe0a1578980/mkdocstrings-1.0.6-py3-none-any.whl name: mkdocstrings version: 1.0.6 @@ -6752,17 +6768,10 @@ packages: - mkdocstrings-python-legacy>=0.2.1 ; extra == 'python-legacy' - mkdocstrings-python>=1.16.2 ; extra == 'python' requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-nvshmem-cu12 - version: 3.7.2 - sha256: 02b6d3482d90ea8ba214bc400362297f573257763990a98c9846ec5c786a7227 - requires_dist: - - nvidia-cuda-cccl-cu12 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/67/21/6649c5843586eebf9e2fb762926c615b9c441d251de8ebe2da1ca8bcbda7/mujoco-3.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl +- pypi: https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl name: mujoco - version: 3.12.0 - sha256: 3915a59674f90e490bf7d7e0bfd26e618cbe39c43795aa985c43c63b5ee69592 + version: 3.10.0 + sha256: a4d35e9d0b13ff9ad3196294a7dac363f1d0cdaa988832d0b687d42d98f4ee29 requires_dist: - absl-py - etils[epath] @@ -6782,6 +6791,13 @@ packages: - usd-core ; extra == 'usd' - pillow ; extra == 'usd' requires_python: '>=3.10' +- pypi: https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl + name: nvidia-nvshmem-cu12 + version: 3.7.2 + sha256: 02b6d3482d90ea8ba214bc400362297f573257763990a98c9846ec5c786a7227 + requires_dist: + - nvidia-cuda-cccl-cu12 + requires_python: '>=3' - pypi: https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl name: glfw version: 2.10.2 @@ -6867,22 +6883,6 @@ packages: name: farama-notifications version: 0.0.6 sha256: f84839188efa1ce5bb361c2a84881b2dc2c0d0d7fb661ff00421820170930935 -- pypi: https://files.pythonhosted.org/packages/7d/ea/cf8c48e71a8059f686e447b5c9722a44291ad7913e1593838504f324e6dc/mujoco_mjx-3.12.0-py3-none-any.whl - name: mujoco-mjx - version: 3.12.0 - sha256: afc31b6296a2d5ae5fe20663fd7348ca4fb5bc3159c3ace945f2c5f8224789f1 - requires_dist: - - absl-py - - etils[epath] - - jax - - jaxlib - - mujoco>=3.12.0.dev0 - - scipy - - trimesh - - warp-lang==1.16.0 ; extra == 'warp' - - isort ; extra == 'dev' - - pyink ; extra == 'dev' - requires_python: '>=3.10' - pypi: https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl name: urllib3 version: 2.7.0 @@ -7918,10 +7918,10 @@ packages: - zstandard ; python_full_version < '3.14' and extra == 'test-full' - tqdm ; extra == 'tqdm' requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/fe/eb/c167f660219f72595ea1f89f5bc22b2b9bd43dad6779ce77895d679896bd/mujoco-3.12.0-cp312-cp312-macosx_11_0_arm64.whl +- pypi: https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl name: mujoco - version: 3.12.0 - sha256: 8e9cfd259b2614e858c6d86a2d095c627d11083203ecad7a42286122e671aac9 + version: 3.10.0 + sha256: 04b57f3f9a58b35b99d5966f8be12f68158ad0c7393b219378af0686f3856dbe requires_dist: - absl-py - etils[epath] diff --git a/pyproject.toml b/pyproject.toml index 49d6c614..1d25db8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "numpy>=2.0.0", "scipy>=1.17.0", "jax>=0.7.0,!=0.10.2", # 0.10.2 XLA:GPU miscompiles for larger swarms. See https://github.com/jax-ml/jax/issues/39157 - "mujoco>=3.3.0", + "mujoco>=3.3.0,<3.11", # 3.11 breaks the viewer. Undo when gymnasium 1.4 lands: https://github.com/Farama-Foundation/Gymnasium/pull/1678 "mujoco-mjx>=3.3.0", "gymnasium[mujoco]>=1.2.0", "imageio", diff --git a/tests/integration/test_interfaces.py b/tests/integration/test_interfaces.py index 34b1036a..131da498 100644 --- a/tests/integration/test_interfaces.py +++ b/tests/integration/test_interfaces.py @@ -1,4 +1,5 @@ import jax +import jax.numpy as jnp import numpy as np import pytest from scipy.spatial.transform import Rotation as R @@ -55,6 +56,32 @@ def test_attitude_interface(dynamics: Dynamics): assert distance < 0.05, f"Failed to maintain hover with {dynamics} ({dpos})" +@pytest.mark.integration +def test_body_rate_interface(): + sim = Sim(dynamics=Dynamics.first_principles, control=Control.body_rate) + # Disable the attitude terms of the firmware controller to track body rates directly + body_rate = sim.data.controls.body_rate + params = body_rate.params | {"kR": jnp.zeros(3), "ki_m": jnp.zeros(3)} + controls = sim.data.controls.replace(body_rate=body_rate.replace(params=params)) + # Spawn the drone in the air so that it can roll freely without hitting the ground + states = sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(2.0)) + sim.data = sim.data.replace(controls=controls, states=states) + + ang_vel_des = np.array([2.0, 0.0, 0.0]) # Roll rate command + thrust = sim.data.params.mass[0] * np.linalg.norm(sim.data.params.gravity_vec) + cmd = np.concatenate([ang_vel_des, [thrust]])[None, None, :] + sim.body_rate_control(cmd) + + errors = [] + for i in range(sim.control_freq): + sim.step(sim.freq // sim.control_freq) + if i >= int(0.25 * sim.control_freq): # Give the controller time to settle on the command + errors.append(sim.data.states.ang_vel[0, 0] - ang_vel_des) + + err_max = np.max(np.abs(errors)) + assert err_max < 0.02, f"Failed to track the body rate command (max error {err_max:.2e})" + + @pytest.mark.integration def test_rotor_vel_interface(): sim = Sim(dynamics=Dynamics.first_principles, control=Control.rotor_vel) diff --git a/tests/unit/control/test_core.py b/tests/unit/control/test_core.py index 1fd698cf..4b2f2dbd 100644 --- a/tests/unit/control/test_core.py +++ b/tests/unit/control/test_core.py @@ -9,12 +9,18 @@ from crazyflow.control import load_params, parametrize from crazyflow.control.mellinger import ( attitude2force_torque, + body_rate2force_torque, force_torque2rotor_vel, state2attitude, ) from crazyflow.drones import available_drones -_MELLINGER_FNS = [state2attitude, attitude2force_torque, force_torque2rotor_vel] +_MELLINGER_FNS = [ + state2attitude, + attitude2force_torque, + body_rate2force_torque, + force_torque2rotor_vel, +] @pytest.mark.unit diff --git a/tests/unit/control/test_mellinger.py b/tests/unit/control/test_mellinger.py index 60298948..b4a518db 100644 --- a/tests/unit/control/test_mellinger.py +++ b/tests/unit/control/test_mellinger.py @@ -4,10 +4,12 @@ import numpy as np import pytest +from scipy.spatial.transform import Rotation as R from crazyflow.control import load_params, parametrize from crazyflow.control.mellinger import ( attitude2force_torque, + body_rate2force_torque, force_torque2rotor_vel, state2attitude, ) @@ -59,6 +61,27 @@ def test_attitude2force_torque(drone: str) -> None: assert r_int_error.shape == (5, 4, 3) +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque(drone: str) -> None: + controller = parametrize(body_rate2force_torque, drone) + # Single input + _, quat, _, ang_vel = create_rnd_states() + cmd = np.array([0.1, 0.1, 0.1, 1.0]) # roll rate, pitch rate, yaw rate, thrust command + force_des, torque_des, r_int_error = controller(quat, ang_vel, cmd) + assert force_des.shape == (1,) + assert torque_des.shape == (3,) + assert r_int_error.shape == (3,) + # Batch input + _, quat, _, ang_vel = create_rnd_states((5, 4)) + cmd = np.random.randn(5, 4, 4) + cmd[..., 3] = np.abs(cmd[..., 3]) # Ensure positive thrust + force_des, torque_des, r_int_error = controller(quat, ang_vel, cmd) + assert force_des.shape == (5, 4, 1) + assert torque_des.shape == (5, 4, 3) + assert r_int_error.shape == (5, 4, 3) + + @pytest.mark.unit @pytest.mark.parametrize("drone", available_drones) def test_force_torque2rotor_vel(drone: str) -> None: @@ -147,6 +170,88 @@ def test_attitude2force_torque_zero_thrust(drone: str): assert np.allclose(torque_des, 0.0, atol=1e-6) +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_at_setpoint(drone: str) -> None: + # Level drone with measured rates equal to the commanded rates → zero corrective torque. + controller = parametrize(body_rate2force_torque, drone) + quat = R.from_euler("xyz", [0.0, 0.0, 0.7]).as_quat() # Any yaw is level + ang_vel = np.array([0.3, -0.2, 0.1]) + cmd = np.array([0.3, -0.2, 0.1, 0.5]) + force_des, torque_des, _ = controller(quat, ang_vel, cmd, prev_ang_vel=ang_vel) + assert np.allclose(torque_des, 0.0, atol=1e-6), ( + f"Torque at setpoint should be ~0, got {torque_des}" + ) + assert force_des[0] > 0.0, "Force must be positive for positive thrust command" + + +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_zero_thrust(drone: str): + # Zero thrust command → firmware zeros torque; outputs are all zero. + controller = parametrize(body_rate2force_torque, drone) + quat = np.array([0.0, 0.0, 0.0, 1.0]) + ang_vel = np.zeros(3) + cmd = np.array([0.1, 0.1, 0.1, 0.0]) # non-zero rates but zero thrust + force_des, torque_des, _ = controller(quat, ang_vel, cmd) + assert np.allclose(force_des, 0.0, atol=1e-6) + assert np.allclose(torque_des, 0.0, atol=1e-6) + + +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_sign(drone: str): + # A positive rate error about one axis must produce a positive torque about that axis only. + controller = parametrize(body_rate2force_torque, drone) + quat = np.array([0.0, 0.0, 0.0, 1.0]) + ang_vel = np.zeros(3) + for axis in range(3): + cmd = np.array([0.0, 0.0, 0.0, 0.5]) + cmd[axis] = 1.0 + _, torque_des, _ = controller(quat, ang_vel, cmd, prev_ang_vel=ang_vel, prev_cmd=cmd) + assert torque_des[axis] > 0.0, f"Torque about axis {axis} must be positive: {torque_des}" + others = np.delete(torque_des, axis) + assert np.allclose(others, 0.0, atol=1e-6), f"Cross-axis torque for axis {axis}: {others}" + + +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_matches_attitude(drone: str): + # A zero body rate command is equivalent to commanding a level attitude at the current yaw. + att_controller = parametrize(attitude2force_torque, drone) + rate_controller = parametrize(body_rate2force_torque, drone) + quat = R.from_euler("xyz", [0.2, -0.1, 0.7]).as_quat() + ang_vel = np.array([0.1, -0.2, 0.05]) + prev_ang_vel = np.array([0.05, -0.1, 0.0]) + att_cmd = np.array([0.0, 0.0, 0.7, 0.5]) + rate_cmd = np.array([0.0, 0.0, 0.0, 0.5]) + force_att, torque_att, err_att = att_controller( + quat, ang_vel, att_cmd, prev_ang_vel=prev_ang_vel + ) + force_rate, torque_rate, err_rate = rate_controller( + quat, ang_vel, rate_cmd, prev_ang_vel=prev_ang_vel + ) + assert np.allclose(force_att, force_rate, atol=1e-6) + assert np.allclose(torque_att, torque_rate, atol=1e-6) + assert np.allclose(err_att, err_rate, atol=1e-6) + + +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_leveling(drone: str): + # The firmware levels a tilted drone even at the rate setpoint. Zero attitude gains disable it. + controller = parametrize(body_rate2force_torque, drone) + params = load_params(body_rate2force_torque, drone) + quat = R.from_euler("xyz", [0.2, 0.0, 0.0]).as_quat() # Rolled by 0.2 rad + ang_vel = np.zeros(3) + cmd = np.array([0.0, 0.0, 0.0, 0.5]) + _, torque_des, _ = controller(quat, ang_vel, cmd) + assert torque_des[0] < 0.0, f"Leveling torque must oppose the roll, got {torque_des}" + params["kR"], params["ki_m"] = np.zeros(3), np.zeros(3) + _, torque_des, _ = body_rate2force_torque(quat, ang_vel, cmd, **params) + assert np.allclose(torque_des, 0.0, atol=1e-6), f"Torque with zero attitude gains {torque_des}" + + # Batch consistency (batch result == sequential result) @@ -182,6 +287,33 @@ def test_attitude2force_torque_batch_consistency(drone: str): assert np.allclose(err_batch[i, j], err_s, atol=1e-5) +@pytest.mark.unit +@pytest.mark.parametrize("drone", available_drones) +def test_body_rate2force_torque_batch_consistency(drone: str): + controller = parametrize(body_rate2force_torque, drone) + batch = (3, 2) + _, quat, _, ang_vel = create_rnd_states(batch) + _, _, _, prev_ang_vel = create_rnd_states(batch) + cmd = np.random.randn(*batch, 4) + cmd[..., 3] = np.abs(cmd[..., 3]) + prev_cmd = np.random.randn(*batch, 4) + force_batch, torque_batch, err_batch = controller( + quat, ang_vel, cmd, prev_ang_vel=prev_ang_vel, prev_cmd=prev_cmd + ) + for i in range(batch[0]): + for j in range(batch[1]): + force_s, torque_s, err_s = controller( + quat[i, j], + ang_vel[i, j], + cmd[i, j], + prev_ang_vel=prev_ang_vel[i, j], + prev_cmd=prev_cmd[i, j], + ) + assert np.allclose(force_batch[i, j], force_s, atol=1e-5) + assert np.allclose(torque_batch[i, j], torque_s, atol=1e-5) + assert np.allclose(err_batch[i, j], err_s, atol=1e-5) + + @pytest.mark.unit @pytest.mark.parametrize("drone", available_drones) def test_force_torque2rotor_vel_batch_consistency(drone: str): diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index e246a28e..6d5cabda 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -58,7 +58,7 @@ def test_sim_init(dynamics: Dynamics, device: str, control: Control, n_worlds: i n_drones = 1 if dynamics != Dynamics.first_principles: - if control in (Control.force_torque, Control.rotor_vel): + if control in (Control.body_rate, Control.force_torque, Control.rotor_vel): with pytest.raises(ConfigError): Sim(n_worlds=n_worlds, dynamics=dynamics, device=device, control=control) return @@ -89,9 +89,16 @@ def test_sim_init(dynamics: Dynamics, device: str, control: Control, n_worlds: i array_meta_assert(sim.data.controls.attitude.cmd, (n_worlds, n_drones, 4), device) else: assert sim.data.controls.attitude is None + # Test body rate buffer shapes + if control == Control.body_rate: + assert isinstance(sim.data.controls.body_rate, ControlData) + array_meta_assert(sim.data.controls.body_rate.staged_cmd, (n_worlds, n_drones, 4), device) + array_meta_assert(sim.data.controls.body_rate.cmd, (n_worlds, n_drones, 4), device) + else: + assert sim.data.controls.body_rate is None # Test force torque buffer shapes - if control in (Control.state, Control.attitude, Control.force_torque): + if control in (Control.state, Control.attitude, Control.body_rate, Control.force_torque): ft_ctrl = sim.data.controls.force_torque assert isinstance(ft_ctrl, ControlData) array_meta_assert(ft_ctrl.cmd, (n_worlds, n_drones, 4), device) @@ -104,6 +111,7 @@ def test_sim_init(dynamics: Dynamics, device: str, control: Control, n_worlds: i def test_sim_data_buffers_are_distinct(dynamics: Dynamics, control: Control, device: str): """Every leaf of SimData must own its buffer, or XLA refuses to donate the pytree.""" if dynamics != Dynamics.first_principles and control in ( + Control.body_rate, Control.force_torque, Control.rotor_vel, ): @@ -210,7 +218,7 @@ def test_reset_masked(device: str, dynamics: Dynamics): @pytest.mark.parametrize("control", Control) def test_sim_step(n_worlds: int, n_drones: int, dynamics: Dynamics, control: Control, device: str): if dynamics != Dynamics.first_principles: - if control in (Control.force_torque, Control.rotor_vel): + if control in (Control.body_rate, Control.force_torque, Control.rotor_vel): pytest.skip(f"{control} is not supported with non-first-principles dynamics") sim = Sim(