🤖 AI

Q-Learning With World Models

본 논문은 월드 모델(world models)을 표준 Q-러닝과 결합하여 행동 선택을 위한 상상된 궤적(imagined trajectories)에 대한 테스트 타임 탐색을 수행함으로써, 실제 전이(real transitions)로만 학습하여 모델 편향의 누적을 방지하는 동시에 로봇 조작 벤치마크에서 우수한 샘플 효율성과 성능을 달성하는 프레임워크인 QWM을 제안한다.

원저자: Perry Dong, Yueru Jia, Chelsea Finn, Dorsa Sadigh

게시일 2026-08-19

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Perry Dong, Yueru Jia, Chelsea Finn, Dorsa Sadigh

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

직접 해보며 배우는 로봇들이 현실이 되고 있지만, 이들은 아주 간단한 기술을 익히는 데도 종종 엄청난 양의 연습을 필요로 합니다. 테이블에서 선반으로 컵을 옮기기 위해, 로봇은 마침내 성공하기 전까지 수천 번의 시도를 반복하며 계속해서 실패해야 할 수도 있습니다. 이러한 시행착오 과정은 강화 학습(reinforcement learning)이라 불리는 분야의 핵심으로, 여기서 에이전트는 환경과 상호작용하며 자신의 행동에 대한 피드백을 받으며 학습합니다. 최근의 발전 덕로 이 학습 과정이 빨라지기도 했지만, 도구를 조립하거나 복잡한 방을 통과하는 것과 같은 가장 복잡한 작업들은 여전히 방대한 양의 데이터를 요구합니다. 연구자들은 이러한 로봇들을 더 효율적으로 만들기 위해 끊임없이 방법을 찾고 있으며, 이를 통해 훈련에 필요한 시간과 에너지를 줄이기를 희망하고 있습니다. 유망한 아이디어 중 하나는 로봇에게 '세계 모델(world model)', 즉 특정 행동을 취했을 때 어떤 일이 일어날지 단순히 추측하는 대신 상상할 수 있게 하는 정신적 시뮬레이션을 부여하는 것이었습니다. 하지만 이러한 정신적 시뮬레이션을 사용하여 로봇을 직접 훈련시키는 것은 역사적으로 위험했습니다. 만약 로봇의 상상이 조금이라도 틀린다면, 그 오류들이 쌓여 로봇이 실제 세계에서는 실패하는 왜곡된 현실을 학습하게 될 수 있기 때문입니다.

스탠퍼드 대학교와 베이징 대학교의 연구진은 로봇이 상상을 사용하는 시점을 바꿈으로써 이 위험을 피하는 새로운 접근 방법을 개발했습니다. 로봇에게 어떻게 움직여야 하는지를 가르치기 위해 세계 모델에 의존하는 대신, 그들은 오직 결정의 순간에만 이를 사용합니다. QWM이라고 불리는 이 방법은 로봇이 실제 세계와의 상호작용으로부터 핵심 기술을 배울 수 있게 하여, 물리 법칙과 인과관계에 대한 이해가 현실에 기반을 두도록 보장합니다. 하지만 로봇이 새로운 상황에 직면하여 다음에 무엇을 할지 결정해야 할 때, 로봇은 잠시 멈추어 빠른 정신적 시뮬레이션을 실행합니다. 로봇은 각 잠재적 행동에 대해 몇 가지 가능한 미래를 상상하며, 세상이 어떻게 변할지 예측한 다음, 가장 좋은 결과를 가져올 행동을 선택합니다. 이 과정은 체스 선수가 수를 두기 전에 몇 수 앞을 내다보는 것과 유사하지만, 여기에서 로봇은 어떤 물리적 행동이 가장 높은 보상을 가져다줄지를 보기 위해 앞을 내다봅니다.

연구진은 물체를 들어 올리거나, 캔을 배치하거나, 도구를 조립하는 것과 같은 도전적인 로봇 조작 작업에서 이 시스템을 테스트했습니다. 그들은 정신적 모델 없이 학습하거나 완전히 시뮬레이션된 세계 안에서 학습하는 다른 최고 수준의 알고리즘들과 이들의 방법을 비교했습니다. 결과는 QWM이 이러한 다른 방법들보다 성능이 훨씬 뛰어남을 보여주었습니다. QWM은 더 빠르게 학습하여 높은 성공 수준에 도달하기 위해 더 적은 실세계 시도를 필요로 했으며, 다양한 작업에 걸쳐 더 안정적으로 수행되었습니다. 결정적으로, 이 시스템은 내부 모델이 너무 결함이 생겨 로봇에게 나쁜 습관을 가르치는 흔한 함정을 피했습니다. 학습 과정을 실제 데이터에 묶어두고 상상을 오직 마지막 순간의 선택을 정교화하는 데만 사용함으로써, 로봇은 실세계 학습의 안전성과 예측 모델의 선견지명이라는 두 마리 토끼를 모두 잡을 수 있었습니다.

