AnyGoal: On World Model Generalization Beyond Expert Rollouts

Heejeong Nam; Divyat Mahajan; Quentin Le Lidec; Randall Balestriero; David Klindt

cs.AIOct 19, 2026

AnyGoal: On World Model Generalization Beyond Expert Rollouts

Authors: Heejeong Nam, Divyat Mahajan, Quentin Le Lidec, Randall Balestriero, David Klindt

Organizations: Brown University · Cold Spring Harbor Laboratory · Mila, Université de Montréal · New York University

Abstract

The goal of a world model is to capture the dynamics of an environment. This has the great appeal that it allows agents to plan and simulate a course of actions with the help of a world model. For these uses to be reliable, however, the learned dynamics must behave well beyond the specific state-action combinations observed in expert rollouts. Despite this, the relationship between data coverage, dynamics generalization, and downstream planning performance remains underexplored. We study this relationship through the lens of compositional generalization: can a world model recombine familiar state and action factors into valid dynamics? To answer this, we introduce AnyGoal, a benchmark of goal-shifted task variants built from widely used world-model environments, together with an open-source library. Each task is designed to be within reach of a world model that has learned the underlying dynamics: individual state factors and actions remain within their training support, but their combinations are underrepresented in expert trajectories. Across multiple environments and world model families, we find that current models consistently exhibit failures under these recombinations. By examining dynamics and representation failures across model families, interaction regimes, and evaluation metrics, we study where current world models generalize and where challenges remain for goal-shifted prediction and planning. Together, this framework, benchmark, and analysis provide a basis for evaluating progress toward world models that generalize to new planning tasks.

Figures & tables

Figure 2: Overview of AnyGoal. Each task preserves the environment dynamics of the corresponding base tasks and keep states and actions individually within training coverage. Throughout the paper, AnyGoal is used solely as a test set for evaluating generalization.

Figure 3: Compositional support structure. $\mathcal{C}^{\mathrm{te}}\subseteq\mathcal{C}^{\times}$ and $C^{te}\not\subset\mathcal{C}^{\mathrm{tr}}$ is necessary.

Figure 4: AnyGoal targets the combinatorial gap. (a) Example factor-wise support check for Avoid-T relative to Push-T using KDE. Darker regions indicate higher density, and gray regions mark areas whose density is below $1\%$ of the peak reference density. (b) Estimated overlap between each target variant and the base-task training distribution via KDE.

Model	Family	Representation	Action conditioning
End-to-End World Models
LeWorldModel	JEPA	Learned latent	AdaLN
Sensorimotor WM	JEPA	Learned latent	AdaLN
DINO-WM	JEPA	Frozen DINO features	Feature concatenation
DriftWorld	Generative	Pixels	Conditional generator
IRASim	Generative	Frozen SDXL VAE	DiT conditioning

Table 1: End-to-End World models and oracle state predictors evaluated in our experiments.

Figure 6: Oracle-state predictors reveal dynamics-level failures. Even with task-relevant state factors provided directly, goal shifts still degrade prediction and planning, showing that the AnyGoal gap cannot be reduced to visual representation failure. See Appendix I.1 for absolute metrics.

Figure 7: End-to-end degradation reflects representation effects. Compared with oracle-state predictors, end-to-end models degrade more under goal shifts, suggesting additional failures from observation-level shift and learned representation geometry. See Appendix I.1 for absolute metrics.

Figure 8: Representation geometry matters under goal shifts. (a) The observations remain visually in support, but end-to-end models degrade while oracle-state predictors remain stable. (b) The 2D t-SNE projections illustrate that learned representations can organize the same underlying agent positions with substantially different geometry, which can affect prediction and planning.

