Behavior Uncloning: Distilling Mode Redirection into Policy Weights without Inference-Time Steering

기술 요약: Behavior Uncloning: 추론 시 스티어링 없이 정책 가중치로 모드 재지정(Mode Redirection) 증류하기

문제 정의

행동 복제(Behavior Cloning, BC)는 대규모의 다양하고 다양한 데모 데이터셋으로부터 유능한 시각-운동 정책(visuomotor policies)을 학습할 수 있게 해줍니다. 그러나 이러한 데이터셋은 서로 다른 작업자나 맥락에서 얻은 데모를 집계하는 경우가 많으며, 이로 인해 정책이 동일한 작업에 대해 여러 개의 "행동 모드(behavior modes)"를 학습하게 됩니다. 모든 모드가 작업을 성공적으로 수행할 수는 있지만, 일부는 안전하지 않거나, 바람직하지 않거나, 특정 배포 제약 조건과 호환되지 않을 수 있습니다 (예: 칼날을 먼저 보여주는 대신 손잡이를 먼저 보여주며 건네주는 경우).

기존의 솔루션들은 다음과 같은 상당한 트레이드오프를 안고 있습니다:

1. 필터링된 재학습 (Filtered Retraining): 원치 않는 데모를 제거하고 처음부터 다시 학습하는 것은 직접적인 제어를 제공하지만, 원래 데이터에 대한 접근이 필요하며 새로운 타겟 모드마다 높은 계산 비용이 발생합니다.
2. 추론 시 스티어링 (Inference-Time Steering): 배포 중에 행동을 재순위화하거나, 확산 샘플링(diffusion sampling)을 가이드하거나, 은닉 상태(hidden states)에 개입하는 방법은 재학습을 피할 수 있지만, 외부 모듈이나 샘采样 루프를 필요로 하여 저지연 제어를 복잡하게 만들고 상당한 추론 시간 오버헤드를 발생시킵니다.

본 논문은 이미 학습된 정책에서 원본 데이터셋을 요구하지 않으면서도, 추론 시의 오버헤드 없이 원치 않는 모드를 억제할 수 있는 방법에 대한 공백을 식별했습니다.

방법론: MoRE (Mode Redirection)

저자들은 모드 재지정 신호를 정책 가중치로 직접 증류하는 사후 정책 편집(post-hoc policy-editing) 프레임워크인 MoRE를 제안합니다. 이 과정은 세 가지 주요 단계로 구성됩니다:

1. 문제 정식화

혼합 모드 정책 $\pi_{\theta_0}$가 모드 $D_k$로 분할된 데이터셋 $D_{BC}$로 학습되었을 때, 목표는 정책을 원하는 배포 집합 $S \subseteq \{1, \dots, K\}$ 내의 폐쇄 루프 롤아웃(closed-loop rollouts)이 발생하도록 새로운 파라미터 집합 $\theta^*$로 편집하는 것입니다. 편집 데이터 $D_{edit}$는 행동 모드가 레이블링된 원본 데모 또는 혼합 정책으로부터의 폐쇄 루프 롤아웃으로 구성될 수 있습니다.

2. 미분 가능한 모드 재지정 에디터 (Differentiable Mode-Redirection Editor)

MoRE는 원본 정책에서 추출된 특징(features)을 기반으로 학습된 일시적이고 경량화된 행동 모드 분류기 $g_\phi$를 활용합니다.

• 특징 추출 (Feature Extraction): 시각-언어-행وذ(VLA) 정책의 경우, 특징은 평균 풀링된 은닉 상태(mean-pooled hidden states)입니다. 확산 정책(Diffusion Policies)의 경우, 특징에는 관측 조건(observation condition)과 예측된 깨끗한 액션 청크(clean action chunk)가 포함됩니다.
• 동결 (Freezing): 분류기 파라미터 $\phi$는 학습 후 동결됩니다. 편집 단계 동안 정책 파라미터 $\theta$는 업데이트되지만 분류기는 고정된 상태를 유지합니다. 이를 통해 분류기로부터의 그래디언트가 정책 의존적 특징을 통해 정책 가중치로 흐를 수 있게 합니다.

3. 유지 및 재지정 손실을 이용한 최적화 (Optimization with Retain and Redirect Losses)

편집 목적 함수는 행동 억제와 작업 능력 사이의 균형을 맞추기 위해 두 가지 손실 항을 결려합니다:
$$L_{MoRE}(\theta) = \underbrace{\mathbb{E}_{x \sim D_{des}} [L_{BC}(x; \theta)]}_{\text{Retain}} + \gamma \underbrace{\mathbb{E}_{x \sim D_{undes}^\theta} [L_{red}(x; \theta, \phi, S)]}_{\text{Redirect}}$$