이 접근 방식의 성공은 로봇이 어떻게 정신적 탐색을 구축하느냐에 달려 있습니다. 로봇이 행동해야 할 때, 단순히 하나의 무작위 추측을 선택하거나 단일한 예측에 의존하는 것이 아닙니다. 대신, 로봇은 가능성의 트리를 생성합니다. 현재 위치에서 로봇은 몇 가지 다른 행동들을 고려합니다. 각 행동에 대해, 로봇은 자신의 세계 모델을 사용하여 다음 상태의 세계를 예측합니다. 그런 다음 로봇은 그 새로운 상태에서 또 다른 행동을 취했을 때 어떤 일이 일어날지 상상하는 과정을 반복하며, 이렇게 하여 잠재적인 미래의 분기되는 경로를 만들어냅니다. 로봇은 즉각적인 결과만을 보는 것이 아니라, 상상된 전체 시퀀스의 가치를 종합함으로써 이 경로들을 평가합니다. 로봇은 어떤 행동이 가치가 있는지에 대한 직접적인 지식과 미래 단계들로부터 예측된 보상을 결합합니다. 이를 통해 로봇은 즉각적인 순간을 넘어, 지금 당장은 어려워 보일 수 있지만 나중에 훨씬 더 나은 결과를 가져올 행동을 선택할 수 있습니다.

가장 중요한 발견 중 중 하나는, 로봇이 카메라를 통해 세상을 바라보며 단순한 숫자가 아닌 고차원의 시각 정보를 다루는 상황에서도 이 방법이 작동한다는 점입니다. 이러한 픽셀 기반 환경에서 완벽한 세계 모델을 학습하는 것은 매우 어렵지만, 연구진은 이 데이타 기반의 테스트 타임 탐색(test-time search)을 사용할 때 다소 불완전한 모델이라도 성능을 크게 높일 수 있다는 것을 발견했습니다. 로봇은 완벽한 예측자가 될 필요는 없었으며, 단지 좋은 행동과 나쁜 행동을 구별할 수 있을 정도로만 충분하면 되었습니다. 또한 연구는 로봇이 얼마나 멀리 미래를 내다봐야 하는지도 탐구했습니다. 너무 멀리 내다보는 것은 예측이 덜 신뢰할 수 있게 되어 불확실성을 높였고, 너무 얕게 보는 것은 행동의 장기적인 결과를 포착하지 못했습니다. 연구진은 로봇이 예측 오류의 소음 속에서 길을 잃지 않으면서도 효과적으로 계획할 수 있도록, 몇 단계 앞을 내다보는 최적의 지점을 찾아냈습니다.

연구진은 또한 시뮬레이션 내부에서 로봇을 완전히 훈련시키는 것과 같은 다른 방식의 세계 모델 활용법과 비교했습니다. 그들은 그러한 방식들이 특히 로봇이 성공했을 때만 신호가 발생하는 희소한 보상(sparse rewards) 작업에서 어려움을 겪는다는 것을 발견했습니다. 그런 경우 시뮬레이션의 오류가 빠르게 누적되어 로봇이 유용한 것을 배우기 어렵게 만듭니다. 반면, QWM 방식은 훈련을 현실에 기반하게 유지했습니다. 로봇의 핵심 정책과 가치에 대한 이해는 오직 실제 세계의 데이터로만 업데이트되었습니다. 세계 모델은 단지 선택을 위한 도구, 즉 근본적인 학습을 손상시키지 않으면서 의사결정을 날카롭게 만드는 수단으로 사용되었습니다. 이러한 차이는 결정적이었으며, 이전의 모델 기반 방식들이 실패했던 복잡한 실세계 로봇 문제로 시스템을 확장할 수 있게 해주었습니다.

결국, 이 연구는 세계 모델의 힘이 반드시 실제 세계를 대체하는 데 있는 것이 아니라, 실제 세계를 항해하는 능력을 강화하는 데 있음을 보여줍니다. 상상을 경험의 대체물이 아닌 선택의 필터로 사용함으로써, 연구진은 샘플 효율적이면서도 견고한 시스템을 만들어냈습니다. 로봇은 현실의 혹독한 교훈으로부터 배우지만, 일어나기 전에 실수를 피하기 위해 자신의 마음을 사용합니다. 이 접근 방식은 계획하고 적응하는 능력이 학습하는 능력만큼이나 중요한 비정형 환경에 로봇을 배치하는 실질적인 경로를 제시합니다. 이 연구 결과는 로봇 학습의 미래가 완벽한 시뮬레이션을 구축하는 것이 아니라, 언제 자신의 경험을 믿고 언제 자신의 상상을 믿어야 하는지를 아는 로봇을 구축하는 데 있음을 시사합니다.

