PaperlayerView exact revision on arXiv

Source-linked AI summary

How do World Models and Policies Compose in LLM Agents? A Joint Spectral and Behavioral Account

Ruize Xu, Xiao Yu, Yujin Tang, Chenming Shang, Nikhil Singh

arXiv:2608.30067v1cs.LGcs.AIcs.CL

TL;DR

The paper asks how LLM agents can combine environmental knowledge with task-directed ability, and whether policy training preserves world knowledge. Through controlled world-model and policy training experiments, it analyzes additive updates spectrally and behaviorally, then tests merging and online auxiliary training. The results show complementary update geometry, greater robustness and state coverage for sequential training, and improved preservation from both tested interventions.

  • Problem

    The paper studies the relationship between world knowledge and task-directed ability in LLM agents, including whether policy training preserves world knowledge sufficiently.

  • Method

    The authors combine controlled world-model and policy training with additive parameter-update dissection, spectral interventions, behavioral measurements, merging, and an online world-model loss.

  • Results

    World-model and policy updates share input-feature directions while writing to largely distinct output directions; sequential training also yields redundant input pathways and expanded state coverage, while both preservation interventions improve over baseline.

  • Takeaways & Limitations

    World knowledge and task-directed ability can be learned in geometrically complementary forms, making their interface a relevant post-training design consideration.

  • Takeaways & Limitations

    The empirical analysis is confined to two text-only agent benchmarks and two Qwen2.5/Qwen3 models, and its preservation experiments cover only a small, lightly tuned part of the design space.

Abstract

from arXiv · show

How do LLM agents come to both understand environments they act in and master tasks set within them? Through controlled experiments combining world-model training (next-state prediction) and policy training (reward maximization), we investigate this question. We dissect the resulting models through their additive parameter updates. Geometrically, we find effective world-model updates are low-rank and share an input-feature subspace with policy updates while writing to nearly orthogonal output directions, whether trained separately or sequentially. However, we find that, in projection interventions, the sequential update induces more robustness than separate policy RL when removing the world model's leading input directions, suggesting that it has learned alternative input pathways. Behaviorally, we find the sequentially trained agent explores a wider range of states and actions. Based on this, we ask: does policy training preserve world knowledge as well as it could? We probe this with training-free merging built on the geometrically motivated input basis plus an online world-model loss during policy RL, and show both improve over the untreated baseline. Our findings suggest world knowledge and task-directed ability can be learned in geometrically complementary forms, and that future post-training pipelines should consider how best to engineer the interface between them.

1 Introduction

The paper investigates how world-model and policy training interact in LLM agents, combining behavioral evaluation with spectral analysis of parameter updates. It finds complementary update geometry and tests merging and auxiliary world-model training as preservation strategies.

  • The work positions stage composition as a design issue because structurally different contributions may need to be preserved, combined, and scheduled.Prior work documented pipeline benefits but rarely examined their mechanism directly.
  • The authors dissect additive parameter updates across controlled post-training checkpoints using spectral interventions and behavioral measurements.The analyses include subspace projection, rank truncation, magnitude pruning, next-state prediction entropy, state coverage, and success patterns.
  • The two-stage agent preserves world-model calibration and expands state coverage relative to the policy-only version.
  • The study asks whether world-model training has geometric and behavioral signatures and how its parameter updates relate to policy training.
  • World-model and policy updates share input-feature directions but write to nearly orthogonal output directions, including under sequential training.Sequential training also gives the policy update redundant input pathways unavailable to base-initialized policy training.
  • Training-free KnOTS/TIES merging and an online world-model SFT auxiliary loss during policy learning both improve over the untreated baseline.

2 Preliminaries

The preliminaries define the text-based agent setting, world-model and policy-training stages, benchmarks, and replicated evaluation setup. World-model training predicts next observations, while policy RL optimizes task-success reward.

  • 2.1 Problem Setup: The agent operates as a ReAct-style policy in a Markov Decision Process, mapping observation history and prior actions to actions and receiving task-success rewards.
  • 2.1 Problem Setup: A world model is a learned surrogate for environment transitions that predicts the next state from recent observations and actions.
  • 2.2 Two-stage Training: The two-stage pipeline first grounds the policy in environment dynamics and then uses GRPO to optimize end-task reward.
  • 2.2 Two-stage Training: World-model SFT trains on transition triplets but inherits token-level supervision, data-quality sensitivity, and model-collapse or catastrophic-forgetting risks.
  • 2.2 Two-stage Training: World-model RL uses embedding-space proximity between predicted and realized next states as a reward for semantic next-state alignment.A frozen embedding model and threshold define the similarity-based reward, and GRPO optimizes the policy on its own rollouts.
  • 2.3 Benchmarks and Models: Experiments use ALFWorld as the primary text-based embodied benchmark and τ 2-Bench as a secondary interleaved tool-use environment.The study evaluates Qwen2.5-7B-Instruct on ALFWorld and Qwen3-8B on τ 2-Bench, including base, world-model-only, policy-only, and sequential checkpoints.

