기술 요약: EmbodiedVAE

문제 정의

최근 잠재 확산 모델(Latent Diffusion Models, LDMs)의 발전은 에이전트가 미래의 상호작용을 예측하는 월드 모델을 구축할 수 있게 함으로써 체화된 학습(Embodied Learning)을 크게 향상시켰습니다. 그러나 기존의 LDM은 주로 자연스러운 장면(Natural Scenes)에 최적화된 변분 오토인코더(Variational Autoencoders, VAEs)에 의존하고 있으며, 이는 로봇 조작 시나리오에는 적합하지 않습니다. 이러한 표준 VAE는 두 가지 결정적인 한계에 직면해 있습니다:

1. 압축성 부족: 2D 이미지 VAE는 시간적 연속성은 보존하지만 프레임 간의 중복성을 활용하지 못하여, 효율적인 LDM 학습을 방해하는 큰 잠재 공간(Latent Space)을 생성합니다.
2. 제어 가능성 부족: 기존의 비디오 VAE는 시공간적 압축을 통해 압축성을 달성하지만, 종종 로봇의 동작에 대한 명시적 모델링을 놓칩니다. 이로 인해 운동 의미론(Motion Semantics)이 손실되어, 다운스트림 제어를 위한 정밀한 액션 조건 주입이 어려워집니다.

결과적으로, 현재의 방식들은 압축 효율성과 운동 충실도(Motion Fidelity) 사이의 균형을 맞추는 데 어려움을 겪고 있으며, 이는 로봇 월드 모델의 액션 제어 정확도를 제한합니다.

방법론

저자들은 로봇 조작에 특화된 새로운 비디오 VAE인 EmbodiedVAE를 제안합니다. 이 아키텍처는 로봇의 움직임을 배경 환경으로부터 자동으로 분리함으로써 압축적이면서도 제어 가능한 잠재 표현을 생성하도록 설계되었습니다.

1. 아키텍처: 이중 인코더, 단일 디코더 (Dual-Encoder, Single-Decoder)

EmbodiedVAE는 이중 인코더, 단일 디코더 구조를 사용하여 비대칭적인 시공간 압축 전략을 채택합니다:

• 로봇 팔 인코더 (Robotic Arm Encoder): 전경의 매니퓰레이터에 집중합니다. 로봇의 제한된 공간적 점유 면적에 맞춰 미세한 운동 역학을 포착하기 위해 더 강력한 공간 압축(4× 시간적, 16× 공간적)을 적용합니다.
• 환경 컨텍스트 인코더 (Environmental Context Encoder): 배경 장면을 모델링합니다. 정적이거나 변화가 느린 환경에서의 시간적 중복성을 활용하기 위해 더 강력한 시간적 압축(8× 시간적, 8× 공간적)을 적용합니다.
• 통합 디코더 (Unified Decoder): 두 개의 잠재 스트림(팔과 환경)을 독립적으로 업샘플링한 후, 이를 결합하여 공유된 테일 디코더(Tail Decoder)를 통과시켜 전체 비디오를 재구성합니다.

2. OT 기반 운동 일관성 모듈 (OT-based Motion Consistency Module)

시간적 압축으로 인한 운동 충실도 손실 문제를 해결하기 위해, 저자들은 최적 운송(Optimal-Transport, OT) 기반 일관성 모듈을 도입했습니다.

• 이 모듈은 연속된 프레임 간의 잠재 분포 사이의 정보 운송 비용을 최소화함으로써 로봇 팔의 잠재 공간에서 시간적 일관성을 강제합니다.
• 이 모듈은 안정적인 움직임 패턴이 시각적 불변성을 유지한다는 관찰을 활용합니다. 운송 비용을 최소화함으로써, 모델은 잠재 표현이 일관된 운동 역학을 보존하도록 유도하여 압축된 잠재 상태가 물리적 동작에 충실하게 유지되도록 합니다.

3. 2단계 학습 전략

모델은 다음과 같은 특정 2단계 파이프라인을 통해 학습됩니다:

