一日坊主

雰囲気でやっている

Task-Relevant Adversarial Imitation Learning を読んだ

この記事は強化学習苦手の会 Advent Calendar 2020の22日目の記事です.

本記事では,Task-Relevant Adversarial Imitation LearningというDeepMindの論文を紹介します. この論文はCoRL2020のSpotlight Talkに選ばれたようです. 本記事の構成は落合先生の論文まとめフォーマットに従っています.

なお私は強化学習が本当に苦手なので,論文をちゃんと理解できている自信がありません.有識者の皆様からのご指摘を心待ちにしています.

どんなもの?

敵対的模倣学習(GAIL)のパフォーマンスを低下させる原因として,GAILのDiscriminator(識別器)がspurious associations(本記事ではこれを疑似関連と呼ぶことにします)に引っ張られてしまうことを指摘し,それを解決する手法であるTask-Relevant Adversarial Imitation Learning (TRAIL)を提案しています.


模倣学習とは,エキスパートの方策を模倣する方策を学習する問題です.模倣学習手法の一つであるGAILは,敵対的生成ネットワーク(GANs)の枠組みで,エキスパートの方策の生成する状態行動対と,Generator(生成器)の生成する状態行動対を区別する識別器を敵対的に学習します.こうすることで,報酬関数を陽に定義することなく,エキスパートの状態行動対の集合(軌跡)からターゲットの方策を獲得できます.

疑似関連とは何かというと,「エキスパートのラベル」と「タスク無関係の特徴」との間に生じる疑似相関のようなものです.イメージ図が論文中のFigure 1にありますので引用します.

f:id:twakamori:20201221051518p:plain
疑似関連のイメージ (元論文のFigure 1より引用)

図の縦軸はタスクに関係する特徴,横軸はspuriousな特徴(疑似特徴)を表しています.緑色の点はエキスパートの観測,赤色は学習中のエージェントの観測を表します. 学習が進むにつれてエージェントはエキスパートに近づこうとするため,生成器はタスクに関連する特徴を出力するようになります(学習の経過は赤色の濃さで表されています). その結果,識別器の決定境界(図中の青色の破線)は,タスクに関連する特徴の代わりに,擬似特徴に依存してエキスパートとエージェントを区別するように学習されてしまいます. このような識別器がタスク無関係の特徴に依存してしまう問題を,この論文では疑似関連と呼んでいます.

論文では疑似関連の例として,画像を状態とする制御タスクにおける初期状態や画像の輝度を挙げています.エージェントはそれらを制御できないため,いかなる方策を生成しても識別器を騙すことができず,その結果,生成器の学習が進まなくなります.

先行研究と比べてどこがすごい?

Stadieらの先行研究*1は,画像において,エキスパートとエージェントの観測が異なる視点から来るような三人称の問題にGAILを適用しました.先行研究は識別器が視点を利用しないようにするために,gradient flipping*2というテクニックを用いてドメイン不変の特徴を学習します.しかし,本論文の著者らによると,Stadieらの先行研究は疑似関連の問題に対してドメイン不変の特徴をうまく抽出できなかったそうです.

また,疑似関連の問題は,識別器の過学習の問題と似ています.先行研究*3*4*5は,GAILの識別器の過学習を防ぐために正則化を行いますが,前述のような画像特徴の疑似関連をもつ場合,正則化ではそれを防ぐことができず,場合によっては疑似関連を強調してしまうことさえあります.

技術や手法のキモはどこ?

提案手法(TRAIL)は,constraining sets(制約集合) \mathcal{I} _ E \mathcal{I} _ Aを導入します. \mathcal{I} _ E \mathcal{I} _ Aは,それぞれエキスパートの観測の集合とエージェントの観測の集合を表します. \mathcal{I} _ E \mathcal{I} _ Aを識別する特徴を疑似特徴として区別することで,識別器がタスクに関連する特徴をうまく捉えることができるようになります.


f:id:twakamori:20201224033707p:plain
 \mathcal{I} _ E \mathcal{I} _ Aの導入イメージ (元論文のFigure 3より引用)

上図の左側は,冒頭の図で学習された決定境界がタスクに関連する境界と擬似的な境界の組み合わせであることを示しています.上図の右側のように, \mathcal{I} _ E \mathcal{I} _ Aを導入し,これらを識別する擬似的な境界を学習し,識別器がタスクに関連する境界をうまく捉えられるようにすることが本研究の課題となります.

