[WIP] Add drone-models. Switch to scipy PR branch. Add new controllers

Author: amacati

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "submodules/drone-models"]
+	path = submodules/drone-models
+	url = https://github.com/utiasDSL/drone-models.git

crazyflow/__init__.py

@@ -1,3 +1,7 @@
+import os
+
+os.environ["SCIPY_ARRAY_API"] = "1"
+
 import crazyflow.envs  # noqa: F401, ensure gymnasium envs are registered
 from crazyflow.control import Control
 from crazyflow.sim import Physics, Sim

crazyflow/control/control.py

@@ -50,6 +50,23 @@ class Control(str, Enum):
     default = attitude
 
 
+@jax.jit
+def controllable(step: Array, freq: int, control_steps: Array, control_freq: int) -> Array:
+    """Check which worlds can currently update their controllers.
+
+    Args:
+        step: The current step of the simulation.
+        freq: The frequency of the simulation.
+        control_steps: The steps at which the controllers were last updated.
+        control_freq: The frequency of the controllers.
+
+    Returns:
+        A boolean mask of shape (n_worlds,) that is True at the worlds where the controllers can be
+        updated.
+    """
+    return ((step - control_steps) >= (freq / control_freq)) | (control_steps == -1)
+
+
 KF: float = 3.16e-10
 KM: float = 7.94e-12
 P_F: Array = np.array([0.4, 0.4, 1.25])
@@ -111,7 +128,7 @@ def attitude2rpm(
 ) -> tuple[Array, Array]:
     """Convert the desired collective thrust and attitude into motor RPMs."""
     rot = R.from_quat(quat)
-    target_rot = R.from_euler("xyz", controls[1:])
+    target_rot = R.from_euler("xyz", controls[:3])
     drot = (target_rot.inv() * rot).as_matrix()
     # Extract the anti-symmetric part of the relative rotation matrix.
     rot_e = jnp.array([drot[2, 1] - drot[1, 2], drot[0, 2] - drot[2, 0], drot[1, 0] - drot[0, 1]])
