UniWALM이 언어를 작업기억으로 재사용해 로봇 조작을 개선한다

무엇인가

이 논문은 로봇이 언어·영상·행동을 하나의 표현으로 통합해 조작을 수행하는 문제를 다룬다. 기존 VLA는 언어를 고정 조건으로만 쓰고, CoT형 VLA는 언어를 생성해도 다음 단계에서 재사용하지 않으며, 월드-액션 모델은 언어를 아예 빼거나 외부 플래너·검색 버퍼를 붙인다. 저자들은 언어를 행동과 같은 패스에서 생성하고 다음 제어 단계의 프롬프트로 다시 소비하면, 공유 표현이 풍부해지고 기억 의존 작업도 외부 모듈 없이 풀 수 있는지 묻는다.

어떻게 동작하나

UniWALM은 비디오·행동·언어를 각각 전용 DiT 헤드로 두되 하나의 Mixture-of-Transformers 백본 안에서 동시에 디노이징한다. 비디오 백본은 Wan2.2 사전학습 모델을 쓰고, 행동 DiT는 1B(히든 1024), 언어 DiT는 0.5B(히든 512)로 각각 30층이며 처음부터 학습한다. 각 모달리티 토큰은 통합 self-attention으로 특징을 공유하고, task instruction과 proprioception에는 cross-attention한다. 언어 헤드는 이산 확산으로 subtask와 textual memory를 같은 패스에서 생성하고, 생성 문자열은 다음 제어 단계의 프롬프트로 재입력된다. 선택적 시각 기억은 하나 이상의 앵커 프레임을 받아 가림·이동으로 잃는 장면 상태를 보존하고, 행동 벡터에 붙인 1차원 task-completion 신호 β가 subtask 완료 시점에만 텍스트 기억 갱신을 게이트한다.

무엇과 다른가

학습 목표는 비디오 흐름 매칭 손실, 행동 흐름 매칭 손실, 언어 이산 확산의 cross-entropy를 합친 것이다. 비디오와 행동에는 L1 흐름 매칭을 써서 L2보다 안정적이었다고 한다. 여기에 Flow Matching with Correction(FMC)을 추가한다. 학습 때 깨끗한 행동 x0에 잡음 ε1을 더해 궤적 밖 점을 만들고, 표준 속도 (ε−x0)에 (1−t)/Δt·ε1 보정항을 더한 목표 속도로 원래 궤적으로 돌아오게 학습한다. 보정 분기는 t≥0.5에서만 적용해 보정 크기를 제한하며, 추론 시 추가 조건·샘플링 단계·모델 용량이 들지 않는다. 비디오 백본은 LoRA rank 256으로 미세조정해 전체 미세조정보다 학습 파라미터를 줄이면서 성능을 개선했다. 행동 청크 32, 비디오 호라이즌 5, 221K 궤적 혼합 데이터에서 15k 스텝, 12개 H200 노드로 44시간 학습했다. 추론에서 비디오 생성은 끄고 언어 DiT는 행동 DiT와 동기적으로 돈다.

어떻게 쓰나

MolmoSpaces 시뮬레이션에서는 Pick, Pick-RandCam, Pick&Place, Pick&Place-NextTo, Pick&Place-Color의 221K 궤적을 학습하고 과제별 1000개 OOD 에피소드로 평가했다. UniWALM은 세 가지 Pick&Place 변형 모두에서 가장 높은 성공률을 냈고, 같은 Wan2.2 백본과 행동 전문가 크기를 쓰는 FastWAM보다 평균 14.8%포인트 높았다. 비디오 히스토리 없이도 MolmoBot을 Pick&Place 변형에서 평균 7%포인트 앞섰다. 다만 303스텝 예산이 걸린 짧은 Pick·Pick-RandCam에서는 더 신중한 궤적 때문에 밀렸고, 조합성이 큰 과제에서 이득이 커졌다.

전제와 한계

실제 양팔 조작은 Flexiv Rizon 두 팔과 dexterous hands로 Cube Pick&Place, Empty Container, Find the Item을 평가했다. Cube Pick&Place는 95.0%로 다음 최고 베이스라인 80.0%를 앞섰고, Empty Container는 FastWAM과 50.0%로 동률이었다. 기억 의존 Find the Item에서는 물체가 중간에 가려지는데, UniWALM이 생성한 텍스트 기억을 다음 단계에 재소비해 성공률이 FastWAM의 11.1%에서 61.1%로 올랐고, π0.5·GROOT-N1.7·MolmoBot 등 VLA/WAM 베이스라인은 33.3% 이하로 우연 수준에 머물렀다. 이 실험에서는 공정성을 위해 앵커 프레임과 β를 쓰지 않아 기억 신호가 전적으로 자체 생성 언어에서 나왔다.

RMBench의 기억 의존 M(n) 하위 집합(Press Button, Blocks Ranking Try, Cover Blocks, Battery Try)에서는 과제당 50개 데모와 100개 롤아웃으로 평가했다. UniWALM은 사전학습 언어·추론 모델 없이 평균 34.5% 성공률을 기록해 같은 백본 FastWAM의 13.3%를 21.2%포인트 앞섰고, Mem-0의 키프레임 지원 설정 28.5%도 단일 앵커 프레임만으로 넘었다. Mem-0의 키 메모리를 제거한 버전은 4.8%로 떨어졌다. 동시대 WLA-0은 더 높은 점수를 냈지만 사전학습 VLM과 이미지 생성 모델을 메타쿼리로 연결한 파이프라인이라, UniWALM처럼 하나의 공유 표현을 학습한 결과는 아니다.

Pick&Place 단일 분할(38K 궤적) 절제 실험에서 언어 생성은 OOD 성공률을 4.9%포인트(상대 15.8%) 올렸고, FMC는 3.0%포인트(상대 9.7%)를 더했다. 기준 WAM 31.0%에서 UniWALM은 41.0%로 10%포인트(상대 32%) 올랐으며, 전체 벤치마크 평균으로는 FastWAM 대비 약 14.8%포인트(상대 37%) 개선이다. 이는 0.5B 언어 전문가를 추가한 효과로 정리된다. Press Button 기억 절제에서는 앵커 프레임만으로는 1%, 언어 생성만으로는 15%, 둘을 함께 쓰면 51%가 나와 언어가 기억 신호의 대부분을 나르고 시각 앵커가 목표 안정화를 돕는다고 해석한다.

개발자 관점에서 이 논문은 외부 플래너나 검색 버퍼 없이 긴 호라이즌·기억 의존 조작을 단일 정책에 넣는 설계를 보여준다. 다만 언어 생성이 추가되어 최적화 전 기준 32개 행동 청크 생성이 76ms에서 123ms로 늘고, 학습에는 12개 H200 노드로 44시간이 들며, 언어·행동 전문가는 사전학습 언어 모델 없이 처음부터 학습한다. 저자들은 복잡한 형상 파지와 혼잡한 장면에서 실패하고, RMBench에서는 subtask와 텍스트 기억은 맞게 추적해도 task-completion 타이밍 지연으로 행동 전문가가 경계를 넘거나 못 미치는 경우, 그리고 거의 같은 관측에 다른 기억이 붙은 경우 모델이 혼동한다고 밝힌다. 실제 적용 전에는 언어 출력 품질, β 갱신 타이밍, 시각 기억 필요성, 추론 지연을 함께 검증해야 한다.