기술 요약: NoiseGate – 세계 행동 모델에서 정보 게이트로서 잠재 변수별 시간 단계 스케줄 학습

1. 문제 정의

세계 행동 모델 (WAMs) 은 로봇 행동 생성이 미래 관측 모델링과 결합된 패러다임으로, 현재 관측치 ($v_0$) 와 언어 지시 ($l$) 를 기반으로 미래 잠재 비디오 프레임 ($v$) 의 일부와 행동 블록 ($a$) 을 공동 생성합니다. Motus 와 Fast-WAM 과 같은 대표적인 구현체들은 Mixture-of-Transformers (MoT) 아키텍처에 기반한 결합 비디오 - 행동 디노이징 백본을 활용합니다. 이러한 시스템에서 비디오 토큰과 행동 토큰은 공유된 자기 주의 (self-attention) 계층을 통해 상호작용합니다.

기존 WAM 의 중요한 한계는 확산 시간 단계 (diffusion timestep) 의 처리 방식에 있습니다. 아키텍처 이론상 각 예측된 잠재 프레임이 독립적인 시간 단계 ($t_f$) 를 가질 수 있음에도 불구하고, 현재 시스템들은 모든 예측 프레임에 대해 각 디노이징 단계에서 이 자유도를 단일 공유 스칼라 시간 단계 ($t$) 로 축소합니다. Diffusion Forcing 의 "노이즈 - 마스킹" 해석 하에서, 이 공유 스케줄은 정당화되지 않은 사전 확률을 부과합니다. 즉, 모든 예측된 잠재 프레임이 어떤 단계에서도 행동 생성 과정에 대해 동등하게 신뢰할 수 있고 "가시적"이라고 가정합니다.

그러나 미래 프레임이 행동 예측에 미치는 영향은 본질적으로 이질적이며 작업에 따라 다릅니다. 예를 들어, 잡기 (grasping) 와 같은 중요한 사건 동안, 모델은 불확실한 예측에 대한 과도한 의존을 피하기 위해 해당 상상된 잠재 프레임을 더 "흐릿하게"(높은 노이즈) 유지하는 것이 유리할 수 있습니다. 수동으로 설계된 스케줄 (예: 단조로운 노이즈 감소) 은 단순히 하나의 고정된 사전 확률을 다른 것으로 대체할 뿐이며, 서로 다른 작업이 요구하는 특정 정보 밀도에 적응하지 못합니다.

2. 방법론: NoiseGate

저자들은 NoiseGate를 제안하는데, 이는 잠재 변수별 시간 단계 스케줄을 학습 가능한 정보 게이트 정책으로 재해석하는 프레임워크입니다. 고정된 스케줄 대신, NoiseGate 는 각 잠재 프레임의 키/값 (Key/Value) 기여도가 행동 토큰에 미치는 신뢰도를 동적으로 노이즈 수준을 조정함으로써 조절하는 법을 학습합니다.

이 방법은 세 가지 결합된 구성 요소로 이루어집니다:

A. 결합 비디오 - 행동 MoT 백본

이 시스템은 청크 양방향 비디오 Diffusion Transformer (DiT) 와 행동 전문가 (Action-Expert) DiT 를 통합하는 Mixture-of-Transformers 백본을 사용합니다.

• 공유 자기 주의: 비디오 및 행동 토큰이 공유 자기 주의 블록을 공유하여, 행동 토큰이 비디오 잠재 변수에 주의를 기울일 수 있게 합니다.
• 모달리티별 전문가: 각 모달리티는 자체 피드포워드 계층을 유지합니다.
• 게이트 메커니즘: 비디오 잠재 변수의 노이즈 수준 $t_f$는 행동 토큰의 숨겨진 표현에 대한 그 기여도의 신뢰도를 직접 조절합니다.

B. 독립적 잠재 변수별 시간 단계 샘플링 (학습 기저)

추론 중 임의의 프레임별 스케줄에 대해 백본이 일반화할 수 있도록, 저자들은 Diffusion Forcing 에서 영감을 받은 학습 전략을 채택합니다.

• 학습 중에는 공유 스칼라 대신 각 잠재 프레임에 대해 독립적인 시간 단계 $t_v = (t_1, \dots, t_F)$가 샘플링됩니다.
• 이는 "이질적으로 노이즈가 추가된" 입력을 생성하여, 백본이 비디오 청크 전반에 걸친 다양한 불확실성 수준을 효과적으로 "읽을" 수 있도록 가르칩니다.

C. 게이트 정책 네트워크 (GPN) 및 GRPO 최적화

핵심 혁신은 최적의 시간 단계 스케줄을 학습하는 경량 네트워크인 게이트 정책 네트워크 ($\pi_\phi$) 입니다.

• 입력: 각 디노이징 단계 $k$에서 GPN 은 현재 이질적으로 노이즈가 추가된 비디오 청크 ($v_t^{(k)}$), 깨끗한 관측치 ($v_0$), 그리고 현재 시간 단계 ($t_v^{(k)}$) 를 관찰합니다.
• 출력: 각 프레임의 시간 단계가 얼마나 감소해야 하는지 ($\Delta t_f$) 를 결정하기 위해 상대적 스케일 증가분 ($r_f$) 을 예측합니다. 관측 프레임은 $t_0=0$에 고정되며, 행동 시간 단계는 고정된 전역 스케줄을 따릅니다.
• 학습: GPN 은 그룹 상대적 정책 최적화 (GRPO) 를 사용하여 학습됩니다. 백본은 고정된 채로 유지되며, 정책은 희소 이진 작업 성공 보상 ($r \in \{0, 1\}$) 에 대해 최적화됩니다. 이는 수동으로 설계된 형태 사전 확률 없이 실제 작업 유용성에 게이트 정책을 기반을 두게 합니다.

