Predicción de Recompensas con Estados del Mundo Factorizados

Resumen

Los agentes deben inferir los resultados de las acciones y seleccionar aquellas que maximicen una señal de recompensa que indica cuán cerca se está de alcanzar el objetivo. El aprendizaje supervisado de modelos de recompensa podría introducir sesgos inherentes a los datos de entrenamiento, limitando la generalización a objetivos y entornos novedosos. En este artículo, investigamos si las representaciones bien definidas del estado del mundo, por sí solas, pueden permitir una predicción precisa de recompensas en distintos dominios. Para abordar esto, presentamos StateFactory, un método de representación factorizada que transforma observaciones no estructuradas en una estructura jerárquica de objetos-atributos utilizando modelos de lenguaje. Esta representación estructurada permite estimar las recompensas de forma natural como la similitud semántica entre el estado actual y el estado objetivo bajo una restricción jerárquica. En general, la estructura de representación compacta inducida por StateFactory permite capacidades sólidas de generalización de recompensas. Evaluamos en RewardPrediction, un nuevo conjunto de datos de referencia que abarca cinco dominios diversos y comprende 2.454 trayectorias únicas de acción-observación con recompensas reales paso a paso. Nuestro método muestra resultados prometedores de cero disparos frente a los modelos de recompensa VLWM-critic y LLM-as-a-Judge, logrando una distancia EPIC un 60% y un 8% menor, respectivamente. Además, esta calidad superior de la recompensa se traduce exitosamente en un mejor rendimiento de planificación del agente, produciendo ganancias en la tasa de éxito de +21.64% en AlfWorld y +12.40% en ScienceWorld sobre las políticas reactivas del sistema 1 y mejorando la planificación del agente del sistema 2. Página del proyecto: 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

Predicción de Recompensas con Estados del Mundo Factorizados

Reward Prediction with Factorized World States

Resumen

Support