今日はベルマン方程式について。強化学習を勉強していると最初の方に出てくるやつ。専門書を読みながら式変形を追うのに精いっぱいになっていると、いつの間にかその式がもつ意味を見失っちゃいますよね。

ってことで本記事では、細かい式変形は追わずに「ベルマン方程式の気持ち」を説明することを目指す。説明の流れを重視するので残念ながら厳密性は大いに欠いている。解の一意性とか極限の収束性とか。ちなみに、ベルマン方程式は強化学習だけでなく経済学などにも応用があるらしい。どんな経済の問題を解くのに使われるんだろうか。

本稿の目次は以下の通りである。

準備：Fixed-Point Iteration
マルコフ決定過程
方策から評価関数の計算（ベルマン期待方程式）
価値反復法（ベルマン最適方程式）
ベルマン方程式の気持ちを整理する
さいごに
追記

なお本稿の内容や数式の記法は Richard S. Sutton 氏らの Reinforcement Learning: An Introduction, second edition に基づいている。この本の表記は、確率変数は大文字、その他の変数は小文字で表記されているのがよい。

準備：Fixed-Point Iteration

まずは関数の不動点を求める数値計算のテクニックから紹介する。

関数 $f$ の不動点（fixed point）とは、 $x = f(x)$ を満たすような $x$ である。つまり、関数を適用しても値が変わらない入力が不動点である。

ある "よい性質" をもつ関数 $f$ に対して不動点、すなわち方程式 $x = f(x)$ の解を求めるテクニックとして Fixed-Point Iteration というやり方が知られている。Fixed-Point Iteration では、 $n= 0, 1, 2, \dots$ に対して $x_{n+1} = f(x_{n})$ を適用してゆくことで数列 $x_0, x_1, x_2, \dots$ を得る。このとき $x_0$ の値は適当に設定する。この数列の極限が $x$ に収束するとき、 $x$ が関数 $f$ の不動点になることが知られている。要するに、関数 $f$ を適当な $x_0$ に適用しまくることで不動点が得られるわけである。非常に簡単な話である。ただし、この不動点の求め方ができるのは関数 $f$ が "よい性質" をもつ場合のみであった。この条件について詳しくは Wikipedia のページを参照のこと。

Wikipedia のページに載っている例を紹介する。ここでは方程式 $x = \text{sin}(x)$ の解、すなわち関数 $\text{sin}(x)$ の不動点を求めることを考える。たとえば $x_0 = 2$ を初期値として $x_1 = \text{sin}(x_0), x_2 = \text{sin}(x_1), x_3 = \dots$ を求めてゆく。この数列の極限は 0 に収束するため、求めたい不動点は 0 であることが分かる。実際、 $\text{sin}(0) = 0$ でありこの解は正しい。この極限を求める過程を示したのが以下の図である。

f:id:t-keita:20210628232800p:plain:w600

青色の線で描かれているのが曲線 $y = \text{sin}(x)$ と直線 $y = x$ である。 $\text{sin}(2)$ などの関数の出力を次の関数の入力にするために、直線 $y = x$ によって y 軸の値を x 軸の値へ変換している。赤色の線が関数 $\text{sin}(x)$ の入出力の値を示しており、徐々に不動点の $x = 0$ 上の点（原点）に近づいている。この計算はコンピュータで計算可能であり、関数適用の回数を増やすほど精度の高い解が得られる。

重要なこととして、ベルマン方程式はまさに $x = f(x)$ という形をもった方程式である。しかも、ベルマン方程式における関数 $f$ は "よい性質" を持っている。したがって、関数 $f$ を適用しまくるだけでベルマン方程式の解が求まるのである。

マルコフ決定過程

次にマルコフ決定過程（Markov decision process）について説明する。とはいっても真面目に説明すると長くなるので、ここでは本稿の表記を明確にするくらいにしておく。このモデルの意味や背景など詳しくは世の中の書籍を参照のこと。

