기술 요약: 대규모 결정 모델(Large Decision Models)을 통한 확장 가능한 멀티태스크 강화학습을 향하여

문제 정의

강화학습(RL)은 게임, 로보틱스, 제어와 같은 특정 영역에서 상당한 성공을 거두었습니다. 그러나 실제 환경에 적용하는 것은 광범한 환경 상호작용, 세심한 보상 설계, 도메인 특화 아키텍처, 그리고 상당한 하이퍼파라미터 튜닝이 필요하기 때문에 여전히 도전적인 과제로 남아 있습니다. 오프라인 RL 및 오프라인-투-온라인 접근 방식이 온라인 상호작용의 비용을 완화해주기는 하지만, 새로운 환경에 적합한 알고리즘을 선택하고 적응시키는 과정은 여전히 어렵습니다.

이러한 과제와 병행하여, 대규모 시퀀스 모델은 자연어 처리와 컴퓨터 비전 분야를 변화시켰습니다. RL 분야에서는 Decision Transformer와 같은 시퀀스 모델링 접근 방식이 정책을 궤적(trajectories)에 대한 자기회귀 모델로 표현할 수 있음을 입증했습니다. 이는 다음과 같은 중요한 질문을 제기합니다: 다양한 도메인의 오프라인 RL 궤적들을 하나의 확장 가능한 트랜스포머 정책으로 통합하면서도, 많은 도메인에 걸쳐 강력한 태스크 성능을 유지할 수 있는가? 주요 장애물은 멀티 도메인 RL 생태계가 파편화되어 있고, 이질적이며, 통합하기 어렵다는 점입니다.

방법론

저자들은 다양한 도メイン의 궤적들로부터 오프라인 학습된 통합 멀티태스크, 멀티모달 트랜스포머 정책인 LDM-v0(Large Decision Model v0)를 소개합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어져 있습니다.

1. 자동 참조 정책 감독 (Automated Reference-Policy Supervision)

수동적인 보상 설계 없이 다양한 도메인에 걸쳐 고품질의 학습 데이터를 생성하기 위해, 저자들은 데이터셋 구축을 AutoRL 문제로 취급합니다:

