UMA: Unified Motion-Action Modeling for Heterogeneous Robot Learning

TL;DR

UMA (Cornell; Yunhao Cao, Andrew Owens, Kuan Fang) is a masked generative model that treats 3D object motion and robot actions as co-evolving variables under one objective, with the mask pattern selecting both the supervision regime during pretraining and the inference mode at deployment. A single pretrained model serves as a motion-conditioned visuomotor policy, a motion-based dynamics predictor for model-predictive control, and an adapter target for few-shot task tuning, depending on which tokens are masked. Pretraining mixes action-free human videos, real-robot trajectories, and simulation under the same objective, with object motion as the embodiment-agnostic common ground; task tokens are made consistent across scenes by hindsight relabeling plus a contrastive objective so they capture intent rather than scene geometry. The project page reports outperforming the strongest baseline by 20–25 percentage points on novel zero-shot manipulation tasks and beating the SOTA motion-based dynamics model on a held-out MSE benchmark.

Key claims

3D object motion (monocular keypoints in the robot’s coordinate frame) is the chosen common interface across human video, real-robot, and simulation data because it is embodiment-independent and lives in the same space as robot actions [§How does it work].
Object motion and robot actions are modeled jointly under one masked generative objective; varying the mask pattern selects supervision regime in pretraining and inference mode at deployment (motion-conditioned control, dynamics prediction, task adaptation) [§How does it work].
A learned task token is trained via hindsight relabeling plus a contrastive objective so that it captures abstract task intent while remaining invariant to keypoint sampling, timing, and scene placement, enabling task-token swap-in across scenes and across task-specification modalities at deployment [§How does it work, §Pretraining].
On novel-task motion-conditioned visuomotor control, UMA outperforms the strongest baseline by 20–25 percentage points without task-specific finetuning [§Motion-Conditioned Visuomotor Control].
Used as a dynamics predictor by conditioning on a candidate action sequence rather than generating one, UMA forecasts 3D object motion with MSE 0.042, beating PointWorld (0.054) and the no-human / no-sim ablations (0.208 / 0.044) [§Motion-Based Dynamics Modeling, Table].
Task adaptation reduces to optimizing only the low-dimensional task token (soft prompt tuning, all other parameters frozen): action-free human demos in the test scene supervise via the motion-prediction term only, while teleoperated robot demos supervise via both motion- and action-prediction terms [§Task Adaptation from Few-Shot Demonstrations].
Because the contrastive objective decouples intent from specification, the same pretrained model can be driven at deployment by a reference motion, a natural-language task description encoded into the same latent space, or a goal image whose start-to-end reference motion is recovered via correspondence — all by swapping the task token, with no finetuning [§Versatile Task Specifications].

Method

The model is a denoising transformer that operates jointly over two token streams: 3D object-motion keypoints sampled from monocular video, and robot actions. A single masked generative objective reconstructs randomly masked motion and action tokens; the random mask at training time exposes the network to every supervision regime — action-only (real-robot teleop), motion-only (action-free human video), and both (sim) — without needing per-source heads. At inference, the mask pattern flips the model between roles: mask the actions to use UMA as a visuomotor controller (given the reference motion and the current observation, predict the action chunk and re-plan from each new observation in closed loop); mask the motion to use it as a dynamics predictor (given a candidate action sequence, forecast the resulting 3D motion, scoring candidate actions against the reference for MPC).

The task token is the second non-obvious piece. A task encoder reads the initial observation of the test scene and a reference object motion and emits a task token that conditions the denoiser. Hindsight relabeling treats each segment of an observed motion as a valid task whose intent is “produce this motion,” and a contrastive objective enforces that two encodings of the same intent under different keypoint samplings, timing, and scene placements collapse to the same token while different intents repel. This is what makes the token a clean control surface — it can be replaced at deployment by a language-encoder output aligned to the same latent space, or by an inferred motion derived from goal-image correspondences, without retraining the denoiser.

Few-shot adaptation freezes the denoiser and the task encoder and optimizes only the task token on the new task’s demos. When the demos are action-free human videos, only the motion term backpropagates; when they are teleoperated robot trajectories, both motion and action terms do.

Results

The project page reports two top-line numbers. On novel-task motion-conditioned visuomotor control, UMA outperforms the strongest baseline by 20–25 percentage points without any task-specific finetuning; qualitative videos overlay the predicted end-effector trajectory in closed loop. On motion-based dynamics modeling, UMA scores MSE 0.042 vs. PointWorld 0.054 (the cited SOTA point-cloud dynamics baseline), UMA-w/o-Sim 0.208, and UMA-w/o-Human 0.044 — the Sim ablation is the dominant degradation, the Human ablation is small. The page also reports that UMA scales with both data and model size, and that residual errors trace to execution rather than perception, though the absolute numbers behind those statements are not on the public page at filing time. Specific benchmark identities (suite, splits, evaluation protocol) are not disclosed on the project page; only the relative gap to the named baselines is shown.

Why it’s interesting

UMA is the closest parallel filed to μ₀: A Scalable 3D Interaction-Trace World Model — both papers commit to 3D motion of interaction points as the embodiment-agnostic interface between heterogeneous data sources and downstream control, and both report being competitive with action-supervised VLAs despite pretraining largely action-free. The two split on which abstraction layer carries the weight: µ₀ frames trace prediction as a world-model layer that downstream embodiment-specific action experts sit on top of (frozen-WM + small action head); UMA folds both directions into the same masked generative model and toggles role via mask pattern, so the same parameters serve as policy, dynamics model, and adapter target. The action-supervised-baseline gap claims (~20–25 pp on novel tasks here; “competitive with π₀” for µ₀) make these two papers the strongest co-filed evidence that the 3D-trace interface is a real third path between pixel-rollout world models and end-to-end VLAs (VLA Models cluster).

It also slots cleanly into the MolmoMotion: Language-guided 3D motion forecasting line — MolmoMotion is the language-conditioned-3D-motion-forecasting variant on a VLM backbone, with two heads (AR and flow-matching) and a 1.16M-video auto-annotated dataset; UMA is the joint motion-action variant with mask-pattern role switching, with comparable transfer claims (76.3% pick-and-place from MolmoMotion init vs Molmo 2; 20–25pp novel-task lift for UMA). The third closely-related sibling is DynaFLIP: Rethinking Robotics Perception via Tri-Modal-Dynamics Guided Representation — DynaFLIP pushes (image-transition, language, 3D-flow) into the encoder, UMA pushes (motion, action, task) into the generative model. Together this is a coherent emerging recipe across at least four labs (UMD/Maryland for µ₀+DynaFLIP, Ai2 for MolmoMotion, Cornell for UMA) that 3D motion is the unifying interface for heterogeneous-source pretraining.

The task-token-as-soft-prompt design is also distinctive: it is the same kind of decoupling between what the task is and how it’s specified that Learning What to Say to Your VLA: Mostly Harmless VLA Steering (Language Feedback Policy) advocates at test time on frozen VLAs, but baked into the pretraining objective via the contrastive consistency constraint, rather than discovered post-hoc via interactive search.