Add drone models to envs. Fix so_rpy_rotor_drag physics

Author: amacati

crazyflow/envs/drone_env.py

@@ -63,6 +63,7 @@ def __init__(
         num_envs: int = 1,
         max_episode_time: float = 10.0,
         physics: Literal["so_rpy", "first_principles"] | Physics = Physics.so_rpy,
+        drone_model: str = "cf2x_L250",
         freq: int = 500,
         device: str = "cpu",
         reset_randomization: Callable[[SimData, Array], SimData] | None = None,
@@ -85,7 +86,9 @@ def __init__(
         assert Physics(physics) in Physics, f"Invalid physics type {physics}"
 
         # Initialize the simulation
-        self.sim = Sim(n_worlds=num_envs, n_drones=1, device=device, physics=physics)
+        self.sim = Sim(
+            n_worlds=num_envs, n_drones=1, drone_model=drone_model, device=device, physics=physics
+        )
         assert self.sim.freq >= self.sim.control_freq, "Sim freq must be higher than control freq"
         if not self.sim.freq % self.freq == 0:
             # We can handle other cases, but it's not recommended

crazyflow/envs/figure_8_env.py

@@ -31,6 +31,7 @@ def __init__(
         num_envs: int = 1,
         max_episode_time: float = 10.0,
         physics: Literal["so_rpy", "first_principles"] | Physics = Physics.so_rpy,
+        drone_model: str = "cf2x_L250",
         freq: int = 500,
         device: str = "cpu",
     ):
@@ -43,13 +44,15 @@ def __init__(
             num_envs: Number of environments to run in parallel.
             max_episode_time: Maximum episode time in seconds.
             physics: Physics backend to use.
+            drone_model: Drone model of the environment.
             freq: Frequency of the simulation.
             device: Device to use for the simulation.
         """
         super().__init__(
             num_envs=num_envs,
             max_episode_time=max_episode_time,
             physics=physics,
+            drone_model=drone_model,
             freq=freq,
             device=device,
         )
@@ -60,9 +63,9 @@ def __init__(
         n_steps = int(np.ceil(trajectory_time * self.freq))
         t = np.linspace(0, 2 * np.pi, n_steps)
         radius = 1  # Radius for the circles
-        y = np.zeros_like(t)  # x is 0 everywhere
         x = radius * np.sin(t)  # Scale amplitude for 1-meter diameter
-        z = radius * np.sin(2 * t) + 1.2  # Scale amplitude for 1-meter diameter
+        y = np.zeros_like(t)  # x is 0 everywhere
+        z = radius / 2 * np.sin(2 * t) + 1  # Scale amplitude for 1-meter diameter
         self.trajectory = np.array([x, y, z]).T
 
         # Define trajectory sampling parameters

crazyflow/sim/physics.py

@@ -263,7 +263,7 @@ class SoRpyRotorDragData:
     """Force constant of the drone."""
     KM: Array  # (N, M, 1)
     """Torque constant of the drone."""
-    rotor_coef: Array  # (N, M, 1)
+    thrust_time_coef: Array  # (N, M, 1)
     """Rotor coefficient of the drone."""
     acc_coef: Array  # (N, M, 1)
     """Acceleration coefficient of the drone."""
@@ -277,8 +277,8 @@ class SoRpyRotorDragData:
     """Roll pitch yaw command coefficient of the drone."""
     drag_linear_coef: Array  # (N, M, 1)
     """Linear drag coefficient of the drone."""
-    drag_abs_linear_coef: Array  # (N, M, 1)
-    """Absolute linear drag coefficient of the drone."""
+    drag_square_coef: Array  # (N, M, 1)
+    """Square drag coefficient of the drone."""
 
     def create(
         n_worlds: int, n_drones: int, drone_model: str, device: Device
@@ -293,27 +293,27 @@ def create(
             J_inv=jnp.linalg.inv(J),
             KF=jnp.asarray(p.KF, device=device),
             KM=jnp.asarray(p.KM, device=device),
-            rotor_coef=jnp.asarray(p.rotor_coef, device=device),
+            thrust_time_coef=jnp.asarray(p.thrust_time_coef, device=device),
             acc_coef=jnp.asarray(p.acc_coef, device=device),
             cmd_f_coef=jnp.asarray(p.cmd_f_coef, device=device),
             rpy_coef=jnp.asarray(p.rpy_coef, device=device),
             rpy_rates_coef=jnp.asarray(p.rpy_rates_coef, device=device),
             cmd_rpy_coef=jnp.asarray(p.cmd_rpy_coef, device=device),
             drag_linear_coef=jnp.asarray(p.drag_linear_coef, device=device),
-            drag_abs_linear_coef=jnp.asarray(p.drag_abs_linear_coef, device=device),
+            drag_square_coef=jnp.asarray(p.drag_square_coef, device=device),
         )
 
 
 def so_rpy_rotor_drag_physics(data: SimData) -> SimData:
     """Compute the forces and torques from the so_rpy_rotor_drag physics model."""
     params: SoRpyRotorDragData = data.params
-    vel, _, acc, ang_acc, _ = so_rpy_rotor_drag_dynamics(
+    vel, _, acc, ang_acc, rotor_acc = so_rpy_rotor_drag_dynamics(
         pos=data.states.pos,
         quat=data.states.quat,
         vel=data.states.vel,
         ang_vel=data.states.ang_vel,
         cmd=data.controls.attitude.cmd,
-        rotor_vel=None,  # TODO: Add rotor_vel once the parameters are available
+        rotor_vel=data.states.rotor_vel,
         dist_f=data.states.force,
         dist_t=data.states.torque,
         mass=params.mass,
@@ -322,17 +322,16 @@ def so_rpy_rotor_drag_physics(data: SimData) -> SimData:
         J_inv=params.J_inv,
         KF=params.KF,
         KM=params.KM,
-        rotor_coef=params.rotor_coef,
+        thrust_time_coef=params.thrust_time_coef,
         acc_coef=params.acc_coef,
         cmd_f_coef=params.cmd_f_coef,
         rpy_coef=params.rpy_coef,
         rpy_rates_coef=params.rpy_rates_coef,
         cmd_rpy_coef=params.cmd_rpy_coef,
         drag_linear_coef=params.drag_linear_coef,
-        drag_abs_linear_coef=params.drag_abs_linear_coef,
+        drag_square_coef=params.drag_square_coef,
     )
-    # TODO: Add rotor_vel to the states_deriv
     states_deriv = data.states_deriv.replace(
-        vel=vel, ang_vel=data.states.ang_vel, acc=acc, ang_acc=ang_acc
+        vel=vel, ang_vel=data.states.ang_vel, acc=acc, ang_acc=ang_acc, rotor_acc=rotor_acc
     )
     return data.replace(states_deriv=states_deriv)

crazyflow/sim/sim.py

@@ -16,7 +16,6 @@
 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.control.control import Control, controllable
 from crazyflow.exception import ConfigError, NotInitializedError