Fix incorrect torque conversion in collective_torque2rpy_rates_deriv. Add tests.

Author: amacati

crazyflow/sim/physics.py

@@ -87,7 +87,8 @@ def collective_force2acceleration(force: Array, mass: Array) -> Array:
 @partial(vectorize, signature="(3),(4),(3,3)->(3)")
 def collective_torque2rpy_rates_deriv(torque: Array, quat: Array, J_INV: Array) -> Array:
     """Convert torques to rpy_rates_deriv."""
-    return R.from_quat(quat).apply(J_INV @ torque)
+    rot = R.from_quat(quat)
+    return rot.apply(J_INV @ rot.apply(torque, inverse=True))
 
 
 @partial(vectorize, signature="(4),(4),(3),(3,3)->(3),(3)")

tests/integration/test_interfaces.py

@@ -1,5 +1,6 @@
-import jax.numpy as jnp
+import numpy as np
 import pytest
+from scipy.spatial.transform import Rotation as R
 
 from crazyflow.control.control import Control, state2attitude
 from crazyflow.sim import Physics, Sim
@@ -11,43 +12,43 @@ def test_state_interface(physics: Physics):
     sim = Sim(physics=physics, control=Control.state)
 
     # Simple P controller for attitude to reach target height
-    cmd = jnp.zeros((1, 1, 13), dtype=jnp.float32)
-    cmd = cmd.at[0, 0, 2].set(1.0)  # Set z position target to 1.0
+    cmd = np.zeros((1, 1, 13), dtype=np.float32)
+    cmd[0, 0, 2] = 1.0  # Set z position target to 1.0
 
     for _ in range(int(2 * sim.control_freq)):  # Run simulation for 2 seconds
         sim.state_control(cmd)
         sim.step(sim.freq // sim.control_freq)
-        if jnp.linalg.norm(sim.data.states.pos[0, 0] - jnp.array([0.0, 0.0, 1.0])) < 0.1:
+        if np.linalg.norm(sim.data.states.pos[0, 0] - np.array([0.0, 0.0, 1.0])) < 0.1:
             break
 
     # Check if drone reached target position
-    distance = jnp.linalg.norm(sim.data.states.pos[0, 0] - jnp.array([0.0, 0.0, 1.0]))
+    distance = np.linalg.norm(sim.data.states.pos[0, 0] - np.array([0.0, 0.0, 1.0]))
     assert distance < 0.1, f"Failed to reach target height with {physics} physics"
 
 
 @pytest.mark.integration
 @pytest.mark.parametrize("physics", Physics)
 def test_attitude_interface(physics: Physics):
     sim = Sim(physics=physics, control=Control.attitude)
-    target_pos = jnp.array([0.0, 0.0, 1.0])
+    target_pos = np.array([0.0, 0.0, 1.0])
 
-    i_error = jnp.zeros((1, 1, 3))
+    i_error = np.zeros((1, 1, 3))
 
     for _ in range(int(2 * sim.control_freq)):  # Run simulation for 2 seconds
         pos, vel, quat = sim.data.states.pos, sim.data.states.vel, sim.data.states.quat
-        des_pos = jnp.array([[[0, 0, 1.0]]])
+        des_pos = np.array([[[0, 0, 1.0]]])
         dt = 1 / sim.data.controls.attitude_freq
         cmd, i_error = state2attitude(
-            pos, vel, quat, des_pos, jnp.zeros((1, 1, 3)), jnp.zeros((1, 1, 1)), i_error, dt
+            pos, vel, quat, des_pos, np.zeros((1, 1, 3)), np.zeros((1, 1, 1)), i_error, dt
         )
         sim.attitude_control(cmd)
         sim.step(sim.freq // sim.control_freq)
-        if jnp.linalg.norm(sim.data.states.pos[0, 0] - target_pos) < 0.1:
+        if np.linalg.norm(sim.data.states.pos[0, 0] - target_pos) < 0.1:
             break
 
     # Check if drone maintained hover position
     dpos = sim.data.states.pos[0, 0] - target_pos
-    distance = jnp.linalg.norm(dpos)
+    distance = np.linalg.norm(dpos)
     assert distance < 0.1, f"Failed to maintain hover with {physics} ({dpos})"
 
 
@@ -56,17 +57,40 @@ def test_attitude_interface(physics: Physics):
 def test_swarm_control(physics: Physics):
     n_worlds, n_drones = 2, 3
     sim = Sim(n_worlds=n_worlds, n_drones=n_drones, physics=physics, control=Control.state)
-    target_pos = sim.data.states.pos + jnp.array([0.2, 0.2, 0.2])
+    target_pos = sim.data.states.pos + np.array([0.2, 0.2, 0.2])
 
-    cmd = jnp.zeros((n_worlds, n_drones, 13))
-    cmd = cmd.at[..., :3].set(target_pos)
+    cmd = np.zeros((n_worlds, n_drones, 13))
+    cmd[..., :3] = target_pos
 
     for _ in range(int(3 * sim.control_freq)):  # Run simulation for 2 seconds
         sim.state_control(cmd)
         sim.step(sim.freq // sim.control_freq)
-        if jnp.linalg.norm(sim.data.states.pos[0, 0] - target_pos) < 0.1:
+        if np.linalg.norm(sim.data.states.pos[0, 0] - target_pos) < 0.1:
             break
 
     # Check if drone maintained hover position
-    max_dist = jnp.max(jnp.linalg.norm(sim.data.states.pos - target_pos, axis=-1))
+    max_dist = np.max(np.linalg.norm(sim.data.states.pos - target_pos, axis=-1))
     assert max_dist < 0.05, f"Failed to reach target, max dist: {max_dist}"
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("physics", Physics)
+def test_yaw_rotation(physics: Physics):
+    if physics == Physics.sys_id:  # TODO: Remove once yaw is supported for sys_id
+        pytest.skip("Yaw != 0 currently not supported for sys_id")
+
+    sim = Sim(physics=physics, control=Control.state)
+    sim.reset()
+
+    cmd = np.zeros((sim.n_worlds, sim.n_drones, 13))
+    cmd[..., :3] = 0.2
+    cmd[..., 9] = np.pi / 2  # Test if the drone can rotate in yaw
+
+    sim.state_control(cmd)
+    sim.step(sim.freq * 2)
+    pos = sim.data.states.pos[0, 0]
+    rot = R.from_quat(sim.data.states.quat[0, 0])
+    distance = np.linalg.norm(pos - np.array([0.2, 0.2, 0.2]))
+    assert distance < 0.1, f"Failed to reach target, distance: {distance}"
+    angle = rot.as_euler("xyz")[2]
+    assert np.abs(angle - np.pi / 2) < 0.1, f"Failed to rotate in yaw, angle: {angle}"