Update state interface to current cflib

README.md

@@ -18,9 +18,11 @@ Crazyflow is a research simulator for quadrotors. It runs batched, differentiabl
 import numpy as np
 from crazyflow.sim import Sim
 from crazyflow.control import Control
+from scipy.spatial.transform import Rotation as R
 
 sim = Sim(n_worlds=4096, n_drones=1, control=Control.state)
-cmd = np.zeros((4096, 1, 13))
+cmd = np.zeros((4096, 1, 16))
+cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
 cmd[..., 2] = 0.5  # hover at 0.5 m across all worlds
 
 for _ in range(100):

benchmark/performance.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
+os.environ["SCIPY_ARRAY_API"] = "1"
+
 import gymnasium
 import jax
 import numpy as np
 from ml_collections import config_dict
 from pyinstrument import Profiler
 from pyinstrument.renderers.html import HTMLRenderer
+from scipy.spatial.transform import Rotation as R
 
 import crazyflow  # noqa: F401, ensure gymnasium envs are registered
 from crazyflow.sim import Sim
@@ -19,9 +23,11 @@
 def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
     sim = Sim(**sim_config)
     device = jax.devices(device)[0]
-    ndim = 13 if sim.control == "state" else 4
+    ndim = 16 if sim.control == "state" else 4
     control_fn = sim.state_control if sim.control == "state" else sim.attitude_control
     cmd = np.zeros((sim.n_worlds, sim.n_drones, ndim))
+    if sim.control == "state":
+        cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
     # Ensure JIT compiled dynamics and control
     sim.reset()
     control_fn(cmd)

benchmark/splat.py

@@ -24,12 +24,14 @@
 # splax rasterizes with warp, which needs GPU memory outside JAX's pool. Disable JAX preallocation
 # before it initializes so both share the device. Must run before the first jax import.
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+os.environ["SCIPY_ARRAY_API"] = "1"
 
 import fire
 import jax
 import jax.numpy as jnp
 import numpy as np
 from jax.errors import JaxRuntimeError
+from scipy.spatial.transform import Rotation as R
 from splax.io import fetch
 
 from crazyflow.sim import Sim