@@ -123,7 +140,7 @@ def attitude2rpm(
     # PID target torques.
     target_torques = -P_T * rot_e + D_T * rpy_rates_e + I_T * rpy_err_i
     target_torques = jnp.clip(target_torques, -3200, 3200)
-    thrust_per_motor = jnp.atleast_1d(controls[0]) / 4
+    thrust_per_motor = jnp.atleast_1d(controls[3]) / 4
     pwm = jnp.clip(thrust2pwm(thrust_per_motor) + MIX_MATRIX @ target_torques, MIN_PWM, MAX_PWM)
     return pwm2rpm(pwm), rpy_err_i
 

crazyflow/control/mellinger.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import jax.numpy as jnp
+from drone_models.controller.mellinger import MellingerStateParams
+from flax.struct import dataclass, field
+from jax import Array, Device
+
+
+@dataclass
+class MellingerStateData:
+    cmd: Array  # (N, M, 13)
+    """Full state control command for the drone.
+
+    A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
+    We currently do not use the acceleration and angle rate components. This is subject to change.
+    """
+    steps: Array  # (N, 1)
+    """Last simulation steps that the state control command was applied."""
+    freq: int = field(pytree_node=False)
+    """Frequency of the state control command."""
+    pos_err_i: Array  # (N, M, 3)
+    """Integral errors of the state control command."""
+    # Parameters for the state controller
+    params: MellingerStateParams
+
+    @staticmethod
+    def create(
+        n_worlds: int, n_drones: int, freq: int, drone_model: str, device: Device
+    ) -> MellingerStateData:
+        """Create a default set of state data for the simulation."""
+        cmd = jnp.zeros((n_worlds, n_drones, 13), device=device)
+        steps = jnp.zeros((n_worlds, 1), dtype=jnp.int32, device=device)
+        pos_err_i = jnp.zeros((n_worlds, n_drones, 3), device=device)
+        params = MellingerStateParams.load(drone_model)
+        return MellingerStateData(
+            cmd=cmd, steps=steps, freq=freq, pos_err_i=pos_err_i, params=params
+        )
+
+
+# @dataclass
+# class MellingerAttitudeData:
+#     cmd: Array  # (N, M, 4)
+#     """Full attitude control command for the drone.
+
+#     A command consists of [collective thrust, roll, pitch, yaw].
+#     """
+#     steps: Array  # (N, 1)
+#     """Last simulation steps that the attitude control command was applied."""
+#     freq: int = field(pytree_node=False)
+#     """Frequency of the attitude control command."""
+#     pos_err_i: Array  # (N, M, 3)
+#     """Integral errors of the attitude control command."""
+#     # Parameters for the attitude controller
+#     params: MellingerAttitudeParams
+
+#     @staticmethod
+#     def create(
+#         n_worlds: int, n_drones: int, freq: int, drone_model: str, device: Device
+#     ) -> MellingerAttitudeData:
+#         """Create a default set of attitude data for the simulation."""
+#         cmd = jnp.zeros((n_worlds, n_drones, 4), device=device)
+#         steps = jnp.zeros((n_worlds, 1), dtype=jnp.int32, device=device)
+#         pos_err_i = jnp.zeros((n_worlds, n_drones, 3), device=device)
+#         params = MellingerAttitudeParams.load(drone_model)
+#         return MellingerAttitudeData(
+#             cmd=cmd, steps=steps, freq=freq, pos_err_i=pos_err_i, params=params
+#         )

crazyflow/sim/sim.py

@@ -8,13 +8,15 @@
 import jax.numpy as jnp
 import mujoco
 import mujoco.mjx as mjx
+from drone_models.controller.mellinger import state2attitude
 from einops import rearrange
 from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 from jax import Array, Device
 from jax.scipy.spatial.transform import Rotation as R
 
 from crazyflow.constants import J_INV, MASS, J
-from crazyflow.control.control import Control, attitude2rpm, pwm2rpm, state2attitude, thrust2pwm
+from crazyflow.control.control import Control, attitude2rpm, controllable, pwm2rpm, thrust2pwm
+from crazyflow.control.control import state2attitude as state2attitude_legacy
 from crazyflow.exception import ConfigError, NotInitializedError
 from crazyflow.sim.integration import Integrator, euler, rk4, symplectic_euler
 from crazyflow.sim.physics import (
@@ -24,13 +26,23 @@
     rpms2collective_wrench,
     surrogate_identified_collective_wrench,
 )
-from crazyflow.sim.structs import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
+from crazyflow.sim.structs import (
+    SimControls,
+    SimControlsNew,
+    SimCore,
+    SimData,
+    SimParams,
+    SimState,
+    SimStateDeriv,
+)
 from crazyflow.utils import grid_2d, leaf_replace, pytree_replace, to_device
 
 if TYPE_CHECKING:
     from mujoco.mjx import Data, Model
     from numpy.typing import NDArray
 
+    from crazyflow.control.mellinger import MellingerStateData
+
 Params = ParamSpec("Params")  # Represents arbitrary parameters
 Return = TypeVar("Return")  # Represents the return type
 
@@ -141,7 +153,11 @@ def state_control(self, controls: Array):
         assert controls.shape == (self.n_worlds, self.n_drones, 13), "controls shape mismatch"
         assert self.control == Control.state, "State control is not enabled by the sim config"
         controls = to_device(controls, self.device)
-        self.data = self.data.replace(controls=self.data.controls.replace(state=controls))
+        self.data = self.data.replace(
+            new_controls=self.data.new_controls.replace(
+                state=self.data.new_controls.state.replace(cmd=controls)
+            )
+        )
 
     def thrust_control(self, cmd: Array):
         """Set the desired thrust for all drones in all worlds."""
@@ -182,7 +198,7 @@ def seed(self, seed: int):
         Args:
             seed: The seed for the JAX rng.
         """
-        self.data = seed_sim(self.data, seed, self.device)
+        self.data: SimData = seed_sim(self.data, seed, self.device)
 
     def close(self):
         if self.viewer is not None:
@@ -282,7 +298,7 @@ def build_mjx(self):
 
     def init_data(
         self, state_freq: int, attitude_freq: int, thrust_freq: int, rng_key: Array
-    ) -> tuple[SimData, SimData]:
+    ) -> SimData:
         """Initialize the simulation data."""
         drone_ids = [self.mj_model.body(f"drone:{i}").id for i in range(self.n_drones)]
         N, D = self.n_worlds, self.n_drones
@@ -292,6 +308,9 @@ def init_data(
             controls=SimControls.create(N, D, state_freq, attitude_freq, thrust_freq, self.device),
             params=SimParams.create(N, D, MASS, J, J_INV, self.device),
             core=SimCore.create(self.freq, N, D, drone_ids, rng_key, self.device),
+            new_controls=SimControlsNew.create(
+                N, D, self.control, state_freq, attitude_freq, thrust_freq, self.device
+            ),
         )
         if D > 1:  # If multiple drones, arrange them in a grid
             grid = grid_2d(D)
@@ -417,23 +436,6 @@ def integrate(data: SimData) -> SimData:
     return integrate
 
 
-@jax.jit
-def controllable(step: Array, freq: int, control_steps: Array, control_freq: int) -> Array:
-    """Check which worlds can currently update their controllers.
-
-    Args:
-        step: The current step of the simulation.
-        freq: The frequency of the simulation.
-        control_steps: The steps at which the controllers were last updated.
-        control_freq: The frequency of the controllers.
-
-    Returns:
-        A boolean mask of shape (n_worlds,) that is True at the worlds where the controllers can be
-        updated.
-    """
-    return ((step - control_steps) >= (freq / control_freq)) | (control_steps == -1)
-
-
 @jax.jit
 def contacts(geom_start: int, geom_count: int, data: Data) -> Array:
     """Filter contacts from MuJoCo data."""
@@ -461,18 +463,53 @@ def sync_sim2mjx(data: SimData, mjx_data: Data, mjx_model: Model) -> tuple[SimDa
 
 def step_state_controller(data: SimData) -> SimData:
     """Compute the updated controls for the state controller."""
-    states, controls = data.states, data.controls
-    mask = controllable(data.core.steps, data.core.freq, controls.state_steps, controls.state_freq)
-    des_pos, des_vel = controls.state[..., :3], controls.state[..., 3:6]
-    des_yaw = controls.state[..., [9]]  # Keep (N, M, 1) shape for broadcasting
-    dt = 1 / data.controls.state_freq
-    attitude, pos_err_i = state2attitude(
-        states.pos, states.vel, states.quat, des_pos, des_vel, des_yaw, controls.pos_err_i, dt
+    states, ctrl_state = data.states, data.new_controls.state
+    assert ctrl_state is not None, "Using state controller without initialized state control data"
+    ctrl_state: MellingerStateData
+    mask = controllable(data.core.steps, data.core.freq, ctrl_state.steps, ctrl_state.freq)
+    jax.debug.print("Ctrl cmd: {cmd}", cmd=ctrl_state.cmd)
+    attitude, (pos_err_i,) = state2attitude(
+        states.pos,
+        states.quat,
+        states.vel,
+        states.ang_vel,
+        ctrl_state.cmd,
+        ctrl_freq=ctrl_state.freq,
+        ctrl_errors=(ctrl_state.pos_err_i,),
+        **ctrl_state.params._asdict(),
     )
-    controls = leaf_replace(
-        controls, mask, state_steps=data.core.steps, staged_attitude=attitude, pos_err_i=pos_err_i
+    jax.debug.print("Attitude: {attitude}", attitude=attitude)
+    ctrl_state = leaf_replace(ctrl_state, mask, steps=data.core.steps, pos_err_i=pos_err_i)
+    data = data.replace(
+        controls=data.controls.replace(staged_attitude=attitude),
+        new_controls=data.new_controls.replace(state=ctrl_state),
     )
-    return data.replace(controls=controls)
+    return data
+
+
+# def step_state_controller(data: SimData) -> SimData:
+#     """Compute the updated controls for the state controller."""
+#     states, ctrl_state = data.states, data.new_controls.state
+#     assert ctrl_state is not None, "Using state controller without initialized state control data"
+#     mask = controllable(data.core.steps, data.core.freq, ctrl_state.steps, ctrl_state.freq)
+#     attitude, pos_err_i = state2attitude_legacy(
+#         states.pos,
+#         states.vel,
+#         states.quat,
+#         ctrl_state.cmd[..., :3],
+#         ctrl_state.cmd[..., 3:6],
+#         ctrl_state.cmd[..., [9]],
+#         ctrl_state.pos_err_i,
+#         1 / ctrl_state.freq,
+#     )
+#     attitude = jnp.roll(attitude, -1, axis=-1)
+#     jax.debug.print("Attitude: {attitude}", attitude=attitude)
+#     ctrl_state = leaf_replace(ctrl_state, mask, steps=data.core.steps, pos_err_i=pos_err_i)
+#     data = data.replace(
+#         controls=data.controls.replace(staged_attitude=attitude),
+#         new_controls=data.new_controls.replace(state=ctrl_state),
+#     )
+#     return data
 
 
 def step_attitude_controller(data: SimData) -> SimData:
@@ -533,14 +570,6 @@ def identified_wrench(data: SimData) -> SimData:
 identified_derivative = analytical_derivative  # We can use the same derivative function for both
 
 
-def identity(data: SimData, *args: Any, **kwargs: Any) -> SimData:
-    """Identity function for the simulation pipeline.
-
-    Used as default function for optional pipeline steps.
-    """
-    return data
-
-
 def clip_floor_pos(data: SimData) -> SimData:
     """Clip the position of the drone to the floor."""
     clip = data.states.pos[..., 2] < -0.001

crazyflow/sim/structs.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import typing
+
 import jax
 import jax.numpy as jnp
 from flax.struct import dataclass, field
 from jax import Array, Device
 
+from crazyflow.control import Control
+from crazyflow.control.mellinger import MellingerStateData
+
 
 @dataclass
 class SimState:
@@ -146,6 +151,54 @@ def create(
         )
 
 
+class ControlData(typing.Protocol):
+    cmd: Array  # (N, M, X)
+    """Control command for the drone."""
+    steps: Array  # (N, 1)
+    """Last simulation steps that the state control command was applied."""
+    freq: int
+    """Frequency of the state control command."""
+    # Parameters for the controller
+    params: typing.Any
+
+
+@dataclass
+class SimControlsNew:
+    state: ControlData | None = None
+    """State control data."""
+    attitude: ControlData | None = None
+    """Attitude control data."""
+    thrust: ControlData | None = None
+    """Thrust control data."""
+
+    @staticmethod
+    def create(
+        n_worlds: int,
+        n_drones: int,
+        control: Control,
+        state_freq: int | None,
+        attitude_freq: int | None,
+        thrust_freq: int | None,
+        device: Device,
+    ) -> SimControlsNew:
+        """Create a default set of controls for the simulation."""
+        match control:
+            case Control.state:
+                state = MellingerStateData.create(n_worlds, n_drones, state_freq, "", device)
+                attitude = None  # MellingerAttitudeData.create(n_worlds, n_drones, device)
+                thrust = None  # MellingerThrustData.create(n_worlds, n_drones, device)
+                return SimControlsNew(state=state, attitude=attitude, thrust=thrust)
+            case Control.attitude:
+                attitude = None  # MellingerAttitudeData.create(n_worlds, n_drones, device)
+                thrust = None  # MellingerThrustData.create(n_worlds, n_drones, device)
+                return SimControlsNew(attitude=attitude, thrust=thrust)
+            case Control.thrust:
+                thrust = None  # MellingerThrustData.create(n_worlds, n_drones, device)
+                return SimControlsNew(thrust=thrust)
+            case _:
+                raise ValueError(f"Control mode {control} not implemented")
+
+
 @dataclass
 class SimParams:
     mass: Array  # (N, M, 1)
@@ -217,6 +270,8 @@ class SimData:
     """Derivative of the state of the simulation."""
     controls: SimControls
     """Drone control values."""
+    new_controls: SimControlsNew
+    """New style control data TODO improve this"""
     params: SimParams
     """Drone parameters."""
     core: SimCore