Reward Prediction with Factorized World States

This paper introduces StateFactory, a method that leverages language models to transform unstructured observations into hierarchical, factorized world states, enabling accurate zero-shot reward prediction via semantic similarity and significantly improving agent planning performance across diverse domains.

The Gist

Imagine you are teaching a robot to cook a complex meal, like a soufflé. The robot needs to know if it's doing a good job at every single step: Did it crack the eggs right? Is the oven hot enough? Did it fold the batter gently?

In the world of AI, this "knowing how well you're doing" is called Reward Prediction.

For a long time, we taught AI to guess its own score by showing it thousands of examples of "good" and "bad" cooking videos. But this is like teaching a student only by showing them past exams. If you give the student a new type of recipe they've never seen before, they get confused because they just memorized the old answers, not the logic. They can't generalize.

This paper introduces a new way to teach AI how to score itself, using a method called StateFactory. Here is the breakdown in simple terms:

1. The Problem: The "Black Box" vs. The "Lego Set"

Most AI agents look at the world like a blurry photograph. They see a jumbled mess of text: "You are in the kitchen. There is a red mug on the table. The stove is on. You are holding a spoon."

If you ask the AI, "How close are you to making coffee?" it has to guess based on that blurry photo. It's hard to tell if the "red mug" is actually the right mug, or if the "stove on" is actually hot.

The Paper's Solution: Instead of a blurry photo, StateFactory turns the world into a Lego set.
It breaks the messy text down into tiny, organized blocks:

• Object: Mug
• Attribute: Color = Red, Location = Table, Temperature = Cold.
• Object: Stove
• Attribute: Status = On, Heat = High.

By turning the world into a structured list of "Things" and their "Properties," the AI can see exactly what is happening, step by step.

2. The Method: The "Checklist" Analogy

Once the AI has this Lego-like structure, it doesn't need to guess the score. It just needs to compare two checklists.

• Checklist A (The Goal): "I need a Hot Mug on the Table."
• Checklist B (Current State): "I have a Cold Mug on the Table."

The AI simply calculates the "distance" between these two lists.

• If the mug is cold, the score is low.
• If the mug is hot, the score goes up.
• If the mug is on the floor, the score goes down.

Because the AI is comparing clear facts (Hot vs. Cold) rather than guessing from a blurry picture, it can figure out the score for any new task, even one it has never seen before. It's like a chef who understands the principles of cooking (heat + time = cooked) rather than just memorizing one specific recipe.

3. The Benchmark: The "Grand Tournament"

To prove this works, the authors built a giant testing ground called RewardPrediction. Imagine a video game tournament with five different levels:

1. AlfWorld: A robot doing household chores (folding laundry, making coffee).
2. ScienceWorld: A robot doing science experiments (mixing chemicals, measuring temperature).
3. WebShop: A robot shopping online (finding a specific blue shoe under $50).
4. TextWorld: A robot playing a text-based adventure game (finding a key to unlock a chest).
5. BlocksWorld: A robot stacking blocks like a puzzle.

They tested their "Lego" method against other AI methods. The results were impressive:

