Lifting Embodied World Models for Planning and Control

Abstract: World models of embodied agents predict future observations conditioned on an action taken by the agent. For complex embodiments, action spaces are high-dimensional and difficult to specify: for example, precisely controlling a human agent requires specifying the motion of each joint. This makes the world model hard to control and expensive to plan with as search-based methods like CEM scale poorly with action dimensionality. To address this issue, we train a lightweight policy that maps high-level actions to sequences of low-level joint actions. Composing this policy with the frozen world model produces a lifted world model that predicts a sequence of future observations from a single high-level action. We instantiate this framework for a human-like embodiment, defining the high-level action space as a small set of 2D waypoints annotated on the current observation frame, each specifying a near-term goal position for a leaf joint (pelvis, head, hands). Waypoints are low-dimensional, visually interpretable, and easy to specify manually or to search over. We show that the lifted world model substantially outperforms searching directly in low-level joint space ($3.8\times$ lower mean joint error to the goal pose), while remaining more compute-efficient and generalizing to environments unseen by the policy.

• The paper introduces a modular abstraction that maps interpretable 2D waypoints to joint actions, reducing planning complexity significantly.
• The methodology leverages a conditional diffusion model and waypoint masking to achieve a 36% reduction in mean joint error on the Nymeria dataset.
• The integration of a lifted world model with a frozen low-level model enhances planning efficiency and generalizes robustly to novel environments.

Lifted Embodied World Models: Abstraction for Efficient Planning and Control

Introduction

As embodied world models scale to richer, human-like agents, specifying actions in high-dimensional joint spaces presents well-acknowledged planning and control challenges. "Lifting Embodied World Models for Planning and Control" (2604.26182) proposes a modular abstraction that reduces planning complexity by operating over high-level, interpretable actions—waypoints—while leveraging a frozen, low-level world model. This essay details the methodology, empirical results, and broader implications of the research, with a focus on the Lifted World Model (LWM) instantiated on the Nymeria dataset using PEVA as the base world model.

Problem Formulation and Motivation

World models conditioned on per-joint actions face slow planning due to the curse of dimensionality: methods such as the Cross-Entropy Method (CEM) scale poorly in both action dimensionality and sequence length. For human-like embodiments, naive sampling or optimization in joint space is computationally prohibitive and prone to unrealistic or physically implausible sequences.

The authors introduce "lifting" as a compositional abstraction: a lightweight, trainable policy maps compact, high-level waypoint actions to sequences of low-level joint actions. Composing this policy with a (frozen) low-level world model produces a lifted interface that enables efficient search, goal-conditioning, and counterfactual reasoning, all grounded in the agent’s egocentric observational frame.

Methodology

High-Level Action Space: Waypoints

High-level actions are defined as sets of 2D waypoints, each representing the image-plane projection of a 3D goal position for a leaf joint (pelvis, head, left hand, right hand). These waypoints are overlaid as annotations on the current egocentric observation (Figure 1). The process of projecting 3D joints from a goal pose onto the current image establishes a visually interpretable and low-dimensional goal specification suitable for both manual entry and gradient-free search.

Figure 1: Waypoints specify the desired locations of selected joints over the egocentric observation, providing an interpretable and compact goal signal.

Policy for Low-Level Action Generation

The policy $\pi_\theta$ maps the context—recent observations, poses, and annotated waypoints—to a sequence of low-level joint actions. This policy is implemented as a conditional diffusion model that denoises initial action sequences to generate physically plausible trajectories. Waypoint masking during training enables the policy to handle varying numbers of specified goals, supporting both full and partial action supervision.

Notably, replacing goal images $o_g$ as conditioning signals (as in prior work) with waypoints yields significantly stronger alignment between generated actions and the intended pose—empirically evidenced by a substantial drop in mean joint error (MJE) versus baselines.

Lifted World Model Composition

The LWM $f^{HL}$ is formed by chaining the policy and the frozen low-level world model: the policy outputs a sequence of joint actions, which are fed autoregressively to the world model to generate future observations. This decouples the low-level dynamics model from the problem of intention specification and search, allowing for modular upgrades or substitution of either component.

Figure 2: The lifted world model produces a future observation sequence given a high-level action, by composing a learned mapping from waypoints to joint actions with the underlying world model.

