Update drone-controllers to 0.2.0

crazyflow/control/mellinger.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 import jax.numpy as jnp
-from drone_controllers.mellinger.params import AttitudeParams, ForceTorqueParams, StateParams
+from drone_controllers.core import load_params
+from drone_controllers.mellinger import (
+    attitude2force_torque,
+    force_torque2rotor_vel,
+    state2attitude,
+)
 from flax.struct import dataclass, field
 from jax import Array, Device
 
-from crazyflow.utils import named_tuple2device
-
 
 @dataclass
 class MellingerStateData:
@@ -25,7 +28,7 @@ class MellingerStateData:
     pos_err_i: Array  # (N, M, 3)
     """Integral errors of the state control command."""
     # Parameters for the state controller
-    params: StateParams
+    params: dict[str, Array]
 
     @staticmethod
     def create(
@@ -35,7 +38,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 = named_tuple2device(StateParams.load(drone_model), device)
+        params = load_params(state2attitude, drone_model, xp=jnp, device=device)
         return MellingerStateData(
             cmd=cmd, staged_cmd=cmd, steps=steps, freq=freq, pos_err_i=pos_err_i, params=params
         )
@@ -59,7 +62,7 @@ class MellingerAttitudeData:
     last_ang_vel: Array  # (N, M, 3)
     """Last angular velocity of the drone."""
     # Parameters for the attitude controller
-    params: AttitudeParams
+    params: dict[str, Array]
 
     @staticmethod
     def create(
@@ -69,7 +72,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 = named_tuple2device(AttitudeParams.load(drone_model), device)
+        params = load_params(attitude2force_torque, drone_model, xp=jnp, device=device)
         return MellingerAttitudeData(
             cmd=cmd,
             staged_cmd=cmd,
@@ -95,15 +98,15 @@ class MellingerForceTorqueData:
     freq: int = field(pytree_node=False)
     """Frequency of the force and torque control command."""
     # Parameters for the force and torque controller
-    params: ForceTorqueParams
+    params: dict[str, Array]
 
     @staticmethod
     def create(
         n_worlds: int, n_drones: int, freq: int, drone_model: str, device: Device
     ) -> MellingerForceTorqueData:
         zero_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = named_tuple2device(ForceTorqueParams.load(drone_model), device)
+        params = load_params(force_torque2rotor_vel, drone_model, xp=jnp, device=device)
         return MellingerForceTorqueData(
             cmd=zero_4d, staged_cmd=zero_4d, steps=steps, freq=freq, params=params
         )

crazyflow/envs/drone_env.py

@@ -5,7 +5,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from drone_controllers.mellinger.params import ForceTorqueParams
+from drone_controllers.core import load_params
+from drone_controllers.mellinger import force_torque2rotor_vel
 from gymnasium import spaces
 from gymnasium.vector import AutoresetMode, VectorEnv
 from gymnasium.vector.utils import batch_space
@@ -31,8 +32,8 @@ def action_space(control_type: Control, drone_model: str) -> spaces.Box:
     """
     match control_type:
         case Control.attitude:
-            params = ForceTorqueParams.load(drone_model)
-            thrust_min, thrust_max = params.thrust_min * 4, params.thrust_max * 4
+            params = load_params(force_torque2rotor_vel, drone_model)
+            thrust_min, thrust_max = params["thrust_min"] * 4, params["thrust_max"] * 4
             return spaces.Box(
                 np.array([-np.pi / 2, -np.pi / 2, -np.pi / 2, thrust_min], dtype=np.float32),
                 np.array([np.pi / 2, np.pi / 2, np.pi / 2, thrust_max], dtype=np.float32),

crazyflow/sim/data.py

@@ -97,7 +97,7 @@ class ControlData(typing.Protocol):
     freq: int
     """Frequency of the state control command."""
     # Parameters for the controller
-    params: tuple[typing.Any, ...]
+    params: dict[str, Array]
 
 
 @dataclass

crazyflow/sim/sim.py

@@ -533,11 +533,10 @@ def step_state_controller(data: SimData) -> SimData:
         states.pos,
         states.quat,
         states.vel,
-        states.ang_vel,
         state_ctrl.cmd,
         ctrl_errors=(state_ctrl.pos_err_i,),
         ctrl_freq=state_ctrl.freq,
-        **state_ctrl.params._asdict(),
+        **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)
@@ -552,15 +551,13 @@ def step_attitude_controller(data: SimData) -> SimData:
     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(
-        states.pos,
         states.quat,
-        states.vel,
         states.ang_vel,
         attitude_ctrl.cmd,
         ctrl_errors=(attitude_ctrl.r_int_error,),
         ctrl_freq=attitude_ctrl.freq,
         prev_ang_vel=attitude_ctrl.last_ang_vel,
-        **attitude_ctrl.params._asdict(),
+        **attitude_ctrl.params,
     )
     attitude_ctrl = leaf_replace(
         attitude_ctrl,
@@ -592,7 +589,7 @@ def step_force_torque_controller(data: SimData) -> SimData:
     mask = controllable(data.core.steps, data.core.freq, ft_ctrl.steps, ft_ctrl.freq)
     ft_ctrl = leaf_replace(ft_ctrl, mask, cmd=ft_ctrl.staged_cmd)
     rotor_vel = force_torque2rotor_vel(
-        ft_ctrl.cmd[..., [0]], ft_ctrl.cmd[..., 1:], **ft_ctrl.params._asdict()
+        ft_ctrl.cmd[..., [0]], ft_ctrl.cmd[..., 1:], **ft_ctrl.params
     )
     ft_ctrl = leaf_replace(ft_ctrl, mask, steps=data.core.steps)
     return data.replace(controls=data.controls.replace(rotor_vel=rotor_vel, force_torque=ft_ctrl))

crazyflow/utils.py

@@ -68,11 +68,6 @@ 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,