Previsão de Recompensa com Estados do Mundo Fatorizados

Reward Prediction with Factorized World States

March 10, 2026

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

cs.AI

Resumo

Os agentes devem inferir os resultados das ações e selecionar ações que maximizem um sinal de recompensa que indica a proximidade de atingir o objetivo. A aprendizagem supervisionada de modelos de recompensa pode introduzir vieses inerentes aos dados de treino, limitando a generalização para novos objetivos e ambientes. Neste artigo, investigamos se representações bem definidas do estado do mundo, por si só, podem permitir uma previsão precisa de recompensas em diversos domínios. Para abordar esta questão, introduzimos o StateFactory, um método de representação fatorizada que transforma observações não estruturadas numa estrutura hierárquica de objetos e atributos utilizando modelos de linguagem. Esta representação estruturada permite que as recompensas sejam estimadas naturalmente como a semelhança semântica entre o estado atual e o estado objetivo sob uma restrição hierárquica. No geral, a estrutura de representação compacta induzida pelo StateFactory permite fortes capacidades de generalização de recompensas. Avaliamos o método no RewardPrediction, um novo conjunto de dados de referência que abrange cinco domínios diversos e compreende 2.454 trajectórias únicas de ação-observação com recompensas verdadeiras passo a passo. O nosso método apresenta resultados promissores *zero-shot* contra os modelos de recompensa VLWM-critic e LLM-as-a-Judge, alcançando uma distância EPIC 60% e 8% menor, respetivamente. Além disso, esta qualidade superior da recompensa traduz-se com sucesso numa melhoria do desempenho do planeamento do agente, resultando em ganhos na taxa de sucesso de +21,64% no AlfWorld e +12,40% no ScienceWorld em comparação com políticas reativas do sistema 1 e melhorando o planeamento do agente do sistema 2. Página do Projeto: 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