• 유지 손실 (Retain Loss): 원하는 모드 집합($D_{des}$)의 샘플에 적용되며, 이 항은 원래의 모방 손실(예: 노이즈 예측 또는 속도 매칭)을 사용하여 작업 능력을 보존하고 정책이 작업을 완료하는 방법을 잊어버리는 것을 방지합니다.
• 재지정 손실 (Redirect Loss): 원치 않는 모드 샘플($D_{undes}$)에 적용되며, 이 항은 서브셋 확률(subset-probability) 손실을 사용하여 분류기의 예측 확률 질량을 원하는 집합 $S$ 쪽으로 밀어냅니다.
• 소스 모드 마스킹 (Source-Mode Masking): 안정성을 보장하기 위해, 재지정 손실은 소스 모드의 확률이 임계값 $\tau$(예: 0.5) 미만인 원치 않는 샘플에 대해서만 적용됩니다. 이는 재지정이 모드 간에 공유되는 상태-행동 영역에서만 발생하도록 하여, 정책이 이미 궤적에 전념한 모드별 실행 분기(mode-specific execution branches)에 대한 교란을 방지합니다.

편집 후, 분류기는 폐기됩니다. 결과적인 정책 $\pi_{\theta^*}$는 추가적인 오버헤드 없이 원래의 추론 경로를 통해 실행됩니다.

주요 기여

1. Behavior Uncloning의 정식화: 본 논문은 사전 학습된 정책에서 배포에 원치 않는 모드를 억제하는 데 초점을 맞춘 새로운 정책 편집 설정을 정의하며, 작업 완료와 허용 가능한 행동 모드 준수를 모두 요구하는 "배포 성공률(Deployment Success Rate, SR)" 지표를 통해 평가합니다.
2. MoRE 프레임워크: 유지 손실을 사용하여 작업 능력을 유지하면서 분류기 가이드 재지정을 정책 가중치로 증류하는 방법을 제공합니다. 이는 단일 및 다중 타겟 편집을 지원하며, 데모 또는 롤아웃 모두에서 작동하고 편집 후 분류기를 제거합니다.
3. 포괄적인 검증: MoRE는 Diffusion Policy 및 $\pi_{0.5}$ VLA 백본, 이진 및 다중 방식 모드 제어, 시뮬레이션 및 실제 로봇 배포, 그리고 조작 및 사족 보행 네비게이션 작업을 가로질러 검증되었습니다.

실험 결과

실험은 칼 전달(knife handover), 다중 위치 배치(multi-location placement), 사족 보행 네비게이션과 같은 작업을 포함하는 8가지 작업(시뮬레이션 4개, 실제 로봇 4개)에 대해 수행되었습니다.

• 성능 향상: MoRE는 원본 혼합 모드 정책에 비해 평균 배포 성공률(SR)을 44 퍼센트 포인트 개선했습니다. 예를 들어, Push-T 작업의 경우 SR이 50.0%에서 98.3%로 증가했습니다.
• 베이스라인과의 비교:
  • 선호도 학습(CPL) 대비: MoRE는 더 강력한 모드 전환과 더 높은 SR을 달{% achieved. CPL은 부분적인 전환만을 제공했으며 때때로 작업 완료율(TCR)을 감소시켰습니다.
  • 추론 시 스티어링(DynaGuide) 대비: DynaGuide는 일부 작업에서 모드 정렬을 개선했지만, 8.47배의 추론 시간 지연을 초래했습니다. MoRE는 원래의 추론 속도를 유지했습니다.
  • 필터링된 재학습 대비: MoRE는 원본 데이터셋에 대한 접근이나 전체 재학습 없이도 필터링된 데이터로 처음부터 재학습된 정책의 성능에 근접했습니다.
• 일반화: Mo-RE는 서로 다른 정책 백본(Diffusion Policy에서 $\pi_{0.5}$ VLA로), 서로 다른 형태(팔 조작에서 사족 보행으로), 시뮬레이션에서 실제 세계 배포까지 성공적으로 전이되었습니다.
• 데이터 효율성: MoRE는 편집 데이터로 혼합 정책으로부터의 폐쇄 루프 롤아웃만을 사용했을 때도 효과적이었으나, 원본 데모를 사용했을 때 약간 더 나은 결과를 보였습니다.

의의 및 주장

본 논문은 MoRE를 "데이터 다양성"과 "안전/수용 가능성 제약" 사이의 "배포 긴장(deployment tension)"에 대한 실질적인 해결책으로 포지셔닝합니다. 가중치로 재지정 신호를 증류함으로써, MoRE는 외부 스티어링 모듈이나 처음부터의 재학습 없이도 독립적이고 안전한 정책을 생성할 수 있게 합니다.

저자들은 MoRE를 인증된 안전 계층이라기보다 **모드 제어 편집 방법(mode-control editing method)**으로 간주해야 한다고 명시적으로 밝힙니다. 그 효과는 다음의 특정 가정에 의존합니다:

• 모드 레이블이 있는 편집 데이터(또는 신뢰할 수 있게 추론 가능한 롤아웃)의 가용성.
• 궤적이 소스 모드로 완전히 전념하기 전에 특징이 모드 정보를 노출하는 분류기 인터페이스.

이 연구는 타겟팅된 행동이 모델 가중치로 직접 재지정될 수 있음을 시사하며, 이는 추론 시의 개입에 따른 계산적 또는 지연 시간 비용 없이 더 안전하고 제어 가능한 로봇 배포를 위한 경로를 제공합니다. 향ặc 연구 방향으로는 노이즈가 많은 레이블 처리, 계층적 모드 분류 체계 발견, 그리고 더 강력한 분포 변화(distribution shifts) 대응 등이 포함됩니다.