Skip to content
AI.info

Research

AnyGoal: On World Model Generalization Beyond Expert Rollouts

AnyGoal: On World Model Generalization Beyond Expert Rollouts Overview Research area: World models, model-based reinforcement learning, and compositional generalization (arXiv:2610.26844v1 [cs.AI], 19

AnyGoal: On World Model Generalization Beyond Expert Rollouts
arXiv
2610.26844
Published
2026-10-19
Authors
Heejeong Nam, Divyat Mahajan, Quentin Le Lidec, Randall Balestriero, David Klindt

AI summary

AnyGoal: On World Model Generalization Beyond Expert Rollouts

Overview

Research area: World models, model-based reinforcement learning, and compositional generalization (arXiv:2610.26844v1 [cs.AI], 19 Oct 2026).

Technical level: Advanced. The paper combines a formal treatment of support and distribution shift with an empirical benchmark and a multi-model evaluation suite.

Scope: The paper introduces AnyGoal, a benchmark of goal-shifted task variants plus an open-source library, to test whether world models trained on expert rollouts can recombine familiar states and actions into valid dynamics for new goals.

What This Paper Is About

World models learn how an environment evolves under actions so that agents can simulate and plan. When training data comes only from expert, goal-directed rollouts, states and actions become strongly correlated, and a model can fit those trajectories by learning the behavior policy's habits instead of the actual effect of the supplied action. The paper asks a controlled version of this question: can a world model generalize to new goals when every individual state factor and action stays inside its training support, but their joint combinations are rare or absent in the expert data?

Key Contributions

  1. A formal criterion for reusable dynamics. The authors formalize the desired behavior as compositional generalization: test trajectories must satisfy C^te ⊆ C^× (inside the Cartesian product of marginal training supports) while C^te ⊄ C^tr (outside the joint training support). They show that an action-ignorant predictor T_deg(z, a) = T*(z, π_E(z)) achieves zero prediction error on an expert dataset while ignoring the queried action, so expert-rollout accuracy cannot distinguish a true dynamics model from a behavior-cloning shortcut.

  2. The AnyGoal benchmark. Eight goal-shifted variants built on four widely used world-model environments: TwoRoom (Sobal et al., 2025), Push-T (Chi et al., 2025), DMC Manipulator Insert-Peg (Tassa et al., 2018), and OGBench Cube (Park et al., 2025). Each base task gets one free-space variant, and all but TwoRoom also get one or two contact-rich variants. The release contains 400 episodes per variant plus data-collection tools, and AnyGoal is used solely as a test set.

  3. A statistically controlled support-verification protocol. Three factor-wise checks (range, Mahalanobis at the reference 0.99-quantile requiring at least 95% of query samples below it, and KDE at the 0.01-quantile requiring at least 95% of query samples above it) confirm that target variants stay within the combinatorial space. A separate full-dimensional Gaussian KDE at the 0.01-quantile of log-densities measures how much of each target distribution falls in the combinatorial gap rather than in the base-task joint support.

  4. A multi-model diagnostic study and an open-source library. The paper evaluates JEPA-style and generative world models, end-to-end and oracle-state predictors, on both rollout prediction error and CEM-based MPC planning success from Stable-WorldModel (Maes et al., 2026a). The AnyGoal Python library provides a unified CLI and Python APIs for dataset access, data collection, training, and evaluation.

Main Findings

  • Goal shifts break current world models broadly. Failures appear across modeling paradigms, including generative world models (DriftWorld, IRASim) and JEPA-style architectures (LeWM, SMWM, DINO-WM). The authors state that current models consistently exhibit failures under factor recombination.

  • Oracle-state predictors still fail, so this is not only a vision problem. Even when task-relevant state factors are provided directly, goal shifts degrade prediction and planning. Contact-rich variants degrade consistently, while free-space variants such as TR-Route, Avoid-T, Avoid-Peg, and Cube-Hop preserve or even improve MPC performance relative to the base task.

  • Contact-rich tasks show the largest dynamics errors. With oracle states, contact-rich variants show rollout error increases often by one to three orders of magnitude, whereas free-space variants degrade far less and sometimes match or improve on base-task error.

  • Prediction error and planning success can disagree. In free-space variants, planning tolerates imperfect rollout prediction because many trajectories reach the goal. In contact-rich tasks the two metrics align more closely, since fewer action sequences succeed and success depends on correct object interaction.

  • End-to-end models degrade more than oracle-state models. Recombining marginally familiar factors can produce rendered observations outside the joint visual support of expert rollouts, adding an encoder-inference challenge on top of dynamics prediction.

  • Latent geometry matters independently of visual shift. In TwoRoom, agent position is the only changing factor and TR-Route positions remain marginally within training support, so observations stay visually in support. End-to-end models still degrade while oracle-state predictors remain stable, and 2D t-SNE projections show the same underlying agent positions organized with substantially different geometry across models.

  • Models learn behavior-policy shortcuts. An interventional action-sensitivity score S_θ(s), computed by substituting actions from a fixed interventional set A_cf and measuring the change in the predicted representation, is systematically lower in goal-biased regions. Examples: motion toward the narrow passage in TwoRoom, approaching the block before pushing in Push-T, carrying the grasped peg toward the target in Insert-Peg, and prevalent vertical actions at low end-effector heights in Cube. The action-sensitivity plot in Figure 1 shows reduced sensitivity in biased regions of Push-T.

  • A theoretical route forward. The paper sketches regime-dependent dynamics z_{t+1} = T_{r_t}(z_t, a_t) where each next-state factor depends on only a small subset of coordinates (‖∇x T{k,r_t}(x)‖_0 ≪ dim(x)). Under this sparsity, a test composition can be unseen (C^te ⊄ C^tr) while every mechanism input stays in support (U_k^te ⊆ U_k^tr for all k), so no individual mechanism needs extrapolation. A formal statement appears as Proposition 1 in Appendix J.

Methodology in Plain English

The authors start by defining what "generalizing beyond expert rollouts" should mean. They separate the state of the environment into semantic factors (object position, end-effector pose, and so on) and describe training and test distributions by the sets of factor and action values they contain. A test task qualifies only if every individual factor and action is familiar but the full combination is not.

They then build such tasks from existing environments by changing the goal while keeping the environment's physics and dynamics identical. Push-T becomes Avoid-T: instead of pushing the block to a target, the agent must avoid the block using the same 2D navigation. Because the underlying dynamics are unchanged, any failure must come from the model, not from new physics.

Before running any model, they verify each variant statistically: range checks, Mahalanobis distance checks against the reference 0.99-quantile, and KDE density checks at the 0.01-quantile, all requiring at least 95% of query samples to pass. They then run a separate joint KDE test to confirm that the combinations themselves are underrepresented in the training data.

Models are trained only on the base task and evaluated on both base and goal-shifted trajectories. Prediction error is measured as MSE in each model's native prediction space, autoregressively over horizons of 5 and 25 environment steps with frame skip of 5, sampling 100 episodes per seed over seeds 0, 1, 2. Planning is measured with the same CEM-based MPC configuration for every model, using a 25-step planning horizon where the full action sequence is executed before replanning.

To separate representation failure from dynamics failure, they also train lightweight predictors directly on oracle states. To diagnose behavioral shortcuts, they intervene on the action input while holding the state fixed and measure how much the model's representation changes.

Why This Matters

Impact on research. The paper reframes world-model evaluation around compositional generalization rather than in-distribution accuracy, and it shows why the distinction is not academic: a model can achieve zero prediction error on expert data while ignoring the action entirely. It also supplies a statistically grounded protocol for constructing valid generalization tests, so that OOD evaluations can be distinguished from tests that require extrapolation to genuinely novel factors.

Real-world applications (as implied by the framing):

  • Robot manipulation. Contact-rich tasks such as peg insertion and pushing are exactly where the oracle-state experiments show the largest degradation, which matters for planning with learned dynamics in manipulation pipelines.

  • Offline model-based control. The findings connect directly to support mismatch in offline RL, where agents must plan with models trained on fixed behavior datasets.

  • Simulation and predictive modeling for physical systems. Any setting where a dynamics model is trained on goal-directed trajectories and then queried for a different objective faces the combinatorial gap the paper isolates.

  • Evaluation infrastructure for learned world models. The AnyGoal library and standardized CLI give practitioners a way to test a base-trained model on unseen goals without building task-specific harnesses.

Industry relevance. Teams deploying model-based planning on top of pretrained world models need to know whether their model responds to the action they supply or to the action an expert would have taken. The action-sensitivity diagnostic is cheap to compute and targets that question directly.

Future Directions

  • Identifying reusable mechanisms from offline data. The authors frame this as open: fitting training trajectories does not guarantee that a predictor has learned reusable, action-conditioned mechanisms, and this becomes harder when action effects vary with state and interaction regime.

  • Extrapolating beyond the training support in interaction-rich dynamics. Even with the right mechanisms, behavior on novel compositions is undetermined without structural assumptions. The sparse, regime-dependent formulation in Section 6 is offered as a starting point rather than a solution.

  • Building architectures and objectives that encode the sparsity assumption. The paper argues for world models that need coverage only of local mechanism inputs rather than of the full global state–action space, which implies new training objectives or architectural constraints.

  • Extending the benchmark beyond the eight current variants. AnyGoal ships with data-collection tools for generating additional datasets, and the paper's five interpretive dimensions (failure source, evaluation metric, task complexity, representation quality, prediction space) suggest axes along which the benchmark could be widened.

Target Audience

Researchers and graduate students working on world models, model-based reinforcement learning, and offline RL who want a rigorous account of what "generalization" should mean for learned dynamics. It is also useful for practitioners deploying learned simulators for planning, and for anyone designing OOD or compositional benchmarks, since the support-verification protocol transfers to other domains. Readers need familiarity with distribution shift, latent-variable models, and model-predictive control to follow the formal sections and the planning experiments.

Note: the main text reports qualitative comparisons and refers to Appendix I.1 for absolute metric values; specific numerical results are not included in the provided content.

Authors’ 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.

Read the original paper