요인화된 세계 상태를 활용한 보상 예측
Reward Prediction with Factorized World States
March 10, 2026
저자: Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung
cs.AI
초록
에이전트는 행동 결과를 추론하고 목표 달성에 대한 근접도를 나타내는 보상 신호를 최대화하는 행동을 선택해야 합니다. 보상 모델의 지도 학습은 훈련 데이터에 내재된 편향을 도입할 수 있어, 새로운 목표와 환경에 대한 일반화를 제한합니다. 본 논문에서는 명확히 정의된 세계 상태 표현만으로도 다양한 도메인에서 정확한 보상 예측이 가능한지 연구합니다. 이를 위해 언어 모델을 사용하여 비정형 관측값을 계층적 객체-속성 구조로 변환하는 인수분해 표현 방법인 StateFactory를 제안합니다. 이 구조화된 표현을 통해 계층적 제약 조건 하에서 현재 상태와 목표 상태 간의 의미론적 유사성으로 자연스럽게 보상을 추정할 수 있습니다. 전반적으로 StateFactory가 유도하는 간결한 표현 구조는 뛰어난 보상 일반화 능력을 가능하게 합니다. 우리는 5개의 다양한 도메인을 아우르고 단계별 실제 보상이 포함된 2,454개의 고유 행동-관측 궤적으로 구성된 새로운 벤치마크 데이터셋인 RewardPrediction에서 평가를 수행했습니다. 우리의 방법은 VLWM-critic 및 LLM-as-a-Judge 보상 모델 대비 각각 60%, 8% 더 낮은 EPIC 거리를 달성하며 유망한 제로샷 결과를 보여줍니다. 나아가, 이 우수한 보상 품질은 에이전트 계획 성능 향상으로 효과적으로 이어져, 반응형 시스템-1 정책 대비 AlfWorld에서 +21.64%, ScienceWorld에서 +12.40%의 성공률 향상을 달성하고 시스템-2 에이전트 계획을 개선합니다. 프로젝트 페이지: 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