WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems

Yuchen Wang; Jiangtao Kong; Sizhe Wei; Xiaochang Li; Haohong Lin; Hongjue Zhao; Tianyi Zhou; Lu Gan

cs.LGMar 15, 2026

WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems

Authors: Yuchen Wang, Jiangtao Kong, Sizhe Wei, Xiaochang Li, Haohong Lin, Hongjue Zhao, Tianyi Zhou, Lu Gan,

Organizations: William & Mary · Georgia Institute of Technology · Carnegie Mellon University · University of Illinois at Urbana-Champaign ·

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.

Figures & tables

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.

Method	Walker2d ( $\times 10^{-2}$ )		Hopper ( $\times 10^{-2}$ )		Franka ( $\times 10^{-2}$ )
Method	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$
MLPEnsemble	26.006	12.028	19.987	7.216	12.164	4.271
TDM	20.122	6.428	17.634	5.076	23.686	8.435
Trajworld	22.261	8.623	17.388	5.441	13.102	5.127
Ours	16.350	5.064	13.731	3.368	7.737	2.539

Table 1: Zero-shot generalization performance of different models on three dynamical systems. Errors are computed in the normalized space and reported as MAE and MSE (

\times 10^{-2}

); lower is better.

Method	Cassie ( $\times 10^{-2}$ )		A1 ( $\times 10^{-2}$ )		UR5 ( $\times 10^{-2}$ )
Method	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$
MLPEnsemble	14.369 $\pm$ 0.523	4.416 $\pm$ 0.368	14.357 $\pm$ 1.009	3.707 $\pm$ 0.523	15.181 $\pm$ 0.597	4.936 $\pm$ 0.322
TDM	14.510 $\pm$ 0.642	3.404 $\pm$ 0.300	10.624 $\pm$ 0.312	2.151 $\pm$ 0.109	18.578 $\pm$ 0.324	5.510 $\pm$ 0.175
Trajworld	7.834 $\pm$ 0.167	1.697 $\pm$ 0.109	5.138 $\pm$ 0.200	0.900 $\pm$ 0.050	8.066 $\pm$ 0.799	2.117 $\pm$ 0.433
Ours	5.316 $\pm$ 0.108	0.808 $\pm$ 0.025	4.227 $\pm$ 0.120	0.628 $\pm$ 0.040	4.925 $\pm$ 0.317	0.831 $\pm$ 0.150

Table 2: Few-shot generalization performance on three robotic systems. Results are computed in the normalized space and reported as MAE and MSE (

\times 10^{-2}

) averaged over three random seeds (mean

\pm

standard deviation); lower is better.

Figure 3: 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 4: Comparison between our method against the best performing SOTA by scaling the number of environments.

Method	Pretrain	Walker2d	Hopper	Go1
Method	Pretrain	Accumulate Reward $\uparrow$	Accumulate Reward $\uparrow$	Accumulate Reward $\uparrow$
MLPEnsemble	✗	119.18	147.37	-1.07
MLPEnsemble	✓	190.51	200.56	-0.31
TDM	✗	122.65	242.34	-0.72
TDM	✓	207.61	154.64	0.03
Trajworld	✗	395.23	366.26	0.05

Table 3: Downstream model-based control performance using MPPI on Walker2D, Hopper, and Unitree Go1. We report accumulated episode reward (higher is better) averaged over the evaluation episodes. All methods are evaluated with fixed random seeds and identical MPPI hyperparameters for fair comparison.

Method	Sys-MoE	Structural embedding	Walker2D ( $\times 10^{-2}$ )		Hopper ( $\times 10^{-2}$ )		Franka ( $\times 10^{-2}$ )
Method	Sys-MoE	Structural embedding	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	MSE $\downarrow$
Ours w/o Sys-MoE	✗	✓	18.707	6.797	15.978	4.630	9.392	2.873
Ours w/o Structural embedding	✓	✗	21.156	7.872	16.227	4.990	7.897	2.707
Ours	✓	✓	16.350	5.064	13.731	3.368	7.737	2.539

Table 4: Ablation results under the zero-shot setting. We ablate the Sys-MoE layer by replacing it with a dense SSM, and ablate the structural encoding by removing it during pretraining. Results are computed in the normalized space and reported as MAE and MSE (

\times 10^{-2}

); lower is better.

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. A side-by-side comparison with TrajWorld is provided in the project website at https://westworldrobot.github.io/ .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Definition
$\mathcal{S}$	state space
$\mathcal{A}$	action space
${\bm{s}}$	state
${\bm{a}}$	action
$n\in\mathbb{N}^{+}$	number of robotics systems
$N$	number of task-level environments in scalability evaluation

Appendix

Table 5: Summary of notations

Body	Global idx	$\pi_{\mathrm{pre}}$	$\pi_{\mathrm{in}}$	$\pi_{\mathrm{post}}$
torso	0	0	6	6
right_thigh	1	1	3	5
right_leg	2	2	5	4
right_foot	3	3	4	3
left_thigh	4	4	2	2
left_leg	5	5	1	1

Appendix

Table 6: Walker example of LCRS traversal ranks used to construct structural embeddings. Here

\pi_{\mathrm{obj}}=0

since the scene contains only the robot.

Dataset	Robot morphology categories
ExORL ( Yarats et al., 2022 )	Cartpole, Jaco, Quadruped, Walker †
RL Unplugged ( Gulcehre et al., 2020 )	Cartpole, Fish, Humanoid, Manipulator, Walker †
JAT ( Gallouédec et al., 2024 )	Double Pendulum, Pendulum, Pusher, Reacher, Swimmer*
DB-1 ( Wen et al., 2022 )	Acrobat, Ball In Cup, Cartpole, Cheetah-2-back, Cheetah-2-front, Cheetah-3-back, Cheetah-3-balanced, Cheetah-3-front, Cheetah-4-allback, Cheetah-4-allfront, Cheetah-4-back, Cheetah-4-front, Cheetah-5-back, Cheetah-5-balanced, Cheetah-5-front, Cheetah-6-back, Cheetah-6-front, Finger, Fish, Hopper † , Hopper-3, Hopper-5, Humanoid, Humanoid-2d-7-left-arm, Humanoid-2d-7-left-leg, Humanoid-2d-7-lower-arms, Humanoid-2d-7-right-arm, Humanoid-2d-7-right-leg, Humanoid-2d-8-left-knee, Humanoid-2d-8-right-knee, Humanoid-2d-9-full, Manipulator, Reacher, Swimmer6, Swimmer15, Walker † , Walker-2-flipped, Walker-2-main, Walker-3-flipped, Walker-3-main, Walker-4-flipped, Walker-4-main, Walker-5-flipped, Walker-5-main, Walker-6-flipped, Walker-6-main
TD-MPC2 ( Hansen et al., 2024 )	Acrobot, Ball In Cup, Cartpole, Cheetah † , Finger, Fish, Hopper † , Pendulum † , Reacher † , Walker †
Modular RL ( Huang et al., 2020 )	Cheetah-2-back, Cheetah-2-front, Cheetah-3-back, Cheetah-3-balanced, Cheetah-4-allback, Cheetah-4-back, Cheetah-4-front, Cheetah-5-back, Cheetah-5-balanced, Cheetah-5-front, Cheetah-6-back, Cheetah-6-front, Hopper-3, Hopper-5, Walker-2-flipped, Walker-3-flipped, Walker-4-flipped, Walker-5-flipped, Walker-6-flipped, Walker-7-flipped

Appendix

Table 7: A detailed list of diverse robotics used in the pretraining dataset. For robotics sharing the same name, we mark those from OpenAI Gym with an asterisk (*) and those from DeepMind Control Suite with a dagger (†).

Component	Task-level environments	#
TD-MPC2	walker-stand, walker-walk, walker-run, cheetah-run, reacher-easy, reacher-hard, acrobot-swingup, pendulum-swingup, cartpole-balance, cartpole-balance-sparse, cartpole-swingup, cartpole-swingup-sparse, cup-catch, finger-spin, finger-turn-easy, finger-turn-hard, fish-swim, hopper-stand, hopper-hop, walker-walk-backwards, walker-run-backwards, cheetah-run-backwards, cheetah-run-front, cheetah-run-back, cheetah-jump, hopper-hop-backwards, reacher-three-easy, reacher-three-hard, cup-spin, pendulum-spin	30
ExORL	jaco-reach, quadruped-locomotion	2
RL-Unplugged	humanoid_run, manipulator-insert-ball, manipulator-insert-peg	3
JAT	inverted_double_pendulum, inverted_pendulum, pusher, reacher, swimmer	5
DB-1	acrobot, cheetah_2_back, cheetah_2_front, cheetah_3_back, cheetah_3_balanced, cheetah_3_front, cheetah_4_allback, cheetah_4_allfront, cheetah_4_back, cheetah_4_front, cheetah_5_back, cheetah_5_balanced, cheetah_5_front, cheetah_6_back, cheetah_6_front, hopper_3, hopper_5, humanoid, humanoid_2d_7_left_arm, humanoid_2d_7_left_leg, humanoid_2d_7_lower_arms, humanoid_2d_7_right_arm, humanoid_2d_7_right_leg, humanoid_2d_8_left_knee, humanoid_2d_8_right_knee, humanoid_2d_9_full, manipulator-bring_ball, swimmer6, swimmer15, walker_2_flipped, walker_2_main, walker_3_flipped, walker_3_main, walker_4_flipped, walker_4_main, walker_5_flipped, walker_5_main, walker_6_flipped, walker_6_main	39
Modular-RL	walker_7_flipped	1

Appendix

Table 8: Task-level environments used in the pretraining dataset.

$N$	Selected task-level environments
1	Cartpole Swingup
2	Cartpole Swingup; Acrobot Swingup
5	Walker Walk Backwards; Finger Turn Easy; Cartpole Swingup; Reacher Three Easy; Cheetah Run Front
10	Walker Walk Backwards; Finger Turn Easy; Cartpole Swingup; Reacher Three Easy; Cheetah Run Front; Acrobot Swingup; Reacher Hard; Cup Catch; Cartpole Swingup Sparse; Finger Turn Hard
20	Walker Walk Backwards; Finger Turn Easy; Cartpole Swingup; Reacher Three Easy; Cheetah Run Front; Acrobot Swingup; Reacher Hard; Cup Catch; Cartpole Swingup Sparse; Finger Turn Hard; Cartpole Balance Sparse; Hopper Hop Backwards; Pendulum Spin; Cheetah Run Backwards; Fish Swim; Reacher Three Hard; Walker Walk; Finger Spin; Hopper Hop; Walker Run
30	Walker Walk Backwards; Finger Turn Easy; Cartpole Swingup; Reacher Three Easy; Cheetah Run Front; Acrobot Swingup; Reacher Hard; Cup Catch; Cartpole Swingup Sparse; Finger Turn Hard; Cartpole Balance Sparse; Hopper Hop Backwards; Pendulum Spin; Cheetah Run Backwards; Fish Swim; Reacher Three Hard; Walker Walk; Finger Spin; Hopper Hop; Walker Run; Hopper Stand; Cup Spin; Reacher Easy; Cheetah Jump; Pendulum Swingup; Cartpole Balance; Cheetah Run Back; Walker Stand; Cheetah Run; Walker Run Backwards

Appendix

Table 9: Detailed task-level environment subsets and data split protocols for the environment-scaling study. For

N\leq 50

, each selected environment uses 1000/500/500 episodes for train/validation/test. For

N=60

, each selected environment uses 500/250/250 episodes due to data availability. All compared methods use the same split for each

N

#Experts $P$	Val. MAE $\downarrow$	Best step $s^{\star}$ (k updates)
2	0.0409	108.6
4 †	0.0379	137.2
6	0.0390	176.7
8	0.0340	214.0

Appendix

Table 10: Best validation MAE and the optimizer update step

s^{\star}

at which the best MAE is achieved during pretraining. We report

s^{\star}

in thousands of updates (k). MAE is computed in the normalized space. Lower values indicate better performance.

\dagger

denotes the setting used in all experiments.

Number of bins	MAE $\downarrow$	MSE $\downarrow$
128	1.455	0.098
256	1.415	0.095
512	1.404	0.094

Appendix

Table 11: Effect of the number of discretization bins on the Walker system. Errors are reported as MAE and MSE (

\times 10^{-2}

). Lower values are better.

Method	Walker root velocity $(\mathrm{m/s})$	Walker foot angle $(\mathrm{rad})$	Hopper leg angle $(\mathrm{rad})$	Hopper foot angle $(\mathrm{rad})$	Franka $y$ position $(\mathrm{m})$	Franka pitch $(\mathrm{rad})$
	MSE $\downarrow$	MSE $\downarrow$	MSE $\downarrow$	MSE $\downarrow$	MSE $\downarrow$	MSE $\downarrow$
MLPEnsemble	4.623	1.339	0.089	0.685	0.068	0.547
TDM	4.654	0.615	0.134	0.372	0.653	0.390
TrajWorld	3.202	0.779	0.093	0.468	0.538	0.162
Ours	2.544	0.183	0.045	0.244	0.039	0.130

Appendix

Table 12: Zero-shot prediction errors in the original physical space. We report MSE on representative state dimensions from three robotic systems. Lower is better.

Figure 6: The effect of pre-training on few-shot learning for three different robotic systems: (a) Cassie, (b) A1 and (c) UR5 Robots.

Method	Tuned Layers	Tuned Params	MSE ( $\times 10^{-2}$ ) $\downarrow$
Method	Tuned Layers	Tuned Params	Cassie	A1	UR5
Trajworld	Full	10.0 M	1.674	0.931	2.616
Ours	Last 2 Sys-MoE	3.2 M (21.91%)	1.249	0.775	1.052
	Last 4 Sys-MoE	6.3 M (43.37%)	1.074	0.762	1.031
	Last 6 Sys-MoE	9.5 M (64.83%)	1.055	0.744	0.895
	Full	14.6 M (100%)	0.785	0.636	0.761

Appendix

Table 13: We compare fine-tuning the last 2/4/6 Sys-MoE layers against full fine-tuning and the best-performing SOTA baseline (TrajWorld). All experiments use a fixed random seed. Lower is better.

Method	Latency (ms) $\downarrow$
Method	$H=10$	$H=30$	$H=50$	$H=70$	$H=100$
TDM	388.85 $\pm$ 0.41	697.89 $\pm$ 0.64	1099.63 $\pm$ 1.41	1602.34 $\pm$ 1.33	2558.94 $\pm$ 1.64
TrajWorld	94.62 $\pm$ 0.30	236.63 $\pm$ 0.33	379.43 $\pm$ 0.59	522.37 $\pm$ 0.43	736.67 $\pm$ 0.91
Ours	113.61 $\pm$ 0.71	138.03 $\pm$ 1.45	159.90 $\pm$ 0.97	183.96 $\pm$ 1.01	238.52 $\pm$ 0.47

Appendix

Table 14: Wall-clock latency (ms, mean

\pm

std) for multi-step prediction on a single NVIDIA RTX A6000. All methods take

T_{\mathrm{hist}}=50

history steps as input (

B=4

D_{o}=78

D_{a}=21

). Results are averaged over 10 repeated runs after warm-up. Lower is better.

Robot Motion Prediction

See all 9 papers

Oct 7, 2026

RoboJEPA: Scaling Robotic Latent World Models

Latent world models have shown a remarkable ability to predict future states and to plan in the real world. In practice, however, we lack a principled way to estimate how their capabilities scale with model size, data, and compute, an open problem that slows progress in the field. In this work we present RoboJEPA, a world model based on the Joint Embedding Predictive Architecture (JEPA) and trained on a large-scale dataset spanning 12 robotic embodiments. We show that RoboJEPA's imagination error, the error of its latent rollouts, follows a second-order power law in compute, allowing us to predict model quality well beyond the scale at which the law is fit. We further show that downstream robotic planning performance improves predictably with compute, and that imagination error is strongly correlated with it, making it a reliable proxy for real-robot evaluation. Finally, we demonstrate that latent world models can be deployed zero-shot as robotic agents, planning toward a single goal image to solve tasks requiring long-horizon planning on real hardware. We release all model checkpoints together with our training and robot deployment code. To our knowledge, this is the first work to establish scaling laws for multi-embodiment robotic world models trained on real robot data, and RoboJEPA, at 8B parameters, is the largest JEPA predictor model trained to date.

Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan +9

FAIR at Meta · Chandar Research Lab · Mila - Quebec AI Institute +2

World Models Latent World Models World Model-Based Planning World Model Evaluation Joint-Embedding Predictive Architecture Cross-Embodiment Robot Learning

Ask PDF

Aug 6, 2026

XEWorld: Can Action-Conditioned World Models Generalize to Unseen Robot Embodiments?

Action-conditioned world models are promising learned simulators for robotic manipulation, yet evaluating them exclusively on training robots fails to reveal whether they capture physical dynamics or merely memorize visual patterns. To answer whether a model can faithfully render a robot it has never seen, we introduce XEWorld, a controlled cross-embodiment testbed for world models that isolates embodiments by evaluating held-out robots within physically identical scenes. Our systematic analysis uncovers a shared architectural bottleneck: current models act primarily as 2D visual pattern matchers whose generalization is governed by visual similarity rather than physical kinematic similarity. Driven by this limitation, they struggle to translate abstract numeric joint actions into coherent visual trajectories, and fail to predict dynamic visual changes from static initial observations. Consequently, successfully rendering an unseen embodiment zero-shot strictly requires heavily grounded cues, specifically pixel-space actions and explicit spatial-temporal alignment. Even when bypassing this zero-shot barrier via few-shot adaptation, the forced appearance recovery triggers catastrophic forgetting of seen embodiments. Together, these failures expose a critical inability to apply learned physical dynamics to novel visual appearances, highlighting that achieving true cross-embodiment generalization requires architectural innovations that decouple visual appearance from underlying physical dynamics.

Yixiang Chen, Jiabing Yang, Yuan Xu +10

New Laboratory of Pattern Recognition (NLPR), Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Shandong University +1

World Model Learning Action-Conditioned World Models Robot Manipulation Benchmarks Cross-Embodiment Transfer

Ask PDF

Apr 30, 2026

World Model for Robot Learning: A Comprehensive Survey

World models, which are predictive representations of how environments evolve under actions, have become a central component of robot learning. They support policy learning, planning, simulation, evaluation, data generation, and have advanced rapidly with the rise of foundation models and large-scale video generation. However, the literature remains fragmented across architectures, functional roles, and embodied application domains. To address this gap, we present a comprehensive review of world models from a robot-learning perspective. We examine how world models are coupled with robot policies, how they serve as learned simulators for reinforcement learning and evaluation, and how robotic video world models have progressed from imagination-based generation to controllable, structured, and foundation-scale formulations. We further connect these ideas to navigation and autonomous driving, and summarize representative datasets, benchmarks, and evaluation protocols. Overall, this survey systematically reviews the rapidly growing literature on world models for robot learning, clarifies key paradigms and applications, and highlights major challenges and future directions for predictive modeling in embodied agents. To facilitate continued access to newly emerging works, benchmarks, and resources, we will maintain and regularly update the accompanying GitHub repository alongside this survey.

Bohan Hou, Gen Li, Jindou Jia +15

Nanyang Technological University · University of California, Berkeley · Stanford University +6

World Model Learning World Models for Robotics Embodied World Models Robot Learning Video World Models

Ask PDF

Abstract

Robot Motion Prediction

Explore similar work

RoboJEPA: Scaling Robotic Latent World Models

XEWorld: Can Action-Conditioned World Models Generalize to Unseen Robot Embodiments?

World Model for Robot Learning: A Comprehensive Survey