StateFactory: Zero-Shot Reward Prediction

• StateFactory is a zero-shot reward prediction framework that converts unstructured observations into explicit object–attribute graphs for compositional state representation.
• It uses prompt-based large language model calls to extract world and goal states, leveraging semantic similarity for reward estimation in partially observable, goal-augmented MDPs.
• Empirical results on the RewardPrediction benchmark demonstrate significant improvements, including a 60% reduction in EPIC distance and enhanced planning performance in diverse domains.

StateFactory is a zero-shot reward prediction framework based on factorized representations that transforms unstructured observations into explicit object–attribute graphs, enabling robust, generalizable reward estimation in partially observable, goal-augmented Markov Decision Processes (MDPs). By structuring the world and goal states hierarchically and using semantic similarity as the core metric, StateFactory provides strong zero-shot reward generalization and improved downstream planning performance, as demonstrated on diverse domains in the RewardPrediction benchmark (Shen et al., 10 Mar 2026).

1. Factorized World-State Representation

StateFactory operates in a partially observable, goal-augmented MDP, denoted as $$\mathcal{M} = \langle \mathcal{S}, \mathcal{A}, \mathcal{O}, \mathcal{T}, \mathcal{R}, \mathcal{G}, \phi, \Omega \rangle$$, where at each time $t$, the agent receives an observation $o_t = \Omega(s_t) \in \mathcal{O}$ and the goal $g \in \mathcal{G}$ is provided as a language string. StateFactory defines two mappings—world-state extraction and goal-state interpretation—both implemented via prompt-based LLM calls:

• World-state extraction:

$$\hat{s}_t = f_{\mathrm{state}}(g, \hat{g}_{t-1}, \hat{s}_{t-1}, o_t, a_{t-1})$$

• Goal-state interpretation:

$$\hat{g}_t = f_{\mathrm{goal}}(g, \hat{g}_{t-1}, \hat{s}_t, o_t, a_{t-1})$$

Each factorized state $\hat{s}_t$ is a set of object instances:

$$\hat{s}_t = \{ e_i \}_{i=1}^N, \quad e_i = \langle d_i, \{ (\alpha_{i,l}, v_{i,l}) \}_{l=1}^{L_i} \rangle$$

where $d_i$ is the object identity (e.g., “Mug”) and $(\alpha_{i,l}, v_{i,l})$ are attribute–value pairs (e.g., “location$t$0on table”). This induces a bipartite object–attribute graph structure suitable for hierarchical reasoning.

2. Reward Estimation as Hierarchical Semantic Similarity

Rewards are defined as semantic similarity between the current state $t$1 and the interpreted goal state $t$2. The per-step reward is:

$t$3

with attribute satisfaction

$t$4

where $t$5 is selected by maximum key similarity:

$t$6

Similarity is computed via cosine distance in the all-MiniLM-L6-v2 text embedding space (384 dimensions). Only object-attribute matches with sufficient identity similarity contribute nontrivially, enforcing hierarchical, context-sensitive reward attribution.

Performance is evaluated against human-annotated ground-truth rewards using the EPIC distance (policy-invariant):

$t$7

where $t$8 is the Pearson correlation.

3. Architecture and Inference Protocol

StateFactory is a zero-shot method that leverages a pipeline of LLM prompt modules:

• Module A: Extracts world state into a structured JSON object–attribute graph.
• Module B: Applies filtering to maintain only task-relevant state history.
• Module C: Builds a dynamic, interpreted goal blueprint encoding static entities and permissible state evolutions.

No finetuning or parameter updates are performed. The system uses GPT-OSS-20B (temperature 0.01, 8k context) for all prompt-based reasoning, and the all-MiniLM-L6-v2 embedding model for similarity computations. Reward estimation is carried out in inference-only mode, avoiding correlation with specific training data.

4. Benchmarking and Empirical Results

StateFactory is evaluated on the RewardPrediction benchmark, which consists of 2,454 stepwise, human-verified trajectories across five domains: AlfWorld, ScienceWorld, WebShop, BlocksWorld, and TextWorld.

Model	Overall EPIC ↓
VLWM-critic (Llama3.2)	0.738
LLM-as-Judge	0.322
StateFactory	0.297

StateFactory achieves a 60% reduction in EPIC distance relative to VLWM-critic, and an 8% lower EPIC distance than LLM-as-a-Judge (both are prior SOTA reward modeling baselines). On ablations, hierarchical object-attribute states (EPIC 0.30) outperform both object-centric (0.35) and flat denoised text (0.43) variants. Oracle (“offline”) vs. online goal representation is stable (EPIC 0.28 vs 0.30).

5. Integration with Reactive and Planning Agents

StateFactory directly augments downstream agents by providing differential reward signals:

For reactive system-1 (ReAct), at every step, action value is scored as:

$t$9

where reward increments and penalties are injected into the prompt to influence the action proposal.

Integration results (success rate increases):

• AlfWorld: 34.33% → 55.97%
• BlocksWorld: 85.00% → 93.00%
• ScienceWorld: 22.63% → 35.03%

For system-2 (planning), StateFactory rewards are used for node evaluation in Monte Carlo Tree Search (MCTS) over world model rollouts, providing reward guidance at each simulated state.

6. Analysis, Limitations, and Future Directions

StateFactory’s compact, structured representation enables generalizable, policy-invariant reward estimation and robust transfer across domains, as corroborated by ablation analyses: the embedding probe shows that higher triplet discrimination accuracy strongly correlates with lower reward error.

Limitations include high dependency on the quality of LLM extraction—very long or noisy observations can cause degradation—and significant inference costs due to three sequential LLM calls and embedding similarity computation per step. Future work is expected to address efficiency and robustness by training lightweight extractors on generated object-attribute labels, incorporating visual encoders for raw sensory inputs, and developing learned key-alignment modules.

In summary, StateFactory provides an explicit, compositional, and semantically robust framework for general reward prediction through LLM-mediated, factorized world-state construction, outperforming both supervised and direct LLM-judgment baselines in both zero-shot reward accuracy and downstream agent performance (Shen et al., 10 Mar 2026).