強化学習に出てくるベルマン方程式を理解しよう

はじめに

強化学習の基礎に置かれている「ベルマン方程式」について、言葉は知っているが実はちゃんと理解していないという方は意外と多いのではないかと思われます。これを知っていようが知っていまいが、正直世の中の便利なフレームワークを活用すれば強化学習を実行することは可能であるためだと推測されます。

しかし、ある種の出発点になっているはずの基礎方程式を無視して、ガチャガチャ色々試してみても、なんだかフワついたままでモヤモヤしてしまうのではないでしょうか。少なくとも自分はそうです。

なので今回はベルマン方程式を基本から丁寧に解説していきたいと思います。

ベルマン方程式の概要

細かい話をする前に、ベルマン方程式がどういうものであるのか、そのゴールを示しておきましょう。 ベルマン方程式とは状態 $\bf x$ を外部からの入力 $\bf u$ で制御できると考えているときに、ある種の評価 $J$ の下で $\bf u$ を色々変えてみて、いざ評価 $J$ を最も良くするような $\bf u$ が見つかったときに成り立っているべき方程式を表します。

ここで重要なポイントが2つあります。

まず、我々は勝手に 状態 $\bf x$ なるものが $\bf u$ で制御できるのだと考えている点です。例えば、ある回路に流れる電流を適当な印加電圧によって自在に制御できると、勝手に考えているのです。ではいざその回路の構成を見てみると、実は配線がちゃんとできていなくて、印加電圧をどんな風に変えてみても電流を操作することができなかったとしましょう。そうしたならば当然、制御できるわけもないのですから、ベルマン方程式だのなんだの言っても意味がありません。

次に、ベルマン方程式はある評価が最も良くなったときに成り立っている方程式であるという点です。私達は状態 $\bf x$ を制御するような外部からの入力 $\bf u $ を求めたいと考えているわけですが、当然どんな $\bf u$ が最適であるのかを知りません。しかし、「もしも最適な $\bf u$ が見つかった暁には、その $\bf u$ はベルマン方程式なる式を満たしているはずである」ということだけを知っているのです。 つまり、$\bf u$ を色々変えてみて最適なものを見つけようとしているその過程において（すなわち今保持している最適でない $\bf u$ という制御則について）、ベルマン方程式が成り立っているということではないということです。

これは言い換えれば、今持っている $\bf u$ がベルマン方程式を満たしているか、それが最適なものであるのかそうでないのかを確認できるということです。そうして、いざベルマン方程式を満たしていない $\bf u$ を使ってしまっているのだとすれば、どのようにしてそれを修正して、最適なものを見つけていけばよいのかの手段があれば良いということになります。これが強化学習におけるベルマン方程式の役割です。

最適制御と評価関数

最適制御

ここで一旦寄り道をさせてください。状態 $\bf x$ を外部からの入力 $\bf u$ によって思い通りに操作してやりたい、という問題を制御問題と呼び、工学システムが複雑化していった産業革命あたりで大きく発展してきました。その中でも、$\bf x$ を制御できるような $\bf u$ が色々考えられる場合に、その中で最も良さそうな $\bf u$ を選定する制御問題を、最適制御問題と呼びます。

ここに制御問題を簡単に見ておきます。状態 $\bf x$ に外部入力 $\bf u$ を与えたときの時間発展が下記の微分方程式に従うような場合に、

$$ \bf \dot x = Ax + Bu $$

この微分方程式の解、すなわち $\bf x$ の時間波形を $\bf u$ を色々変えてみて自在に操りたいというのが制御問題です。もしも入力 $\bf u$ を外部から与えなかった場合には、状態 $\bf x$ は微分方程式 $\bf \dot x = Ax$ という単純な式に従って勝手に時間発展していきます（これは、行列指数関数を使えば簡単に一般解が求まります）。