• 1단계 (분리 - Disentanglement): 두 개의 별도 인코더를 학습시킵니다. 로봇 팔 인코더는 매니퓰레이터에 집중하도록 마스크(SAM2 및 Grounding DINO를 통해 생성)를 사용하여 감독되며, 환경 인코더는 배경에 집중합니다. 손실 함수에는 재구성, KL 발산, 적대적(Adversarial) 손실, 그리고 OT 기반 운동 일관성 손실이 포함됩니다.
• 2단계 (통합 재구성 - Unified Reconstruction): 1단계의 인코더들을 고정(Freeze)합니다. 분리된 잠재 표현들을 융합하여 전체 비디오를 재구성하도록 통합 디코더를 학습시킵니다. 추론 시 모델은 명시적인 마스크 입력 없이 작동합니다.

4. LDM과의 통합

EmbodiedVAE는 기존 LDM 백본(예: IRASim)과 원활하게 통합되도록 설계되었습니다. 분리된 잠재 상태($Z_{arm}$ 및 $Z_{env}$)는 토큰화되어 확산 모델을 통해 처리됩니다. 크로스 어텐션(Cross-attention) 레이어가 삽로되어 팔과 배경 표현 간의 적응적 정보 교환을 허용하며, 이를 통해 전역적으로 일관된 장면 역학과 국부적으로 정확한 움직임을 보장합니다.

주요 기여

1. EmbodiedVAE 아키텍처: 로봇 조작을 위해 특별히 설계되어 높은 압축률(0.39%)을 달성하면서도 재구성 품질을 유지하는 새로운 비디오 VAE를 제공합니다.
2. 분리된 잠재 공간: 비대칭 압축률을 가진 이중 인코더를 사용하여 로봇의 움직임을 환경 컨텍스트와 분리함으로써, 배경 노이즈의 간섭 없이 정밀한 액션 주입을 가능하게 합니다.
3. OT 기반 일관성 손실: 최적 운송 기반 모듈을 통해 잠재 공간 내에서 운동 충실도와 프레임 간 일관성을 명시적으로 강제하며, 이는 고도로 압축된 비디오 VAE에서 나타나는 운동 손실 문제를 해결합니다.
4. 2단계 학습: 추론 시 마스크 입력 없이도 자동 분리가 가능하도록 하는 학습 체계를 제공합니다.

실험 결과

Agibot-2025, Bridge, RT-1 및 약 100만 개의 비디오로 구성된 자체 수집 데이터셋을 포함한 광범위한 실험이 수행되었습니다.

• 재구성 품질: EmbodiedVAE는 최첨단 베이스라인(예: OpenSoraPlan, Wan-VAE, CMD)에 비해 우수한 재구성 성능을 달owed했습니다. Agibot-2025 데이터셋에서 훨씬 높은 압축률을 가진 경쟁 모델들(예: CMD 6.85%)보다 훨씬 낮은 압축률(0.39%)에도 불구하고 31.67의 PSNR을 달성하며 성능을 입증했습니다.
• 액션 제어 생성: IRASim-L 백본을 사용한 다운스트림 로봇 조작 작업에서, EmbodlyVAE는 가장 성능이 좋은 비디오 VAE 베이스라인인 Wan-VAE보다 PSNR을 약 2 dB 개선했습니다. 이는 다른 비디오 VAE들이 흔히 물리적으로 일관되지 않은 움직임을 생성하는 것과 달리, 로봇-환경 상호작용을 정확하게 예측하는 능력을 보여주었습니다.
• 효Efficiency: 압축적인 잠재 표현은 이미지 기반 VAE(예: SDXL) 및 다른 비디오 VAE와 비교했을 때 다운스트림 LDM 학습 시간과 메모리 소비를 크게 줄였습니다.
• Ablation 연구: OT 기반 운동 일관성 손실이나 마스크 가이드 분리를 제거했을 때 재구성과 액션 제어 작업 모두에서 측정 가능한 성능 저하가 발생하여, 이러한 구성 요소들의 필요성을 검증했습니다.

의의

본 논문은 기존의 월드 모델이 높은 압축 효율성과 정밀한 액션 제어 가능성을 동시에 달성하지 못하는 임바디드 AI의 핵심적인 병목 현상을 EmbodiedVAE가 해결한다고 주장합니다. 움직임을 컨텍스트로부터 명시적으로 분리하고 최적 운송을 통해 시간적 일관성을 강제함으로써, EmbodiedVAE는 잠재 확산 모델의 효율적인 학습과 더 정확한 로봇 제어를 가능하게 합니다. 저자들은 이 연구를 복잡한 물리적 상호작용을 신뢰성 있게 시뮬레이션하고 예측할 수 있는 확장 가능하고 지능적인 체화된 시스템을 향한 기초적인 단계로 제시합니다.