Predizione della Ricompensa con Stati del Mondo Fattorizzati

Reward Prediction with Factorized World States

March 10, 2026

Autori: Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung

cs.AI

Abstract

Gli agenti devono inferire gli esiti delle azioni e selezionare quelle che massimizzano un segnale di ricompensa che indica quanto si è vicini al raggiungimento dell'obiettivo. L'apprendimento supervisionato di modelli di ricompensa potrebbe introdurre distorsioni insite nei dati di addestramento, limitando la generalizzazione a nuovi obiettivi e ambienti. In questo articolo, indaghiamo se rappresentazioni ben definite dello stato del mondo possano da sole abilitare una previsione accurata della ricompensa attraverso diversi domini. A tal fine, introduciamo StateFactory, un metodo di rappresentazione fattorizzata che trasforma osservazioni non strutturate in una struttura gerarchica oggetto-attributo utilizzando modelli linguistici. Questa rappresentazione strutturata permette di stimare le ricompense in modo naturale come similarità semantica tra lo stato corrente e lo stato obiettivo sotto vincoli gerarchici. In sintesi, la struttura compatta di rappresentazione indotta da StateFactory abilita solide capacità di generalizzazione della ricompensa. Valutiamo il metodo su RewardPrediction, un nuovo dataset di benchmark che copre cinque domini diversi e comprende 2.454 traiettorie azione-osservazione uniche con ricompense reali passo-passo. Il nostro metodo mostra promettenti risultati zero-shot rispetto sia ai modelli di ricompensa VLWM-critic che LLM-as-a-Judge, raggiungendo rispettivamente una distanza EPIC inferiore del 60% e dell'8%. Inoltre, questa qualità superiore della ricompensa si traduce con successo in una migliore performance di pianificazione dell'agente, producento guadagni nel tasso di successo del +21,64% su AlfWorld e del +12,40% su ScienceWorld rispetto a politiche reattive di sistema-1 e potenziando la pianificazione degli agenti di sistema-2. Pagina del progetto: 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