기술 요약: Q-Learning with World Models (QWM)

문제 정의

오프-폴리시(Off-policy) 강화 학습(RL)은 상당한 샘플 효율성을 달면하여, 시각-언어-행동(VLA) 모델을 신뢰할 수 있는 정책으로 미세 조정하는 것을 가능하게 했습니다. 그러나 RL이 실제 로보틱스와 같이 점점 더 복잡하고 고차원적인 시나리오에 적용됨에 따라, 추가적인 샘야 효율성 향상이 필수적입니다. 월드 모델(World models)은 단순히 행동을 예측하는 것을 넘어 상태 변화를 예측하는 메커니 당연한 메커니즘을 제공하지만, 이들의 성공적인 적용은 주로 지도 학습 기반의 정책 학습에 국한되어 왔습니다.

전통적인 모델 기반 RL 방법들은 대개 학습된 역학 모델(dynamics model)에서 생성된 "상상된(imagined)" 롤아웃(rollout) 상에서 직접 정책이나 가치 함수를 최적화합니다. 이러한 접근 방식은 누적 편향(compounding bias) 문제에 취과하기 쉬운데, 이는 월드 모델의 오류가 시간이 지남에 따라 축적되어 성능을 저하시키는 현상을 의미합니다. 이 문제는 작업 호라이즌(task horizon)이 길어지고 시각적 복잡성이 증가할수록 악화되며, 이는 실제 로보틱스로의 확장을 어렵게 만듭니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: 전통적인 모델 기반 접근 방식처럼 모델 편향을 훈련 과정에 주입하지 않고도, 월드 모델을 활용하여 온라인 RL 성능을 개선할 수 있는가?

방법론: QWM

저자들은 **QWM (Q-Learning with World Models)**을 제안합니다. 이 프레임워크는 학습된 월드 모델을 **테스트 타임 탐색(test-time search)**을 통해 표준 Q-러닝에 통합합니다. 기존의 모델 기반 RL과 달리, QWM은 월드 모델을 사용하여 훈련 데이터를 생성하거나 정책을 직접 최적화하지 않습니다. 대신, 정책과 가치 함수(critic)는 오로지 **실제 환경 전이(real, online environment transitions)**만을 사용하여 훈련됩니다. 월드 모델은 오직 결정 시점에 높은 가치를 가진 행동을 선택하기 위한 상상된 궤적(imagined trajectories) 탐색을 수행하는 데에만 사용됩니다.

핵심 구성 요소

