행동 기반 세계 모델링으로부터 전이 가능한 동역학 프라이어 학습

로봇 조작에서 행동은 장면 변화를 인과적으로 규정하는 신호이므로 행동을 조건으로 한 대규모 사전학습은 외형적 영상 생성 능력보다 재사용 가능한 상호작용 동역학 지식을 직접 학습한다. 이 논문은 실제 로봇 동작 어노테이션이 포함된 2,156k개의 궤적으로 멀티뷰 확산 기반 세계 모델을 사전학습하여 그 사전(prior)을 시뮬레이터와 정책 양쪽으로 재사용하는 실용적 경로를 제시했다. 결과적으로 현실 롤아웃을 대체하는 장기 시뮬레이션과 직접 액션 예측 정책의 두 흐름에서 전이 이득이 관측되었다.

로봇 조작 문제는 장면 외형과 로봇 본체, 카메라 배치가 크게 달라지더라도 행동이 유발하는 상태 변화 규칙은 비교적 보편적이다. 기존의 텍스트-또는 비디오 기반 사전학습은 외형적 패턴이나 언어-비디오 정렬에 최적화되는 경우가 많아 직접적인 행동-동역학 규칙을 포착하는 데 한계가 있다. 행동을 조건으로 미래 장면 변화를 예측하는 접근은 행동이 유발하는 물리적 상호작용과 접촉 이벤트, 운동 궤적의 결과를 직접 학습하도록 유도하므로 재사용 가능한 동역학 표현을 형성한다.

A2World는 이 직관을 실현하기 위해 DiT 기반 확산 블록 내부에 행동 임베딩을 타임스텝 임베딩에 합산하고 AdaLN을 통해 블록별로 동적 모듈레이션을 적용한다. 멀티뷰 입력은 시간 차원으로 병합된 잠재 시퀀스로 처리되며 뷰 식별을 위한 learnable view embeddings가 채널 차원에 연결된다. 각 블록에는 cross-view attention이 삽입되어 서로 다른 카메라 뷰들 간의 공간·시간적 일관성을 강제하고 뷰 특유의 세부를 보존하도록 설계되었다.

이러한 행동 조건화와 멀티뷰 일관성 메커니즘을 학습한 모델은 두 가지 방식으로 재사용될 수 있다. 하나는 A2World-sim처럼 과거 이력을 효과적으로 주입하고 자기회귀적으로 롤아웃하여 장기 시뮬레이터로 동작시키는 것이고, 다른 하나는 A2World-policy처럼 비디오와 액션 토큰을 공동 생성하는 MoE-like 구조로 손쉽게 전이되어 즉시 정책 생성 성능을 제공하는 것이다. 이 전이는 별도의 액션 전용 사전학습 없이도 행동-동역학 prior를 정책 학습과 시뮬레이션 양쪽에 동시에 전달한다.

전체 접근은 액션-조건화된 확산 기반 잠재 비디오 생성기를 대규모 로봇 조작 데이터로 사전학습하고 동일 체크포인트를 두 가지 다운스트림으로 적응시키는 방식이다. 베이스 A2World는 WAN2.1 토크나이저로 인코딩된 비디오 잠재를 입력으로 받고, DiT 블록을 통해 EDM denoising score matching 목표로 학습된다. 행동 청크는 MLP로 임베딩되어 확산 타임스텝 임베딩에 더해지고 AdaLN을 통해 각 블록의 scale·shift·gate를 동적으로 조정하여 행동 신호가 시간적 모듈레이션으로 투영되도록 한다.

멀티뷰 처리는 V개의 뷰를 시간 차원에 연속으로 쌓아 단일 확산 패스로 처리하고, 각 뷰에 대한 learnable view embedding을 채널에 concat하여 카메라 정체성을 부여한다. 블록 내부에 삽입된 cross-view attention은 한 뷰의 토큰이 다른 뷰의 토큰을 참조하도록 하여 시공간 일관성을 강화한다. 손실은 잠재공간에서의 L2 denoising으로 구성되며 노이즈 수준 σ에 대한 기대값으로 최적화된다.

A2World-sim으로의 전환에서는 과거 이력의 압축 표현을 도입하기 위해 pose-guided arc-uniform sampling을 사용하여 운동적으로 중요한 프레임을 균등 배치하였고, 선택된 이력은 Tok_hist로 토크나이즈되어 cross-attention 조건으로 투입된다. 추가로 동일 이력을 key/value 메모리로 투사하여 self-attention 경로에 병합함으로써 글로벌 역사 메모리를 제공한다. 롤아웃은 청크 단위의 자기회귀로 수행되며 self-forcing 학습으로 모델이 자체 생성 오류에 견고해지도록 했다.

