기술 요약: BiTrajDiff

문제 정의

오프라인 강화 학습 (RL) 은 추가적인 환경 상호작용 없이 정적이며 사전 수집된 데이터셋으로부터 효과적인 정책을 학습하는 것을 목표로 합니다. 이 분야의 주요 과제는 Out-of-Distribution(OOD) 문제이며, 이는 학습된 정책이 기록된 데이터에 존재하지 않는 상태 - 행동 쌍을 마주할 때 발생하여 외삽 오류와 가치 과대평가를 초래합니다. 이를 완화하기 위해 기존 방법들은 종종 보수적인 제약을 부과하는데, 이는 학습을 데이터셋의 구조적 편향으로 우연히 제한하여 행동의 다양성과 일반화 능력을 제한합니다.

최근 데이터 증강 (DA) 접근법들은 생성 모델 (특히 확산 모델) 을 사용하여 새로운 전이를 합성함으로써 이를 해결하려 시도합니다. 그러나 현재의 DA 방법들은 주로 데이터셋 내 상태에 조건부로 작용하는 단일 롤아웃 패러다임(전방향 - 미래 또는 후방향 - 과거) 내에서 작동합니다. 그 결과, 이러한 방법들은 데이터셋의 원래 궤적 수준의 연결성을 유지하는 경향이 있습니다. 기존 행동 패턴 내에서 국소적 변이를 생성하지만, 서로 다른 분리된 행동 패턴 간의 연결을 복원하지는 못합니다 (예: 원래 데이터에 문간을 가로지르는 궤적이 존재하지 않는 내비게이션 작업에서 두 개의 분리된 방을 연결하는 경우).

방법론

저자들은 이전에 분리되었던 행동 패턴들을 연결하는 궤적을 합성함으로써 전역적 궤적 수준의 연결성을 명시적으로 복원하도록 설계된 새로운 DA 프레임워크인 BiTrajDiff(양방향 궤적 확산) 를 제안합니다.

핵심 프레임워크

BiTrajDiff 는 궤적 합성을 두 개의 독립적이면서도 상호 보완적인 확산 과정으로 분해합니다:

1. 전방향 - 미래 확산 (Forward-Future Diffusion): 중간 상태로부터 시작하는 미래 상태 시퀀스의 생성을 모델링합니다.
2. 후방향 - 과거 확산 (Backward-History Diffusion): 동일한 중간 상태에서 끝나는 과거 상태 시퀀스의 재구성을 모델링합니다.

두 과정 모두 오프라인 데이터셋에서 샘플링된 동일한 중간 앵커 상태($s_t$) 와 목표 누적 보상 신호에 조건부로 수행됩니다.

주요 구성 요소

1. 양방향 확산 훈련:

   • 전방향용 ($\epsilon_{\theta_f}$) 과 후방향용 ($\epsilon_{\theta_b}$) 두 개의 별도 확산 모델이 오프라인 데이터셋에서 훈련됩니다.
   • 전방향 모델은 초기 상태 $s_t$ 와 목표 반환 $R_f$ 를 조건으로 하여 $\{s_t, \dots, s_{t+H-1}\}$ 궤적을 생성합니다.
   • 후방향 모델은 종료 상태 $s_t$ 와 목표 반환 $R_b$ 를 조건으로 하여 $\{s_{t-H+1}, \dots, s_t\}$ 궤적을 생성합니다.
   • 훈련은 조건부 및 무조건부 생성을 처리하기 위해 Classifier-Free Guidance(CFG) 를 활용하며, 누적 보상이 주요 조건으로 작용합니다.

2. 양방향 궤적 생성 및 연결 (Stitching):

   • 추론 중에는 데이터셋에서 앵커 상태 $s_t$ 가 샘플링됩니다.
   • 전방향 모델은 미래 구간을 생성하고, 후방향 모델은 과거 구간을 생성합니다.
   • 이러한 구간들은 공유된 앵커 상태 $s_t$ 에서 연결 (stitched) 되어 완전한 양방향 궤적을 형성합니다: $\hat{\tau}_s(s_t) = \{ \hat{s}_{t-H+1}, \dots, \hat{s}_{t-1}, s_t, \hat{s}_{t+1}, \dots, \hat{s}_{t+H-1} \}$.
   • 이 연결 메커니즘을 통해 원래 데이터셋에서 도달 불가능하거나 분리되어 있던 상태들을 연결하는 궤적을 합성할 수 있습니다.

3. 궤적 완성:

   • 확산 모델이 상태 시퀀스를 생성하므로, 상태 전이 ($s_t, s_{t+1}$) 가 주어졌을 때 행동 ($a_t$) 을 예측하기 위해 역동역학 모델 (Inverse Dynamics Model, IDM) 이 사용됩니다.
   • 보상 모델 (Reward Model, RM) 은 상태 - 행동 쌍에 대한 보상 ($r_t$) 을 예측합니다.
   • 이를 통해 행동 및 보상 신호로 궤적이 완성되어 완전한 전이 튜플이 형성됩니다.

