BUG: control.mellinger: fix xp.all() targeting all worlds instead of per-world

crazyflow/control/mellinger/control.py

@@ -233,7 +233,8 @@ def attitude2force_torque(
     torque_pwm = xp.where((force_des > 0)[..., None], torque_pwm, 0.0)
     force_des_pwm = force2pwm(force_des / 4, thrust_max, pwm_max)
     pwms = force_torque_pwms2pwms(force_des_pwm, torque_pwm, mixing_matrix)
-    pwms = xp.where(xp.all(pwms == 0), 0.0, xp.clip(pwms, pwm_min, pwm_max))
+    idle = xp.all(pwms == 0, axis=-1, keepdims=True)
+    pwms = xp.where(idle, 0.0, xp.clip(pwms, pwm_min, pwm_max))
 
     # Info: The Mellinger controller in the firmware ends here. However, we enforce a standardized
     # interface in the simulation from states -> attitude -> force_torque. We therefore need this
@@ -304,7 +305,8 @@ def force_torque2rotor_vel(
     torque_forces = (torque * xp.asarray([1 / L, 1 / L, 1 / thrust2torque])) @ mixing_matrix
     motor_forces = (torque_forces + force) / 4
     # Clip motor forces on the thrust instead of PWM level.
-    motor_forces = xp.where(xp.all(force == 0), 0.0, xp.clip(motor_forces, thrust_min, thrust_max))
+    idle = xp.all(force == 0, axis=-1, keepdims=True)
+    motor_forces = xp.where(idle, 0.0, xp.clip(motor_forces, thrust_min, thrust_max))
     # Assume perfect battery compensation and calculate the desired motor speeds directly
     return motor_force2rotor_vel(motor_forces, rpm2thrust)
 

tests/unit/control/test_mellinger.py

@@ -196,6 +196,35 @@ def test_force_torque2rotor_vel_batch_consistency(drone: str):
             assert np.allclose(rpm_batch[i, j], rpm_s, atol=1e-5)
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize("drone", available_drones)
+def test_attitude2force_torque_batch_zero_thrust(drone: str):
+    # Drones with zero thrust must stay at zero force, independent of other drones
+    controller = parametrize(attitude2force_torque, drone)
+    quat = np.tile(np.array([0.0, 0.0, 0.0, 1.0]), (2, 1))
+    ang_vel = np.zeros((2, 3))
+    cmd = np.array([[0.1, 0.1, 0.1, 0.0], [0.1, 0.1, 0.1, 0.5]])
+    force_batch, torque_batch, _ = controller(quat, ang_vel, cmd)
+    force_single, torque_single, _ = controller(quat[0], ang_vel[0], cmd[0])
+    assert np.allclose(force_single, 0.0, atol=1e-6), "Drone with zero thrust must have zero force"
+    assert np.allclose(force_batch[0], force_single, atol=1e-6)
+    assert np.allclose(torque_batch[0], torque_single, atol=1e-6)
+    assert force_batch[1] > 0.0, "Drone thrust must be positive"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("drone", available_drones)
+def test_force_torque2rotor_vel_batch_zero_force(drone: str):
+    # Drones with zero desired force must not be clipped because other drones have non-zero force.
+    controller = parametrize(force_torque2rotor_vel, drone)
+    force = np.array([[0.0], [0.2]])
+    torque = np.zeros((2, 3))
+    rotor_vel_batch = controller(force, torque)
+    rotor_vel_single = controller(force[0], torque[0])
+    assert np.allclose(rotor_vel_batch[0], rotor_vel_single, rtol=1e-5)
+    assert np.all(rotor_vel_batch[1] > rotor_vel_single)
+
+
 # Symmetric force check