ATC: 教師なし学習を強化学習から切り離す！！

3つの要点
✔️ RL Agentの学習から切り離すことが出来る表現学習の手法の提案
✔️ 新しい教師なし学習のタスク Augmented Temporal Contrastを提案
✔️ ATCで事前学習されたencoderを用いて高いパフォーマンスを様々なタスクで達成

Decoupling Representation Learning from Reinforcement Learning
written by Adam Stooke, Kimin Lee, Pieter Abbeel, Michael Laskin
(Submitted on 9 Mar 2020)
Comments: Accepted at arXiv
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computer Vision and Pattern Recognition (cs.CV); Machine Learning (stat.ML)
Paper  Official Code COMM Code

はじめに

本記事では、Decoupling Representation Learning From Reinforcement Learningという、Pixel-based RL (入力が画像) 従来の、End-to-EndでRLと一緒にEncoderを学習するというrewardドリブンな特徴量を学習するのではなく、unsupervised learning (教師なし学習) を使って強化学習から切り離した上で、Encoderを事前学習しても従来の手法と同様、もしくはより良いパフォーマンスを出すことを示した論文を紹介します。過去のいくつかの論文では、RL本来のloss関数に加えて、特徴量を学習するようなunsupervised taskを用いてauxiliary lossを導入することによって学習させることにより、パフォーマンスの向上をさせました。また、コンピュータービジョンの界隈では、最近、unsupervisedやself-supervised learningを使うことによってImageNetのクラス分類などで有効なことを示す論文が出てきています。よって、RLも報酬を使わない特徴量を学習する方法があれば、パフォーマンス向上においてとても柔軟になると考えられます。そこで、この論文では、Augmented Temporal Contrast (ATC)と呼ばれる、convolutional encoderをある時間の画像とそれより少し、先の時間の画像をペアとして、image augmentationとcontrastive lossを使うことで学習する、新しいunsupervised learningのタスクを提案しました。この新しいunsupervised learningを利用した、事前学習の手法と、それによる結果を合わせて詳しく説明してきます。

手法 

この章では、新しく提案されたAugmented Temporal Contrast (ATC)について紹介します。ATCは下図のような構造となっており、入力となるobservation $o_{t}$と、近い未来におけるobservation $o_{t+k}$のペアが入力として必要となります。そして、それぞれの学習バッチに対して、stochatic data augmentationをobservationに対してかけます。この論文では、とてもシンプルで実装が簡単である、random shiftを採用しています。そして、augmentationがかけられたobservationはencoderによって、小さい潜在空間に移され、その特徴量に対してcontrastive lossが利用されlossが計算されます。このタスクは、encoderがobservationからMDPの構造に関する意味のある情報を抽出することを目的としています。

