PaperlayerView exact revision on arXiv

Source-linked AI summary

Reward Prediction with Factorized World States

Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung

arXiv:2603.09400v1cs.CL

TL;DR

The paper asks whether structured world-state representations alone can support accurate reward prediction across domains despite limited generalization from supervised reward models. StateFactory factorizes observations into hierarchical object–attribute structures and estimates rewards from constrained semantic similarity between current and goal states. Across the RewardPrediction benchmark, it achieves strong zero-shot reward generalization and improves agent planning performance.

  • Problem

    The paper investigates whether well-defined world-state representations alone can enable accurate reward prediction across domains when supervised reward models may generalize poorly to novel goals and environments.

  • Method

    StateFactory transforms unstructured observations into hierarchical object–attribute structures and estimates rewards using semantic similarity between current and goal states under hierarchical routing.

  • Results

    StateFactory achieves strong zero-shot reward generalization across five domains, reducing EPIC distance by 60% versus VLWM-critic and 8% versus LLM-as-a-Judge.

  • Takeaways & Limitations

    StateFactory's state-driven signals improve system-1 planning success by up to 21.64% and effectively guide structured search for system-2 agent planning.

Abstract

from arXiv · show

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

1 Introduction

The paper addresses zero-shot reward prediction across goals and environments by introducing a structured state representation and a benchmark for fine-grained evaluation. StateFactory factorizes observations into hierarchical object–attribute structures, achieving strong reward generalization and improved planning performance.

  • Motivation: Zero-shot agents need world-state-based reward signals because real-world settings may not permit massive trial-and-error for explicit rewards.Semantic distance between current and desired goal states offers an alternative for estimating task progress.
  • Motivation: World-state reward estimation is bottlenecked by constructing representations whose geometry reflects task progress and by evaluating fine-grained step-wise reward quality.The latter is difficult because existing datasets mainly use sparse, outcome-oriented rewards.
  • Method: StateFactory transforms unstructured observations into hierarchical object–attribute structures and estimates rewards through semantic similarity between current and goal states.Large language models decompose observations into atomic semantic factors, producing compact abstractions for zero-shot reward generalization.
  • Benchmark: RewardPrediction spans five interactive environments and 2,454 unique trajectories with step-wise action–observation pairs and scalar rewards.Predicted rewards can be evaluated against ground truth using EPIC distance.
  • Results: StateFactory reduces EPIC distance by 60% versus VLWM-critic and 8% versus LLM-as-a-Judge in zero-shot reward prediction.The method also improves planning success by +21.64% on AlfWorld and +12.40% on ScienceWorld over reactive system-1 policies while strengthening system-2 planning.

2 The RewardPrediction Benchmark

RewardPrediction is a benchmark for evaluating step-wise reward prediction from textual goals, action-observation trajectories, and ground-truth progress across five text-based environments. It uses EPIC distance to measure alignment between predicted and ground-truth rewards.

  • Benchmark formulation: RewardPrediction evaluates reward predictors on offline trajectories containing goals, observations, actions, and step-wise ground-truth rewards.Predictors estimate a reward sequence from the goal, past observations, and action history.
  • Evaluation: EPIC distance measures the policy-invariant difference between predicted and ground-truth reward functions while preserving fine-grained magnitude information.Average EPIC distance across trajectories is reported as final performance.
  • Benchmark construction: Paired positive-negative instances are constructed for the same task to reduce heuristic reward hacking, using expert trajectories with deterministic progress gradients.The positive anchor follows an expert trajectory whose reward increases as r_t = t/T.
  • Benchmark construction: The benchmark spans AlfWorld, ScienceWorld, TextWorld, WebShop, and BlocksWorld, covering robotics, science, games, navigation, and classical planning.These environments provide diverse physical and semantic task settings.

3 Reward Prediction Methods

