[논문 리뷰] TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning

Ruijie Zheng, Xiyao Wang외 6명|arXiv (Cornell University)|2023. 06. 22.

Reinforcement Learning in Robotics인용 수 5

한 줄 요약

TACO는 현재 상태와 행동 시퀀스, 그리고 그에 대응하는 향후 상태 간 상호정보를 최대화하여 상태 및 행동 표현을 동시에 학습하는 시간적 대비 학습 프레임워크를 도입한다. 원자료 픽셀에서의 시각적 연속 제어 작업에서 100만 번의 환경 단계 이후 SOTA 대비 40%의 성능 향상을 달성하며, 다양한 품질의 데이터셋에서 오프라인 RL 성능을 향상시킨다.

ABSTRACT

Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO: Temporal Action-driven Contrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.

연구 동기 및 목표

• 원자료 픽셀에서의 시각적 강화학습에서의 샘플 비효율성 문제를 해결하기 위해.
• 최적의 정책 및 가치 함수를 위한 제어 관련 정보를 포착하는 통합된 상태 및 행동 표현을 학습하기 위해.
• 이전의 대비 학습 방법이 정책에 의존하는 양성 쌍에 의존하거나 연속 제어에서 행동 표현을 忽略하는 한계를 극복하기 위해.
• 기존의 온라인 및 오프라인 RL 알고리즘과 호환되는 즉시 사용 가능한 모듈을 제공하기 위해.
• 이론적으로 탄탄한 상호정보 최대화를 통해 고차원 연속 제어 작업에서의 안정성과 일반화 능력을 향상시키기 위해.

제안 방법

• TACO는 현재 상태 표현과 행동 시퀀스 표현, 그리고 그에 대응하는 향후 상태 표현 간의 상호정보를 최대화하기 위해 InfoNCE 손실 기반의 대비 학습 목표를 사용한다.
• 시간적 동역학과 제어 관련 정보를 포착하는 잠재 표현을 생성하기 위해 상태 인코더와 행동 인코더를 함께 훈련한다.
• 행동 시퀀스와 그에 대응하는 향후 상태를 양성 쌍으로 간주하고, 다른 시퀀스는 음성 샘플로 간주한다.
• 대비 목표는 정책에 종속되지 않게 설계되어, 정책에 의존하는 양성 쌍으로 인한 불안정성을 방지한다.
• TACO는 어떤 오프더쉐프 RL 알고리즘과도 호환되며, 온라인 및 오프라인 RL 훈련 파이프라인에 원활하게 통합될 수 있다.
• 이론적으로 탄탄한 접근으로, 학습된 표현이 최적 제어를 위한 충분한 정보를 포함하고 있음을 보여준다.

실험 결과

연구 질문

• RQ1 통합된 상태 및 행동 표현 학습이 시각적 강화학습에서 샘플 효율성을 향상시킬 수 있는가?
• RQ2 행동 시퀀스를 통합한 정책에 종속되지 않는 대비 목표는 더 안정적이고 일반화 가능한 표현을 이끌어내는가?
• RQ3 TACO는 온라인 연속 제어 벤치마크에서 SOTA 모델리스 및 모델기반 시각적 강화학습 알고리즘을 초월할 수 있는가?
• RQ4 TACO는 다양한 품질의 데이터셋에서 오프라인 시각적 강화학습의 성능을 어느 정도 향상시킬 수 있는가?
• RQ5 TACO의 상호정보 최대화 방식은 안정성과 표현 품질 측면에서 직접적인 잠재 동역학 모델링 방식과 비교해 어떻게 다를까?

주요 결과

• TACO는 9개의 DeepMind Control Suite 작업에서 100만 번의 환경 상호작용 단계 이후 SOTA 모델리스 시각적 강화학습 알고리즘 대비 평균 40%의 성능 향상을 달성한다.
• DrQ-v2와 결합했을 때, TACO는 온라인 시각적 강화학습 환경에서 가장 강력한 모델기반 베이스라인인 Dreamer-v3를 능가한다.
• TD3+BC 및 CQL와 통합했을 때 TACO는 다양한 데이터 품질의 오프라인 데이터셋에서 성능을 향상시켜, 오프라인 데이터셋 전반에서 새로운 SOTA를 수립한다.
• 정책에 종속되지 않는 대비 목표 덕분에 안정적인 훈련을 보이며, 정책에 의존하는 양성 쌍으로 인한 불안정성을 피한다.
• TACO의 통합된 상태-행동 표현 학습은 연속 제어 작업에서 더 나은 일반화 능력과 장기 계획 능력을 가능하게 한다.
• 효과적임에도 불구하고 TACO는 큰 배치 크기를 요구하여 계산 효율성에 영향을 미치며, 향후 최적화의 방향성이 된다.