Batch of fixes and changes

- Remove constants
- Rename structs to data
- Add drag and rotor physics
- Decrease interface tolerances
- Add gradient tests
- Bump drone-models version

Author: amacati

crazyflow/control/control.py

@@ -59,15 +59,3 @@ def controllable(step: Array, freq: int, control_steps: Array, control_freq: int
         updated.
     """
     return ((step - control_steps) >= (freq / control_freq)) | (control_steps == -1)
-
-
-# TODO: Remove these constants
-KF: float = 3.16e-10
-PWM2RPM_SCALE: float = 0.2685
-PWM2RPM_CONST: float = 4070.3
-MIN_PWM: float = 20000
-MAX_PWM: float = 65535
-MIN_RPM: float = PWM2RPM_SCALE * MIN_PWM + PWM2RPM_CONST
-MAX_RPM: float = PWM2RPM_SCALE * MAX_PWM + PWM2RPM_CONST
-MIN_THRUST: float = KF * MIN_RPM**2
-MAX_THRUST: float = KF * MAX_RPM**2

crazyflow/envs/drone_env.py

@@ -5,33 +5,37 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+from drone_controllers.mellinger.params import ForceTorqueParams
 from gymnasium import spaces
 from gymnasium.vector import AutoresetMode, VectorEnv
 from gymnasium.vector.utils import batch_space
 from jax import Array
 from numpy.typing import NDArray
 
-from crazyflow.control.control import MAX_THRUST, MIN_THRUST, Control
+from crazyflow.control.control import Control
 from crazyflow.sim import Sim
+from crazyflow.sim.data import SimData
 from crazyflow.sim.physics import Physics
-from crazyflow.sim.structs import SimData
 from crazyflow.utils import leaf_replace
 
 
-def action_space(control_type: Control) -> spaces.Box:
+def action_space(control_type: Control, drone_model: str) -> spaces.Box:
     """Select the appropriate action space for a given control type.
 
     Args:
         control_type: The desired control mode.
+        drone_model: Drone model of the environment.
 
     Returns:
         The action space.
     """
     match control_type:
         case Control.attitude:
+            params = ForceTorqueParams.load(drone_model)
+            thrust_min, thrust_max = params.thrust_min * 4, params.thrust_max * 4
             return 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),
+                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),
             )
         case Control.force_torque:
             return spaces.Box(-1.0, 1.0, shape=(6,))
@@ -99,7 +103,7 @@ def __init__(
         self._marked_for_reset = jnp.zeros((self.sim.n_worlds), dtype=jnp.bool_, device=self.device)
 
         # Define action and observation spaces
-        self.single_action_space = action_space(self.sim.control)
+        self.single_action_space = action_space(self.sim.control, self.sim.drone_model)
         self.action_space = batch_space(self.single_action_space, self.sim.n_worlds)
         self.single_observation_space = spaces.Dict(
             {

crazyflow/envs/figure_8_env.py

@@ -8,8 +8,8 @@
 from jax import Array
 
 from crazyflow.envs.drone_env import DroneEnv
+from crazyflow.sim.data import SimData
 from crazyflow.sim.physics import Physics
-from crazyflow.sim.structs import SimData
 from crazyflow.sim.visualize import draw_line, draw_points
 from crazyflow.utils import leaf_replace
 

crazyflow/envs/landing_env.py

@@ -9,8 +9,8 @@
 from jax import Array
 
 from crazyflow.envs.drone_env import DroneEnv
+from crazyflow.sim.data import SimState
 from crazyflow.sim.physics import Physics
-from crazyflow.sim.structs import SimState
 
 
 class LandingEnv(DroneEnv):

crazyflow/envs/reach_pos_env.py

@@ -10,8 +10,8 @@
 from jax import Array
 
 from crazyflow.envs.drone_env import DroneEnv
+from crazyflow.sim.data import SimData
 from crazyflow.sim.physics import Physics
-from crazyflow.sim.structs import SimData
 from crazyflow.utils import leaf_replace
 
 

crazyflow/randomize/randomize.py

@@ -3,7 +3,7 @@
 from jax import Array
 
 from crazyflow.sim import Sim
-from crazyflow.sim.structs import SimData
+from crazyflow.sim.data import SimData
 from crazyflow.utils import leaf_replace
 
 

crazyflow/sim/structs.py crazyflow/sim/data.pycrazyflow/sim/structs.py renamed to crazyflow/sim/data.py

@@ -13,7 +13,13 @@
     MellingerForceTorqueData,
     MellingerStateData,
 )
-from crazyflow.sim.physics import FirstPrinciplesData, Physics, SoRpyData
+from crazyflow.sim.physics import (
+    FirstPrinciplesData,
+    Physics,
+    SoRpyData,
+    SoRpyRotorData,
+    SoRpyRotorDragData,
+)
 
 
 @dataclass
@@ -173,6 +179,10 @@ def create(
                 return FirstPrinciplesData.create(n_worlds, n_drones, drone_model, device)
             case Physics.so_rpy:
                 return SoRpyData.create(n_worlds, n_drones, drone_model, device)
+            case Physics.so_rpy_rotor:
+                return SoRpyRotorData.create(n_worlds, n_drones, drone_model, device)
+            case Physics.so_rpy_rotor_drag:
+                return SoRpyRotorDragData.create(n_worlds, n_drones, drone_model, device)
             case _:
                 raise ValueError(f"Physics mode {physics} not implemented")
 

crazyflow/sim/integration.py

@@ -8,7 +8,7 @@
 from jax.numpy import vectorize
 from jax.scipy.spatial.transform import Rotation as R
 
-from crazyflow.sim.structs import SimData
+from crazyflow.sim.data import SimData
 
 
 class Integrator(str, Enum):