Planning via CEM in Lifted Action Space

With the abstraction in place, the CEM planner samples, evaluates, and iteratively refines waypoint sequences (instead of low-level joint actions). The cost function aggregates perceptual dissimilarity (e.g., DreamSIM) and MJE between achieved and goal poses. Empirical evidence demonstrates pronounced scaling advantages, with lower computational expense and improved convergence across horizon, sample, and environment settings.

Figure 3: CEM search operates in a low-dimensional, high-level action space for the LWM (green), contrasting with high-dimensional joint-space search for the base model (blue).

Experimental Results

Task and Dataset

Experiments are conducted on the Nymeria dataset (50 indoor/outdoor environments, first-person video, and full-body motion capture), with PEVA as the low-level world model.

Quantitative Results

Waypoint-conditioned policies outperform observation-conditioned and unconditioned baselines by large margins. Adding waypoint conditioning reduces goal-conditioned all-joint MJE from $0.323$ to $0.208$ (meters)—an $\approx$36% improvement. Crucially, learning with waypoint masking enables generalization to missing/invisible joints in the egocentric frame, with only $8.8$cm performance drop for non-visible joints versus $26.5$cm for baselines.

For search-based planning, lifted models deliver $3.8\times$ lower mean joint error (MJE) and superior compute efficiency. Direct CEM in low-level joint space achieves an all-joint MJE of $0.616$; planning in 2D waypoint space with the LWM achieves $o_g$0, a substantial reduction. The scaling behavior is strictly favorable as sampling budgets or planning horizons increase.

Figure 4: Lifted CEM in high-level action space consistently outperforms baseline CEM in joint space across CEM budgets, with 3D waypoints incurring higher search cost due to dimensionality.

LWM-based planning also generalizes to unseen environments, with only minor degradation compared to environments seen during policy training.

Qualitative Results

The policy responds sensibly to counterfactual waypoints, generating realistic and context-aware whole-body actions not present in the training data. Scene context is leveraged for waypoint interpretation, yielding differentiated behavioral outputs for identical waypoint configurations in different settings.

Figure 5: Waypoint-driven policy generates plausible, counterfactual actions such as climbing/stepping, even with sparse or novel goal specifications.

Figure 6: Identical waypoint configurations elicit distinct actions depending on scene context, such as grasping or walking, revealing contextual sensitivity.

Figure 7: CEM-derived planning solutions using LWM closely match the ground-truth actions for navigation and manipulation tasks.

Analysis and Discussion

Theoretical and Practical Implications

This work establishes that, for high-dimensional embodied agents, planners benefit immensely from explicit abstract representations of intention (waypoints), both for efficiency and for the production of physically consistent behavior. Constraining the solution space with a learned policy regularizes action selection, mitigating adversarial exploitability of perceptual similarity metrics and increasing the plausibility of generated motion.

Practically, the adopted interface—2D waypoints—is directly accessible to user input, compatible with vision-based intent inference, and efficiently searchable. The modular lifting approach enables reuse of existing/future world models and layered augmentation with richer high-level interfaces.

Comparative Claims

The authors provide empirical evidence that goal images are ineffective conditioning signals for egocentric, articulated agents, as much of the agent’s body is not visible. In contrast, sparsely specified waypoints suffice for robust control, even under severe partial observability.

Abstraction via lifting produces a dramatic reduction in planning complexity—a $o_g$1 decrease in MJE compared to low-level search, at reduced computational cost.

Limitations and Future Work

Explicit control of head orientation remains unaddressed, impacting the completeness of goal specification. Further, waypoints are limited to within-field-of-view poses; future models with improved visual consistency could leverage waypoints beyond the current image plane. Exploring alternative high-level abstractions and robustifying policies to unseen scenes are also proposed future directions.

Conclusion

"Lifting Embodied World Models for Planning and Control" demonstrates a methodologically rigorous and empirically validated approach to planning in embodied agents via modular abstraction. By training a policy to map human-interpretable waypoints to joint actions and composing it with an existing world model, the authors show that planning is not only more efficient and effective but is also more robust to partial specification and generalizes to novel environments. This lifting paradigm sets a precedent for scalable, interpretable, and controllable planning interfaces for embodied AI, suggesting clear avenues for future expansion in abstraction design, multi-agent coordination, and interaction-rich tasks.