4. 양방향 궤적 필터링:
   증강된 데이터의 품질을 보장하기 위해 2 단계 필터링 메커니즘이 적용됩니다:

   • OOD 필터: Isolation Forest 를 사용하여 원래 데이터 분포에서 크게 벗어난 (새로움 탐지) 궤적을 탐지하고 제거합니다.
   • Greedy 필터: 가장 높은 누적 보상을 기반으로 상위 궤적을 선택하여, 증강된 데이터가 정책이 고성능 행동으로 이끌도록 보장합니다.

주요 기여

1. 한계점 식별: 저자들은 단일 롤아웃 DA 방법들이 원래 궤적 수준의 연결성을 유지하여 서로 다른 행동 패턴 간의 연결을 탐색하지 못한다는 점을 식별했습니다.
2. 양방향 합성: 공유 앵커에 조건부로 전방향 및 후방향 과정으로 생성을 명시적으로 분해하는 BiTrajDiff 를 제안하여, 분리된 행동 모드들을 연결하는 궤적의 합성을 가능하게 합니다.
3. 실증적 검증: D4RL 벤치마크에 대한 광범위한 실험을 통해 BiTrajDiff 가 다양한 오프라인 RL 백본 (IQL, TD3BC, CQL, DT) 에서 Synther, DiffStitch, RTDiff 를 포함한 고급 DA 베이스라인보다 일관되게 우수한 성능을 발휘함을 입증했습니다.

실험 결과

이 논문은 이동 (locomotion), 내비게이션 (Maze2D, AntMaze), 조작 (Kitchen) 작업을 포함하는 D4RL 벤치마크에서 BiTrajDiff 를 평가합니다.

• 성능 향상: BiTrajDiff 는 대부분의 작업과 알고리즘에서 가장 높은 평균 보상을 달성합니다. 예를 들어, 이동 작업에서는 단일 방향 베이스라인보다 일관되게 우수한 성능을 보입니다. Maze 및 Kitchen 과 같이 데이터 품질이 중요한 희소 보상 환경에서 BiTrajDiff 는 안정성과 성능에서 상당한 개선을 보여줍니다.
• Ablation 연구:
  • 방향성: BiTrajDiff(양방향) 는 전방향 전용 및 후방향 전용 변형 모두보다 일관되게 우수하여, 전역 연결성을 복원하기 위해 과거 및 미래 구간을 모두 연결하는 것의 필요성을 검증합니다.
  • 필터링: OOD 및 Greedy 필터의 조합은 어느 하나만 사용하거나 필터를 사용하지 않는 경우보다 우월한 것으로 나타나 다양성과 신뢰성 사이의 균형을 맞춥니다.
  • 증강 비율: 원본 데이터 대비 30% 의 증강 데이터 비율 ($\sigma=30\%$) 이 최적의 안정성과 성능을 제공하며, 과도한 증강 (100%) 은 성능을 저하시키는 노이즈를 도입합니다.
• n-step TD 추정기: BiTrajDiff 는 n-step 시야가 증가함에 따라 베이스라인에 대한 성능 우위를 유지하며, 이는 생성된 궤적이 다른 DA 방법들의 궤적보다 더 높은 품질을 가지며 기본 역학에 더 일관적임을 시사합니다.
• 다양성 vs 정확도: 시각화 지표 (Dynamic Error 및 L2 Distance) 는 BiTrajDiff 가 원래 데이터와의 L2 거리가 더 큰 (더 큰 다양성) 궤적을 생성하면서도 경쟁력 있는 역동적 정확도를 유지하는 반면, 단일 방향 방법들은 원래 분포에 과적합되는 경향이 있음을 보여줍니다.

중요성 및 주장

이 논문은 BiTrajDiff 가 오프라인 RL 데이터 증강의 근본적인 한계, 즉 단일 과정 롤아웃이 분리된 행동 패턴을 연결하지 못하는 문제를 해결한다고 주장합니다. 공유 앵커 상태 주변에서 전방향 - 미래 및 후방향 - 과거 궤적을 명시적으로 모델링하고 연결함으로써, BiTrajDiff 는 원래 데이터셋에 부재했던 전역 궤적 수준의 연결성을 복원합니다.

저자들은 이 접근법이 오프라인 RL 알고리즘이 물리적 일관성을 희생하지 않고 더 다양한 행동 패턴 세트에서 학습할 수 있게 하여, 다양한 작업과 알고리즘에서 상당한 성능 개선을 이끌어낸다고 주장합니다. 이 연구는 생성적 데이터 증강을 통해 보수적 오프라인 학습의 구조적 편향을 극복하는 한 걸음으로 제시됩니다. 향후 연구는 더 넓은 행동 패턴을 가진 멀티태스크 환경에서의 양방향 생성을 탐구할 것으로 언급됩니다.