Beloningsvoorspelling met Gefactoriseerde Wereldtoestanden

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

Samenvatting

Agenten moeten de uitkomsten van acties afleiden en handelingen selecteren die een beloningssignaal maximaliseren dat aangeeft hoe dicht het doel bij bereiking is. Supervised learning van beloningsmodellen kan vooroordelen introduceren die inherent zijn aan trainingsdata, wat de generalisatie naar nieuwe doelen en omgevingen beperkt. In dit artikel onderzoeken we of goed gedefinieerde wereldtoestandrepresentaties op zichzelf accurate beloningsvoorspelling over domeinen heen mogelijk kunnen maken. Om dit aan te pakken, introduceren we StateFactory, een gefactoriseerde representatiemethode die ongestructureerde observaties transformeert naar een hiërarchische object-attribuutstructuur met behulp van taalmodel(len). Deze gestructureerde representatie maakt het mogelijk om beloningen natuurlijk in te schatten als de semantische gelijkenis tussen de huidige toestand en de doeltoestand onder hiërarchische beperking. Over het geheel genomen maakt de compacte representatiestructuur, geïnduceerd door StateFactory, sterke beloningsgeneraliseermogelijkheden mogelijk. We evalueren op RewardPrediction, een nieuwe benchmarkdataset die vijf diverse domeinen omvat en bestaat uit 2.454 unieke actie-observatietrajecten met stapsgewijze ground-truth beloningen. Onze methode toont veelbelovende zero-shot resultaten tegen zowel VLWM-critic- als LLM-as-a-Judge beloningsmodellen, met respectievelijk 60% en 8% lagere EPIC-afstand. Bovendien vertaalt deze superieure beloningskwaliteit zich succesvol in verbeterde planningsprestaties van agenten, wat successratio-winsten oplevert van +21,64% op AlfWorld en +12,40% op ScienceWorld ten opzichte van reactieve systeem-1-beleidsregels en de planning van systeem-2-agenten verbetert. Projectpagina: 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