HaWMPO: Hallucination-Aware World Model-based Policy Optimization for Generalist Robot Policy

기술 요약: HaWMPO

문제 정의

범용 로봇 정책, 특히 시각-언어-행동(VLA) 모델은 강력한 다중 작업 일반화 능력을 보여주었으나, 복잡하고 긴 호흡(long-horizon)의 시나리오에서는 성공률이 저하되는 문제를 겪고 있습니다. 물리적 로봇을 이용한 온라인 강화 학습(RL)은 이러한 정책을 개선할 수 있지만, 높은 샘플 비효효율성, 비용이 많이 드는 물리적 상호작용, 그리고 상당한 안전 위험으로 인해 제약이 따릅니다.

월드 모델(World models)은 가상 공간에서의 "상상된" 롤아웃(rollouts)을 통해 정책 최적화를 가능하게 함으로써 유망한 대안을 제시합니다. 그러나 생성형 월드 모델은 긴 호흡의 롤아웃 과정에서 누적된 예측 오류와 환각(hallucinations) 현상에 취약합니다. 이러한 환각은 편향된 상태 전이와 잘못된 학습 신호를 생성합니다. 만약 VLA 정책이 이러한 환각된 궤적을 직접 학습하게 된다면, 실제 환경에서 효과적인 행동을 배우는 대신 월드 모델의 편향을 악용하게 되어, 정책 업데이트가 불안정해지고 성능이 저하될 수 있습니다.

방법론: HaWMPO

저자들은 VLA 사후 학습(post-training) 중 월드 모델의 환각으로 인한 부정적인 영향을 완화하기 위해 설계된 폐쇄 루프 RL 파이프라인인 **환각 인지 월드 모델 기반 정책 최적화(Hallucination-Aware World Model-based Policy Optimization, HaWMPO)**를 제안합니다. 이 프레임워크는 세 가지 핵심 구성 요소와 하나의 새로운 최적화 메커니즘으로 이루어져 있습니다.

1. 행동 조건부 환각 인지 모델 (HAM)

롤아웃의 신뢰성 문제를 해결하기 위해, HaWMPO는 생성된 이미지 시퀀스의 신뢰도를 추정하는 전용 모델을 도입합니다.

• 입력: HAM은 8개 프레임의 예측 비디오 청크($\hat{V}_c$), 조건부 행동($a_c$), 현재 관측값($I_c$), 그리고 앵커 이미지($I_0$)를 입력으로 받습니다.
• 구조: 시각적 특징 추출을 위해 동결된 DINOv3 인코더를 사용하며, 입력을 공유 공간으로 매핑하기 위한 Context 및 Action Projector, 그리고 표현을 통합하기 위한 Fusion Transformer(교차 주의 집중 및 자기 주의 집중 포함)를 활용합니다. 학습 가능한 쿼리와 MLP 헤드는 연속적인 **환각 점수(hallucination score, $\hat{h}_c$)**를 출력합니다.
• 감독(Supervision): 이 모델은 DINOv3 유사도, 깊이 일관성, 광학 흐름(optical-flow) 기반 궤적 일관성, 그리고 MUSIQ 이미지 품질라는 네 가지 지표에서 유도된 복합 타겟을 사용하여 학습됩니다. 점수가 높을수록 환각이 심함을 의미합니다.

2. Reward-Soft 메커니즘

모든 월드 모델 생성 롤아웃을 동일하게 취급하는 대신, HaWMPO는 예측된 환각 수준에 따라 기여도를 가중 조절합니다.

• 메커니즘: 특정 행동 청크에 대해 보상 모델이 제공하는 원시 보상($R$)은 해당 청크의 평균 환각 점수($H$)에 의해 조절됩니다.
• 공식: $\hat{R} = (1 - \alpha \cdot H) \cdot R$ (여기서 $\alpha$는 튜닝 가능한 하이퍼파라미터입니다).
• 효과: 이 메커니즘은 환각이 심한 궤적 내에서 발생하는 겉보기에 성공적인 RL 단계에 페널티를 부여함으로써, 정책이 편향된 신호로부터 학습하는 것을 방지합니다.

3. 정책 최적화 파이프라인

시스템은 VLA 정책, 행동 조건부 월드 모델, HAM, 그리고 보상 모델을 그룹 상대 정책 최적화(GRPO) 프레임워크에 통합합니다.

• 프로세스: VLA 에이전트는 월드 모델에 의해 구축된 병렬적인 상상 환경에서 롤아웃을 실행합니다.
• 학습: 각 행동 청크에 대해 HAM은 환각 점수를 예측하고, Reward-Soft 메커니즘은 보상을 조절합니다. 이후 GRPO 알고리즘은 환각 인지 보상을 사용하여 VLA 정책을 업데이트하며, 학습을 안정화하기 위해 그룹 단위로 정규화된 어드밴티지(advantage)를 계산합니다.

