[논문리뷰] Acting Through Registers

링크: 논문 PDF로 바로 열기

저자: Yiwei Zhang, Yantai Yang, Jin Gao, Jin Lou, et al.

1. Key Terms & Definitions (핵심 용어 및 정의)

• World Action Models (WAMs): Generative visual representation을 활용하여 로봇 Action을 Conditioning하는 모델입니다. 이는 로봇 Policy가 대규모 비디오 Pretraining 및 물리 역학의 Predictive Representation으로부터 이점을 얻도록 합니다.
• Spatial Outliers: WAM Representation 내에서 지속적으로 나타나는 높은 Hidden-state Norm을 가진 Spatial Token을 지칭합니다. 이 Token들은 극단적인 Magnitude를 가지지만 Action Relevance를 유지하며, 그 기능적 역할이 공간적 위치로부터 부분적으로 분리되었음을 시사합니다.
• Registers: World Model에 추가되는 학습 가능한 Non-spatial Token입니다. 이들은 이미지 Region에 명시적인 대응 관계가 없으며, World Model 내에서 Spatial Token과 상호작용하고 Action Model을 직접 Conditioning하여 Spatial Identity에 묶이지 않는 Representation Capacity를 제공합니다.
• Observation Flow-Matching (FM) Loss Supervision: Dense Clean Observation Token의 Output에 적용되는 Loss Supervision 기법입니다. 이는 Spatial Outlier를 억제하고 높은 Norm Response를 Registers에 집중시키며, Visual Observation에 대한 예측력을 유지하도록 장려합니다.
• Action DiT (Action Diffusion Transformer): Video DiT에서 생성된 Representation에 Conditioning되어 로봇 Action을 Denoising하는 Diffusion Transformer 기반의 Action Model입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 World Action Models (WAMs)의 Visual Representation이 Action-relevant Information을 어떻게 조직하는지에 대한 이해 부족 문제를 해결하고자 합니다. 기존 WAM에서는 Visual Representation의 Middle-to-Late Layer에서 극단적으로 높은 Hidden-state Norm을 보이는 'Spatial Outliers'가 지속적으로 관찰됩니다. 이러한 Outlier들은 Magnitude가 정규화된 후에도 Action Relevance를 유지하지만, 그 기능적 역할이 Nominal Spatial Location과 부분적으로 Decouple되어 있다는 점이 문제로 지적됩니다. 즉, Local Image Region에 초점을 맞춘 Spatial Token들이 Action Generation에 유용하지만 특정 공간에 자연스럽게 연결되지 않는 정보를 표현하는 역할을 수행하며 Architectural Tension이 발생합니다. 이는 Spatial Capacity가 아닌 Non-spatial Capacity를 통해 이러한 Representational Role을 더 잘 수용할 수 있다는 가설로 이어집니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 연구는 WAM의 Spatial Outlier 문제를 해결하고 Action Learning 성능을 향상시키기 위해 RegisterWAM을 제안합니다. RegisterWAM은 World Model에 학습 가능한 Non-spatial Registers를 도입하고, 이 Registers를 Action DiT의 Action Conditioning을 위한 주요 인터페이스로 사용합니다. Registers 단독으로는 Spatial Outliers를 완전히 제거하지 못하기 때문에, 저자들은 Observation Flow-Matching (FM) Loss Supervision을 Dense Clean Observation Token의 Output에 추가합니다. 이 Joint Design은 Spatial Token에서 높은 Norm을 가진 Outliers를 억제하고, 그 Response를 Registers에 집중시켜 Representational Reorganization을 유도합니다. 또한, Observation FM Loss Supervision이 Registers 없이도 Action Learning을 개선한다는 점에 착안하여, Current Scene Modeling을 Corrupted Input으로부터의 Recovery Task로 전환하는 Noised-observation FM Loss Supervision을 Auxiliary Objective로 도입하여 World Model에 추가적인 Supervision을 제공합니다.

이러한 제안 방법론의 핵심 결과는 다음과 같습니다:

• LIBERO-Plus 테스트 세트에서 RegisterWAM은 추가 Embodied Pretraining Data나 Foundation Model 없이 81.64%의 높은 전체 Success Rate를 달성하여, 기존 FastWAM 및 Faster-WAM 대비 우수한 성능을 보였습니다.
• RoboTwin 2.0 벤치마크에서는 93.0%의 평균 Success Rate를 기록하여, FastWAM의 91.9% 및 Faster-WAM의 92.6%를 능가했습니다.
• 실제 환경(Real-world)의 6가지 로봇 Manipulation Task에서 FastWAM 대비 평균 Task Score를 0.3483에서 0.5238로 크게 향상시켰습니다.
• Ablation Study 결과, Noised-observation FM Loss Supervision 단독 적용 시 LIBERO-Plus Subset에서 Success Rate가 76.29%에서 78.59%로 개선되었으며, 여기에 8개의 Registers와 Observation FM Loss Supervision을 결합했을 때 82.52%로 추가 향상되었습니다. 최적의 성능은 8 Registers 및 Noised-observation Noise Level σ=0.5 설정에서 확인되었습니다.

4. Conclusion & Impact (결론 및 시사점)

본 논문은 World Action Models (WAMs)에서 발견되는 Persistent High-norm Spatial Outliers가 Action-relevant Information을 Nominal Spatial Location과 분리된 형태로 인코딩한다는 점을 밝혀냈습니다. 저자들은 이를 해결하기 위해 Non-spatial Registers와 Observation Flow-Matching (FM) Loss Supervision을 결합한 RegisterWAM을 제안하며, 이 방법이 Spatial Outliers를 효과적으로 억제하고 높은 Norm Response를 Registers에 집중시켜 Action Conditioning을 위한 효율적인 인터페이스를 제공함을 입증했습니다. 또한, Noised-observation FM Loss Supervision의 도입은 World Model의 Action Learning 성능을 추가적으로 강화하는 효과를 가져왔습니다. 이러한 연구 결과는 LIBERO-Plus, RoboTwin 2.0, 그리고 Real-world Manipulation Task 등 다양한 로봇 벤치마크에서 기존 WAM 대비 우월한 성능을 달성함으로써 로봇 Manipulation 분야에 상당한 발전을 가져왔습니다. 궁극적으로 이 연구는 Robust하고 효율적인 WAM 설계를 위해 명시적인 Non-spatial Capacity와 적절한 Auxiliary Supervision의 중요성을 강조하며, 해당 분야의 향후 연구 방향에 중요한 시사점을 제공합니다.

Figure 1 (c) — FastWAM의 Spatial Outlier 시각화

Figure 7 — 제안하는 RegisterWAM 아키텍처

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.