Skeleton-Grounded Latent Action Model for Robot World Models

AI summary

Overview

Research area: Robotics — action-conditioned world models and latent action learning from video.

Technical level: Advanced. The paper assumes familiarity with variational autoencoders, diffusion/video-generation backbones, and the latent-action paradigm for world modeling.

Scope: SkeLAM introduces a skeleton-supervised decomposition of latent actions into a kinematic component and a residual component, then adapts robot actions into the kinematic space for controllable video-based robot world modeling on DROID.

What This Paper Is About

Robot world models need to predict how a scene changes in response to robot commands. A common shortcut is to learn "latent actions" from video by training a model to reconstruct the next frame; but because that objective only asks for accurate reconstruction, the resulting latent ends up mixing true robot motion with object reactions and unrelated scene changes. SkeLAM fixes this by splitting the latent action into a part that is explicitly tied to the robot's skeleton (pure motion) and a leftover part that carries everything else, so that real robot actions can be mapped onto the motion part alone.

Key Contributions

1. Skeleton-grounded latent decomposition. SkeLAM encodes each RGB transition into a kinematic latent and a residual latent, and grounds the kinematic latent in robot motion using paired skeleton images and a frozen skeleton encoder-decoder. This gives a dedicated kinematic target for adapting robot actions, while the residual retains complementary visual dynamics for reconstruction.

2. Cross-modal swap reconstruction. During training, the RGB kinematic latent is replaced with its paired skeleton latent (via posterior mean) while keeping the same RGB residual sample, and the decoder must still reconstruct the target frame. This pushes motion conditioning from the skeleton branch into the RGB decoder rather than only aligning the two in latent space.

3. Structured world-model conditioning with branch dropout. The world model (Cosmos-Predict2.5-2B backbone) receives the two components through separate MLP projections. Branch dropout uses four whole-branch mask modes — full/keep, partial (drop kinematic or residual), and null — so the model is trained for the zero-residual condition actually used at inference.

