Reward Prediction with Factorized World States

• URL: http://arxiv.org/abs/2603.09400v1
• Date: Tue, 10 Mar 2026 09:12:20 GMT
• Title: Reward Prediction with Factorized World States
• Abstract summary: Supervised learning of reward models could introduce biases inherent to training data, limiting generalization to novel goals and environments.<n>We introduce StateFactory, a factorized representation method that transforms unstructured observations into a hierarchical object-attribute structure using language models.<n>Overall, the compact representation structure induced by StateFactory enables strong reward generalization capabilities.
• Score: 26.23676250908859
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Agents must infer action outcomes and select actions that maximize a reward signal indicating how close the goal is to being reached. Supervised learning of reward models could introduce biases inherent to training data, limiting generalization to novel goals and environments. In this paper, we investigate whether well-defined world state representations alone can enable accurate reward prediction across domains. To address this, we introduce StateFactory, a factorized representation method that transforms unstructured observations into a hierarchical object-attribute structure using language models. This structured representation allows rewards to be estimated naturally as the semantic similarity between the current state and the goal state under hierarchical constraint. Overall, the compact representation structure induced by StateFactory enables strong reward generalization capabilities. We evaluate on RewardPrediction, a new benchmark dataset spanning five diverse domains and comprising 2,454 unique action-observation trajectories with step-wise ground-truth rewards. Our method shows promising zero-shot results against both VLWM-critic and LLM-as-a-Judge reward models, achieving 60% and 8% lower EPIC distance, respectively. Furthermore, this superior reward quality successfully translates into improved agent planning performance, yielding success rate gains of +21.64% on AlfWorld and +12.40% on ScienceWorld over reactive system-1 policies and enhancing system-2 agent planning. Project Page: https://statefactory.github.io

Related papers

• ConfAL-WM: Confidence-Guided Active Learning for Action-Conditioned World Models [33.57978728341114]
  This paper presents ConfAL-WM, a confidence-guided active learning framework for post-training embodied world models.<n>Built upon EVAC, we attach a lightweight confidence to probe UNet decoder features and predict dense confidence maps in the latent space.<n>These maps are aggregated into task-, frame-, patch-level scores, enabling both efficient data selection and localized training enhancement.
  arXiv  Detail & Related papers  (2026-08-26T09:29:08Z)
• World Action Verifier: Self-Improving World Models via Forward-Inverse Asymmetry [82.93104394404781]
  General-purpose world models promise scalable policy evaluation, optimization, and planning.<n>We propose World Action Verifier (WAV), a framework that enables world models to identify their own prediction errors and self-improve.
  arXiv  Detail & Related papers  (2026-04-02T12:48:36Z)
• Self-Correcting VLA: Online Action Refinement via Sparse World Imagination [55.982504915794514]
  We propose Self-Correcting VLA (SC-VLA), which achieve self-improvement by intrinsically guiding action refinement through sparse imagination.<n>SC-VLA achieve state-of-the-art performance, yielding the highest task throughput with 16% fewer steps and a 9% higher success rate than the best-performing baselines.
  arXiv  Detail & Related papers  (2026-02-25T06:58:06Z)
• SRPO: Self-Referential Policy Optimization for Vision-Language-Action Models [42.89413870143421]
  Vision-Language-Action (VLA) models excel in robotic manipulation but are constrained by their heavy reliance on expert demonstrations.<n>Current VLA-RL methods, including group-based optimization approaches, are crippled by severe reward sparsity.<n>We propose Self-Referential Policy Optimization (SRPO), a novel VLA-RL framework.
  arXiv  Detail & Related papers  (2025-11-19T16:52:23Z)
• 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)
• What Fundamental Structure in Reward Functions Enables Efficient Sparse-Reward Learning? [6.908972852063454]
  Policy-Aware Matrix Completion (PAMC) is a first concrete step toward a structural reward learning framework.<n>Our results highlight PAMC as a practical and principled tool when structural rewards exist, and as a concrete first instantiation of a broader structural reward learning perspective.
  arXiv  Detail & Related papers  (2025-09-04T00:53:02Z)
• 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)
• VARP: Reinforcement Learning from Vision-Language Model Feedback with Agent Regularized Preferences [13.337649128532307]
  Preference-based RL mitigates some of these pitfalls by learning rewards from comparative feedback.<n>A single final-state image generally fails to capture the agent's full motion.<n>We present a two-part solution that both improves feedback accuracy and better aligns reward learning with the agent's policy.
  arXiv  Detail & Related papers  (2025-03-18T01:51:27Z)
• Uncovering Untapped Potential in Sample-Efficient World Model Agents [51.65485693709418]
  Simulus is a highly modular TBWM agent that integrates a multi-modality tokenization framework, intrinsic motivation, prioritized WM replay, and regression-as-classification.<n>Simulus achieves state-of-the-art sample efficiency for planning-free WMs across three diverse benchmarks.
  arXiv  Detail & Related papers  (2025-02-17T08:06:10Z)
• PcLast: Discovering Plannable Continuous Latent States [24.78767380808056]
  We learn a representation that associates reachable states together for effective planning and goal-conditioned policy learning. Our proposals are rigorously tested in various simulation testbeds.
  arXiv  Detail & Related papers  (2023-11-06T21:16:37Z)
• When Demonstrations Meet Generative World Models: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning [62.00672284480755]
  This paper aims to recover the structure of rewards and environment dynamics that underlie observed actions in a fixed, finite set of demonstrations from an expert agent. Accurate models of expertise in executing a task has applications in safety-sensitive applications such as clinical decision making and autonomous driving.
  arXiv  Detail & Related papers  (2023-02-15T04:14:20Z)

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.