Прогнозирование вознаграждения с факторизованными состояниями мира

Reward Prediction with Factorized World States

March 10, 2026

Авторы: Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung

cs.AI

Аннотация

Агенты должны предсказывать результаты действий и выбирать операции, максимизирующие сигнал вознаграждения, который указывает на близость достижения цели. Обучение с учителем моделей вознаграждения может привнести смещения, присущие обучающим данным, ограничивая обобщаемость для новых целей и сред. В данной работе мы исследуем, могут ли хорошо определенные представления о состоянии мира сами по себе обеспечить точное предсказание вознаграждения в различных областях. Для решения этой задачи мы представляем StateFactory — метод факторизованного представления, который преобразует неструктурированные наблюдения в иерархическую объектно-атрибутную структуру с использованием языковых моделей. Такое структурированное представление позволяет естественным образом оценивать вознаграждение как семантическое сходство между текущим и целевым состоянием при иерархическом ограничении. В целом, компактная структура представления, порождаемая StateFactory, обеспечивает мощные возможности обобщения вознаграждения. Мы проводим оценку на RewardPrediction — новом наборе данных, охватывающем пять различных доменов и содержащем 2 454 уникальные траектории "действие-наблюдение" с пошаговыми истинными вознаграждениями. Наш метод демонстрирует перспективные результаты в режиме zero-shot по сравнению с моделями вознаграждения VLWM-critic и LLM-as-a-Judge, достигая на 60% и 8% меньшего расстояния EPIC соответственно. Более того, это превосходное качество оценки вознаграждения успешно транслируется в улучшенную производительность планирования агента, обеспечивая прирост процента успешных завершений на +21,64% в AlfWorld и +12,40% в ScienceWorld по сравнению с реактивными политиками Системы 1 и усиливая планирование агентов Системы 2. Страница проекта: https://statefactory.github.io

English

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