[論文解説] AQL: Q-Learning in enormous action spaces via amortized approximate maximization

この記事は，以下の論文の解説です．

Q-Learning in enormous action spaces via amortized approximate maximization (2020)

記事内容では，強化学習の基礎的な知識を前提としています．
また，記事中の図は全て論文からの引用です．
不備がございましたら，ご指摘頂けると幸いです．

概要

この論文は，「提案分布からのサンプル上でmax操作を行うことで，離散や連続，もしくはその混合行動空間上でQ学習を行うことのできる手法(AQL)を提案した」論文になります．

Q学習は，実装が容易であることに加えて，off-policyで学習可能なため，リプレイメモリや分散学習などのテクニックと容易に併用することができます．一方，Q学習では，行動空間上でmax操作を行う必要があるため，連続行動空間や高次元の離散行動空間に適用しにくいという課題があります．

この論文では，正確なmax操作を行うよりも，最適な行動の検索の仕方(=argmax)を学習することで，max操作による行動の選択のコストを削減することを目標としています．すなわち，状態で条件づけられた行動空間上の提案分布を学習し，提案分布から行動の候補をサンプルし，その候補の中でmax操作を行うことで計算量を削減します．

Deep Deterministic Policy Gradient(DDPG)などの手法では，決定的な方策がmax操作と同等の処理を行なっていると考えられ，Q学習の一種とみなすことができます．これに対し，提案手法 Amortized Q-Learning(AQL) では，教師あり学習により提案分布を学習するため，提案分布の学習はQ関数の勾配に依存しないという大きな特徴があります．これにより，連続行動しか扱えないDDPGと異なり，離散や連続，もしくはその混合行動空間を扱うことができます．

Amortized Q-Learning (AQL)

提案手法 AQL では，Q関数を推定するネットワーク(重み: $\theta^Q$)に加えて，状態$s$で条件付けられた，行動の提案分布$\mu$のパラメータを出力するネットワーク(重み: $\theta^\mu$)を学習します．分布のパラメータとは，ガウス分布の平均やカテゴリカル分布のロジットなどを指します．その後，正確なmax操作を行う代わりに，この提案分布$\mu$からのサンプルの中でmax操作を行います．

提案分布の学習

状態$s$に関して，提案分布から$N$個のサンプルを，一様分布から$M$個のサンプルを取ります．この中で現在の価値$Q(s,a)$が最も高い行動を$a^{*}(s)$とします．このとき，提案分布に関する損失関数は，以下のようになります．

\mathcal L (\theta^\mu; s) = - \log (a^*(s)|s; \theta^\mu) - \lambda H(\mu(a|s; \theta^\mu))

1項目は，最適な行動をサンプルしやすくなるように，対数尤度を最大化する項です．これにより，提案分布からのサンプルが，より高い行動価値を持ちやすくなると考えられます．2項目は，提案分布が決定的にならないように多様性を推進する，提案分布のエントロピーを最大化する項です．

このように提案分布の学習では，$M+N$個のサンプルの中で最も価値の高い行動を正解とした，教師あり学習(対数尤度最大化+正則化)を行なっていると考えることができます．

Q関数の学習

Q関数の損失関数は，以下のようにになります．

\mathcal L(\theta^Q) = \mathop{\mathbb{E}}_{\pi, p_\epsilon}[( r_t + \gamma Q(s_{t+1}, a^*(s_{t+1})) -  Q(s_t, a_t; \theta^Q))^2]

これは通常のQ学習とほぼ同じものですが，max操作を全ての行動からではなく，提案分布からのサンプルでのみ行なっている点のみが異なります．また，提案分布$\mu$のクラスを変えることで，離散(Categorical)や連続(Gaussian)の行動空間に対応することが可能です．

アルゴリズム

まとめると，以下のようなアルゴリズムになります．

検証

DM Control Suite(状態入力)とDM Lab(画像入力)で検証しています．

DM Control Suite

ここでは，比較的低次元の状態空間として，画像ではなく生の状態を扱います．DM Control Suiteの行動空間は，低次元のものから高次元(21次元，humanoidタスク)のものまで存在します．

比較対象としては，D3PG，QT-Opt，IMPARAを用いています．また，AQLの提案分布では，

• AQL
  • 連続行動空間を$\{-1, -0.5, 0, 0.5,1.0\}$の5通りに離散化し，カテゴリカル分布のロジットを出力する場合
  • $5^D$個のユニークな行動
• AQL_CONT
  • 連続行動でガウス分布の平均(分散 $\sigma^2=0.25$ は定数)を出力する場合

の2通りを検証しています．また，提案分布からのサンプリングでは，次元順に自己回帰的にサンプリングを行っています(下式)．これにより，各次元間の依存関係を考慮していると主張しています．その他の設定は，論文を参照してください．

\mu(a|s;\theta) = \prod_{d=1}^D \mu_d(a_d|s, a_{<D};\theta^\mu)

DM Control Suiteの結果

図右上のグラフはDM Control Suiteの結果の平均を示しており，AQLが最も良い性能であることがわかります．

また，図下の2つのグラフでは，中次元の行動空間(左)と高次元の行動空間(右)のタスクでの結果を示しています．高次元行動空間では，AQLとD3PGが他の手法を明らかに上回っていることがわかります．また，AQL_CONTは中次元行動空間では学習できているものの，高次元ではうまく学習できていないことがわかります．これは，連続行動空間より離散行動空間の方が探索が容易であることが原因であると指摘されています．

まとめると，DM Control Suiteでの検証から，提案分布の学習は複雑・高次元な行動空間において有益であるということが言えます．

DM Lab

DM Labでは，高次元の状態空間($84\times 84\times 3$ 次元の画像)と，以下のような複雑で構造的な7次元の行動空間を扱います．

• 1人称視点の角速度(2次元)
  • left/right，up/downの2次元
  • $[-512, 512]$の整数値
• 動き(2次元)
  • left/right/no-op，forward/backward/no-opの2次元
  • 3つの離散値
• その他(3次元)
  • fire/no-op，jump/no-op，crouch/no-op
  • バイナリ値

以下では，2つの行動集合を考えます．

• curated
  • 全てのタスクで良い性能を示すための最低限の行動で構成された行動集合
  • 11個のユニークな行動
• large
  • 角速度をそれぞれ7個に離散化
  • $3528 = 7^2 \cdot 3^2\cdot 2^2$個のユニークな行動

また，比較対象として以下の手法を検証します．

• 分散バージョンの提案手法 AQL (離散行動空間なのでAQL_CONTはなし)
• AQLのAblationとして，厳密なQ学習(全行動からのmax操作)
• IMPARA

DM Labの結果

上図では，左右それぞれ2つのタスクでの結果を示しています．両方のタスクにおいて，large行動集合で学習したAQLが，他の手法を上回る最終性能を示していることが読み取れます．

またAQLと厳密なQ学習の比較では，curated行動集合を用いた際には大きな差異は見られなかったものの，より探索が難しいlarge行動集合を用いた際には，AQLの方がより良いサンプル効率を示しています．このことは，AQLの近似的な行動選択によって，探索が促進されているからだと考えられます．

まとめ

この論文では，連続行動空間や高次元離散行動空間上でQ学習を行う手法AQLを提案しました．

個人的には，行動の次元順に自己回帰的にサンプリングを行っており，先頭次元の行動は他の行動に依存していない点が気になりました．行動の次元間の依存関係を表現する，もっと洗練された手法がありそうな気がします...