SAPO、単一の自己回帰モデルにactorとcriticを兼任させ、エージェントRLの1イテレーション当たりの学習を33.2%高速化

エージェント強化学習では通常、2種類のコストの間でトレードオフが生じる。PPOはcriticを使って各ステップに信用を割り当てられる一方、policyとほぼ同規模のモデルを別途維持する必要がある。GRPOはcriticを省けるものの、同じプロンプトから複数の軌跡を生成しなければならず、長いプロセスや疎な報酬、あるいはグループ内の結果がすべて同じ場合には、advantageがゼロになる可能性がある。8月20日に公開されたSAPOは、両者を単一の自己回帰モデルに統合しようとする手法だ。

SAPOは、エージェントが各ステップで出力する際の2つの因果的境界から価値を読み取る。行動を生成する前にV(s)を推定し、行動全体を生成した後にQ(s,a)を推定する。研究チームは語彙内の2つのtokenを予約し、そのlogit差を有界スカラーへ写像する。この2つのtokenがデコードされたり、環境へ送られたりすることはない。モデルが各タスクでサンプリングする軌跡は1本だけで、λ-return、Generalized Advantage Estimation（GAE）、batchをまたぐepisode単位の正規化を用いて、最終報酬をより早い段階の意思決定へ逆伝播する。policy、V、Qの損失は、共有バックボーンを1回のbackpropagationで更新する。

Qwen2.5-1.5BとQwen2.5-7Bを3つのrandom seedでテストした結果、SAPOはPPO、GRPOと比べて平均でそれぞれ15.1ポイント、12.1ポイント改善した。7B版はALFWorldで総合成功率94.0%を記録し、WebShopのスコアと成功率はそれぞれ88.6%と82.4%だった。同じ設定におけるPPOの3指標は80.4%、81.4%、68.7%だった。1.5Bを使ったALFWorldのテストでは、1イテレーション当たりの所要時間がPPOの451.2秒から301.4秒へ短縮された。主な要因は、生成する軌跡の削減と、criticの推論および更新経路の除去にある。

この設計は、価値推定をvalue headや第2のモデルに切り出さず、因果言語モデル自体へ戻すため、長いプロセスを扱うツール利用エージェントにとって魅力的だ。ただし、速度測定の対象は1種類のモデルと1つの環境に限られ、学習にはH200を4基、A40を8基使用している。論文はまだコードを公開しておらず、メモリ削減効果、予約tokenの安定性、より大規模なモデルでのpolicy–value干渉は確認できない。今後は独立した再現結果に加え、検索やコード実行など、報酬がより不規則な環境へ移植できるかが注目される。