[논문 리뷰] Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

Rishabh Agarwal, Marlos C. Machado외 2명|arXiv (Cornell University)|2021. 01. 13.

Reinforcement Learning in Robotics참고 문헌 64인용 수 27

한 줄 요약

본 논문은 정책 유사성 지표(PSM)와 대조 학습 절차를 도입하여 정책 유사성 임베딩(PSE)을 만들어 보지 않은 환경에서의 RL 정책 일반화를 향상시킨다.

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의 순차적 구조를 표현 학습에 통합하여 환경 간 일반화를 개선한다.
• 상태/환경 전반의 장기 최적 행동의 유사성을 포착하는 보상 무관 정책 유사성 지표(PSM)를 정의한다.
• PSM을 기반으로 한 상태 유사성을 행동적으로 클러스터링하는 임베딩(PSE)을 학습하기 위해 대조적 objectives를 사용한다.
• 픽셀 기반 점프 작업, 스푸리어스 상관관계가 있는 LQR, Distracting DM Control Suite를 포함한 다양한 벤치마크에서 일반화 개선을 시연한다.

제안 방법

• 최적 정책 간의 Dist와 다음 상태 분포에 대한 1-Wasserstein 항을 사용하여 정책 유사성 지표(PSM)를 정의하고, x와 y에 대한 재귀적 식 d*(x,y)를 도출한다.
• Dist(이산 행동의 경우 TV, 연속 행동의 경우 평균 행동에 대한 L1 등)와 Wasserstein 항을 사용하여 환경 간 정책 전이 오차를 하한으로 bounding 한다(정리 1).
• d를 Gauss 커널을 통해 유사도 Gamma로 변환하고, 환경 간 상태 유사성을 존중하는 임베딩 z_theta(t)를 학습하기 위한 대조 학습 objective(CMEs)을 구성한다.
• PSM으로 CMEs를 구현하여 긴 시점의 최적 행동이 유사한 상태를 그룹화하는 정책 유사성 임베딩(PSE)을 얻고, RL imitation/training loss와 CME 보조 손실을 결합한다.
• 인코더, 프로젝터, 정책 헤드가 공동으로 학습되는 end-to-end 아키텍처를 제공한다(SimCLR에서 영감을 받은 대조 손실 + RL imitation 손실).
• Train MDP 쌍으로부터 CMEs를 계산하는 실용 알고리즘(Algorithm 1)을 제공하며, Gamma 하의 이웃-positive 샘플과 소프트 대조 손실을 사용한다.

실험 결과

연구 질문

• RQ1 보상 무관 정책 유사성 지표(PSM)가 RL에서 환경 간 일반화를 견고하게 제공할 수 있는가?
• RQ2 PSM 기반 상태 유사성을 인코딩하는 대조적으로 학습된 임베딩(PSE)이 표준 정규화나 데이터 증강 이상의 일반화를 개선하는가?
• RQ3 PSE가 다양한 일반화 벤치마크(pixel 기반 점프 작업, 방해 요소가 있는 LQR, Distracting DM Control Suite)에서 기존 방법과 비교하여 어떤 성능을 보이는가?

주요 결과

데이터 증강	방법	넓은	좁은	무작위
✗ 드롭아웃 및 ℓ2 정규화	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는 여러 벤치마크에서 제로샷 일반화를 향상시키며, 여러 설정에서 정규화 및 이합 전이 기반 기준보다 우수하다.
• PSEs는 RandConv 데이터 증강과 결합할 때 일반화를 크게 향상시키며, 특정 구성에서는 증강만 사용하는 경우보다 우수할 수 있다.
• PSM과 함께 CMEs를 사용할 때의 차원 축소 연구에서 l2-임베딩이나 π*-bisimulation과 같은 대안보다 우수한 결과를 보이며, 학습된 불변성의 중요성을 강조한다.
• PSE는 ε-서브옵티멀 정책과 같은 상태에서도 견고하게 작동하며 정확한 π*가 없을 때도 근사적으로 활용될 수 있다.
• Jumping Task의 색상 수정 버전에서 PSE는 작업 의존적 불변성을 포착하고 RandConv 단독보다 일반화를 계속 우수하게 유지하며, 증강은 작업별 구분을 흐리게 할 수 있다.
• 방해 요소가 있는 LQR 및 Distracting DM Control Suite에서 PSM 기반 표현은 진정한 최적 정책에 접근하지 못해도 강력한 일반화를 가능하게 한다.