Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models

Kurtland Chua; Roberto Calandra; Rowan McAllister; Sergey Levine

Source-linked AI summary

Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models

Kurtland Chua, Roberto Calandra, Rowan McAllister, Sergey Levine

arXiv:1805.12114v2May 30, 2018cs.LGcs.AIcs.ROstat.ML

TL;DR

Model-based reinforcement learning is sample-efficient but can lag in asymptotic performance, especially with neural-network dynamics models. PETS combines uncertainty-aware deep ensembles with trajectory sampling, matching model-free performance on benchmark tasks while using substantially fewer samples.

Problem
Model-based reinforcement learning can be sample-efficient, but neural-network dynamics models struggle with overfitting and long-horizon prediction, limiting asymptotic performance.
Method
PETS combines probabilistic neural-network dynamics models, bootstrapped ensembles, and trajectory sampling to propagate aleatoric and epistemic uncertainty.
Results
PETS matches model-free asymptotic performance across benchmark tasks, reaching PPO’s asymptotic performance in fewer than 100 trials on all four tasks.
Takeaways & Limitations
Uncertainty estimation can help neural-network model-based reinforcement learning bridge the asymptotic-performance gap with model-free methods at substantially lower sample complexity.
Takeaways & Limitations
The analysis focuses on dynamics estimation and planning, leaving policy-learning alternatives for future work.

Abstract

from arXiv · show

Model-based reinforcement learning (RL) algorithms can attain excellent sample efficiency, but often lag behind the best model-free algorithms in terms of asymptotic performance. This is especially true with high-capacity parametric function approximators, such as deep networks. In this paper, we study how to bridge this gap, by employing uncertainty-aware dynamics models. We propose a new algorithm called probabilistic ensembles with trajectory sampling (PETS) that combines uncertainty-aware deep network dynamics models with sampling-based uncertainty propagation. Our comparison to state-of-the-art model-based and model-free deep RL algorithms shows that our approach matches the asymptotic performance of model-free algorithms on several challenging benchmark tasks, while requiring significantly fewer samples (e.g., 8 and 125 times fewer samples than Soft Actor Critic and Proximal Policy Optimization respectively on the half-cheetah task).

1 Introduction

Model-free RL is expensive to train, while neural-network model-based RL struggles with overfitting and long-horizon prediction on complex dynamics. PETS addresses this gap with high-capacity probabilistic ensemble models and trajectory sampling, achieving model-free-like benchmark performance with fewer samples.

Contribution: PETS uses trajectory sampling to propagate uncertainty through the learned dynamics model.The propagation procedure re-samples particles using the dynamics model.
Motivation: Neural-network model-based RL is challenging because small datasets cause overfitting and poor predictions far into the future.Efficient models can learn quickly but struggle to represent complex and discontinuous dynamical systems.
Contribution: PETS combines high-capacity neural-network dynamics models with bootstrapped ensembles and probabilistic predictions to incorporate uncertainty.The ensemble captures uncertainty through disagreement among bootstraps, while each model represents distributions rather than point predictions.
Contribution: The approach narrows the gap between model-based and model-free RL by rivaling model-free performance on standard benchmark control tasks at a fraction of the sample complexity.This contribution targets the high training cost of current model-free RL algorithms.

2 Related work

Prior MBRL work balances Gaussian processes’ low-data strengths against neural networks’ scalability and expressiveness, while probabilistic neural-network modeling remains promising but unresolved. Deterministic neural-network models can overfit early, motivating uncertainty-aware approaches.

Neural-network dynamics: Most neural-network MBRL methods use deterministic models, which suffer from early-stage overfitting; the proposed approach therefore improves data efficiency over deterministic MBRL such as Nagabandi et al. [2017].
Probabilistic neural networks: Constructing effective Bayesian neural networks remains open [MacKay, 1992, Neal, 1995, Osband, 2016, Guo et al., 2017], despite progress using dropout, ensembles [Osband et al., 2016, Lakshminarayanan et al., 2017], and α-divergence.Probabilistic neural networks have also been applied to control using dropout [Gal et al., 2016; Higuera et al., 2018] and α-divergence [Depeweg et al., 2016].

