Performance Asymmetry in Model-Based Reinforcement Learning

1. 문제 정의 (Problem)

최근 모델 기반 강화학습 (MBRL) 은 Atari100k 벤치마크에서 평균적으로 인간을 능가하는 성능을 달성했습니다. 그러나 저자들은 기존의 평균 지표 (Arithmetic Mean) 가 **성능 비대칭성 (Performance Asymmetry)**이라는 치명적인 문제를 숨기고 있음을 발견했습니다.

• 성능 비대칭성: MBRL 에이전트들은 특정 태스크 (Agent-Optimal, 인간 점수가 낮은 태스크) 에서는 인간을 압도적으로 능가하지만, 다른 태스크 (Human-Optimal, 인간 점수가 높은 태스크) 에서는 인간보다 현저히 낮은 성능을 보입니다.
• 현황: 최신 SOTA 모델인 DIAMOND(픽셀 기반 확산 모델) 는 전체 평균 HNS(Human Normalized Score) 가 1.46 으로 최상위권이지만, Human-Optimal 태스크에서는 모든 베이스라인 중 가장 낮은 성능을 기록했습니다. 두 태스크 군 간의 성능 격차는 무려 21 배에 달했습니다.
• 기존 평가 지표의 한계: 산술 평균 (Arithmetic Mean) 은 극단적인 값에 둔감하며, HNS 의 값 범위 (음수에서 8.3 까지) 가 넓어 높은 점수를 받은 소수 태스크가 낮은 점수의 다수 태스크를 과대평가하게 만듭니다. 이로 인해 특정 태스크에 과적합 (Overfitting) 된 문제가 드러나지 않습니다.

2. 방법론 (Methodology)

저자들은 이 문제를 해결하기 위해 데이터 분석, 새로운 평가 지표 제안, 그리고 새로운 아키텍처 (JEDI) 를 도입했습니다.

A. 태스크 분류 및 새로운 평가 지표 (Sym-HNS)

• 태스크 분할: Atari100k 의 26 개 태스크를 기존 베이스라인들의 평균 성능을 기준으로 **Agent-Optimal(13 개)**과 **Human-Optimal(13 개)**로 균등하게 분할했습니다.
  • Human-Optimal 태스크 특징: 행동 공간 (Action Space) 이 크고, 'Shooter'(총격) 태스크가 많아 복잡도가 높음.
• Sym-HNS 도입: 기존의 산술 평균 대신 **조화 평균 (Harmonic Mean)**을 기반으로 한 새로운 집계 지표인 Sym-HNS를 제안했습니다.
  • 공식: $Sym-HNS = \frac{2 \cdot \mu_{AO} \cdot \mu_{HO}}{\mu_{AO} + \mu_{HO}}$
  • 효과: 한쪽 태스크 군의 성능이 극도로 낮으면 전체 점수가 크게 떨어지도록 설계되어, 균형 잡힌 성능 향상을 유도하고 성능 비대칭성을 명확히 드러냅니다.

B. 원인 분석

• 픽셀 기반 에이전트의 실패 원인: Human-Optimal 태스크에서 픽셀 기반 에이전트 (DIAMOND 등) 가 실패하는 주된 원인은 **차원의 저주 (Curse of Dimensionality)**입니다. 고해상도 이미지 입력과 복잡한 행동 공간 (특히 Shooter 태스크) 이 결합되어 학습 효율이 급격히 떨어집니다.
• 확산 모델 (Diffusion) 의 양면성: 확산 모델은 시각적 디테일이 풍부한 Agent-Optimal 태스크 (예: Breakout) 에서는 뛰어난 성능을 내지만, 복잡한 논리적 추론이 필요한 Human-Optimal 태스크에서는 오히려 성능을 저하시킵니다.

C. 제안 모델: JEDI (Joint Embedding DIffusion)

성능 비대칭성을 역전시키고 효율성을 높이기 위해 JEDI라는 새로운 잠재 공간 (Latent Space) 기반의 세계 모델을 제안했습니다.