このATCは、下図が表すように、4つのコンポーネントに分けられます。１つ目は、convolutional encoder $f_{\theta}$でobservation $o_{t}$を入力とし、潜在特徴量 $z_{t} = f_{\theta}(\text{AUG}(o_{t}))$を出力します。２つ目は、linear global compressor $g_{\phi}$と呼ばれる、小さい潜在コードベクトル (latent code vector)を出力するものです。3つ目は、residual predictor MLP $h_{\psi}$という、$p_{t} = h_{\psi}(c_{t})+c_{t}$のような操作を行うforward modelを表すものです。そして最後に、contrastive transformation matrix $W$です。$o_{t+k}$ は$\bar{c}_{t+k}=g_{\bar{\phi}}\left(f_{\bar{\theta}}\left(\operatorname{AUG}\left(o_{t+k}\right)\right)\right.$という処理を経て、target codeとして表現されます。このobservation $o_{t+k}$を処理するパラメーター $\bar{\theta}$と$\bar{\psi}$は、momentum encoderと呼ばれる、学習しているパラメーターである$\theta$や$\psi$のmoving averageによって計算されます。

$$\bar{\theta} \leftarrow(1-\tau) \bar{\theta}+\tau \theta ; \quad \bar{\phi} \leftarrow(1-\tau) \bar{\phi}+\tau \phi$$

そして、このタスクを利用してencoderを事前学習した後、RL agentでは$f_{\theta}$のみが使われます。このタスクのloss関数は、InfoNCEを利用して以下のように表すことが出来ます。

$$\mathcal{L}^{A T C}=-\mathbb{E}_{\mathcal{O}}\left[\log \frac{\exp l_{i, i+}}{\sum_{o_{j} \in \mathcal{O}} \exp l_{i, j+}}\right]$$

ここで$l$は、$l=p_{t} W \bar{c}_{t+k}$と表されます。

 

  

実験

評価に用いる環境とアルゴリズム

本論文では、ATCをDeepMind control suite (DMControl)、Atari games、そしてDeepMind Labの３種類の環境において評価しました。Atariは、discrete control（右、左などの離散的なアクション）を用いるゲームで、DMControlはcontinuous controlを用いたゲームです。そして、DMLabは、DMLabは、RL Agentが視覚的に複雑な3Dの迷路の環境において行動するタスクです。

また、アルゴリズムに関しては、ATCをon-policyとoff-policy RLアルゴリズム両方に対して利用し評価しました、DMControlに対してはRAD-SACと呼ばれるrandom shift augmentationを利用したsoft actor criticがベースのアルゴリズムを使いました。AtariとDMLabに対してはPPOを使って評価しました。 

Online RL とUnsupervised Learning

DMControl:この実験を通して、encoderをRL agentから切り離し、encoderはATCのタスクに関してonlineで学習しつつ、同時にRL agentをencoderからのrepresentationを受け取り学習を行った場合、end-to-end RLの場合 (ATCを使わず従来の方法での学習)と比べて、パフォーマンスが同等もしくは、良い結果を出していることが下図の結果から分かります。

DMLab: DMLabを使った実験では、EXPLORE_GOAL_LOCATIONSという迷路の中を動き回りゴールを探すタスクと、LASERTAG_THREE_OPPONENTSと呼ばれる素早く敵を追いかけ、敵に対してタグを付けるタスクの２種類において、迷宮の大きさが異なる２つのレベル、計４つのタスクに関して実験を行いました。EXPLOREに関するタスクにおいて、ゴールを見つける頻度が低く、学習が進みにくいので、$p \propto 1+R_{abs}$ (ここで$R_{a b s}=\sum_{t=0}^{n} \gamma^{t} \mid r_{t}$)の値に関して優先順位をつけて、優先順位が高いものを優先的にsamplingをして学習に利用する方法を実験では使っています。この方法を使ったものをprioritized-UL（下図 オレンジ線）とすると、これは従来のRL (下図　赤線) と比べてパフォーマンスが同等もしくは良い結果となっており、またauxiliary taskとしてUnsupervised Learningを使う、つまりencoderをRL agentから切り離して学習しない場合と比べても精度がほぼ同じになっています（緑のライン)。

またLASERTAGに関しては、学習するagentは敵を良く見かけることは出来ますが、tagを敵に対してつけるのはなかなか起きないイベントです。従来のRLの場合 (下図右 赤線)は、特に大きい迷宮の場合、パフォーマンスがとても低く、ATCを使いRL agentから切り離したもののほうがパフォーマンスが高くなっています。また、encoderを更新する頻度を２倍に上げることで、encoderをRL agentから切り離さないで学習した場合と比較した時に、パフォーマンスがほぼ同様になることが分かりました。

 Atari: Atariを使った実験では８種類のゲームにおいて評価をし、そのうちの5つのゲームでend-to-end RLと比べて、EncoderをRL Agentから切り離してATCのタスクを使って学習した場合のほうが、パフォーマンスが高いことを示しました。ただし、BREAKOUTとSPACE INVADERSにおいては、逆にパフォーマンスが上がりませんでした。ただし、ATCをauxiliary taskとしてRL Agentから切り離さずに学習させた場合はこれらのタスク、そして他のタスクも良いパフォーマンスを示しました。また、その他にも最初の100kのtransition (遷移)に関してのみATCを使ってencoderを学習し、weightを初期化する目的で使った場合も、いくつかのゲームで良い結果を示したことから、ATCが有効的なことを表しています。

