Control as Inference ざっくり入門

はじめに

この記事は、強化学習の理論 Advent Calendar 2023の21日目の記事です。
https://qiita.com/advent-calendar/2023/rl-theory

Control as Inference (CaI)とは、最適制御や強化学習を確率グラフィカルモデル上の推論として定式化する理論の枠組みです。
CaIから導かれる代表的なアルゴリズムにSoft Actor-Critic (SAC)があります。
最近ではSACやその派生手法が広く知られ、その理論や実装を日本語で説明する記事も多い印象がある一方で、このCaIについては軽い気持ちで読める日本語の記事が少なそうだったので入門記事を書いてみることにしました。

この記事では、主に（通常の）強化学習とControl as Inferenceそれぞれが持つ目的関数の違いに着目して説明してみたいと思います。

対象読者は、強化学習は知っているがCaIにはあまり馴染みがない...という方です。
また、あくまでこの記事の内容はCaIの枠組みの入り口を少し丁寧めに解説しているだけですので、より深い内容やアルゴリズムの導出が気になる方は次の参考資料などをご覧ください。

参考資料は以下のS. Levine先生のチュートリアル論文です。
ただ、この記事内で参考資料にはない独自のnotationを使っている部分があるので注意してください。
Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review

（通常の）強化学習のおさらい

CaIの特徴を比較しながら説明するために、まずは通常の強化学習の枠組みを軽く振り返っておきます。
強化学習は、$\langle\mathcal{S},\mathcal{A},p_{0},p,r\rangle$で定義されるMDP上の期待割引報酬和

$$\mathbb{E}^\pi\left[\sum^T_{t=1}\gamma^{t-1}r(s_t,a_t)\right]$$

を最大化する方策を求める最適化問題として定式化されます。ここで、$\mathcal{S},\mathcal{A}$はそれぞれ状態・行動空間、$r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$は報酬関数です。
また、CaIと合わせるために終端時刻$T$の有限時刻MDPを採用し、初期状態確率と状態遷移確率をどちらも$p$で表すことにします（つまり、$p(s_1),p(s_{t+1}|s_t,a_t)$）。

このとき、方策$\pi(a_t|s_t)$に従った時の状態行動軌跡$\tau=(s_1,a_1,...,s_T,a_T)$が従う分布は、

$$p^\pi(\tau)=p(s_1)\prod^{T-1}_{t=1}p(s_{t+1}|s_t,a_t)\prod^{T}_{t=1}\pi(a_t|s_t)$$

となり、それをグラフィカルモデルにしてみると次のようになります。

通常の強化学習（の多くの問題設定）では報酬関数は決定的です。
つまり、上の分布に従って状態行動軌跡$\tau$が一つ定まると、その軌跡に対する割引報酬和も一意に定まります。
そこで、ある軌跡に対して割引報酬和を計算する関数を次のように定義してみます。

$$C(\tau)=\sum^T_{t=1}\gamma^{t-1}r(s_t,a_t)$$

すると、先ほどの強化学習の目的関数は

$$\mathbb{E}_{\tau\sim p^\pi(\tau)}\left[C(\tau)\right]$$

と書けますね。

さて、ここまで通常の強化学習について主に目的関数に焦点を当てて振り返りました。
実は、CaIはMDP上での意思決定問題を、通常の強化学習とは異なる目的関数を最大化する最適化問題として定式化していると解釈することができます。

Control as Inferenceの概要と強化学習との違い

最適性変数の導入

続いて、Control as Inference (CaI)が扱う最適化問題を見ていきます。

まず、CaIが導入されたモチベーションの部分を紹介します。
先ほど確認した通り、通常の強化学習では状態遷移確率と方策から成る系列モデルによって、意思決定を行うシステムのダイナミクスをモデリングしています。
そのため、そのダイナミクスだけに着目すると、今まで統計・機械学習分野で研究されてきた様々な確率モデル上の推論手法が利用できそうに見えます。
しかし目的関数及びそれを構成する報酬関数については（明示的には）そのモデリング内に含まれていないため、確率的推論手法によって強化学習をend-to-endで解けません^[1]。

そこで、確率モデルとして扱いやすい形で、状態・行動の最適性を表す確率変数を導入したくなってきます。　
CaIにおいて、最適性を表す指標として導入されるのが最適性変数$\mathcal{O}_t$です。
この最適性変数$\mathcal{O}_t$は、時刻$t$の状態と行動ペアに対してそれが最適であるかを示す変数であり、0か1のどちらかの値を確率的にとる二値変数です。
具体的には、報酬関数$r(s_t,a_t)$を用いて

$$p(\mathcal{O}_t=1|s_t,a_t)=\exp\left(r(s_t,a_t)\right)$$

のように定義される分布に従います（ただし、$r(s_t,a_t)<0$）。

