PaperlayerView exact revision on arXiv

Source-linked AI summary

AnyGoal: On World Model Generalization Beyond Expert Rollouts

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

arXiv:2610.26844v1cs.AI

TL;DR

World models are intended to support planning by learning environment dynamics, but evidence that they generalize beyond expert state–action combinations remains limited. This paper introduces AnyGoal, a benchmark of goal-shifted tasks designed around familiar factors with underrepresented combinations. Across models, these recombinations expose persistent failures whose severity varies with task and evaluation conditions.

  • Problem

    How well world models generalize to new state–action combinations matters for planning, but the relationship between data coverage, dynamics generalization, and planning performance remains underexplored.

    Source: §1 · pp. 1–2

  • Method

    AnyGoal evaluates goal-shifted variants that preserve dynamics and marginal factor support while targeting under-covered combinations of familiar states and actions.

    Source: §1 · p. 2, §3.2 · p. 4, §4 · p. 5

  • Results

    Current world models show broad failures on controlled goal shifts across modeling paradigms, with degradation varying by task complexity, evaluation metric, and model family.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Goal-shift generalization remains difficult even when relevant state and action factors are marginally covered by training data.

    Source: §7 · p. 11

    Compositional extrapolation is harder in interaction-rich dynamics, where action effects vary with joint physical configuration and regime.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper argues that in-distribution evaluation offers limited evidence that world models have learned reusable dynamics, since goal-directed rollouts leave state–action combinations under-covered. AnyGoal tests recombinations of familiar factors and reveals broad failures across model paradigms, with severity varying by task complexity, metric, and model family.

1 INTRODUCTION

  • AnyGoal preserves base-task dynamics and individual state and action coverage while testing underrepresented combinations as a generalization test set.The benchmark is designed to avoid requiring extrapolation to unseen factors.Source: Figure 2 · p. 2, §1 · p. 2
  • In Push-T, expert rollouts couple far-from-block states with actions toward the block, and a model may fail on Avoid-T despite familiar navigation actions.The model may learn rollout regularities rather than the action-conditioned effect of the supplied action.Source: §1 · p. 2
  • Goal-directed offline datasets can induce strong state–action correlations, making relevant joint-space coverage difficult to assess or fill.This motivates studying generalization under incomplete coverage rather than assuming data scale alone will resolve it.Source: §1 · p. 2
  • Prior OOD evaluations that introduce unsupported states or factors test robustness to novelty, whereas this paper targets recombinations within observed marginal support.The authors distinguish this controlled setting from evaluations that may exceed the generalization expected of reusable dynamics.Source: §1 · p. 2
  • AnyGoal exposes broad failures across generative world models and JEPA-style architectures on controlled goal shifts, with degradation varying by task complexity, evaluation metric, and model family.The failures can arise from dynamics prediction and representation learning.Source: §1 · pp. 2–3

2 RELATED WORKS

Prior work studies several forms of world-model generalization and support mismatch, but whether expert-trained models handle new goal-induced combinations of familiar states and actions is less understood. AnyGoal evaluates reusable dynamics under those goal shifts, complementing offline-RL approaches that constrain policies or model learning in poorly covered regions.

Compositional generalization diagnostics and benchmarks

  • Compositional benchmarks test whether familiar components can be composed correctly, and AnyGoal brings this diagnostic to world models trained on goal-directed trajectories.Related benchmarks span sequence, vision, multimodal, and mechanistic-interpretability settings.Source: §2 · p. 3

World-model generalization and support mismatch

  • AnyGoal addresses the less-studied question of whether expert-trained world models generalize to goal-induced recombinations of familiar states and actions.It complements offline-RL methods that constrain policies or use uncertainty- or pessimism-based objectives for poorly covered regions.Source: §2 · p. 3

3 COMPOSITIONAL GENERALIZATION FOR WORLD MODELS

The section frames reusable world-model dynamics as compositional generalization: predicting transitions for combinations of familiar factors and actions that training behavior under-covers. It explains how behavior-policy data can weaken action conditioning and defines the support shift AnyGoal evaluates.

3 COMPOSITIONAL GENERALIZATION FOR WORLD MODELS

  • Compositional generalization asks whether a world model can recombine familiar state and action factors into valid dynamics beyond observed behavior.The state is represented by semantic factors, with observations mapped from the underlying physical state.Source: §3 · p. 3

