# Reward Prediction with Factorized World States

arXiv: https://arxiv.org/abs/2603.09400 · HTML: https://mlsanity.com/paper/2603.09400/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Process reward models for reasoning, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 82 of its 96 references that are in the map.
- First posted on arXiv 10 Mar 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 30 of 82)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Generative Agents: Interactive Simulacra of Human Behavior](https://mlsanity.com/paper/2304.03442/) (2023)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [WebArena: A Realistic Web Environment for Building Autonomous Agents](https://mlsanity.com/paper/2307.13854/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [Mind2Web: Towards a Generalist Agent for the Web](https://mlsanity.com/paper/2306.06070/) (2023)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [Qwen2.5-Coder Technical Report](https://mlsanity.com/paper/2409.12186/) (2024)
- [OSWorld: Benchmarking Multimodal Agents for Open-Ended Tasks in Real Computer Environments](https://mlsanity.com/paper/2404.07972/) (2024)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Towards Understanding Sycophancy in Language Models](https://mlsanity.com/paper/2310.13548/) (2023)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [Open Problems and Fundamental Limitations of Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2307.15217/) (2023)
- [Eureka: Human-Level Reward Design via Coding Large Language Models](https://mlsanity.com/paper/2310.12931/) (2023)
- [ORPO: Monolithic Preference Optimization without Reference Model](https://mlsanity.com/paper/2403.07691/) (2024)

## Built on by (most cited first, 2 of 2)

- [Self-Evolving World Models for LLM Agent Planning](https://mlsanity.com/paper/2606.30639/) (2026)
- [HyperWorld: Hypergraph-Structured State Serialization Improves Learned Textual World Models](https://mlsanity.com/paper/2609.00002/) (2026)

## Abstract (arXiv)

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