A2World-policy 적응은 비디오와 액션을 공동으로 모델링하는 joint diffusion 포뮬레이션을 사용한다. 비디오와 액션 잠재에 독립적 가우시안 노이즈를 추가하고 shared-timestep 변형에서는 기본 노이즈 레벨을 모달리티별 스케일로 조정하여 정렬을 개선한다. 모델 블록은 공유 self-attention과 모달리티별 AdaLN 및 MLP 분기를 결합한 MoE-like 구조를 채택하여 상호작용 표현을 공유하면서도 액션 전용 덴어이징 경로를 유지한다. 추론 시에는 modality-wise classifier-free guidance를 적용하여 비디오 품질과 액션 정확성 사이의 트레이드오프를 제어한다.

사전학습과 미세조정 실험에서 A2World 기반 모델은 롤아웃 품질, 행동-유사성 지표, 정책 성공률 등에서 강한 성능을 보였다. LIBERO 평가에서 A2World-sim은 PSNR 26.64, SSIM 0.8957, tSSIM 0.7862를 기록하였고 행동 일치도 지표로 EPE 0.3498과 cos 0.6045를 달성하여 비교군 대비 전반적으로 우수한 롤아웃 품질을 보였다. 실험실 실험의 실제 로봇 데이터에서도 A2World-sim은 PSNR 25.95와 SSIM 0.8511을 달성하여 현실 롤아웃 품질을 개선했다.

시뮬레이터 유효성 측정에서는 A2World-sim의 롤아웃 기반 정책 성공률이 실제 로봇 성공률과 높은 상관을 보였다(Spearman ρ=0.916, Pearson r=0.965, R²=0.930, N=8), 이는 학습된 세계 모델이 정책 평가용 시뮬레이터로 실용적임을 시사한다. 정책 관점에서 A2World-policy는 LIBERO 기준 평균 성공률 98.6%를 달성하여 여러 최신 방법들과 경쟁 또는 우위를 보였고, OOD 검증(LIBERO-Plus Spatial)에서 action-pretrained 초기화(A-pre)가 평균 88.5%로 C-init(80.2%)과 T-pre(85.8%)보다 우수하였다.

절제된 분석으로는 이력 샘플링의 효과가 확인되었고, pose-guided 샘플링이 없는 경우 PSNR 25.41, sliding-window는 25.63인 반면 pose-guided는 26.64로 개선되었다. 또한 비디오 분기를 고정(frozen)한 정책 변형은 성공률이 86.2%로 떨어져 비디오 예측과 액션 생성 간의 공동 학습이 상호 이득을 준다는 증거가 확보되었다. RoboNet 평가에서도 A2World-sim은 FVD 146.1을 기록하여 일부 기존 자동회귀 기반 방법 대비 우수한 정량 성능을 보였다.

전체 아키텍처는 DiT 기반 확산 블록을 중심으로 구성되며 입력은 WAN2.1 토크나이저로 인코딩된 잠재 비디오 시퀀스이다. 행동은 각 청크를 MLP로 임베딩하여 확산 타임스텝 임베딩과 합쳐지고, 이 합성 임베딩은 AdaLN을 통해 각 블록의 scale, shift, gate 파라미터로 변환되어 블록 내 연산에 동적 영향을 준다. 멀티뷰 처리는 V뷰를 시간 축으로 연속적으로 배치한 잠재 텐서 형태로 입력하고, 각 뷰에 대한 view embedding을 채널에 붙여 DiT가 카메라 정체성을 인식하도록 한다.

Cross-view attention은 각 DiT 블록에 삽입되어 한 뷰의 토큰이 다른 뷰의 토큰을 참조하게 함으로써 뷰 간의 시공간 일관성을 확보한다. 손실 함수는 EDM denoising score matching 기반의 L2 재구성 손실로, 잠재 z와 잡음 추가된 입력 z+n을 비교하는 형태로 정의된다. A2World-sim에서는 pose-guided history sampling으로 이력 프레임을 선택하고, 선택된 이력은 Tok_hist로 토크나이즈되어 cross-attn 조건으로 사용되며 동시에 self-attention의 key/value 메모리로도 연결된다.

A2World-policy의 joint diffusion은 비디오·액션 잠재에 별도 노이즈 스케일을 적용하고 shared-timestep 샘플링을 통해 정렬을 보장한다. 모델 블록은 공유 self-attention과 모달리티별 AdaLN·MLP 분기 구조를 가지며 최종 액션은 마지막 레이어의 액션 토큰에 선형 헤드를 적용해 예측한다. 추론 단계에서는 modality-wise classifier-free guidance를 사용하여 비디오 품질과 액션 정확성 사이의 가이던스 스칼라를 별도로 설정할 수 있다.

논문에서 명시한 한계는 크게 세 가지이다. 첫째, A2World의 장기 롤아웃 안정성은 self-forcing으로 개선되었으나 매우 긴 호라이즌이나 극단적 행동 노이즈 상황에서는 여전히 드리프트 위험이 존재한다. 둘째, 멀티뷰와 다양한 로봇 형태를 지원하기 위해 행동을 공통의 듀얼암 포맷으로 통일하였으나 이 통일화가 일부 로봇 고유의 제어 세부를 희생할 수 있다. 셋째, 대규모 사전학습은 H200 GPU 64대와 같은 고성능 인프라를 요구하므로 리소스 제약이 있는 기관에서 재현하기 어렵다.