DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

기술 요약: DAWM – 행동 추론 전이를 통한 오프라인 강화학습을 위한 확산 행동 세계 모델

1. 문제 제기

확산 기반 세계 모델은 오프라인 강화학습 (RL) 을 위한 현실적인 장기 시나리오 궤적을 합성하는 데 뛰어난 능력을 입증해 왔습니다. 그러나 기존 방법론들은 단일 시간차 (TD) 학습에 의존하는 표준 가치 기반 오프라인 RL 알고리즘 (예: TD3BC, IQL) 과의 중요한 호환성 문제를 안고 있습니다.

• 행동 없는 모델의 한계: 확산 세계 모델 (DWM) 과 같은 방법론들은 미래 상태와 보상을 생성하지만 행동을 명시적으로 생성하지는 않습니다. 결과적으로 이들은 완전한 전이 $(s, a, r, s')$를 구성할 수 없으며, 이는 효과적인 비평가 정규화를 위해 상태, 행동, 보상 간의 일관성이 필요한 단일 시간차 학습과 양립할 수 없습니다.
• 연합 모델링의 한계: 상태, 보상, 행동을 동시에 모델링하려는 접근 방식은 종종 증가된 학습 복잡성과 불안정성을 겪어 실제 성능이 저하됩니다.
• 계획 기반 방법의 한계: 역동역학 모델 (IDM) 이나 연합 모델링을 통해 행동을 복원하는 확산 기반 계획기 (예: Diffuser, Decision Diffuser) 는 일반적으로 반복적 샘플링과 사후 보상 선택에 의존합니다. 이는 높은 추론 비용을 초래하고 단일 시간차 학습의 효율성을 우회하게 만듭니다.

핵심적인 과제는 연합 모델링의 학습 불안정성이나 반복적 계획의 추론 비용을 초래하지 않으면서도 안정적인 단일 시간차 학습을 지원하는 합성 데이터를 생성하는 것입니다.

2. 방법론: DAWM

저자들은 오프라인 RL 에 적합한 완전한 레이블이 지정된 전이를 생성하기 위해 궤적 생성과 행동 추론을 분리하는 모듈식 프레임워크인 **DAWM(확산 행동 세계 모델)**을 제안합니다. 이 프레임워크는 두 가지 주요 단계로 구성됩니다:

1 단계: 조건부 확산 궤적 생성

• 확산 모델 $p_\theta$는 현재 상태 $s_t$, 행동 $a_t$, 그리고 목표 **회수 (return-to-go, $R_{tg}$)**에 조건부로 미래 상태 및 보상 시퀀스를 생성하도록 학습됩니다.
• 이 모델은 조건부 분포 $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$를 정의합니다.
• 학습 목적은 분류기 없는 가이드 (classifier-free guidance) 를 사용하는 스코어링 탈노이즈 목적 함수를 통해 예측된 노이즈와 실제 노이즈 간의 기대 오차를 최소화하는 것입니다.
• 결정적으로, 이 모델은 확산 과정 내에서 행동을 연합 모델링하는 복잡성을 피함으로써 생성 작업을 단순화하는 행동 없는 궤적 (상태와 보상만) 을 생성합니다.

2 단계: 역동역학 모델 (IDM) 을 통한 행동 완성

• 행동 없는 궤적을 완전한 전이로 변환하기 위해 별도로 학습된 역동역학 모델 (IDM), $f_\phi$가 사용됩니다.
• IDM 은 대각 가우시안 분포로 모델링된 분포 $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$를 통해 누락된 행동을 재구성합니다.
• IDM 은 상태 전이가 주어졌을 때 관측된 행동의 가능도를 최대화하도록 실제 오프라인 데이터로 학습됩니다.
• 추론 시, IDM 은 확산으로 생성된 궤적에 적용되어 행동을 추론하고, 완전한 튜플 $(s, a, r, s')$을 산출합니다.

오프라인 정책 학습

• 합성된 완전한 전이는 표준 가치 기반 알고리즘, 특히 TD3 에 행동 복제 정규화기를 추가한 TD3BC와 이득 가중 감독 회귀를 수행하는 IQL을 사용하여 오프라인 RL 에이전트를 학습시키는 데 활용됩니다.
• Q 함수는 재구성된 전이에서 계산된 벨만 타겟을 사용하여 업데이트되어 효율적인 단일 시간차 TD 업데이트를 가능하게 합니다.

3. 주요 기여

1. 모듈식 프레임워크: 저자들은 확산 기반 상태 - 보상 생성과 별도로 학습된 IDM 을 결합한 효율적인 프레임워크를 제안합니다. 이 설계는 완전한 전이를 생성하면서도 연합 모델링의 학습 불안정성을 피합니다.
2. 단일 시간차 학습과의 호환성: 완전한 $(s, a, r, s')$ 전이를 생성함으로써 DAWM 은 단일 시간차 학습에 의존하는 보수적인 오프라인 RL 알고리즘 (TD3BC, IQL) 의 사용을 가능하게 하여, 반복적 확산 계획기에 대한 계산 효율적인 대안을 제공합니다.
3. 실증적 성능: 이 논문은 보수적인 오프라인 RL 알고리즘이 이러한 증강된 궤적으로 학습할 때 크게 혜택을 받아 여러 D4RL 이동 작업에서 기존 확산 기반 베이스라인 (DWM, DD) 을 일관되게 능가함을 보여줍니다.
4. 합성 데이터의 타당성: 결과는 DAWM 으로 생성된 합성 데이터로 학습된 에이전트가 실제 오프라인 데이터셋으로 학습된 에이전트와 동등하거나 더 나은 성능을 달성할 수 있음을 나타내며, 합성 궤적이 실용적인 대안이 될 수 있는 잠재력을 강조합니다.

4. 실험 결과

실험은 다양한 데이터 품질 (Medium, Medium-Replay, Medium-Expert) 을 가진 9 개의 표준 D4RL 이동 벤치마크 (Hopper, Walker2d, HalfCheetah) 에서 수행되었습니다.

• DWM 과의 비교: DAWM 은 행동 없는 DWM 베이스라인을 일관되게 능가했습니다. TD3BC 를 사용할 때 DAWM 은 평균 9.3% 개선되었고, IQL 을 사용할 때는 9.5% 개선되었습니다. 특히 Walker2d 와 같은 복잡한 제어 작업에서 개선 폭이 두드러졌습니다.
• 계획 기반 방법 (DD) 과의 비교: DAWM 은 정규화된 회수 측면에서 Decision Diffuser (DD) 를 능가했습니다. 또한, DAWM 은 추론 중 반복적 샘플링을 피함으로써 DD 대비 약 4.0 배 빠른 추론 효율성을 보여주었습니다.
• 실제 데이터와의 비교: DAWM 으로 생성된 데이터로 학습된 에이전트는 표준 TD3BC/IQL 을 사용하여 실제 데이터로 학습된 에이전트와 비교 가능한 정규화된 회수를 달성했으며, 일부 경우 (예: TD3BC 를 사용한 Hopper-MR) 에는 이를 능가했습니다. SSORL(실제 데이터에 IDM 을 적용한 베이스라인) 은 잘 수행되었으나, 저자들은 SSORL 이 가장 긴 궤적을 필터링하는 혜택을 받았음을 지적하며, 이는 DAWM 의 고정 길이 생성에는 존재하지 않는 요소라고 언급했습니다.
• 절대성 연구 (Ablation Studies):
  • 데이터 양 (DAWM-T): 조건부 쌍의 수를 8 배 줄인 것 (DAWM-T) 은 성능이 약간만 저하되어, 행동 완성 메커니즘이 순수한 데이터 양보다 더 중요한 요소임을 시사합니다.
  • 생성 지평: DAWM 은 다양한 생성 지평 ($H=1, 3, 7$) 에서 견고한 성능을 유지하여, 롤아웃 길이에 관계없이 단일 시간차 학습 프레임워크가 안정적임을 나타냈습니다.

5. 중요성과 주장

이 논문은 DAWM 이 효율적인 생성 모델링과 TD 기반 정책 학습 사이의 간극을 메운다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 단일 시간차 학습 가능: 이전에는 행동 없는 세계 모델과 양립할 수 없었던 표준적이고 안정적인 오프라인 RL 알고리즘을 위해 확산 생성 데이터를 사용할 수 있는 메커니즘을 제공합니다.
• 계산 효율성: 반복적 계획 방법에 비해 추론 비용을 줄임으로써 대규모 생성 오프라인 RL 을 위한 실용적인 솔루션을 제공합니다.
• 실용적 대안: 결과는 이 모듈식 접근 방식을 통해 생성된 합성 궤적이 실제 오프라인 데이터셋에 대한 타당하고 고품질의 대안이 될 수 있음을 시사하며, 탐사가 비용이 많이 들거나 위험한 실제 응용 분야에서 데이터 부족 문제를 완화할 잠재력을 보여줍니다.

저자들은 이 방법이 유망함을 보여주지만, 향후 연구는 시각적 제어, 고차원 로봇 공학, 그리고 불완전한 역동성 하에서의 견고성으로의 확장 가능성을 조사할 것이라고 결론지었습니다.