오프라인 강화학습을 위한 Decision Transformers 훈련

Decision Transformers가 RL을 시퀀스 모델링으로 전환

Decision Transformers는 강화학습(RL)을 전통적인 정책 최적화 작업이 아니라 조건부 시퀀스 모델링 문제로 간주합니다. 누적 보상을 최대화하기 위해 가치 함수를 맞추는 대신, Decision Transformer는 시퀀스 모델링 알고리즘(Transformer)을 사용하여 원하는 반환값, 과거 상태, 과거 행동을 기반으로 미래 행동을 생성합니다.

이 접근법은 생성적 궤적 모델링(상태, 행동, 보상의 결합 분포를 모델링)을 활용하여 기존 RL 알고리즘을 대체함으로써 RL에서 패러다임 전환을 의미합니다. 모델은 전통적인 의미에서 반환값을 최대화하지 않으며, 대신 특정 목표 반환값을 달성할 것으로 예측되는 행동을 생성합니다.

모델 아키텍처 및 프로세스

  1. Input Sequence: 모델은 마지막 $K$ 타임스텝을 받으며, 각 타임스텝은 Return-to-go, State, Action의 세 구성 요소로 이루어집니다.
  2. Embedding: 토큰은 벡터 상태에 대해 선형 레이어를, 이미지 프레임에 대해 CNN 인코더를 사용하여 임베딩됩니다.
  3. Prediction: GPT-2 기반 모델이 이러한 입력을 처리하고 자동회귀 모델링을 통해 미래 행동을 예측합니다.

오프라인 Decision Transformers 구현

오프라인 Decision Transformers는 훈련 중 환경과 직접 상호작용하지 않고 다른 에이전트 또는 인간 시연으로부터 수집된 데이터셋을 사용해 학습됩니다.

데이터셋 준비 및 커스텀 Data Collator

Decision Transformer를 훈련하기 위해서는 원시 RL 데이터를 시퀀스 모델링에 적합한 형식으로 전처리해야 합니다. Hugging Face Hub의 halfcheetah-expert-v2 데이터셋을 사용하여 다음 전처리 단계가 필요합니다:

  • Normalization: 각 특성은 평균을 빼고 표준편차로 나누어 정규화됩니다.
  • Return Computation: 각 궤적에 대해 할인된 반환값이 사전 계산됩니다.
  • Scaling: 보상과 반환값은 1000배로 스케일링됩니다.
  • Sampling Distribution: 샘플링 분포는 전문가 에이전트 궤적의 길이를 고려하도록 보강됩니다.

이 단계들은 커스텀 Data Collator를 통해 구현되며, 이는 상태, 행동, 보상, returns-to-go, 타임스텝, 어텐션 마스크를 포함하는 배치를 반환합니다.

Hugging Face Trainer를 사용한 훈련

훈련은 Hugging Face Trainer 클래스를 사용하여 수행됩니다. 표준 DecisionTransformerModel은 기본적으로 손실 값을 반환하지 않기 때문에, 래퍼 클래스(TrainableDT)를 사용하여 모델의 행동 예측과 데이터셋의 목표 행동 사이의 L-2 노름(평균 제곱 오차)을 계산합니다.

핵심 훈련 하이퍼파라미터:

  • 에폭: 120
  • 배치 크기: 64
  • 학습률: $1e-4$
  • 가중치 감소: $1e-4$
  • 워밍업 비율: 0.1
  • 옵티마이저: AdamW
  • 최대 그래디언트 노름: 0.25

Transformer 기반 RL의 향후 방향

Hugging Face는 여러 계획된 이니셔티브를 통해 딥 강화학습 커뮤니티에 대한 지원을 확대할 예정입니다:

  • Online Training: 온라인 환경에서 학습 또는 파인튜닝된 Decision Transformer 모델을 포함하도록 저장소를 확장합니다.
  • Tool Integration: sample-factory 버전 2.0을 에코시스템에 통합합니다.

Sources