Encoderの事前学習に関するベンチマーク

本論文では、RLのための異なるunsupervised learning algorithmに対する評価として、新しい評価方法を提案しています。この評価は以下の手順で行われます。

1. それぞれの環境から、expert demonstrationを集める
2. CNN encoderを集めたデータを利用し、unsupervised learningを利用して事前学習を行う。
3. Encoderのweightを固定した上で、RL agentを学習し、そのパフォーマンスを評価する。

では、それぞれの環境に関して結果を紹介していきます。

DMControl: この実験では、ATCと比較するunsupervised learning algirithmとして主にCURLで利用されるAugmented Contrast (AC)、そしてVAEを用いました。下図の結果の通り、 ATCを利用した場合、基本的に他のunsupervised algorithmと比較して同等もしくはより良いパフォーマンスを示しています。また、ATCが全てのタスクにおいて、end-to-endのRLのパフォーマンスと比較して同等もしくはより良いパフォーマンスを示しています。

DMLab: 下図はDMLabに関する結果です。この実験ではATCに対して、DMLabの環境で効果的とされるCPCとPixel Controlという手法を用いて比較しました。EXPLOREのタスクにおいては、パフォーマンスがどれも同じような結果ですが、Lasertagに関しては、ATCを使ったものが一番良い結果となっていることが下図から分かります。

Atari: Atariに関しては、4つのゲームに関して、ATCをPixel Control、VAE-T、Inverse model (２つの連続したobservationが与えた時に、そのペアのobservationに遷移するactionを予測することで学習するモデル)、Spatio-Temporal Deep InfoMax (ST-DIM)と呼ばれる手法と比較をしました。結果は下図が示すように、ATCだけが、3つのゲームにおいて、end-to-end RLよりも高いパフォーマンスを表していることが分かります。

  

マルチタスクにおけるEncoder

この実験では、ATCがmulti-task encoderを学習することができるかどうかを評価するための実験で、いくつかの環境に関するexpert demonstrationを利用してencoderを事前学習をし、そのencoderのweightを固定した上でRL agentを別のタスクに対して学習をし、そのパフォーマンスを評価するものです。

DMControl: この実験では、下図の上の段にあるタスク Ball-in-Cup: Catch, Cartpole: Swingup, Cheetah: Run, Walker: Walkのタスクのdemonstrationを用いてencoderを事前学習します。そのうえで、学習されたencoderのweightを固定し下図の下の段のタスクに関してそれぞれRL agentを学習した結果を表しています。予め別のタスクを用いてencoderを事前学習し、その学習されたencoderを利用することで、事前学習なしのRLと比較して、より効率的にRL agentを学習することが出来るということを下図の下の段の結果が示しています。

Ablation studyとEncoderに関する分析

本記事では、RL agentからencoderを切り離して事前学習を行った場合でも、end-to-end RLのパフォーマンスと比べて、同様もしくは向上するATCと呼ばれるunsupervised learningのタスクを紹介しました。実験を通して、いくつかの環境では良い結果を出していますが、効果があまり現れない環境も見受けられるので、より様々な環境で実験を行った上でさらなる分析が必要だと考えさせられる論文でした。ですが、このRL agentからencoderを切り離すことができ、事前学習が可能になると、より柔軟にそして効率的にRL agentを学習することが出来ると考えられるので、今後の研究に期待していきたいです。

この記事に関するカテゴリー

• 強化学習
• 表現学習