[wip,broken] Temporary workaround for lsy_models compatibility

Author: amacati

benchmark/op_count.py

@@ -8,8 +8,8 @@ def main():
 
     compiled_reset = sim._reset.lower(sim.data, sim.default_data, None).compile()
     compiled_step = sim._step.lower(sim.data, 1).compile()
-    op_count_reset = compiled_reset.cost_analysis()[0]["flops"]
-    op_count_step = compiled_step.cost_analysis()[0]["flops"]
+    op_count_reset = compiled_reset.cost_analysis()["flops"]
+    op_count_step = compiled_step.cost_analysis()["flops"]
     print(f"Op counts:\n Reset: {op_count_reset}\n Step: {op_count_step}")
 
 

crazyflow/constants.py

@@ -8,6 +8,6 @@
 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]])
 SIGN_MIX_MATRIX: NDArray = np.sign(MIX_MATRIX)
-MASS: float = 0.027
+MASS: float = 0.03253
 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)

crazyflow/sim/structs.py

@@ -6,9 +6,11 @@
 import jax.numpy as jnp
 from flax.struct import dataclass, field
 from jax import Array, Device
+from lsy_models.utils.constants import Constants
 
 if TYPE_CHECKING:
     from mujoco.mjx import Data, Model
+    from numpy.typing import NDArray
 
 
 @dataclass
@@ -161,6 +163,16 @@ class SimParams:
     J_INV: Array  # (N, M, 3, 3)
     """Inverse of the inertia matrix of the drone."""
 
+    # TODO: Remove duplicate definition of constants. Move into constants from lsy_models
+    THRUST_TAU: float = field(pytree_node=False)
+    SIGN_MATRIX: NDArray = field(pytree_node=False)
+    L: float = field(pytree_node=False)
+    KF: float = field(pytree_node=False)
+    KM: float = field(pytree_node=False)
+    GRAVITY_VEC: NDArray = field(pytree_node=False)
+    MASS: float = field(pytree_node=False)
+    J_inv: NDArray = field(pytree_node=False)
+
     @staticmethod
     def create(
         n_worlds: int, n_drones: int, mass: float, J: Array, J_INV: Array, device: Device
@@ -170,7 +182,21 @@ def create(
         j, j_inv = jnp.array(J, device=device), jnp.array(J_INV, device=device)
         J = jnp.tile(j[None, None, :, :], (n_worlds, n_drones, 1, 1))
         J_INV = jnp.tile(j_inv[None, None, :, :], (n_worlds, n_drones, 1, 1))
-        return SimParams(mass=mass, J=J, J_INV=J_INV)
+        constants = Constants.from_config("cf2x_L250")
+
+        return SimParams(
+            mass=mass,
+            J=constants.J,
+            J_INV=J_INV,
+            THRUST_TAU=constants.THRUST_TAU,
+            SIGN_MATRIX=constants.SIGN_MATRIX,
+            L=constants.L,
+            KF=constants.KF,
+            KM=constants.KM,
+            GRAVITY_VEC=constants.GRAVITY_VEC,
+            MASS=constants.MASS,
+            J_inv=constants.J_inv,
+        )
 
 
 @dataclass