[Paper Review] Uncertainty-aware Model-based Policy Optimization

Tung-Long Vuong, Kenneth Tran|arXiv (Cornell University)|Jun 25, 2019

Reinforcement Learning in Robotics26 references9 citations

TL;DR

This paper proposes Uncertainty-aware Model-based Policy Optimization (MBPO), a novel framework that jointly learns an uncertainty-aware dynamics model and optimizes a policy via differentiable planning. By combining model-based planning with policy gradient learning through automatic differentiation, MBPO achieves superior sample efficiency and competitive asymptotic performance compared to state-of-the-art model-free and model-based baselines on MuJoCo continuous control benchmarks.

ABSTRACT

Model-based reinforcement learning has the potential to be more sample efficient than model-free approaches. However, existing model-based methods are vulnerable to model bias, which leads to poor generalization and asymptotic performance compared to model-free counterparts. In addition, they are typically based on the model predictive control (MPC) framework, which not only is computationally inefficient at decision time but also does not enable policy transfer due to the lack of an explicit policy representation. In this paper, we propose a novel uncertainty-aware model-based policy optimization framework which solves those issues. In this framework, the agent simultaneously learns an uncertainty-aware dynamics model and optimizes the policy according to these learned models. In the optimization step, the policy gradient is computed by automatic differentiation through the models. With respect to sample efficiency alone, our approach shows promising results on challenging continuous control benchmarks with competitive asymptotic performance and significantly lower sample complexity than state-of-the-art baselines.

Motivation & Objective

• Address the high sample complexity of model-free reinforcement learning in real-world control settings.
• Overcome model bias in model-based RL by explicitly modeling uncertainty in dynamics predictions.
• Enable efficient, low-latency policy inference by replacing model predictive control (MPC) with an explicit policy representation.
• Support transfer learning by enabling reuse of both learned dynamics models and policies across similar tasks.
• Achieve high sample efficiency while maintaining competitive asymptotic performance compared to model-free methods.

Proposed method

• Train a deep neural network dynamics model with ensemble methods to estimate predictive uncertainty, using bootstrapping for scalable uncertainty quantification.
• Use automatic differentiation to compute policy gradients through the dynamics model, enabling end-to-end policy optimization.
• Replace traditional MPC with an explicit policy network, allowing direct policy evaluation and transferability.
• Perform online policy updates after each environment interaction, using off-policy data for efficient learning.
• Incorporate uncertainty estimates into the planning objective to reduce compounding errors over long horizons.
• Apply risk-sensitive planning by weighting uncertainty in value function estimation to improve robustness.

Experimental results

Research questions

• RQ1 Can uncertainty-aware dynamics modeling reduce model bias and improve long-horizon planning in model-based RL?
• RQ2 Does differentiable policy optimization through learned dynamics lead to better sample efficiency than MPC-based approaches?
• RQ3 Can the proposed framework support transfer learning of both dynamics models and policies across related tasks?
• RQ4 How does the sample efficiency of the proposed method compare to state-of-the-art model-free and model-based algorithms on continuous control benchmarks?
• RQ5 What is the impact of uncertainty quantification on asymptotic performance and training stability?

Key findings

• MBPO achieves significantly lower sample complexity than state-of-the-art baselines, including PPO, DDPG, SAC, and STEVE, on Pendulum-v0, Swimmer-v2, and HalfCheetah-v2.
• The method attains asymptotic performance competitive with or better than model-free methods like SAC and PPO, despite being trained with far fewer environment interactions.
• Model error and compounding effects are effectively mitigated through principled uncertainty quantification, as shown by reduced value function estimation errors over increasing planning horizons.
• The explicit policy representation enables transfer learning, allowing the agent to bootstrap from existing policies and dynamics models in similar environments.
• Despite strong performance, MBPO's results show sensitivity to random seed, suggesting instability in training dynamics due to aggressive online updates and lack of initial random exploration.
• The framework demonstrates that model-based RL can achieve both high sample efficiency and strong asymptotic performance when uncertainty is properly modeled and integrated into policy optimization.