Figure 9: Goal-biased training support is associated with reduced action sensitivity. Each panel plots model action sensitivity against an environment-specific bias coordinate induced by the expert policy. The analysis suggests that the model can partially rely on a goal-directed prior rather than the supplied action.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Reference dataset $P$	Query dataset $Q$	Factor groups checked
TwoRoom	TR-Route	Agent position $(2\mathrm{D})$ ; action position delta $(2\mathrm{D})$
PushT	Avoid-T; Stack-T; Push-T Antigoal	Agent position $(2\mathrm{D})$ ; block position $(2\mathrm{D})$ ; action $(2\mathrm{D})$ ; block angle.
Insert-Peg	Avoid-Peg; Extract-Peg	Arm root angle; Arm shoulder angle; Arm elbow angle; Arm wrist angle; Finger angle; Fingertip angle; Thumb angle; Thumbtip angle; hand, peg, and target position (2D); hand, peg, and target orientation (2D) ; 5 motor commands ( $1\mathrm{D}$ ) independently.
Cube	Cube-Hop; Cube-Diagonal	End-effector position $(3\mathrm{D})$ ; block position $(3\mathrm{D})$ ; block quaternion $(4\mathrm{D})$ ; action position delta $(3\mathrm{D})$ ; gripper opening; gripper contact; end-effector yaw; action yaw delta; action gripper delta. ; 6 joint pos independently

Appendix

Table 2: Every independent factor is in-distribution with respect to the reference dataset. Factor groups used for support verification. For each reference dataset

P

, we compare the corresponding goal-shifted variant datasets

Q

against the listed task-relevant factors..

$P$	$Q$	Factor of variants	Mahalanobis 0.99	KDE	Range	Final verdict
TwoRoom	TR-Route	agent pos (2D)	0.970	1.000	1.000	pass
TwoRoom	TR-Route	action (2D)	1.000	0.957	1.000	pass
PushT	Avoid-T	block angle	0.970	0.963	–	pass
		agent pos (2D)	0.968	0.984	1.000	pass
		block pos (2D)	1.000	0.990	1.000	pass
		action (2D)	1.000	0.995	1.000	pass

Appendix

Table 3: Distribution-support verification between each reference dataset (

P

) and query dataset (

Q

). A dash indicates that the range criterion is not applied to circular variables.( All Passed )

Family	$s_{\mathrm{agent}}$	$s_{\mathrm{all}}$	$[a,s_{\mathrm{all}}]$
TwoRoom	position (2D)	position (2D)	4D
PushT	agent position (2D)	$+$ block pose (5D total)	7D
Cube	EE/gripper state (6D)	$+$ block pose (13D total)	18D
Manipulator	joints $+$ hand pose (20D)	$+$ peg/target poses (28D total)	33D

Appendix

Table 4: Components and dimensionalities of the full-dimensional joint KDE.

Base environment	Dim.	Oracle state variables
TwoRoom	2D	Agent position $(x,y)$ .
Push-T	10D	Goal position $(x,y)$ , goal orientation $(\sin\theta,\cos\theta)$ , block position $(x,y)$ , block orientation $(\sin\theta,\cos\theta)$ , and agent position $(x,y)$ .
Cube	15D	End-effector position $(x,y,z)$ , end-effector yaw $(\cos\theta,\sin\theta)$ , gripper opening, gripper contact, gripper velocity, block position $(x,y,z)$ , and block quaternion $(w,x,y,z)$ .
Manipulator Insert-Peg	33D	Arm joint positions, touch sensors, hand pose $(x,z,q_{w},q_{y})$ , peg pose $(x,z,q_{w},q_{y})$ , and target pose $(x,z,q_{w},q_{y})$ .

Appendix

Table 5: Oracle state variables used by predictor-only models for each base environment.

Dataset	Model	LR	Configured duration	Selected checkpoint
Push-T	LeWM	-	-	official LeWM checkpoint
Push-T	Sensorimotor-WM	$1\times 10^{-4}$	10 epochs	epoch 10
Push-T	DINO-WM	-	-	official DINO-WM checkpoint
Push-T	DriftWorld	$5\times 10^{-4}$	max 500k steps	step 500k
Push-T	IRASim	$5\times 10^{-5}$	max 300k steps	step 300k
Cube	LeWM	-	-	official LeWM checkpoint