The paper compares representation-free and representation-based reward predictors, including supervised and zero-shot regimes. StateFactory builds recurrent, goal-conditioned factorized states and estimates progress through hierarchical semantic alignment between state and goal representations.

  • Method families: Reward prediction methods are categorized by whether they construct explicit world and goal representations and whether they use RewardPrediction training data.This yields representation-free versus representation-based methods and supervised versus zero-shot regimes.
  • StateFactory: StateFactory decomposes observations into hierarchical object-attribute states and iteratively grounds textual goals into dynamic goal states.The framework integrates world-state transition, state extraction, and goal interpretation.
  • Hierarchical routing: StateFactory derives rewards by measuring semantic similarity between the current factorized state and the dynamically interpreted goal state.The similarity evaluation proceeds from fine-grained attribute alignment to global progress aggregation.
  • StateFactory: StateFactory separates object identity from evolving attributes and updates state representations recurrently under goal-conditioned context.This design filters task-irrelevant details while tracking temporal changes.
  • Results: StateFactory establishes the best zero-shot result in the benchmark and approaches supervised baselines.Table 1 reports comparisons across the RewardPrediction benchmark.

4 Experiments

The experiments compare StateFactory with supervised and representation-free reward models, then analyze how its factorization, goal interpretation, language-model backbone, embeddings, and planning integration affect performance.

  • Main Results: Supervised reward models increase prediction error by 138% on unseen domains despite high within-domain precision, indicating weak zero-shot generalization.The reported within-domain AlfWorld precision is 0.212.
  • Main Results: Training-free representation-free methods perform slightly better in checklist-like TextWorld and reverse-engineerable BlocksWorld but break down in open-ended environments.Their implicit reward signals are described as discrete and volatile because of probabilistic generation.
  • Main Results: StateFactory converts reward estimation into continuous semantic distance over hierarchical object-attribute states, addressing reward generalization and signal continuity.The representation reduces noise in flat observations and captures subtle progress that implicit methods may miss.
  • Main Results: StateFactory achieves average D_EPIC of 0.297 without parameter updates, outperforming the best representation-free baseline at 0.322 and approaching the fully supervised upper bound.The comparison concerns cross-domain generalization.
  • Ablation Studies: State extraction lowers D_EPIC from 0.57 to 0.43, while explicit object-attribute factorization reaches 0.30 by reducing distractors and semantic interference.Object-centric states alone reach 0.35 because entangled attributes limit precision.
  • Ablation Studies: Online goal interpretation increases average D_EPIC by only 0.02 relative to the offline oracle, from 0.28 to 0.30.The offline setting uses the expert trajectory’s final state as the goal; the online setting infers goals from interaction history.
  • Ablation Studies: Reasoning-enhanced language-model modes improve alignment, including gpt-oss-20b thinking medium from 0.36 to 0.30 and qwen3-30B-thinking at 0.41 versus 0.55 standard.Performance generally also improves with parameter count within a model family.
  • Ablation Studies: Embedding discriminative power strictly correlates with StateFactory performance: all-MiniLM achieves D_EPIC approximately 0.30, while gemma-300m reaches 0.45 at triplet accuracy 0.67.Triplet accuracy measures whether positive pairs receive higher similarity than negative pairs.

5 Related Works

Related work covers reward learning, state representations, and language-based agents, emphasizing the challenge of measuring progress while avoiding hand-engineered or noisy reward signals.

  • Goal Conditioned Reinforcement Learning: Goal-conditioned reinforcement learning treats accurate task-advancement measurement as essential for policy optimization, while manual rewards are labor-intensive and vulnerable to reward hacking.Recent approaches use human feedback, demonstrations, preferences, and related supervision.
  • State Representation in World Models: State-representation work contrasts raw observations, which retain task-irrelevant noise, with approaches imposing stronger structural assumptions on world states.The passage frames representation fidelity as important for language-based agent learning.
  • Language-based Agent Learning: Language-based agent learning formulates web navigation, operating-system control, and text-based games as sequential control problems with comparable dynamics.The passage situates these agents alongside physical-control settings and model-free policies.
  • Language-based Agent Learning: Recent approaches refine policies from suboptimal offline data using implicit value frameworks and extend this principle to self-generated agent experiences.These methods aim to leverage existing trajectories without extensive online interaction.