3 Methodology

The paper dissects additive parameter updates from world-model and policy training, then applies spectral projections, rank truncation, pruning, and behavioral measurements to compare their structure and effects.

  • ∆-decomposition: Additive updates ∆X are defined relative to the shared base model, with ∆∆PRL isolating the policy contribution after WMRL.The decomposition compares WMSFT, WMRL, base-trained PRL, and sequential WMRL → PRL checkpoints.
  • Spectral interventions: Singular-value decomposition provides effective-rank and truncated-rank measures of how broadly each update spreads across directions.Effective rank uses normalized singular-value mass, while the Eckart–Young approximation retains the top k components.
  • Spectral interventions: Subspace projections retain either the components parallel or orthogonal to another update’s leading singular-vector subspace, with random orthonormal projections as controls.Projected checkpoints are reconstructed and re-evaluated end-to-end.
  • Subspace overlap: Module-wise principal-angle overlap measures alignment between corresponding leading left or right subspaces, using a random baseline of approximately k/d.The analysis distinguishes output-direction alignment from input-feature alignment.
  • Additional analyses: Magnitude pruning retains only the top-P% update entries by absolute value, complementing singular-vector analysis with coordinate-space analysis.Behavioral readouts include held-out next-state-token entropy and downstream success-set overlap.

4 What makes a good world model? Spectral signatures of effective two-stage training (RQ1)

The study finds that effective WMRL world-model updates are low-rank and functionally concentrated, while WMSFT’s similarly low rank does not predict success across training regimes.

  • Effective-rank signatures: Within the WMRL regime, stronger world models combine lower effective rank with higher performance across checkpoints.Figure 1 varies reward threshold, judge type, data source, and prediction format.
  • Effective-rank signatures: WMSFT reaches the lowest effective rank yet performs catastrophically, and WMSFT → PRL trails WMRL → PRL by more than 20 AVG points.The authors therefore treat SFT and RL world-model updates as spectrally incomparable regimes.
  • Sequential training: WMRL initialization produces substantially higher per-layer effective rank in the downstream policy update than policy RL from the base model.The largest gap appears in early-to-mid layers, consistent with broader directions explored after WMRL initialization.
  • Functional concentration: Rank-1 truncation of ∆WMRL recovers within approximately 10 points of the full sequential pipeline while analogous policy-update truncations collapse performance toward WMRL-only levels.The world-model stage’s downstream value is concentrated in a few directions, whereas policy value is more diffuse.

5 Weight-space Relationship between ∆WMRL and the Policy Stage (RQ2)

The world-model and policy stages are complementary across parameter subspaces: they share leading input-feature directions but occupy nearly distinct output directions, with sequential training adding input redundancy.

  • Experimental comparison: The study tests shared and orthogonal accounts through left/right projections, rank-controlled interventions, and module-wise principal-angle overlap.Projected sequential checkpoints retain WMRL plus modified policy-stage updates before end-to-end evaluation.
  • Left-space output directions: Removing WMRL’s leading output directions from ∆∆PRL has essentially no cost, indicating that policy-relevant signal lies largely outside WMRL’s principal output subspace.At k=500, the orthogonal variant reaches ID 91.7 and OOD 84.9, at or above the un-projected WMRL → PRL baseline.
  • Left-space output directions: Keeping only WMRL-aligned output directions reaches ID 60.4 / OOD 56.8 at k=500, only modestly above the matched random projection at ID 50.3 / OOD 49.5.This suggests WMRL’s populated output directions contain only marginal policy-relevant information.
  • Left-space output directions: The same output-direction orthogonality appears for independently trained ∆PRL, whose orthogonal projection preserves performance while parallel projection performs poorly.This suggests U-disjointness is intrinsic to policy training rather than caused only by sequential conditioning.
  • Right-space input features: Right-space parallel projection of ∆∆PRL reaches ID 83.6 / OOD 79.7 at k=500, close to the WMRL → PRL baseline of ID 87.0 / OOD 81.8.The corresponding left-space projection reaches only ID 60.4 / OOD 56.8, while random control remains near the WMRL baseline.
  • Right-space input features: Removing WMRL’s leading input directions hurts ∆PRL but leaves ∆∆PRL comparatively robust, indicating redundant input pathways in sequential training.For ∆∆PRL, performance changes only mildly from ID 91.4 → 89.6 and OOD 82.8 → 77.6 across k=1 to k=500.
  • Cross-benchmark replication: Module-level overlap confirms near-zero rank-1 left overlap and strong right overlap, with mean values of 0.04 and 0.63 respectively.The pattern replicates on τ 2-Bench, although behavioral recovery is weaker there: 71.8% on ALFWorld versus 46.9% on τ 2-Bench.

