[WIP, broken] Switch to mellinger from drone-models

Author: amacati

crazyflow/constants.py

@@ -6,9 +6,16 @@
 
 # Drone constants
 ARM_LEN: float = 0.0325 * np.sqrt(2)
-MIX_MATRIX: NDArray = np.array([[-0.5, -0.5, -1], [-0.5, 0.5, 1], [0.5, 0.5, -1], [0.5, -0.5, 1]])
+# fmt: off
+MIX_MATRIX: NDArray = np.array([[-0.5, -0.5, -1],
+                                [-0.5,  0.5,  1],
+                                [ 0.5,  0.5, -1],
+                                [ 0.5, -0.5,  1]])
+# fmt: on
 SIGN_MIX_MATRIX: NDArray = np.sign(MIX_MATRIX)
 # Crazyflie 2.1 mass as measured in the lab with battery included
 MASS: float = 0.033
 J: NDArray = np.array([[2.3951e-5, 0, 0], [0, 2.3951e-5, 0], [0, 0, 3.2347e-5]])
 J_INV: NDArray = np.linalg.inv(J)
+KF: float = 8.701227710666256e-10
+KM: float = 7.94e-12

crazyflow/control/control.py

@@ -19,7 +19,7 @@
 from jax import Array
 from jax.scipy.spatial.transform import Rotation as R
 
-from crazyflow.constants import GRAVITY, MASS, MIX_MATRIX
+from crazyflow.constants import MIX_MATRIX
 
 
 class Control(str, Enum):
@@ -41,8 +41,8 @@ class Control(str, Enum):
     Note:
         Recommended frequency is >=100 Hz.
     """
-    thrust = "thrust"
-    """Thrust control takes [thrust1, thrust2, thrust3, thrust4] for each drone motor.
+    force_torque = "force_torque"
+    """Force and torque control takes [fx, fy, fz, tx, ty, tz].
 
     Note:
         Recommended frequency is >=500 Hz.
@@ -90,38 +90,6 @@ def controllable(step: Array, freq: int, control_steps: Array, control_freq: int
 THRUST_CURVE_C: float = 0.0209
 
 
-@partial(jnp.vectorize, signature="(3),(3),(4),(3),(3),(1),(3)->(4),(3)", excluded=[7])
-def state2attitude(
-    pos: Array,
-    vel: Array,
-    quat: Array,
-    des_pos: Array,
-    des_vel: Array,
-    des_yaw: Array,
-    i_error: Array,
-    dt: float,
-) -> tuple[Array, Array]:
-    """Compute the next desired collective thrust and roll/pitch/yaw of the drone."""
-    pos_error, vel_error = des_pos - pos, des_vel - vel
-    # Update integral error
-    i_error = jnp.clip(i_error + pos_error * dt, -I_F_RANGE, I_F_RANGE)
-    # Compute target thrust
-    thrust = P_F * pos_error + I_F * i_error + D_F * vel_error
-    thrust = thrust.at[2].add(MASS * GRAVITY)
-    # Update z_axis to the current orientation of the drone
-    z_axis = R.from_quat(quat).as_matrix()[:, 2]
-    # Project the thrust onto the z-axis
-    thrust_desired = jnp.clip(thrust @ z_axis, 0.3 * MASS * GRAVITY, 1.8 * MASS * GRAVITY)
-    # Update the desired z-axis
-    z_axis = thrust / jnp.linalg.norm(thrust)
-    yaw_axis = jnp.concatenate([jnp.cos(des_yaw), jnp.sin(des_yaw), jnp.array([0.0])])
-    y_axis = jnp.cross(z_axis, yaw_axis)
-    y_axis = y_axis / jnp.linalg.norm(y_axis)
-    x_axis = jnp.cross(y_axis, z_axis)
-    euler_desired = R.from_matrix(jnp.vstack([x_axis, y_axis, z_axis]).T).as_euler("xyz")
-    return jnp.concatenate([jnp.atleast_1d(thrust_desired), euler_desired]), i_error
-
-
 @partial(jnp.vectorize, signature="(4),(4),(3),(3)->(4),(3)", excluded=[4])
 def attitude2rpm(
     controls: Array, quat: Array, last_rpy: Array, rpy_err_i: Array, dt: float

crazyflow/control/mellinger.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
 import jax.numpy as jnp
-from drone_models.controller.mellinger import MellingerStateParams
+from drone_models.controller.mellinger import (
+    MellingerAttitudeParams,
+    MellingerForceTorqueParams,
+    MellingerStateParams,
+)
 from flax.struct import dataclass, field
 from jax import Array, Device
 
@@ -14,6 +18,8 @@ class MellingerStateData:
     A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
     We currently do not use the acceleration and angle rate components. This is subject to change.
     """
