RTC adjustments. Bug fix & Alex Soare optimization

src/lerobot/policies/pi0/modeling_pi0.py

@@ -842,6 +842,7 @@ def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
                     time=time,
                     original_denoise_step_partial=denoise_step_partial_call,
                     execution_horizon=execution_horizon,
+                    num_flow_matching_steps=num_steps,
                 )
             else:
                 v_t = denoise_step_partial_call(x_t)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -814,6 +814,7 @@ def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
                     time=time,
                     original_denoise_step_partial=denoise_step_partial_call,
                     execution_horizon=execution_horizon,
+                    num_flow_matching_steps=num_steps,
                 )
             else:
                 v_t = denoise_step_partial_call(x_t)

src/lerobot/policies/rtc/configuration_rtc.py

@@ -47,6 +47,9 @@ class RTCConfig:
     debug: bool = False
     debug_maxlen: int = 100
 
+    use_soare_optimization: bool = True
+    variance_clipping_factor: float = 0.2
+
     def __post_init__(self):
         """Validate RTC configuration parameters."""
         if self.max_guidance_weight <= 0:

src/lerobot/policies/rtc/modeling_rtc.py

@@ -121,6 +121,7 @@ def denoise_step(
         time,
         original_denoise_step_partial,
         execution_horizon=None,
+        num_flow_matching_steps=None,
     ) -> Tensor:
         """RTC guidance wrapper around an existing denoiser.
 
@@ -163,6 +164,9 @@ def denoise_step(
         # So we need to invert the time
         tau = 1 - time
 
+        if self.config.use_soare_optimization and num_flow_matching_steps is None:
+            raise ValueError("num_flow_matching_steps must be provided when use_soare_optimization is True")
+
         if prev_chunk_left_over is None:
             # First step, no guidance - return v_t
             v_t = original_denoise_step_partial(x_t)
@@ -217,10 +221,23 @@ def denoise_step(
             grad_outputs = err.clone().detach()
             correction = torch.autograd.grad(x1_t, x_t, grad_outputs, retain_graph=False)[0]
 
+        max_guidance_weight = self.rtc_config.max_guidance_weight
+
+        # Check the following paper - https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html
+        # num of steps could be used as clipping parameter without requirements on hyperparameters tuning
+        if self.config.use_soare_optimization:
+            max_guidance_weight = num_flow_matching_steps
+
         max_guidance_weight = torch.as_tensor(self.rtc_config.max_guidance_weight)
         tau_tensor = torch.as_tensor(tau)
         squared_one_minus_tau = (1 - tau_tensor) ** 2
-        inv_r2 = (squared_one_minus_tau + tau_tensor**2) / (squared_one_minus_tau)
+        if self.config.use_soare_optimization:
+            variance_clipping_factor = torch.as_tensor(self.rtc_config.variance_clipping_factor)
+            inv_r2 = (squared_one_minus_tau + tau_tensor**2 * variance_clipping_factor) / (
+                squared_one_minus_tau * variance_clipping_factor
+            )
+        else:
+            inv_r2 = (squared_one_minus_tau + tau_tensor**2) / (squared_one_minus_tau)
         c = torch.nan_to_num((1 - tau_tensor) / tau_tensor, posinf=max_guidance_weight)
         guidance_weight = torch.nan_to_num(c * inv_r2, posinf=max_guidance_weight)
         guidance_weight = torch.minimum(guidance_weight, max_guidance_weight)

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -814,6 +814,7 @@ def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
                     time=time,
                     original_denoise_step_partial=denoise_step_partial_call,
                     execution_horizon=execution_horizon,
+                    num_flow_matching_steps=self.config.num_steps,
                 )
             else:
                 v_t = denoise_step_partial_call(x_t)