CaIでは一番シンプルなモデルとして、各時刻の最適性変数を導入した次図のグラフィカルモデルを考えます。
このとき、行動は状態に依存しないモデルを考え、特に行動の事前分布として行動空間上の一様分布を用います。

このモデルでは、各時刻の状態・行動を潜在変数、最適性変数を観測変数としています。
イメージとしては、見えない水面下でランダムに行動が選ばれながらシステムが変化していき、外からは各時刻でその動きが最適であったかどうかだけが観測できるモデル、といったところでしょうか。
このモデルを使って、常に最適性変数が1となるような行動を推論し、それを選択することでうまく意思決定が出来そうです。

強化学習の目的関数との比較

CaIでは、報酬関数の代わりに最適性変数を用いて価値や方策の最適性を考えていきます。
ただ、最適性変数がなぜこのような定義になっているのかピンと来ない方も多いのではないでしょうか。
そこで、もう少しその性質を見てみましょう。

各時刻の最適性変数は同一時刻の状態・行動のみに依存するため、ある状態行動軌跡$\tau$が与えられた際に、常に最適性変数が1となる対数尤度は定義より次のように書けます。

$$\begin{aligned} \log p(\mathcal{O}_{1:T}=1|\tau)&=\log\left(\prod^T_{t=1}\exp(r(s_t,a_t))\right)\\ &=\sum^T_{t=1}r(s_t,a_t) \end{aligned}$$

つまり、この対数尤度は強化学習において軌跡$\tau$が与えられたときの報酬和$C(\tau)$（ただし$\gamma=1$）に一致しており、最尤推定が収益最大化に対応していることがわかります。
最尤推定は統計論的な機械学習の常套手段の一つなので、こう見ると統計・機械学習における既存の知見を生かしやすい定式化になっていると感じられるのではないでしょうか。

（通常の）強化学習との本質的な違い

一方で、CaIは通常の強化学習とは異なる定式化をしていることから、本質的に異なる部分も存在します。
この違いを調べるために、CaIで用いるグラフィカルモデルにおいて各時刻の行動が状態に依存した方策で決まる次図のモデルを考えてみます。
（このモデルで方策を一様分布にすると、元のCaIのモデルに戻るので一般化になっています）

このモデルの同時分布は

$$\begin{aligned} p(\tau,\mathcal{O}_{1:T}=1)&=p(s_1)\prod^{T-1}_{t=1}p(s_{t+1}|s_t,a_t)\prod^T_{t=1}\pi(a_t|s_t)\prod^T_{t=1}p(\mathcal{O}_{t}=1|s_t,a_t)\\ &=p(s_1)\prod^{T-1}_{t=1}p(s_{t+1}|s_t,a_t)\prod^T_{t=1}\pi(a_t|s_t)\exp\left(\sum^T_{t=1}r(s_t,a_t)\right)\end{aligned}$$

と書けます。この分布はある方策$\pi(a_t|s_t)$に従って行動選択した際のダイナミクスとその最適性のモデルと解釈できるので、周辺対数尤度$\log p(\mathcal{O}_{1:T}=1)$は強化学習で言うところの期待報酬和になってほしい気がします。

しかし、実際に計算してみると、

$$\begin{aligned} \log p(\mathcal{O}_{1:T}=1)&=\log\int p(s_1)\prod^{T-1}_{t=1}p(s_{t+1}|s_t,a_t)\prod^T_{t=1}\pi(a_t|s_t)\exp\left(\sum^T_{t=1}r(s_t,a_t)\right)d\tau\\ &=\log\mathbb{E}_{\tau\sim p^\pi(\tau)}\exp\left(\sum^T_{t=1}r(s_t,a_t)\right) \end{aligned}$$

となり、一般的な期待報酬和では軌跡に対する期待値だった部分が、$\log\mathbb{E}\exp$になります。
この違いが、（おそらく）CaIと通常の強化学習の本質的な違いであり、CaIから導かれるアルゴリズムが持つ方策エントロピー最大化の特性もここに由来するものです。

おわりに

この記事では、Control as Inferenceの枠組みの入り口部分を（通常の）強化学習との違いに着目しながら解説しました。

繰り返しになりますが、この記事の内容はCaIの枠組みの概要のみに留まっていて、ここから更にCaI特有のソフト価値関数の定義や、ナイーブに定義したソフト価値関数の欠点を補うための変分推論の導入などの内容に発展していきます。
それらの内容が気になる方は記事冒頭のチュートリアル論文などをご覧ください。
（もし、需要がありそうであれば続きを書くかもしれません。）

急いで書き上げた記事なので、もし間違い等があればご連絡ください。
それでは最後までご覧頂きありがとうございました。

脚注

1. 各時刻で受け取る報酬を確率変数$R_t$として定めてやるとグラフィカルモデルに導入できる気もしますが、その累積和が従う分布を考えるために畳み込みが必要になって面倒ということだと思います...