Lifting Ego World Models for Planning and Control

World models of embodied agents predict future observationsconditioned 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 ofeach joint. This makes the world model hard to control and expensiveto plan with as search-based methods like CEM scale poorly with actiondimensionality. To address this issue, we train a lightweight policy thatmaps high-level actions to sequences of low-level joint actions. Composingthis policy with the frozen world model produces a lifted world model thatpredicts a sequence of future observations from a single high-level action.We instantiate this framework for a human-like embodiment, defining thehigh-level action space as a small set of 2D waypoints annotated on thecurrent observation frame, each specifying a near-term goal position for aleaf joint (pelvis, head, hands). Waypoints are low-dimensional, visuallyinterpretable, and easy to specify manually or to search over. We showthat the lifted world model substantially outperforms searching directlyin low-level joint space (3.8× lower mean joint error to the goal pose),while remaining more compute-efficient and generalizing to environmentsunseen by the policy.