Appendix

Table 6: Training details and selected checkpoints for end-to-end world models.

Dataset	Model	selected LR	Configured duration	Selected checkpoint
Push-T	Toy-AdaLN	$3\times 10^{-4}$	100 epochs	epoch 100
Push-T	Toy-Action-Duplicate	$3\times 10^{-4}$	100 epochs	epoch 100
Push-T	Toy-DriftWorld	$1\times 10^{-4}$	50 epochs, max 500k steps	step 500k
Cube	Toy-AdaLN	$3\times 10^{-4}$	100 epochs	epoch 100
Cube	Toy-Action-Duplicate	$3\times 10^{-4}$	100 epochs	epoch 100
Cube	Toy-DriftWorld	$1\times 10^{-4}$	50 epochs, max 500k steps	step 500k

Appendix

Table 7: Training details for the selected oracle-state predictor checkpoints.

Environment	Planning success criterion
TwoRoom	The 2D Euclidean distance between the agent position and the goal agent position is less than 16 pixels. Target and door positions are not used.
PushT	The 4D Euclidean error over agent and block positions is less than 20 pixels, and the block angle error is less than $\pi/9$ radians. The angle error accounts for object symmetries.
Cube	The rollout must simultaneously satisfy: end-effector 3D position error $\leq 0.04$ m, end-effector yaw error $\leq 0.10$ rad, and cube 3D position error $\leq 0.04$ m. Cube orientation, gripper opening, contact, and velocity are not used.
Insert Peg	The rollout must simultaneously satisfy: hand 2D position error $\leq 0.04$ m, hand planar angle error $\leq 0.10$ rad, peg 2D position error $\leq 0.04$ m, and peg planar angle error $\leq 0.10$ rad. Arm joints, touch, velocity, and the physical target or slot pose are not used.

Appendix

Table 8: Planning success criteria. Goals are defined by the dataset future state 25 raw steps after the initial state. An episode is successful if the criterion is satisfied at least once within 50 environment steps.

Base-Env	Horizon	LeWM	SMWM	DINO-WM	DriftWorld	IRASim	AdaLN	Action Dup.	Toy-Drift
TwoRoom	5	0.065188	0.000058	0.033841	0.000154	0.003838	0.00000466	0.0000144	0.000051
TwoRoom	25	0.511843	0.000754	0.23277	0.002164	0.023646	0.000145	0.000636	0.00111671
PushT	5	0.00743	0.02933	0.05124	0.001331	0.005813	0.000005	0.000028	0.000066
PushT	25	0.16739	0.23039	0.48649	0.018599	0.050163	0.000229	0.001068	0.00388751
Insert Peg	5	0.000717	0.002913	0.172737	0.001448	0.00939	0.003526	0.009356	0.054486
Insert Peg	25	0.010678	0.024729	1.410328	0.013113	0.0593	0.053553	0.149251	0.404549

Appendix

Table 9: Absolute base-task rollout MSE at 5- and 25-step prediction horizons.

	End-to-End World Models						Oracle State Predictors
Env	Variant	LeWM	SMWM	DINO-WM	DriftWorld	IRASim	AdaLN	Action Dup.	Toy-Drift
TwoRoom	TR-Route	14.56	19.11	6.56	0.012	0.9	0.045	1.03	0.014
	AvoidT	4.14	5.2	5.77	0.71	0.77	1.26	1.25	1.32
	StackT	11.33	3.13	3.38	2.54	1.36	501.14	113.55	70.24
PushT	Anti-goal	6.23	2.66	2.83	1.64	1.09	54.74	12.85	19.96
	Avoid-Peg	6.83	2.54	1.3	7.01	7.1	1.45	0.62	0.08