6 Conclusion

The paper argues that structured text-based world states can support accurate, generalizable reward prediction and improve planning across diverse environments.

  • Conclusion: StateFactory factorizes observations into hierarchical object-attribute structures to produce zero-shot reward predictions across domains.The conclusion presents this as the paper’s central method-level finding.
  • Conclusion: StateFactory outperforms supervised models and LLM-as-a-Judge baselines, while its state-driven signals improve system-1 success rates by up to 21.64%.Qualitative analysis also supports system-2 planning by guiding structured search grounded in physical evidence.

A. Benchmark Details and Data Construction

The benchmark appendix describes a five-environment evaluation suite and a data-construction process combining expert positives, randomized augmentation, hybrid rewards, filtered negatives, and native environment modalities.

  • Data Construction: The appendix documents a four-stage pipeline for generating positive and negative trajectories and assigning hybrid ground-truth rewards.The pipeline is designed to mitigate reward sparsity and preserve native environmental feedback.
  • Evaluation Environments: RewardPrediction evaluates agents in AlfWorld, ScienceWorld, TextWorld, WebShop, and BlocksWorld, covering navigation, scientific reasoning, games, web shopping, and classical planning.The environments span household manipulation, causal reasoning, partially observable puzzles, noisy e-commerce navigation, and spatial sequencing.
  • Appendix Scope: The appendix also describes statistical summaries, timestep-wise sample visualizations, semantic alignment, baselines, implementation details, ablations, and planning utilities.These materials cover trajectory statistics, reward derivation, model configurations, and system-1/system-2 evaluation procedures.
  • Data Construction: Positive trajectories derive from expert demonstrations and receive random padding of k in [0, 3] interaction steps for temporal robustness.Padding is distributed stochastically at trajectory initiation and termination.
  • Data Construction: Hybrid labeling assigns expert-segment rewards by linear interpolation, r_t = t/T_expert, while appended random steps inherit native environmental rewards.This combines dense monotonic progress with sharp grounding for random interactions.
  • Data Construction: Negative trajectories come from a stochastic random policy and are discarded when they overlap with any expert sub-segment, reducing accidental-success label noise.The filtering protocol targets false failure labels.
  • Data Construction: The benchmark preserves native observations and the complete admissible action space rather than manually modifying modalities.This minimal-intervention principle is presented as supporting empirical fidelity.
  • Domain-Specific Construction: Domain-specific construction includes uniform AlfWorld sampling, ScienceWorld validity filtering, WebShop human demonstrations, TextWorld procedural generation, and BlocksWorld PDDL-to-text conversion.The appendix tailors sampling, filtering, or generation to each environment’s task structure.

A.4 Dataset Statistics

The RewardPrediction benchmark reports trajectory-pair counts, native reward signals, and task types across domains.

  • Table 3 summarizes positive and negative trajectory pairs across different domains.

A.5 Data Sample Visualization

The paper visualizes positive and negative AlfWorld samples and details reward construction and semantic state-to-goal alignment. Positive expert progress receives increasing rewards, whereas non-progressing random behavior receives zero rewards.

  • Positive Trajectory: Expert segments receive monotonically increasing dense rewards from 0.17 to 1.00 via linear interpolation.The gradient represents task progress across steps t=0 to 5.
  • Positive Trajectory: A final random augmentation inherits the native success signal R=1.0 to preserve the successful state.
  • Negative Trajectory: Random-policy actions that fail to advance the task state receive a uniformly flat ground truth of 0.This defines a clean non-progressing negative baseline.
  • Hierarchical Routing: StateFactory aligns goal attributes with candidate-state keys dynamically rather than assuming fixed linguistic correspondences.Aligned values are then compared semantically for attribute satisfaction.
  • Hierarchical Routing: Joint object selection combines identity and attribute consistency, suppressing scores when either the object or its state is incorrect.The system selects the candidate maximizing the composite alignment score.

