UMI-on-Air: Embodiment-Aware Guidance for Embodiment-Agnostic Visuomotor Policies (2510.02614v1)

Abstract: We introduce UMI-on-Air, a framework for embodiment-aware deployment of embodiment-agnostic manipulation policies. Our approach leverages diverse, unconstrained human demonstrations collected with a handheld gripper (UMI) to train generalizable visuomotor policies. A central challenge in transferring these policies to constrained robotic embodiments-such as aerial manipulators-is the mismatch in control and robot dynamics, which often leads to out-of-distribution behaviors and poor execution. To address this, we propose Embodiment-Aware Diffusion Policy (EADP), which couples a high-level UMI policy with a low-level embodiment-specific controller at inference time. By integrating gradient feedback from the controller's tracking cost into the diffusion sampling process, our method steers trajectory generation towards dynamically feasible modes tailored to the deployment embodiment. This enables plug-and-play, embodiment-aware trajectory adaptation at test time. We validate our approach on multiple long-horizon and high-precision aerial manipulation tasks, showing improved success rates, efficiency, and robustness under disturbances compared to unguided diffusion baselines. Finally, we demonstrate deployment in previously unseen environments, using UMI demonstrations collected in the wild, highlighting a practical pathway for scaling generalizable manipulation skills across diverse-and even highly constrained-embodiments. All code, data, and checkpoints will be publicly released after acceptance. Result videos can be found at umi-on-air.github.io.

• The paper introduces a novel embodiment-aware diffusion policy that bridges the gap between demonstration and deployment by leveraging gradient feedback from low-level controllers.
• It utilizes a modular framework with shared action/observation spaces, enabling real-time adaptation without the need for embodiment-specific retraining.
• Experimental results demonstrate significant improvements in task success and robustness, particularly in precision aerial manipulation under dynamic conditions.

Embodiment-Aware Guidance for Cross-Embodiment Visuomotor Policy Deployment: An Analysis of UMI-on-Air

Introduction

UMI-on-Air addresses a central challenge in scalable robot learning: deploying general visuomotor policies, trained from human demonstrations, across diverse and physically constrained robotic embodiments. The framework leverages the Universal Manipulation Interface (UMI) for large-scale, in-the-wild data collection, and introduces Embodiment-Aware Diffusion Policy (EADP) to bridge the embodiment gap at deployment. EADP enables two-way communication between a high-level, embodiment-agnostic policy and a low-level, embodiment-specific controller, using gradient-based guidance to steer trajectory generation toward dynamically feasible actions. This essay provides a technical analysis of the methodology, experimental results, and implications for future research in cross-embodiment policy transfer.

Problem Formulation and Motivation

The deployment of manipulation policies across heterogeneous robots is impeded by the embodiment gap: the mismatch between the action space and dynamics of the demonstration interface (e.g., a handheld gripper) and the target robot (e.g., an aerial manipulator). While UMI enables efficient, hardware-agnostic data collection, policies trained on such data often fail when executed on robots with significant physical constraints, such as underactuated aerial manipulators subject to disturbances and actuation limits.

UMI-on-Air targets this gap by introducing a plug-and-play adaptation mechanism at inference time, obviating the need for embodiment-specific retraining or large-scale finetuning datasets. The approach is validated on challenging aerial manipulation tasks, including long-horizon and high-precision operations previously inaccessible to standard UMI-based deployments.

Figure 1: Aerial manipulation tasks enabled by UMI-on-Air, including lemon harvesting, high-precision peg insertion, and long-horizon light bulb installation.

System Architecture and Data Collection

The data collection pipeline utilizes a modified UMI setup: a lightweight, hand-held gripper with egocentric vision and accurate 6-DoF pose tracking via iPhone-based SLAM. The action and observation spaces are carefully aligned between demonstration and deployment, minimizing the embodiment gap at the interface level.

Figure 2: Data collection and deployment setup, emphasizing shared observation/action spaces and lightweight, compliant hardware.

Demonstrations consist of synchronized RGB images, EE pose trajectories, and gripper widths, forming input-output pairs for policy learning. A conditional UNet-based diffusion policy is trained to generate multimodal action sequences from these demonstrations.

Embodiment-Aware Diffusion Policy (EADP)

Controller Abstraction

At deployment, the high-level policy outputs EE-centric reference trajectories, which are tracked by a low-level controller tailored to the embodiment. Two controller classes are considered:

• Inverse Kinematics (IK) with velocity limits for fixed-base arms.
• Model Predictive Control (MPC) for aerial manipulators, incorporating full-body dynamics, actuation constraints, and disturbance rejection.