@@ -118,7 +120,8 @@ def benchmark(
 
             # Hold a constant target so the drone keeps moving and each frame renders a distinct
             # pose. A static scene would let XLA hoist the render out of the loop.
-            cmd = np.zeros((sim.n_worlds, sim.n_drones, 13), dtype=np.float32)
+            cmd = np.zeros((sim.n_worlds, sim.n_drones, 16), dtype=np.float32)
+            cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
             cmd[..., 2] = 0.5
             sim.reset()
             sim.state_control(jnp.asarray(cmd, device=sim.device))

crazyflow/control/core.py

@@ -38,10 +38,12 @@ def parametrize(
     import numpy as np
     from crazyflow.control import parametrize
     from crazyflow.control.mellinger import state2attitude
+    from scipy.spatial.transform import Rotation as R
 
     ctrl = parametrize(state2attitude, "cf2x_L250")
     pos, quat = np.zeros(3), np.array([0.0, 0.0, 0.0, 1.0])
-    vel, cmd = np.zeros(3), np.zeros(13)
+    vel, cmd = np.zeros(3), np.zeros(16)
+    cmd[9:13] = R.from_euler("z", 0.0).as_quat()
     rpyt, int_pos_err = ctrl(pos, quat, vel, cmd)
     ```
 
@@ -85,14 +87,14 @@ class Control(StrEnum):
     """Control type of the simulated onboard controller."""
 
     state = "state"
-    """State control takes [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
-    
+    """State control takes [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz].
+
     Note:
         Recommended frequency is >=20 Hz.
 
     Warning:
-        Currently, we only use positions, velocities, and yaw. The rest of the state is ignored.
-        This is subject to change in the future.
+        Only the yaw of the attitude quaternion is used, as in the firmware. The so_rpy family
+        ignores the body rate setpoint.
     """
     attitude = "attitude"
     """Attitude control takes [roll, pitch, yaw, collective thrust].
@@ -101,7 +103,7 @@ class Control(StrEnum):
         Recommended frequency is >=100 Hz.
     """
     body_rate = "body_rate"
-    """Body rate control takes [roll_rate, pitch_rate, yaw_rate, collective thrust].
+    """Body rate control takes [wx, wy, wz, collective thrust].
 
     Note:
         Recommended frequency is >=200 Hz.

crazyflow/control/mellinger/control.py

@@ -58,8 +58,9 @@ def state2attitude(
         pos: Drone position with shape (..., 3).
         quat: Drone orientation as xyzw quaternion with shape (..., 4).
         vel: Drone velocity with shape (..., 3).
-        cmd: Full state command in SI units and rad with shape (..., 13). The entries are
-            [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
+        cmd: Full state command in SI units with shape (..., 16). The entries are
+            [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz]. Only the yaw of the
+            attitude quaternion is used. The body rates are forwarded to the attitude controller.
         pos_err_i: Position integral error (..., 3) from the previous call. If None, it is
             initialised to zero.
         ctrl_freq: Control frequency in Hz
@@ -83,9 +84,8 @@ def state2attitude(
     setpoint_pos = cmd[..., 0:3]
     setpoint_vel = cmd[..., 3:6]
     setpoint_acc = cmd[..., 6:9]
-    setpoint_yaw = cmd[..., 9]
+    setpoint_quat = cmd[..., 9:13]
     dt = 1 / ctrl_freq
-    # setpointRPY_rates = cmd[..., 10:13]
     # From firmware controller_mellinger
     pos_err = setpoint_pos - pos  # l. 145 Position Error (ep)
     vel_err = setpoint_vel - vel  # l. 148 Velocity Error (ev)
@@ -100,7 +100,7 @@ def state2attitude(
     )
     # l. 178 Rate-controlled YAW is moving YAW angle setpoint
     # => only one case here, since the setpoint is always in absolute mode
-    desired_yaw = setpoint_yaw
+    desired_yaw = R.from_quat(setpoint_quat).as_euler("xyz")[..., 2]
     # l. 189 Z-Axis [zB]
     rot = R.from_quat(quat).as_matrix()
     z_axis = rot[..., -1]  # 3rd column or roation matrix is z axis
@@ -489,13 +489,12 @@ def force_torque2rotor_vel(
 
 @dataclass
 class MellingerStateData:
-    cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 13)
+    cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 16)
     """Full state control command for the drone.
 
-    A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
-    We currently do not use the acceleration and angle rate components. This is subject to change.
+    A command consists of [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz].
     """
-    staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 13)
+    staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 16)
     """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 state control command was applied."""
@@ -512,12 +511,12 @@ def create(
     ) -> MellingerStateData:
         """Create a default set of state data for the simulation."""
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
-        zeros_13d = jnp.zeros((n_worlds, n_drones, 13), device=device)
+        cmd = jnp.zeros((n_worlds, n_drones, 16), device=device).at[..., 12].set(1.0)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
         params = load_params(state2attitude, drone, xp=jnp, device=device)
         return MellingerStateData(
-            cmd=zeros_13d.copy(),
-            staged_cmd=zeros_13d.copy(),
+            cmd=cmd,
+            staged_cmd=cmd.copy(),
             steps=steps,
             freq=freq,
             pos_err_i=zeros_3d.copy(),
@@ -528,20 +527,21 @@ def create(
 @dataclass
 class MellingerAttitudeData:
     cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
-    """Full attitude control command for the drone.
-
-    A command consists of [roll, pitch, yaw, collective thrust].
-    """
+    """Attitude control setpoint consisting of [roll, pitch, yaw, 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."""
+    ang_vel_des: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
+    """Body rate setpoint [wx, wy, wz] of the attitude controller."""
+    staged_ang_vel_des: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
+    """Staging buffer to store the most recent body rate setpoint until the next controller tick."""
     steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Last simulation steps that the attitude control command was applied."""
     freq: int = field(pytree_node=False)
     """Frequency of the attitude control command."""
     r_int_error: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Integral errors of the attitude control command."""
-    last_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
-    """Last angular velocity of the drone."""
+    prev_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
+    """Previous angular velocity of the drone."""
     # Parameters for the attitude controller
     params: dict[str, Array]
 
@@ -557,10 +557,12 @@ def create(
         return MellingerAttitudeData(
             cmd=zeros_4d.copy(),
             staged_cmd=zeros_4d.copy(),
+            ang_vel_des=zeros_3d.copy(),
+            staged_ang_vel_des=zeros_3d.copy(),
             steps=steps,
             freq=freq,
             r_int_error=zeros_3d.copy(),
-            last_ang_vel=zeros_3d.copy(),
+            prev_ang_vel=zeros_3d.copy(),
             params=params,
         )
 
@@ -570,7 +572,7 @@ 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].
+    A command consists of [wx, wy, wz, 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."""
@@ -580,8 +582,8 @@ class MellingerBodyRateData:
     """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."""
+    prev_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
+    """Previous angular velocity of the drone."""
     # Parameters for the body rate controller
     params: dict[str, Array]
 
@@ -600,7 +602,7 @@ def create(
             steps=steps,
             freq=freq,
             r_int_error=zeros_3d.copy(),
-            last_ang_vel=zeros_3d.copy(),
+            prev_ang_vel=zeros_3d.copy(),
             params=params,
         )
 
@@ -650,7 +652,9 @@ def control_state2attitude(data: SimData) -> SimData:
         **state_ctrl.params,
     )
     state_ctrl = leaf_replace(state_ctrl, mask, steps=data.core.steps, pos_err_i=pos_err_i)
-    attitude_ctrl = leaf_replace(data.controls.attitude, mask, staged_cmd=rpyt)
+    attitude_ctrl = leaf_replace(
+        data.controls.attitude, mask, staged_cmd=rpyt, staged_ang_vel_des=state_ctrl.cmd[..., 13:16]
+    )
     return data.replace(controls=data.controls.replace(state=state_ctrl, attitude=attitude_ctrl))
 
 
@@ -660,21 +664,30 @@ def control_attitude2force_torque(data: SimData) -> SimData:
     attitude_ctrl: MellingerAttitudeData = data.controls.attitude
     assert attitude_ctrl is not None, "Using attitude controller without initialized data"
     mask = controllable(data.core.steps, data.core.freq, attitude_ctrl.steps, attitude_ctrl.freq)
-    attitude_ctrl = leaf_replace(attitude_ctrl, mask, cmd=attitude_ctrl.staged_cmd)
-    force, torque, r_int_error = attitude2force_torque(
+    prev_ang_vel_des = attitude_ctrl.ang_vel_des
+    attitude_ctrl = leaf_replace(
+        attitude_ctrl,
+        mask,
+        cmd=attitude_ctrl.staged_cmd,
+        ang_vel_des=attitude_ctrl.staged_ang_vel_des,
+    )
+    force, torque, r_int_error = _attitude2force_torque(
         states.quat,
         states.ang_vel,
-        attitude_ctrl.cmd,
-        r_int_error=attitude_ctrl.r_int_error,
-        ctrl_freq=attitude_ctrl.freq,
-        prev_ang_vel=attitude_ctrl.last_ang_vel,
+        attitude_ctrl.cmd[..., :3],
+        attitude_ctrl.ang_vel_des,
+        attitude_ctrl.cmd[..., 3],
+        attitude_ctrl.prev_ang_vel,
+        prev_ang_vel_des,
+        attitude_ctrl.r_int_error,
+        attitude_ctrl.freq,
         **attitude_ctrl.params,
     )
     attitude_ctrl = leaf_replace(
         attitude_ctrl,
         mask,
         r_int_error=r_int_error,
-        last_ang_vel=states.ang_vel,
+        prev_ang_vel=states.ang_vel,
         steps=data.core.steps,
     )
     ft_ctrl = leaf_replace(
@@ -697,7 +710,7 @@ def control_body_rate2force_torque(data: SimData) -> SimData:
         states.quat,
         states.ang_vel,
         body_rate_ctrl.cmd,
-        prev_ang_vel=body_rate_ctrl.last_ang_vel,
+        prev_ang_vel=body_rate_ctrl.prev_ang_vel,
         prev_cmd=prev_cmd,
         r_int_error=body_rate_ctrl.r_int_error,
         ctrl_freq=body_rate_ctrl.freq,
@@ -707,7 +720,7 @@ def control_body_rate2force_torque(data: SimData) -> SimData:
         body_rate_ctrl,
         mask,
         r_int_error=r_int_error,
-        last_ang_vel=states.ang_vel,
+        prev_ang_vel=states.ang_vel,
         steps=data.core.steps,
     )
     ft_ctrl = leaf_replace(

crazyflow/sim/functional.py

@@ -16,7 +16,7 @@
 def state_control(data: SimData, controls: Array) -> SimData:
     """State control function."""
     assert data.controls.mode == Control.state, f"control type {data.controls.mode} not enabled"
-    assert controls.shape == (data.core.n_worlds, data.core.n_drones, 13), "controls shape mismatch"
+    assert controls.shape == (data.core.n_worlds, data.core.n_drones, 16), "controls shape mismatch"
     controls = jnp.asarray(controls)
     data = data.replace(
         controls=data.controls.replace(state=data.controls.state.replace(staged_cmd=controls))