Flexible Multitask Learning with Factorized Diffusion Policy

"Factorized Diffusion Policy 를 활용한 유연한 멀티태스크 학습 (FDP)" 논문에 대한 상세한 기술적 요약은 다음과 같습니다.

1. 문제 제기

이 논문은 로봇 공학 분야에서 **멀티태스크 모방 학습 (multitask imitation learning)**이 직면한 과제, 특히 고도로 다중 모달 (multimodal) 이고 다양한 행동 분포를 모델링하는 어려움을 다룹니다.

• 단일 모델 (Monolithic Models) 의 한계: 전통적인 단일 네트워크 정책 (예: 표준 Diffusion Policy) 은 다양한 작업에 걸쳐 일반화하는 데 어려움을 겪으며, 복잡한 분포를 과소적합하거나 서로 다른 행동 모드 (예: 서로 다른 물체를 잡기 위한 다양한 그리핑 전략) 를 표현하지 못하는 경우가 많습니다.
• 기존 모듈형 모델의 한계: 전문가 혼합 (Mixture-of-Experts, MoE) 아키텍처는 작업을 분해하려 시도하지만 다음과 같은 문제점을 겪습니다:
  • 학습 불안정성: 전문가 활용 균형을 맞추기 어렵습니다.
  • 라우팅 문제: 이산적 전문가 선택으로 인해 "라우팅 불균형" (일부 전문가는 전혀 사용되지 않고 다른 전문가는 과도하게 사용됨) 이 발생할 수 있습니다.
  • 낮은 해석 가능성: 전문가들의 역할이 종종 중복되거나 불명확하며, 기술 (skills) 이 계층 간에 분리되지 않고 분산되어 있을 수 있습니다.
  • 파괴적 망각 (Catastrophic Forgetting): 새로운 작업에 적응하려면 종종 전체 모델을 재학습해야 하므로, 이전에 학습된 기술을 잃을 위험이 있습니다.

2. 방법론: Factorized Diffusion Policy (FDP)

저자들은 행동을 분산된 특수한 확산 모델들의 조합으로 분해하는 모듈형 정책 아키텍처인 FDP를 제안합니다.

A. 핵심 아키텍처

• 분해 (Factorization): 단일 정책 $p(a_t|o_t)$ 대신, FDP 는 행동 분포를 $N$개의 구성 요소 분포의 가중 곱으로 모델링합니다:
  $$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$$
  여기서 각 $p_i$는 서로 다른 행동 모드 (하위 기술) 를 포착하는 특수화된 확산 모델 (전문가) 이며, $w_{t,i}$는 관측치에 의존하는 가중치입니다.
• 구성적 샘플링 (Compositional Sampling):
  • 라우터: 경량 MLP 가 관측치 $o_t$를 입력받아 모든 구성 요소에 대한 연속 가중치 $\{w_{t,i}\}$를 출력합니다.
  • 스코어 집계: 이산적 MoE(하나의 전문가를 선택) 와 달리, FDP 는 모든 구성 요소로부터 **스코어 추정치 (로그 확률의 기울기)**를 집계합니다:
    $$\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$$
  • 디노이징 과정: 반복적인 디노이징 단계는 이 가중 합산된 스코어를 사용하여 행동 궤적 생성을 안내합니다. 이는 분포의 곱에서 샘플링하는 것에 해당하며, 행동 생성을 **제약 충족 문제 (constraint satisfaction problem)**로 프레임합니다.

B. 학습 및 적응

• 연합 학습 (Joint Training): 모든 확산 구성 요소와 라우터는 집계된 스코어 예측에 대한 평균 제곱 오차 (MSE) 손실 함수를 사용하여 엔드투엔드로 학습됩니다.
• 효율적 적응 (Upcycling): 전체 모델을 재학습하지 않고 새로운 작업에 적응하기 위해:
  1. 새로운 확산 구성 요소가 도입됩니다.
  2. 기존 구성 요소에서 가중치를 **업사이클링 (복사)**하여 초기화됩니다.
  3. 새로운 구성 요소와 라우터만 미세 조정되며, 모든 이전 구성 요소는 **동결 (frozen)**됩니다.
  4. 이는 파괴적 망각을 방지하고 학습 가능한 매개변수 수를 크게 줄입니다.

3. 주요 기여

