OneWM, 통합 인과 학습으로 세계 모델과 행동 학습을 연결하다

들어가며

로봇 지능에서는 환경의 변화를 예측하는 세계 모델과 다음 행동을 결정하는 정책 모델을 별도로 개발하는 경우가 많습니다. OneWM은 두 기능을 하나의 비디오-행동 아키텍처에서 훈련하면 서로에게 도움이 되는지 검증합니다. 핵심은 행동 예측과 미래 영상 예측을 통합된 인과 학습 문제로 다루는 것입니다.

OneWM의 구성

OneWM에서 행동은 두 가지 역할을 합니다. 먼저 특정 행동을 조건으로 넣어, 해당 행동 이후 나타날 수 있는 시각적 결과를 예측합니다. 동시에 행동 자체는 정책이 학습하고 생성해야 하는 목표가 됩니다. 따라서 모델은 “이 행동을 하면 무엇이 달라지는가”와 “현재 관측에서 어떤 행동을 선택해야 하는가”를 함께 학습합니다.

연구는 세계의 전방 동역학을 학습하는 것이 행동 학습을 개선하는지, 반대로 행동 생성 학습이 세계 모델을 더 유용하게 만드는지를 살펴봅니다. 실험에는 LIBERO, LIBERO-Plus, RoboCasa365가 사용됐으며, 정책 학습과 세계 모델링을 모두 평가했습니다. 여기에 실제 로봇을 이용한 정책 실험도 추가됐습니다.

핵심 결과

• 행동 결과 예측이 정책을 개선: 지정된 행동이 만들어내는 시각적 변화를 학습하면 장면이 달라지는 상황에서 정책 성능과 일반화가 좋아졌습니다.
• 행동 학습이 세계 예측을 개선: 행동을 생성하는 능력을 학습한 모델은 동일한 구조에서 여러 시점의 미래 장면을 더 충실하게 예측했습니다.
• 양방향 이득 확인: 세계 모델은 정책이 행동의 결과를 고려하도록 돕고, 정책 학습은 제어에 중요한 환경 변화를 모델이 보존하도록 유도합니다.
• 실제 로봇으로 확장: 색상 블록 쌓기, 물체 수납, 간단한 아침 식사 준비와 같은 시연이 포함됐습니다.

의미와 한계

시각적 패턴만 모방하는 정책은 배경이나 물체 배치가 바뀌면 성능이 떨어질 수 있습니다. 행동 이후의 장면을 예측하게 하면 관측, 개입, 미래 상태 사이의 관계를 더 인과적으로 학습할 여지가 생깁니다. 반대로 행동 생성 과제는 세계 모델이 단순히 그럴듯한 영상을 만드는 데 그치지 않고, 실제 제어와 관련된 변화를 표현하도록 압박할 수 있습니다.

다만 제공된 자료에는 구체적인 성능 향상 수치가 없습니다. 긴 시간의 롤아웃, 서로 다른 로봇 형태, 복잡한 접촉 역학에서의 안정성도 확인되지 않았습니다. 프로젝트의 ‘Endless Dream’은 초기 관측에서 시작해 모델이 자체적으로 생성한 전개입니다. 예측된 제어 신호는 시뮬레이터에서 실행된 궤적이 아니므로 실제 제어 성능과 혼동해서는 안 됩니다. 따라서 이번 결과는 통합 인과 학습의 가능성을 뒷받침하지만, 장기 계획 문제가 해결됐다는 의미는 아닙니다.

OneWM의 중요한 점은 세계 모델과 정책을 분리된 모듈이 아니라 서로 강화할 수 있는 학습 목표로 제시했다는 데 있습니다. 앞으로는 불확실성 추정, 장기 예측, 실제 환경과의 지속적인 상호작용을 통해 이 상호 이득의 범위를 더 검증해야 합니다.

출처: Hugging Face Daily Papers