DiSPo: Diffusion-SSM based Policy Learning for Coarse-to-Fine Action Discretization

1. 문제 정의 (Problem)

로봇 조작 (Manipulation) 작업은 종종 **다중 해상도 (Multi-granularity)**의 행동을 요구합니다. 예를 들어, 큰 위치 이동 (Coarse) 과 나사 조임, 용접, 삽입과 같은 정밀한 미세 조작 (Fine) 이 혼합되어 있습니다. 기존 학습 방법들은 다음과 같은 한계를 가집니다:

• 데이터 의존성: 정밀한 제어를 위해서는 고주파수 (High-frequency) 의 세밀한 데모 데이터가 필수적이지만, 이를 수집하는 것은 비용이 많이 들고 비효율적입니다.
• 해상도 고정: 기존 확산 정책 (Diffusion Policy) 이나 행동 트랜스포머 (Behavior Transformer) 는 특정 주파수나 시간 척도에서 학습된 데이터를 기반으로 하여, 학습된 주파수와 다른 해상도의 행동을 생성하거나 다양한 해상도의 데모를 통합하여 학습하는 데 어려움을 겪습니다.
• 보간 및 동역학 모델의 한계: 기존 방법은 단순 보간이나 동역학 모델에 의존하여 다양한 해상도의 행동을 인코딩/디코딩하는 데 한계가 있습니다.

핵심 목표: coarse (거친) 데모 데이터로부터 학습하여, 사용자의 의도에 따라 fine-grained (세밀한) 행동을 생성할 수 있는 다중 해상도 학습 (Multi-granularity Learning) 및 **재생 (Reproduction)**을 가능하게 하는 정책 개발.

---

2. 제안 방법론: DiSPo (Methodology)

저자들은 **Diffusion-SSM 기반 정책 (DiSPo)**을 제안합니다. 이는 확산 모델 (Diffusion Model) 의 표현력과 상태 공간 모델 (SSM, 구체적으로 Mamba) 의 유연한 이산화 (Discretization) 능력을 결합한 하이브리드 아키텍처입니다.

A. 핵심 아키텍처

• Mamba 기반 확산 과정: 기존 확산 모델이 CNN 또는 Transformer 를 백본으로 사용하는 반면, DiSPo 는 **Mamba(최신 SSM)**를 기반으로 합니다. 이는 긴 시퀀스 처리에 효율적이며 메모리 사용량이 적습니다.
• Step-scaling Mechanism (단계 스케일링): DiSPo 의 가장 큰 혁신입니다. Mamba 의 이산 파라미터 (특히 단계 크기 $\Delta$) 를 사용자 정의 **단계 스케일 팩터 (Step-scale factor, $r_t$)**에 따라 동적으로 조정합니다.
  • 입력: 확산 단계 ($k$), 관찰 ($o_t$), 노이즈가 있는 행동 ($a^{(k)}_t$), 그리고 단계 스케일 벡터 ($r_t$).
  • $r_t$는 관찰 및 행동 시퀀스의 시간적 해상도 (Granularity) 를 조절합니다. 예를 들어, $r_t=0.5$는 2 배 더 정밀한 행동을, $r_t=2$는 2 배 더 거친 행동을 의미합니다.
• 아키텍처 구성:
  • DiSPo Block: Mamba 블록의 변형으로, 적응형 레이어 정규화 (adaLN) 를 사용하여 확산 단계 임베딩에 조건부 (Conditioned) 로 노이즈 관련 특징을 정제합니다.
  • Step-scaled SSM: 입력 시퀀스와 $r_t$를 기반으로 SSM 의 단계 크기 ($\Delta$) 를 예측하고, 이를 통해 이산화된 SSM 파라미터 ($\bar{A}, \bar{B}$) 를 업데이트합니다.

B. 다중 해상도 학습 전략 (Multi-granularity Learning)

고주파수 데이터가 없는 상황에서도 정밀한 행동을 생성하기 위해 두 단계의 학습 전략을 사용합니다.

1. 샘플률 증강 (Sample-rate Augmentation) 프리트레이닝:
   • 기존 데모 데이터에 무작위 단계 스케일 팩터를 적용하여 다양한 주파수의 시퀀스를 인위적으로 생성하고, 이를 통해 모델이 다양한 해상도를 학습하도록 합니다.