放っておいたら勝手な振る舞いをする $\bf x$ を $\bf u$ を使ってなんとかしたいということです。このときの「なんとかしたい」というのは、$\bf x$ の収束先を決めてやりたいというだけの場合もありますし、各時刻でとある値を取ってほしいという場合もあります。前者ならば、指定した値にどれだけ精確にどれだけ速く収束するかが評価対象になりますし、後者の場合は各時刻全体に渡ってその値を評価する必要があるかもしれません。

評価関数

ここで、自分が持ち合わせている $\bf u$ の良さを評価するために評価関数 $J$ なるものを考えます。 例えば、典型的には以下のような評価関数が使われます。

$$ J = \int _ 0 ^ \infty \left( {\bf (x _ {ref} - x )^ T Q(x _ {ref} - x) + u ^ T R u} \right) {\rm d} t $$

ここで $\bf x _ {ref}$ というのは $\bf x$ が各時刻にどういう値を取っておいてほしいかを表す目標値です。 このような評価関数では、状態 $\bf x$ と自分が定めた目標値との二乗誤差（みたいなもの）を第一項が表現しており、入力 $\bf u$ の値の大きさを第二項が表しており、式全体では、それらの時間全体での総和を表しています。

その心は「各時刻で状態は目標値に近い値を取っているべきであり、かつ、そのような状態を実現するに当たって外部からの入力はできるだけ小さな値の方が良い」というものになります。

なぜこんな評価関数が典型となっているのかというと、多くの工学的な応用においては $\bf u$ という外部入力は、電流であったり電圧であったり、外部から作用される力であったり、何かしらのエネルギーを要するものであるケースが多いからです。ですから $\bf u$ は小さいほうが良いということが組み込まれます。

また単純に二乗ではなく、 $\bf Q, R$ を間に挟んだ二次形式になっているのは、状態ベクトルや入力ベクトルのどの成分に重きを置くのか（あるいはその相関に重きをおくのか）を表現するためです。

もちろんこの評価関数は自分が達成したい制御に合わせて設計されるべきであり、この形式に限る必要はありません。入力や状態を受け取り、スカラー値で良し悪しを評価するような汎関数の形式をしていればひとまず良いということです。ですので、以降は

$$ J = \int _ 0 ^ \infty l({\bf x}, {\bf u}, t ){\rm d} t $$

ということにしておきます。（積分の中身は、なるべく小さいほうが嬉しいような値という意味を込めて、loss の頭文字 $l$ を使いました）

価値関数

ここで、上記の評価関数を別の視点で見てみることにします。 どういう視点であるかというと、実は制御が開始してから既に時刻が $t$ だけ経過しているようなケースです。過去にはもう戻れないので、時刻 $t$ 以降についてのみの評価関数を考えます。それは単純に下記のような形になるでしょう。

$$ J _ t = \int _ t ^ \infty l({\bf x}, {\bf u}, t ){\rm d} t $$

単に積分の開始が $t$ からになっているだけです。既に時刻が $t$ だけ経過したあとに、これからできる限りの努力をしようと考えた場合には上記の評価関数を考えるのが自然です（もちろん積分の中身については、問題に応じて別のものでも構いません）。

さて、既に時刻が経過しているので、本来の評価関数（積分範囲が $0 \rightarrow \infty$）に関しての最適な $\bf u$ は求められないかもしれませんが、これ以降の時刻についてのみ考えた $J _ t$ を最小化するようなことを目指そうということです。このときに、$J _ t$ の最小値が$ \bf u $ を変えてみることで見つかったとしましょう。そのような最小値を

$$ V _ t = \min _ {\bf u} J _ t $$

と書き直し、これを 価値関数 と呼ぶことにします。言い換えると価値関数とは、時刻 $t$ から開始した評価関数 $J _ t$ の 最小値そのもの、ということになります。ここで、 $t = 0$ としてみます。要するに時刻は実は一切経過していないという場合を考えるのです。すると

$$ V _ 0 = \min _ {\bf u} J _ 0 = \min _ {\bf u} J $$

となっており、時刻 $ t $ における価値関数はもともとの評価関数の最小値そのものであるということになります。

