기술 요약: 라그랑주 섭동 확산 스티어링 (Lagrangian Perturbation Diffusion Steering, LP-DS)

1. 문제 정의

확산 모델(diffusion models) 및 플로우 매칭(flow-matching)에 기반한 고용량 생성 정책은 연속 제어 및 로보틱스 분야에서 강력한 패러다임으로 자리 잡았습니다. 이러한 모델은 오프라인 데이터로부터 학습된 표현력이 풍부하고 다중 모드(multimodal)인 행동 분포를 나타내는 데 탁월합니다. 그러나 이들은 두 가지 주요 한계점에 직면해 있습니다:

1. 행동 복제(Behavior Cloning, BC)의 제약: BC를 통해 학습된 정책은 불완전한 상태 공간 커버리지, 희귀한 회복 동작의 누락, 분포 변화(distribution shift) 하에서의 취약성 등 데모 데이터셋의 한계를 그대로 물려받습니다.
2. 강화 학습(RL)의 불안정성: RL은 온라인 개선을 위한 메커즘을 제공하지만, 대규모 생성 디코더를 직접 미세 조정하는 것은 빈번하게 불안정하고 샘플 효율성이 낮습니다. 최적화 목적 함수와 장기 호라이즌(long-horizon) 디노이징 또는 적분 역학 사이의 상호작용은 종종 학습 불안정성을 초래합니다.

**DSRL(Diffusion Steering via Reinforcement Learning)**과 같은 기존의 대안들은 학습된 잠재 노이즈 공간 정책을 학습함으로써 동결된 정책을 개선하려고 시사합니다. 그러나 DSRL은 사전 학습된 프라이어(prior)를 학습된 잠재 정책으로 사실상 대체합니다. 이 접근 방식은 두 가지 실패 모드를 겪습니다:

• 오프 매니폴드 드리프트(Off-Manifold Drift): 스티어링된 노이즈가 디코더를 학습시킬 때 사용된 표준 가우시안 서포트(support)에서 벗어나, 불안정한 디코딩을 유발할 수 있습니다.
• 모드 붕괴(Mode Collapse): 학습된 잠재 정책이 백본의 다중 모드 행동을 좁은 하위 집합으로 붕괴시켜 행동의 다양성을 감소시킬 수 있습니다.

2. 방법론: 라그랑주 섭동 확산 스티어링 (LP-DS)

저자들은 동결된 생성 백본을 수정하거나 프라이어를 완전히 교체하는 대신, 상태 조건부 잔차(state-conditioned residual)를 잠재 노이즈 공간에서 학습함으로써 동결된 생성 정책을 개선하는 경량 적응 프레임워크인 LP-DS를 제안합니다.

핵심 메커니즘

LP-DS는 동결된 생성 백본 $\Phi$를 잠재 노이즈 $w$와 상태 $s$를 행동으로 매핑하는 블랙박스 디코더로 취급합니다. 노이즈를 프라이어 $p_0$ (예: $\mathcal{N}(0, I)$)에서 직접 샘플링하는 대신, LP-DS는 학습 가능한 섭동 네트워크 $\Delta_\theta(s)$를 도입합니다:
$$w = \epsilon + \Delta_\theta(s), \quad \text{단, } \epsilon \sim \mathcal{N}(0, I)$$
$\Delta_\theta(s) = 0$일 때, 정책은 원래의 행동 복제 분포를 정확하게 복구합니다.

라그랑주 완화를 통한 제약 최적화

섭동이 노이즈를 디코더가 학습되었던 매니폴드 밖으로 밀어내어 발생하는 문제를 방지하기 위해, LP-DS는 학습 문제를 제약 최적화 문제로 정식화합니다. 목표는 섭동 크기에 대한 신뢰 영역(trust-region) 제약을 준수하면서 생성된 행동의 기대 Q-값을 최대화하는 것입니다.

저자들은 스티어링된 노이즈 분포와 프라이어 사이의 킬백-라이블러(Kullback–Leibler, KL) 발산을 이동된 크기의 제곱으로 근사합니다:
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
여기서 $\delta$는 허용 가능한 편차를 제어하는 하이퍼파라미터입니다.

이는 라그랑주 완화(Lagrangian relaxation) 접근 방식으로 해결됩니다. 제약을 동적으로 강제하기 위해 듀얼 변수(라그랑주 승수) $\alpha$가 도입됩니다:
$$\mathcal{L}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
알고리즘은 다음을 교대로 수행합니다:

1. 액터 업데이트: $\mathcal{L}$에 대해 $\theta$를 최대화합니다 (경사 상승법).
2. 듀얼 업데이트: $\alpha \geq 0$을 강제하기 위해 $\mathcal-L$에 대해 $\alpha$를 최대화합니다 (투영 경사 상승법). 만약 섭동이 신뢰 영역을 위반하면 $\alpha$가 증가하여 액터를 억제하고 프라이어에 더 가깝게 유지하도록 강제합니다.

알고리즘 구성 요소

학습 절차(Algorithm 1)는 다음을 포함합니다:

• 데이터 수집: 스티어링된 정책 $w = \epsilon + \Delta_\theta(s)$를 사용하여 행동을 샘플링합니다.
• 크리틱 업데이트: TD(Temporal Difference) 학습을 통한 행동 가치 크리틱 $Q^A$ 업데이트와, 기본 노이즈 분포 상의 $Q^A$로부터의 증류(distillation)를 통한 잠재 가치 크리틱 $Q^W$ 업데이트를 수행합니다.
• 프라이멀-듀얼 업데이트: 보상 최대화와 프라이어 보존 사이의 균형을 맞추기 위해 섭동 네트워크 $\theta$와 승수 $\alpha$를 업데이트합니다.

