WorldDiT: A Unified Diffusion Architecture for World and Action Modeling

로봇이 단순히 고정된 명령 목록을 따르는 것이 아니라, 다음에 일어날 일에 대해 실제로 '꿈을 꾸는' 세상을 상상해 보십시오. 이것은 기계에게 우리의 복잡하고 실제적인 세상을 헤쳐 나가는 법을 가르치려는 분야인 로봇 학습의 최전선입니다. 오랫동안 가장 큰 돌파구들은 특정한 레시피에서 나왔습니다. 즉, 언어와 이미지에 대해 이미 모든 것을 알고 있는 거대한 사전 학습된 '두뇌'(이미지를 볼 줄 아는 매우 똑똑한 백과사전 같은 것)를 가져와서 로봇 팔을 부착하는 방식입니다. 이 방식은 작동하긴 하지만, 마치 강아지에게 공 가져오기를 가르치기 위해 먼저 그 강아지를 교수님으로 훈련시키는 것과 같습니다. 이는 무겁고 비용이 많이 들며, 로봇이 과업을 배우고 있는 것인지 아니면 단순히 교수의 노트를 베끼고 있는 것인지 구분하기 어렵게 만듭니다. 과학자들이 던지는 핵심 질문은 이것입니다. "우리는 거대한 사전 로드된 백과사전 없이도, 더 작고 가벼우며 움직임과 미래 예측을 동시에 배울 수 있는 로봇 두뇌를 만들 수 있을까?"

여기 WorldDiT가 있습니다. Bagel Labs의 연구진이 설계한 새로운 로봇 두뇌입니다. WorldDT를 텍스트북을 베끼는 학생이 아니라, 자신의 미래가 담긴 영화를 보며 그림 그리는 법을 배우는 젊은 예술가라고 생각하십시오. WorldDT는 이미지를 설명하고 텍스트를 이해하는 거대한 '시각-언어 모델(Vision-Language Model, 이미지를 설명할 줄 아는 거대한 AI)'에 의존하는 대신, 두 가지 일을 동시에 수행하는 단일한 통합 엔진을 사용합니다. 첫째, 로봇 팔의 다음 움직임을 파악합니다. 둘째, 카메라가 몇 초 후에 무엇을 보게 될지 예측합니다. 이는 마치 체스 플레이어가 다음 수를 계획할 뿐만 아니라, 자신의 계획이 타당한지 확인하기 위해 세 수 앞의 보드 상황을 시각화하는 것과 같습니다.

연구진은 로봇이 블록을 쌓거나 물체를 옮기는 등의 과업을 수행해야 하는 LIBERO라는 시뮬레이션 세계에서 이 시스템을 훈련시켰습니다. 그들은 WorldDT가 (다른 최고 수준의 방법들이 사용하는 수십억 개의 파라미터와 비교했을 때) 놀라울 정도로 작음에도 불구하고(4억 개 미만의 파라미터), 미래의 카메라 뷰와 로봇의 행동을 믿기 힘들 정도로 정확하게 예측할 수 있다는 것을 발견했습니다. 실제로 테스트했을 때, 이 로봇은 네 가지 유형의 도전 과제 전반에서 약 95%의 성공률을 보였습니다. 가장 흥兴奋되는 부분은 무엇일까요? 이 로봇은 대부분의 성공적인 로봇들이 사용하는 무거운 '백본(backbone)' 없이도 이 일을 해냈다는 점입니다. 이는 훌륭한 로봇이 되기 위해 거대한 사전 로드된 두뇌가 필요한 것이 아니라, 움직이면서 미래를 보는 법을 배우는 두뇌가 필요하다는 것을 증명했습니다.

"꿈꾸는" 로봇

그렇다면 WorldDT는 실제로 어떻게 작동할까요? 여러분이 저글링을 배우려고 한다고 상상해 보십시오. 오늘날 대부분의 로봇은 저글링 명수들의 수천 개 영상을 암기하고 그것을 똑같이 흉내 내려고 노력하는 학생들과 같습니다. WorldDT는 다릅니다. 그것은 눈을 가린 채 수정구슬을 가진 저글링 학생과 같습니다.

비결은 이렇습니다. WorldDT가 움직임을 계획할 때마다, 그것은 카메라가 미래에 무엇을 보게 될지도 함께 "꿈을 꿉니다". 로봇의 현재 상태, 주어진 지시 사항(예: "빨간 블록을 집어라"), 그리고 지난 몇 초간의 영상을 살펴봅니다. 그런 다음 동시에 두 가지 질문을 스스로에게 던집니다:

1. "로봇이 다음에 무엇을 해야 하는가?"
2. "만약 내가 그렇게 한다면, 잠시 후의 모습은 어떠할 것인가?"

