Lifting Embodied World Models for Planning and Control

ListenWaypoints Over Joints: Making World Models Planable10:19

Show notes

A look at *Lifting Embodied World Models for Planning and Control* by Wang, Darrell, Izmailov, Bai, and Bar. The authors leave a high-dimensional egocentric world model (PEVA) frozen and wrap it in a lightweight diffusion policy that turns a handful of 2D image-space waypoints into low-level joint actions — making search-based planning tractable and cutting mean joint error to the goal by 3.8× versus searching in joint space directly. We dig into why goal images fail for egocentric agents, why waypoints are the load-bearing trick, and what this deployment pattern means for controlling big generative world models. **Sources** - Lifting Embodied World Models for Planning and Control

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

CEM planning through a high-dimensional embodied world model (like PEVA, which predicts egocentric video from per-joint XSens actions with 15 upper-body joints × 3 Euler angles ≈ 48-dim actions per step) scales poorly because search cost grows exponentially with action dimensionality. The authors train a lightweight diffusion policy that maps a small, low-dim, image-space “waypoint” action (2D projected goal positions for the pelvis, head, and hands — 8 dims total) into a sequence of low-level joint actions, then compose that policy with the frozen PEVA world model to obtain a Lifted World Model (LWM) that consumes waypoints and produces future observations. Running CEM in waypoint space gives 3.8× lower mean joint error to the goal than searching directly in joint-action space on Nymeria hybrid navigation-plus-interaction tasks, at lower compute, with generalization to environments unseen by the policy. The recipe leaves the world model untouched — lifting is entirely a change of action interface.

Key claims

Cross-Entropy Method (CEM) planning scales poorly with action dimensionality, and reaching-out with a human-like embodiment requires coordinated per-joint control across shoulder / elbow / wrist — motivating an abstraction layer between high-level goals and the world model’s low-level action interface [§1].
Waypoints — 2D image-space goal positions for the four leaf joints (pelvis, head, left hand, right hand) — are a low-dimensional (8-dim total), visually interpretable, search-friendly high-level action space; ground-truth waypoints at training time are obtained by projecting future 3D leaf-joint positions into the current image via the camera matrix [§3.1, Eq. 7–8, Fig. 3].
Goal observations (a target egocentric frame) are a poor conditioning signal for egocentric policies: the agent’s own body is rarely in frame, so a goal image carries little info about the target pose — a NoMaD-style diffusion policy conditioned on goal-image barely beats the unconditioned baseline (0.392 vs 0.405 all-MJE), and adding architecture and pose-context tweaks doesn’t change this pattern [§4.1.1, Table 1].
Waypoint conditioning is the load-bearing addition: it drops goal-conditioned all-MJE from 0.323 (pose+image, unconditioned) to 0.243, a 25% relative reduction [§4.1.1, Table 1].
Waypoint masking during training (half the batches: no masking; half: each waypoint independently masked with p=0.5) further improves goal-conditioned MJE to 0.226 all-MJE and makes the policy robust to sparse / out-of-frame goals, at the cost of unconditioned MJE — an acceptable tradeoff for the downstream planner [§3.2.1, §4.1.1, Table 1].
The waypoint-conditioned policy is robust to counterfactual and sparse waypoint sets (a single head waypoint suffices to steer the agent) and uses image context to interpret identical waypoints differently depending on the scene (stove context → grasp pot; open room → walk forward) [§4.1.2, Fig. 6, Fig. 7].
On planning: Lifted CEM (search in 8-dim waypoint space per step, planning horizon T=3) drops all-MJE from 0.704 (no action) → 0.374, while PEVA CEM (search in 48-dim joint space directly) only reaches 0.616 — a 3.8× larger reduction, with the LWM also outperforming the unconditioned policy and the observation-conditioned policy [§4.2.2, Table 4].
Lifted CEM strictly dominates PEVA CEM on the compute × MJE Pareto frontier (all iteration and sample-count settings); PEVA CEM actually gets worse after one CEM iteration, evidencing the naive-search failure mode [§4.2.4, Fig. 9].
Waypoint 3D-augmentation (adding depth to specify an exact 3D goal, dims 8→12) improves policy MJE only marginally (~0.5 cm), and hurts planning after the first CEM step because the extra depth dimensionality is hard to search over — 2D is the sweet spot [§4.1.4, §4.2.5, Table 1, Fig. 9].
The lifted policy trained with three Nymeria environments (locations 6, 19, 34) held out still outperforms PEVA CEM on those held-out environments (0.362 vs 0.553 all-MJE), only slightly worse than the same-environment policy (0.333) — the lifting mechanism generalizes to unseen scenes [§4.2.7, Table 5].
Decomposition experiment: predicting an action sequence given the goal pose is much easier (0.105 all-MJE) than predicting the goal pose from a goal image (0.279) — egocentric observations are insufficient to identify the goal pose because the embodiment is seldom self-visible, formally justifying why goal-image conditioning fails and waypoints work [§4.1.5, Table 3].

