[Paper Review] When to Trust Your Model: Model-Based Policy Optimization

Michael Jänner, Justin Fu+2 more|arXiv (Cornell University)|Jun 19, 2019

Reinforcement Learning in Robotics44 references119 citations

TL;DR

MBPO uses short model rollouts branched from real data to achieve fast learning: it matches model-free asymptotic performance while improving data efficiency and avoiding long-horizon model pitfalls.

ABSTRACT

Designing effective model-based reinforcement learning algorithms is difficult because the ease of data generation must be weighed against the bias of model-generated data. In this paper, we study the role of model usage in policy optimization both theoretically and empirically. We first formulate and analyze a model-based reinforcement learning algorithm with a guarantee of monotonic improvement at each step. In practice, this analysis is overly pessimistic and suggests that real off-policy data is always preferable to model-generated on-policy data, but we show that an empirical estimate of model generalization can be incorporated into such analysis to justify model usage. Motivated by this analysis, we then demonstrate that a simple procedure of using short model-generated rollouts branched from real data has the benefits of more complicated model-based algorithms without the usual pitfalls. In particular, this approach surpasses the sample efficiency of prior model-based methods, matches the asymptotic performance of the best model-free algorithms, and scales to horizons that cause other model-based methods to fail entirely.

Motivation & Objective

• Motivate and analyze how to best use predictive models for policy optimization in reinforcement learning.
• Provide monotonic improvement guarantees under model-based updates while accounting for model error and distribution shift.
• Introduce a practical, empirically driven approach (MBPO) that uses short, branched model rollouts to improve data efficiency.
• Demonstrate that carefully controlled model usage can outperform prior model-based methods while maintaining strong asymptotic performance.

Proposed method

• Formulate a monotonic model-based policy improvement framework with generalization and distribution-shift errors psilon_m and psilon_pi and derive a bound on true returns in terms of model returns.
• Introduce branched rollouts that start from the data-collecting policy distribution and run k steps under the learned model to limit error accumulation.
• Propose MBPO: train an ensemble of probabilistic dynamics models, use SAC for policy optimization, and generate short model rollouts from replay buffer states.
• Use short, repeated model rollouts to create a large volume of model-generated data while mitigating model exploitation and horizon-coupling issues.
• In practice, measure model generalization empirically and adapt rollout usage to balance model-based and model-free updates.

Experimental results

Research questions

• RQ1 How can model-based updates guarantee monotonic improvement in policy performance given model error and distribution shift?
• RQ2 Under what conditions do short model rollouts provide practical benefits without exacerbating model exploitation or compounding error?
• RQ3 Can branched, short-horizon model rollouts yield faster learning while retaining best model-free asymptotic performance?
• RQ4 How does model generalization to unseen policy distributions affect the usefulness of model-based data?
• RQ5 What design choices (model ensembles, rollout length, optimization algorithm) optimize sample efficiency in model-based policy optimization?

Key findings

• MBPO achieves substantially faster learning than prior model-based methods while matching the final performance of leading model-free algorithms.
• On continuous control benchmarks, MBPO can reach model-free performance with an order of magnitude less data and steps (e.g., Ant task: 300k steps vs SAC at 3M steps).
• Short (even single-step) model rollouts provide significant benefits and longer rollouts can hurt due to error accumulation.
• A branched rollout strategy (starting from real data distribution and then simulating k steps) mitigates compounding errors and scales to longer horizons.
• Ensembles of probabilistic dynamics models help capture uncertainty and reduce model exploitation.
• Empirical measurements show model generalization improves with more data, enabling a more realistic bound that justifies model usage in practice.