3. 주요 기여

본 논문은 네 가지 주요 기여를 설명합니다:

1. LP-DS 프레임워크: 전체 디코더 미세 조정의 불안정성과 비용을 피하면서, 상태 조건부 잔차 노이즈 제어를 통해 동결된 생성 정책을 개선하는 경량 적응 방법입니다.
2. 라그랑주 신뢰 영역 목적 함수: 섭동 크기를 동적으로 제약하여, 보상 최대화와 사전 학습된 잠재 프라이어 보존 사이의 균형을 맞추는 조절 가능한 메커니즘을 제공하는 새로운 목적 함수입니다.
3. 모드 붕괴 완화: LP-DS가 기존의 프라이어 교체형 노이즈 공간 스티어링 방식(DSRL 등)에 비해 더 높은 행동 공간 다양성을 유지하고 모드 붕개를 완화함을 실증적으로 보여줍니다.
4. 일반화 가능성: 압축된 확산 정책을 넘어 플로우 매칭 백본, 대규모 시각-언어-행동 모델($\pi_0$), 그리고 물리적인 Franka 로봇 실험까지 검증 범위를 확장했습니다.

4. 실험 결과

저자들은 RoboMimic (조작), OpenAI Gym (이동), Adroit (정교한 조작) 벤치마크에서 LP-DS를 DSRL, DPPO (정책 경사 미세 조정), 그리고 오프라인-투-온라인 확산 Q-러닝 베이스라인과 비교 평가했습니다.

• 성능: LP-DS는 성공률과 에피소드 리턴을 일관되적으로 향상시켰습니다. 특히, WALKER2D 환경에서 가장 강력한 베이스라인 대비 약 25%의 리턴 향상을 달お성했습니다.
• 샘플 효율성: 전체 디코더 대신 경량 섭동 모듈만을 업데이트함으로써, LP-DS는 우수한 샘플 효율성과 더 안정적인 학습 역학을 입증했습니다.
• 다양성 보존:
  • 토이 다중 목표 내비게이션 태스크에서, LP-DS는 신뢰 영역 경계 $\delta$를 통해 성능과 다양성 사이의 트레이드오프를 명시적으로 제어할 수 있었습니다. 작은 $\delta$는 다중 모드 행동을 보존한 반면, 큰 $\delta$는 특성화를 허용했습니다. 반면, DSRL은 즉시 단일 모드로 붕괴되었습니다.
  • AVOIDING 환경에서, LP-DS는 작은 $\delta$를 통해 광범위한 장애물 회피 경로(궤적 수준의 다중 모드)를 유지했으나, DSRL은 단일 지배적 경로로 붕괴되었습니다.
  • 엔트로피 분석: LP-DS는 온라인 적응 중 DSRL보다 일관되게 높은 행동 공간 엔트로피를 유지했으며, 이는 라그랑주 제약이 조기 붕괴를 효과적으로 억제하면서도 성능을 개선함을 나타냅니다.
• 어블레이션 연구: 라그랑주 페널티를 제거하면 학습이 불안정해지고 성공률이 저하되었습니다. 라그랑주 페널티와 노이즈 경계 모두를 제거하면 매우 불안정한 동작과 낮은 성공률이 나타나, 신뢰 영역 제약이 오프-매니폴드 드리프팅을 방지하는 핵심 메커니즘임을 확인했습니다.
• 민감도: 방법론은 신뢰 영역 타겟 $\delta$의 특정 값에 견고하며, 0.1 이상의 값에서 일반적으로 강력한 성능을 보였습니다.

5. 의의 및 주장

본 논문은 LP-DS가 생성 정책을 위한 안정적이고 효율적인 풀 미세 조정(full fine-tuning)의 대안임을 주장합니다. 동결된 백본을 고정된 프라이어로 취급하고 제한된 잔차 이동만을 학습함으로써, 이 방법은 다음과 같은 이점을 제공합니다:

• 오프 매니폴드 드리프트 방지: 잠재 쿼리를 디코더의 유효 작동 범위 내에 유지하여, 제약 없는 노이즈 공간 스티어링에서 보이는 불규칙한 동작을 방지합니다.
• 다중 모드 보존: 사전 학습된 프라이어를 교체할 때 흔히 발생하는 모드 붕괴를 완화하여, 동결된 백본이 가진 행동 다양성을 유지합니다.
• 아키텍처 간 일반화: 이 접근 방식은 확산 모델에 국한되지 않고, 플로우 매칭 모델 및 대규모 시각-언어-행동 모델을 포함하여 잠재-행동 매핑을 가진 모든 생성 정책에 적용 가능합니다.
• 물리적 배포 가능성: 시뮬레이션 기반 잠재 적응을 사용하여 물리적인 Franka 로봇 실험에서 성공적으로 검증되었으며, 이는 시뮬레이션 벤치마크를 넘어 실제 적용 가능성을 입증합니다.

저자들은 LP-DS를 작업 특화 최적화의 필요성과 사전 학습된 모델의 견고한 다중 모드 구조 유지의 필요성 사이에서 균형을 맞추는, 고용량 생성 정책의 "온라인 개선"을 위한 실용적인 솔루션으로 제시합니다.