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 (see our demo on the website: https://westworldrobot.github.io/). The code will be available upon publication.

Figures (6)

• Figure 1: The overall architecture of our proposed WestWorld, consisting of two core components: (a) a Knowledge-Encoded Embedding Modular that injects structural embeddings as an inductive bias into trajectory representations, and (b) a System-aware MoE block that models diverse system dynamics via system-aware expert routing.
• Figure 2: Trajectory plot comparison of our method and three baselines for 100-step rollout prediction on three robots: Walker2D foot joint angle, Hopper foot angular velocity, and Franka end-effector $y$ position, given a 50-step history window as input. We can observe that our method tracks the ground-truth dynamics substantially more closely than the baselines over the 100-step horizon.
• Figure 3: Comparison between our method against the best performing SOTA by scaling the number of environments.
• Figure 4: Sys-MoE routing weights across six layers (L1--L6), each containing four experts (E1--E4), for three robotic systems. Color indicates the router weight, where brighter values correspond to higher expert activation. The router exhibits near-sparse, system-dependent expert specialization, suggesting that different systems are modeled by different combinations of experts to capture their distinct dynamics.
• Figure 5: Real-world deployment on Unitree Go1. The distilled-and-fine-tuned WestWorld serves as the dynamics predictor in MPPI and enables the robot to walk straight toward the target goal.