C.2.3 Training Setup and Hyperparameters

The training setup uses memory-efficient implementation choices and largely shared configurations, while hardware, batch size, and training duration vary by task.

  • Evaluations use an 8,192-token maximum sequence length with mixed-precision BF16 training and gradient checkpointing.
  • Core optimization hyperparameters remain constant despite task-specific hardware, batch-size, and training-duration differences.Most models train for 3 epochs; WebShop and the Combined model train for 2 epochs.

C.3 LLM-as-a-Judge

The LLM-as-a-Judge evaluation compares intuitive and analytical scoring modes under controlled inference settings, alongside structured StateFactory representations and representation-granularity controls.

  • Evaluation Protocol: Evaluators receive full uncompressed interaction histories within an 8,192-token context window and use temperature 0.01 for near-deterministic outputs.
  • Judge Modes: Intuitive mode requests immediate scores without analysis, whereas Analytical Mode requests step-by-step trajectory review and justification.
  • Compute Scaling: Test-time reasoning effort is compared between low- and medium-effort configurations for reasoning-specialized models.
  • Representation: StateFactory factorizes raw observations into structured object-attribute states through semantic state factorization and dynamic goal interpretation.The first pipeline half grounds, filters, and refines task-relevant state history.
  • Representation Ablation: The representation ablation varies structural granularity from raw observations to textual, object-centric, and fully factorized object-attribute states.The same backbone and input context are retained while output schemas and description constraints change.

C.5.2 LLM Backbone

The ablation shows that reasoning-enhanced and larger language-model backbones improve StateFactory’s reward alignment, while embedding discriminability determines whether semantic contradictions are recognized.

  • Reasoning-enhanced backbones substantially improve alignment, with GPT-OSS-20B improving average EPIC distance from 0.36 to 0.30 when moving from Low to Medium thinking.
  • Qwen3-30B-Thinking achieves an EPIC distance of 0.41 versus 0.55 for its standard instruction-tuned counterpart.
  • Qwen3-30B improves over Qwen3-14B within the same family, reducing EPIC distance from 0.58 to 0.55.
  • StateFactory’s scaling behavior suggests that future gains can come from both larger models and stronger reasoning strategies.
  • Semantic Embeddings: Triplet-based Accuracy measures whether embeddings rank a semantic paraphrase above a negative distractor using cosine similarity.
  • Semantic Embeddings: Embedding models can fail when they cluster states by shared topic rather than truth conditions, whereas discriminative models reject contradictory distractors.

C.5.4 Comprehensive Ablation Study Ablation Results

The ablation results emphasize structured state representations, robust dynamic goal interpretation, semantically discriminative embeddings, and reward-guided candidate selection as key components of the system.

  • Representation: Structured object-attribute representations achieve the lowest EPIC distance among the compared state-representation granularities.
  • Goal Interpretation: Dynamic goal interpretation remains robust relative to offline oracle goals, with only a minimal performance gap and no privileged goal-state information.
  • Semantic Embeddings: Embedding backbones with higher Triplet Accuracy consistently produce lower downstream EPIC distances.
  • LLM Backbone: Reasoning-enhanced and larger LLM backbones consistently improve reward precision, measured by lower EPIC distance.
  • Reward-Guided Planning: StateFactory augments ReAct with dense, state-aware rewards so candidate actions can be evaluated for progress before selection.
  • Reward-Guided Planning: The framework ranks candidate actions using predicted relative gains, helping prioritize progress and avoid redundant or cyclical behavior without expensive lookahead simulations.
Loading 2603.09400v1…