3.1 ACTION-CONDITIONED PREDICTION CAN DEGENERATE INTO BEHAVIOR PREDICTION

  • When offline data follow a behavior policy, a predictor may fit observed transitions while becoming action-ignorant in state regions with restricted conditional action support.The supplied action's effect can be effectively lost.Source: §3.1 · p. 4

3.2 COMPOSITIONAL DISTRIBUTION SHIFTS IN WORLD MODELS

  • Compositional distribution shift keeps transition dynamics invariant but tests state–action pairs absent from training joint support and within the product of training marginal supports.Training accuracy cannot determine how predictors extend to novel combinations in that gap.Source: §3.2 · p. 4
  • AnyGoal targets the combinatorial gap: test trajectories need not cover all of it, but goals are chosen to place substantial mass in under-covered regions.Figure 4 illustrates factor-wise support checks and estimated overlap with base-task training distributions.Source: §3.2 · p. 4, Figure 4 · p. 5

4 ANYGOAL: CONTROLLED GOAL-SHIFT EVALUATION

AnyGoal evaluates whether world models generalize to goal-shifted tasks that recombine familiar factors in underrepresented ways. It pairs factor-wise support checks with joint-support analysis, base-versus-shifted prediction and planning evaluation, and a reproducible library.

4 ANYGOAL: CONTROLLED GOAL-SHIFT EVALUATION

  • AnyGoal provides eight goal-shifted variants of four world-model tasks, with 400 episodes per variant and tools for collecting more data.The benchmark includes free-space variants for each base task and one or two contact-rich variants, except for TwoRoom.Source: §4 · p. 5

4.1 ANYGOAL TARGETS THE COMBINATORIAL GAP

  • All variants pass three factor-wise support checks, while a separate joint-density estimate measures overlap with base-task support and coverage of the combinatorial gap.The factor-wise verdict combines range, Mahalanobis, and KDE checks; the joint estimate uses full-dimensional KDE over concatenated actions and states.Source: §4.1 · p. 5

4.2 RELATIVE EVALUATION AGAINST BASE TASKS

  • Models trained on base data are evaluated on base and goal-shifted trajectories using both rollout prediction error and downstream MPC planning success.Prediction error is measured in each model’s native space at 5- and 25-step horizons; MPC executes a 25-step sequence before replanning.Source: §4.2 · pp. 5–6

4.3 ANYGOAL LIBRARY AND EVALUATION INTERFACE

  • The open-source library provides dataset access, task-specific collection, and standardized prediction and planning evaluation, including evaluation of the same base-trained model on Push-T and Avoid-T.Researchers can select models, datasets, and seeds through a unified CLI, or use Python APIs for collection, training, and evaluation.Source: §4.3 · p. 6, Figure 5 · p. 6

5 EXPERIMENT AND RESULTS

The experiments compare diverse end-to-end and oracle-state world models across goal shifts, prediction and planning metrics, and task types. Results show dynamics-level failures in contact-rich variants, additional representation-related degradation in end-to-end models, and reduced action sensitivity in goal-biased training regions.

5 EXPERIMENT AND RESULTS

  • The evaluation spans JEPA-style and generative objectives, learned and frozen visual encoders, and multiple action-conditioning mechanisms, with oracle-state predictors used to separate dynamics from visual representation effects.All models use frame skip 5 and one-step prediction; the analysis considers failure source, metric, task complexity, representation quality, and prediction space.Source: §5 · pp. 6–7, §5.1 · p. 7

5.2 DYNAMICS GENERALIZATION BEYOND REPRESENTATION FAILURE

  • Contact-rich goal shifts consistently degrade oracle-state MPC performance and produce rollout-error increases often one to three orders of magnitude, while free-space performance is preserved, improves, or degrades less.This asymmetry shows that supplying task-relevant state factors does not make goal-shifted dynamics prediction uniformly reliable.Source: Figure 6 · p. 7, §5.2 · p. 8
  • Rollout error and planning performance can diverge in free-space tasks, where many action sequences reach the goal despite imperfect intermediate predictions.Contact-rich tasks offer fewer viable sequences because success depends on correct object interaction.Source: §5.2 · p. 8