The controller exposes a differentiable tracking cost $L_{\text{track}}(a)$, quantifying the feasibility of a candidate trajectory $a$ under embodiment constraints.

Diffusion Guidance Mechanism

EADP integrates the controller's gradient feedback into the diffusion sampling process. At each denoising step $k$, the noisy trajectory $a^k$ is nudged in the direction that reduces the tracking cost:

$$\tilde{a}^k = a^k - \lambda \cdot \bar{\omega}_k \cdot \nabla_{a^k} L_{\text{track}}(a^k)$$

where $\lambda$ is a global guidance scale and $\bar{\omega}_k$ is a time-dependent scheduler. This mechanism is analogous to classifier guidance in diffusion models, but the guidance signal is the embodiment-specific tracking cost rather than a semantic class score.

Figure 3: EADP architecture, showing the integration of controller gradients into the diffusion denoising process for embodiment-aware trajectory generation.

The full procedure is summarized in a modified DDIM sampling loop, with guidance applied at each step. This enables real-time, test-time adaptation of the policy to the deployment embodiment, without retraining or access to embodiment-specific data during training.

Experimental Evaluation

Simulation Benchmarks

A suite of four manipulation tasks (Open-and-Retrieve, Peg-In-Hole, Rotate-Valve, Pick-and-Place) is evaluated across three embodiments: an oracle flying gripper, a UR10e arm, and a UAM (with and without disturbances). The embodiment gap is quantified by the performance drop from the oracle to the baseline diffusion policy (DP).

Figure 4: Visualization of EADP's adaptation across embodiments, showing how action samples are steered toward feasible regions for each robot.

Figure 5: Simulation results demonstrating that EADP consistently outperforms unguided DP, with the largest gains in less "UMI-able" embodiments.

Key findings:

• EADP reduces the embodiment gap: On UAMs, EADP recovers over 9% (no disturbance) and 20% (with disturbance) in success rates compared to DP.
• Robustness to disturbances: EADP maintains high performance under injected noise, while DP collapses.
• Task-specific improvements: In precision tasks (Peg-In-Hole), EADP enables reliable execution even when the embodiment's tracking error exceeds the task tolerance.

Real-World Deployment

UMI-on-Air is deployed on a fully actuated hexarotor with a 4-DoF manipulator, evaluated on lemon harvesting, peg-in-hole, and lightbulb insertion tasks. EADP achieves:

• 100% success in peg-in-hole and lightbulb insertion.
• 80% success in lemon harvesting, with failures attributable to perception errors rather than control infeasibility.

  Figure 6: Real-world results for DP and EADP, with colored borders indicating trial outcomes.

A cross-environment generalization test demonstrates that EADP, trained on demonstrations from varied environments, can adapt to unseen deployment settings with high reliability.

Ablation and Hyperparameter Analysis

The guidance scale $\lambda$ is critical: too little guidance fails to correct infeasible trajectories, while excessive guidance leads to conservative, out-of-distribution behaviors. The optimal regime balances task-oriented generation with embodiment feasibility.

Figure 7: Ablation of guidance scale $\lambda$ for UAM under disturbance, illustrating the trade-off between feasibility and task performance.

Theoretical and Practical Implications

UMI-on-Air demonstrates that embodiment-aware guidance at inference time is a viable alternative to large-scale, embodiment-specific finetuning. The approach is modular: any differentiable controller (analytical or learned) can provide the guidance signal, and the policy architecture remains agnostic to the deployment embodiment.

Theoretically, this decouples policy generalization from embodiment-specific adaptation, enabling scalable transfer of manipulation skills across a wide range of robots. Practically, it enables rapid deployment of new skills on novel hardware, provided a suitable controller and shared action/observation interface.

Limitations and Future Directions

• Temporal mismatch: The current system operates with a low-frequency policy (1-2 Hz) and high-frequency control (50 Hz), which may limit responsiveness in highly dynamic tasks. Streaming diffusion or continuous guidance could address this.
• Controller dependence: The quality of adaptation depends on the fidelity and differentiability of the low-level controller. Extending to learned or RL-based controllers is a promising direction.
• Perception errors: Failures in real-world deployment are often due to perception, not control. Integrating robust visual processing and uncertainty estimation could further improve reliability.

Conclusion

UMI-on-Air provides a principled framework for embodiment-aware deployment of general visuomotor policies, leveraging gradient-based guidance from embodiment-specific controllers during diffusion-based trajectory generation. The method achieves strong empirical results in both simulation and real-world aerial manipulation, particularly in embodiments with significant physical constraints. By decoupling policy learning from embodiment adaptation, UMI-on-Air advances the scalability and generality of robot learning from demonstration, with broad implications for cross-embodiment transfer and universal manipulation skill deployment.