MAGE: Multi-scale Autoregressive Generation for Offline Reinforcement Learning

이 논문은 희소한 보상을 가진 장기 과제를 해결하기 위해 계층적 궤적 표현 학습과 다중 스케일 자기회귀 생성 방식을 결합하여 일관되고 제어 가능한 궤적을 생성하는 오프라인 강화학습 방법인 MAGE 를 제안합니다.

핵심

이 논문은 **'MAGE'**라는 새로운 인공지능 학습 방법을 소개합니다. 이 방법은 로봇이나 게임 캐릭터가 과거의 경험 데이터만 보고도 복잡한 미션을 잘 수행할 수 있도록 도와줍니다.

기존의 방법들은 긴 시간 동안 목표를 달성해야 하는 미션 (예: 로봇이 주방에서 요리를 하거나, 미로를 빠져나가는 것) 에서 자주 실패했습니다. 특히 보상이 드물게 주어지는 상황에서는 길을 잃기 쉽죠. MAGE 는 이 문제를 '거시적 관점과 미시적 관점을 동시에 잡는' 독특한 방식으로 해결합니다.

이해를 돕기 위해 여행 계획과 건축에 비유해서 설명해 드리겠습니다.

---

1. 기존 방법들의 한계: "지도 없이 걷기" vs "모든 걸 한 번에 그리기"

기존의 인공지능 학습 방법들은 크게 두 가지 문제가 있었습니다.

• 한 번에 모든 걸 결정하는 방법 (Decision Transformer 등): 마치 여행 계획 없이 "지금부터 100 시간 뒤까지의 모든 발걸음을 한 번에 결정하라"는 명령을 받은 것과 같습니다. 처음 몇 발은 잘 걸어도, 100 시간 뒤의 목표는 잊어버리고 엉뚱한 곳으로 갈 확률이 높습니다.
• 확률적으로 무작위로 만들어내는 방법 (Diffusion 모델 등): 안개 속에서 길을 찾는 것과 비슷합니다. 주변은 잘 보이지만, 전체적인 큰 그림 (전체 경로) 이 흐릿해서 벽을 뚫고 지나가거나, 목표와 반대 방향으로 가는 실수를 자주 합니다.

2. MAGE 의 핵심 아이디어: "건축가처럼 단계별로 설계하기"

MAGE 는 **'거시 (Macro)'**와 **'미시 (Micro)'**를 나누어 생각합니다. 마치 건축가가 건물을 지을 때 다음과 같이 작업하는 것과 같습니다.

1. 대략적인 청사진 그리기 (거시적 스케일):
   먼저 "이건 3 층짜리 빌딩이고, 1 층은 로비, 2 층은 식당, 3 층은 사무실이야"라고 큰 틀만 잡습니다. 이 단계에서는 벽의 정확한 위치나 창문의 크기 같은 세부 사항은 중요하지 않습니다. 중요한 건 **'전체적인 흐름과 목표'**입니다.

   • 비유: 여행 계획에서 "서울에서 부산까지 기차로 가고, 부산에서 제주도로 배를 타고, 제주도에서 해변으로 간다"는 큰 루트만 먼저 정하는 것입니다.

2. 세부 사항 채워 넣기 (미시적 스케일):
   큰 틀이 잡히면, 이제 그 틀 안에서 세부 사항을 채웁니다. "1 층 로비의 문은 어디에 있고, 계단은 몇 칸인가?"를 하나씩 정해 나갑니다. 이때 중요한 건, 아까 정한 큰 틀 (루트) 을 절대 잊지 않는다는 점입니다.

   • 비유: "부산에서 제주도로 가는 배를 타야 한다"는 큰 계획 아래에서, "어떤 배를 타고, 몇 시에 출발할지, 어떤 선실로 갈지"를 구체적으로 결정하는 것입니다.

3. MAGE 가 사용하는 두 가지 핵심 도구

이 아이디어를 실현하기 위해 MAGE 는 두 가지 기술을 사용합니다.

• 다중 스케일 자동 인코더 (Multi-scale Autoencoder):
  과거의 경험 데이터를 여러 개의 레이어로 쪼개는 도구입니다.

  • 레이어 1 (거시): "어디로 가야 하는가?"라는 큰 방향성만 담습니다.
  • 레이어 2, 3... (미시): "왼쪽으로 10 걸음, 오른쪽으로 5 걸음" 같은 구체적인 행동들을 담습니다.
    마치 사진을 확대해 보면 픽셀 (세부) 이 보이고, 축소해 보면 전체적인 모양 (전체) 이 보이는 것과 같습니다.