Appendix

Table 10: Base-relative one-step (5 environment step) prediction error. Values report the error ratio with respect to the corresponding base task.

	end-to-end world models						oracle state predictors
Env	Variant	LeWM	SMWM	DINO-WM	DriftWorld	IRASim	AdaLN	Action Dup .	Toy-Drift
TwoRoom	TR-Route	11.2	41.6	9.66	0.01	0.82	0.062	0.816	0.01
	StackT	9.83	4.56	3.8	2.34	1.44	158.5	47.57	22.09
PushT	Anti-goal	8.49	4.27	3.94	1.72	1.28	29.97	12.68	8.96
	Avoid-Peg	9.73	16.69	2.8	1.14	1.12	0.67	0.42	0.09
Insert Peg	Extract-Peg	4.08	7.15	1.78	2.75	1.19	34.39	12.52	4.49

Appendix

Table 11: Base-relative five-step (25 environment step) prediction error. Values report the error ratio with respect to the corresponding base task.

Task	LeWM	SMWM	DINO-WM	DriftWorld	IRASim	AdaLN	Action Dup.	Drift Toy
TwoRoom	89.30	100.00	100.00	90.00	76.00	100.00	99.30	100.00
TR-Route	60.70	64.70	68.00	76.00	63.30	100.00	100.00	100.00
Push-T	84.70	84.70	74.00	43.30	5.30	92.70	92.00	89.30
AvoidT	3.30	4.70	6.70	0.70	0.00	100.00	100.00	99.30
StackT	20.70	20.00	8.70	10.70	2.00	70.70	66.70	54.00
Push-T Antigoal	34.00	44.00	27.30	14.70	5.30	79.30	81.30	75.30

Appendix

Table 12: Absolute planning success rate (%) across different models and tasks.

World Models

World models represent and predict how environments change, helping agents plan, act, and reason about possible futures.

See all 359 papers

stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation
May 20, 2026Lucas Maes, Quentin Le Lidec, Luiz Facury +9World Models ML Reproducibility
Mila & Université de Montréal · New York University · Universidade Federal de Minas Gerais +4
World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models
Jul 30, 2026Xiangcheng Zhang, Yilun DuWorld Models for Robotics World Model-Based Planning
Harvard University
Qwen-AgentWorld: Language World Models for General Agents
Jun 23, 2026Yuxin Zuo, Zikai Xiao, Li Sheng +30LLM World Models Agentic RL
Qwen Team

May 20, 2026J/K move · Enter open · S save

stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation

Lucas Maes, Quentin Le Lidec, Luiz Facury, Nassim Massaudi+8

Mila & Université de Montréal · New York University · Universidade Federal de Minas Gerais+4

World Models ML Reproducibility World Model Evaluation

Open paper AskPDF arXiv

World models are central to building agents that can reason, plan, and generalize beyond their training data. However, research on world models is currently fragmented, with disparate codebases, data pipelines, and evaluation protocols hindering reproducibility and fair comparison. Current practice is further limited by three key bottlenecks: fragile one-off codebases, slow video data loading, and the lack of standardized generalization benchmarks. We present stable-worldmodel (swm), an open-source platform for standardized and reproducible world modeling research and evaluation. It delivers (1) a high-performance Lance-based data layer with native support and conversion tools for MP4, HDF5, and LeRobot datasets, (2) clean, well-tested implementations of modern world model baselines and planning solvers, and (3) a broad suite of environments and tasks extended with controllable visual, geometric, and physical factors of variation for systematic in-silico evaluation of dynamics understanding, control performance, representation quality, and out-of-distribution generalization. By unifying the full pipeline under a single, scalable framework, swm dramatically reduces research overhead and accelerates trustworthy progress toward reliable world models.

Abstract

World Models

Explore similar work

stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

Qwen-AgentWorld: Language World Models for General Agents