1. 새로운 모듈형 아키텍처: 이산적 전문가 선택 대신 연속 스코어 집계를 사용하는 FDP 를 도입했습니다. 이는 안정적인 학습을 보장하고 라우팅 불균형을 방지하며 구성 요소 간 명확한 전문화를 촉진합니다.
2. 확률론적 해석: 에너지 기반 모델 (EBM) 과 스코어 매칭에 기반하여, 결합된 정책이 행동 제약의 교집합을 나타내는 원칙적인 형식을 제공합니다.
3. 확장 가능한 적응 전략: 업사이클링과 동결을 사용하는 "선택적 튜닝" 메커니즘을 제안하여, 최소한의 데이터로 효율적인 기술 확장을 가능하게 하고 파괴적 망각을 방지합니다.
4. 실증적 검증: 시뮬레이션 및 실제 환경 모두에서 강력한 베이스라인 (단일 DP, 희소 확산 정책, MoDE) 보다 우수한 성능을 입증했습니다.

4. 실험 결과

저자들은 MetaWorld, RLBench, LIBERO, 그리고 **실제 로봇 조작 (UR5e 암)**에서 FDP 를 평가했습니다.

• 멀티태스크 성능:

  • FDP 는 30 개 이상의 작업 전반에 걸쳐 단일 확산 정책 (DP) 및 모듈형 베이스라인 (SDP, MoDE) 보다 일관되게 우수한 성능을 발휘했습니다.
  • MetaWorld: FDP 는 평균 성공률 **74.8%**를 기록했으며, DP 는 70.8%, SDP 는 69.8% 였습니다.
  • RLBench: FDP 는 평균 성공률 **63.9%**를 기록했으며, DP 는 45.6%, SDP 는 42.1% 였습니다.
  • 실제 환경: 베이스라인이 과적합이나 다중 모달 분포 포착 불가로 실패한 복잡한 작업 (예: 머그컵 걸기, 색상 큐브 집기) 에서 FDP 는 더 높은 성공률을 보였습니다.

• 작업 적응 및 전이:

  • 데이터 효율성: 제한된 데이터 (10~25 회 시연) 로 새로운 작업에 적응할 때, "새 모듈 추가" 전략을 사용한 FDP 는 전체 매개변수 미세 조정 및 부분 미세 조정보다 우수한 성능을 보였습니다.
  • 연속 학습: 12 개 작업 연속 학습 벤치마크 (LIBERO) 에서 FDP 는 새로운 전문가가 추가됨에 따라 높은 성능을 유지한 반면, 다른 방법들은 성능이 저하되었습니다.
  • 지식 유지: 이전 구성 요소를 동결함으로써 기존 기술을 효과적으로 보존했습니다. 최소한의 재생 버퍼 (작업당 5 회 시연) 를 사용하면 적응 후에도 원래 성능의 거의 100% 를 유지할 수 있었습니다.

• 분석 및 확장:

  • 구성 요소 확장: 구성 요소 수가 2 개에서 4 개로 증가함에 따라 성능이 향상된 후 정점에 도달했습니다.
  • 작업 확장: 작업 수가 증가함에 따라 FDP 와 단일 모델 간의 성능 격차가 벌어져 FDP 의 확장성을 강조했습니다.
  • 해석 가능성: 시각화를 통해 개별 구성 요소가 노이즈 수준이나 임의의 계층이 아닌, 명확한 하위 기술 (예: 정렬, 접근, 그리핑) 에 전문화되었음을 확인했습니다.
  • 추론 속도: FDP 는 "부분 재구성 (상위 $k$개 구성 요소만 사용)"을 지원하여, 약 19% 의 성능 저하만으로도 2 배의 추론 속도 향상을 달성했습니다.

5. 의의

• 로봇 학습: FDP 는 로봇이 새로운 기술을 빠르게 학습하면서도 기존 기술을 잊지 않고 동적 환경에 배포할 수 있는 실용적인 솔루션을 제공합니다.
• 안정성: 이산적 라우팅을 연속 스코어 구성으로 대체함으로써, MoE 기반 확산 정책이 겪어 온 학습 불안정성 문제를 해결합니다.
• 해석 가능성: 분해 구조를 통해 정책 내에서 특정 하위 기술 (예: "그리핑" 대 "접근") 을 검사하고 이해할 수 있어, 디버깅 및 안전 검증을 지원합니다.
• 향후 방향: 이 프레임워크는 이질적인 모듈 설계 (아키텍처 혼합) 와 로봇이 지속적으로 전문화된 기술 라이브러리를 축적하는 평생 학습 시스템을 위한 새로운 길을 엽니다.

요약하자면, FDP는 확산 모델의 생성적 힘과 견고하고 해석 가능하며 적응 가능한 모듈형 아키텍처를 결합함으로써 로봇 모방 학습 분야에서 중요한 발전을 이루었습니다.