PocketDP3: Efficient Pocket-Scale 3D Visuomotor Policy

기술 요약: PocketDP3: 효율적인 포켓 규모의 3D 시각-운동 정책 (Visuomotor Policy)

문제 정의
최근 로봇 조작 분야의 발전은 복잡한 기술을 학습하기 위해 3D 비전 기반의 확산 정책(diffusion policies)을 활용해 왔습니다. 그러나 저자들은 DP3와 같은 최신 기술(SOTA) 모델에서 발생하는 결정적인 구조적 불일치를 지적합니다. 이러한 모델들은 일반적으로 가볍고 효율적인 포인트 클라우드 인코더와 과도하게 큰 디코더 백본(주로 조건부 U-Net)을 결কে합니다. U-Net은 미세한 픽셀 재구성이 필요한 고해상도 이미지 합성에는 적합하지만, 저자들은 이 설계가 3D 시각-운동 제어에는 비효율적이라고 주장합니다. 이 문맥에서 인코더는 이미 압축적이고 정보가 풍부한 장면 표현을 생성하며, 예측 대상은 저주파수의 액션 궤적(action trajectory)이기 때문입니다. 결과적으로, 거대한 디코더는 상당한 파라미터 중복을 초래하고, 메모리 점유율을 높이며, 제어 성능에 비례하는 이득을 제공하지 못한 채 자원이 제한된 로봇에서의 실시간 배포를 방해합니다. 또한, 표준 확산 샘플링은 많은 반복적인 디노이징 단계를 요구하여 실시간 제어의 레이턴시 병목 현상을 일으키며, 이를 가속화하기 위해 복잡한 일관성 증류(consistency distillation) 기술이 필요할 때가 많습니다.

방법론
본 논문은 구조적 불균형을 해결하고 추론을 가속화하기 위해 설계된 포켓 규모의 3D 확산 정책인 PocketDP3를 제안합니다.

• 아키텍처 (Diffusion Mixer - DiM): 무거운 조건부 U-Net 대신, PocketDP3는 MLP-Mixer 블록을 기반으로 구축된 가벼운 Diffusion Mixer (DiM) 디코더를 사용합니다. 이 아키텍처는 효율적인 포인트 클라우드 인코더(Farthest Point Sampling과 MLP를 사용하여 압축된 3D 표현을 생성)는 유지하되, 디코더를 $K$개의 스택된 DiM 블록으로 교체합니다.

  • 융합 메커니즘: DiM 블록은 시간적 믹싱(궤적의 시간 차원을 가로지르는 정보 집계)과 채널 믹싱(특징 차원을 가로지르는 정보 집계)을 교대로 수행함으로써 효율적인 정보 융합을 수행합니다.
  • 컨디셔닝 (Conditioning): 모델은 선형 투영을 통해 장면 표현과 로봇 상태를 통합하며, FiLM (Feature-wise Linear Modulation) 레이어와 게이트형 잔차 연결(gated residual connections)을 사용하여 이 컨디셔닝 정보를 DiM 블록에 주입합니다.
  • 파라미터 효율성: 저자들은 "Tiny" (0.53M 파라미터)와 "Base" (1.73M 파라미터)라는 두 가지 구성을 제안하며, 이는 DP3(~250M 파라미터)와 같은 기존 방식의 파라미터 수의 1% 미만입니다.

• 추론 전략: 이 방법은 DDIM 스타일의 샘플링을 활용합니다. 특히, 저자들은 PocketDP3가 추가적인 일관성 증류나 특수한 학습 목표(예: One-Step Diffusion 또는 Consistency Models) 없이도 2단계 추론(2번의 함수 평가)을 지원함을 입증했습니다. 이는 성능을 유지하면서 추론 레이턴시를 크게 줄여줍니다.

• 학습: 모델은 $x_0$-예측 목표를 사용하여 학습되며, 관측값 및 타임스텝에 조건화된 예측 클린 궤적과 실제 전문가 궤적 사이의 평균 제곱 오차(MSE)를 최소화합니다.

주요 기여