6 How does WMRL help in practice? A rollout-time behavioral account (RQ1)

WMRL changes downstream behavior beyond simply refining policy RL: the sequential agent combines distinct task successes, broader state exploration, policy efficiency, and better world-model calibration.

  • Success-set overlap: The composite WMRL→PRL model has non-nested success sets, solving tasks that each WMRL and PRL baseline misses.This rules out interpreting the composite as merely a refined version of PRL.
  • Exploration statistics: WMRL preserves broader state exploration, while PRL contracts action and state coverage in exchange for shorter trajectories.The rollout analysis counts unique valid actions, unique states, and total turns, with runs capped at 30 turns.
  • Exploration statistics: The composite inherits PRL’s similar turn count while recovering some of WMRL’s state-coverage advantage.This combines rollout efficiency with a broader range of visited states.
  • Next-state prediction entropy: Held-out next-state entropy shows WMRL improves world-model calibration, whereas PRL-from-base degrades it.Lower entropy indicates better calibration on the ground-truth next-state tokens.
  • Summary: Across task, trajectory, and representational readouts, WMRL→PRL behaves as a combination of qualitatively different stages.The authors relate these behavioral correlates to updates sharing input subspaces while writing to disjoint output subspaces.

7 Preserving the World Model (RQ3)

The paper tests whether preserving world-model structure during policy training improves performance, using geometry-guided merging and an online world-model loss.

  • Motivation: The preservation interventions are motivated by asymmetric geometry: world-model updates provide a low-rank scaffold while policy updates use output-orthogonal directions.The tests include one training-free and one training-time intervention.
  • Training-free merging: KnOTS/TIES merges world-model and policy checkpoints using either the default U-sharing basis or a V-sharing variant.The V-sharing variant tests the predicted shared input-feature basis.
  • Training-based preservation: Online WM-SFT adds a next-state SFT loss from 32 realized rollout transitions to the policy-gradient objective.This provides a training-time world-model signal during PRL.
  • Results: +2.0 points (+4.2 OOD) separates KnOTS-V from KnOTS-U on the (MWMRL, MPRL) pair.The result matches the prediction that factoring out the shared V basis preserves signal lost by the U-basis variant.
  • Results: Both merge variants remain ∼2 points above the untreated baseline on ALFWorld, while online WMSFT reaches ID 90.4 but OOD 79.5.The interventions therefore provide complementary preservation benefits.

8 Conclusion

The paper finds complementary geometric and behavioral forms of world knowledge and policy ability, then identifies important limits on the scope and causal interpretation of those findings.

  • Conclusion: World-model and policy training share input-feature directions but write to largely distinct output directions.Sequential training also makes the resulting policy less dependent on the world model’s leading input directions.
  • Conclusion: These findings provide an initial pathway for investigating and engineering how world-model and policy stages compose.The paper frames this as both a scientific question about internal composition and an engineering question about combining stages.
  • Limitations: The analysis is confined to two text-only agent benchmarks and two Qwen2.5/Qwen3 models.The authors note that multimodal agents introduce vision-language perception, grounding, and interaction requirements.
  • Limitations: The framework analyzes additive updates from one base model and applies the same module-wise intervention across modules.LoRA-style decompositions and stronger cross-module, layer-local, activation, or training-time causal tests remain future work.
  • Limitations: Preservation experiments cover only KnOTS/TIES merging and online WM-SFT, with limited hyperparameter tuning and incomplete τ 2-Bench coverage.The authors identify scalable preservation methods and broader benchmark/model coverage as next steps.

A.5 Robustness Check: rank-k Pruning Experiments

