グラフ機械学習と強化学習について

主にグラフ機械学習や強化学習手法を記載します。

強化学習入門(理論)

強化学習に必要な理論についてまとめていきます。様々な数式表記がありますが、「これからの強化学習」にできるだけ統一していきます。間違った記述があれば、ご指摘いただければ幸いです。日々アップデートしていきます。よろしくお願いいたします。

序文

強化学習機械学習のサブトピックの1つです。機械学習には、

の3つに大別されます。強化学習モデルは環境からの観測と報酬によってのみ学習されます。長期的な報酬を最大化させるようにモデルを構築していきます。

f:id:udnp:20180708214421p:plain

強化学習ですが、教師あり、教師なし学習とは異なり環境探索型の学習です。上図に示すように「環境」における「価値」を最大化させるようにエージェントを学習させていきます。強化学習の課題でよく見かけるのはゲームです。例えば、ブロック崩しでは、エージェントはプレーヤで、その行動はブロックを壊すことです。

出典: https://openai.com/blog/universe/

Markov Decision Process (MDP)

MDPが一般的な強化学習の大枠だと思います。MDPは現在の状態 sと行動 aから次の状態 s'が決まるということです。状態遷移確率 \mathcal{P}は次式のように表されます。

\begin{align} \mathcal{P_{ss'}^a}=\mathcal{P}(s, a, s')=\mathcal{p}\{s_{t+1}=s'| s_t=s, a_t=a\} \end{align} ここで、状態集合を \mathcal{S}= \{ s_1, s_2, ..., s_N \} 、行動集合を \mathcal{A}= \{ a_1, a_2, ..., a_M \} とすれば、 t+1ステップ目における状態 s_{t+1}  tステップ目からの遷移確率によって決まります。

\begin{align} S_{t+1} \sim \mathcal{P}(s' | S_t, A_t) \end{align}

上式からもわかるように、MDPは1つ前の状態および行動から次の状態が決まっています。状態遷移確率が既知のものをモデルベース強化学習と呼ばれます。例えば、囲碁は次の状態遷移が確実にわかるため、モデルベース強化学習の手法を適用することができます。現在の状態 S_tで行動 A_tを選択したとき、報酬 R_tが得られ、次の状態 S_{t+1}に遷移します。すなわち、報酬関数を下式のように表現することができます。

\begin{align} R_{t} = r(S_t, A_t, S_{t+1}) \end{align}

この報酬の履歴を管理する際は、$(S_t, A_t, R_t, S_{t+1})$という形で保存したりします。

累積報酬

教師あり学習」との違いになりますが、強化学習では最終的な報酬である累積報酬和が最大となるように行動するようなモデルになります。そのため次に示すような「累積報酬」 Gを導入します。

\begin{align} G_t = \sum_{τ=0}^{T} \gamma^{τ} R_{t+τ} \end{align}

ここで、 \gammaは時間割引率です。ゴールから遠いステップになるほど報酬の価値は低くなる経済学の理論に基づいています。この値は \gamma=0.99が多いです。1に近いほどより長期的に有益な行動を評価するようになります。割引率が0に近いほど収益はステップ tにおいて得られる報酬に近くなります。

価値関数

価値関数とは、ある状態・行動における累積報酬の予測関数です。この値を基に、ある状態が与えられたときの行動指針である方策(Policy) \piを決定します。具体的には、行動価値関数 Q^{\pi}、状態価値関数 V^{\pi}を用いて、直接行動を選択することができます。

状態価値関数(state value function)

\begin{align} V^\pi(s) = \mathbb{E}[\sum_{τ=0}^\infty \gamma^{τ} R_{τ} | S_τ=s] \end{align}

ある方策 \piに従って行動を選択すると、環境により状態 sが決まります。状態価値とは、この時の状態で得られた累積報酬和の推定量となります。この状態価値が最大となるような行動を選択するための方策を決めるアルゴリズムはon policyや方策ベースと呼ばれます。

行動価値関数(action value function)

\begin{align} Q^\pi(s, a) = \mathbb{E}[\sum_{τ=0}^\infty \gamma^{τ} R_{τ} | S_τ=s, A_τ=a] \end{align}

状態 sのとき行動 aをしたとき期待される報酬値です。ここでの期待価値と実際に得られた報酬を比較することで、価値を最大化させます。この行動価値が最大となるように行動を決める方法はoff policyや価値ベースと呼ばれます。

最適ベルマン方程式

MDP成立時の最適価値を求めるための式になります。上記の価値関数は再帰的な関係があり、ステップ t+2以降で式変形することと次のようになります。

\begin{align} V^\pi(s) &= \mathbb{E}^{\pi} [\sum_{k=0}^\infty \gamma^{k} r_{t+1+k} | s_{t}=s ] \\ &= \mathbb{E} [r_{t+1} | s_t = s ] + \gamma\mathbb{E} [\sum_{k=0}^\infty \gamma^{τ} r_{t+2+k} | s_{t+1}=s' ] \\ \end{align} ここで、状態 sのときの累積期待報酬は、それまで全ての行動( \sum_a \pi(a|s))と状態遷移したときの報酬和で求まります。

\begin{align} \mathbb{E^\pi}[r_{t+1}|s_t=s] = \sum_a \pi(a|s)\sum_{s'}\mathcal{P}(s, a, s')R(s,a,s') \end{align} 上式と、再帰的な式を利用すると

\begin{align} V^\pi(s) &= \sum_a \pi(a|s)\sum_{s'}\mathcal{P}(s, a, s')R(s,a,s') + \gamma \sum_a \pi(a|s)\sum_{s'}\mathcal{P} V^\pi(s') \\ &= \sum_a \pi(a|s)\sum_{s'}\mathcal{P}(s, a, s')\big(R(s,a,s')+\gamma V^\pi(s')\big) \end{align}

Qについても同様です。

\begin{align} Q^\pi(s,a) &= \sum_{s'}\mathcal{P}(s, a, s')\big(R(s,a,s')+\gamma \sum_a \pi(a'|s')Q^\pi(s',a')\big) \end{align}

これらの二つの式を合わせると次のようになります。

\begin{align} Q^\pi(s,a) &= \sum_{s'}\mathcal{P}(s, a, s')\big(R(s,a,s')+\gamma V^{\pi}(s')\big) \end{align}

状態遷移確率が既知ならベルマン方程式を解くことで状態や行動の価値がわかります。しかし、基本的には状態遷移確率はわからないので試行錯誤で決めていきます。ベルマン方程式を解くことができ、価値関数を基にして行動を決定すれば、最適な結果となります。

Policy(方策)とは

環境からの観察により状態が決まり、エージェントは行動を選択します。環境から状態 sを観測したときに、行動 aを取る確率分布が方策(policy)です。この方策を \pi(a|s)とすると A_t \sim \pi(a|s_t)となります。

方策関数には様々なものが提案されています。微分可能な方策関数は後述する方策勾配法で用いられます。活用と探索の例を見ていきます。

greedy 法

{} \begin{align} \pi(a|s) = \left\{ \begin{array}{} 1 & (a = \arg \max_a Q(s, a)) \\ 0 & ( \text{otherwise}) \end{array} \right. \end{align}

greedy方策(貪欲法)は最適行動価値関数が既知であれば最適な行動をとります。しかし、このQ関数は未知であるため、推定していく必要があります。また、初めの状態ではQ関数は間違っているため、この値を基にした方策では間違った行動をし続ける危険性があります。遷移確率が既知であれば、後述するベルマン方程式により価値関数は解析的に求めることができます。しかしながら通常、遷移確率は未知であるため、価値関数を近似するアルゴリズム(Q-learning, Sarsa, DQNなど)を用いてエージェントの行動が決まります。

epsilon-greed法

一方、 \epsilongreedでは乱数の確率が \epsilonより大きければQ関数の最大値を利用、小さければ利用を行う方法です。貪欲法の同じ行動を起こすことを防ぐ方法です。

\begin{align} \pi(a|s) = \left\{ \begin{array}{} 1-\epsilon + \frac{\varepsilon}{|\mathcal{A}(s)|} & (a = \arg \max_a Q(s, a)) \\ \frac{\varepsilon}{|\mathcal{A}(s)|} & ( \text{otherwise}) \end{array} \right. \end{align}

Softmax関数

ボルツマン分布がよく使用されます。

\begin{align} \pi_\theta (a | s) = \frac{\exp(\frac{1}{\beta} Q_\theta(s, a))}{\sum_{j=1}^{T}\exp(\frac{1}{\beta} Q_\theta(s, a))} \end{align}

方策を \thetaでパラメタライズしています。

Dynamic Programming (DP)

動的計画法は、部分問題を交互に記録しながら解いていく方法です。強化学習の枠組みでは、価値反復法と方策反復法の大きく2通りが考えれます。上記ベルマン方程式を求める際に用いられます。環境が既知の場合は、この式は反復法により直接的に解くことができます。

モンテカルロ法(MC)

モンテカルロ法では報酬の期待値を平均とすることで近似しています。状態のパスを通るたびに回数をカウントしていき、それまでの累積報酬和との平均を取ることで状態価値関数を推定します。もしくは一度取ったパスの累積報酬を用いる場合もあります。モンテカルロ法では、エピソードが終了するまで価値の更新を待つ必要があります。

エピソード$i$においてTステップ目までの累積報酬は以下の式で表されます。遠くのステップの報酬ほど時間割引率γによって価値を減衰させていきます。

$$ G_{i, t} = r_{i, t} + \gamma r_{i, t+1} + \gamma^{2} r_{i, t+2} + \dots + \gamma^{T_{i}-1} r_{i, T_t} $$

Every-Visit Monte Carlo

状態価値$V$は次のように更新されます。

$$ \begin{align} N(s) &= N(s) + 1 \\ G(s) &= G(s) + G_{i, t} \\ V^{\pi}(s) &= \frac{G(s)}{N(s)} \end{align} $$

パスカウントを1つ増やし、以前の累積報酬和に対して今回得られたものを足し合わせます。最後にそれをパスカウントで割り、状態$s$の価値とします。Every visit MCではパスを通るたびに状態価値の更新を行います。この場合ではVの推定量は偏ったものになります(バイアス)。しかし収束性が良く、より良いMSEとなります。

First-Visit Monte Carlo

First visit MC法では、Every-Visit MCと異なり、1度取ったパスのみ更新を行います。 直接報酬を使うため状態価値関数の推定量にはバイアスがなく、真の$\mathbb{E}_{\pi} [G_t | s_t = s]$となります。回数が増えるにつれて一定のところに収束します。

Incremental Monte Carlo

一方、パスの途中でも更新する場合は、後述するTD学習に類似する式変形を行います。

$$ \begin{align} V^{\pi}(s_t) := V^{\pi}(s_t) + \alpha[ G_{i, t} - V^{\pi}(s_t)] \end{align} $$

  • $\alpha = \frac{1}{N(s)}$: every visit MCと同一
  • $\alpha \gt \frac{1}{N(s)}$: 古いデータを忘れます。非定常のドメインで役立つ

繰り返し累積報酬和と状態価値の誤差を足し合わせることで状態価値推定します。

Qiitaに書いた記事の多腕バンディット問題でも、状態価値は$\alpha = \frac{1}{n}, G_{i, t} = r_{i, t}$とした場合の逐次モンテカルロによって推定したものです。多腕バンディットでは状態遷移は行われませんが、MC法によって状態価値を推定でき、最適な行動を行えるようになっていきます。

Monte-Carlo Tree Search (MCTS)

モンテカルロ法による価値推定に類似する考え方にMCTSがあります。これはヒューリスティックアルゴリズムです。ボードゲーム、チェスや囲碁のようなモデルベース強化学習で用いられています。

  • Selection -> Expansion -> Simulation -> Backpropagation

という流れでノード(状態)を拡大していくことで最適な行動を探索していくことができます。

Temporal Difference (TD) Learning

TD learningはMCとDPを組み合わせた考え方です。最終的な結果を待たずに価値を推測できます。TD(0)、one-step TDの更新式は下式のように表せます。

\begin{align} V(s_t) := V(s_t) + \alpha[ R_{t} + \gamma V(s_{t+1}) - V(s_t)] \end{align}

TD学習ではTDターゲットと呼ばれる$R_{t} + \gamma V(s_{t+1})$ と $V(s_t)$との差(TD誤差)を使って更新していきます。

SARSA (on-policy TD control)

価値関数のベルマン方程式を解く(近似する)ためのアルゴリズムがSARSAです。一般的に環境の遷移確率や方策は未知です。そこで実際の観測値(s,a,r,s',a')に基づき、行動価値関数を求めます。SARSAの元のベルマン方程式には方策 \piを使用するため、on-policy型と呼ばれています。

\begin{align} Q(s_t, a_t) := Q(s_t, a_t ) + \alpha ( {R_{t} + \gamma Q(s_{t+1}, \pi(s_{t})) - Q(s_t, a_t))} \end{align}

Q-learning (off-policy TD control)

一方Q-leraningでは、次の行動におけるQ関数式に方策が含まれていないのでoff-policyです。(Watkins, 1989)

\begin{align} Q(s_t, a_t) := Q(s_t, a_t ) + \alpha ( {R_{t} + \gamma \max_{a} Q(s_{t+1}, a) - Q(s_t, a_t))} \end{align}

方策勾配法 ( Policy Gradient )

Policy gradient methodsは報酬の勾配を繰り返し推定することで期待報酬を最大化する方法です。

方策勾配定理

 \thetaでパラメタライズされた確率的な方策 \pi_\thetaを考えます。この方策は最適な行動をするための確率分布となります。

\begin{align} \pi_\theta = p_\theta (a_t | s_t) \end{align}

ここで、 pは状態 s_tにおける行動 a_tに対する確率分布です。  τをステップ0からHまでの(state, action)系列  τ = (s_0, a_0, ..., s_H, a_H ) とし、このとき方策の評価関数を下式のように定義します。

\begin{align} J(\theta) &= \mathbb{E}_{\pi_{\theta}} [\sum_{t=0}^H \mathcal{R} (s_t, a_t)] \\ &= \sum_{τ} \mathcal{P}(τ ; \theta) \mathcal{R}(τ) \end{align}

ここで、  \mathcal{R}(τ)=\sum_{t=0}^H \mathcal{R} (s_t, a_t)としています。  \mathcal{P}(τ ; \theta)は状態遷移確率モデルです。すなわち、

\begin{align} \mathcal{P}(τ ; \theta) = \prod_{t=0}^H \mathcal{P}(s_{t+1} | s_t, a_t) \pi_\theta (a_t | s_t) \end{align}

です。これは全ての方策での行動とその結果として得られる状態遷移確率の尤度となります。方策を学習するためには、この尤度関数を最大化させれば良いです。

\begin{align} \max_{\theta} J(\theta) = \max_{\theta} \sum_{τ} \mathcal{P} (τ ; \theta) \mathcal{R}(τ) \end{align}

状態遷移 \mathcal{P}が既知であるものはモデルベース強化学習と呼ばれています。

尤度の最大化

この目的関数を最大にするためにJに対する勾配 \nabla_\theta = \frac{dJ}{d\theta}を求めることで方策を更新していきます。

\begin{align} \nabla_\theta J(\theta) &= \nabla_{\theta}\sum_{τ} \mathcal{P} (τ ; \theta) \mathcal{R}(τ) \\ &=\sum_{τ} \nabla_{\theta} \mathcal{P} (τ ; \theta) \mathcal{R}(τ) \\ &=\sum_{τ} \nabla_{\theta}\frac{ \mathcal{P} (τ ; \theta)}{ \mathcal{P} (τ ; \theta)} \mathcal{R}(τ) \\ &=\sum_{τ} \mathcal{P} (τ ; \theta) \nabla_{\theta} \log{ \mathcal{P} (τ ; \theta)} \mathcal{R}(τ) \\ &= \mathbb{E}_{\pi_\theta} [ \nabla_{\theta} \log{ \mathcal{P} (τ ; \theta)} \mathcal{R}(τ) ] \end{align}

この勾配式より、報酬Rが高い遷移の場合に方策が大きく更新されます。次に方策と遷移確率について立式します。

\begin{align} \nabla_{\theta} \log{ \mathcal{P} (τ ; \theta)} &= \nabla_{\theta} \log{ \prod_{t=0}^H \mathcal{P}(s_{t+1} | s_t, a_t) \pi_\theta (a_t | s_t)} \\ &= \nabla_{\theta} [ \sum_{t=0}^H \log{\mathcal{P}(s_{t+1} | s_t, a_t)] + \nabla_\theta [ \sum_{t=0}^H \log{ \pi_\theta (a_t | s_t)}}] \\ &= \nabla_\theta \sum_{t=0}^H \log{ \pi_\theta (a_t | s_t)} \\ &= \sum_{t=0}^H \nabla_\theta \log{ \pi_\theta (a_t | s_t)} \end{align}

以上により、方策勾配は下式です。

\begin{align} \nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} [\sum_{t=0}^H \nabla_\theta \log{ \pi_\theta (a_t | s_t)} \mathcal{R}(s_t, a_t)] \end{align}

Hステップ目まで方策$\pi_{\theta}$に基づいて行動し、これまでの報酬を取得します。方策ネットワークの出力値はsoftmaxで正規化され、行動の出力分布になっています。これを離散値の場合はカテゴリカル分布でサンプリングすることで行動を決定します。方策の勾配を計算するために確率密度関数から対数尤度を計算し、それとこれまで得られた報酬関数のベクトル値をかけて平均を取り最大化させます。

交差エントロピーとの関係性

方策勾配法はクロスエントロピーと密接な関係があります。クロスエントロピーの式は以下となります。

\begin{align} H(p,q) = \mathbb{E}_p [- \log q] = H(p) + D_{KL} (p || q ) \end{align}

ここで、pとqが離散ならば、

\begin{align} H(p, q) = - \sum_x p(x) \log q(x) \end{align}

となります。クロスエントロピーを使った方法でも方策を学習することができます。この場合では、(state, action, reward)を収集し、良かったペアを選別します。そのような行動を取れるように方策を学習させます。すなわち、方策のlogit値と行動ベクトルとのクロスエントロピー誤差を計算し最小となるように方策を学習します。これは方策勾配法での良いエピソードであったらR=1にし、悪いエピソードであった場合0にしたものと同じです。

方策勾配法のアルゴリズム

報酬の決め方によってアルゴリズムがやや異なります。

  1.  \sum_{t=0}^{\infty} r_t: トラジェクトリーの総報酬.
  2.  \sum_{t'=t}^{\infty} r_t': 行動 a_tに続く報酬
  3.  \sum_{t'=t}^{\infty} r_t' - b(s_t): 2. ベースラインを導入したもの
  4.  Q^{\pi}(s_t, a_t): state-action value function
  5.  A^{\pi}(s_t, a_t): advantage function
  6.  r_t + V^{\pi}(s_{t+1}) - V^{\pi}(s_t): TD誤差

https://arxiv.org/pdf/1506.02438.pdf

Baseline

方策勾配の収束性を上げるために、baselineを導入する方法が提案されています。これにより分散を減少させることができます。

\begin{align} \nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} [\sum_{t=0}^H \nabla_\theta \log{ \pi_\theta (a_t | s_t)} \big(\mathcal{R}(s_t, a_t) -b \big)] \end{align}

vanilla policy gradient

 \thetaだけでなくbaseline bも最適化します。ここでのベースラインbは、報酬との最小2乗法により求めます。

REINFORCE

複数の遷移情報を使って計算した \mathcal{R}_tの代わりに、その時々の報酬[r_t]を使う方法。baselineは報酬の平均値を用いることが多いです。

  • REINFORCEの問題点

全ての行動における累積報酬は平均値を用いるため、一部の行動により悪い結果となったとしても、その行動が考慮されないことになります。

Actor-Critic

上記アルゴリズムは全て方策ベースのアルゴリズムです。REINFORCEではQ関数を報酬の平均値で近似していますが、上記で述べたように学習に用いられていません。一方、Actor Criticでは価値ベースも加えたhybrid法となっています。

  • Criticでは、どの程度行動が良いかを評価する。 (policy - based)
  • Actorでは、Criticの基準で更新された方策に基づき行動する。(value - based)

  • Q関数との関係

全ての状態を考慮し、時間割引率を \gamma、行動価値関数をQとすると、Actorモデルの更新は下式の方策勾配法で表現できます。

\begin{align} \nabla_{\theta} J(\theta) &= \mathbb{E}_{\pi_\theta} [\sum_{t=0}^{\infty} \nabla_{\theta} \log{\pi_{\theta} (a_t | s_t)} \sum_{k=t}^{\infty} \gamma^{k} r(s_k, a_k) ] \\ &= \mathbb{E}_{\pi_{\theta}} [\sum_{t=0}^{\infty} \nabla_{\theta} \log{ \pi_{\theta} (a_t | s_t)} ] \hat{Q}_w(s_t, a_t) \end{align}

これで、各ステップにおいて方策を更新することができる。その代わりに、Criticモデルはこの価値関数 Qを近似するように学習させる必要がある。 Qをパラメータ w微分し、この勾配を用いてQ関数を更新するが、TD誤差も用いる。

\begin{align} \text{TD}_{Error} = R(s_t, a_t) + \gamma \hat{Q}_w (s_{t+1}, a_{t+1}) - \hat{Q}_w (s_t, a_t) \\ \nabla w = \text{TD}_{Error} \nabla_{w} \hat{Q}_w (s_t, a_t) \end{align}

ActorとCriticモデルにおいて学習率は変え、それぞれ別々に最適化されなければならないです。これは収束性の問題です。

Actor Critic

ステップ tにおいて環境から状態 s_tが得られる。インプットとしてActorとCriticによりそれぞれ、方策 \pi_{\theta}と価値 \hat{Q}_wが得られる。 この方策 \pi_{\theta}に基づき、Actorからのアウトプットとして a_tが得られる。新しい状態 s_{t+1}と報酬 r_{t+1}が得られる。そのおかげで、

Criticはその状態での行動の価値(Q値)を計算。 ActorはこのQ値を用いて方策のパラメータを更新する。

\begin{align} \nabla_\theta = \alpha \nabla_\theta \log{ \pi_\theta (a_t | s_t)} \hat{Q}_w(s_{t}, a_t) \end{align}

パラメータ \thetaが更新されたおかげで、Actorは状態 s_{t+1}において、行動 a_{t+1}を取り、環境から状態 s_{t+2}と報酬 r_{t+1}が得られます。Criticはこの値を元にパラメータ wを更新します。

\begin{align} \text{TD}_{Error} = R(s_t, a_t) + \gamma \hat{Q}_w (s_{t+1}, a_{t+1}) - \hat{Q}_w (s_t, a_t) \\ \nabla w = \text{TD}_{Error} \nabla_{w} \hat{Q}_w (s_t, a_t) \end{align}

A2C

Advantage Actor Criticでは、baselineに状態価値関数 V^{\pi} (s)を、報酬部分にQ(s, a)を用いる方法です。これはそれぞれ、その状態における平均価値Vと状態sを取ったときの行動価値Qであり、その差を取ったAdvantage関数を利用しています。

\begin{align} A(s, a) = Q^{\pi}(s, a) - V_{\theta}(s) \\ \nabla_\theta J(\theta) = \alpha \nabla_\theta \log{ \pi_\theta (a_t | s_t)} A(s, a) \end{align} Advantage関数は、ある状態でとった行動の平均と比較し改善することを意味しています。

  • A(s, a) > 0 : 勾配は正しい方向へ。
  • A(s, a) < 0 : 行動が状態平均よりも悪い。

\begin{align} A(s, a) &= Q(s, a) - V(s) \\ &= r + \gamma V(s') - V(s) \\ &= \text{TD}_{Error} \end{align}

遠くまで先読みするとREINFORCEに近くなります。A(s, a)の二乗が最小となるように更新します。

from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.evaluation.postprocessing import discount_cumsum


def discount_cumsum(x: np.ndarray, gamma: float) -> np.ndarray:
    """Calculates the discounted cumulative sum over a reward sequence `x`.
    y[t] - discount*y[t+1] = x[t]
    reversed(y)[t] - discount*reversed(y)[t-1] = reversed(x)[t]
    Args:
        gamma: The discount factor gamma.
    Returns:
        The sequence containing the discounted cumulative sums
        for each individual reward in `x` till the end of the trajectory.
    Examples:
        >>> x = np.array([0.0, 1.0, 2.0, 3.0])
        >>> gamma = 0.9
        >>> discount_cumsum(x, gamma)
        ... array([0.0 + 0.9*1.0 + 0.9^2*2.0 + 0.9^3*3.0,
        ...        1.0 + 0.9*2.0 + 0.9^2*3.0,
        ...        2.0 + 0.9*3.0,
        ...        3.0])
    """
    return scipy.signal.lfilter([1], [1, float(-gamma)], x[::-1], axis=0)[::-1]


rollout= SampleBatch(
    {
        SampleBatch.OBS: np.array(
            [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]]
        ),
        SampleBatch.ACTIONS: np.array([0, 1, 1]),
        SampleBatch.REWARDS: np.array([1.0, 1.0, 1.0]),
        SampleBatch.DONES: np.array([False, False, True]),
        SampleBatch.EPS_ID: np.array([1234, 1234, 1234]),
        SampleBatch.AGENT_INDEX: np.array([0, 0, 0]),
    }
)

last_r = 0.0
gamma = 0.99
rewards_plus_v = np.concatenate([rollout[SampleBatch.REWARDS], np.array([last_r])])
discounted_returns = discount_cumsum(rewards_plus_v, gamma)[:-1].astype(np.float32)
> array([2.9701, 1.99  , 1.    ], dtype=float32)
# A = [0.99^2 * 1.0 + 0.99 * 1.0 + 1.0, 0.99 * 1.0 + 1.0, 1.0] = [2.9701, 1.99, 1.0]

GAE

Generalized advantage estimatorでは、方策勾配法学習中のバリアンスを減らすために、上記の価値関数に時間割引率 \gamma \in [0, 1] と パラメータ \lambda \in [0, 1]を導入したものです。

# V(s)と最後の報酬を結合
vpred_t = np.concatenate([rollout[SampleBatch.VF_PREDS], np.array([last_r])])

# TD誤差の計算
delta_t = rollout[SampleBatch.REWARDS] + gamma * vpred_t[1:] - vpred_t[:-1]

# Advantage計算(総和)
rollout[Postprocessing.ADVANTAGES] = discount_cumsum(delta_t, gamma * lambda_)
rollout[Postprocessing.VALUE_TARGETS] = (
            rollout[Postprocessing.ADVANTAGES] + rollout[SampleBatch.VF_PREDS]
        ).astype(np.float32)