モデルベース強化学習
前回までの記事で紹介した方法では、環境との相互作用を通してデータをサンプリングし、その結果を用いて価値を推定していた。この環境との相互作用にはコストがかかる。たとえばロボットの行動学習の場合は、未熟な方策による試行錯誤によってハードウェアが破損するかもしれないし、実験に人的・物的コストがかかる。シミュレーションでも、精度の高いものが必要ならば計算コストがかかる。もし、行動を選択したときに得られる未来の状態や報酬を予測することができれば、これらのコストを抑えることができる。このような手法は環境をモデリングすることから モデルベース強化学習（RL） と呼ばれる。対して、モデリングを行わない手法は モデルフリー強化学習（RL） と呼ばれる。
モデルベース強化学習は、状態遷移確率 \( p(s_{t+1}|s_t, a_t) \) をニューラルネットワークやベイジアンなどの近似によって学習しつつ、強化学習の価値推定と方策学習を行う。状態 \( s_t \) が画像やセンサー情報などで高次元である場合には、深層学習による次元削減が効果的であり、この次元削減に変分オートエンコーダーを用いて行い、その状態の潜在特徴量と行動を入力とするリカレントニューラルネットワークで遷移確率をモデル化するのが World Models である。このようにオートエンコーダーを用いて潜在空間上での状態遷移のモデリングを行う手法は多く提案されている。たとえば、Dreamerの系譜 (v1, v2, v3) は有名だ。
オートエンコーダーは明示的に状態を復元するが、これはメモリ消費が大きく、価値や方策の学習に必要な情報以外の情報も復元する必要がある。MuZero は次元削減を行った特徴量空間上でモンテカルロ木探索を行うが、状態の復元を明示的には行わない。同様に TD-MPC は復元を行わないが、MuZero とは異なり、現在の状態からエンコーダーによって抽出された（潜在）特徴量と、対応する実際に訪れた未来の状態をエンコードした結果の誤差を最小化することで、特徴量空間のダイナミクスに制約を加えている。さらに、この TD-MPC にモデルアーキテクチャや損失関数の設計の改良を加えた TD-MPC2 も提案されている。
個人的な感覚として言えば、状態または状態遷移のモデリングに復元可能性や一貫性などの制約を加えると学習が容易になる一方、必ずしも価値や方策にとって利用しやすい表現になるとは限らない。逆に MuZero のように制約がほとんどない場合には、全体的な性能が向上する可能性がある反面、大量のサンプルが必要となる。
行動計画
状態遷移確率が得られると行動選択に工夫ができる。もちろん、方策が学習できていれば行動選択自体は可能であるが、さまざまな未来の可能性を実際に計算して考慮することで、より高い価値が得られるような選択が可能になる。たとえば、MuZero は学習後の対戦時にも木探索を実行する。一方、TD-MPC は一定ステップ先まで近似した状態遷移確率を用いて行動の軌道を複数生成し、別途学習したモデルが推定する価値の高いものだけを残して更新する Cross-Entropy Method (CEM) を用いている。
このような行動計画はモデルベース強化学習以外にも用いられている。事前に状態遷移確率が分かっている場合に、報酬を最大化するように行動自体を最適化する Model Predictive Control(MPC) はモデルベースRLと深い関係がある。MPCによる軌道最適化（Trajectory Optimization）は状態遷移確率を近似するわけではないが、与えられた状態遷移確率を用いた動的計画法などで行動計画を立てる。Guided Policy Searchは軌道最適化で得られたデータを使って方策を学習するRLの手法だ。この方法はオンポリシーであるが、決定的であった状態遷移を環境の不確実性としてのノイズが介入する確率的な過程に置き換え、ノイズの事後確率を推定して補正することでオフポリシーにする Counterfactually-Guided Policy Search (CF-GPS) が提案されている。
最近では Transformer や拡散過程モデルを用いた時系列モデルによるサンプリング過程に工夫を加えることで、特に実際に行動選択するオンラインでの行動計画を実現する試みもある。Diffuser は拡散過程モデルで状態と行動の軌道を学習しておき、実際に行動選択する際の生成において Classifier-Free Guidance を用いて報酬和が大きくなる軌道を生成する。拡散過程モデルは事前に決定された長さの軌道しか生成できないが、Diffusion Forcing は軌道の生成をノイズレベルが最大（無情報）のトークンをデノイジングしてデータ分布に従うトークンへ変換するタスクであると捉え、任意のレベルのデノイジングを学習することで自己回帰による生成を可能にしている。この Diffusion Forcing でも Classifier Guidance が適用可能であり、さらに入力のノイズレベルが高いままで生成することで、遠い不確実な未来のみを不確実なまま予測生成できるという興味深い特徴がある。このようなノイズレベルのスケジューリングや Guidance の適用方法をオンラインで探索することもでき、実際に Monte Carlo Tree Diffusion は Guidance を適用するタイミングをモンテカルロ木探索で決定する。
オンラインでの行動計画といえば、昨今の大規模言語モデル（LLM）における Test-Time Scaling や Reasoning といった、回答前の生成における工夫が挙げられる。LLM はテキストの軌道を生成するように学習しており、いわばテキストベースのモデルベース強化学習と言えなくもない。一般的な Reasoning はテキスト上で行われるが、連続空間で行う手法 なども登場しており、よりモデルベース強化学習や軌道生成を行う強化学習の手法を取り入れたり組み合わせたりすることができそうだ。