Pruning experiments support the geometric account: the world-model update is highly compressible, whereas policy updates require more rank or retained magnitude, with the pattern replicated on τ 2-Bench.

  • Rank-k truncation: The experiments truncate ΔWMRL, ΔΔPRL, and ΔPRL separately and evaluate sequential, jointly pruned, and from-base compositions.The four regimes are WkP, WPk, WkPk, and Pk.
  • Rank-k truncation: Rank-1 ΔWMRL truncation yields ID 77.6 / OOD 79.7, versus full WMRL→PRL at ID 87.0 / OOD 81.8.The nearly flat WkP trend indicates that downstream world-model utility is concentrated in few directions.
  • Rank-k truncation: At k=1, truncating ΔΔPRL yields ID 43.8 / OOD 46.4, and ΔPRL is worst at every rank.The results indicate that policy-stage task value is distributed across more singular components, especially for the from-base update.
  • Magnitude pruning: Top-p magnitude pruning reaches best performance with p=0.1 for ΔΔPRL, while ΔPRL continues improving before p=0.5.Bottom-p sharply degrades both updates, and random-p requires substantially more retained mass to recover performance.
  • τ 2-Bench replication: The spectral ordering is preserved across all 36 Qwen3-8B layers on τ 2-Bench, with WMRL above RL-stage checkpoints and WMSFT trailing in deeper layers.This replicates the qualitative Qwen2.5/ALFWorld pattern on a different model and benchmark.

B.3 Cross-model Replication of the Subspace Asymmetry

Cross-model experiments reproduce the geometric U-disjoint/V-shared asymmetry between policy and incremental sequential updates, although behavioral separation is weaker on the harder, lower-headroom benchmark.

  • Geometric replication: Mean rank-1 left overlap is 0.04 on both models, while mean rank-1 right overlap is 0.81 on Qwen3 versus 0.63 on Qwen2.5.The left side corresponds to output directions and the right side to input-feature directions.
  • Geometric replication: The Frobenius cosine is 0.015 on Qwen3 versus 0.028 on Qwen2.5, confirming near-orthogonality of the two updates as matrices.Figure 12 shows most modules near zero, with only mild similarity in a small subset of late attention projections.
  • Behavioral replication: Qwen3 preserves the behavioral asymmetry at k=1: orthogonal projection reaches 40.3% success versus 34.0% for parallel projection, a +6.3-point gap.The sign matches Qwen2.5, where the corresponding gap is +50 points, but the Qwen3 margin is compressed.
  • Behavioral replication: On Qwen3, ΔPRL gains only +5.6 points over base, from 33.7% to 39.3%, limiting the behavioral signal available to projection interventions.The corresponding Qwen2.5 gain is +60.6 points, from 17.7% to 78.3%.
  • Scope: The replication spans different models, benchmarks, and depths, but the experiments cannot establish generality across models and agent types.The authors also report that stronger end-to-end effects require substantial headroom in the policy update.

C Preservation Method Details

The paper preserves world-modeling ability through either training-free checkpoint merging or an online world-model SFT loss during policy learning, with merging adapted to the shared input basis.

  • Preservation strategies: The preservation approaches produce a single deployable policy checkpoint using either checkpoint merging or online world-modeling SFT during PRL.The online loss is added while adapting the policy with PRL.
  • Training-free merging: Training-free merging combines the base model with a WMRL checkpoint and either a separate PRL checkpoint or a sequential WMRL→PRL checkpoint.Task vectors are formed for each target weight matrix before merging.
  • Training-free merging: KnOTS represents module-wise task vectors in a shared low-rank SVD frame before applying TIES masking and aggregation.The default U-sharing variant concatenates task matrices along the input dimension.
  • Training-free merging: The V-sharing variant transposes matrices so the shared basis lies on the input side, then transposes the reconstructed matrices back.This variant is motivated by the stronger right-subspace overlap observed in the geometric analysis.
  • Online WM-SFT: Online WM-SFT reuses on-policy GRPO rollouts as next-observation prediction examples, requiring no additional environment interaction.Each example maps recent observation history and the executed action to the next observation.
  • Replication: On Qwen3 plus τ 2-Bench, V-sharing broadly outperforms U-sharing in training-free KnOTS-TIES merges and attains the best reported AVG@3 within the training-free comparisons.The replication uses top-K=40 and the same merging settings as the ALFWorld experiments.

D WMRL Data Example

The WMRL data example trains next-observation prediction from an observation history and action, illustrated in both embodied and tool-use environments.

  • Data format: Each WMRL example contains a model prompt and the ground-truth next observation used as the supervision target.The displayed examples are drawn from ALFWorld and τ 2-Bench.
  • ALFWorld example: In ALFWorld, the prompt asks the model to predict the immediate next observation after an action in a household-object environment.The example includes recent observations, corresponding actions, and a required next-state output format.
  • τ 2-Bench example: In τ 2-Bench, the example combines a customer-service policy, tool schemas, dialogue history, a potential action, and task-completion status.The ground-truth response asks for another way to locate an order and says the task is not yet completed.
  • Prediction template: The online WMSFT template uses the previous raw observation and extracted action as input, supervising only the next environment observation.The concrete example maps “Nothing happens” plus “go to desk 1” to arrival at desk 1 with visible objects.
Loading 2608.30067v1…