• 후보 선택: Stable-Baselines3 및 SB3-Contrib에서 제공하는 A2C, PPO, SAC, DQN 등을 포함한 후보 알고리즘과 설정(configuration) 풀을 정의합니다.
• 순위 매기기: 각 환경 라이브러리에 대해, 호환 가능한 후보들을 5개의 환경 서브셋에서 학습시킵니다. 설정들은 최종 성능에 대한 단측 웰치 t-검정(one-tailed Welch's t-test)을 기반으로 한 쌍 비교를 통해 순위가 매겨집니다.
• 데이터 생성: 상위 순위의 설정들을 3백만 번의 전이(transitions) 동안 학습시킵니다. 궤적이 기록되며, 최종 학습된 정책을 재생(replay)하여 참조 액션 주석(reference action annotations)을 생성합니다.
• 필터링: 통계적으로 유의미한 성능 향상이 없거나 최적의 정책 성능의 95% 미만인 저품질 실행 건들은 제거됩니다.
• 결과: 약 3,000개의 서로 다른 환경으로부터 추출된 약 93억 개의 전이 데이터셋이 생성되었으며, 이는 탐색적 액션이 아닌 참조 정책 액션으로 주석이 달렸습니다.

2. 통합 환경 인프라

RL 라이브러리의 파편화(서로 다른 Python 버전, 의존성, 시뮬레이터 등)를 처리하기 위해 저자들은 내부 오케스트레이션 프레임워크를 개발했습니다:

• 환경들은 격리된 Docker 컨테이너 내에 캡슐화됩니다.
• 현대적인 Gymnasium 인터페이스와 호환되는 통합 상호작용 API가 노출됩니다.
• 이 파이프라인은 약 15,000개의 개별 환경을 아우르는 146개의 환경 라이브러리를 통합했습니다.

3. 모델 아키텍처 및 학습

• 아키텍처: LDM-v0는 처음부터 학습된 Llama 아키텍처 기반의 디코더 전용(decoder-only) 트랜스포머입니다.
• 입력 처리: 모델은 상호작용 히스토리(과거 관측값, 액션, 보상, 종료 신호)와 현재 관측값을 받습니다.
  • 토큰화: 연속적인 입력은 $\mu$-law 인코딩을 통해 1024개의 빈(bin)으로 이산화됩니다. 이산적 입력은 룩업 테이블을 사용합니다. 이미지는 64x64로 크기가 조정되고 컨볼루션 인코더를 통해 인코딩됩니다.
  • 임베딩: 관측값, 액션, 보상, 종료 신호는 정렬되어 **전이 수준 메타 토큰(transition-level meta-tokens)**으로 스택됩니다. 이러한 압축된 표현 방식은 (Gato와 같은) 차원별 토큰화에 비해 시퀀스 길이를 줄여주어, 고정된 컨텍스트 예산 내에서 더 긴 히스토리를 가질 수 있게 합니다.
• 학습 목표: 모델은 이산화된 액션 빈에 대한 교차 엔트로피 손실(cross-entropy loss)을 사용하여 지도 학습 기반의 다음 액션 예측(next-action prediction)을 통해 학습됩니다. 모델은 궤적 히스토리 $(o_i, a_i, r_i, d_i)$를 조건으로 하여 참조 액션 $a_t$를 예측하는 법을 배웁니다.
• 규모: 기본 모델인 LDM-v0는 3억 8백만 개의 파라미터, 12개의 레이어, 2048개의 전이 컨텍스트 길이를 가집니다. 이는 16개의 NVIDIA H200 GPU가 장착된 2개의 노드에서 6일 동안 학습되었습니다.

주요 기여

1. LDM-v0: 146개 라이브러리와 약 3,000개의 환경을 아우르는 통합 데이터셋으로 학습된 대규모 멀티모달 트랜스포머 정책입니다.
2. 확장 가능한 인프라: 매우 이질적인 RL 생태계 전반에서 참조 정책 생성 및 환경 오케스트레이션을 자동화하는 강력한 파이프라인을 구축했습니다.
3. 압축된 전이 수준 모델링: 관측, 액션, 보상, 종료 신호를 단일 메타 토큰으로 묶는 토큰화 전략을 통해 효율적인 장기 시퀀스 모델링을 가능하게 했습니다.
4. 경험적 검증: 단일 공유 파라미터화가 태스크별 미세 조정 없이도 다양한 도메인에서 경쟁력 있는 성능을 달랄 수 있다는 증거를 제시했습니다.

결과

저자들은 학습 환경(In-distribution)에 대해 LDM-v0를 평가하고, 이를 독립적으로 학습된 태스크 특화 참조 정책들과 비교했습니다.

• 성능 임계값: LDM-v0는 1,600개 이상의 환경에서 참조 정책 성능의 80% 이상을 달성했으며, 약 1,000개의 환경에서는 참조 성능과 일치하는 성능을 보였습니다.
• 도메인 커버리지: 다음과 같은 분야에서 강력한 성능이 관찰되었습니다:
  • 로보틱스 및 제어 (예: gymnasium_robotics, panda-gym).
  • 자율 주행 (highway-env).
  • 금융 트레이딩 (gym_trading_env) 및 재고 최적화.
  • 사이버 보안 (cymnasium).
  • 게임형 환경 (Atari 스타일, procogen, Super Mario Bros).
• 스케일링 법칙 (Scaling Laws): 32M에서 736M 파라미터 범위의 모델 크기에 따른 실험 결과, 32M에서 308M까지는 상당한 성능 향상이 나타났습니다. 308M과 736M 사이에서는 이득이 정체되는 양상을 보였는데, 이는 더 큰 용량이 도움이 되긴 하지만 308M 모델이 이미 중요한 통계적 규칙성을 포착하고 있음을 시사합니다.

의의 및 주장

본 논문은 단일 트랜스포머 정책을 사용하여 이질적인 강화학습 환경 전반에 걸쳐 대규모 오프라인 사전 학습을 수행하는 것이 가능하다는 것을 입증한다고 주장합니다.

• 통합된 표현: 결과는 관측 모달리티, 액션 공간, 보상 스케일, 시간적 지평(temporal horizons)의 상당한 이질성에도 불구하고, 대규모 시퀀스 모델이 환경의 하위 집합 전반에 걸쳐 공유되는 통계적 규칙성을 활용할 수 있음을 시사합니다.
• 자동화: 본 연구는 현재의 RL 생태계 파편화를 극복하기 위해 전문가의 지식과 태스크별 설계 선택에 대한 의존도를 낮추는 자동화된 통합 파이프라인의 잠재력을 강조합니다.
• 향후 방향: 저자들은 LDM-v0를 확장 가능한 사전 학습된 강화학습 시스템을 향한 단계로 간주합니다. 이들은 본 연구의 주된 목적이 분포 외(out-of-distribution) 일반화를 연구하는 것이 아니라, 단일 정책이 다양한 RL 행동을 규모 있게 공동으로 모델링할 수 있는지 조사하는 것이었음을 명시했습니다.

저자들은 일반화 능력에 대해 겸허한 태도를 유지하며, 평가가 주로 학습 분포 내의 환경에서 수행되었음을 인정했습니다. 또한, 전이 가능한 의사결정 전략을 이해하고, (인컨텍스트 적응 또는 파인튜닝을 통해) 보지 못한 환경에 대한 일반화 능력을 개선하는 것을 중요한 향후 과제로 꼽았습니다.