• 엔드 투 엔드 잠재 확산 (End-to-End Latent Diffusion):
  • 픽셀 입력을 직접 처리하는 대신, 인코더를 통해 저차원 잠재 공간 (Latent Space) 으로 압축한 후 확산 모델을 학습합니다.
  • JEPA(Joint Embedding Predictive Architecture) 스타일 학습: 기존 재구성 (Reconstruction) 손실이나 은닉 상태 추출, 증류 (Distillation) 없이, 잠재 확산 모델의 데노이징 (Denoising) 손실 그라디언트를 직접 사용하여 인코더를 학습시킵니다.
  • JEPA 와 확산의 호환성 증명: 기존 JEPA 는 맞춤형 손실 함수가 필요하다는 통념을 깨고, 오프더셸 (Off-the-shelf) 확산 데노이징 손실이 JEPA 프레임워크와 직접 호환됨을 처음 보였습니다.
• 기술적 세부 사항:
  • 인코더와 확산 모델 (Dynamics) 을 결합하여 학습하며, 표현 붕괴 (Representation Collapse) 를 방지하기 위해 Stop-grad 와 학습률 스케일링을 적용했습니다.
  • 잠재 공간에서의 안정성을 위해 Tanh 클램핑 (Clamping) 과 랜덤 스위칭 기법을 도입했습니다.

3. 주요 기여 (Key Contributions)

1. 성능 비대칭성 발견: MBRL 에이전트가 특정 태스크에 과적합되어 인간이 잘하는 태스크에서는 실패한다는 사실을 최초로 규명했습니다.
2. 균형 잡힌 평가 체계: Atari100k 태스크를 Agent/Human-Optimal 로 분할하고, 조화 평균을 활용한 Sym-HNS를 제안하여 성능 비대칭성을 정량화할 수 있는 기준을 마련했습니다.
3. JEDI 모델 제안: JEPA 스타일의 학습 프레임워크와 잠재 공간 확산 모델을 결합한 새로운 세계 모델을 개발했습니다. 이는 재구성 손실 없이 인코더를 학습시키는 최초의 방법입니다.
4. 성능 역전 및 효율성 증대: JEDI 는 SOTA 픽셀 기반 모델의 성능 비대칭성 추세를 역전시켰으며, 계산 효율성 (GPU 메모리 43% 감소, 학습 속도 2 배 향상) 을 크게 개선했습니다.

4. 결과 (Results)

• Sym-HNS 및 Human-Optimal 성능: JEDI 는 Sym-HNS 에서 SOTA 를 기록했으며, Human-Optimal 태스크에서 기존 모델들을 압도했습니다. 특히 인간이 잘하는 BankHeist 태스크에서 인간을 능가하는 성능을 처음 달성했습니다.
• Agent-Optimal 및 전체 성능: JEDI 는 시각적으로 복잡한 태스크 (Breakout 등) 에서도 경쟁력 있는 성능을 유지하여 전체 Atari100k 벤치마크에서 최상위권 (Mean, IQM, Optimality Gap 기준) 을 기록했습니다.
• 효율성:
  • GPU 메모리: DIAMOND 대비 43% 감소 (12 배 잠재 압축 효과).
  • 학습/추론 속도: 학습 속도는 2 배, 추론 속도는 3 배 빨라졌습니다.
• Qualitative 분석: JEDI 는 복잡한 Shooter 태스크 (BankHeist, Hero, DemonAttack) 에서 픽셀 기반 모델 (DIAMOND) 이 자주 저지르는 자살 (Self-destruction) 행동을 줄이고, 적을 효과적으로 추적 및 제압하는 정책을 학습했습니다.

5. 의의 및 결론 (Significance)

이 논문은 MBRL 의 발전이 단순히 평균 점수의 상승이 아니라, 모든 태스크에 대한 균형 잡힌 일반화 능력이 필요함을 강조합니다.

• 평가 지표의 패러다임 전환: 기존의 산술 평균 중심 평가가 가진 맹점을 지적하고, 조화 평균을 활용한 새로운 평가 기준 (Sym-HNS) 을 제시함으로써 향후 연구 방향을 제시했습니다.
• 알고리즘적 통찰: 픽셀 입력의 차원의 저주가 복잡한 태스크에서의 실패 원인이었음을 규명하고, 잠재 공간 확산 모델이 이를 해결할 수 있음을 증명했습니다.
• 미래 연구 방향: 인간이 잘하는 태스크와 에이전트가 잘하는 태스크의 근본적인 차이와 알고리즘 내 성공/실패 메커니즘에 대한 심층 연구를 촉발시켰습니다.

결론적으로, JEDI 는 모델 기반 강화학습이 특정 태스크에 편향되지 않고, 인간 수준의 복잡한 태스크에서도 견고한 성능을 발휘할 수 있는 새로운 방향성을 제시한 획기적인 연구입니다.