LeWAM이 재구성 없는 월드액션모델로 계획 실패를 해결한다

무엇인가

로봇 제어는 점점 예측 모델에 의존한다. 기존 월드모델은 이미지 재구성으로 잠재공간을 만들고 그 위에 정책을 얹었고, 최근의 월드액션모델(WAM)은 미래 상태 예측과 행동 생성을 묶어 제어를 학습된 동역학에 직접 접지시킨다. 문제는 이 계열이 대부분 픽셀 수준 재구성을 표현 학습의 기준으로 삼는다는 점이다. 그러면 잠재공간이 디코더가 렌더링할 수 있는 모든 시각적 디테일을 담아야 하고, 과제와 무관한 잡음에 모델 용량을 낭비해 하위 예측 작업이 어려워진다. JEPA 계열은 디코더를 버리고 잠재공간에서만 동역학을 모델링해 이 문제를 피하지만, 동역학 모델과 정책이 분리되어 있어 최신 WAM의 통합 구조가 없다. 이 논문은 그 간극을 메우는 LeWAM을 제안한다.

어떻게 동작하나

LeWAM은 디코더 없는 JEPA 잠재 위에 네 가지 모드를 동시에 수행하는 양방향 트랜스포머다. 각 프레임은 224픽셀에서 ViT-tiny/14로 독립 인코딩되는데, DINOv2 small 아키텍처를 사전학습 가중치 없이 처음부터 학습한다. 인코더의 [CLS] 토큰이 BatchNorm MLP 프로젝터를 지나 384차원 잠재 벡터 z_t 하나가 되고, 이 벡터가 프로브가 읽는 표현이자 플래너가 비교하는 대상이자 동역학 모델이 회귀하는 대상이다. 표현 붕괴는 SIGReg로 막는데, 배치의 잠재를 등방 가우시안 쪽으로 밀어 넣는 정규화다. EMA 타깃 네트워크나 stop-gradient는 쓰지 않는다. 백본은 깊이 8, 폭 512, 헤드 16개의 양방향 트랜스포머로, 세 개의 프레임 잠재와 세 개의 액션 토큰을 모달리티·프레임·σ 임베딩과 함께 읽는다. 학습 창은 프레임 4장(프레임스킵 5)과 그 사이 세 개의 5스텝 액션 청크로 구성된다.

무엇과 다른가

네 모드는 입력을 마스킹해 구분한다. Forward는 마지막 프레임 잠재를 마스크 토큰으로 바꾸고 액션은 깨끗하게 두어 z_{t+1}을 살아 있는 인코더 출력에 대해 MSE로 회귀한다. Backward는 첫 프레임 잠재를 마스킹해 이후 프레임과 모든 액션으로부터 복원한다. Inverse는 모든 프레임을 깨끗하게 두고 타깃 청크 a_t를 노이즈 낀 flow-matching 샘플로 교체해 백본이 속도를 예측하게 하며, 과거와 미래 상태를 잇는 행동을 추론한다. Policy는 마지막 프레임을 마스킹하고 a_t를 노이즈 샘플로 바꿔 두 손실을 모두 적용하되, 과거 액션 토큰을 학습된 unknown-action 토큰으로 대체해 정책이 시연자의 최근 동작을 베끼지 못하게 한다. 매 배치마다 네 모드를 전부 실행하고, 인코더는 한 번만 돌리고 백본은 모드별로 독립 실행한 뒤 손실을 합산한다. Forward·Backward와 SIGReg가 JEPA 구조를 만들고, Inverse·Policy가 그 공간을 과제에 정렬시킨다. 행동 생성은 rectified flow를 쓴다. 학습 시 σ~U(0,1)을 뽑아 a_σ=(1-σ)a+σε를 만들고 속도 타깃 ε-a를 예측하며, 테스트 시에는 ε~N(0,I)에서 8번의 오일러 스텝으로 적분해 결정적 매핑 a=g(ε; z_t)를 얻는다. 픽셀 디코더도, 토크나이저도, 픽셀·토큰 공간 손실도 없다.

어떻게 쓰나

계획은 잠재 거리 기반 목표 조건부 방식이다. 시연에서 가져온 목표 잠재 z_g에 대해 청크 a의 비용을 전방 모드 예측과의 L2 거리로 정의하고, 샘플링 기반 플래너가 H개 청크에 대해 이를 최소화한다. 문제는 원래의 샘플링 MPC가 raw 액션 공간 a_{1:H}∈R^{AH}를 탐색한다는 점이다. 저자들은 데이터가 적은 인간 시연 데이터셋에서 이 방식이 실패한다고 보고한다. 동역학 모델이 불완전하면 가우시안 액션 샘플링 MPC가 환각 미래를 최적화하고, 그 계획이 배포에서 무너진다는 것이다. 그래서 결정 변수를 flow-matching 정책의 노이즈 공간으로 옮긴다. 청크 h는 이전 청크들이 상상한 상태에서 디코딩되고(a_h=g(ε_h; ẑ_{h-1})), 그 상태에서 다음 잠재를 예측해 컨텍스트를 shift한다. 추가로, 정책 헤드는 N(0,I)에서 출발하도록 학습되어 노름 √A 구에 집중하는데 탐색이 좁아지면 제안이 작은 노름 쪽으로 표류해 저온 샘플링처럼 모드 추구적으로 변한다. 이를 막기 위해 제안을 구면에 재투영한다(ε←√A ε/‖ε‖). 결과적으로 탐색은 방향에 대해 이루어지고, 디코딩되는 모든 후보는 헤드가 학습한 분포 안에 있다. DS-CEM, DS-MPPI, DS-CMA-ES는 업데이트 규칙만 다르고 예산 64×3, 호라이즌 5를 공유한다.

