[Paper Review] Trust the Model When It Is Confident: Masked Model-based Actor-Critic

Feiyang Pan, Jia He+2 more|arXiv (Cornell University)|Oct 10, 2020

Reinforcement Learning in Robotics21 references19 citations

TL;DR

This paper proposes Masked Model-based Actor-Critic (M2AC), a novel model-based reinforcement learning algorithm that restricts policy optimization to state-action pairs where the model is confident, using uncertainty-based masking to reduce reliance on inaccurate model rollouts. M2AC achieves superior sample efficiency and robustness—especially in noisy environments with long rollouts—outperforming state-of-the-art methods like MBPO across continuous control benchmarks.

ABSTRACT

It is a popular belief that model-based Reinforcement Learning (RL) is more sample efficient than model-free RL, but in practice, it is not always true due to overweighed model errors. In complex and noisy settings, model-based RL tends to have trouble using the model if it does not know when to trust the model. In this work, we find that better model usage can make a huge difference. We show theoretically that if the use of model-generated data is restricted to state-action pairs where the model error is small, the performance gap between model and real rollouts can be reduced. It motivates us to use model rollouts only when the model is confident about its predictions. We propose Masked Model-based Actor-Critic (M2AC), a novel policy optimization algorithm that maximizes a model-based lower-bound of the true value function. M2AC implements a masking mechanism based on the model's uncertainty to decide whether its prediction should be used or not. Consequently, the new algorithm tends to give robust policy improvements. Experiments on continuous control benchmarks demonstrate that M2AC has strong performance even when using long model rollouts in very noisy environments, and it significantly outperforms previous state-of-the-art methods.

Motivation & Objective

• Address the fundamental challenge in model-based RL: deciding when to trust the model, especially in noisy or low-data regimes.
• Overcome the performance degradation of existing model-based methods when using long model rollouts due to accumulated model errors.
• Improve sample efficiency and robustness by selectively using model-generated data only when the model is confident.
• Develop a practical algorithm that balances model-based rollouts with real environment experience through uncertainty-aware masking.
• Demonstrate that restricting model use to low-error regions leads to tighter bounds on the true value function and better policy performance.

Proposed method

• Propose a masking mechanism that filters model-generated transitions based on model uncertainty, allowing only high-confidence predictions to be used in policy learning.
• Introduce a model-based lower-bound of the true value function that accounts for uncertainty, enabling safe policy optimization using masked rollouts.
• Use an ensemble of models with one-vs-rest (OvR) uncertainty estimation to quantify prediction disagreement and guide masking decisions.
• Implement a linearly decaying masking rate $ w_h = \frac{H_{\text{max}} - h}{2(H_{\text{max}} + 1)} $ to progressively reduce reliance on model rollouts as rollout depth increases.
• Combine masked model rollouts with real experience replay in a hybrid actor-critic framework to stabilize training and improve generalization.
• Apply a penalty coefficient $ \alpha $ to balance exploration and exploitation, with $ \alpha = 0.001 $ found optimal in ablation studies.

Experimental results

Research questions

• RQ1 Can restricting model-based rollouts to only high-confidence predictions reduce the performance gap between model and real environment returns?
• RQ2 How does uncertainty-based masking affect sample efficiency and robustness in noisy or complex environments?
• RQ3 What is the optimal masking strategy (e.g., masking rate, mode) for long-horizon model rollouts in continuous control tasks?
• RQ4 How does the choice of uncertainty estimation method (e.g., OvR vs. OvA vs. negative likelihood) impact algorithm performance?
• RQ5 Can a model-based RL algorithm maintain strong performance even when using long rollouts in environments with high dynamics noise?

Key findings

• M2AC significantly outperforms MBPO and other SOTA methods in HalfCheetah-noisy2 environments, achieving higher returns even with 20-step model rollouts where MBPO fails.
• In HalfCheetah and Walker2d tasks, M2AC maintains stable performance across varying noise levels and rollout lengths, while MBPO’s performance degrades rapidly with increased noise.
• The masking rate $ w = 0.25 $ (via linear decay) yields robust performance across environments and rollout lengths, with only minor degradation at $ w = 0.9 $ in long rollouts.
• The non-stop masked rollout mode outperforms the hard-stop mode due to a richer distribution of imaginary transitions, enhancing policy generalization.
• OvR uncertainty estimation outperforms OvA disagreement and negative likelihood in low-data regimes, as it better captures model disagreement.
• With $ \alpha = 0.001 $, M2AC achieves the best balance between exploration and exploitation, avoiding overfitting to the model while still leveraging its predictions effectively.