행동을 구분하는 세계모델 'AD-WM', 로봇 계획 성능 대폭 향상

로봇 계획(planning)에 쓰이는 대부분의 세계모델은 현재 상태와 행동이 주어졌을 때 다음에 일어날 일을 예측하고, 실제 결과와의 오차를 줄이는 방식으로 학습됩니다. 그러나 모델예측제어(MPC, 앞으로 취할 여러 행동의 결과를 미리 시뮬레이션해 최적 행동을 고르는 제어 기법)에는 단순히 정확한 예측만으로는 부족합니다. 같은 상태에서 출발하는 여러 대안 행동들을 서로 비교해 그중 최선을 골라내야 하기 때문입니다. 최근 arXiv에 게재된 한 논문은 이 두 가지 목표가 실제로는 다르며, 예측 정확도는 높지만 행동 간 차이를 잘 구분하지 못하는 모델이 존재할 수 있다고 지적합니다.

Jiabin Qiu 등 연구진이 제출한 이 논문은 이러한 문제를 정면으로 다루는 'AD-WM(Action-Discriminative World Model, 행동 구별형 세계모델)'을 제안합니다. AD-WM은 잔차 잠재 동역학(residual latent dynamics)에 '예측기 단계의 행동 복원 정규화(predictor-level action-recovery regularization)'라는 학습 기법을 결합합니다. 이는 역동역학(inverse dynamics)과, 조건부 상호정보량(conditional mutual information) 개념에서 착안한 정규화된 복원 목표를 활용해, 내부 예측 표현만 보고도 원래 어떤 행동이었는지 복원할 수 있을 만큼 충분한 정보를 유지하도록 모델을 유도하는 방식입니다. 결과적으로 서로 다른 행동은 상상된 결과에서도 의미 있게 다르게 나타나게 됩니다. 이 보조 학습 목표들은 실제 테스트 시에는 제거되므로, 계획 알고리즘(플래너) 자체는 그대로 유지되고 오직 내부 세계모델만 달라집니다.

시뮬레이션 벤치마크인 OGBench-Cube에서 AD-WM은 어려운 시작 조건(hard-start)에서의 성공률을 기존 LeWM 기반 비교 모델의 3.7%에서 52.0%로 끌어올렸습니다. 다섯 개 시뮬레이션 환경 중 네 곳에서 평균 성공률도 향상되었습니다. 연구진은 계획 성능을 진단하는 실험도 진행했는데, 흔히 쓰이는 지표인 '단순 사실 예측 오차'와 '전체 행동 후보군에 대한 순위 정확도'는 실제 폐루프(closed-loop) 과제 성공률과 잘 맞아떨어지지 않았습니다. 반면 상위 후보 행동들만을 대상으로 오차를 측정하는 'CEM(교차엔트로피법) 정렬 엘리트 후회(CEM-aligned elite regret)' 지표는 실제 성공률과 훨씬 밀접하게 일치했습니다.

실제 환경 전이 성능을 확인하기 위해 연구진은 고정된 V-JEPA 2 시각 인코더와 DROID 로봇 조작 데이터셋으로 후속 학습을 진행한 뒤, 별도의 실험실 맞춤 조정 없이 프랑카(Franka) 로봇팔에서 제로샷(zero-shot) 평가를 실시했습니다. 그 결과 기본적인 물건 집기·놓기(pick-and-place) 성공률이 기존 방식의 42.2%에서 AD-WM 적용 시 71.1%로 상승했습니다.

연구진은 계획을 위한 세계모델은 단순 예측 정확도만을 좇을 것이 아니라, 의사결정에 필요한 '행동에 따른 차이'를 명시적으로 보존하도록 학습해야 한다고 결론짓습니다. 이는 학습된 잠재 세계모델 위에 MPC 기반 로봇 제어기를 구축하려는 연구자와 개발자들에게 중요한 시사점이 될 수 있습니다. 추가 영상과 코드는 논문 프로젝트 페이지에서 확인할 수 있습니다.