手順は以下のとおりです.まず,GAILの識別器の目的関数を変形して,状態行動対 (s, a)の代わりに状態 (s)のみを扱うようにします(状態のみを扱う模倣学習は先行研究*6があります).

\displaystyle{
\underset{\psi}{\max}\mathbb{E}_{s\sim\pi_E}\left[\log D_\psi(s)\right]+\mathbb{E}_{s\sim\pi_A}\left[\log \left(1-D_\psi(s)\right)\right] \tag{1}
}

ここで, \pi _ A \pi _ Eはそれぞれエージェントの方策とエキスパートの方策を表し, Dは識別器のネットワークを表します.学習した識別器から報酬関数が得られます(本研究ではシンプルに R(s) = -\log(1-D _ \psi(s))と置いています). \pi _ Aはこの報酬関数を最大化するよう学習されます.

TRAILの識別器は,式 (1)と同様に交差エントロピーを用いますが, \mathcal{I} _ E \mathcal{I} _ Aからの観測に対して制約をかけます.具体的には, \mathcal{I} _ E \mathcal{I} _ Aからの観測に対して, accuracy(\mathcal{I} _ E, \mathcal{I} _ A)\geq\frac{1}{2}の場合は負の制約をかけます.ここで, accuracy(\mathcal{I} _ E, \mathcal{I} _ A)は,2つの制約集合からの観測の均等な集合における平均識別器精度として次式で定義されます.

\displaystyle{
accuracy(\mathcal{I}_E, \mathcal{I}_A)=\frac{1}{2}\mathbb{E}_{s\in\mathcal{I}_E}\left[\mathbf{1}_{D_\psi(s)\geq\frac{1}{2}}\right]+\frac{1}{2}\mathbb{E}_{s\in\mathcal{I}_A}\left[\mathbf{1}_{D_\psi(s)\lt\frac{1}{2}}\right] \tag{2}
}

直感的には,この制約は「識別器が疑似特徴を識別・利用することを抑制し,もし識別器の学習に使われた場合はそれを忘却させる」という働きをします.

エキスパートとエージェントからの Nサンプルのバッチをそれぞれ s _ E \sim \pi _ E, s _ A \sim \pi _ A,制約観測のサンプルを  \hat{s} _ E \subset \mathcal{I} _ E, \hat{s} _  A \subset \mathcal{I} _ Aとすると,TRAILの識別器は以下を最大化します.

\displaystyle{
\mathcal{L}_\psi(s_E,s_A,\hat{s}_E,\hat{s}_A)=G_\psi(s_E,s_A)-\mathbf{1}_{accuracy(\hat{s}_E,\hat{s}_A)\geq\frac{1}{2}}G_\psi(\hat{s}_E,\hat{s}_A)
}

ここで, G _ \psi(s _ E,s _ A) D _ \psiに対するGAILの識別器の損失のサンプルによる推定値を表し,次式で定義されます.

\displaystyle{
G_\psi(s_E,s_A)=\sum_{i=1}^{N}\log D_\psi\left(s_E^{(i)}\right)+\log \left[1-D_\psi\left(s_A^{(i)}\right)\right] \tag{3}
}

また, \mathbf{1} _ {accuracy(\hat{s} _ E,\hat{s} _ A)\geq\frac{1}{2}}は指示関数であり,制約を違反するかどうかを表します.

以上が手法の全体像となります. さて,ここで大きな疑問として「 \mathcal{I} _ E \mathcal{I} _ Aをどうやって用意するのか」という点が挙げられます. これについて,論文では以下のように述べられていました.

  • エキスパートとエージェントの観測に非定常な分布シフトがある場合:不可能
  • そうでない場合:
    • エキスパートとエージェントの両方の設定でランダムな方策を実行する…ランダム方策はタスク依存の情報を持たないため有効
    • エキスパートとエージェントのエピソードの初期のフレームを使用する...初期のフレームはタスク依存の行動がほとんど見られないため有効

どうやって有効だと検証した?

シミュレータ上のロボットマニピュレーションタスク(積み上げ,持ち上げ等)で他手法(GAIL,GAIL+AES,D4PGfD)と性能を比較し,エキスパートとエージェントの見た目が異なるタスクにおいて,TRAILが最も良い性能であることを確認しています.


