Fix most tests

Author: amacati

crazyflow/control/mellinger.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
 import jax.numpy as jnp
-from drone_models.controller.mellinger import (
-    MellingerAttitudeParams,
-    MellingerForceTorqueParams,
-    MellingerStateParams,
-)
+from drone_models.controller.mellinger.params import AttitudeParams, ForceTorqueParams, StateParams
 from flax.struct import dataclass, field
 from jax import Array, Device
 
+from crazyflow.utils import named_tuple2device
+
 
 @dataclass
 class MellingerStateData:
@@ -27,7 +25,7 @@ class MellingerStateData:
     pos_err_i: Array  # (N, M, 3)
     """Integral errors of the state control command."""
     # Parameters for the state controller
-    params: MellingerStateParams
+    params: StateParams
 
     @staticmethod
     def create(
@@ -37,7 +35,7 @@ def create(
         cmd = jnp.zeros((n_worlds, n_drones, 13), device=device)
         steps = -jnp.ones((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)
+        params = named_tuple2device(StateParams.load(drone_model), device)
         return MellingerStateData(
             cmd=cmd, staged_cmd=cmd, steps=steps, freq=freq, pos_err_i=pos_err_i, params=params
         )
@@ -61,7 +59,7 @@ class MellingerAttitudeData:
     last_ang_vel: Array  # (N, M, 3)
     """Last angular velocity of the drone."""
     # Parameters for the attitude controller
-    params: MellingerAttitudeParams
+    params: AttitudeParams
 
     @staticmethod
     def create(
@@ -71,7 +69,7 @@ def create(
         cmd = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
-        params = MellingerAttitudeParams.load(drone_model)
+        params = named_tuple2device(AttitudeParams.load(drone_model), device)
         return MellingerAttitudeData(
             cmd=cmd,
             staged_cmd=cmd,
@@ -85,40 +83,27 @@ def create(
 
 @dataclass
 class MellingerForceTorqueData:
-    cmd_force: Array  # (N, M, 1)
-    """Force command for the drone.
-
-    A command consists of [fz].
-    """
-    cmd_torque: Array  # (N, M, 3)
-    """Torque command for the drone.
+    cmd: Array  # (N, M, 4)
+    """Force-torque command for the drone.
 
-    A command consists of [tx, ty, tz].
+    A command consists of [fz, tx, ty, tz].
     """
-    staged_cmd_force: Array  # (N, M, 1)
-    staged_cmd_torque: Array  # (N, M, 3)
+    staged_cmd: Array  # (N, M, 4)
     """Staging buffer to store the most recent command until the next controller tick."""
     steps: Array  # (N, 1)
     """Last simulation steps that the force and torque control command was applied."""
     freq: int = field(pytree_node=False)
     """Frequency of the force and torque control command."""
     # Parameters for the force and torque controller
-    params: MellingerForceTorqueParams
+    params: ForceTorqueParams
 
     @staticmethod
     def create(
         n_worlds: int, n_drones: int, freq: int, drone_model: str, device: Device
     ) -> MellingerForceTorqueData:
-        zero_1d = jnp.zeros((n_worlds, n_drones, 1), device=device)
-        zero_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
+        zero_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = MellingerForceTorqueParams.load(drone_model)
+        params = named_tuple2device(ForceTorqueParams.load(drone_model), device)
         return MellingerForceTorqueData(
-            cmd_force=zero_1d,
-            cmd_torque=zero_3d,
-            staged_cmd_force=zero_1d,
-            staged_cmd_torque=zero_3d,
-            steps=steps,
-            freq=freq,
-            params=params,
+            cmd=zero_4d, staged_cmd=zero_4d, steps=steps, freq=freq, params=params
         )

crazyflow/sim/__init__.py

@@ -1,5 +1,5 @@
 from crazyflow.sim.physics import Physics
 from crazyflow.sim.sim import Sim
-from crazyflow.sim.symbolic import symbolic_attitude, symbolic_from_sim, symbolic_thrust
+from crazyflow.sim.symbolic import symbolic_from_sim
 
-__all__ = ["Sim", "Physics", "symbolic_attitude", "symbolic_from_sim", "symbolic_thrust"]
+__all__ = ["Sim", "Physics", "symbolic_from_sim"]

crazyflow/sim/sim.py

@@ -74,6 +74,7 @@ def __init__(
         self,
         n_worlds: int = 1,
         n_drones: int = 1,
+        drone_model: str = "cf2x_L250",
         physics: Physics = Physics.default,
         control: Control = Control.default,
         integrator: Integrator = Integrator.default,
@@ -93,6 +94,7 @@ def __init__(
             raise ConfigError("Double precision mode is required for high frequency simulations")
         self.physics = physics
         self.control = control
+        self.drone_model = drone_model
         self.integrator = integrator
         self.device = jax.devices(device)[0]
         self.n_worlds = n_worlds
@@ -179,9 +181,7 @@ def force_torque_control(self, cmd: Array):
         controls = to_device(cmd, self.device)
         self.data = self.data.replace(
             controls=self.data.controls.replace(
-                force_torque=self.data.controls.force_torque.replace(
-                    staged_cmd_force=controls[..., [0]], staged_cmd_torque=controls[..., 1:]
-                )
+                force_torque=self.data.controls.force_torque.replace(staged_cmd=controls)
             )
         )
 
@@ -325,7 +325,14 @@ def init_data(
             states=SimState.create(N, D, self.device),
             states_deriv=SimStateDeriv.create(N, D, self.device),
             controls=SimControls.create(
-                N, D, self.control, state_freq, attitude_freq, force_torque_freq, self.device
+                N,
+                D,
+                self.control,
+                self.drone_model,
+                state_freq,
+                attitude_freq,
+                force_torque_freq,
+                self.device,
             ),
             params=SimParams.create(N, D, MASS, J, J_INV, self.device),
             constants=SimConstants.create(
@@ -501,7 +508,7 @@ def step_state_controller(data: SimData) -> SimData:
         **state_ctrl.params._asdict(),
     )
     state_ctrl = leaf_replace(state_ctrl, mask, steps=data.core.steps, pos_err_i=pos_err_i)
-    attitude_ctrl = data.controls.attitude.replace(staged_cmd=rpyt)
+    attitude_ctrl = leaf_replace(data.controls.attitude, mask, staged_cmd=rpyt)
     data = data.replace(controls=data.controls.replace(state=state_ctrl, attitude=attitude_ctrl))
     return data
 
@@ -520,7 +527,6 @@ def step_attitude_controller(data: SimData) -> SimData:
         states.ang_vel,
         attitude_ctrl.cmd,
         ctrl_errors=(attitude_ctrl.r_int_error,),
-        ctrl_info=(attitude_ctrl.last_ang_vel,),
         ctrl_freq=attitude_ctrl.freq,
         **attitude_ctrl.params._asdict(),
     )
@@ -531,16 +537,14 @@ def step_attitude_controller(data: SimData) -> SimData:
         last_ang_vel=states.ang_vel,
         steps=data.core.steps,
     )
-    force_torque_ctrl = data.controls.force_torque.replace(
-        staged_cmd_force=force, staged_cmd_torque=torque
-    )
+    ft_ctrl = data.controls.force_torque.replace(staged_cmd=jnp.concat([force, torque], axis=-1))
     # TODO: Remove. Set the force and torque directly into the physics step.
     r = R.from_quat(states.quat)
     torque = r.apply(torque)
     force = r.apply(jnp.zeros_like(torque).at[..., 2].set(force[..., 0]))
-    data = data.replace(states=data.states.replace(force=force, torque=torque))
+    states = leaf_replace(states, mask, force=force, torque=torque)
     return data.replace(
-        controls=data.controls.replace(attitude=attitude_ctrl, force_torque=force_torque_ctrl)
+        states=states, controls=data.controls.replace(attitude=attitude_ctrl, force_torque=ft_ctrl)
     )
 
 
@@ -550,19 +554,17 @@ def step_force_torque_controller(data: SimData) -> SimData:
     assert ft_ctrl is not None, "Using force torque controller without initialized data"
     ft_ctrl: MellingerForceTorqueData
     mask = controllable(data.core.steps, data.core.freq, ft_ctrl.steps, ft_ctrl.freq)
-    ft_ctrl = leaf_replace(
-        ft_ctrl, mask, cmd_force=ft_ctrl.staged_cmd_force, cmd_torque=ft_ctrl.staged_cmd_torque
-    )
+    ft_ctrl = leaf_replace(ft_ctrl, mask, cmd=ft_ctrl.staged_cmd)
     rotor_vel = force_torque2rotor_vel(
-        ft_ctrl.cmd_force, ft_ctrl.cmd_torque, **ft_ctrl.params._asdict()
+        ft_ctrl.cmd[..., [0]], ft_ctrl.cmd[..., 1:], **ft_ctrl.params._asdict()
     )
     ft_ctrl = leaf_replace(ft_ctrl, mask, steps=data.core.steps)
     data = data.replace(controls=data.controls.replace(rotor_vel=rotor_vel, force_torque=ft_ctrl))
     # TODO: Remove
     r = R.from_quat(data.states.quat)
-    cmd_force = jnp.zeros_like(ft_ctrl.cmd_torque)
-    cmd_force = cmd_force.at[..., 2].set(ft_ctrl.cmd_force[..., 0])
-    force, torque = r.apply(cmd_force), r.apply(ft_ctrl.cmd_torque)
+    cmd_force = jnp.zeros_like(ft_ctrl.cmd[..., 1:])
+    cmd_force = cmd_force.at[..., 2].set(ft_ctrl.cmd[..., 0])
+    force, torque = r.apply(cmd_force), r.apply(ft_ctrl.cmd[..., 1:])
     data = data.replace(states=data.states.replace(force=force, torque=torque))
     return data
 

crazyflow/sim/structs.py

@@ -71,6 +71,7 @@ def create(n_worlds: int, n_drones: int, device: Device) -> SimStateDeriv:
         return SimStateDeriv(dpos=dpos, drot=drot, dvel=dvel, dang_vel=dang_vel)
 
 
+@typing.runtime_checkable
 class ControlData(typing.Protocol):
     staged_cmd: Array  # (N, M, X)
     """Staged control command for the drone.
@@ -106,6 +107,7 @@ def create(
         n_worlds: int,
         n_drones: int,
         control: Control,
+        drone_model: str,
         state_freq: int | None,
         attitude_freq: int | None,
         force_torque_freq: int | None,
@@ -115,29 +117,31 @@ def create(
         rotor_vel = jnp.zeros((n_worlds, n_drones, 4), device=device)
         match control:
             case Control.state:
-                state = MellingerStateData.create(n_worlds, n_drones, state_freq, "", device)
+                state = MellingerStateData.create(
+                    n_worlds, n_drones, state_freq, drone_model, device
+                )
                 attitude = MellingerAttitudeData.create(
-                    n_worlds, n_drones, attitude_freq, "", device
+                    n_worlds, n_drones, attitude_freq, drone_model, device
                 )
                 force_torque = MellingerForceTorqueData.create(
-                    n_worlds, n_drones, force_torque_freq, "", device
+                    n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
                     state=state, attitude=attitude, force_torque=force_torque, rotor_vel=rotor_vel
                 )
             case Control.attitude:
                 attitude = attitude = MellingerAttitudeData.create(
-                    n_worlds, n_drones, attitude_freq, "", device
+                    n_worlds, n_drones, attitude_freq, drone_model, device
                 )
                 force_torque = MellingerForceTorqueData.create(
-                    n_worlds, n_drones, force_torque_freq, "", device
+                    n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
                     state=None, attitude=attitude, force_torque=force_torque, rotor_vel=rotor_vel
                 )
             case Control.force_torque:
                 force_torque = MellingerForceTorqueData.create(
-                    n_worlds, n_drones, force_torque_freq, "", device
+                    n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
                     state=None, attitude=None, force_torque=force_torque, rotor_vel=rotor_vel
@@ -183,9 +187,11 @@ def create(
         L: float, mixing_matrix: Array, KF: float, KM: float, device: Device
     ) -> SimConstants:
         """Create a default set of constants for the simulation."""
-        return SimConstants(
-            L=L, MIXING_MATRIX=jnp.array(mixing_matrix, device=device), KF=KF, KM=KM
-        )
+        L = jnp.array(L, device=device, dtype=jnp.float32)
+        mixing_matrix = jnp.array(mixing_matrix, device=device, dtype=jnp.float32)
+        KF = jnp.array(KF, device=device, dtype=jnp.float32)
+        KM = jnp.array(KM, device=device, dtype=jnp.float32)
+        return SimConstants(L=L, MIXING_MATRIX=mixing_matrix, KF=KF, KM=KM)
 
 
 @dataclass

crazyflow/sim/symbolic.py

@@ -224,73 +224,6 @@ def symbolic_attitude(dt: float, params: dict | None = None) -> SymbolicModel:
     return SymbolicModel(dynamics=dynamics, cost=cost, dt=dt)
 
 
-def symbolic_thrust(mass: float, J: NDArray, dt: float) -> SymbolicModel:
-    """Create symbolic (CasADi) models for dynamics, observation, and cost of a quadcopter.
-
-    This model is based on the analytical model of Luis, Carlos, and Jérôme Le Ny. "Design of a
-    trajectory tracking controller for a nanoquadcopter." arXiv preprint arXiv:1608.05786 (2016).
-
-    Returns:
-        The CasADi symbolic model of the environment.
-    """
-    # Define states.
-    z = MX.sym("z")
-    z_dot = MX.sym("z_dot")
-
-    # Set up the dynamics model for a 3D quadrotor.
-    nx, nu = 12, 4
-    Ixx, Iyy, Izz = J.diagonal()
-    J = cs.blockcat([[Ixx, 0.0, 0.0], [0.0, Iyy, 0.0], [0.0, 0.0, Izz]])
-    Jinv = cs.blockcat([[1.0 / Ixx, 0.0, 0.0], [0.0, 1.0 / Iyy, 0.0], [0.0, 0.0, 1.0 / Izz]])
-    gamma = KM / KF
-    # System state variables
-    x, y = MX.sym("x"), MX.sym("y")
-    x_dot, y_dot = MX.sym("x_dot"), MX.sym("y_dot")
-    phi, theta, psi = MX.sym("phi"), MX.sym("theta"), MX.sym("psi")
-    p, q, r = MX.sym("p"), MX.sym("q"), MX.sym("r")
-    # Rotation matrix transforming a vector in the body frame to the world frame. PyBullet Euler
-    # angles use the SDFormat for rotation matrices.
-    Rob = csRotXYZ(phi, theta, psi)
-    # Define state variables.
-    X = cs.vertcat(x, y, z, phi, theta, psi, x_dot, y_dot, z_dot, p, q, r)
-    # Define inputs.
-    f1, f2, f3, f4 = MX.sym("f1"), MX.sym("f2"), MX.sym("f3"), MX.sym("f4")
-    U = cs.vertcat(f1, f2, f3, f4)
-
-    # Defining the dynamics function.
-    # We are using the velocity of the base wrt to the world frame expressed in the world frame.
-    # Note that the reference expresses this in the body frame.
-    pos_ddot = Rob @ cs.vertcat(0, 0, f1 + f2 + f3 + f4) / mass - cs.vertcat(0, 0, GRAVITY)
-    pos_dot = cs.vertcat(x_dot, y_dot, z_dot)
-    # We use the spin directions (signs) from the mix matrix used in the simulation.
-    sx, sy, sz = SIGN_MIX_MATRIX[..., 0], SIGN_MIX_MATRIX[..., 1], SIGN_MIX_MATRIX[..., 2]
-    Mb = cs.vertcat(
-        ARM_LEN / cs.sqrt(2.0) * (sx[0] * f1 + sx[1] * f2 + sx[2] * f3 + sx[3] * f4),
-        ARM_LEN / cs.sqrt(2.0) * (sy[0] * f1 + sy[1] * f2 + sy[2] * f3 + sy[3] * f4),
-        gamma * (sz[0] * f1 + sz[1] * f2 + sz[2] * f3 + sz[3] * f4),
-    )
-    rate_dot = Jinv @ (Mb - (cs.skew(cs.vertcat(p, q, r)) @ J @ cs.vertcat(p, q, r)))
-    ang_dot = cs.blockcat(
-        [
-            [1, cs.sin(phi) * cs.tan(theta), cs.cos(phi) * cs.tan(theta)],
-            [0, cs.cos(phi), -cs.sin(phi)],
-            [0, cs.sin(phi) / cs.cos(theta), cs.cos(phi) / cs.cos(theta)],
-        ]
-    ) @ cs.vertcat(p, q, r)
-    X_dot = cs.vertcat(pos_dot, ang_dot, pos_ddot, rate_dot)
-
-    Y = cs.vertcat(x, x_dot, y, y_dot, z, z_dot, phi, theta, psi, p, q, r)
-
-    # Define cost (quadratic form).
-    Q, R = MX.sym("Q", nx, nx), MX.sym("R", nu, nu)
-    Xr, Ur = MX.sym("Xr", nx, 1), MX.sym("Ur", nu, 1)
-    cost_func = 0.5 * (X - Xr).T @ Q @ (X - Xr) + 0.5 * (U - Ur).T @ R @ (U - Ur)
-    # Define dynamics and cost dictionaries.
-    dynamics = {"dyn_eqn": X_dot, "obs_eqn": Y, "vars": {"X": X, "U": U}}
-    cost = {"cost_func": cost_func, "vars": {"X": X, "U": U, "Xr": Xr, "Ur": Ur, "Q": Q, "R": R}}
-    return SymbolicModel(dynamics=dynamics, cost=cost, dt=dt)
-
-
 def symbolic_from_sim(sim: Sim) -> SymbolicModel:
     """Create a symbolic model from a Sim instance.
 
@@ -308,8 +241,7 @@ def symbolic_from_sim(sim: Sim) -> SymbolicModel:
         case Control.attitude:
             return symbolic_attitude(1 / sim.control_freq)
         case Control.force_torque:
-            mass, J = sim.default_data.params.mass[0, 0], sim.default_data.params.J[0, 0]
-            return symbolic_force_torque(mass, J, 1 / sim.control_freq)
+            raise NotImplementedError("Symbolic model for force torque control is not implemented")
         case _:
             raise ValueError(f"Unsupported control type for symbolic model: {sim.control}")
 

crazyflow/utils.py

@@ -68,6 +68,11 @@ def to_device(data: Array, device: str) -> Array:
     return jnp.array(data, device=device)
 
 
+def named_tuple2device(data: T, device: str) -> T:
+    """Turn a named tuple into a jax array on the specified device."""
+    return data._replace(**{f: jnp.asarray(getattr(data, f), device=device) for f in data._fields})
+
+
 def enable_cache(
     cache_path: Path = Path("/tmp/jax_cache"),
     min_entry_size_bytes: int = -1,

examples/force_torque.py

@@ -11,7 +11,7 @@ def main():
     duration = 5.0
     fps = 60
 
-    cmd = np.ones((sim.n_worlds, sim.n_drones, 4))  # [fz, tx, ty, tz]
+    cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))  # [fz, tx, ty, tz]
     cmd[..., 0] = (MASS + 1e-4) * GRAVITY  # Plus a small margin to accelerate slightly
     for i in range(int(duration * sim.control_freq)):
         sim.force_torque_control(cmd)

examples/gradient.py

@@ -19,7 +19,7 @@ def step(cmd: NDArray, data: SimData) -> jax.Array:
         )
         data = sim_step(data, sim.freq // sim.control_freq)
         # Quadratic cost to reach 1m height
-        return (data.states.pos[0, 0, 2] - 1.0) ** 2 - 1e-3 * jnp.sum(cmd**2)
+        return (data.states.pos[0, 0, 2] - 1.0) ** 2
 
     step_grad = jax.jit(jax.grad(step))