2. 가짜 데모 (Pseudo Demonstration) 를 통한 파인튜닝:
   • 학습된 모델을 사용하여 저주파수 데모를 고주파수 타겟으로 변환합니다.
   • 반복적 노이즈 제거 및 교체: 확산 과정에서 생성된 행동 시퀀스 중, 원본 데모와 일치하는 부분 (저주파수 성분) 은 실제 데모 데이터로 교체하고, 나머지 부분만 모델이 생성하도록 하여 고주파수 정밀도를 높입니다.
   • 이를 통해 추가 데이터 수집 없이 고주파수 행동 생성 능력을 향상시킵니다.

---

3. 주요 기여 (Key Contributions)

1. Coarse-to-Fine 행동 생성: coarse 데모로부터 fine-grained 행동을 생성하는 최초의 확산-SSM 기반 정책 제안.
2. 동적 해상도 조절 (Online Modulation): 학습된 예측기 (Predictor, $\phi_r$) 를 통해 작업 상황에 따라 실시간으로 행동의 이산화 수준 (Step size) 을 조절할 수 있습니다. 이는 불필요한 구간에서는 coarse motion 을, 정밀한 구간에서는 fine motion 을 생성하여 추론 효율성을 극대화합니다.
3. 데이터 효율성: 고주파수 데이터 없이도 coarse 데이터만으로 정밀한 제어를 학습할 수 있으며, 가짜 데모 생성 기법을 통해 성능을 추가적으로 향상시킵니다.
4. 새로운 벤치마크: 다양한 주파수의 데모에서 coarse-to-fine 성능을 평가하는 새로운 시뮬레이션 및 실세계 벤치마크를 구축했습니다.

---

4. 실험 결과 (Results)

시뮬레이션 벤치마크

• 작업: Clamp passing (파이프 통과), Passage passing (좁은 통로 통과), Button touch (버튼 터치).
• 성능: 2.5Hz ~ 20Hz 의 다양한 주파수 데모로 학습하여 20Hz 타겟으로 테스트한 결과, DiSPo 는 81% 이상의 성공률을 기록했습니다.
• 비교: 기존 Diffusion Policy (DP-C, DP-T), Mamba 기반 모델 (D-Ma), VQ-BeT 등 최첨단 베이스라인보다 성능이 월등히 높았습니다. 특히 2.5Hz 같은 저주파수 데모에서는 베이스라인이 급격히 성능이 떨어지는 반면, DiSPo 는 안정적인 성능을 유지했습니다.
• 혼합 주파수 학습: 2.5Hz 와 5Hz 가 섞인 데이터로 학습했을 때, DiSPo 는 93% 의 성공률로 모든 베이스라인을 압도했습니다.

실세계 실험 (Real-world Experiments)

• 하드웨어: UR5e 로봇 팔 사용.
• 작업: Clamp passing (원형 링 삽입), Button touch (셔터 버튼 터치).
• 결과:
  • DiSPo 는 D-Ma 보다 높은 성공률을 보였습니다.
  • D-Ma 는 거친 움직임은 잘 따라가지만, 정밀한 삽입 시 충돌 (Pipe scratching) 이나 버튼 근처에서 멈추는 문제가 발생했습니다.
  • DiSPo 는 2.5mm 의 여유 공간 (Radial clearance) 을 가진 정밀한 삽입과 충돌 없는 버튼 터치를 성공적으로 수행했습니다.
  • 효율성: 적응형 단계 스케일링을 적용했을 때, 버튼 터치 작업에서 필요한 행동 단계 수를 32% 감소시키면서도 성공률은 크게 떨어지지 않았습니다.

---

5. 의의 및 결론 (Significance)

• 효율성과 정밀도의 균형: DiSPo 는 중요한 구간에서는 정밀한 제어를, 그렇지 않은 구간에서는 거친 제어를 수행하여 **추론 비용 (Inference overhead)**을 줄이면서도 작업 성공률을 높이는 새로운 패러다임을 제시했습니다.
• 데이터 수집 비용 절감: 고주파수 데모 데이터의 수집은 로봇 학습에서 큰 병목 현상이었는데, DiSPo 는 coarse 데이터만으로도 정밀한 제어를 가능하게 함으로써 실제 로봇 적용의 장벽을 낮춥니다.
• SSM 의 로봇 적용 확장: Mamba 와 같은 최신 SSM 이 로봇 제어의 '행동 이산화 (Action Discretization)' 문제를 해결하는 데 효과적임을 입증했습니다.

결론적으로, DiSPo 는 다양한 해상도의 데모를 통합적으로 학습하고, 사용자의 필요에 따라 실시간으로 행동의 정밀도를 조절할 수 있는 강력한 정책 학습 프레임워크입니다.