Better World Models Can Lead to Better Post-Training Performance

• URL: http://arxiv.org/abs/2512.03400v1
• Date: Wed, 03 Dec 2025 03:13:20 GMT
• Title: Better World Models Can Lead to Better Post-Training Performance
• Abstract summary: We study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers.<n>We compare standard next-token prediction to two explicit world-modeling strategies.<n>We find that explicit world-modeling yields more linearly decodable and causally steerable state representations.
• Score: 9.713688760042544
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: In this work we study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers across different training stages. We use a controlled 2x2x2 Rubik's Cube and ask: (1) how does explicitly pretraining a world model affect the model's latent representations, and (2) how does world-model quality affect the model's performance after reinforcement learning post-training? We compare standard next-token prediction to two explicit world-modeling strategies -- (i) state-prediction pretraining and (ii) a joint state-prediction + next-token objective -- and assess task performance after Group Relative Policy Optimization (GRPO) is applied as post-training. We evaluate the representation quality with linear probes and causal interventions. We find that explicit world-modeling yields more linearly decodable and causally steerable state representations. More importantly, we find that improved state representations lead to higher gains for GRPO, especially on harder cube states. Our results indicate that sharpening state representations can improve the effectiveness of post-training for sequence-planning tasks.

Related papers

• WorldAgen: Unified State-Action Prediction with Test-Time World Model Training [19.89058127984308]
  We present WorldAgen, a unified framework that jointly learns world modeling and action prediction.<n>WorldAgen samples exploratory actions, collects ground-truth state transitions, and performs lightweight TTT updates to refine its world model.<n>Our method surpasses existing state-of-the-art models, highlighting the effectiveness of adapting world models at inference time.
  arXiv  Detail & Related papers  (2026-09-08T02:46:13Z)
• Orbis 2: A Hierarchical World Model for Driving [21.534159792492186]
  We present a hierarchical driving world model that factorizes future prediction across two levels operating at distinct temporal and abstraction scales.<n>This decomposition yields high perceptual fidelity while also capturing strong spatial and semantic representations.<n>We introduce a generic two-stage training paradigm that pretrains the model with diffusion forcing and fine-tunes with teacher forcing.
  arXiv  Detail & Related papers  (2026-07-17T12:12:33Z)
• Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL [0.8546394944138487]
  We formalize text-based world modeling as a steerable transition-dynamics problem decomposed into initial state.<n>We curate 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families.<n>We introduce a plug-and-play GRPO training framework with deterministic state checks, and perform zero-shot transfer ablations on three OOD environments.
  arXiv  Detail & Related papers  (2026-05-07T00:40:32Z)
• FRAPPE: Infusing World Modeling into Generalist Policies via Multiple Future Representation Alignment [45.21358158991569]
  VLA models can predict environmental dynamics, known as world modeling.<n>Current approaches face two main issues: 1) The training objective forces models to over-emphasize pixel-level reconstruction, which constrains semantic learning and generalization, and 2) Reliance on predicted future observations during inference often leads to error accumulation.<n>We introduce Future Representation Alignment via Parallel Progressive Expansion (FRAPPE) to address these challenges.
  arXiv  Detail & Related papers  (2026-02-19T11:00:46Z)
• VAGEN: Reinforcing World Model Reasoning for Multi-Turn VLM Agents [130.70999337445468]
  Key challenge in training Vision-Language Model (VLM) agents, compared to Language Model (LLM) agents, is shift from textual states to complex visual observations.<n>We ask: Can VLM agents construct internal world models through explicit visual state reasoning?<n>We architecturally enforce and reward the agent's reasoning process via reinforcement learning (RL)<n>We find that the agent's reasoning into State Estimation and Transition Modeling is critical for success.
  arXiv  Detail & Related papers  (2025-10-19T16:05:07Z)
• Internalizing World Models via Self-Play Finetuning for Agentic RL [65.96875390986655]
  Large Language Models (LLMs) as agents often struggle in out-of-distribution (OOD) scenarios.<n>We show how to encode this world model by decomposing it into two components: state representation and transition modeling.<n>We introduce SPA, a simple reinforcement learning framework that cold-starts the policy via a Self-Play supervised finetuning stage to learn the world model.
  arXiv  Detail & Related papers  (2025-10-16T18:03:39Z)
• Enhancing Training Data Attribution with Representational Optimization [57.61977909113113]
  Training data attribution methods aim to measure how training data impacts a model's predictions.<n>We propose AirRep, a representation-based approach that closes this gap by learning task-specific and model-aligned representations explicitly for TDA.<n>AirRep introduces two key innovations: a trainable encoder tuned for attribution quality, and an attention-based pooling mechanism that enables accurate estimation of group-wise influence.
  arXiv  Detail & Related papers  (2025-05-24T05:17:53Z)
• ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning [90.41852663775086]
  ACT-JEPA is a novel architecture that integrates imitation learning and self-supervised learning.<n>We train a policy to predict action sequences and abstract observation sequences.<n>Our experiments show that ACT-JEPA improves the quality of representations by learning temporal environment dynamics.
  arXiv  Detail & Related papers  (2025-01-24T16:41:41Z)
• Learning Goal-Conditioned Representations for Language Reward Models [10.94845204766088]
  We propose training reward models (RMs) in a contrastive, $textitgoal-conditioned$ fashion. We show this way of training RM representations enables improved $textitsteerability$ because it allows us to evaluate the likelihood of an action achieving a particular goal-state. We additionally find that these representations can perform fine-grained control by conditioning on desired future goal-states.
  arXiv  Detail & Related papers  (2024-07-18T20:23:11Z)
• Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models [19.05224410249602]
  We propose a novel approach for offline reinforcement learning with closed-loop policy evaluation and world-model adaptation. We analyzed the performance of the proposed method and provided an upper bound on the return gap between our method and the real environment under an optimal policy.
  arXiv  Detail & Related papers  (2024-05-30T09:34:31Z)
• ReCoRe: Regularized Contrastive Representation Learning of World Model [21.29132219042405]
  We present a world model that learns invariant features using contrastive unsupervised learning and an intervention-invariant regularizer. Our method outperforms current state-of-the-art model-based and model-free RL methods and significantly improves on out-of-distribution point navigation tasks evaluated on the iGibson benchmark.
  arXiv  Detail & Related papers  (2023-12-14T15:53:07Z)
• Bi-Drop: Enhancing Fine-tuning Generalization via Synchronous sub-net Estimation and Optimization [58.90989478049686]
  Bi-Drop is a fine-tuning strategy that selectively updates model parameters using gradients from various sub-nets. Experiments on the GLUE benchmark demonstrate that Bi-Drop consistently outperforms previous fine-tuning methods.
  arXiv  Detail & Related papers  (2023-05-24T06:09:26Z)
• Large Language Models with Controllable Working Memory [64.71038763708161]
  Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP) What further sets these models apart is the massive amounts of world knowledge they internalize during pretraining. How the model's world knowledge interacts with the factual information presented in the context remains under explored.
  arXiv  Detail & Related papers  (2022-11-09T18:58:29Z)
• Value-Consistent Representation Learning for Data-Efficient Reinforcement Learning [105.70602423944148]
  We propose a novel method, called value-consistent representation learning (VCR), to learn representations that are directly related to decision-making. Instead of aligning this imagined state with a real state returned by the environment, VCR applies a $Q$-value head on both states and obtains two distributions of action values. It has been demonstrated that our methods achieve new state-of-the-art performance for search-free RL algorithms.
  arXiv  Detail & Related papers  (2022-06-25T03:02:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.