WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems
Abstract
Trajectory world models play a crucial role in robotic dynamics learning, planning, and control. While recent works have explored trajectory world models for diverse robotic systems, they struggle to scale to a large number of distinct system dynamics and overlook domain knowledge of physical structures. To address these limitations, we introduce WestWorld, a knoWledge-Encoded Scalable Trajectory World model for diverse robotic systems. To tackle the scalability challenge, we propose a novel system-aware Mixture-of-Experts (Sys-MoE) that dynamically combines and routes specialized experts for different robotic systems via a learnable system embedding. To further enhance zero-shot generalization, we incorporate domain knowledge of robot physical structures by introducing a structural embedding that aligns trajectory representations with morphological information. After pretraining on 89 complex environments spanning diverse morphologies across both simulation and real-world settings, WestWorld achieves significant improvements over competitive baselines in zero- and few-shot trajectory prediction. Additionally, it shows strong scalability across a wide range of robotic environments and significantly improves performance on downstream model-based control for different robots. Finally, we deploy our model on a real-world Unitree Go1, where it demonstrates stable locomotion performance. The code is available at https://github.com/511205787/WestWorld.
Lay Summary
Robots have very different bodies: some walk on two legs, some hop, some have four legs, and others use arms to interact with objects. Because their sensors, joints, and movements are different, a model trained for one robot often struggles to work well on another. This makes it difficult to build a general prediction model that can support planning and control across many robotic systems. To address this problem, we develop WestWorld, a model that learns how different robots move from large collections of trajectory data. Instead of treating all robots in the same way, WestWorld learns to use different internal components for different types of robot dynamics. It also uses simple knowledge about each robot’s body structure, such as how its joints and limbs are connected, so that the model can better understand the relationship between body shape and motion. Our experiments show that WestWorld predicts the future motion of unseen robots more accurately, adapts better when only limited new data are available, and improves model-based control. We further test it on a real Unitree Go1 robot, where it helps the robot walk toward a target. These results suggest that WestWorld is a step toward more general robot learning systems that can transfer across different robot bodies and tasks.