Whale-X: Learning Scalable Embodied World Models with Enhanced Generalizability

Authors: Zhilong Zhang, Ruifeng Chen, Junyin Ye, Yihao Sun, Pengyuan Wang, Jing-Cheng Pang, Kaiyuan Li, Tian-Shuo Liu, Haoxin Lin, Yang Yu
Primary area: reinforcement learning
Keywords: world model, sequential decision-making, embodied control

Abstract:
World models play a crucial role in decision-making within embodied environments, enabling cost-free explorations that would otherwise be expensive in the real world. However, to support faithful imagination in out-of-distribution (OOD) regions, world models must possess significant generalizability, which poses substantial challenges for previous scalable approaches. This paper addresses two primary sources of the world model generalization error: the \emph{policy distribution shift} caused by the divergence between test and data-collection policies, and the \emph{compounding error} arising from long-horizon autoregressive rollout. To tackle these issues, we introduce the \emph{policy-conditioning} and the \emph{retracing-rollout} techniques, respectively. Incorporating these two techniques, we present Whale, a scalable spatial-temporal transformer-based world model with enhanced generalizability. We first demonstrate the effectiveness of the two techniques, showcasing their consistent superiority over previous baselines in both trajectory generation quality and value estimation accuracy. Furthermore, we propose Whale-X, a 414M parameter world model trained on 970K trajectories from Open X-Embodiment datasets. We show that Whale-X exhibits promising scalability and strong generalizability in real-world manipulation scenarios using minimal demonstrations.

Decision: reject

Reviews and discussion:

Metareview: