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

RoboJEPA: Scaling Robotic Latent World Models
Oct 7, 2026Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan +9World Models Latent World Models
FAIR at Meta · Chandar Research Lab · Mila - Quebec AI Institute +2
XEWorld: Can Action-Conditioned World Models Generalize to Unseen Robot Embodiments?
Aug 6, 2026Yixiang Chen, Jiabing Yang, Yuan Xu +10World Model Learning Action-Conditioned World Models
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 for Robot Learning: A Comprehensive Survey
Apr 30, 2026Bohan Hou, Gen Li, Jindou Jia +15World Model Learning World Models for Robotics
Nanyang Technological University · University of California, Berkeley · Stanford University +6

Oct 7, 2026J/K move · Enter open · S save

RoboJEPA: Scaling Robotic Latent World Models

Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan, Sarath Chandar+8

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

Open paper AskPDF arXiv

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.

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