연구진은 이를 **디퓨전 트랜스포머(diffusion transformer)**라고 부릅니다. 만약 이것이 공상과학 용어처럼 들린다면, 이를 "노이즈 제거(denoising)" 과정이라고 생각하십시오. 사진이 정전기 노이즈로 뒤덮여 있다고 상상해 보세요. WorldDT는 로봇의 다음 움직임과 미래의 모습에 대한 완전히 무작위적이고 노이즈가 섞인 추측에서 시작합니다. 그런 다음 단계별로 그 노이즈를 걷어내며, 목표를 향한 명확하고 매끄러운 경로를 찾을 때까지 추측을 정교하게 다듬어 나갑니다. 이 과정은 로봇의 행동과 미래의 이미지 모두에 대해 수행됩니다.

마법은 이 두 가지 작업이 서로를 돕는 데서 일어납니다. 미래의 이미지를 예측하려고 노력함으로써, 로봇은 세상의 물리 법칙을 이해하도록 강요받습니다. 만약 로봇이 블록을 밀려고 계획했는데, 그 "꿈" 속의 미래에서 블록이 공중에 떠 있다면, 로봇은 자신의 계획이 잘못되었다는 것을 알게 됩니다. 로봇은 "꿈"이 현실과 일치할 때까지 행동을 조정해야 합니다. 이는 로봇이 다음에 올 것을 예측함으로써 세상의 규칙을 배우는 피드백 루프를 생성합니다.

결과: 작은 두뇌, 큰 성공

연구진은 네 가지 다른 시뮬레이션 스위트(공간, 물체, 목표, 장기 과제로 분류되는 네 가지 다른 장애물 코스라고 생각하십시오)에서 WorldDT를 테스트했습니다. 그들은 이를 24개의 다른 유명한 로봇 학습 방법들과 비교했습니다.

핵심은 이렇습니다. WorldDT는 매우 작습니다. 약 3억 9,900만 개의 파라미터(두뇌의 '뉴런')를 가지고 있습니다. 다른 많은 상위 성능의 로봇들은 10억 개 또는 심지어 100억 개의 파라미터를 사용하는 모델을 사용합니다. 보통 더 큰 두뇌가 더 나은 성능을 의미하지만, WorldDT는 이 규칙을 깼습니다.

이 시뮬레이션에서 WorldDT는 네 가지 코스 전체에서 평균 94.9%의 성공률을 달랐습니다.

• Spatial (공간): 98.0% 성공
• Object (물체): 97.0% 성공
• Goal (목표): 92.8% 성공
• Long (장기): 91.8% 성공

이 결과를 그래프로 그리면, WorldDT는 "파레토 프런티어(Pareto frontier)"에 위치합니다. 쉬운 말로, 이것은 이 로봇이 그 방에서 가장 효율적인 로봇이라는 뜻입니다. 만약 여러분이 WorldDT만큼 똑똑하면서도 더 작은 두뇌를 사용하는 로봇을 원한다면, 찾을 수 없을 것입니다. 반대로 WorldDT만큼 작은 두면서 더 똑똑한 로봇을 원해도 찾을 수 없습니다. 94.9%보다 더 높은 성과를 낸 다른 모든 로봇은 훨씬 더 크고 무거웠습니다.

이것이 중요한 이유

이 논문은 좋은 로봇 제어를 위해 반드시 거대한 사전 학습된 시각-언어 모델(VLA)을 사용해야 한다는 생각에 명시적으로 반박합니다. 한동안 이 분야는 "세상을 이해하려면 거대한 두뇌가 필요하다"고 생각했습니다. WorldDT는 "꼭 그렇지는 않다"고 말합니다.

연구진은 행동 생성(움직임)과 월드 모델링(미래 예측)을 결합함으로써, 단일한 통합 아키텍처가 로봇을 매우 효과적으로 제어하는 법을 배울 수 있음을 보여주었습니다. 그들은 단순히 제안만 한 것이 아니라, 수천 번의 시뮬레이션 에피소드를 통해 이를 측정했습니다. 로봇은 단순히 "운이 좋았던" 것이 아니라, 훨씬 더 큰 모델들을 지속적으로 능가했습니다.

하지만 주의할 점이 있습니다. 이러한 결과는 시뮬레이션(로봇이 사는 컴퓨터 게임)에서 나온 것입니다. 논문은 성능은 강력하지만, 이것이 향후 연구를 위한 기준점(baseline)임을 언급합니다. 이 "꿈꾸는" 로봇이 실제 중력과 실제 미끄러운 바닥이 있는 실제의 복잡한 주방에서도 완벽하게 작동할지는 아직 알 수 없습니다. 그러나 시뮬레이션 결과는 우리가 미래를 단순히 암기하는 것이 아니라, 미래를 상상함으로써 배우는 더 작고, 저렴하며, 효율적인 로봇을 미래에 구축할 수 있다는 강력한 힌트를 줍니다.

결국, WorldDT는 똑똑한 로봇의 비결이 단순히 방대한 지식의 도서관을 갖는 것이 아니라, 지금 이 순간 더 나은 결정을 내리기 위해 마음속으로 미래를 그려낼 수 있는 두뇌를 갖는 것임을 시사합니다.