今、混乱を避けるために、価値関数 $V _ t$ が何を引数にとる関数であるのかを明示することにします。

$$ V ({\bf x}(t), t) = \min _ {\bf u} J({\bf x}(t), {\bf u}(t), t) $$

となっています。価値関数は既に $\bf u$ によって評価関数を最適化したあとなので、 もはや $\bf u$ の関数にはなっていないことに注意してください（これはすなわち、厳密な価値関数というものは、最適な $\bf u$ が求まった後にわかるということ。これが冒頭に述べた、ベルマン方程式は最適な入力が既に得られてとして、その最適入力が満たしているべき方程式である、という話につながる）。

そして $V ({\bf x}, 0)$ のケースを考えれば、これは時刻 $t$ から $\infty$ までの評価値になるということです。

ベルマンの最適性原理

さて、価値関数 $V ({\bf x}(t), t)$ とは時刻 $t$ 以降の評価関数の（$\bf u$ を調整した末の）最小値となっている関数でした。今、この関数を次のように書き直します。それは「時刻 $t$ から $t + {\rm d} t$ までの僅かな時間の価値関数と時刻 $t + {\rm d} t$ 以降の価値関数の和」というものです。

$$ \begin{align} V ({\bf x}(t), t) & = \min _ {\bf u} J({\bf x}(t), {\bf u}(t), t) \\ & = \min _ {\bf u} \int _ t ^ \infty l({\bf x}, {\bf u}, t ){\rm d} t \\ & = \min _ {\bf u} \int _ t ^ {t+{\rm d} t} l({\bf x}, {\bf u}, t ){\rm d} t + \min _ {\bf u} \int _ {t + {\rm d} t}^ \infty l({\bf x}, {\bf u}, t ){\rm d} t \\ & = \min _ {\bf u} \int _ t ^ {t+{\rm d} t} l({\bf x}, {\bf u}, t ){\rm d} t + V ({\bf x}(t + {\rm d}t ), t + {\rm d}t) \end{align} $$

ここで最後の式変形は、時刻の開始地点が $t$ から $t + {\rm d} t$ に変わっているだけであることに着目したものです。 これは一言でいえば、積分区間を途中で区切って分けただけなのですが、これはベルマンの最適性原理と呼ばれる非常に重要なが因縁に基づいています。（全体の時刻における最適な関数 $ {\bf x} (t) $があったとすれば、それをとある部分的な時刻で切り取った場合にも最適な関数となっている）

ベルマン方程式

価値関数の離散化

価値関数に関して我々は

$$ V ({\bf x}(t), t) = \min _ {\bf u} \int _ t ^ {t+{\rm d} t} l({\bf x}, {\bf u}, t ){\rm d} t + V ({\bf x}(t + {\rm d}t ), t + {\rm d}t) $$

という重要な表現を獲得しました。今時刻は連続的に変化するものとしていますが、これをあるサンプル時間のもとで $t, t+1, t+2$ と数えられる数に変形したとします。すると上記の式は

$$ V ({\bf x}(t), t) = \min _ {\bf u} l({\bf x}(t), {\bf u}(t), t ) + V ({\bf x}(t + 1 ), t + 1) $$

と近似できます（ここで、$l$ という時間関数において、離散化されたために第一項の積分が評価できない。そこで大雑把に、区間 $[t, t + {\rm d}t) $の適当な時刻における $l$ の値を代表して採用している）

以降は、離散化されていることを意識するために ${\bf x} (t)$ などを $ {\bf x} _ t$ と表記します。 また、時刻 $t$ が少々煩いので、$x _ t$ $V$ と $l$ に関しては時刻 $t$ を引数に取ることを省略する（これは、状態 $\bf x$ らが時刻を引数に取っていることを見れば分かるので問題はないはずだ）。

すなわち、ここまでで以降の価値関数を

$$ V ({\bf x} _ t) = \min _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + V ({\bf x} _ {t + 1}) $$

と書くことにした。

状態の時間発展再訪

