因子分解された世界状態を用いた報酬予測

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報酬モデルと比較して、有望なゼロショット結果を示し、それぞれEPIC距離を60％および8％低減した。さらに、この優れた報酬品質は、エージェントの計画性能の向上に成功裏に変換され、反応的なシステム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