5.3 REPRESENTATION EFFECTS IN END-TO-END WORLD MODELS

  • End-to-end models can degrade on both metrics even when oracle-state predictors remain stable, indicating additional representation-level failures under goal shifts.Recombined familiar factors may produce observations outside the joint visual support of expert rollouts, creating an inference challenge before dynamics prediction.Source: §5.3 · pp. 8–9, Figure 7 · p. 8
  • In TwoRoom TR-Route, observations remain visually in support but end-to-end models degrade while oracle-state predictors remain stable, implicating latent geometry that may not preserve task-relevant state structure.The learned representation may obscure agent location, add nuisance variation, or organize states in a geometry unfavorable for prediction and planning.Source: §5.3 · p. 9, Figure 8 · p. 9

5.4 DIAGNOSING BEHAVIOR-POLICY SHORTCUTS THROUGH ACTION SENSITIVITY

  • Action sensitivity is reduced in state regions where expert behavior is strongly constrained by the goal, consistent with models partly relying on behavior-policy regularities.The reported examples include motion toward a narrow passage in TwoRoom, approaching the block in Push-T, and goal-directed peg or cube actions.Source: §5.4 · pp. 9–10, Figure 9 · p. 10

6 TOWARDS WORLD MODELS THAT GENERALIZE COMPOSITIONALLY

Oracle-state experiments show that compositional failures are not solely visual, and contact-rich goal shifts remain difficult. The section identifies mechanism discovery and extrapolation to novel compositions as distinct challenges.

6 TOWARDS WORLD MODELS THAT GENERALIZE COMPOSITIONALLY

  • Even with task-relevant physical state provided directly, goal shifts remain difficult, particularly for contact-rich prediction.This shows compositional failures cannot be attributed only to visual representation learning.Source: §6 · p. 10
  • Fitting training trajectories does not ensure reusable action-conditioned mechanisms, and identifying them from offline trajectories remains open.The challenge is sharper when action effects vary with state and interaction regime.Source: §6 · p. 10
  • Sparse transition mechanisms can generalize to unseen full compositions when each mechanism’s relevant inputs stay within training support and the mechanisms are correctly learned.In interaction-rich dynamics, effects may vary across regimes and joint physical configurations, making extrapolation harder without additional structural assumptions.Source: §6 · pp. 10–11

7 CONCLUSION

AnyGoal evaluates whether world models recombine familiar states and actions for new goals, and the paper finds this remains difficult across diverse models. Its analyses connect these failures to representation, dynamics, task structure, and evaluation metrics.

7 CONCLUSION

  • AnyGoal evaluates generalization to new behavioral goals by recombining familiar states and actions.The benchmark asks for prediction on novel joint state–action combinations within the product of marginal training supports.Source: §7 · p. 11
  • Across diverse models, goal-shift generalization remains difficult even when relevant state and action factors are marginally covered by training data.Failures are heterogeneous across representation, dynamics prediction, task structure, and evaluation metric.Source: §7 · p. 11
  • The paper argues that global coverage of all state–action combinations is neither realistic nor necessary; local mechanism-input coverage and learning action-conditioned mechanisms are needed.AnyGoal provides a controlled basis for evaluating progress toward reusable dynamics beyond expert rollouts.Source: §7 · p. 11

Supplementary Material

The supplementary discussion frames compositional evaluation as testing familiar physical factors under observed actions in combinations not selected by the behavior policy. It explains why this requires action-conditioned dynamics.

A COMPOSITIONAL GENERALIZATION: ESSENTIAL WORLD-MODEL TEST

  • A good action-conditioned world model should predict an observed action’s effect on familiar physical factors even when the behavior policy did not select that combination.This distinguishes recombination of observed factors from extrapolation to unseen factors.Source: A COMPOSITIONAL GENERALIZATION:… · p. 18
  • A predictor can achieve zero error on expert behavior data while ignoring the queried action, because training constrains it only on behavior-policy state–action pairs.The degenerate predictor reproduces the transition associated with the behavior action.Source: A COMPOSITIONAL GENERALIZATION:… · p. 18
  • When actions affect the next state, compositional evaluation rules out action-ignorant predictors and tests whether models can evaluate candidate actions beyond the behavior action.This is the action-conditioning requirement relevant to model-based control.Source: A COMPOSITIONAL GENERALIZATION:… · p. 18

B ANYGOAL TASK DETAILS

The task details pair familiar physical factors and actions with changed goals, while support checks verify factor-level coverage even when joint configurations are novel.

B.1 TWO ROOM AND ITS VARIANTS

  • TwoRoom isolates action-conditioned planar movement, while TwoRoom Route tests familiar positions and actions recombined outside dominant slot-crossing behavior.The route variant uses a waypoint-defined triangular route and mixes clean, noisy, and variable-speed actions within the observed support.Source: §B.1 · p. 18