3 Model-based reinforcement learning

The framework models the system as a Markov decision process, learns forward dynamics from real-system measurements, and uses the learned model to evaluate and select action sequences.

The MBRL framework formulates states, actions, rewards, and transitions within a Markov decision process.
Forward-dynamics learning fits an approximation ef of the true transition function f using measured state-action transitions D from the real system.
The learned dynamics model predicts state-trajectory distributions for candidate action sequences, whose expected rewards determine the selected sequence.

4 Uncertainty-aware neural network dynamics models

The section develops uncertainty-aware neural-network dynamics models for model-based reinforcement learning, motivated by the need to perform well across scarce- and abundant-data regimes. It distinguishes aleatoric from epistemic uncertainty and proposes bootstrapped probabilistic ensembles to represent both.

Ensembles: Bootstrapped probabilistic neural-network ensembles model epistemic uncertainty while retaining expressive neural dynamics models for model-based reinforcement learning.The section presents this ensemble as the paper’s method and contrasts it with prior uncertainty-aware work whose implementation choices were not rigorously analyzed empirically [Gal et al., 2016, Depeweg et al., 2016].
Motivation: Model choice is crucial because even small dynamics-model bias can substantially affect controller quality, while Gaussian-process models scale poorly and often use kernels unsuitable for contact-rich discontinuous dynamics [Atkeson and Santamaría, 1997, Abbeel et al., 2006, Calandra et al., 2016].The design challenge is to perform well in both low-data and high-data regimes.
Uncertainty types: The models represent aleatoric uncertainty through predicted distribution parameters and epistemic uncertainty through uncertainty over plausible neural dynamics models.Aleatoric uncertainty includes observation and process noise, whereas epistemic uncertainty reflects uncertainty not captured by purely discriminative training.
Probabilistic neural networks: Probabilistic neural networks output conditional distribution parameters and are trained with negative log prediction probability, such as Gaussian means and diagonal covariances.This Gaussian formulation captures heteroscedastic, input-dependent noise but does not capture epistemic uncertainty by itself.
Deterministic neural networks: Deterministic neural networks output point predictions trained with mean squared error, but their fixed variance does not provide a meaningful uncertainty estimate for uncertainty-aware propagation.They are defined as a special case of probabilistic networks with delta-distribution outputs.
Ensembles: Each ensemble member is trained on a bootstrap-resampled dynamics dataset, and five models were sufficient for all experiments.Bootstrap datasets are formed by sampling the recorded dynamics dataset with replacement.

5 Planning and control with learned dynamics

This section presents planning with probabilistic dynamics models, using model-predicted trajectories to select high-reward action sequences under uncertainty. PETS propagates plausible trajectories with particle-based trajectory sampling, including variants that represent time-varying or time-invariant model uncertainty.

Planning with learned dynamics: Probabilistic dynamics models predict candidate trajectories, and model-predictive control applies the first action from the sequence maximizing expected reward over a finite horizon.The controller replans from the resulting state at each time step.
Planning with learned dynamics: Random shooting approximates optimal action-sequence search through parallelizable sampling, while PETS improves upon deterministic-model control by incorporating uncertainty into modeling and planning.This extends data-efficient control beyond the deterministic neural-network approach of Nagabandi et al. [2017].
Trajectory sampling: Trajectory sampling represents the probabilistic model’s future-state distribution with particles, enabling uncertainty propagation when closed-form recursive expected rewards are generally intractable.Particle, deterministic, and parametric approaches differ in how they represent propagated state distributions.
Trajectory sampling: TS1 resamples ensemble bootstraps independently at each time step, whereas TS∞ keeps each particle’s bootstrap fixed throughout a trial to represent time-invariant dynamics uncertainty.TS1 restricts trajectory multimodality from bootstrap changes, while TS∞ separates aleatoric and epistemic uncertainty [Depeweg et al., 2018].
Alternative state propagation methods: The method is compared with deterministic expectation, Gaussian moment matching, and distribution sampling, which impose progressively different restrictions on multimodality and epistemic uncertainty.Expectation uses one propagated particle; moment matching enforces a unimodal Gaussian state distribution; distribution sampling moment-matches only across bootstraps.

