EXPO: Stable Reinforcement Learning with Expressive Policies

"EXPO: Stable Reinforcement Learning with Expressive Policies" 논문에 대한 상세한 기술적 요약은 다음과 같습니다.

1. 문제 정의

이 논문은 오프라인 데이터셋이 주어졌을 때, **온라인 강화학습 (RL)**을 사용하여 **표현력 있는 정책 클래스 (예: 확산 정책 및 흐름 매칭 정책)**를 미세 조정하는 과제를 다룹니다.

• 핵심 갈등: 표현력 있는 정책은 모방 학습 (IL) 을 통해 복잡한 다중 모드 행동 분포를 포착하는 데 뛰어나지만, RL 로 직접 최적화하기는 어렵습니다. 단순한 가우시안 정책과 달리, 표현력 있는 정책은 긴 탈노이즈 체인으로 매개변수화됩니다. 이러한 체인을 통해 가치 기울기를 역전파하여 정책 매개변수를 업데이트하는 것은 계산 비용이 많이 들고 불안정하여, 종종 학습 붕괴나 수렴 저하로 이어집니다.
• 격차: 기존 온라인 RL 방법은 단순한 가우시안 정책을 위해 설계되어 사전 훈련된 표현력 있는 모델의 능력을 활용하지 못합니다. 반면, 표현력 있는 정책을 미세 조정하려는 시도들은 종종 안정성 문제를 겪거나 탈노이즈 과정에 복잡한 수정을 요구합니다.
• 목표: 표현력 있는 정책의 매개변수를 가치 함수에 대해 직접 최적화하지 않고, 오프라인 데이터를 사용하여 사전 훈련된 표현력 있는 정책을 미세 조정할 수 있는 샘플 효율적이고 안정적인 온라인 RL 알고리즘을 개발하는 것입니다.

2. 방법론: EXPO (Expressive Policy Optimization)

저자들은 표현력 있는 정책을 직접적인 가치 최대화에서 분리하는 EXPO 프레임워크를 제안합니다. 표현력 있는 정책을 통해 기울기를 역전파하는 대신, EXPO 는 두 가지 구성 요소를 사용하여 Q 값을 최대화하는 "실시간 (On-the-Fly, OTF)" 정책을 구축합니다.

A. 두 구성 요소 정책 아키텍처

1. 기본 표현력 있는 정책 ($\pi_{base}$):
   • 역할: 행동 분포를 기반으로 초기 행동 후보를 생성합니다.
   • 학습: 안정적인 모방 학습 (IL) 목적 함수 (예: 탈노이즈 확산 확률 모델) 를 사용하여 훈련됩니다. 이 정책은 절대 Q 값을 최대화하도록 명시적으로 훈련되지 않아, 안정성과 복잡한 모드 포착 능력을 유지합니다.
2. 경량 편집 정책 ($\pi_{edit}$):
   • 역할: $\pi_{base}$에서 샘플링된 행동 $a$를 입력받아 잔차 편집 $\hat{a}$를 출력하여 정제된 행동 $\tilde{a} = a + \hat{a}$를 생성하는 소형 가우시안 정책입니다.
   • 학습: 엔트로피 정규화와 함께 표준 RL 정책 기울기 ( $Q(s, a + \hat{a})$ 최대화) 를 사용하여 훈련됩니다.
   • 제약: 편집은 하이퍼파라미터 $\beta$로 스케일링되어 원래 행동과 가깝게 유지되도록 제한됩니다. 이를 통해 편집 정책이 전역 최적화 문제가 아닌 국소 최적화 문제를 해결하도록 하여 효율적이고 안정적인 훈련을 가능하게 합니다.

B. 실시간 (OTF) 정책 추출

샘플링 및 시간차 (TD) 백업에 사용되는 실제 정책은 단일 신경망이 아니라 동적 선택 과정입니다.

1. 샘플링: 주어진 상태 $s$에 대해 시스템이 $\pi_{base}$에서 $N$개의 행동을 샘플링합니다.
2. 편집: 각 샘플링된 행동을 $\pi_{edit}$가 정제합니다.
3. 선택: 시스템은 원래 행동 ($a_i$) 과 편집된 행동 ($\tilde{a}_i$) 에 대한 Q 값을 모두 평가합니다.
4. 최대화: 모든 후보 중 가장 높은 Q 값을 가진 행동을 실행용으로 선택하고, 이를 TD 백업의 목표값으로 사용합니다.

C. 데이터 제한 영역을 위한 엔트로피 백업

소규모 오프라인 데이터셋이 있는 시나리오에서 저자들은 엔트로피 백업을 도입합니다. 확산과 같은 표현력 있는 정책은 폐형식 엔트로피가 없으므로, Q 값으로 가중치를 둔 $N$개의 기본 + 편집 행동 후보에 대한 소프트 샘플링 분포를 구성합니다. 이를 통해 분석적 엔트로피 계산 없이도 RL 손실에 엔트로피 보너스를 포함하여 탐색을 장려할 수 있습니다.