월드 모델 훈련: 역학 모델 $M_\psi(s'|s, a)$ 는 환경 전이 데이터셋(오프라인 또는 온라인으로 축적된 데이터)을 통해 사전 훈련됩니다. 저차원 상태의 경우 결정론적 잔차 역학 모델(deterministic residual dynamics model)이 사용됩니다. 픽셀 기반 관측치의 경우, 액션 조건부 비디오 월드 모델(Wan2.2 기반)이 미세 조정됩니다.
행동에 대한 트리 탐색 (Tree Search over Actions): 각 결정 단계에서 QWM은 현재 상태 $s_0$ $s_{0}$ 를 루트로 하는 탐색 트리를 구축합니다.
- 확장 (Expansion): 정책은 $N$ 개의 후보 행동을 제안합니다. 각 행동에 대해, 월드 모델은 $K$ 개의 미래 상태를 예측합니다. 이 과정은 깊이 $D$ 까지 재귀적으로 반복됩니다.
- 가치 집계 (Value Aggregation): 트리의 노드 가치는 두 가지 추정치를 결합하여 계산됩니다:
  - 상태-행동 가치 ( $V_Q$ ): 학습된 Q-함수 $Q_\phi$ 를 사용하여 행동을 직접 점수화합니다. 이는 분산이 낮고 모델 오류로부터 독립적이지만, 미래의 롤아웃을 내다보지 못한다는 단점이 있습니다.
  - 상태 가치 ( $V_r$ ): 월드 모델이 예측한 보상과 후속 상태의 가치를 사용하여 가치를 추정합니다. 이는 전체 상상된 롤아웃을 활용하지만, 깊이가 깊어질수록 모델 오류가 누적됩니다.
  - 결합 가치 (Combined Value): 최종 노드 가치는 $V_Q$ 와 $V_r$ 의 가중 평균(통상 0.5)으로, 두 추정치의 이점을 균형 있게 결합합니다.
탐색 가지치기 (Search Pruning): 계산 복잡도를 관리하기 위해, 각 깊이에서 유망하지 않은 브랜치를 누적된 할인된 Q-점수(discounted Q-scores)를 기준으로 제거하는 휴리스틱을 사용하여 일부 경로만을 확장합니다.
행동 선택: 실행할 행동은 집계된 트리 탐색 가치가 가장 높은 행동으로 선택됩니다. 이 탐색은 온라인 롤아웃(리플레이 버퍼를 위한 고품질 데이터 수집용)과 평가(evaluation) 단계 모두에서 적용됩니다.

주요 기여

QWM 프레임워크: 표준 Q-러닝 상단에서 테스트 타임 탐색을 위해 월드 모델을 활용하는 단순하고 일반적인 프레임워크입니다. 이는 실제 환경 상호작작용을 대체하지 않으면서도 미래 궤적을 상상함으로써 행동 선택을 개선합니다.
누적 편향 회피: 정책과 크리틱을 실제 전이만을 사용하여 훈련함으로써, QWM은 불완전한 월드 모델 내에서 정책을 최적화할 때 발생하는 편향 축적 문제를 피합니다.
테스트 타임 스케일링 (Test-Time Scaling): 이 접근 방식은 후보군이 예측된 다운스트림 가치에 따라 점수화되는 일종의 테스트 타임 스케일링으로 볼 수 있으며, 이를 통해 단일 단계 Q-값 최대화보다 더 나은 최적 행동 추정치를 얻을 수 있습니다.

실험 결과

저자들은 도전적인 로봇 조작 벤치마크인 Robomimic(상태 기반)과 LIBERO(픽셀 기반)에서 QWM을 평가했습니다.

모델 프리(Model-Free) 베이스라인과의 비교: QWM은 모든 작업(Lift, Can, Square, Tool Hang 및 다양한 LIBERO 작업)에서 샘플 효율성과 최종 성능 모두에서 강력한 최신 모델 프리 베이스라인(EXPO, RLPD, IDQL 등 포함)을 크게 능가했습니다.
모델 기반(Model-Based) 베이스라인과의 비교: QWM은 TD-MPC2 및 EfficientZero V2와 같은 최신 모델 기반 방법들보다 일관되게 강력한 성능을 달성했습니다. 특히, TD-MPC2와 EfficientZero V2가 보고된 훈련 단계 내에서 희소 보상(sparse reward) 설정 시 성공률이 0에 가까웠던 것과 대조적으로, QWM은 모든 작업에서 성공을 거두었습니다.
절제 연구 (Ablation Studies):
- 탐색 범위 (Search Scope): 온라인 샘플링과 평가 모두에서 트리 탐색을 적용했을 때 최상의 결과를 얻었으며, 이는 데이터 수집과 실행 모두를 개선하는 상호 보완적 이점을 입증합니다.
- 하이퍼파라미터: 적절한 탐색 깊이( $D=2$ ~ $4$)가 미래의 결과를 포착하는 것과 과도한 모델 오류를 피하는 것 사이의 최적의 절충안을 제공했습니다.
- Q-함수 vs 가치 함수: Q-함수를 사용하여 탐색하는 것이 상태 가치 함수( $V$ )만을 사용하는 것보다 성능이 훨씬 높았으며, 이는 액션 조건부 추정의 중요성을 강조합니다.
픽셀 기반 스케일링: QWM은 고차원 시각적 관측치(LIBERO)로 성공적으로 확장되었으며, 월드 모델이 온라인 훈련 중 샘플링 용도로만 사용될 때에도 학습 속도와 최종 성능에서 명확한 개선을 보여주었습니다.

의의 및 주장

본 논문은 QWM이 예측 가능한 월드 모델의 능력과 온라인 Q-러닝의 안정성을 결합하는 실용적인 경로를 제공한다고 주장합니다. 월드 모델의 역할을 테스트 타임 탐색으로 제한함으로써, 이 방법은 상상된 롤아웃 상에서 정책을 훈련할 때 발생하는 누적 편향을 피하면서도 예측 계획(predictive planning)의 샘플 효율성 이점을 얻습니다.

저자들은 QWM을 Q-함수를 사용하는 모든 RL 미세 조정 알고리즘, 특히 시각-언어-행동(VLA) 모델에 적용 가능한 일반적인 개선책으로 위치시킵니다. 저자들은 롤아웃 시의 트리 탐색에 따른 계산 오버헤드와 고품질 월드 모델 학습의 어려움이라는 한계를 인정하면서도, QWM 프레임워크가 실제 로보틱스를 위한 강력한 대안을 제공한다고 단언합니다.