Prédiction de récompense avec états du monde factorisés
Reward Prediction with Factorized World States
March 10, 2026
Auteurs: Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung
cs.AI
Résumé
Les agents doivent inférer les résultats des actions et sélectionner celles qui maximisent un signal de récompense indiquant la proximité de l'objectif. L'apprentissage supervisé des modèles de récompense pourrait introduire des biais inhérents aux données d'entraînement, limitant la généralisation à de nouveaux objectifs et environnements. Dans cet article, nous étudions si des représentations bien définies de l'état du monde peuvent à elles seules permettre une prédiction précise des récompenses across les domaines. Pour répondre à cette question, nous présentons StateFactory, une méthode de représentation factorisée qui transforme les observations non structurées en une structure hiérarchique objet-attribut à l'aide de modèles de langage. Cette représentation structurée permet d'estimer naturellement les récompenses comme la similarité sémantique entre l'état actuel et l'état objectif sous contrainte hiérarchique. Globalement, la structure de représentation compacte induite par StateFactory permet de solides capacités de généralisation des récompenses. Nous évaluons notre méthode sur RewardPrediction, un nouveau jeu de données de référence couvrant cinq domaines variés et comprenant 2 454 trajectoires action-observation uniques avec des récompenses réelles étape par étape. Notre méthode montre des résultats prometteurs en zero-shot contre les modèles de récompense VLWM-critic et LLM-as-a-Judge, atteignant respectivement une distance EPIC inférieure de 60 % et 8 %. De plus, cette qualité supérieure des récompenses se traduit avec succès par une amélioration des performances de planification des agents, entraînant des gains de taux de réussite de +21,64 % sur AlfWorld et +12,40 % sur ScienceWorld par rapport aux politiques réactives de type système 1, et améliorant la planification des agents de type système 2. Page du projet : 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