さて、制御問題に立ち返ると、 $\bf x$ という状態は

$$ \bf \dot x = Ax + Bu $$

という時間発展をするのでした。今、これは上記のような線形システムである必要はありません。

$$ {\bf \dot x} = f({\bf x, u}) $$

という非線形な何らかの関数で表されてもいいです。そして離散化されていることを考えると、適当なサンプル時刻を使って

$$ {\bf x} _ {t+1} = g({\bf x _ t, u _ t}) $$

となっていることでしょう。また、この離散時間の非線形システムは確率システムになっていてもよく、適当な確率分布によって

$$ {\bf x} _ {t+1} \sim p({\bf x} _ {t+1} \mid {\bf x} _ {t}, {\bf u} _ {t}) $$

となっているかもしれません。ここで、 $\bf p$ は現在の状態 ${\bf x} _ t$ と現在の入力 ${\bf u} _ t $ によって次の時刻の状態 ${\bf x} _ {t + 1}$ がどのような値を取るのかを表す確率分布であり、状態遷移確率と呼ばれます。強化学習問題では通常、このような確率システムを使うことが多いので、以降は状態の時間発展が確率システムであるとします。

ベルマン方程式

強化学習におけるベルマン（最適）方程式は状態 $\bf x$ で行動 $\bf u$ と書いたときには

$$ V ({\bf x} _ t) = \max _ {\bf u} R({\bf x } _ t, {\bf u } _ t) + \gamma \max _ {\bf u} \sum _ {{\bf x} _ { t + 1}} p( {\bf x} _ {t + 1} \mid {\bf x , u}) V ({\bf x} _ {t+1}) $$

と書かれます（項を分けているのは意図的です）。 ここで $\gamma$ は割引率と呼ばれます（こいつの意義は後に説明します）。また $R$ は状態と行動の組に対して与えられるスカラー報酬関数です。

さて、これまで連続システムの制御問題として見てきた価値関数（の離散化）を振り返ります。

$$ V ({\bf x} _ t) = \min _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + V ({\bf x} _ {t + 1}) $$

というものでした。ベルマン方程式と非常に似ている形をしています。

まず、違いを見ると、 $\min$ 演算子と $\max$ 演算子の違いがあります。これは単に符号を入れ替えれば問題としては等価になるので表面上の違いしかありません。最適制御では、入力が大きくなってほしくない、とか、状態が目標値から離れてほしくない、などの「値が小さくなってほしい」ということを評価に盛り込むため $\min$ 演算が用いられます。一方強化学習では、こうなったら嬉しいという報酬を与えるために $\max$ 演算子が使われています。

次の違いは第二項についている $\gamma$ という割引率です。これも最適制御と強化学習で「悪さを評価する」か「良さを評価するか」が異なることに起因していると考えています。まず最適制御では悪さを「正の値」で評価します。すなわち、理想的な最適のときには価値関数が $0$ を取るのです。一方で強化学習では良さを評価するために、問題設定で下手をしたら無限大に値が発散してしまいます。そこで、時刻無限遠方まで総和を取ったときに、値が発散しないように $0 < \gamma < 1$ を掛けておくのです。

また、この $\gamma$ という値を導入することを肯定的に見る場合は、「現在の状態を良くすることに集中し、未来に得られる評価は予想に過ぎないので、少し重みを小さくしておく」と考えた度合いだと捉えられます。実際、強化学習では確率システムを扱うケースが多く、未来の評価値は現在から見たら（計算ができないという不都合とは関係なく）不明であり、確率的にしか見積もれないので、このような考え方をするのは合理的です。

実際、第二項の総和の式は、状態遷移確率によって、次の状態が確率的に定まるため、次の時刻における価値関数の期待値しか見れないことを表しています（確率システムでなければ、時間発展の微分方程式にしたがって次の状態は決まっているので、このような計算は不要）。