6 Algorithm summary

PETS combines probabilistic-ensemble dynamics models with trajectory-sampling planning in model-predictive control. It initializes with random data, repeatedly trains the model, optimizes action sequences with CEM, propagates particles, and executes only the first action.

6 Algorithm summary: PETS uses probabilistic ensembles to represent heteroskedastic aleatoric and epistemic uncertainty, which trajectory sampling can exploit during planning.
6 Algorithm summary: CEM learned faster than random shooting for guiding PETS's MPC action search.This finding is discussed in Appendix A.8.
6 Algorithm summary: PETS initializes data with one random-controller trial, trains an ensemble dynamics model, samples candidate action sequences, propagates particles with trajectory sampling, and evaluates them.
6 Algorithm summary: The MPC controller updates the CEM action distribution and executes only the first action from the selected optimal sequence before replanning.
6 Algorithm summary: Under 100K time steps or 100 trials, PETS learned across all tasks and significantly outperformed baselines except Cartpole, where a Gaussian-process model was efficient.Figure 3 averages maximum rewards over 10 experiments; each time step represents 0.01 seconds, or 0.02 seconds for Cartpole.

7 Experimental results

PETS is evaluated against state-of-the-art model-free and model-based methods on four continuous-control tasks, then analyzed through model and uncertainty-propagation ablations. It matches PPO’s asymptotic performance with far fewer samples, while probabilistic ensembles and model quality drive performance more strongly than advanced propagation techniques.

7.1 Benchmark comparison: PETS matches PPO’s asymptotic performance on all four tasks in fewer than 100 trials, whereas PPO requires several orders of magnitude more samples.Its asymptotic performance also substantially exceeds that of the prior model-based method by Nagabandi et al. [Nagabandi et al., 2017].
7.2 Design analysis: Probabilistic ensembles perform best across tasks except cartpole, followed by probabilistic networks and deterministic ensembles, while deterministic single models perform worst.The evaluated combinations include different dynamics models and uncertainty-propagation techniques.
7.2 Design analysis: Model quality and using uncertainty during learning substantially affect MBRL performance, whereas more advanced uncertainty-propagation techniques provide only minor improvements.Moment matching is competitive in low-dimensional settings.
7.2 Design analysis: The experiments show that combining ensembles, probabilistic networks, and approximate uncertainty-propagation techniques yields substantially improved results over previous approaches.These components existed previously in supervised learning and reinforcement learning, but their particular combination improves the proposed method.

8 Discussion & conclusion

The experiments suggest that uncertainty-aware neural-network dynamics models can narrow the asymptotic-performance gap between model-based and model-free reinforcement learning while converging more efficiently. They also indicate that both epistemic and aleatoric uncertainty are important to this process.

Neural-network model-based reinforcement learning matches model-free algorithms such as PPO and SAC in asymptotic performance while converging substantially more efficiently.
Incorporating uncertainty estimation into model learning can at least partly bridge the asymptotic-performance gap between model-based and model-free reinforcement learning.
The findings indicate that both epistemic and aleatoric uncertainty play crucial roles in this process.

A Appendix · A.1 Well behaved probabilistic networks

Probabilistic network variance outputs can become undefined outside the training distribution, unlike the bounded, Lipschitz-smooth variance of Gaussian processes. The appendix bounds variance using automatic-differentiable transformations and regularizes the learned limits.

