[論文レビュー] Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

Rishabh Agarwal, Marlos C. Machado+2 more|arXiv (Cornell University)|Jan 13, 2021

Reinforcement Learning in Robotics参考文献 64被引用数 27

ひとこと要約

本論文は policy similarity metric (PSM) を導入し、対比学習手法を用いて policy similarity embeddings (PSEs) を作成し、 unseen environments での強化学習ポリシーの一般化を改善する。

ABSTRACT

Reinforcement learning methods trained on few environments rarely learn policies that generalize to unseen environments. To improve generalization, we incorporate the inherent sequential structure in reinforcement learning into the representation learning process. This approach is orthogonal to recent approaches, which rarely exploit this structure explicitly. Specifically, we introduce a theoretically motivated policy similarity metric (PSM) for measuring behavioral similarity between states. PSM assigns high similarity to states for which the optimal policies in those states as well as in future states are similar. We also present a contrastive representation learning procedure to embed any state similarity metric, which we instantiate with PSM to obtain policy similarity embeddings (PSEs). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, a jumping task from pixels, and Distracting DM Control Suite.

研究の動機と目的

• RLの逐次構造を表現学習に組み込み、環境を跨ぐ一般化を向上させる。
• 報酬に依存しない policy similarity metric (PSM) を定義し、状態/環境間で長期的な最適挙動の類似性を捉える。
• PSMに基づく対比目的で、行動的に類似する状態をクラスタリングする埋め込み（PSEs）を学習する。
• ピクセルベースのジャンプタスク、偽の相関を含むLQR、Distracting DM Control Suite などの多様なベンチマークで一般化の改善を実証する。

提案手法

• 最適ポリシー間の距離 Dist と次状態分布上の1-Wasserstein項を用いてポリシー類似度PSMを定義し、d*(x,y)の再帰方程式を導出する。
• Distを用い（離散行動では総変動距離TV、連続行動では平均作用のL1など）、Wasserstein項と合わせて環境横断のポリシー転送誤差を上界化する（定理1）。
• dをガウス核を用いて類似度Gammaへ変換し、環境間で状態の類似性を保つ埋め込みz_theta(t)を学習する対比学習目的（CMEs）を構築する。
• CMEsをPSMで具体化し、長期的な最適挙動が類似する状態をグルーピングする policy similarity embeddings (PSEs) を得る。RLの模倣/訓練損失とCME補助損失を組み合わせる。
• エンコーダ、プロジェクター、ポリシーヘッドを共同訓練するエンドツーエンドのアーキテクチャを提供する（SimCLR風の対比損失＋RL模倣損失）。
• Gamma下の最近傍ポジティブを用い、ソフト対比損失で訓練する訓練MDPの対からCMEsを計算する実用的なアルゴリズム（Algorithm 1）を提供する。

実験結果

リサーチクエスチョン

• RQ1 報酬に依存しないポリシー類似度指標（PSM）は、RLにおける環境横断の頑健な一般化を提供できるか。
• RQ2 PSMベースの状態類似性を符号化する対比学習埋め込み（PSEs）は、標準的な正則化やデータ拡張を超える一般化を改善するか。
• RQ3 既存手法と比較して、PSEsは多様な一般化ベンチマーク（ピクセルベースのジャンプタスク、妨害因子を含むLQR、Distracting DM Control Suite）でどのように性能を示すか。

主な発見

Data Augmentation	Method	Wide	Narrow	Random
✗ Dropout and ℓ2 reg.	Bisimulation Transfer	17.8 (2.2)	10.2 (4.6)	9.3 (5.4)
✓ RandConv	PSEs	33.6 (10.0)	9.3 (5.3)	37.7 (10.4)
✗ RandConv + π*-bisimulation	RandConv + π*-bisimulation	41.4 (17.6)	17.4 (6.7)	33.4 (15.6)
✓ RandConv + PSEs	RandConv + PSEs	87.0 (10.1)	52.4 (5.8)	83.4 (10.1)

• PSEsは複数のベンチマークでゼロショット一般化を改善し、いくつかの設定で正則化および bisimulation-transfer のベースラインを上回る。
• PSEsはRandConvデータ拡張と組み合わせたとき一般化を大幅に向上させ、いくつかの構成では拡張だけを超えることがある。
• PSMと共にCMEsを用いると、l2-embeddingやπ*-bisimulationのような代替手法を上回るアブレーションが得られ、学習された不変性の重要性を浮き彫りにする。
• PSEsは適度なポリシーのサブ最適性（epsilon-suboptimal policies）に対して頑健で、厳密なポリシーが利用できない場合でも近似的なπ*を活用できる。
• Jumping Taskのカラー変更バリアントでは、PSEsがタスク依存の不変性を捉え、RandConvだけの場合よりも優れた一般化を維持する。拡張はタスク固有の差異をぼかしてしまうことがある。
• 妨害因子を含むLQRとDistracting DM Control Suite全体で、PSMベースの表現は真の最適政策にアクセスできなくても強い一般化を可能にする。