Technical Summary: UniReflex

Problem Statement

Generative imitation learning policies, such as Diffusion Policies (DP), Vision-Language-Action (VLA) models, and World Action Models (WAM), have achieved state-of-the-art performance in visuomotor trajectory planning. However, these models typically lack closed-loop force regulation capabilities. They predict kinematic motion targets and track them with fixed low-level gains, making them unreliable in contact-rich scenarios (e.g., wiping, insertion, peeling) where small geometric variations or disturbances can destabilize contact.

Existing solutions to incorporate force control generally fall into three categories: treating force as an additional observation, coupling policies with hybrid force-position control, or employing fast-slow hierarchical architectures. A critical limitation of these approaches is that they typically require end-to-end training, architectural refits, or joint optimization. In the era of foundation models, where powerful policies are pre-trained on massive, vision-only datasets, discarding these weights or subjecting them to heavy fine-tuning to accommodate force modalities is computationally inefficient and risks destroying well-generalized distributions. There is a lack of a principled solution that preserves the original pre-trained distribution of a foundation model while independently endowing the system with robust physical reflexes.

Methodology: UniReflex

UniReflex is a universal, plug-and-play framework designed to equip frozen generative policies with variable impedance control (VIC) for contact regulation without further fine-tuning of the slow backbone. The system operates on a "Fast-Slow" reflex architecture.

1. Non-Invasive Representation Interception

UniReflex intercepts deep latent representations from the action head of a frozen generative policy ($\pi_{gen}$) without modifying the network architecture or weights.

• Mechanism: Regardless of whether the policy uses diffusion (predicting noise) or flow matching (predicting velocities), the network forms a temporally structured hidden representation ($z$) prior to the final action projection. UniReflex hooks into this pre-projection state.
• Interface: This provides a backend-agnostic "slow intent" interface ($z \in \mathbb{R}^{B \times T_{slow} \times H}$) that encapsulates the policy's high-level trajectory intent, preserving task semantics.

2. Active Effort Recovery and Force-Stiffness Decoupling

The framework learns to predict normalized anisotropic stiffness directions and reference forces, decoupling active force exertion from external interaction responses.

• Supervision: Using teleoperation logs, the system derives a proxy for interaction-force intent ($F_{intent} \approx K_{virtual} \Delta x$).
• Labeling: It generates supervision signals for normalized anisotropic stiffness directions ($K_{label}$) based on the teleoperation mismatch direction. This allows the system to increase stiffness along axes with higher commanded effort, improving contact stability.
• Decoupling: The reflex network predicts three coupled outputs:
  1. $x_{reflex}$: A single-step pose correction to limit geometric drift.
  2. $F_{ref}$: The desired contact wrench for sustained normal-pressure tracking.
  3. $\Phi_K$: End-effector stiffness parameters to support decoupling of effort direction and magnitude.

3. Adaptive Gating for Hybrid Control

UniReflex integrates a lightweight gating mechanism to manage the transition between execution phases:

• Phase I (Pre-contact): The controller prioritizes tracking the frozen generative policy's pose targets, preserving the original planner's alignment quality.
• Phase II (Contact): Upon detecting contact via proprioceptive sensing, the system switches to a force-dominant regime. It regulates $F_{ref}$ and impedance parameters via the VIC interface while streaming single-step $x_{reflex}$ corrections.
• Architecture: The fast reflex network ($\phi_{reflex}$) is a causal temporal model (implemented as a GRU) that takes real-time feedback ($e_t$) and the aligned slow intent ($z'_t$) as input to generate high-frequency updates.

Key Contributions

1. Universal Plug-and-Play Framework: UniReflex equips supervised fine-tuned generative policies with closed-loop force control by non-invasively intercepting deep action-head latent representations. It requires no additional training or fine-tuning of the slow backbone.
2. Force-Stiffness Decoupling: The method jointly supervises normalized anisotropic stiffness directions and reference forces. This enables VIC-based closed-loop regulation that yields more physically consistent wrench behavior compared to passive or fixed-stiffness approaches.
3. Adaptive Gating: A lightweight gating mechanism enables seamless transitions between position-dominant planning and force-dominant execution, ensuring the robot maintains the frozen planner's alignment quality while gaining robust regulation during contact.

Experimental Results

The authors evaluated UniReflex on five contact-rich benchmarks (Wipe, Plug Charger, Peel Note, Flip Chip Box, Box Lift) using three frozen backbones: DP, $\pi_0$, and DreamZero.

• Performance Improvement: UniReflex significantly improved success rates in the contact phase (Phase II) by 20–60 percentage points across all backbones compared to vanilla policies. For example, on the "Peel Note" task, UniReflex + $\pi_0$ achieved 90% success in Phase II, compared to 30% for the vanilla $\pi_0$.
• Preservation of Non-Contact Skills: The addition of UniReflex altered Phase I (non-contact) success rates by no more than 10 percentage points, demonstrating that the original visuomotor behavior is largely preserved.
• Comparison with Baselines: UniReflex-augmented policies outperformed monolithic force-aware baselines (RDP, ForceVLA, TA-VLA) in both Phase I and Phase II. Notably, end-to-end force-aware models often degraded in position-control quality (Phase I), whereas UniReflex maintained or improved it.
• Ablation Studies:
  • Removing learned stiffness modulation or using fixed isotropic stiffness caused Phase II success to drop from 70–90% to 20–50%.
  • Removing the adaptive gating mechanism resulted in Phase I/II success rates dropping to 0–30%, highlighting the necessity of protecting the frozen planner before contact.
• Disturbance Recovery: Under dynamic disturbances (e.g., changing surface inclination), UniReflex achieved 80–100% recovery rates, whereas frozen DP failed to maintain contact in most trials.
• Efficiency: UniReflex achieved 25–66$\times$ lower per-step backward latency compared to joint training strategies. It trains only 0.188–0.375% of the parameters relative to the backbone, significantly reducing memory and compute overhead.

Significance and Claims

The paper claims that UniReflex offers a scalable alternative to monolithic end-to-end training pipelines. By treating force feedback as an independent module, it allows the robotics community to leverage powerful, pre-trained generative policies (which lack force data) without the inefficiency of retraining or the risk of catastrophic forgetting.

The authors position UniReflex as a solution to the "dilemma" of integrating specialized, sparse force observations into foundation models. They demonstrate that it is possible to add closed-loop contact regulation as a lightweight module that benefits from the high-level intent of frozen policies while independently providing robust physical reflexes. The work suggests that for contact-rich manipulation, decoupling the slow visual reasoning from high-rate contact feedback via a plug-and-play reflex is a more efficient and effective strategy than co-training perception and action for large-scale models.

Limitations noted by the authors:

• Gains are smaller for long-horizon world-action backbones (like DreamZero) due to the latency of their action chunks.
• The system is better suited for sustained contact than brief, high-impulse rigid insertions (e.g., Plug Charger).
• Responsiveness is bounded by the quality of force sensing, low-level control, and hardware bandwidth.