3. 주요 기여

1. 스케줄의 재해석: 저자들은 결합 비디오 - 행동 WAM 의 기존 공유 스칼라 시간 단계를 모델 유연성을 제한하는 강력한 암묵적 사전 확률로 식별합니다. 그리고 잠재 변수별 스케줄링을 공유 자기 주의에 대한 학습 가능한 정보 게이트 정책으로 바라볼 것을 제안합니다.
2. NoiseGate 아키텍처: 백본 학습 중 독립적 잠재 변수별 시간 단계 샘플링, 잠재 변수별 시간 증가분을 방출하는 경량 GPN, 그리고 작업 보상에 의해 주도되는 GRPO 최적화를 결합한 NoiseGate 를 소개합니다.
3. 실증적 검증: 무작위 장면 조건 하의 RoboTwin 벤치마크에서 NoiseGate 는 공유 스칼라 베이스라인과 수동 설계 스케줄 변형 모두보다 일관된 개선을 보여주어, 게이트 정책을 학습하는 것이 잠재 변수별 자유도를 활용하는 데 필수적임을 확인시켰습니다.

4. 실험 결과

평가 작업은 RoboTwin벤치마크에서 수행되었으며, 특히 물체 자세, 색상, 배경이 무작위화되는 무작위 장면 조건 하에서 수행되었습니다.

• 주요 비교: NoiseGate 는 강력한 베이스라인 (Fast-WAM, LingBot-VA, $\pi_{0.5}$, Motus) 과 Stage-1 WAM(독립 노이즈로 학습되었으나 공유 스칼라 추론을 사용하는 백본) 과 비교되었습니다.
  • NoiseGate 는 평균 성공률 **94.28%**를 달성하여 Stage-1 WAM(92.58%) 및 기타 모든 베이스라인을 능가했습니다.
  • 정밀한 타이밍이 필요하거나 불확실성을 처리해야 하는 작업에서 유의미한 향상이 관찰되었습니다 (예: Hanging Mug +9%, Blocks Ranking Size +7%, Put Object Cabinet +7%).
• Ablation 연구:
  • GPN 을 제거하고 공유 스칼라 ($Shared\text{-}t$) 로 되돌리는 것은 가장 큰 성능 저하를 초래했습니다 (통제된 ablation 설정에서 67.5% $\to$ 57.5%).
  • 수동 설계된 단조 스케줄을 사용하는 것은 공유 스칼라보다 성능을 향상시켰지만 학습된 GPN 보다 여전히 열등했습니다 (63.4% 대 67.5%), 이는 작업 적응형 제어가 필요함을 시사합니다.
• 정성적 분석:
  • 마스킹으로서의 노이즈: 주의 맵 (attention maps) 은 노이즈 수준이 증가함에 따라 행동에서 비디오로의 주의가 단조롭게 감소함을 확인시켜 게이트 메커니즘을 검증했습니다.
  • 작업 적응형 스케줄: 시각화 결과, GPN 이 서로 다른 작업에 대해 고유하고 비균일한 시간 단계 궤적을 학습함을 보여주었습니다. 예를 들어, 잡기 작업에서 GPN 은 다른 프레임들보다 중요한 잡기 프레임을 더 오랫동안 더 높은 노이즈 수준 (더 불확실한 상태) 에 유지하여 조기 행동 결정을 방지했습니다.

5. 중요성 및 주장

이 논문은 NoiseGate 가 세계 행동 모델 내의 잠재 변수별 시간 단계 스케줄이 고정된 하이퍼파라미터나 암묵적 사전 확률이 아닌 **1 차 설계 객체 (first-class design object)**임을 성공적으로 입증했다고 주장합니다.

• 메커니즘: 시간 단계를 학습 가능한 게이트로 취급함으로써, 모델은 행동 정책에 미래 정보를 선택적으로 노출시킬 수 있습니다. 이를 통해 시스템은 예측하기 어려운 중요한 미래 사건에 대해서는 더 높은 불확실성 (흐림) 을 유지하면서 덜 중요한 정보는 명확히 하여 복잡한 조작 작업에서의 견고성을 향상시킵니다.
• 효율성: 이 접근법은 복잡한 수동 설계 사전 확률이나 거대한 생성 백본의 재학습 없이 경량 정책 네트워크와 희소 보상을 사용하여 스케줄을 최적화합니다.
• 일반화 가능성: 로봇 공학에 초점을 맞추고 있지만, 저자들은 잠재 변수별 노이즈를 학습 가능한 정보 게이트로 보는 관점은 하류 소비자가 이질적 유용성을 가진 결합 디노이징 토큰을 읽는 다른 생성 설정 (예: 자율 주행, 비디오 예측) 에도 적용될 수 있다고 지적합니다.

저자들은 현재 학습이 희소 시뮬레이션 롤아웃에 의존하여 데이터 비효율적일 수 있으며, 학습된 스케줄 정책을 대규모 다중 작업 세트로 확장하는 것은 향후 과제로 남아있음을 지적하며 한계에 대해 겸손한 태도를 유지합니다.