実験設定の詳細については説明を割愛します.本研究で比較対象としているアルゴリズムは以下です.

  • BC: Behavior Cloning
  • GAIL: ベースライン(式 (1)の目的関数で識別器を訓練)
  • TRAIL: 提案手法(ベースラインの識別器の目的関数を式 (3)に変えたもの+early stopping)
  • GAIL+AES: ベースライン+early stopping
  • D4PGfD: D4PGによるエージェントの学習において,経験再生にエキスパートの軌跡を加えたもの(報酬関数には別途定義したスパースな報酬関数を用いる)

1つ目の実験は,Lift(持ち上げ),Stack(積み上げ),Stack banana(バナナ形状のオブジェクトの積み上げ),Insertion(挿入)の4つのマニピュレーションタスクにおいて,各アルゴリズムを比較しています.この実験では,エキスパートとエージェントの観測の間で初期状態や輝度などを特にいじっていないので,疑似特徴がどのような特徴を捉えるかは分かりません.

f:id:twakamori:20201224220458p:plain
マニピュレーションタスクにおける性能比較結果 (元論文のFigure 5より引用)

結果は上図のとおりで,TRAILは唯一すべてのタスクでエキスパートの性能に到達しています.(1つのアルゴリズムに対して複数の線が描画されていて分かりづらいですね…)

2つ目の実験は,Liftタスクにおいて視覚的な疑似特徴を挿入した場合の性能を評価しています.具体的には,グリップの色について,エキスパートは明るい色,エージェントは暗い色にそれぞれ設定し,両者に変化をつけています.

f:id:twakamori:20201224221859p:plain
視覚的な疑似特徴がある場合の性能比較 (元論文のFigure 6より引用)

結果は上図のとおりです.こちらは,GAIL+AESとTRAILの2つで比較していて,TRAILの性能がGAIL+AESをはっきりと上回り,エキスパートの性能に達していることが分かります.両者のアルゴリズムの差分は制約集合の有無だけですので,制約集合の導入が性能向上に寄与したということが分かります.

3つ目の実験は,2つ目の実験と同じようにエキスパートとエージェントの観測に変化を付けますが,今度はグリップの色ではなく,背景にタスクと関係ないブロックを配置します.ここでは,Reedら*7によって提案された,ランダムに初期化された畳み込みネットワークあるいは畳み込みcriticネットワークを使用して視覚的な特徴を捉え,その上に小さな識別器を訓練する手法と比較します(それぞれ,GAIL+random, GAIL+criticと呼びます).

f:id:twakamori:20201224224115p:plain
背景にタスクと関係ないブロックがある場合の性能比較 (元論文のFigure 7より引用)

結果は上図のとおり,通常のLiftタスクは全ての手法でそこそこ良い性能が出ています(とはいえ,GAIL+AESとTRAILの性能が最も良い)が,背景にタスクと関係ないブロックがある場合(Lift distracted)は,TRAILのみがエキスパートの性能を達成しています.このように,明らかにタスクと関係ない(グリップの色のような)特徴でなくても,TRAILは疑似特徴を見つけ出し,区別することができているようです.

3つ目の実験では,制約集合を構築するために初期フレームを使用していました.4つ目の実験は,制約集合を構築するもう一つの戦略である,ランダム方策を用いる戦略について評価します.制約集合を構築する2つの戦略の違いを明らかにするために,2つ目の実験と3つ目の実験を組み合わせた難しいタスク(グリップの色が異なる+背景にタスクと関係ないブロックがある)を用います.

f:id:twakamori:20201224225240p:plain
制約集合を構築する2つの戦略の性能比較 (元論文のFigure 8より引用)

結果は上図のとおり,このような困難なタスクではTRAILも性能向上に時間がかかっています.また2つの戦略の性能はよく似ています.このため,著者らは追加データ収集の必要のない,初期のフレームを使用する方法を推奨しています.

議論はある?

特にないようですが,early stoppingや汎化性能,最適なDiscriminatorを用いた場合との比較等について,ablation studiesを実施しています. また,supplementary materialに実験設定やデータ拡張などのチューニング,ネットワーク構造やハイパーパラメータが詳細に記載されています.

次に読むべき論文は?


以上,Task-Relevant Adversarial Imitation LearningというDeepMindの論文を紹介いたしました.

個人的な所感としては,タスクが画像に限定されているものの,問題提起やアプローチがシンプルで分かりやすく,実装しやすそうな点が良いと思いました(本当は実装しようと思ったんですが,実装力と時間が足りなくて間に合いませんでした). ただ,GAILの改善系アルゴリズムは複数提案されているにもかかわらず,素のGAILとしか比較していないところがちょっと気になりました(実際,そのような理由でICLR2020をRejectされているようでした).