• Old Methods: When given a new level, they got confused and failed (like a student who memorized answers but can't do new math problems).
• StateFactory: It figured out the scoring rules instantly and helped the robot plan better, succeeding in tasks it had never seen before.

4. Why This Matters

Think of it like upgrading from a GPS that only knows one city to a GPS that understands the concept of "roads" and "destinations."

• Before: If you asked the old AI to navigate a new city, it would get lost because it didn't have a map for that specific city.
• Now: With StateFactory, the AI understands the structure of the world. It knows that "putting a hot mug in a cabinet" is a specific sequence of steps, regardless of whether the kitchen is in New York or Tokyo.

The Bottom Line

This paper shows that if you teach an AI to organize its thoughts (breaking the world into objects and attributes) rather than just memorize examples, it becomes much smarter at figuring out what it's doing right or wrong. This allows robots and digital agents to tackle new, complex challenges without needing to be retrained from scratch every time.

In short: They gave the AI a better way to take notes, which helped it understand the game rules so it could win, even on levels it had never played before.

Technical Summary

Here is a detailed technical summary of the paper "Reward Prediction with Factorized World States".

1. Problem Statement

The paper addresses the challenge of zero-shot reward prediction for agents operating in diverse, text-based environments.

• Limitations of Current Approaches:
  • Supervised Reward Models: Training reward models on specific tasks introduces biases and overfitting, limiting their ability to generalize to novel goals or environments (the "domain gap").
  • LLM-as-a-Judge: While flexible, prompting LLMs to judge progress often results in volatile, discrete, and noisy reward signals due to the lack of explicit state tracking.
  • Representation-Free Methods: Directly regressing rewards from raw text observations fails to filter out task-irrelevant noise, making it difficult to measure fine-grained progress.
• Core Challenge: How to construct a world state representation that is compact, structured, and semantically aligned with goals to enable accurate, generalizable reward estimation without task-specific training.

2. Methodology: StateFactory

The authors propose StateFactory, a framework that transforms unstructured observations into hierarchical object-attribute structures using Large Language Models (LLMs). The reward is then estimated as the semantic similarity between the current state and the goal state.

The framework consists of three integrated layers:

A. State Extraction (World State Transition)

Instead of treating the state as a flat text string, StateFactory decomposes observations into a set of object instances, where each object is defined by an identity and a set of dynamic semantic attributes.

• Process: A recurrent function $f_{state}$ updates the state $\hat{s}_t$ by integrating raw observations ($o_t$), action history, and the current goal context.
• Structure: The state is represented as $e_i = \{d_i, \{(\alpha_{i,l}, v_{i,l})\}\}$, where $d_i$ is the object identity (e.g., "Mug") and $(\alpha, v)$ are attribute-value pairs (e.g., "location: on table", "temperature: hot").
• Benefit: This factorization separates entity identity from evolving attributes, allowing the model to track specific changes (e.g., a mug moving from a counter to a table) without being confused by irrelevant text.

B. Goal Interpretation

The framework treats goal interpretation as a dynamic, iterative process rather than a static initialization.

• Process: A function $f_{goal}$ updates the goal state $\hat{g}_t$ based on the textual goal description and the current context.
• Mechanism: It distinguishes between the "blueprint" of success (static requirements) and intermediate plans. It anchors abstract instructions to specific physical entities observed in the environment, ensuring the goal representation evolves as the agent discovers relevant objects.

C. Hierarchical Routing (Reward Calculation)

The reward signal $\hat{r}_t$ is derived by computing the semantic similarity between the factorized current state $\hat{s}_t$ and the dynamic goal state $\hat{g}_t$.

1. Object Matching: For each goal object, the system finds the corresponding physical object in the current state that maximizes identity and attribute consistency.
2. Attribute Matching: It calculates the satisfaction score by averaging the semantic similarity of aligned attribute values (e.g., matching "on the table" in the goal with "on the table" in the state).
3. Aggregation: The global reward is the mean of the local fulfillment scores across all goal objects.

3. The RewardPrediction Benchmark

To rigorously evaluate reward prediction, the authors introduce RewardPrediction, a new benchmark dataset.

• Scope: Covers five diverse domains:
  1. AlfWorld: Embodied household tasks (robotics planning).
  2. ScienceWorld: Scientific reasoning and tool use.
  3. TextWorld: Procedural text-adventure games.
  4. WebShop: E-commerce website navigation.
  5. BlocksWorld: Classical spatial planning.
• Data: Contains 2,454 unique trajectories with step-wise ground-truth rewards.
• Construction Strategy: Uses a "paired data" approach where positive trajectories (expert demonstrations) are augmented with random steps, and negative trajectories (random policies) are strictly filtered to ensure they contain no accidental progress. This forces the model to distinguish meaningful semantic progress from stochastic noise.
• Metric: Uses EPIC Distance (Equivalent Policy-Invariant Comparison), which measures the alignment between predicted and ground-truth reward functions based on Pearson correlation, preserving fine-grained magnitude information.

4. Key Contributions

1. RewardPrediction Benchmark: A comprehensive dataset spanning five domains with 2,454 step-wise labeled trajectories, enabling the first rigorous evaluation of zero-shot reward generalization.
2. StateFactory Framework: A novel representation method that factorizes flat text into hierarchical object-attribute structures, enabling dense reward estimation via semantic similarity rather than regression.
3. Demonstrated Generalization: Proved that structured state representations alone can achieve superior zero-shot performance compared to supervised models and LLM-as-a-Judge baselines.
4. Planning Enhancement: Showed that StateFactory's reward signals significantly improve agent planning performance (System-1 ReAct and System-2 MCTS) by breaking reasoning deadlocks.

5. Experimental Results

The paper evaluates StateFactory against Supervised Reward Models (trained on specific domains), VLWM-critic, and LLM-as-a-Judge baselines.

• Reward Prediction Accuracy (EPIC Distance):
  • Supervised Models: Show a 138% increase in error when transferred to unseen domains, confirming severe overfitting.
  • StateFactory (Zero-Shot): Achieves the best overall performance with an average EPIC distance of 0.297.
  • Comparison: Outperforms the best representation-free baseline (LLM-as-a-Judge) by 8% and VLWM-critic by 60% in terms of EPIC distance reduction.
• Ablation Studies:
  • Granularity: Moving from unstructured text to Object-Attribute factorization is the most critical factor, reducing error significantly compared to object-centric or flat text representations.
  • Reasoning: Models with "Thinking" modes (extended inference) perform significantly better, confirming that decomposing complex states requires deep reasoning.
  • Embeddings: High triplet-based accuracy in the semantic embedding model correlates strictly with better reward prediction.
• Agent Planning Performance:
  • Integrating StateFactory into ReAct agents improved success rates by +21.64% on AlfWorld and +12.40% on ScienceWorld compared to standard reactive policies.
  • In System-2 planning (MCTS), the dense reward signal guided the agent out of local optima and "deadlocks" where pure logic failed.

6. Significance

This work demonstrates that explicit, structured world state representations are a more effective foundation for reward prediction than either supervised learning or raw text prompting. By decoupling state representation from the reward prediction task, StateFactory achieves robust zero-shot generalization across vastly different environments.

The findings suggest a paradigm shift in agent design:

1. From Regression to Similarity: Reward prediction is better framed as a semantic distance calculation in a factorized space rather than a regression task.
2. From Flat to Hierarchical: Agents require explicit tracking of object attributes to understand progress in procedural tasks.
3. From Sparse to Dense: Factorized states enable the generation of dense, continuous reward signals from sparse environments, which is crucial for guiding long-horizon planning without massive trial-and-error.

This approach offers a scalable path toward building agents that can plan effectively in novel, complex environments without requiring extensive domain-specific reward engineering.