A.1 Well behaved probabilistic networks: Variance outputs implemented by exponentiating predicted log variance are well behaved within training data but undefined outside that distribution.The issue arises because the transformed variance is constrained to be nonnegative but not controlled out of distribution.
A.1 Well behaved probabilistic networks: Gaussian-process variance provides a contrasting model because it is bounded and Lipschitz-smooth.
A.1 Well behaved probabilistic networks: The proposed remedy bounds probabilistic-network variance between its minimum and maximum values observed during training.
A.1 Well behaved probabilistic networks: The implementation first constrains logvar to an upper bound using max_logvar and softplus.
A.1 Well behaved probabilistic networks: It then constrains logvar to a lower bound using min_logvar and softplus before computing var = tf.exp(logvar).
A.1 Well behaved probabilistic networks: Regularization keeps max_logvar from exceeding the training maximum and prevents min_logvar from falling below the training minimum.

A.2 Fitting PE model to toy function

The toy-function experiment fits the evaluated models to 2,000 heteroscedastically noisy sine-function samples. Figure A.5 illustrates that the probabilistic ensemble captures aleatoric uncertainty and uses bootstrap disagreement to represent epistemic uncertainty.

A.2 Fitting PE model to toy function: The models were evaluated by fitting 2,000 points sampled from a sine function over [−2π, −π] ∪[π, 2π] after introducing heteroscedastic noise.The samples are generated as {(x_i, y_i)} with x_i drawn uniformly from the two specified intervals.
A.2 Fitting PE model to toy function: The fit corresponding to equation (2) is shown in Figure A.5, reproduced from Figure 1 for convenience.
A.2 Fitting PE model to toy function: The probabilistic ensemble captures observation noise while representing epistemic uncertainty through disagreement among bootstrap models, especially far from observed data.The ensemble uses two bootstraps for visualization, although the authors normally use five; the bootstraps agree near data and disagree far from it.

A.3 One-step predictions of learned models … A.10 Linear model comparison:

The appendix validates PETS across prediction accuracy, uncertainty propagation, control-design choices, stochastic systems, and benchmark settings. Probabilistic propagation and guided action selection improve robustness or performance, while linear models fail on nonlinear half-cheetah dynamics.

A.3 One-step predictions of learned models: After 100 training trials, one-step predictions are most certain where training data are dense and less certain at extreme values.Figure A.6 shows cartpole angular-velocity predictions with two standard deviations, a dimension described as harder to predict.
A.4 Uncertainty propagation methods: Particle-based, moment-matching, and other uncertainty-propagation methods are compared on a learned cartpole model solving a pole-swing-up action sequence.Figure A.7 visualizes how the discussed propagation methods evolve particles from an intermediate state.
A.5 Forward Dynamics Model: The forward model predicts state differences, and angular states are represented using sine and cosine inputs to capture rotational structure.This formulation follows the suggestion in [Deisenroth et al., 2014].
A.6 Experimental setting: The experiments use four MuJoCo continuous-control benchmarks spanning complexity, dimensionality, and contact forces, with repeated seeds and reported mean and standard-deviation costs.Tasks include cartpole swing-up, PR2 reaching and pushing, and half-cheetah; neural dynamics models use three fully connected layers with swish activations.
A.7 Additional considerations: Probabilistic propagation tolerates overly long MPC horizons by separating particles and reducing reliance on unpredictable future costs.Deterministic methods lack this mechanism and therefore remain more exposed to model bias from horizon choice.
A.8 MPC action selection: CEM significantly outperforms random shooting on half-cheetah ground-truth dynamics using the same 2500 action samples.CEM uses five iterations of 500 candidates, whereas random shooting samples 2500 actions directly.
A.9 Stochastic systems:: PETS methods that propagate uncertainty generally provide more consistent cartpole performance under stochastic action noise.The evaluation compares modeling decisions with additive action noise equal to 10% of each action range.
A.10 Linear model comparison:: A linear model cannot capture half-cheetah dynamics well enough, motivating nonlinear learned dynamics models.This comparison is presented in Figure A.13.