B.2 PUSHT AND ITS VARIANTS

  • AvoidT tests navigation around the block without contact; StackT and PushT Anti-goal test contact and rotation under changed or reversed goal configurations.Anti-goal requires pushing away from the original goal region, contrary to the dominant expert behavior.Source: §B.2 · p. 19

B.3 INSERT-PEG AND ITS VARIANTS

  • Insert-Peg centers on structured insertion, Avoid-Peg changes to non-contact navigation, and Extract-Peg reverses insertion while retaining familiar objects and workspace.Extract-Peg balances left- and right-side box placements.Source: §B.3 · p. 19

B.4 OGBENCH CUBE AND ITS VARIANTS

  • Cube tests manipulation beyond pick-and-place behavior; CubeHop avoids object interaction, while CubeDiagonal changes the grasp configuration to diagonal contact.CubeDiagonal keeps the object, objective, and action space unchanged while introducing a different grasp geometry.Source: §B.4 · pp. 19–20

C.1 DISTRIBUTION SUPPORT CRITERIA

  • Support verification applies Mahalanobis, KDE, and coordinate-range checks to semantic factor groups, allowing full joint configurations to remain novel.The final verdict conjuncts all non-ignored criteria; circular factors omit the range check.Source: §C.1 · pp. 20–21

C.2 VERIFIED FACTOR GROUPS

  • All checked factors pass the support criteria, confirming the AnyGoal variants remain marginally in-distribution relative to their base datasets.The criteria include Mahalanobis, KDE, and observed-range checks.Source: Table 2 · p. 21

C.3 SUPPORT-VERIFICATION RESULTS

The support checks verify that AnyGoal variants use familiar state factors and actions, while separate analyses test whether their full combinations are supported. The appendix also specifies controlled interventions and a diverse set of end-to-end and oracle-state model baselines.

C.3 SUPPORT-VERIFICATION RESULTS

  • All checked factors pass the Mahalanobis, KDE, and observed-range support criteria, confirming marginal support in the corresponding base datasets.Table 3 reports factor-group pass rates and a final conjunction verdict for each reference–query pair.Source: §C.3 · p. 21, Table 3 · p. 22

D STATISTICAL TEST 2: FULL-DIMENSIONAL JOINT DISTRIBUTION SUPPORT

  • The full-dimensional KDE test evaluates concatenated state–action vectors, so familiar individual factors can still form a combination rare under the reference distribution.It preserves action–state coordinate correlations rather than marginalizing coordinates or reducing across pairs.Source: D STATISTICAL TEST 2: FULL-DIMENSIONAL… · pp. 23–24

E INTERVENTIONAL ACTION SET

  • Interventions are chosen within empirical action support and spread across task-relevant directions, with Insert-Peg actions selected by their realized hand-endpoint effects.For Insert-Peg, candidates are filtered by a 20-nearest-neighbor threshold before selecting a nine-action grid.Source: E INTERVENTIONAL ACTION SET · pp. 24–25

F MODEL EXPLANATION

  • The model suite spans jointly learned and frozen-encoder JEPAs as well as direct action-conditioned generation and latent video diffusion.This diversity tests whether goal-induced combinatorial bias is tied to a particular representation or generative objective.Source: F MODEL EXPLANATION · p. 26, §F.1 · pp. 26–27

F.2 PREDICTOR-ONLY MODELS FROM TRUE STATES

  • Oracle-state predictors replace visual observations with task-relevant physical variables to isolate dynamics prediction from visual representation learning.The controlled families retain AdaLN conditioning, action duplication, or DriftWorld-style generative prediction.Source: §F.2 · pp. 27–28

H DETAILED PLANNING AND EVALUATION CONFIGURATION

H.1 PLANNING EVALUATION CRITERIA

Planning success is defined by environment-specific distance or pose thresholds against dataset-future goals. The evaluation uses a shared receding-horizon CEM setup to compare goal-conditioned planning across models.

H.1 PLANNING EVALUATION CRITERIA

  • Episodes use dataset-future goals, and success requires meeting the environment-specific criterion at least once within 50 steps.The criteria use different task variables: agent position in TwoRoom, agent and block pose in PushT, end-effector and cube position in Cube, and hand and peg poses in Insert Peg.Source: §H.1 · pp. 29–30, Table 8 · p. 30