また、第二項において $\max$ 演算子が登場しているのは、「価値関数というのは$\bf u$ に関して評価関数を最適化した場合に得られる値を返す関数」であることを思い出しましょう。状態遷移確率に現れる $p$ の条件 $\bf u _ {t}$ については、その次の時刻の価値の期待値が $\bf u$ に関して最大化されるように選ばれているべきであることを表しています。

まとめ

まとめてしまえば、強化学習で出てくるベルマン（最適）方程式とは、最適制御での価値関数を離散化し、確率システムへ拡張したものであると言えます。評価関数をベルマンの最適性原理で、時間による分割を行い

$$ V ({\bf x} _ t) = \min _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + V ({\bf x} _ {t + 1}) $$

を得ます。評価の仕方は、 $\min$ でも $\max$ でも問題に応じて自分が設計するものであるので、

$$ V ({\bf x} _ t) = \max _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + V ({\bf x} _ {t + 1}) $$

と書いても良いはずです。そして、この時刻に関する再帰式で定義される価値が、正の値で和を取られ続けた場合に発散しないように $\gamma$ という定数を再帰の部分に乗算しておきます。

$$ V ({\bf x} _ t) = \max _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + \gamma V ({\bf x} _ {t + 1}) $$

更に、今扱っている状態 $\bf x$ というのは確率的に値が決まるシステムであり、

$$ {\bf x} _ {t+1} \sim p({\bf x} _ {t+1} \mid {\bf x} _ {t}, {\bf u} _ {t}) $$

と与えられるものとします。そうすると、次の時刻の価値関数は決定論的には決まっておらず、次に取りうる $x _ {t+1} $ すべてを考慮した値（すなわち期待値） $\sum _ {{\bf x} _ { t + 1}} p( {\bf x} _ {t + 1} \mid {\bf x , u}) V ({\bf x} _ {t+1}) $ としておく必要があり、

$$ V ({\bf x} _ t) = \max _ {\bf u} l({\bf x} _ t, {\bf u} _ t) + \gamma \sum _ {{\bf x} _ { t + 1}} p( {\bf x} _ {t + 1} \mid {\bf x , u}) V ({\bf x} _ {t+1}) $$

と書けます。そして、価値関数というのはもともと時刻全体に渡る状態を評価する評価関数が $\bf u$ に対して最適化されたときの値を返す関数であったことを思い出します。右辺第二項、状態遷移確率によって次の時刻の価値期待値を計算する際、現在の状態 $\bf u$ は計算結果である期待値を最大化しているものが選ばれているべきであり、

$$ V ({\bf x} _ t) = \max _ {\bf u} R({\bf x } _ t, {\bf u } _ t) + \gamma \max _ {\bf u} \sum _ {{\bf x} _ { t + 1}} p( {\bf x} _ {t + 1} \mid {\bf x , u}) V ({\bf x} _ {t+1}) $$

となります。

最後に

今回は最適制御の価値関数から強化学習のベルマン方程式を見てきました。そのために、強化学習で使われる「状態 $\bf s$ 」や「行動 $\bf a$ 」をそれぞれ $\bf x$ と $\bf u$ で書いてきました。若干見辛さはあったかもしれませんが、（最適制御の式に馴染みがある人にとっては？）、それが確率システムであるということを使って自然に拡張できたと思います。

が、本当は、連続時間の価値関数をテイラー近似してしまえば、価値関数に関する偏微分方程式「ハミルトン・ヤコビ・ベルマン方程式」と呼ばれる連続時間版のベルマン方程式が得られます。いや、実際にはハミルトン・ヤコビ・ベルマン方程式の離散時間版が、強化学習でのベルマン最適方程式になっており、本来はその対応を見る方が良さそうでしたが、都合により割愛しました（そのため、価値関数の離散化はえらくご都合主義的な変形になっている）。

この記事を書くに当たっては下記を大いに参考にしました。

https://www.jstage.jst.go.jp/article/jrsj/36/2/36_36_140/_pdf

強化学習 (機械学習プロフェッショナルシリーズ)

www.hellocybernetics.tech