1. 아키텍적 통찰: 본 논문은 압축된 3D 인코더와 거대한 U-Net 디코더를 결합하는 것이 3D 확산 정책에서 광범위한 비효율성을 초래하여 파라미터를 낭비한다는 점을 식별했습니다.
2. PocketDP3 아키텍처: U-Net을 DiM(Diffusion Mixer)으로 대체하여 최소한의 파라미터 예산 내에서 효과적인 시간적 및 채널 융합을 가능하게 하는 포켓 규모의 정책을 도입했습니다.
3. 증류 없는 빠른 샘플링: 일관성 증류 없이도 고성능의 2단계 추론이 가능하다는 것을 입증하여, 학습 파이프라인을 단순화하고 저지연 배포를 가능하게 했습니다.
4. 실험적 검증: PocketDP3가 SOTA 성능에 부합하거나 이를 능가하면서도 모델 크기와 추론 시간을 획기적으로 줄였음을 종합적인 평가를 통해 보여주었습니다.

실험 결과
이 방법은 세 가지 시뮬레이션 벤치마크(RoboTwin2.0, Adroit, MetaWorld)와 실제 환경 실험에서 평가되었습니다.

• 시뮬레이션 성능:

  • RoboTwin2.0: PocketDP3-Base는 **71.6%**의 평균 성공률을 달성하여, 이전 SOTA인 DP3보다 20.8% 높은 성능을 보였습니다. PocketDP3-Tiny 또한 DP3보다 15.2% 높은 성능을 기록했습니다.
  • Adroit & MetaWorld: PocketDP3-Base는 **77.4%**의 평균 성공률을 달성하여, DP3(73.0%)와 FlowPolicy(72.6%)를 앞질렀습니다.
  • 절제 연구 (Ablation Studies): "Vanilla-MLP"(수렴 실패) 및 축소된 "Vanilla-U-Net"(성능 미달)과의 비교를 통해, 엄격한 파라미터 예산 하에서 DiM 블록의 특정 융합 메커니즘이 성능에 결정적임을 확인했습니다.
  • 추론 단계: 함수 평가 횟수(NFE)에 대한 절제 연구 결과, 성능이 2단계에서 정체됨을 보였습니다. 단계를 5단계나 10단계로 늘려도 유의미한 이득은 없었습니다.

• 효율성:

  • 모델 크기: DP3(255.1M)와 비교했을 때 PocketDP3-Base는 1.73M 파라미터를 사용하여 99% 이상의 감소를 달精度했습니다.
  • 레이턴시: 2단계 추론 시, PocketDP3-Base는 RTX 5880 Ada에서 4.80ms의 레이턴시를 달성했습니다. 이는 DP3(10단계 시 51.4ms)보다 훨씬 빠르며, FlowPolicy(1단계 시 7.04ms)와 유사하지만 파라미터 수는 두 자릿수 차이로 훨씬 적습니다.

• 실제 환경 배포:

  • AgileX Piper 로봇과 RealSense D455 카메라를 사용하여 테스트되었습니다.
  • PocketDP3-Base는 물체 배치 및 블록 쌓기와 같은 작업에서 DP3 베이스라인보다 평균 15.3% 높은 성능을 보이며 강력한 Sim-to-Real 전이 능력을 입증했습니다.

의의 및 주장
본 논문은 PocketDP3가 효율적인 3D 시각-운동 제어를 향한 실용적인 경로를 제공한다고 주장합니다. 입력 표현이 이미 압축되어 있을 때 거대한 디코더가 반드시 필요하다는 가설에 도전함으로써, 저자들은 이 특정 문맥에서 "작은 것이 더 낫다(smaller is better)"는 것을 보여주었습니다. 이 방법은 기존 방식의 1% 미만 파라미터로 SOTA 성능을 달성하며, 복잡한 증류 학습 없이 실시간 추론을 가능하게 합니다. 저자들은 이 연구를 자원이 제한된 하드웨어에 정교한 로봇 정책을 배포하기 위한 단계로 규정합니다.

한계 및 향후 과제
저자들은 U-Net이 가진 픽셀 정렬된 공간적 편향(spatial bias)의 부재로 인해, 미세한 국소 기하학적 추론(예: 서브 밀리미터 단위 삽입)이 필요한 작업에서 DiM 아키텍처가 "공간적 병목(spatial bottleneck)"을 겪을 수 있음을 인정합니다. 또한, 고품질의 2단계 추론을 가능하게 하는 근본적인 메커니즘이 여전히 불분명하며 추가 연구가 필요하다고 언급했습니다. 향후 연구는 더 효율적이고 정보가 풍부한 컨디셔닝 인코더를 설계하고, 컨디셔닝 용량, 작업 복잡도, 디코더 규모 간의 트레이드오프를 체계적으로 연구하는 데 집중할 것을 제안합니다.