주요 기여

1. 폐쇄 루프 RL 파이프라인: 저자들은 월드 모델, 환각 인지 모델, 보상 모델을 통합하여 빈번한 실세계 로봇 상호작용 없이도 온라인 스타일의 정책 최적화를 가능하게 하는 파이프라인을 구축했습니다.
2. 정량적 신뢰도 추정: 특정 행동 입력 하에서의 월드 모델 예측 신뢰도를 명시적으로 정량화하는 행동 조건부 HAM을 설계했습니다.
3. Reward-Soft 메커니즘: 환각이 심한 행동 세그먼트에 페널티를 부여하여 월드 모델의 악영향인 편향을 완화하는, GRPO 최적화에 환각 점수를 통합하는 새로운 방법을 제안했습니다.

실험 결과

시뮬레이션 (LIBERO 벤치마크)

LIBERO Object, Spatial, Goal 태스크 세트에서 평가되었습니다:

• 성능: HaWMPO는 **63.7%**의 가장 높은 평균 성공률을 달s했습니다.
• 비교: 이는 기본 OpenVLA-OFT 모델 대비 15.0%, WMPO 대비 6.9%, 그리고 가장 강력한 베이스라인인 WoVR* 대비 2.8% 향상된 수치입니다.
• 구체적 이득: Spatial 세트에서 HaWMPO는 77.2%를 달성하여(WoVR보다 8.0포인트 상회), Object 세트에서는 52.2%를 달성했습니다(WoVR보다 2.6포인트 상회).

실세계 실험 (G1 로봇)

"Tissue-to-Box" 및 "Headphone-on-Stand" 두 가지 조작 태스크에 대해 실험이 수행되었습니다.

• 결과: Ha-WMPO는 평균 **80.0%**의 성공률(Tissue-to-Box 85%, Headphone-on-Stand 75%)을 기록했습니다.
• 개선: 이는 기본 정책보다 12.5 퍼센트 포인트, WoVR 베이스라인보다 7.5 퍼센트 포인트 높은 수치입니다.

절제 실험 및 분석 (Ablation and Analysis)

• 환각 탐지: HAM은 인간의 주석과 비교하여 환각된 청크를 탐지하는 데 있어 0.9375의 AUROC를 달성하였으며, 이는 단일 지표 베이스라인(예: DINO-only, 0.9062)보다 우수한 성능입니다.
• Reward-Soft의 영향: 절제 실험 결과, HAM과 Reward-Soft 메커니즘을 제거하면 특히 Spatial 세트에서 상당한 성능 저하(다양한 학습 단계에서 8.0~11.6포인트 하락)가 발생함을 확인했습니다.
• 하이퍼파라미터 민감도: 최적의 페널티 강도($\alpha$)는 태스크에 따라 의존적임을 확인했습니다. Object 및 Spatial에는 $\alpha=0.3$이, Goal에는 $\alpha=1$이 가장 적합했습니다.

의의 및 주장

본 논문은 월드 모델 기반 RL의 핵심 병목 구간인 '예측 환각으로 인한 불안정성'을 해결한다고 주장합니다. 이러한 환각을 명시적으로 모델링하고 신뢰할 수 없는 합성 궤적의 가중치를 낮춤으로써, 본 방법론은 더 견고한 정책 최적화를 가능하게 합니다. 저자들은 이 접근 방식이 VLA 정책이 월드 모델의 효율성을 활용하면서도, 환각된 데이터의 편향으로부터 학습하는 함정을 피할 수 있게 하여 시뮬레이션과 실세계 물리적 조작 태스크 모두에서 우수한 성능을 이끌어낸다고 확언합니다.

한계점

저자들은 다음과 같은 몇 가지 한계를 인정합니다:

• HAM은 수동으로 설계된 감독 신호(특징 불일치, 깊이, 궤적 등)에 의존하므로 모든 유형의 환각을 포착하지 못할 수 있습니다.
• 현재 평가는 LIBERO 벤치마크의 상대적으로 짧은 호흡의 태스크에 국한되어 있으며, 복잡하고 긴 호흡의 실세계 시나리오로의 확장성은 아직 탐구되지 않았습니다.
• 환각 인지 모델과 월드 모델이 별도로 학습되므로, 온라인 RL 과정 중의 시스템 적응성을 제한할 수 있습니다. 향론적인 과제로 엔드 투 엔드(end-to-end) 공동 최적화가 제안되었습니다.