전제와 한계

실험은 PushT와 robomimic이다. PushT는 Chi 등의 인간 시연 206개에, 행동과 무관한 두 개의 distractor(리사주 곡선 위의 점, 상단 가장자리를 따라 미끄러지는 막대)를 그려 넣었다. robomimic은 Lift, Can, Square, ToolHang, Transport의 proficient-human 시연 200개를 robosuite에서 재생하고, 과제 물체보다 큰 회색 distractor 상자 두 개를 모든 장면에 흔들리게 넣었다. 모든 모델은 픽셀만 입력받고 proprioception이나 상태는 쓰지 않으며, 300 에폭, 5 시드, 90% 분할로 학습하고 나머지 10%에서 평가한다. 베이스라인은 Dreamer(이미지 재구성 RSSM + MLP 액션 헤드), LeWM(동결한 뒤 그 잠재 위에 LeWAM의 rectified-flow 헤드를 같은 손실·스케줄로 학습), RecWAM(Cosmos Policy 아키텍처를 LeWAM 규모로 옮긴 것), 그리고 정책 손실만으로 학습한 flow-matching policy다. 모두 약 41.0M 학습 파라미터로 맞췄다. 잠재공간 분석에서는 잠재에서 에이전트·물체·그리퍼·distractor 상태로 ridge 회귀를 적합해 홀드아웃에서 MSE, Pearson r, R²를 봤는데, LeWAM이 평균적으로 과제 관련 상태를 LeWM보다, 심지어 재구성 기반 모델보다 잘 인코딩하면서 distractor는 LeWM만큼 잘 무시했다. 폐루프 평가는 6개 과제에서 250 롤아웃(5 학습 × 50 평가 시드)으로 진행했고, LeWAM의 정책은 참조 BC flow-matching 정책과 동등한 성능을 유지했으며 RecWAM의 정책은 참조보다 약간 낮았다. Dreamer와 LeWM+diffusion 헤드는 매우 나빴다. 노이즈 공간 샘플링은 LeWAM과 RecWAM 모두 추론 시 성능을 끌어올렸지만, 액션 공간 MPC는 모든 모델을 3~8% 성공률로 붕괴시켰다.

절제 실험은 Lift, Can, Square 세 과제에서 진행됐다. 각 모드를 하나씩 꺼서 학습하면 잠재공간의 정보성과 폐루프 성능 모두 평균적으로 떨어졌다. Forward와 Backward는 일관되게 도움이 되고, Inverse는 Square에서 중립이었다. 4.4절의 구면 투영을 제거하면 DS-CEM 성공률이 Lift에서 약 5점, Square에서 12점, Can에서 56점 하락했다.

개발자 관점에서 이 논문이 주는 실무적 신호는 세 가지다. 첫째, 재구성 손실 없이 픽셀에서 끝까지 학습한 잠재가 재구성 기반 표현보다 과제 관련 상태를 더 잘 담을 수 있다는 것, 즉 월드모델을 만들 때 디코더를 굳이 붙일 필요가 없다는 근거다. 둘째, 정책과 월드모델을 한 백본에서 마스킹만으로 네 모드를 돌려 학습하는 구성은 구현 비용이 낮고, 정책 성능을 희생하지 않는다는 것을 같은 인코더·같은 크기의 BC 정책과의 비교로 보여준다. 셋째, 샘플링 기반 MPC를 쓸 때 액션 공간이 아니라 정책 샘플러의 노이즈 공간을 탐색하고 제안을 구면에 투영하는 것이 사실상 필수라는 점이다. 액션 공간 MPC가 3~8%로 붕괴하는 현상은 학습된 동역학이 완벽하지 않은 한 일반적으로 나타날 수 있으므로, 월드모델 기반 계획을 붙일 때는 어떤 공간에서 최적화하는지부터 확인해야 한다.

저자들이 밝힌 한계는 명확하다. 계획 알고리즘이 느리고, 추론 시점에 목표(goal) 개념이 필요하다는 점이다. 또한 실험은 시뮬레이터의 인간 시연 데이터에 국한되며, 실제 로봇 데이터는 동적으로 무관한 특징이 훨씬 많기 때문에 적용해 보는 것이 다음 단계라고 저자들은 적었다. 계획 알고리즘 자체를 개선하는 것도 향후 과제로 남겨 두었다.