WIMLE: 불확실성을 고려한 세계 모델 - IMLE를 활용한 샘플 효율적인 연속 제어

WIMLE: Uncertainty-Aware World Models with IMLE for Sample-Efficient Continuous Control

모델 기반 강화 학습은 뛰어난 샘플 효율성을 제공하지만, 종종 모델 오류의 누적, 다중 모드 동역학을 평균화하는 단일 모드 세계 모델, 그리고 학습을 왜곡하는 과도한 예측으로 인해 실제 성능이 좋지 않은 경우가 많습니다. 본 논문에서는 Implicit Maximum Likelihood Estimation (IMLE)을 모델 기반 강화 학습 프레임워크로 확장하여 확률적이고 다중 모드 세계 모델을 반복적인 샘플링 없이 학습하며, 앙상블과 잠재 샘플링을 통해 예측 불확실성을 추정하는 WIMLE라는 모델 기반 방법을 제안합니다. 학습 과정에서 WIMLE는 각 합성 트랜지션을 예측된 신뢰도에 따라 가중치를 부여하여 유용한 모델 롤아웃을 유지하고, 불확실한 예측에서 비롯되는 편향을 줄이며 안정적인 학습을 가능하게 합니다. DeepMind Control, MyoSuite, 그리고 HumanoidBench에 걸쳐 40개의 연속 제어 작업에서 WIMLE는 강력한 모델 기반 및 모델 프리 기반의 비교 기준보다 우수한 샘플 효율성과 경쟁력 있는 또는 더 나은 최종 성능을 달성했습니다. 특히, 어려운 Humanoid-run 작업에서 WIMLE는 가장 강력한 경쟁 모델보다 샘플 효율성을 50% 이상 향상시켰으며, HumanoidBench에서는 14개의 작업 중 8개를 해결했습니다 (BRO는 4개, SimbaV2는 5개). 이러한 결과는 IMLE 기반의 다중 모드 학습과 불확실성을 고려한 가중치 부여가 안정적인 모델 기반 강화 학습에 얼마나 중요한지를 보여줍니다.

Model-based reinforcement learning promises strong sample efficiency but often underperforms in practice due to compounding model error, unimodal world models that average over multi-modal dynamics, and overconfident predictions that bias learning. We introduce WIMLE, a model-based method that extends Implicit Maximum Likelihood Estimation (IMLE) to the model-based RL framework to learn stochastic, multi-modal world models without iterative sampling and to estimate predictive uncertainty via ensembles and latent sampling. During training, WIMLE weights each synthetic transition by its predicted confidence, preserving useful model rollouts while attenuating bias from uncertain predictions and enabling stable learning. Across $40$ continuous-control tasks spanning DeepMind Control, MyoSuite, and HumanoidBench, WIMLE achieves superior sample efficiency and competitive or better asymptotic performance than strong model-free and model-based baselines. Notably, on the challenging Humanoid-run task, WIMLE improves sample efficiency by over $50$\% relative to the strongest competitor, and on HumanoidBench it solves $8$ of $14$ tasks (versus $4$ for BRO and $5$ for SimbaV2). These results highlight the value of IMLE-based multi-modality and uncertainty-aware weighting for stable model-based RL.