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

Zhilong Zhang, Ruifeng Chen, Junyin Ye, Yihao Sun, Pengyuan Wang, Jing-Cheng Pang, Kaiyuan Li, Tian-Shuo Liu, Haoxin Lin, Yang Yu

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.

Primary Area: reinforcement learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 10516