4. Action adapter plus fine-tuning for deployment. A lightweight MLP adapter maps robot actions into the kinematic latent space (regressing the frozen encoder's kinematic posterior mean), followed by full world-model fine-tuning on these predicted conditions. At inference, generation uses only the current observation and robot actions, with the residual branch zeroed.

Main Findings

• Best overall on DROID among action-conditioned models. SkeLAM ranks first on seven of eight reported metrics, including all four visual fidelity metrics: PSNR 19.35, SSIM 0.8340, LPIPS 0.1310, MSE 0.0151, with ADE 41.15 px, F1_dir 0.5465, L1_mag 0.0510, and r_mag 0.5277.

• Gains over latent-action baselines. Against DreamDojo, ADE drops 19.5% (51.11 → 41.15 px) and LPIPS drops 11.1% (0.1473 → 0.1310). Against CD-LAM, ADE drops 23.7% (53.90 → 41.15) and LPIPS 17.4% (0.1586 → 0.1310).

• Trade-off with Ctrl-World and OSCAR. Those models achieve higher direction F1, but SkeLAM reduces ADE from 42.96 to 41.15 px versus Ctrl-World and LPIPS from 0.4029 to 0.1310.

• Skeleton supervision matters most for motion. Removing it raises ADE from 41.15 to 64.59 and drops magnitude correlation from 0.5277 to 0.2009.

• Swap reconstruction matters for both motion and pixels. Removing it raises ADE to 66.74, drops r_mag to 0.1897, and worsens LPIPS from 0.1310 to 0.1436.

• Residual is cheap but useful. Zeroing the residual latent during world-model pretraining keeps visual fidelity nearly unchanged (LPIPS 0.1305 vs. 0.1310) but raises ADE to 43.37 and lowers r_mag to 0.4737.

• Adapter regression is essential. Omitting action adaptation raises ADE to 65.84 and lowers r_mag to 0.1603.

• Kinematic latents are highly decodable. Linear probes on frozen latents give the kinematic component an EEF MAE of 5.99×10⁻³, R²_EEF 0.3200, and magnitude macro-F1 0.7467, versus 0.01620 / 0.0078 / 0.7183 for DreamDojo and 0.01615 / 0.0079 / 0.7006 for CD-LAM. Concatenating the residual only nudges R² to 0.3277 and F1_mag to 0.7482 — so the 32D kinematic latent nearly matches the 64D concatenation.

• The RGB decoder is sensitive to motion correspondence. Substituting paired skeleton kinematic means raises reconstruction MAE by only 0.53%, while mismatched substitution raises it by 11.67%. Zeroing the kinematic component raises MAE by 5.78% versus 0.61% for the residual.

• Qualitative zero-shot human-to-robot transfer. Using UniSkill human demonstration videos as motion prompts with all parameters frozen, SkeLAM produces a downward reach for a towel-pushing prompt and an upward movement for a tissue-pulling prompt, while DreamDojo shows limited end-effector displacement.

• Open-loop policy evaluation agrees with real robots. Across 7 policies, 15 RoboArena sessions, and 105 episodes, SkeLAM reaches 63.02% balanced accuracy and 33.33% success recall (versus 48.44% and 0.00% for Cosmos-Predict2.5), with success-rate MAE 4.76 vs. 7.62 percentage points, Spearman ρ = 0.5347 and Pearson r = 0.6091. The authors note the modest recall means reproducing successful outcomes remains challenging.

Methodology in Plain English

The pipeline has four stages.

Stage 1 — learn structured latent actions. Two Transformer VAEs are trained on paired RGB frames and pre-rendered skeleton images. The RGB encoder outputs two 32-dimensional latents per transition (kinematic and residual); the skeleton encoder outputs a 32-dimensional kinematic latent. Initialization is conditional reconstruction of the next RGB frame and the next skeleton frame. Then the skeleton branch is frozen and the RGB kinematic latent is forced to be usable by the frozen skeleton decoder (grounding loss) and to align with the skeleton's motion latent (alignment loss). Cross-modal swap reconstruction then swaps in the skeleton mean while keeping the RGB residual and requires the RGB decoder to reconstruct the same target frame. A variance/covariance regularizer keeps the latent dimensions from collapsing.

Stage 2 — train the world model on video latents. A frozen RGB encoder extracts the two latents for each transition in a future clip. Each is projected by its own MLP and summed into the conditioning signal for the Cosmos-Predict2.5-2B backbone. Branch dropout randomly drops whole branches per clip (keep, partial, null), so the model learns to predict well even when the residual is missing.

Stage 3 — adapt robot actions. A small MLP adapter regresses the frozen encoder's kinematic posterior mean from the 7-dimensional action proxy (translation delta in the base frame, SO(3) rotation vector, destination gripper openness). Only the adapter is trained; the encoder and world model are frozen.

Stage 4 — fine-tune and deploy. The full world-model backbone is fine-tuned with adapter-predicted kinematic latents and a curriculum over residual/mask pairs. At inference the residual is set to zero and the mask to zero, and rollouts are generated from the current frame plus actions alone.

Data: DROID plus AgiBot World, with DROID supplying 60,000 Stage-1 clips and 45,358 clips to each of Stages 2A, 2B, 3A, and 3B. DROID windows sample 13 frames at indices f₀+4k from 60 FPS video (15 FPS, 0.8 s span); Stage 1 uses 240×320 inputs and Stage 2/3 uses 480×640 crops. Evaluation uses 400 DROID clips with one conditioning frame and 12 future frames. Training runs on NVIDIA A100 GPUs and takes roughly 2 days for one complete run.

Why This Matters

Research impact. The paper targets a widely acknowledged weakness of reconstruction-driven latent actions: the learned proxy action entangles robot motion with object responses, which blurs the mapping to real controls. By making the motion part a supervised, separately projected quantity, SkeLAM offers a template for structuring latent action spaces rather than just scaling them — and its probing analysis suggests the resulting representation is more linearly decodable for end-effector displacement.

Real-world applications.

• Policy evaluation before hardware time: open-loop rollouts producing success rates that track real executions (ΔSR 4.76 points, ρ = 0.5347) could reduce the cost of ranking candidate policies on physical robots.
• Data-scarce manipulation training: a world model that generalizes across AgiBot and DROID embodiments can synthesize experience where teleoperation data is unavailable.
• Cross-embodiment skill transfer: zero-shot human-to-robot motion prompting could let human demonstration videos serve as reusable motion specifications without paired robot data.
• Model-based planning and simulation: more faithful simulation of how end-effector motion affects objects supports foresight-guided control.

Industry relevance. Robot learning pipelines are bottlenecked by the cost of real-robot rollouts and by the mismatch between video-pretrained dynamics and actual controllers. A lightweight adapter that maps a 7-D action proxy into a pretrained latent space, trained while the backbone stays frozen, is an inexpensive way for industrial teams to specialize a general video world model to their own hardware.

Future Directions

1. Closed-loop evaluation. The authors explicitly plan to extend the current recorded-action replay protocol to closed-loop policy evaluation, asking how prediction errors propagate into subsequent policy decisions and task outcomes over longer interaction horizons.

2. Improving success recall. SkeLAM's 33.33% success recall shows that reproducing actual task successes is still hard; better modeling of contact and manipulation outcomes is an open problem.

3. Local gripper fidelity. The qualitative comparison notes residual local gripper errors even as SkeLAM tracks the reference trajectory, pointing to a finer-grained motion representation need.

4. Better residual handling at inference. The residual branch is informative for motion (ADE rises to 43.37 without it) yet must be zeroed at deployment; closing this train-test gap without leaking future information is a natural extension.

Target Audience

Researchers and engineers working on robot world models, video-based dynamics learning, and latent action representations — particularly those building action-conditioned video generators or evaluating policies in simulation. It is also relevant to practitioners doing cross-embodiment transfer and imitation learning from human video, and to readers interested in structured or disentangled latent representations for control. Readers without background in VAEs and video diffusion backbones will find the methodology sections dense; the ablation and probing tables are the most accessible entry points.

Authors’ abstract

Robot world models simulate how environments evolve in response to robot actions. Latent action models enable learning dynamics from videos by encoding visual transitions as proxy actions. However, reconstruction-driven latent actions mix robot motion with object responses and other scene changes, complicating their correspondence with robot controls. We propose SkeLAM, a skeleton-guided latent action model for controllable robot world modeling. SkeLAM decomposes latent actions into a kinematic component grounded by robot skeleton supervision and a residual component retaining complementary visual dynamics. Cross-modal swap reconstruction encourages RGB-derived and skeleton-derived kinematic representations to be interchangeable during video reconstruction. The world model incorporates both components through separate conditioning pathways, while branch dropout reduces its dependence on future-derived residual information. A lightweight adapter maps robot actions to the kinematic space, and subsequent world-model fine-tuning adapts generation to these predicted conditions. At inference, the model generates future observations from the current observation and robot actions. On DROID, SkeLAM reduces end-effector ADE by 19.5% and LPIPS by 11.1% compared with DreamDojo. Motion probing supports kinematic specialization, while qualitative results demonstrate zero-shot human-to-robot motion transfer in generated rollouts. Policy evaluation shows agreement between predicted success rates and those measured in real robot executions.