+    staged_cmd: Array  # (N, M, 13)
+    """Staging buffer to store the most recent command until the next controller tick."""
     steps: Array  # (N, 1)
     """Last simulation steps that the state control command was applied."""
     freq: int = field(pytree_node=False)
@@ -29,39 +35,90 @@ def create(
     ) -> 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)
+        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)
         return MellingerStateData(
-            cmd=cmd, steps=steps, freq=freq, pos_err_i=pos_err_i, params=params
+            cmd=cmd, staged_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.
+@dataclass
+class MellingerAttitudeData:
+    cmd: Array  # (N, M, 4)
+    """Full attitude control command for the drone.
+
+    A command consists of [roll, pitch, yaw, collective thrust].
+    """
+    staged_cmd: Array  # (N, M, 4)
+    """Staging buffer to store the most recent command until the next controller tick."""
+    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."""
+    r_int_error: Array  # (N, M, 3)
+    """Integral errors of the attitude control command."""
+    last_ang_vel: Array  # (N, M, 3)
+    """Last angular velocity of the drone."""
+    # 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.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)
+        return MellingerAttitudeData(
+            cmd=cmd,
+            staged_cmd=cmd,
+            steps=steps,
+            freq=freq,
+            r_int_error=zeros_3d,
+            last_ang_vel=zeros_3d,
+            params=params,
+        )
+
 
-#     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
+@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.
 
-#     @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
-#         )
+    A command consists of [tx, ty, tz].
+    """
+    staged_cmd_force: Array  # (N, M, 1)
+    staged_cmd_torque: Array  # (N, M, 3)
+    """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
+
+    @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)
+        steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
+        params = MellingerForceTorqueParams.load(drone_model)
+        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,
+        )

crazyflow/envs/drone_env.py

@@ -33,8 +33,8 @@ def action_space(control_type: Control) -> spaces.Box:
                 np.array([4 * MIN_THRUST, -np.pi / 2, -np.pi / 2, -np.pi / 2], dtype=np.float32),
                 np.array([4 * MAX_THRUST, np.pi / 2, np.pi / 2, np.pi / 2], dtype=np.float32),
             )
-        case Control.thrust:
-            return spaces.Box(MIN_THRUST, MAX_THRUST, shape=(4,))
+        case Control.force_torque:
+            return spaces.Box(-1.0, 1.0, shape=(6,))
         case _:
             raise ValueError(f"Invalid control type {control_type}")
 
@@ -128,8 +128,8 @@ def _apply_action(self, action: Array):
                 raise NotImplementedError("State control currently not supported")
             case Control.attitude:
                 self.sim.attitude_control(action)
-            case Control.thrust:
-                self.sim.thrust_control(action)
+            case Control.force_torque:
+                self.sim.force_torque_control(action)
             case _:
                 raise ValueError(f"Invalid control type {self.sim.control}")
 

crazyflow/randomize/randomize.py

@@ -50,4 +50,4 @@ def _randomize_mass_params(data: SimData, mass: Array, mask: Array | None = None
 @jax.jit
 def _randomize_inertia_params(data: SimData, new_j: Array, mask: Array | None = None) -> SimData:
     new_j_inv = jnp.linalg.inv(new_j)
-    return data.replace(params=leaf_replace(data.params, mask, J=new_j, J_INV=new_j_inv))
+    return data.replace(params=leaf_replace(data.params, mask, J=new_j, J_inv=new_j_inv))

crazyflow/sim/physics.py

@@ -48,7 +48,7 @@ def surrogate_identified_collective_wrench(
         J: The drone's inertia matrix.
         dt: The simulation time step.
     """
-    collective_thrust, attitude = controls[0], controls[1:]
+    attitude, collective_thrust = controls[:3], controls[3]
     rot = R.from_quat(quat)
     thrust = rot.apply(jnp.array([0, 0, collective_thrust]))
     drift = rot.apply(jnp.array([0, 0, 1]))
@@ -84,9 +84,9 @@ def collective_force2acceleration(force: Array, mass: Array) -> Array:
 
 
 @partial(vectorize, signature="(3),(4),(3,3)->(3)")
-def collective_torque2ang_vel_deriv(torque: Array, quat: Array, J_INV: Array) -> Array:
+def collective_torque2ang_vel_deriv(torque: Array, quat: Array, J_inv: Array) -> Array:
     """Convert torques to ang_vel_deriv."""
-    return J_INV @ R.from_quat(quat).apply(torque, inverse=True)
+    return J_inv @ R.from_quat(quat).apply(torque, inverse=True)
 
 
 @partial(vectorize, signature="(4),(4),(3),(3,3)->(3),(3)")