[論文レビュー] Reward Prediction with Factorized World States

Yijun Shen, Delong Chen+5 more|arXiv (Cornell University)|Mar 10, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

本論文は observations を object–attribute 階層へと変換して zero-shot で報酬を予測する semantic factorization 手法 StateFactory を提案し、RewardPrediction ベンチマークを用いてドメイン間での報酬品質を評価する。StateFactory は zero-shot EPIC 距離の大幅な低減と計画性能の改善を達成する。

ABSTRACT

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

研究の動機と目的

• 新しい目標や環境に対してタスク固有の監視なしで頑健な報酬予測を一般化させることを動機づける。
• 言語モデルを用いて世界状態を object–attribute 階層に semantic factorization することを提案する。
• 階層的制約の下で current state と goal state の semantic 類似度として報酬を推定できるようにする。
• ゼロショットの報酬一般化とそれが多様なドメインでの計画性能に与える影響を評価する。
• テキストベース環境における step-wise 報酬品質を厳密に評価するベンチマーク（RewardPrediction）を提供する。

提案手法

• Observation を構造化された object–attribute 状態へ分解し、動的なゴール状態を反復的に grounding する StateFactory を導入する。
• 再帰的な状態抽出関数を用いて、識別子を有し属性が変化する一連のオブジェクト実例を生成する。
• 反復的なゴール解釈関数を介してゴールを dynamic なゴール表現へ grounding する。
• 階層的な object–attribute 整列を通じて現在の状態と grounding されたゴール状態との semantic 類似度として報酬を計算する。
• 5ドメインに渡る ground-truth なステップ毎報酬と EPIC 距離を用いて報酬信号を評価する。

実験結果

リサーチクエスチョン

• RQ1 ゼロショット StateFactory はベースラインよりも正確な報酬信号を提供するか。
• RQ2 StateFactory は supervised 報酬モデルより未知ドメインへ一般化できるか。
• RQ3 表現の粒度（object vs. object-attribute）が性能にどのように影響するか。
• RQ4 埋め込み、バックボーン、推論能力の選択に対して StateFactory はどれだけ頑健か。

主な発見

• StateFactory はゼロショット平均 EPIC 距離 0.297 を達成し、最良の表現なしベースラインを上回り、監視付き上限値に接近する。
• supervised 報酬モデルは新規ドメインへの一般化が乏しく、報酬予測誤差が平均で 138% 増加する。
• StateFactory の高精度な object–attribute 状態表現はノイズを低減し、ゴールとの整合性を改善する。
• LLM の推論能力と埋め込みの識別性を高めるほど報酬の整合性が向上する。
• StateFactory 由来の報酬信号は計画の改善へと繋がり、例として AlfWorld で +21.64%、ScienceWorld で +12.40% の ReAct + StateFactory による成功率向上を示す。