Reward Prediction with Factorized World States

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 LLMs. 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

• The paper presents a novel RewardPrediction benchmark and a structured StateFactory method to predict rewards in partially observed, language-driven domains.
• It factorizes unstructured observations into hierarchical, object–attribute states and interprets goals dynamically for enhanced reward alignment.
• StateFactory achieves lower EPIC errors and improved cross-domain performance, significantly boosting agent planning in complex settings.

Reward Prediction with Factorized World States: A Technical Assessment

Problem Statement and Motivation

Reward modeling for RL agents in complex, language-driven and partially observed domains is commonly hindered by the lack of robust, generalizable reward signals. Traditional reward models trained via supervised data or human feedback often overfit to domain-specific patterns, resulting in poor transfer and limited zero-shot capacity. The necessity for dense, step-wise, semantically meaningful progress signals is further aggravated by the absence of benchmarks explicitly targeting this problem. "Reward Prediction with Factorized World States" (2603.09400) addresses this gap by introducing (i) the RewardPrediction benchmark for fine-grained scalar reward evaluation across diverse language-driven domains, and (ii) StateFactory, a representation-based, zero-shot reward prediction method leveraging structured, hierarchical world state factorization.

Figure 1: Overview of the RewardPrediction benchmark. Given a textual goal, step-wise rewards are predicted from action-observation sequences and evaluated with EPIC distance.

RewardPrediction Benchmark

The benchmark operationalizes reward prediction as a function $$f_\text{pred}(g, o_{[0:t]}, a_{[0:t]}) \to \hat{r}_t$$, requiring models to estimate the alignment between current state and goal at each trajectory step. Critical elements of the benchmark include:

• Five Domains: AlfWorld (household manipulation), ScienceWorld (science experiments), TextWorld (text-based logic games), WebShop (e-commerce navigation), and BlocksWorld (symbolic spatial reasoning).
• Data Structure: 2,454 unique action-observation trajectories with step-wise, ground-truth rewards, constructed to decouple exploration length from reward signals.
• Evaluation Metric: EPIC distance, a scalar metric quantifying policy-invariant reward alignment, which is robust to reward shaping.

  Figure 3: Reward prediction paradigms. Top: representation-free direct regression. Bottom: StateFactory-style structured state–goal alignment.

Methods: Representation-Free vs. Representation-Based

Supervised Baselines: LLMs are trained via preference ranking on trajectory pairs. Despite domain-matched or pooled training, these models display pronounced overfitting and poor cross-domain generalization, demonstrating a 138% increase in EPIC error when tested out-of-distribution.

LLM-as-a-Judge: Zero-shot evaluation via large, prompted LLMs (e.g., Qwen3-14B, GPT-OSS-20B), with or without chain-of-thought, is feasible but yields reward predictions that are discrete, noisy, and liable to degenerate into status-quo or checklist-like patterns in open-ended domains.

StateFactory (Representation-Based): The proposed approach factorizes unstructured observations into structured, hierarchical object–attribute world states and dynamically interprets abstract goals in a state-aware manner. At each step, rewards are computed as an explicit function of the semantic alignment between the extracted world state $\hat{s}_t$ and the interpreted goal state $\hat{g}_t$:

Figure 2: StateFactory framework with coupled, recurrent state and goal interpretation. Fine-grained factorization enables semantic similarity-based rewards.

Factorization Pipeline:

1. World State Extraction: LLM-based, recurrent mapping from observations to a set of entities, each with named attributes and current values; temporal consistency is enforced across the episode.
2. Goal Interpretation: Online, context-aware goal parsing, allowing progressive anchoring of abstract instructions to concrete physical entities as they are revealed by the environment.
3. Hierarchical Routing and Reward Computation:
   • For each goal object: find the corresponding object in current state by maximizing identity and attribute similarity.
   • Attribute satisfaction is computed through semantic similarity of corresponding values.
   • Aggregate across all goal entities for the dense, global reward at each step.

Experimental Results and Analysis

Main Results: StateFactory achieves 0.297 EPIC overall, surpassing both supervised (domain-specific LLM reward models) and LLM-as-a-Judge baselines (best: 0.322) in zero-shot settings. In cross-domain generalization, StateFactory demonstrates robust reward transfer without any parameter update.

Agent Performance: Augmenting baseline ReAct agents with StateFactory rewards yields substantial absolute success-rate improvements (+21.64% AlfWorld, +12.40% ScienceWorld). In system-2 planning regimes (i.e., MCTS over LLM-simulated world models), continuous reward gradients generated by StateFactory resolve ambiguous long-horizon goals, supporting effective search and action sequencing.

Ablation Studies: Detailed ablations confirm:

• Representation Granularity: Transitioning from flat observations to factorized object–attribute states yields monotonically lower EPIC (0.57→0.30), confirming that semantic disentanglement and denoising are both necessary and complementary.
• Goal Modality: Dynamic online goal interpretation closely approximates performance of oracle, offline goal states (ΔEPIC ≈ 0.02).
• LLM Reasoning and Scale: Improvements (ΔEPIC = 0.06–0.15) with larger, reasoning-augmented backbones.
• Semantic Embeddings: Triplet-based discriminability of the downstream embedding model is tightly coupled to reward quality.

  Figure 4: StateFactory ablation. Representation granularity and embedding discriminability are key factors in reward prediction accuracy.

Theoretical and Practical Implications

The explicit factorization of world states introduces a modular, interpretable scaffolding for reward prediction, robust to spurious correlations and irrelevant observation noise. Unlike supervised approaches that entangle reward logic with domain-specific data distributions, StateFactory supports direct transfer and systematic abstraction across domains.

Practically, this enables scalable deployment of planners in complex, partially observed and instruction-driven environments, where true environmental rewards are unavailable or are only available in sparse form. The framework is compatible with arbitrary text-embedding models and scales with future improvements in LLM reasoning and representation.

Theoretically, the results emphasize the importance of structured world modeling in language-based agents, offer a compelling argument for decoupled state–goal semantics in reward design, and suggest a promising direction for hierarchical, interpretable RL in text and mixed-modality environments.

Conclusion

"Reward Prediction with Factorized World States" (2603.09400) establishes both a rigorous benchmark and an effective, zero-shot reward modeling approach based on explicit world state factorization. StateFactory demonstrates robust, domain-general reward prediction and directly improves language-agent planning performance, highlighting the value of structured representations and dynamic goal grounding for language-based autonomous agents. Future research may extend this paradigm to multi-modal, open-world settings or further integrate structural priors into unified agent architectures.