Fix thrust clipping

Author: ratheron

crazyflow/sim/sim.py

@@ -158,16 +158,21 @@ def __init__(
         # simulation pipeline.
         for name, fn in build_control_fns(self.control, self.dynamics):
             append_fn(self.step_pipeline, fn, name=name)
+        # Keep the rotor command and state (RPM or thrust, see ``rotor_vel_limits``) within the
+        # physical limits. We need both: the command clip models the motor saturation and is the
+        # only limit for models without a rotor state (so_rpy), the state clip catches states set
+        # from outside the pipeline and integrator overshoot.
+        lower, upper = rotor_vel_limits(self.dynamics, self.drone)
+        clip_cmd_fn = partial(clip_rotor_vel_cmd, lower=lower, upper=upper, dynamics=self.dynamics)
+        append_fn(self.step_pipeline, clip_cmd_fn, name="clip_rotor_vel_cmd")
         integrate_fn = select_integrate_fn(self.integrator, select_dynamics_fn(self.dynamics))
         append_fn(self.step_pipeline, integrate_fn, name="integration")
-        # Keep the rotor state (RPM or thrust, see ``rotor_vel_limits``) within its physical limits
-        lower, upper = rotor_vel_limits(self.dynamics, self.drone)
         clip_fn = partial(clip_rotor_vel, lower=lower, upper=upper)
         append_fn(self.step_pipeline, clip_fn, name="clip_rotor_vel")
-        append_fn(self.step_pipeline, increment_steps)
         # We never drop below -0.001 (drones can't pass through the floor). We use -0.001 to
         # enable checks for negative z sign
         append_fn(self.step_pipeline, clip_floor_pos)
+        append_fn(self.step_pipeline, increment_steps)
 
         self._reset = self.build_reset_fn()
         self._step = self.build_step_fn()
@@ -724,6 +729,18 @@ def clip_rotor_vel(data: SimData, lower: Array | float, upper: Array | float) ->
     return data.replace(states=data.states.replace(rotor_vel=rotor_vel))
 
 
+def clip_rotor_vel_cmd(
+    data: SimData, lower: Array | float, upper: Array | float, dynamics: Dynamics
+) -> SimData:
+    """Clip the rotor command (RPM for first principles, collective thrust otherwise)."""
+    if dynamics == Dynamics.first_principles:
+        rotor_vel = jnp.clip(data.controls.rotor_vel, lower, upper)
+        return data.replace(controls=data.controls.replace(rotor_vel=rotor_vel))
+    attitude = data.controls.attitude
+    cmd = attitude.cmd.at[..., -1].set(jnp.clip(attitude.cmd[..., -1], lower, upper))
+    return data.replace(controls=data.controls.replace(attitude=attitude.replace(cmd=cmd)))
+
+
 @partial(jax.jit, static_argnames="device")
 def seed_sim(data: SimData, seed: int, device: Device) -> SimData:
     """JIT-compiled seeding function."""

docs/examples/index.md

@@ -64,7 +64,7 @@ Because the simulator is built entirely from JAX operations, `jax.grad` can diff
 
 ## Gradients and state clipping
 
-Hard state clips such as the `clip_rotor_vel` stage zero the gradients while the state is saturated. This example ramps a motor command beyond the rotor limits and compares the rotor state and the gradient of the vertical acceleration w.r.t. the command for three options: the default clip, a straight-through clip (clipped forward pass, unclipped gradients), and no clip. Replacing the default clip with the straight-through variant can help gradient-based methods such as trajectory optimization or policy learning, which would otherwise receive zero gradients whenever the motors saturate.
+Hard state clips such as the `clip_rotor_vel` stage zero the gradients while the state is saturated. This example removes the `clip_rotor_vel_cmd` stage, ramps a motor command beyond the rotor limits, and compares the rotor state and the gradient of the vertical acceleration w.r.t. the command for three options: the default clip, a straight-through clip (clipped forward pass, unclipped gradients), and no clip. Replacing the default clip with the straight-through variant can help gradient-based methods such as trajectory optimization or policy learning, which would otherwise receive zero gradients whenever the motors saturate.
 
 <!-- notest: imported script, covered by tests/integration/test_examples.py -->
 ```{ .python notest }

docs/user-guide/pipelines.md

@@ -22,19 +22,23 @@ Both pipelines are constructed at `Sim` initialisation and compiled into a singl
 `sim.step_pipeline` contains multiple stages by default:
 
 1. **Control functions** — convert the staged command through the control hierarchy (state → attitude → force/torque → rotor velocities, depending on the selected mode)
-2. **Integrator** (`integration`) — advance the ODE one dynamics step (Euler, RK4, or symplectic Euler)
-3. **Rotor clip** (`clip_rotor_vel`) — clip the motor speeds (first principles) or the collective thrust (fitted models) to the physical limits of the motors
-4. **Step counter** (`increment_steps`) — increment `data.core.steps`
+2. **Rotor command clip** (`clip_rotor_vel_cmd`) — clip the commanded motor speeds (first principles) or collective thrust (so_rpy models) to the physical limits of the motors
+3. **Integrator** (`integration`) — advance the ODE one dynamics step (Euler, RK4, or symplectic Euler)
+4. **Rotor clip** (`clip_rotor_vel`) — clip the rotor state to the same limits
 5. **Floor clip** (`clip_floor_pos`) — prevent drones from passing through the floor
+6. **Step counter** (`increment_steps`) — increment `data.core.steps`
 
 ```pycon
 >>> from crazyflow.sim import Sim
 >>> sim = Sim()
 >>> print(tuple(sim.step_pipeline.keys()))
-('attitude_controller', 'force_torque_controller', 'integration', 'clip_rotor_vel', 'increment_steps', 'clip_floor_pos')
+('attitude_controller', 'force_torque_controller', 'clip_rotor_vel_cmd', 'integration', 'clip_rotor_vel', 'clip_floor_pos', 'increment_steps')
 
 ```
 
+!!! note "Why two rotor clips?"
+    `clip_rotor_vel_cmd` models the motor saturation and is the only limit for models without a rotor state (`so_rpy`). `clip_rotor_vel` catches states set from outside the default pipeline, such as randomizations, and integrator overshoot. Since it runs after the integration, higher order integrators like RK4 can transiently exceed the limits within a step.
+
 ## The reset pipeline
 
 `sim.reset_pipeline` holds a single `reset` stage that restores `SimData` to the default state. Every stage appended after it runs in order on the restored data. Each reset stage has the signature `(data: SimData, default_data: SimData, mask: Array | None) -> SimData`. The `default_data` argument holds the freshly-restored default state, which is useful for selectively reverting fields.
@@ -131,9 +135,6 @@ remove_fn(sim.step_pipeline, "clip_floor_pos")
 sim.build_step_fn()
 ```
 
-!!! note
-    The rotor limits are enforced by the `clip_rotor_vel` stage, which clips the rotor state after each integration step. Higher order integration methods like RK4 evaluate the dynamics at intermediate, unclipped states, so the `rotor_vel` seen by the model can transiently exceed the limits within a single step.
-
 ## Writing a custom stage
 
 A step pipeline function must have the signature `(SimData) -> SimData`. A reset pipeline function must have the signature `(SimData, SimData, Array | None) -> SimData` where the second argument is the default (freshly-restored) data. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled.

examples/jax/gradient_clipping.py

@@ -47,6 +47,8 @@ def acc_z(cmd: jax.Array, data: SimData) -> tuple[jax.Array, SimData]:
 def main(plot: bool = False):
     sim = Sim(control=Control.rotor_vel)
     lower, upper = rotor_vel_limits(sim.dynamics, sim.drone)
+    # Let the command reach the dynamics unclipped so that only the state clip affects the gradients
+    remove_fn(sim.step_pipeline, "clip_rotor_vel_cmd")
     # Start in the air so that the drone never reaches the floor, where the floor clipping would
     # zero the velocity and kill the gradients (see gradient.py)
     sim.data = sim.data.replace(

tests/unit/test_sim.py

@@ -450,7 +450,7 @@ def test_floor_penetration(dynamics: Dynamics):
 @pytest.mark.unit
 @pytest.mark.parametrize("integrator", Integrator)
 def test_rotor_vel_clip(integrator: Integrator):
-    """Test that the first-principles rotor state saturates at its physical limits."""
+    """Test that the first principles rotor command and rotor state are clipped to the limits."""
     sim = Sim(
         dynamics=Dynamics.first_principles,
         control=Control.rotor_vel,
@@ -460,11 +460,22 @@ def test_rotor_vel_clip(integrator: Integrator):
     lower, upper = rotor_vel_limits(Dynamics.first_principles, sim.drone)
     assert 0.0 < lower < upper
 
-    # States outside the limits are clipped after the integration
     for value, target in ((2 * upper, upper), (-upper, lower)):
-        states = sim.data.states.replace(rotor_vel=jnp.full_like(sim.data.states.rotor_vel, value))
+        # Commands outside the limits are clipped before they enter the dynamics
+        sim.reset()
         sim.rotor_vel_control(np.full((1, 1, 4), target))
+        sim.step()
+        rotor_vel_ref = sim.data.states.rotor_vel
+        sim.reset()
+        sim.rotor_vel_control(np.full((1, 1, 4), value))
+        sim.step()
+        assert jnp.all(sim.data.controls.rotor_vel == target)
+        assert jnp.allclose(sim.data.states.rotor_vel, rotor_vel_ref)
+        # States outside the limits are clipped after the integration
+        sim.reset()
+        states = sim.data.states.replace(rotor_vel=jnp.full_like(sim.data.states.rotor_vel, value))
         sim.data = sim.data.replace(states=states)
+        sim.rotor_vel_control(np.full((1, 1, 4), target))
         sim.step()
         assert jnp.all(sim.data.states.rotor_vel >= lower)
         assert jnp.all(sim.data.states.rotor_vel <= upper)
@@ -477,15 +488,29 @@ def test_rotor_vel_clip(integrator: Integrator):
 )
 @pytest.mark.parametrize("integrator", Integrator)
 def test_thrust_clip(dynamics: Dynamics, integrator: Integrator):
-    """Test that the thrust state of every so_rpy model is clipped to its physical limits."""
+    """Test that the so_rpy thrust command and thrust state are clipped to the limits."""
     sim = Sim(dynamics=dynamics, control=Control.attitude, integrator=integrator, device="cpu")
     lower, upper = rotor_vel_limits(dynamics, sim.drone)
     assert 0.0 < lower < upper
 
     for value, target in ((2 * upper, upper), (-upper, lower)):
-        states = sim.data.states.replace(rotor_vel=jnp.full_like(sim.data.states.rotor_vel, value))
+        # Commands outside the limits are clipped before they enter the dynamics. so_rpy has no
+        # thrust state and applies the command directly, so we compare the velocity as well
+        sim.reset()
         sim.attitude_control(np.array([[[0.0, 0.0, 0.0, target]]]))
+        sim.step()
+        vel_ref, rotor_vel_ref = sim.data.states.vel, sim.data.states.rotor_vel
+        sim.reset()
+        sim.attitude_control(np.array([[[0.0, 0.0, 0.0, value]]]))
+        sim.step()
+        assert jnp.all(sim.data.controls.attitude.cmd[..., -1] == target)
+        assert jnp.allclose(sim.data.states.vel, vel_ref)
+        assert jnp.allclose(sim.data.states.rotor_vel, rotor_vel_ref)
+        # States outside the limits are clipped after the integration
+        sim.reset()
+        states = sim.data.states.replace(rotor_vel=jnp.full_like(sim.data.states.rotor_vel, value))
         sim.data = sim.data.replace(states=states)
+        sim.attitude_control(np.array([[[0.0, 0.0, 0.0, target]]]))
         sim.step()
         assert jnp.all(sim.data.states.rotor_vel >= lower)
         assert jnp.all(sim.data.states.rotor_vel <= upper)