WestWorld

WestWorld: A Knowledge-Encoded Scalable

Trajectory World Model for Diverse Robotic Systems

ICML 2026 Spotlight
Paper Code Dataset

TL;DR

  • We propose WestWorld, a novel system-aware MoE architecture for scaling up the training of trajectory world models across diverse robotic systems.
  • We introduce knowledge-encoded structural embedding that provides an explicit inductive bias to enhance zero- and few-shot generalization to unseen robotic systems.
  • We conduct extensive experiments to verify the scalability and generalizability of our method, and apply it to a downstream model-based control task, demonstrating strong performance in long-horizon planning.

Methodology

WestWorld framework overview

Figure 1: WestWorld Framework Overview. The overall architecture of our proposed WestWorld, consisting of two core components: (a) a Knowledge-Encoded Embedding Module 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.

Experiments

Evaluation on Zero-shot Performance

Method Walker2D (×10⁻²) Hopper (×10⁻²) Franka (×10⁻²)
MAE ↓ MSE ↓ MAE ↓ MSE ↓ MAE ↓ MSE ↓
MLP Ensemble 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
WestWorld (Ours) 16.350 5.064 13.731 3.368 7.737 2.539

Table 1: Zero-shot generalization performance. We evaluate zero-shot performance on three unseen robotic environments: Hopper and Walker2D from D4RL, and a real-world Franka manipulator. WestWorld outperforms all baselines across all systems, demonstrating that our system-aware MoE architecture and structural inductive bias effectively capture distinct dynamics while mitigating task interference.

Trajectory visualization for all robots

Figure 2: Trajectory visualization for zero-shot generalization. WestWorld tracks the ground-truth dynamics substantially more closely than the baselines over the 100-step horizon. In a zero-shot setting on unseen but structurally similar systems, our model selects appropriate experts to produce accurate dynamics predictions, whereas baseline methods lack morphology-aware representations and fail to generalize.

Evaluation on Few-shot Adaptation

Method Cassie (×10⁻²) A1 (×10⁻²) UR5 (×10⁻²)
MAE ↓ MSE ↓ MAE ↓ MSE ↓ MAE ↓ MSE ↓
MLP Ensemble 14.369 ± 0.523 4.416 ± 0.368 14.357 ± 1.009 3.707 ± 0.523 15.181 ± 0.597 4.936 ± 0.322
TDM 14.510 ± 0.642 3.404 ± 0.300 10.624 ± 0.312 2.151 ± 0.109 18.578 ± 0.324 5.510 ± 0.175
TrajWorld 7.834 ± 0.167 1.697 ± 0.109 5.138 ± 0.200 0.900 ± 0.050 8.066 ± 0.799 2.117 ± 0.433
WestWorld (Ours) 5.316 ± 0.108 0.808 ± 0.025 4.227 ± 0.120 0.628 ± 0.040 4.925 ± 0.317 0.831 ± 0.150

Table 2: Few-shot adaptation performance. We evaluate few-shot adaptation (10 episodes per system) on three real-world datasets with significant domain gap: Cassie bipedal jumping, Unitree A1 quadruped locomotion, and UR5 tabletop manipulation. Despite the large morphology and dynamics gap, WestWorld consistently outperforms all baselines, showing that pretraining provides strong initialization even for systems with substantial domain differences.

Evaluation on Scalability

Scalability comparison

Figure 4: Scalability comparison. We pretrain on N = 1, 2, 5, 10, 20, 30, 50, 60, 89 environments. WestWorld's accuracy stays consistently low as N grows, while TrajWorld degrades substantially due to gradient interference, demonstrating our model's scalability across diverse environments.

MoE routing weights heatmap

Figure 3: Sys-MoE routing weights across six layers (L1–L6) with four experts (E1–E4). Brighter colors indicate higher expert activation. The router exhibits sparse, system-dependent expert selection, where different systems use different expert combinations to capture their distinct dynamics. This validates our key insight that complex dynamics can be approximated by composing basis dynamics modules, mitigating interference and enabling scalable pretraining.

Evaluation on Downstream Control Task

Method Pretrain Walker2D Hopper Go1
Reward ↑ Reward ↑ Reward ↑
MLP Ensemble ✗ 119.18 147.37 -1.07
✓ 190.51 200.56 -0.31
TDM ✗ 122.65 242.34 -0.72
✓ 207.61 154.64 0.03
TrajWorld ✗ 395.23 366.26 0.05
✓ 1933.52 534.32 0.49
WestWorld (Ours) ✗ 707.61 554.92 0.43
✓ 2134.60 2253.51 2.20

Table 3: Downstream control performance using MPPI. We report accumulated episode reward averaged over evaluation episodes, comparing fine-tuning from pretrained checkpoint vs. training from scratch on Walker2D, Hopper, and Go1. Pretraining consistently improves control performance across nearly all methods, and WestWorld achieves the best performance under both regimes with particularly large gains after pretraining.

Pretraining Yields Largest Gains for WestWorld

Go1 Quadruped (Walk Straight)

With Pretraining
Without Pretraining

Hopper (Hop Forward)

With Pretraining
Without Pretraining

Walker2D (Walk Forward)

With Pretraining
Without Pretraining

Real-World Deployment Comparison on the Go1 Quadruped

TrajWorld (Baseline)
2x
WestWorld (Ours)
2x
Our WestWorld outperforms the strongest baseline TrajWorld in real-world deployment on the Go1 quadruped. Our method successfully completes the straight-walking task and reaches the goal, whereas TrajWorld fails to reliably stand up and move toward the target. Because MPPI relies on accurate long-horizon rollouts, TrajWorld's inaccurate predictions lead to poor action selection and control failure. In contrast, our method remains accurate enough to support stable real-world control, despite being trained only on an offline simulation dataset.

Ablation Study

Method Sys-MoE Structural Embedding Walker2D (×10⁻²) Hopper (×10⁻²) Franka (×10⁻²)
MAE ↓ MSE ↓ MAE ↓ MSE ↓ MAE ↓ MSE ↓
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
WestWorld (Ours) ✓ ✓ 16.350 5.064 13.731 3.368 7.737 2.539

Table 4: Ablation study under zero-shot setting. We ablate Sys-MoE by replacing it with a dense SSM, and ablate structural encoding by removing it during pretraining. Removing structural embedding degrades performance on complex morphologies (Hopper, Walker2D), indicating that modeling physical connectivity improves generalization. Replacing Sys-MoE with dense SSM also degrades performance, confirming it mitigates inter-task interference. Both components are essential for model generalization and scalable pretraining.

Citation

If you find WestWorld useful in your research, we would appreciate it if you consider citing our work:

@inproceedings{wang2026westworld,
  title={WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems},
  author={Wang, Yuchen and Kong, Jiangtao and Wei, Sizhe and Li, Xiaochang and Lin, Haohong and Zhao, Hongjue and Zhou, Tianyi and Gan, Lu and Shao, Huajie},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=ncRRCG4BfP}
}