• 조건부 가이드 (Condition-Guided Decoder):
  인공지능이 길을 잃지 않도록 나침반 역할을 합니다.

  • "목표는 부산이고, 현재는 서울이야"라는 조건을 계속 주입합니다.
  • 만약 세부적인 행동 (미시) 을 결정할 때, 큰 목표 (거시) 와 어긋나면 "아니야, 부산으로 가려면 이쪽으로 가야 해"라고 바로잡아 줍니다.

4. 왜 MAGE 가 더 잘할까? (실제 성과)

논문의 실험 결과, MAGE 는 다음과 같은 환경에서 다른 방법들보다 압도적으로 잘했습니다.

• 복잡한 미로 찾기: 다른 로봇들은 벽을 뚫고 가거나, 같은 곳을 맴돌다가 지치지만, MAGE 는 "먼저 은색 동전을 줍고, 그 다음 금색 동전을 줍고, 마지막으로 출구로 가라"는 긴 미션을 완벽하게 수행했습니다.
• 정교한 로봇 손 조작: 펜을 돌리거나, 문을 열고, 못을 박는 등 섬세하고 긴 시간의 작업에서도 실수가 거의 없었습니다.

5. 한 줄 요약

MAGE 는 "먼저 큰 그림을 그리고, 그 안에서 세부 사항을 채워 넣는" 건축가 같은 인공지능입니다.

기존의 방법들이 "한 번에 모든 걸 맞히려고" 애쓰다가 길을 잃었던 반면, MAGE 는 큰 목표 (거시) 를 잊지 않으면서 세부 행동 (미시) 을 하나씩 정해나가므로, 길고 복잡한 미션에서도 흔들리지 않고 목표를 달성할 수 있습니다.

이 기술은 로봇이 복잡한 작업을 배우거나, 의료 분야에서 장기적인 치료 계획을 세우는 등, 오래 걸리고 보상이 드문 현실 세계의 문제를 해결하는 데 큰 도움이 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

오프라인 강화학습 (Offline RL) 은 환경과의 추가적인 상호작용 없이 기존에 수집된 데이터셋만으로 정책을 학습하는 방식으로, 로봇 공학 및 의료 등 실제 응용 분야에서 큰 잠재력을 가지고 있습니다. 그러나 기존 생성 기반 (Generation-based) 오프라인 RL 방법론들은 다음과 같은 주요 한계를 겪고 있습니다:

• 장기적 과업 (Long-horizon tasks) 과 희소 보상 (Sparse rewards) 의 어려움: 로봇 조작이나 전략적 계획과 같이 보상이 드물게 주어지고 시간이 오래 걸리는 작업에서 기존 방법들은 실패합니다.
• 다중 시간 규모 (Multi-scale temporal structure) 모델링의 부재:
  • Transformer 기반 (Decision Transformer 등): 단방향 자기회귀 (Autoregressive) 특성으로 인해 전역적인 맥락을 양방향으로 이해하는 데 한계가 있습니다.
  • Diffusion 기반 (Decision Diffuser 등): 국소적인 생성 편향 (Local generation bias) 으로 인해 짧은 구간에서는 타당해 보이지만, 긴 시간 범위에서 전체적인 일관성 (Global coherence) 을 잃기 쉽습니다.
• 계층적 방법의 비효율성: 기존 계층적 생성 방법 (HGM) 은 고정된 2 단계 구조 (상위 정책이 하위 정책을 제어) 를 사용하는데, 이는 다중 시간 규모의 추상화를 제한하고 두 정책을 동시에 최적화해야 하는 학습 효율성 문제를 야기합니다.

2. 제안 방법: MAGE (Methodology)

저자들은 이러한 한계를 극복하기 위해 MAGE (Multi-scale Autoregressive GEneration) 를 제안합니다. MAGE 는 거시적인 개요부터 미시적인 세부 사항까지 상향식 (Top-down, coarse-to-fine) 방식으로 궤적을 생성하는 오프라인 RL 프레임워크입니다.

핵심 구성 요소

1. 다중 시간 규모 오토인코더 (Multi-scale Trajectory Autoencoder, MTAE):

   • VQ-VAE(Vector Quantized Variational Autoencoder) 아키텍처를 기반으로 합니다.
   • 궤적 (Trajectory) 을 미리 정의된 시간 규모 스케줄에 따라 계층적으로 잠재 표현 (Latent representation) 으로 인코딩합니다.
   • 코arse 스케일 토큰: 장기적인 의존성과 전역 구조를 포착합니다.
   • Fine 스케일 토큰: 단기적인 세부 동작과 국소적 동역학을 포착합니다.
   • 모든 규모에서 공유되는 코드북 (Codebook) 을 사용하여 토큰을 이산화 (Discretization) 합니다.

