\gdef\argmax{\mathop{\rm argmax}\limits}
\gdef\argmin{\mathop{\rm argmin}\limits}
本記事は東京大学の村山裕和（https://www.linkedin.com/in/裕和-村山-9b42252b1/ ）による寄稿です。
はじめに
前の自分の記事(https://zenn.dev/takuya_fukatsu/articles/74bfc9b4af5ecb )でも書いた通り、今回からiLQR（DDP）を解説します。まずはiLQRの基となるLQR、及び更にその基となるベルマン方程式から、最後はALTROというiLQRのライブラリで使用されているアルゴリズムまで、iLQRの基礎となる部分を網羅出来るように書いて行く予定です。尚、離散システムもしくは連続システムを離散化した場合について扱います。
複数回に渡りますので、どうか続編も楽しみにして頂ければと思います。今回は初回としてiLQRの基礎理論をしっかり説明します。JAXによるコーディング例も載せているので、良ければ参考にしてください。
LQRの理論
まずはLQRから説明します。無限ホライズンの所謂"最適レギュレーター"をLQRと呼ぶ事も多いですが、今回は有限ホライズンのものを説明します。
問題設定
次のような線形システムを考えます。
x_{k+1} = A x_{k} + B u_{k} \left( \triangleq f(x_{k},u_{k}) \right)
この時、以下の評価関数を最小化するような最適制御問題を考えます。
J = \phi(x_{N}) + \sum_{t=0}^{N-1} l_{t}(x_{t},u_{t})\\
\phi(x_{N}) = \frac{1}{2} x_{N}^{T} S x_{N} + s x_{N}\\
l_{t}(x_{t},u_{t}) = \frac{1}{2} x_{t}^{T} Q x_{t} + q x_{t} + \frac{1}{2} u_{t}^{T} R u_{t} + r u_{t}
この様に、線形システムに対し、二次関数で表される評価関数を最小化（もしくは最大化）するのがLQRです。
数式でLQR問題をまとめると
u_{0}, \cdots, u_{N-1} = \argmin_{u_{0}, \cdots, u_{N-1}} J \\
\text{subject to} \ x_{k+1} = A x_{k} + B u_{k}
となります。
文字について説明しておきますと、x_{t}は状態変数、u_{t}は制御入力です。S,s,Q,q,R,r（添え字省略）は評価関数の重みと呼ばれるハイパーパラメータであり、人間が決める定数になります。また、Nはホライズンと呼ばれる定数であり、「どれだけ先までを最適化するか」を決定する定数です。要するに、Nを大きくすれば将来を長く見据えて最適化する事になり、逆もまた然りです。これも人間が決めるハイパーパラメータです。tは時刻ステップを表す変数なのですが、絶対的な時刻ではなく、「現在時刻をt=0とした相対的な時刻」です。この問題設定では、時刻t=0からt=Nまでの間で評価関数を最適化する事になります。モデル予測制御としてLQRを使用する場合、x_{0}を現在の状態変数とし、u_{0}を実際の制御入力として使用するのを繰り返す事になります。
ベルマン方程式
天下り的ではありますが、次のような関数を導入します。
V_{k}(x_{k}) = \min_{u_{k},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right]
このV_{k}(x_{k})は強化学習の文脈で状態価値関数と呼ばれる関数です。時刻ステップが t=k の時に状態変数が x_{k} だった場合、評価関数内でそれ以降の部分の最小値がどうなるのかを表す関数になります。
次に、以下の様な関数を導入します。
Q_{k}(x_{k},u_{k}) = \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right]
こちらの Q_{k}(x_{k},u_{k}) は行動価値関数と呼ばれる関数です。時刻ステップが t=k の時に状態変数が x_{k} だった場合、制御入力を u_{k} した時、評価関数内でそれ以降の部分の最小値がどうなるのかを表す関数になります。
それでは、この二つからベルマン方程式を導いて行きます。まずは状態価値関数 V(x) だけのバージョンを導出します。
\begin{align}
V_{k}(x_{k}) &= \min_{u_{k},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right]\\
&= \min_{u_{k}} \left[ \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right] \right]\\
&= \min_{u_{k}} \left[ l_{k}(x_{k},u_{k}) + \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k+1}^{N-1} l_{t}(x_{t},u_{t}) \right] \right]\\
&= \min_{u_{k}} \left[ l_{k}(x_{k},u_{k}) + V_{k+1}(f(x_{k},u_{k})) \right]
\end{align}
というわけで、以下の様な方程式が出てきました。
V_{k}(x_{k}) = \min_{u_{k}} \left[ l_{k}(x_{k},u_{k}) + V_{k+1}(f(x_{k},u_{k})) \right]
これが状態価値関数で表したベルマン方程式になります。
また、状態価値関数と行動価値関数を両方使って表す事も出来ます。
\begin{align}
V_{k}(x_{k}) &= \min_{u_{k},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right]\\
&= \min_{u_{k}} \left[ \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right] \right]\\
&= \min_{u_{k}} Q_{k}(x_{k},u_{k})
\end{align}
行動価値関数も以下の様に変形できます。
\begin{align}
Q_{k}(x_{k},u_{k}) &= \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k}^{N-1} l_{t}(x_{t},u_{t}) \right]\\
&= l_{k}(x_{k},u_{k}) + \min_{u_{k+1},\cdots, u_{N-1}} \left[ \phi(x_{N}) + \sum_{t=k+1}^{N-1} l_{t}(x_{t},u_{t}) \right]\\
&= l_{k}(x_{k},u_{k}) + V_{k+1}(f(x_{k},u_{k}))
\end{align}
これらをまとめると、次の様になります。
\begin{align}
Q_{k}(x_{k},u_{k}) &= l_{k}(x_{k},u_{k}) + V_{k+1}(f(x_{k},u_{k}))\\
V_{k}(x_{k}) &= \min_{u_{k}} Q_{k}(x_{k},u_{k})
\end{align}
表し方が違うだけでこちらも同じベルマン方程式です。
LQRとベルマン方程式
LQRは超雑に言えば評価関数を最小化する問題です。つまるところ、状態価値関数 V_{0}(x_{0}) を求める問題です。ここから数式で考えてみましょう。
\begin{align}
\min_{u_{0}, \cdots, u_{N-1}} J &= V_{0}(x_{0})\\
&= \min_{u_{0}} Q_{0}(x_{0},u_{0})\\
&= \min_{u_{0}} \left[ l_{0}(x_{0},u_{0}) + V_{1}(f(x_{0},u_{0})) \right]\\
&= \min_{u_{0}} \left[ l_{0}(x_{0},u_{0}) + \min_{u_{1}} Q_{1}(x_{1},u_{1}) \right]\\
&= \min_{u_{0}} \left[ l_{0}(x_{0},u_{0}) + \min_{u_{1}} \left[ l_{1}(x_{1},u_{1}) + V_{2}(f(x_{1},u_{1})) \right] \right]\\
&\ \ \vdots \\
&= \min_{u_{0}} \left[ l_{0}(x_{0},u_{0}) + \min_{u_{1}} \left[ l_{1}(x_{1},u_{1}) \cdots \min_{u_{N-1}} \left[ l_{N-1}(x_{N-1},u_{N-1}) + V_{N}(x_{N}) \right] \cdots \right] \right]
\end{align}
この流れを言葉で説明すると次のようになります。
V_{0} を求めたい
↓
Q_{0} を最小化すればよい（その過程で u_{0} が求まる）
↓
V_{1} を求めておく必要がある
↓
Q_{1} を最小化すればよい（その過程で u_{1} が求まる）
↓
V_{2} を求めておく必要がある
↓
（この流れが続く）
↓
V_{N} を求めておく必要がある
つまるところ、 V_{N} からスタートしてこの流れを逆向きに辿って行けば最適化問題が解けることになります。これはベルマン方程式を t=N から t=0 に向かって逆向きに解いて行く事に相当します。
LQRのアルゴリズム
では、実際に最適化問題を解いてみましょう。まずは V_{N}(x_{N}) からスタートします。当然、
V_{N}(x_{N}) = \frac{1}{2} x_{N}^{T} S x_{N} + s x_{N}
となります。これを使って Q_{N-1}(x_{N-1},u_{N-1}) の最小値を求めます。u_{N-1} で微分して0となれば良いので
\begin{align}
\frac{\partial Q_{N-1}}{\partial u_{N-1}} &=
\frac{\partial l_{N-1}}{\partial u_{N-1}} + \frac{\partial V_{N}}{\partial u_{N-1}}\\
& = u_{N-1}^{T}R + r + \frac{\partial V_{N}}{\partial x_{N}} \frac{\partial x_{N}}{\partial u_{N-1}}\\
& = u_{N-1}^{T}R + r + \left( x_{N}^{T}S + s \right) \frac{\partial f}{\partial u_{N-1}}\\
& = u_{N-1}^{T}R + r + \left( x_{N}^{T}S + s \right) B\\
& = u_{N-1}^{T}R + r + \left\{ \left(Ax_{N-1}+Bu_{N-1}\right)^{T}S + s \right\} B\\
& = u_{N-1}^{T} \left(R + B^{T}SB \right) + x_{N-1}^{T}A^{T}SB + r + sB\\
\\
\therefore 0 &= \left(R + B^{T}SB \right) u_{N-1} + B^{T}SAx_{N-1} + r^{T} + B^{T}s^{T}\\
\therefore u_{N-1} &= - \left(R + B^{T}SB \right)^{-1} B^{T}SAx_{N-1}
- \left(R + B^{T}SB \right)^{-1} \left(r^{T} + B^{T}s^{T} \right)\\
& = K_{N-1}x_{N-1} + d_{N-1}\\
K_{N-1} & \triangleq - \left(R + B^{T}SB \right)^{-1} B^{T}SA\\
d_{N-1} & \triangleq - \left(R + B^{T}SB \right)^{-1} \left(r^{T} + B^{T}s^{T} \right)
\end{align}
となります。x_{N-1} の値は u_{0} から u_{N-2} が求まらないと決まらないので、ここでは K_{N-1} と d_{N-1} の値を記録しておきます。
では、求まった u_{N-1} の式を代入して V_{N-1}(x_{N-1}) を求めます。
\begin{align}
V_{N-1}(x_{N-1}) &= l_{N-1}(x_{N-1},u_{N-1}) + V_{N}(f(x_{N-1},u_{N-1}))\\
& = \frac{1}{2}x_{N-1}^{T} Q x_{N-1} + q x_{N-1}
+ \frac{1}{2}u_{N-1}^{T} R u_{N-1} + r u_{N-1}\\
& \; \; \; \; + \frac{1}{2} \left(Ax_{N-1} + Bu_{N-1} \right)^{T} S \left(Ax_{N-1} + Bu_{N-1} \right) + s\left(Ax_{N-1} + Bu_{N-1} \right)\\
&= \frac{1}{2}x_{N-1}^{T} \left\{ Q + K_{N-1}^{T}(R+B^{T}SB)K_{N-1} + A^{T}SA + A^{T}SBK_{N-1} + K_{N-1}^{T}B^{T}SA \right\} x_{N-1}\\
& \; \; \; \; + \left\{q + (d_{N-1}^{T}R+r)K_{N-1} + (d_{N-1}^{T}B^{T}S+s)(A+BK_{N-1}) \right\} x_{N-1} + \frac{1}{2} d_{N-1}^{T} \left(R + B^{T}S B \right) d_{N-1}
+ sBd_{N-1}\\
& = \frac{1}{2}x_{N-1}^{T}S_{N-1}x_{N-1} + s_{N-1}x_{N-1} + \text{Const.}\\
S_{N-1} & \triangleq Q + K_{N-1}^{T}(R+B^{T}SB)K_{N-1} + A^{T}SA + A^{T}SBK_{N-1} + K_{N-1}^{T}B^{T}SA\\
s_{N-1} & \triangleq q + (d_{N-1}^{T}R+r)K_{N-1} + (d_{N-1}^{T}B^{T}S+s)(A+BK_{N-1})
\end{align}
x_{N-1} に寄らない部分は Const. としてまとめました。これで式の形の上では u_{N}, V_{N-1} が求まりました。つまり、解析的にはベルマン方程式
\begin{align}
Q_{N-1}(x_{N-1},u_{N-1}) &= l_{N-1}(x_{N-1},u_{N-1}) + V_{N}(f(x_{N-1},u_{N-1}))\\
V_{N-1}(x_{N-1}) &= \min_{u_{N-1}} Q_{N-1}(x_{N-1},u_{N-1})
\end{align}
が解けた事になります。ここでのポイントは V_{N-1} が V_{N} 同様に二次形式で表せる事です。なので全く同じプロセスで Q_{N-2} を最小化する事で次は K_{N-2}, d_{N-2} が求まります。
この流れを繰り返して行けば、K_{N-1}, \cdots, K_{0} 及び d_{N-1}, \cdots, d_{0} が求まります。
K_{N-1}, \cdots, K_{0} と d_{N-1}, \cdots, d_{0} が求まったら、次は u_{0}, \cdots, u_{N-1} と x_{1}, \cdots, x_{N} を計算して行きます。x_{0} は現在時刻の状態変数なので、既知の値として扱います。ここでのアルゴリズムはただ次の様に計算して行くだけです。
\begin{align}
u_{0} &= K_{0} x_{0} + d_{0}\\
x_{1} &= A x_{0} + B u_{0}\\
u_{1} &= K_{1} x_{1} + d_{1}\\
& \vdots \\
u_{N-1} &= K_{N-1} x_{N-1} + d_{N-1}\\
x_{N} &= A x_{N-1} + B u_{N-1}
\end{align}
求めておいた K,d と状態方程式から u と x を計算して行っているだけです。
これで晴れて u_{0}, \cdots, u_{N-1} と x_{0}, \cdots, x_{N} が求まりました。次はこのLQRを非線形に拡張します。
iLQRの理論
iLQR は iterative LQR の略であり、簡単に説明すれば「解の周辺で線形化→線形近似出来る微小範囲でLQR」を反復する手法です。
問題設定
次のような非線形システムを考えます。
x_{k+1} = f_{k}(x_{k},u_{k})
この時、以下の評価関数を最小化するような最適制御問題を考えます。
J = \phi(x_{N}) + \sum_{t=0}^{N-1} l_{t}(x_{t},u_{t})\\
LQRと違い、\phi と l_{t} は必ずしも二次形式になっている必要はありません。
iLQRの考え方
何らかの初期解 \bar{x}\_{0}, \cdots, \bar{x}\_{N} 及び \bar{u}\_{0}, \cdots, \bar{u}\_{N-1} があるとします。この初期解は恐らく前の制御ステップの最適解を使う事が多いかと思います。制御入力の初期解は全部0、状態変数は変化無し、つまり \bar{x}\_{0}=\cdots=\bar{x}\_{N} と設定してしまっても良いと思います。
まず、初期解周辺で状態方程式を線形化します。正確には、次の様にテーラー展開の一次までの変分を考えます。
\delta x_{k+1} = A_{k} \delta x_{k} + B_{k} \delta u_{k} \left( \triangleq \delta f_{k}(\delta x_{k}, \delta u_{k}) \right)\\
A_{k} \triangleq \frac{\partial f_{k}}{\partial x_{k}}, \ B_{k} \triangleq \frac{\partial f_{k}}{\partial u_{k}}
この線形近似が成り立つような微小区間での最適化を考えます。
LQRの時と同様にベルマン方程式を導入します。
\begin{align}
Q_{k}(x_{k},u_{k}) &= l_{k}(x_{k},u_{k}) + V_{k+1}(f(x_{k},u_{k}))\\
V_{k}(x_{k}) &= \min_{u_{k}} Q_{k}(x_{k},u_{k})
\end{align}
このベルマン方程式を後ろの時刻から解いて行くのはLQRと同じです。しかし、システムが非線形なので、このままLQRの様に解くことは出来ません。そこで、状態方程式同様に、行動価値関数と状態価値関数も初期解 \bar{x}\_{0}, \cdots, \bar{x}\_{N} 及び \bar{u}\_{0}, \cdots, \bar{u}\_{N-1} 付近で近似して解きます。次の様に、二次までのテーラー展開で近似します。
\begin{align}
Q(x,u) &\fallingdotseq Q(\bar{x}, \bar{u}) + \delta Q(\delta x, \delta u)\\
&= Q(\bar{x}, \bar{u}) + \frac{1}{2}\delta x^{T} \ Q_{xx} \delta x + Q_{x} \delta x
+ \frac{1}{2} \delta u^{T} \ Q_{uu} \delta u + Q_{u} \delta u
+ \delta x^{T} \ Q_{xu} \delta u + \text{Const.}\\
V(x) &\fallingdotseq V(\bar{x}) + \delta V(\delta x) \\
&= V(\bar{x}) + \frac{1}{2}\delta x^{T} \ V_{xx} \delta x + V_{x} \delta x + \text{Const.}
\end{align}
Const.の項は後で正体が明らかになります。今は気にしないでください。ここから、Q_{k} の微小変化\delta Q_{k} を最小化するように u_{k} の微小変化 \delta u_{k} を求めると言う風に捉え方を変えます。無理やり表現すると、次のようなベルマン方程式を考えることになります（ただ変分で表しているだけで、式としての意味は同じです）。
\begin{align}
Q_{k}(\bar{x}_{k},\bar{u}_{k}) + \delta Q_{k}(\delta x_{k}, \delta u_{k}) &= l_{k}(\bar{x}_{k},\bar{u}_{k}) + \delta l_{k}(\delta x_{k}, \delta u_{k}) + V_{k+1}(f(\bar{x}_{k},\bar{u}_{k})) + \delta V_{k+1}(\delta f_{k}(\delta x_{k}, \delta u_{k}))\\
V_{k}(\bar{x}_{k}) + \delta V_{k}(\delta x_{k}) &= \min_{\delta u_{k}} Q_{k}(\bar{x}_{k},\bar{u}_{k}) + \delta Q_{k}(\delta x_{k}, \delta u_{k})
\end{align}
なお、各ヘッシアン及び偏微分は次の様に計算できます。
\begin{align}
Q_{k,uu} &= l_{k,uu}
+ \left( \frac{\partial x_{k+1}}{\partial u_{k}} \right)^{T} V_{k+1,xx} \frac{\partial x_{k+1}}{\partial u_{k}}\\
&= l_{k,uu} + B_{k}^{T} V_{k+1,xx} B_{k} \\
Q_{k,u} &= l_{k,u} + V_{k+1,x} \frac{\partial x_{k+1}}{\partial u_{k}}\\
&= l_{k,u} + V_{k+1,x} B_{k}\\
Q_{k,xx} &= l_{k,xx}
+ \left( \frac{\partial x_{k+1}}{\partial x_{k}} \right)^{T} V_{k+1,xx} \frac{\partial x_{k+1}}{\partial x_{k}}\\
&= l_{k,xx} + A_{k}^{T} V_{k+1,xx} A_{k} \\
Q_{k,x} &= l_{k,x} + V_{k+1,x} \frac{\partial x_{k+1}}{\partial x_{k}}\\
&= l_{k,x} + V_{k+1,x} A_{k}\\
Q_{k,xu} &= l_{k,xu}
+ \left( \frac{\partial x_{k+1}}{\partial x_{x}} \right)^{T} V_{k+1,xx} \frac{\partial x_{k+1}}{\partial u_{x}}\\
&= l_{k,xu} + A_{k}^{T} V_{k+1,xx} B_{k}\\
Q_{k,ux} &= Q_{k,xu}^{T}
\end{align}
ここから、\partial Q_{k}/ \partial \delta u_{k} = 0 となる様に \delta u_{k} を求めます。次の様になります。
\begin{align}
\frac{\partial \delta Q_{k}}{\partial \delta u_{k}} &=
\delta u_{k}^{T} \ Q_{k,uu} + Q_{k,u} + \delta x_{k}^{T} \ Q_{k,xu}\\
\\
\therefore 0 &= Q_{k,uu} \delta u_{k} + Q_{k,u}^{T} + Q_{k,ux} \delta x_{k}\\
\therefore \delta u_{k} &= - Q_{k,uu}^{-1} Q_{k,ux} \delta x_{k} - Q_{k,uu}^{-1}Q_{k,u}^{T}\\
&= K_{k} \delta x_{k} + d_{k}\\
K_{k} & \triangleq - Q_{k,uu}^{-1} Q_{k,ux}\\
d_{k} & \triangleq - Q_{k,uu}^{-1}Q_{k,u}^{T}
\end{align}
これで \delta u_{k} は数式としては求まりました。これを \delta Q_{k} の式に代入すれば、 \delta V_{k} が求まります。
\begin{align}
\delta Q_{k} &= \frac{1}{2}\delta x_{k}^{T} \ Q_{k,xx} \delta x_{k}
+ Q_{k,x} \delta x_{k} + \frac{1}{2} (K_{k} \delta x_{k} + d_{k})^{T} \ Q_{k,uu} (K_{k} \delta x_{k} + d_{k})\\
& \; \; \; + Q_{k,u} (K_{k} \delta x_{k} + d_{k})
+ \delta x_{k}^{T} \ Q_{k,xu} (K_{k} \delta x_{k} + d_{k})\\
&= \frac{1}{2} \delta x_{k}^{T} \left( Q_{k,xx} +K_{k}^{T} Q_{k,uu} K_{k} +K_{k}^{T}Q_{k,ux} + Q_{k,xu}K_{k} \right) \delta x_{k}\\
& \; \; \; + \left( Q_{k,x} + d_{k}^{T} Q_{k,uu}K_{k} + Q_{k,u}K_{k} + d_{k}^{T} Q_{k,ux} \right) \delta x_{k} + \frac{1}{2} d_{k}^{T} Q_{k,uu} d_{k} + Q_{k,u}d_{k}
\end{align}
以上より、\delta Q_{k} の最小値の内、\delta x_{k} によらない部分も \delta V_{k} にまとめてしまえば、
\begin{align}
\delta V_{k}(\delta x_{k}) &= \frac{1}{2} \delta x_{k}^{T} V_{k,xx} \delta x_{k} + V_{k,x} \delta x_{k} + \text{Const.}\\
V_{k,xx} &= Q_{k,xx} +K_{k}^{T} Q_{k,uu} K_{k} +K_{k}^{T}Q_{k,ux} + Q_{k,xu}K_{k}\\
V_{k,x} &= Q_{k,x} + d_{k}^{T} Q_{k,uu}K_{k} + Q_{k,u}K_{k} + d_{k}^{T} Q_{k,ux}\\
\text{Const.} &= \frac{1}{2} d_{k}^{T} Q_{k,uu} d_{k} + Q_{k,u}d_{k}
\end{align}
となります。これで \delta u_{k}, \delta V_{k} が式の形の上では求まったので、ベルマン方程式が微小区間内で解けたことになります。このConst.が最初に出て来たConst.の正体です。
LQR同様、\delta u_{k} の実際の値は \delta x_{k} の値がわかるまで計算できません。なので K_{k}, d_{k} の値を記録して、後で \delta x_{k} が求まってから \delta u_{k} の値を計算します。加えて、勿論これは微小変化内での最適解なので、完全な最適化ではありません。何回もこの流れを繰り返す事でやっと最適解に収束します。
Const. の部分を \delta V_{k} に含めても含めなくても、次に \delta Q_{k-1} の最小値を考える際、Const. は \delta u_{k-1} で偏微分して 0 になるので計算結果は変わりません。ですが、今後の説明の為にここでは含める物とします。
Const. の部分を \delta V_{k} に含める場合、ベルマン方程式
\begin{align}
\delta Q_{k}(\delta x_{k}, \delta u_{k}) &= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \delta V_{k+1}(\delta f_{k}(\delta x_{k}, \delta u_{k}))\\
\delta V_{k}(\delta x_{k}) &= \min_{\delta u_{k}} \delta Q_{k}(\delta x_{k}, \delta u_{k})
\end{align}
を解いたのと同じ事になります。
LQRの視点から
前節の最後に登場した ベルマン方程式を確認します。
\begin{align}
\delta Q_{k}(\delta x_{k}, \delta u_{k}) &= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \delta V_{k+1}(\delta f_{k}(\delta x_{k}, \delta u_{k}))\\
\delta V_{k}(\delta x_{k}) &= \min_{\delta u_{k}} \delta Q_{k}(\delta x_{k}, \delta u_{k})
\end{align}
ここから、\delta Q_{k} と \delta V_{k} を書き下すと、次のようになります。
\begin{align}
\delta Q_{k}(\delta x_{k}, \delta u_{k}) &= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \delta V_{k+1}(\delta x_{k+1})\\
&= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \min_{\delta u_{k+1}} \delta Q_{k+1}(\delta x_{k+1}, \delta u_{k+1}) \\
&= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \min_{\delta u_{k+1}} \left[ \delta l_{k+1}(\delta x_{k+1}, \delta u_{k+1}) + \delta V_{k+2}(\delta x_{k+2}) \right] \\
&= \delta l_{k}(\delta x_{k}, \delta u_{k}) + \min_{\delta u_{k+1}} \left[ \delta l_{k+1}(\delta x_{k+1}, \delta u_{k+1}) + \min_{\delta u_{k+2}} \left[ Q_{k+2}(\delta x_{k+2}, \delta u_{k+2}) \right] \right]\\
& \vdots \\
& = \min_{\delta u_{k+1},\cdots, \delta u_{N-1}} \left[ \delta V_{N}(\delta x_{N}) + \sum_{t=k}^{N-1} \delta l_{t}(\delta x_{t}, \delta u_{t}) \right] \\
& = \min_{\delta u_{k+1},\cdots, \delta u_{N-1}} \left[ \delta \phi(\delta x_{N}) + \sum_{t=k}^{N-1} \delta l_{t}(\delta x_{t}, \delta u_{t}) \right] \\
\\
\delta V_{k}(\delta x_{k}) &= \min_{\delta u_{k}} \delta Q_{k}(\delta x_{k}, \delta u_{k}) \\
&= \min_{\delta u_{k},\cdots, \delta u_{N-1}} \left[ \delta \phi(\delta x_{N}) + \sum_{t=k}^{N-1} \delta l_{t}(\delta x_{t}, \delta u_{t}) \right]
\end{align}
これらは、 初期解周辺で線形化したシステム
\delta x_{k+1} = A_{k} \delta x_{k} + B_{k} \delta u_{k}
に対し、評価関数の微小変化
\delta J = \frac{1}{2} \delta x_{N}^{T} \phi_{xx} \delta x_{N} + \phi_{x} \delta x_{N}
+ \sum_{t=0}^{N-1} \frac{1}{2} \delta x_{t}^{T} l_{t,xx} \delta x_{t} + l_{t,x} \delta x_{t}
+ \frac{1}{2} \delta u_{t}^{T} l_{t,uu} \delta u_{t} + l_{t,u} \delta u_{t}
+ \delta x_{t}^{T} l_{t,xu} \delta u_{t}
を最小化する最適制御問題での行動価値関数及び状態価値関数に当たるのがわかるでしょうか。
つまり、前節で説明したベルマン方程式の解き方は、システムを線形近似可能な微小変化の世界で、評価関数の微小変化を最小化するLQR問題を解いているのと同じという事が確認できます。この微小変化の世界でのLQRを反復（iteration）するからiLQR(iterative LQR)なのです。
iLQRのアルゴリズム
アルゴリズム全体
まずアルゴリズムの全体像を説明します。上で説明した様に「解付近で線形近似→微小範囲でLQR」を繰り返す事で評価関数を少しずつ最適化するのがiLQRです。微小範囲内のLQRには、ベルマン方程式を後ろから解いて行く Backward Pass と、その計算結果から新たな解（少し最適化された解）を求める Forward Pass という二つのプロセスがあります。評価関数が収束するまで、LQR、つまりこの二つのプロセスを繰り返します。
【iLQR全体】
評価関数の値が変化しなくなるまで、Backward Pass と Forward Pass を繰り返す
【Backward Pass】
ベルマン方程式を後ろから解いて、K_{k}, d_{k}, \Delta V_{k}^{1}, \Delta V_{k}^{2} を記録して行く
【Forward Pass】
K_{k}, d_{k} を使って解を更新する
直線探索により更新具合を調整する
Backward Pass
ベルマン方程式を後ろの時刻から、微小変化内で解いて行きます。つまり、以下のベルマン方程式を解くのがスタートです。
\begin{align}
\delta Q_{N-1}(\delta x_{N-1}, \delta u_{N-1}) &= \delta l_{N-1}(\delta x_{N-1}, \delta u_{N-1}) + \delta V_{N}(\delta f_{N-1}(\delta x_{N-1}, \delta u_{N-1}))\\
\delta V_{N-1}(\delta x_{N-1}) &= \min_{\delta u_{N-1}} \delta Q_{N-1}(\delta x_{N-1}, \delta u_{N-1})
\end{align}
尚、V_{N}=\phi(x_{N}) なので、V_{N,xx}, V_{N,x} は次の様になります。
V_{N,xx} = \phi_{xx}\\
V_{N,x} = \phi_{x}
これらの値を用いて、前節で説明した通りベルマン方程式を解けば \delta u_{N-1} 及び V_{N-1,xx}, V_{N-1,x} が次の様に求まります。
\begin{align}
\delta u_{N-1} &= - Q_{N-1,uu}^{-1} Q_{N-1,ux} \delta x_{N-1} - Q_{N-1,uu}^{-1}Q_{N-1,u}^{T}\\
&= K_{N-1} \delta x_{N-1} + d_{N-1}\\
V_{N-1,xx} &= Q_{N-1,xx} +K_{N-1}^{T} Q_{N-1,uu} K_{N-1} +K_{N-1}^{T}Q_{N-1,ux} + Q_{N-1,xu}K_{N-1}\\
V_{N-1,x} &= Q_{N-1,x} + d_{N-1} Q_{N-1,uu}K_{N-1} + Q_{N-1,u}K_{N-1} + d_{N-1}Q_{N-1,ux}
\end{align}
まだ \delta x_{N-1} の値はわからないので、K_{N-1}, d_{N-1} の値を記録しておきます。
次に、この V_{N-1,xx}, V_{N-1,x} の値を用いて次の時刻のベルマン方程式
\begin{align}
\delta Q_{N-2}(\delta x_{N-2}, \delta u_{N-2}) &= \delta l_{N-2}(\delta x_{N-2}, \delta u_{N-2}) + \delta V_{N-1}(\delta f_{N-2}(\delta x_{N-2}, \delta u_{N-2}))\\
\delta V_{N-2}(\delta x_{N-2}) &= \min_{\delta u_{N-2}} \delta Q_{N-2}(\delta x_{N-2}, \delta u_{N-2})
\end{align}
を解きます。この流れを時刻 t=0、つまり K_{0}, d_{0} が求まるまで繰り返します。
この様に、後ろの時刻からベルマン方程式を解いて行き、K_{k}, d_{k} を計算して行くプロセスを Backward Pass と呼びます。
【補足】
Backward Pass において、次二つの値も K_{k}, d_{k} と同様に記録しておきます。これは直線探索をする際に使います。直線探索については後程説明します。
\Delta V_{k}^{1} = Q_{k,u}d_{k}\\
\Delta V_{k}^{2} = \frac{1}{2} d_{k}^{T} Q_{k,uu} d_{k}
Forward Pass
K_{k}, d_{k} が全て求まったら、\delta x_{k}, \delta u_{k} を求めて行きます。つまり、新しい解を求め行きます。ただし、\delta x_{0}=0 です。順番に次の計算をするだけです。
\begin{align}
\delta u_{0} &= K_{0} \delta x_{0} + d_{0}\\
\delta x_{1} &= A_{0} \delta x_{0} + B_{0} \delta u_{0}\\
\delta u_{1} &= K_{1} \delta x_{1} + d_{1}\\
&\vdots\\
\delta u_{N-1} &= K_{N-1} \delta x_{N-1} + d_{N-1}\\
\delta x_{N} &= A_{N-1} \delta x_{N-1} + B_{N-1} \delta u_{N-1}
\end{align}
後は
x_{k} \leftarrow \bar{x}_{k} + \delta x_{k}\\
u_{k} \leftarrow \bar{u}_{k} + \delta u_{k}
とすれば新しい解が求まります。この流れを Forward Pass と呼びます。
これで終わりでも良いのですが、最適化の精度や効率を良くするために直線探索というプロセスを混ぜる場合が多いです。本当は直線探索も Forward Pass の中に入るのですが、説明が長くなるので敢えて分けて説明します。
直線探索ありの Forward Pass
iLQRでやっているのは、行動価値関数を二次形式で近似して、その最小値となる点を求めるというプロセスです。簡単な例で言えば、曲線を二次関数で近似して、その最小値となる点を求めるという事です。分かりやすい図があったので引用します。
出典：RL — LQR & iLQR Linear Quadratic Regulator
現在の解 guess 付近で曲線 original function を二次関数で近似 (quadrtic approximation) し、その二次関数が最小となるような変数 x の値を次の解 next guess とします。これを繰り返せば、いずれ最適解に到達する事が期待されます。
しかし、場合によっては以下の図の様な場合が発生します。
出典：RL — LQR & iLQR Linear Quadratic Regulator
黒い点が現在の解 guess だとすれば、次の解 next guess は二次関数 quadratic approximation 上の赤い点（overshootと書かれている点）になります。しかし、これは最適解を通り過ぎてしまっています。そこで、解の変化量 next guess - guess を小さく加工する事で、この現象を是正する必要があります。そのプロセスが直線探索 (line search) です。
直線探索の手法を説明します。手法は複数あるのですが、今回は AL-iLQR tutorial で提示されている物を説明します。
まず、パラメーター \alpha を設定します。このパラメータを使って、Forward Pass の計算を次の様に変えます。
\begin{align}
\delta u_{0} &= K_{0} \delta x_{0} + \alpha d_{0}\\
\delta x_{1} &= A_{0} \delta x_{0} + B_{0} \delta u_{0}\\
\delta u_{1} &= K_{1} \delta x_{1} + \alpha d_{1}\\
&\vdots\\
\delta u_{N-1} &= K_{N-1} \delta x_{N-1} + \alpha d_{N-1}\\
\delta x_{N} &= A_{N-1} \delta x_{N-1} + B_{N-1} \delta u_{N-1}
\end{align}
d_{k} の部分に \alpha が掛かっているのがわかるでしょうか。ここから
x_{k} \leftarrow \bar{x}_{k} + \delta x_{k}\\
u_{k} \leftarrow \bar{u}_{k} + \delta u_{k}
として新しい解を求めます。この新しい解 x_{k},u_{k} による評価関数値を J、古い解 \bar{x}_{k},\bar{u}_{k} による評価関数値を \bar{J} とします。この上で、
10^{-4} \leqq z \leqq 10\\
z = \frac{J-\bar{J}}{-\Delta V(\alpha)}\\
\Delta V(\alpha) = \sum_{k=0}^{N-1} \frac{1}{2} \alpha^{2} d_{k}^{T} Q_{k,uu} d_{k} + \alpha Q_{k,u}d_{k}
となる様に \alpha の値を調整します。最初は \alpha = 1 として計算して、上の条件を満たさなければ \alpha を半分にして計算し直す、また条件を満たさなければ \alpha を半分にして...と言うのを繰り返します。尚、\Delta V(\alpha) を求めるのに \Delta V_{k}^{1}, \Delta V_{k}^{2} を使います。
ただ、なぜこういう手法になるのかが書かれた書籍や文献が見つかっていない状態です。この直線探索手法に対する自分なりの解釈を Appendix に書いておきましたので、よろしければそちらも読んでみてください。
まとめ
アルゴリズムの見通しが良くなるよう、疑似コードにまとめました。自分でコードを書いてみる際など、是非参考にしてください。
最後に
次回は制約条件を考慮する手法について説明します。その次の記事でいよいよiLQRのライブラリであるALTROのアルゴリズムについて解説しようと思います。
参考文献
スペシャルサンクス
サンプルコードを作成するにあたり、同じく弊社アルバイトである橋本尚典（https://researchmap.jp/Takanori_Hashimoto）から以前教えて頂いたコードが大いに役立ちました。感謝の意味を込め、ここに記させて頂きます。
Appendix
直線探索について
記事内で説明した直線探索手法について、恐らく次の様な疑問が浮かぶと思います。
- 何故 \delta u = \alpha (K \delta x + d) ではなく\delta u = K \delta x + \alpha d なのか
この疑問に対する答えを自分なりに考えましたので、記しておきます。尚、これは現状裏どりが無い解釈であり、本来の解釈とは違う可能性もある事をご了承下さい。
閉ループ項 K \delta x は、状態変数が変化した事による制御入力の変化を表した項です。つまり、制御入力の初期解を真の最適解に直す項というより、状態変数 x が \delta x だけ変化した時、制御入力の最適解がどれだけ変化するかを示す項です。要するに x の影響を示すだけの項であり、そもそも制御入力が最適解からズレていた時に修正してくれる項ではありません。少し考えてみましょう。仮に状態変数 x は最適解だが u は最適解からズレている場合、iLQRに則って最適化計算をしても \delta x=0 となるはずですから、 K \delta x = 0 となります。K \delta x は u を更新してくれませんよね。
じゃあ u を最適解へ収束させてくれるのは何かというと、それが d です。直線探索は最適解では無い所から最適解へ向けて修正する量を調整するというものなので、この d だけ調整しているというわけです。
直線探索終了の条件ですが、これは状態変数のズレ \delta x による影響と、そもそも解を最適解に近づける項 d の影響のバランスが偏り過ぎないようにするというものだと解釈しています。ただ、何故そこのバランスが偏ったらいけないのかまでは良く分かりませんでした。
サンプルコード
参考までに、iLQRのサンプルコードをJAXで作ってみました。何かおかしな所が判明しましたら、コメント等で教えて頂けると助かります。
import jax
import jax.numpy as jnp
import numpy as np
import math
import time
from dataclasses import dataclass
@dataclass
class Model:
R = 0.05
T = 0.2
r = R/2
rT = r/T
model = Model()
@jax.jit
def model_func_risan(x, u, dt):
cos_ = jnp.cos(x[2])
sin_ = jnp.sin(x[2])
dx = jnp.array([model.r * cos_ * (u[0]+u[1]),
model.r * sin_ * (u[0]+u[1]),
model.rT * (u[0]-u[1])])
x_next = x + dx * dt
return x_next
model_dfdx = jax.jit(jax.jacfwd(model_func_risan,0))
model_dfdu = jax.jit(jax.jacfwd(model_func_risan,1))
@dataclass
class Cont_Args:
Ts = 0.1
tf = 1.0
N = int(tf/Ts)
dt = Ts
iter = 10
torelance = 1.0
Q = 100 * jnp.array([[1, 0, 0],
[0, 1, 0],
[0, 0, 0]], dtype=jnp.float32)
R = jnp.eye(2, dtype=jnp.float32)
S = 100 * jnp.array([[1, 0, 0],
[0, 1, 0],
[0, 0, 0]], dtype=jnp.float32)
x_ob = jnp.array([10, 5, 0], dtype=jnp.float32)
u_ob = jnp.array([0, 0], dtype=jnp.float32)
len_x = 3
len_u = 2
x = None
u = None
us = None
args = Cont_Args()
@jax.jit
def stage_cost(x,u):
cost = 0.5 * ( (x-args.x_ob) @ args.Q @ (x-args.x_ob) \
+ (u-args.u_ob) @ args.R @ (u-args.u_ob))
return cost
@jax.jit
def term_cost(x):
cost = 0.5 * (x-args.x_ob) @ args.S @ (x-args.x_ob)
return cost
grad_x_stage = jax.jit(jax.grad(stage_cost,0))
grad_u_stage = jax.jit(jax.grad(stage_cost,1))
hes_x_stage = jax.jit(jax.hessian(stage_cost,0))
hes_u_stage = jax.jit(jax.hessian(stage_cost,1))
hes_ux_stage = jax.jit(jax.jacfwd(jax.grad(stage_cost,1),0))
grad_x_term = jax.jit(jax.grad(term_cost))
hes_x_term = jax.jit(jax.hessian(term_cost))
@jax.jit
def iLQR_control(x,us):
def rollout(x_init,us):
def rollout_body(carry,u):
x = carry
x = model_func_risan(x,u,args.dt)
return x, x
_, xs = jax.lax.scan(rollout_body, x_init, us)
xs = jnp.concatenate([x_init[None], xs])
return xs
def calcJ(xs,us):
J = 0
def calcJ_body(carry,val):
_ = None
J = carry
x, u = val
J += stage_cost(x,u)
return J, _
J, _ = jax.lax.scan(calcJ_body, J, (xs[:-1],us))
J += term_cost(xs[-1])
J = jnp.float32(J)
return J
def Backward(xs,us):
Vxx = hes_x_term(xs[-1])
Vx = grad_x_term(xs[-1])
Q = args.Q
R = args.R
dV1 = 0
dV2 = 0
def Backward_body(carry, val):
Vx, Vxx, dV1, dV2 = carry
x, u = val
Ak = model_dfdx(x,u,args.dt)
Bk = model_dfdu(x,u,args.dt)
Qx = grad_x_stage(x,u) + Vx @ Ak
Qxx = hes_x_stage(x,u) + Ak.T @ Vxx @ Ak
Qu = grad_u_stage(x,u) + Vx @ Bk
Quu = hes_u_stage(x,u) + Bk.T @ Vxx @ Bk
Qux = hes_ux_stage(x,u) + Bk.T @ Vxx @ Ak
try:
kekka = jnp.linalg.cholesky(Quu)
except:
alpa = -jnp.amin(jnp.linalg.eig(Quu))
Quu = Quu + (alpa + 1e-5) * jnp.eye(args.len_u)
K = - jnp.linalg.pinv(Quu) @ Qux
d = - jnp.linalg.pinv(Quu) @ Qu.T
dV1 += Qu @ d
dV2 += 0.5 * d.T @ Quu @ d
Vx = Qx + d.T @ Quu @ K + Qu @ K + d.T @ Qux
Vxx = Qxx + K.T @ Quu @ K + K.T @ Qux + Qux.T @ K
return (Vx, Vxx, dV1, dV2), (K, d)
carry, output_val = jax.lax.scan(Backward_body, (Vx, Vxx, dV1, dV2), (jnp.flip(xs[:-1], 0), jnp.flip(us, 0)))
dV1 = carry[2]
dV2 = carry[3]
Ks, ds = output_val
Ks = jnp.flip(Ks, 0)
ds = jnp.flip(ds, 0)
return Ks, ds, dV1, dV2
def Forward(xs,us,Ks,ds,dV1,dV2,J):
z = 1e5
alpha = 1.0
ls_iteration = 0
def ls_rollout(xs,us,Ks,ds,alpha):
x_init = xs[0]
def ls_rollout_body(carry,val):
x_ = carry
x, u, K, d = val
u_ = u + K @ (x_-x) + alpha * d
x_ = model_func_risan(x_,u_,args.dt)
return x_, (x_,u_)
_, output_val = jax.lax.scan(ls_rollout_body, x_init, (xs[:-1],us,Ks,ds))
xs_, us_ = output_val
xs_ = jnp.concatenate([x_init[None], xs_])
return xs_, us_
def ls_check(val):
z, _, ls_iteration, _, _, _ = val
return jnp.logical_and((jnp.logical_or(1e-4 > z, z > 10)), ls_iteration < 10)
def Forwrd_body(val):
z, alpha, ls_iteration, xs_, us_, J_ = val
xs_, us_ = ls_rollout(xs,us,Ks,ds,alpha)
J_ = calcJ(xs_,us_)
z = (J-J_)/-(alpha*dV1+alpha**2*dV2)
alpha = 0.5*alpha
ls_iteration += 1
return (z, alpha, ls_iteration, xs_, us_, J_)
output_val = jax.lax.while_loop(ls_check, Forwrd_body, (z,alpha,ls_iteration,xs,us,J))
_, _, _, xs, us, J_ = output_val
return xs, us, J_
def conv_check(val):
_, _, J_new, J_old, iteration = val
return jnp.logical_and(abs(J_old-J_new) > args.torelance , iteration < args.iter)
def iLQR_body(val):
xs, us, J_old, _, iteration = val
Ks, ds, dV1, dV2 = Backward(xs,us)
xs, us, J_new = Forward(xs,us,Ks,ds,dV1,dV2,J_old)
iteration += 1
return (xs, us, J_new, J_old, iteration)
xs = rollout(x,us)
J = calcJ(xs,us)
iteration = 0
val = jax.lax.while_loop(conv_check, iLQR_body, (xs,us,J,J+1e5,iteration))
xs, us, _, _, _ = val
return us
args.u = jnp.zeros((args.len_u), dtype=jnp.float32)
args.us = jnp.zeros((args.N, args.len_u), dtype=jnp.float32)
args.x = jnp.zeros((args.len_x), dtype=jnp.float32)
Time = 0
x_log = []
start = time.time()
while Time <= 20:
print("-------------Position-------------")
print(args.x)
print("-------------Input-------------")
print(args.u)
x_log.append(args.x)
x = model_func_risan(args.x,args.u,args.dt)
us = iLQR_control(args.x,args.us)
Time += args.Ts
args.x = x
args.u = us[0]
args.us = us
end = time.time()
loop_time = end - start
print("計算時間：{}[s]".format(loop_time))
Discussion