状態 $s \in \mathcal{S}$ , 行動 $a \in \mathcal{A}$ , 報酬 $r \in \mathcal{R}$
確率を出力とするダイナミクス関数 $p(s', r \mid s, a)$
時刻 $t$ における確率変数：状態 $S_t$ , 行動 $A_t$ , 報酬 $R_t$
目的関数： $G_t = R_{t+1} + \gamma R_{t+2} + \gamma^{2} R_{t+3} + \dots$

典型的なマルコフ決定過程の問題は目的関数 $G_t$ を最大化するための方策 $\pi(s \mid a)$ を求めることである。また、方策 $\pi$ に従ったときの状態 $s$ の価値関数を $v_{\pi}(s) = \mathbb{E}[G_t \mid S_t = s$ ] とする。

おまけとして、方策 $\pi(s \mid a)$ を含めたマルコフ決定過程のベイジアンネットワークを以下に示す。状態 $S_t$ において方策が出力した行動 $A_t$ が、報酬 $R_t$ や次の状態 $S_{t+1}$ に影響を与えることがひと目で分かる。

f:id:t-keita:20210630021132p:plain:w400

方策から評価関数の計算（ベルマン期待方程式）

まずは方策反復評価（iterative policy evaluation）と呼ばれる、方策から価値関数を計算するアプローチについて説明する。

前提として価値関数 $v_{\pi}(s)$ は方策 $\pi$ に依存する。よい方策を選ぶほど価値関数の値は大きくなる。ただし、方策 $\pi$ が与えられたときに価値関数 $v_{\pi}(s)$ の値を計算する方法は自明でない。この計算方法を以下に示す。

まず以下の ベルマン期待方程式（Bellman equation） を考える。"方程式" という名前が付いているが未知数は含んでいないことに注意する。英語の "equation" という言葉は必ずしも未知数を含まない単なる "等式" くらいの意味である。

$v_{\pi}(s) = \sum_{a} \pi(a \mid s) \sum_{s'} \sum_{r} p(s', r \mid s, a)[r + \gamma v_{\pi}(s')$ ]

この式は $v_{\pi}(s)$ の定義から導出されたものである。ポイントは、状態 $s$ から期待される合計報酬 $v_{\pi}(s)$ を、遷移先の状態 $s'$ から期待される合計報酬 $v_{\pi}(s')$ によって表現している点である。すなわち再帰的な構造になっている。

式をスッキリさせるためにベルマン期待方程式の右辺を $v$ の関数 $\textbf{B}_{\pi}(v)$ で置換し以下の式を得る。

$v_{\pi} = \textbf{B}_{\pi}(v_{\pi})$

ただし $(\textbf{B}_{\pi}(v))(s) = \sum_{a} \pi(a \mid s) \sum_{s'} \sum_{r} p(s', r \mid s, a)[r + \gamma v(s')$ ]

プログラミングっぽく言うと、関数 $\textbf{B}_{\pi}(v)$ は関数を引数とする高階関数である。関数 $v$ の値域は状態集合であり有限なので、 $v$ は数学的な関数というよりは Python の連想配列や Java の Map 型オブジェクトをイメージすればよい。連想配列の同値性が同じキーに対して同じ値がマッピングされているかどうかで判定されるように、 $v_{\pi} = \textbf{B}_{\pi}(v_{\pi})$ も同じやり方で関数としての同値が成立することを示している。

ここで $v_{\pi}$ は以下の方程式の解である。

$v = \textbf{B}_{\pi}(v)$

なお、価値関数 $v_{\pi}$ と異なる価値関数 $v'$ はこの方程式の解にならないことが知られている。つまり、ベルマン期待方程式を一般化したこの方程式を解くことで方策 $\pi$ に従う価値関数 $v_{\pi}(s)$ を求めることができる。この方程式をなんとかして解きたい。

この方程式を見れば分かるが、 $v$ は関数 $\textbf{B}_{\pi}$ の不動点になっている。そのためこの方程式を解くには不動点の計算を行えばよい。具体的には、適当な $v_0$ に対して関数 $\textbf{B}_{\pi}$ を繰り返し適用し $v_1, v_2, \dots$ の値が収束すればそれが $v_{\pi}$ である。これにて、与えられた方策 $\pi$ から価値関数 $v_{\pi}(s)$ を導くことができた。めでたし。

以上が方策反復評価と呼ばれるアプローチである。ここまでの話を簡単に整理する。特定の値 $v_{\pi}$ を求めたい状況であり、 $v_{\pi}$ は $v_{\pi} = \textbf{B}(v_{\pi})$ を満たすことは分かっていた。そこで、一般化した方程式 $v = \textbf{B}(v)$ を解くことでその解として $v_{\pi}$ の値を知ろうとした。この方程式の解を不動点の計算によって求めた。

価値反復法（ベルマン最適方程式）

上記の方策反復評価では方策 $\pi$ から価値関数 $v_{\pi}$ を求める方法を示した。ここでマルコフ決定過程の問題を思い出すと、やりたいことは目的関数 $G_t$ を最大化するための方策 $\pi$ を求めることであった。そこで、この最適な方策を 最適方策 と呼び $\pi_{\ast}$ と表すことにする。また、最適方策 $\pi_{\ast}$ のもとでの価値関数を $v_{\ast}(s)$ とする。

実は、最適な価値関数 $v_{\ast}(s)$ さえ求まれば、そこから最適方策 $\pi_{\ast}$ を作ることは簡単である。具体的には、各状態 $s$ から遷移できる状態 $s'$ のうち価値関数 $v_{\ast}(s')$ が最大の状態に常に遷移するような方策が最適方策 $\pi_{\ast}$ である。貪欲（greedy）によい状態を選ぶ方策がよいということである。

すなわち、最適方策の価値関数 $v_{\ast}(s)$ を求めることがマルコフ決定過程問題を解くことに直結する。以下、 $v_{\ast}(s)$ を求めるための価値反復法（value iteration）の流れを説明する。

まずはベルマン期待方程式を立ててみる。

$v_{\ast}(s) = \sum_{a} \pi_{\ast}(a \mid s) \sum_{s'} \sum_{r} p(s', r \mid s, a)[r + \gamma v_{\ast}(s')$ ]

これで不動点を求める計算をすれば価値関数 $v_{\ast}$ が求まりそうであるが１つ問題がある。それは、右辺の方策 $\pi_{\ast}$ こそが今から求めたいものであり未知なのである。よってこの式からは価値関数 $v_{\ast}$ を上手く求められない。困った。

そこで登場するのが ベルマン最適方程式（Bellman optimality equation） である。

$v_{\ast}(s) = \text{max}_{a} \sum_{s'} \sum_{r} p(s', r \mid s, a)[r + \gamma v_{\ast}(s')$ ]

この式のポイントは最適方策 $\pi_{\ast}$ を含まない $v_{\ast}(s)$ の再帰的な式になっていることである。最適な価値関数を考えているので、確率的に状態遷移する方策の代わりに貪欲な方策、すなわち $\text{max}$ を使えばよいのである。非常によくできている。

あとは簡単で、不動点を求めるために関数 $\textbf{B}_{\ast}(v)$ を用いて以下のように変形する。

$v_{\ast} = \textbf{B}_{\ast}(v_{\ast})$

ただし $(\textbf{B}_{\ast}(v))(s) = \text{max}_{a} \sum_{s'} \sum_{r} p(s', r \mid s, a)[r + \gamma v_{\ast}(s')$ ]

例によって $v_{\ast}$ は以下の方程式の解である。

$v = \textbf{B}_{\ast}(v)$

これは関数 $\textbf{B}_{\ast}$ の不動点なので、適当な $v_0$ に関数 $\textbf{B}_{\ast}$ を繰り返し適用することで $v_{\ast}$ に収束する。あとは貪欲に最適方策 $\pi_{\ast}$ を求めればおしまい。マルコフ決定過程が解けた。めでたしめでたし。

ベルマン方程式の気持ちを整理する

最後に２つのベルマン方程式を以下の図に整理する。

f:id:t-keita:20210630023224p:plain:w600

ベルマン期待方程式は、方策 $\pi$ と価値関数 $v_{\pi}$ が満たす一般的な関係である。方策が既知であれば価値関数のみが未知数であるため方程式を解くことで価値関数を求めることができる。しかし最適方策を求めたいときは最適な価値関数も未知であるため、1つの方程式の中に2つの未知数があることになり、最適方策を求めるのには使えなかった。

一方で、ベルマン最適方程式は最適な価値関数のみが満たす特別な関係である。方程式を解くことで最適な価値関数を求めることができる。最適だからこそ方策を含まない等式が立てられるのである。max を含む方程式を解くのは難しそうに見えるが、実際は不動点計算をするだけなので難なく解を得られる。

だいたいの書籍ではこれらの方程式が同じような感じで説明されるが、それぞれがもつ意味合いも使い方も異なるので注意したい。