H.2 CONFIGURATION

  • CEM evaluates 300 candidate sequences for 30 iterations, plans over 25 environment steps, and executes each selected sequence before replanning.Initial states and goals come from the same dataset trajectory with a 25-step temporal offset; episodes last at most 50 steps.Source: §H.2 · p. 30

I.1 FULL RESULTS

The full-results and discussion sections report prediction and planning evaluations, interpret their relationship cautiously, and outline conditional conditions for compositional generalization. They also describe the AnyGoal library’s role in reproducible evaluation.

I.1 FULL RESULTS

  • The full results include absolute MPC success rates and base-relative prediction errors, while Avoid-T is omitted from the 25-step prediction comparison because its episodes are too short.The appendix tables cover one-step and five-step base-relative prediction error and absolute base-task rollout MSE.Source: §I.1 · pp. 30–31, Table 10 · p. 31, Table 11 · p. 31

I.2 FACTOR SUPPORT VERSUS OBSERVATION SUPPORT

  • Factor-wise support does not ensure observation-level support, because rendering unseen factor combinations can create observations outside the training observation support.Oracle-state predictors are analyzed separately to isolate dynamics generalization from encoder inference.Source: §I.2 · pp. 31–32

I.3 ADDITIONAL OBSERVATIONS

  • Prediction robustness and planning utility do not align consistently: stable rollout error can accompany poor MPC success, and greater degradation can coexist with competitive control.Absolute prediction errors are not directly comparable across model families because they are evaluated in different prediction spaces.Source: §I.3 · p. 32
  • Oracle-state access does not guarantee stronger planning; in Cube, several end-to-end models outperform oracle-state counterparts on the base task and CubeDiagonal.The comparison does not isolate representation quality because the end-to-end and oracle predictors differ in more than their inputs.Source: §I.3 · p. 32

J TOWARD COMPOSITIONAL GENERALIZATION IN WORLD MODELS

  • The paper frames reusable dynamics as requiring both identification of mechanisms on training support and extrapolation to novel compositions.The discussion treats this as a conditional question, not a claim that reusable mechanisms are identifiable from expert rollouts alone.Source: J TOWARD COMPOSITIONAL GENERALIZATION… · p. 33

J.1 IDENTIFYING REUSABLE MECHANISMS ON THE TRAINING SUPPORT

  • Goal shifts alter behavior distributions, while the discussion assumes a Markov oracle state and transition laws invariant across training and test goals.It also notes that regime-dependent interactions can make action effects depend on state rather than decompose into independent state-only and action-only terms.Source: §J.1 · p. 33

J.2 EXTRAPOLATING IDENTIFIED MECHANISMS TO NOVEL COMPOSITIONS

  • Sparse local mechanisms can support generalization to globally novel compositions when relevant local inputs remain training-supported and mechanism estimation error vanishes.This is a sufficient condition, not a complete solution: identifying mechanisms, discovering regimes, and handling genuinely new local interactions remain open challenges.Source: §J.2 · pp. 34–35

K ANYGOAL LIBRARY AND REPRODUCIBILITY

  • AnyGoal provides an open-source library for collecting goal-shifted datasets, training models, and evaluating rollout prediction and MPC planning under shared protocols.Its standardized interfaces and released episodes support reproducible evaluation and construction of additional datasets.Source: K ANYGOAL LIBRARY AND REPRODUCIBILITY · p. 35

K.1 INSTALLATION AND DATASET ACCESS

The library supports dataset access and collection, standardized rollout and planning evaluation, and reproducibility through shared interfaces and released episodes.

K.1 INSTALLATION AND DATASET ACCESS

  • The default installation supports dataset access, collection, validation, and simulation, while optional extras add evaluation and training dependencies.The CLI exposes supported dataset and model registries, and artifacts can be downloaded through the library interface.Source: §K.1 · p. 36

K.2 WORLD-MODEL EVALUATION

  • The same model checkpoint can be evaluated on its base task and goal-shifted variants for both rollout prediction and downstream MPC planning.Evaluation commands support explicit random seeds for repeated comparisons.Source: §K.2 · p. 37

K.3 REPRODUCIBILITY AND EXTENSIBILITY

  • The released benchmark contains 400 episodes per goal-shifted variant, and collection tools allow additional evaluation datasets to be built.Shared evaluation interfaces let future models use the same controlled goal shifts without recreating the environment and evaluation infrastructure.Source: §K.3 · p. 37

Related papers

Loading 2610.26844v1…