3. 주요 기여

1. 안정적인 가치 최대화 전략: 이 논문은 표현력 있는 탈노이즈 체인을 통한 직접 기울기 전파를 피하는 새로운 접근법을 제시합니다. 가치 최대화를 경량 편집 정책과 선택 메커니즘으로 위임함으로써 EXPO 는 직접 최적화가 실패하는 상황에서 안정성을 달성합니다.
2. 샘플 효율성: EXPO 는 온라인 RL 및 오프라인 - 온라인 미세 조정 설정 모두에서 기존 최첨단 방법 대비 2~3 배의 샘플 효율성 향상을 달성합니다.
3. 정책 무관성: 이 프레임워크는 기본 정책의 특정 매개변수화 (확산, 흐름 매칭 등) 에 무관하며 사전 훈련된 가치 함수가 필요하지 않아, 모든 사전 훈련된 로봇 정책에 적용 가능합니다.
4. 실시간 TD 백업: 저자들은 TD 목표 (샘플링뿐만 아니라) 에 가치 최대화 행동을 사용하는 것이 학습 속도에 결정적임을 입증하여, 표현력 있는 정책의 맥락에서 SARSA 유사 업데이트와 Q-러닝 유사 업데이트 간의 격차를 효과적으로 해소했습니다.

4. 실험 결과

이 방법은 희소 보상을 특징으로 하는 4 개 도메인 (AntMaze, Adroit, Robomimic, MimicGen) 의 12 가지 까다로운 연속 제어 작업에서 평가되었습니다.

• 온라인 RL (사전 훈련 없음): EXPO 는 RLPD, IDQL, QSM 을 포함한 베이스라인보다 샘플 효율성 측면에서 크게 우위를 보였으며, 오프라인 데이터셋을 사전 훈련에 활용하지 않더라도 종종 더 빠르게 수렴했습니다.
• 오프라인 - 온라인 미세 조정:
  • EXPO 는 다른 방법들에서 흔히 볼 수 있는 성능 붕괴 (예: DAC 붕괴, IDQL 정체) 없이 사전 훈련된 정책을 성공적으로 미세 조정했습니다.
  • 바늘에 실을 꿰기나 블록 쌓기와 같은 정밀 조작 작업에서 다른 방법들이 초기 모방 학습 성능을 넘어서는 개선에 어려움을 겪을 때, EXPO 는 높은 성능을 유지했습니다.
• 절대 실험 (Ablation Studies):
  • OTF TD 백업: TD 목표에 대한 가치 최대화 선택을 제거하고 (샘플링만 사용) 성능이 급격히 저하되어, 백업에서 Q 값 정렬의 중요성이 확인되었습니다.
  • 행동 편집: 편집 정책을 제거하면 좁은 행동 분포를 넘어선 탐색이 필요한 작업에서 특히 성능이 저하되었습니다.
  • 오프라인 데이터 품질: 성능은 오프라인 데이터셋의 품질과 강한 상관관계를 보였습니다. 그러나 엔트로피 백업 변형은 모방 학습만으로는 10% 미만의 성공률만 달성했던 매우 열악한 오프라인 데이터에서도 EXPO 가 성공할 수 있게 했습니다.

5. 의의

EXPO 는 "표현력 있는 정책 미세 조정" 병목 현상을 해결함으로써 로보틱스 및 RL 분야에서 중요한 진전을 이루었습니다.

• 실용적 영향: 인간 데이터를 모방하는 데 뛰어난 대규모 사전 훈련된 생성 모델을, 온라인 상호작용을 통해 적응하고 개선해야 하는 실제 환경에 배포할 수 있게 합니다.
• 이론적 통찰: 복잡한 생성 체인을 통한 역전파가 RL 에 필수적이라는 관념에 도전하여, "국소 편집 + 전역 선택" 전략이 더 안정적이고 효율적임을 보여줍니다.
• 미래 방향: 이 연구는 온라인 설정에서 모든 사전 훈련된 생성 정책 (확산, 흐름 등) 을 미세 조정할 수 있는 문을 열어, 로보틱스 분야에서 기반 모델의 채택을 가속화할 잠재력을 가지고 있습니다.

한계: 이 방법은 각 배치 업데이트마다 여러 행동을 샘플링해야 하므로 계산 비용이 많이 듭니다. 또한, 합리적인 사전 지식 (오프라인 데이터셋 또는 사전 훈련된 정책) 이 존재한다고 가정하며, 사전 지식 없이 처음부터 학습하는 것을 목적으로 하지 않습니다.