2. 다중 시간 규모 조건부 자기회귀 생성기 (Multi-scale Condition-guided Autoregressive Generator):

   • Transformer 기반 생성: 이전 단계에서 생성된 더 거친 (coarser) 스케일의 토큰 맵을 조건으로 하여, 다음 더 정교한 (finer) 스케일의 토큰 맵을 순차적으로 예측합니다.
   • 조건부 가이드: 현재 상태 ($s_0$) 와 목표 회수 (Return-to-Go, $R_0$) 를 조건으로 사용하여 생성 방향을 유도합니다.
   • 조건 유도 정제 (Condition-Guided Refinement): 양자화 과정에서 발생하는 정보 손실로 인해 생성된 궤적이 초기 조건과 달라질 수 있는 문제를 해결하기 위해, 디코더에 경량 Adapter 모듈을 도입합니다. 이는 생성된 궤적의 초기 상태가 실제 조건 ($s_0, R_0$) 과 정확히 일치하도록 손실 함수 ($L_{cond}$) 를 통해 미세 조정합니다.

3. 잠재 역동역학 모델 (Latent Inverse Dynamics Model):

   • 생성된 다중 규모 잠재 표현 ($Z$) 을 입력받아 실제 실행할 행동 ($a$) 을 결정합니다.
   • 생성된 궤적 전체를 사용하는 것보다 잠재 공간에서 역동역학을 학습하는 것이 더 우수한 성능을 보임을 실험을 통해 확인했습니다.

3. 주요 기여 (Key Contributions)

• 새로운 생성 패러다임: 오프라인 RL 에 다중 시간 규모의 자기회귀 생성 (Multi-scale Autoregressive Generation) 을 도입하여, 전역적 일관성과 국소적 정밀도를 동시에 확보했습니다.
• 효율적인 계층 구조: 고정된 2 단계 계층 구조 대신, 단일 정책으로 다중 규모를 처리하는 유연한 구조를 설계하여 학습 효율성을 높였습니다.
• 정밀한 조건 제어: 조건 유도 Adapter 를 통해 생성된 궤적이 초기 상태와 목표 보상에 엄격하게 부합하도록 보장하여, 희소 보상 환경에서의 신뢰성을 높였습니다.
• 광범위한 벤치마크 검증: 5 개의 오프라인 RL 벤치마크 (Adroit, Franka Kitchen, AntMaze, Maze2D 등) 에서 15 가지 기존 알고리즘 대비 최상위 성능을 입증했습니다.

4. 실험 결과 (Results)

• 성능: 5 개 벤치마크, 15 개 베이스라인 알고리즘 (Decision Transformer, Diffuser, CQL, IQL, 계층적 Diffusion 등) 과 비교하여 SOTA(State-of-the-Art) 성능을 달성했습니다.
  • Adroit (정교한 조작): 희소 보상과 고차원 제어의 어려움에도 불구하고 Pen, Door, Hammer 작업에서 압도적인 성능 향상을 보였습니다.
  • Franka Kitchen (조립 작업): 여러 하위 목표를 올바른 순서로 수행해야 하는 복잡한 작업에서 기존 방법들을 크게 능가했습니다.
  • Maze Navigation (AntMaze, Maze2D): 긴 시간 범위의 탐색 작업에서 다른 방법들이 벽을 통과하거나 목표를 찾지 못하는 반면, MAGE 는 일관된 경로로 성공했습니다.
• 추론 속도: Diffusion 기반 방법 (Decision Diffuser 등) 에 비해 약 50~80 배 빠른 추론 속도를 보여주어 (약 27ms/스텝), 실시간 로봇 제어에 적용 가능한 효율성을 입증했습니다.
• Ablation Study:
  • 시간 규모 (K) 가 증가할수록 성능이 향상되지만, 과도한 세분화 (K≥10) 는 노이즈를 유발할 수 있음을 확인했습니다.
  • 조건부 가이드 ($L_{cond}$) 와 RTG(Return-to-Go) 기반 조건화가 궤적의 일관성과 목표 달성에 필수적임을 입증했습니다.

5. 의의 및 결론 (Significance)

MAGE 는 오프라인 강화학습 분야에서 다중 시간 규모 모델링과 조건부 생성을 성공적으로 통합한 획기적인 접근법입니다. 특히, 희소 보상이 주어지는 장기적 과업에서 기존 생성 모델들이 겪던 '전역적 일관성 부재' 문제를 해결했습니다.

이 연구는 복잡한 시퀀스 의사결정 문제를 해결하기 위해 거시적 계획 (Macro-planning) 과 미시적 실행 (Micro-execution) 을 하나의 통합된 자기회귀 프레임워크로 결합할 수 있음을 보여주었으며, 로봇 조작, 자율 주행, 의료 의사결정 등 다양한 실세계 응용 분야에서 신뢰할 수 있는 정책 학습을 위한 강력한 기반을 마련했습니다. 또한, 빠른 추론 속도로 인해 실제 시스템 배포 가능성까지 제시했다는 점에서 중요한 의미를 가집니다.