Method

The pipeline has three parts. (1) High-level action space: for a human-like XSens embodiment, a high-level action is a set of four 2D image-space waypoints — pelvis (green), head (yellow), left hand (blue), right hand (red) — annotated as colored dots on the current observation frame o_t. Ground-truth training waypoints are obtained by running forward kinematics on the current pose representation p_t (pelvis 3D position + per-joint Euler angles) to get 3D leaf-joint positions at the goal time, then projecting them into o_t via the known camera matrix π(·). (2) Waypoint-conditioned diffusion policy: the current observation and the waypoint-annotated goal image g_t are both encoded by a frozen DINOv3-S backbone with tokens preserved (no pooling); context pose vectors are linearly projected and added to the corresponding-timestep image tokens; 3D positional embeddings are added; a vision transformer processes the concatenation; the pooled context vector conditions a NoMaD-style ConditionalUnet1D (dims raised from 256 to 512) that denoises a sequence of H=3 low-level joint actions. Half the training batches mask no waypoints; half independently mask each waypoint with p=0.5. (3) Lifted World Model: at inference the policy predicts a low-level action sequence from one waypoint set; the frozen PEVA world model autoregressively rolls those low-level actions forward H steps to produce future observations. CEM planning samples waypoint sequences of length T=3 from an initial Gaussian prior (μ=0.5, σ=0.25 for pelvis; σ=0.15 for others), rolls each through the LWM, scores rollouts by a DreamSIM perceptual-distance cost against the goal observation on the final predicted frame, and updates the prior on the top-k lowest-cost samples over up to 5 iterations × 128 samples.

Results

Policy MJE (Table 1): Full recipe (arch + pose context + waypoint conditioning + waypoint masking) reaches goal-conditioned all-MJE 0.226, vs 0.392 for a NoMaD baseline with goal-image conditioning — a 42% relative reduction. 3D-conditioning adds negligible benefit (0.226 → 0.208 with additional masking).
Waypoint visibility (Table 2): Waypoint-conditioned + masked policy: all-MJE 0.213 (visible) vs 0.301 (not visible) — only 8.8 cm penalty when the target joint is out of the current frame. NoMaD baseline: 0.314 vs 0.678 (36.4 cm penalty).
Planning MJE (Table 4): Initial distance 0.704 → PEVA CEM 0.616 → Lifted CEM (2D, ours) 0.374. Also beats the unconditioned policy (0.650) and the image-conditioned policy (0.585). Lifted CEM (3D) sits between at 0.420.
Compute-efficiency curves (Fig. 9): Lifted CEM outperforms PEVA CEM across all (samples, iterations) settings tested; PEVA CEM regresses after one iteration.
Held-out environments (Table 5): Lifted CEM (held-out policy) all-MJE 0.362 vs PEVA CEM 0.553 on the same tasks — held-out lift is preserved.
Decomposition (Table 3): Motion generation (goal pose → actions) 0.105 all-MJE vs goal-pose prediction (goal image → pose) 0.279 — a 2.7× gap that quantifies the fundamental limitation of goal-image conditioning for egocentric agents.

Why it’s interesting

This paper adds a distinct deployment pattern for world foundation models to the ledger — call it WFM-as-frozen-planner-backend-with-a-lifted-action-interface — that’s structurally different from the seven patterns already catalogued on World Foundation Models (backbone / predictive-surprise reward / interruptible tool / policy-internal substrate / edit-prior via KV cache / frozen trace-WM + small action expert / value-model / inference-time-guidance signal). Here the WFM (PEVA) is left completely frozen and untouched; the lever is a small learned wrapper that reshapes the input action space the WFM sees, so search-based planning becomes tractable. That’s a sharp counterpoint to the µ₀ / μ₀: A Scalable 3D Interaction-Trace World Model recipe which changes the world model’s interface layer to 3D interaction traces, and to World Pilot: Steering Vision-Language-Action Models with World-Action Priors which steers a VLA through the WFM’s latents — both change the WFM interface; this changes the planner’s interface to an unchanged WFM. The image-space waypoint action space is also a concrete addition to the viewpoint-as-action candidate cluster (LIME: Learning Intent-aware Camera Motion from Egocentric Video, ActiveMimic: Egocentric Video Pretraining with Active Perception): those papers make viewpoint a first-class action learned from egocentric human video; this paper makes 2D image-space goal points for leaf joints a first-class action, with a similar bet that egocentric visual grounding beats absolute-coordinate action specification.