On Effective Scheduling of Model-based Reinforcement Learning

1 Introduction . Deep model-free reinforcement learning ( MFRL ) has achieved great successes in complex decisionmaking problems such as Go [ 26 ] and robotic control [ 10 ] , to name a few . Although MFRL can achieve high asymptotic performance , a tremendous number of samples collected in interactions with the environment are required . In contrast , model-based reinforcement learning ( MBRL ) , which alternately learns a model of the environment and derives an agent with the help of current model estimation , is considered to be more sample efficient than MFRL [ 27 ] . MBRL methods generally fall into different categories according to the specific usage of the learned model [ 25 , 33 ] . Among them , Dyna-style algorithms [ 28 ] adopt some off-the-shelf MFRL methods to train a policy with both real data from environment and imaginary data generated by the model . Since Dyna-style algorithms can seamlessly take advantage of innovations in MFRL literature and have recently shown impressive performance [ 12 , 30 ] , this paper mainly focuses on Dyna-style algorithms . Although some theoretical analysis [ 17 , 12 , 22 ] and thorough empirical evaluation [ 30 ] have been conducted in the previous literature , it remains unclear how to appropriately schedule the hyperparameters to achieve optimum performance when training a Dyna-style MBRL algorithm . In practice , many important hyperparameters may primarily affect performance . Firstly , since Dyna-style MBRL algorithms consist of model learning and policy optimization , how to balance the alternate optimization of these two parts is a crucial problem [ 22 ] . Intuitively , insufficient model training may lead to inaccurate estimation while excessive model training will cause overfitting , and a small amount of policy training may not utilize the model adequately while exhaustive policy training will exploit the model ’ s deficiencies . Moreover , the rollout length of the imaginary trajectory ∗equal contribution . †Weinan Zhang is the corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . is also critical [ 12 ] since too short rollout length fails to sufficiently leverage the model to plan forward , while too long rollout length may bring disastrous compounding error [ 1 ] . Janner et al . [ 12 ] manually design a schedule that linearly increases the rollout length across epochs . Finally , when using both real samples and imaginary samples to train the policy , how to control the ratio of the two datasets remains unclear . Janner et al . [ 12 ] fix the ratio of real data as 5 % , while Kalweit and Boedecker [ 13 ] use the uncertainty to adaptively choose the ratio , which tends to use more imaginary samples initially and gradually use more real samples afterward . According to these existing works , the optimal hyperparameters in model-based methods may be dynamic during the whole training process , which further strengthens the burden of manually scheduling the hyperparameters . Based on these considerations , in this work , we aim to investigate how to appropriately schedule these hyperparameters , i.e. , real data ratio , model training frequency , policy training iteration , and rollout length , to achieve optimal performance of Dyna-style MBRL algorithms . Although real data ratio is an essential factor empirically [ 13 , 12 ] , it has not yet been studied thoroughly in theory . To bridge this gap , we first derive a return discrepancy upper bound for model-based value iteration , which reveals that gradually increasing the ratio of real data yields a tighter bound than choosing a fixed value . Inspired by the analysis , considering the complex interplay between real data ratio and other hyperparameters in practice , we develop AutoMBPO to automatically determine the joint schedule of the above hyperparameters in model-based policy optimization ( MBPO ) , a representative running case of Dyna-style MBRL algorithms . Specifically , AutoMBPO introduces a parametric hyper-controller to sequentially choose the value of hyperparameters in the whole optimization process to maximize the performance of MBPO . We apply AutoMBPO to several continuous control tasks . On all these tasks , the MBPO instance trained with hyperparameters scheduled by AutoMBPO can significantly surpass the one with original configuration [ 12 ] . Furthermore , the hyperparameter schedule found by AutoMBPO is consistent with our theoretical analysis , which yields a better understanding of MBPO and provides insights for the design of other MBRL methods . 2 Related Work . Model-based reinforcement learning ( MBRL ) has been widely studied due to its high sample efficiency . MBRL methods can be roughly categorized into four types according to different model usage : ( i ) Dyna-style algorithms [ 28 , 17 , 4 , 14 , 12 , 15 ] leverage the model to generate imaginary data and adopt some off-the-shelf MFRL algorithms to train a policy using both real data and imaginary data ; ( ii ) shooting algorithms [ 19 , 3 ] use model predictive control ( MPC ) to plan directly without explicit policy ; ( iii ) analytic-gradient algorithms [ 7 , 16 , 5 ] search policies with back-propagation through time by exploiting the model derivatives ; ( iv ) model-augmented value expansion algorithms [ 9 , 2 ] utilize model rollouts to improve the target value for temporal difference ( TD ) updates . This paper mainly focuses on the first category . Dyna-style algorithms [ 28 ] typically alternate between sampling real data , learning dynamics models , generating imaginary data , and optimizing policy . Many efforts have been made to improve one or more of these steps . For model learning , the deep ensemble technique [ 21 , 3 , 14 ] has been leveraged to resist overfitting . As for imaginary data generation , it is suggested to generate short rollouts to reduce compounding model error [ 12 ] , and the model ’ s uncertainty is further incorporated to choose reliable imaginary data dynamically [ 20 ] . When optimizing the policy , we can use merely imaginary data [ 17 , 14 ] or a mixture of real data and imaginary one in a fixed ratio [ 12 ] . Moreover , the ratio can also be dynamically adjusted according to the estimated uncertainty of the Q-function [ 13 ] . Besides , Rajeswaran et al . [ 22 ] present a game-theoretic framework for MBRL by formulating the optimization of model and policy as a two-player game . As for automatic hyperparameter optimization for MBRL , Zhang et al . [ 32 ] utilize population based training ( PBT ) to tune the hyperparameters of a shooting algorithm dynamically . The reinforcement on reinforcement ( RoR ) framework proposed by Dong et al . [ 8 ] applies an RL algorithm to control the sampling and training process of an MBRL method . Our work differs from theirs in four major aspects : ( 1 ) We provide a theoretical analysis for the hyperparameter being scheduled , which is not included in their work . ( 2 ) In terms of problem settings , RoR assumes one can directly access an initial state without further interaction from that state , which is impractical for most RL problems . ( 3 ) As for the method design , our formulated hyper-MDP in Section 5.1 is totally different from RoR ’ s , including state , action and reward . In detail , our state definition contains more training information , and the reward design is based on the return rather than a single-step reward in the target-MDP . ( 4 ) Finally , in the experiments in Section 6 , our method achieves better performance and generalization in all environments , and the empirical findings also show differences . 3 Preliminaries . We first briefly introduce the RL problem with notations used throughout this paper and the previous MBRL method on which our proposed framework is based . 3.1 Reinforcement Learning . In reinforcement learning ( RL ) , there originally exists a Markov decision process ( MDP ) , which we will call target-MDP throughout the paper . The target-MDP is defined by the tuple ( S , A , T , r , γ ) . S and A are the state and action spaces , respectively , and T ( s′ | s , a ) is the transition density of state s′ given action a taken under state s. The reward function is denoted as r ( s , a ) , and γ ∈ ( 0 , 1 ) is the discount factor . The goal of RL is to find the optimal policy π∗ that maximizes the expected return , denoted by η : π∗ : = argmax π η [ π ] = argmax π Eπ [ ∞∑ t=0 γtr ( st , at ) ] , ( 1 ) where st+1 ∼ T ( s | st , at ) and at ∼ π ( a | st ) . Generally , the ground truth transition T is unknown , and MBRL methods aim to construct a model T̂ of the transition dynamics to help improve the policy . To evaluate a reinforcement learning method ( no matter whether it is model-free or model-based ) , we mainly care about two metrics : ( 1 ) asymptotic performance : the expected return of the converged policy in real environments ; ( 2 ) sample efficiency : the number of samples ( s , a , s′ , r ) collected in target-MDP to achieve some performance . Therefore , an effective RL method means it achieves higher expected returns with fewer real samples . 3.2 Model-Based Policy Optimization . Model-based policy optimization ( MBPO ) [ 12 ] is a state-of-the-art MBRL method which will be used as the running case in our framework . MBPO learns a bootstrapped ensemble of probabilistic dynamics models { T̂ 1θ , ... , T̂Eθ } where E is the ensemble size . Each individual dynamics model in the ensemble is a probabilistic neural network that outputs a Gaussian distribution with diagonal covariance T̂ iθ ( s ′|s , a ) = N ( µiθ ( s , a ) , Σ i θ ( s , a ) ) . These models are trained independently via maximum likelihood with different initializations and training batches . The corresponding loss function is LT̂ ( θ ) = N∑ n=1 [ µθ ( sn , an ) − sn+1 ] > Σ−1θ ( sn , an ) [ µθ ( sn , an ) − sn+1 ] + log det Σθ ( sn , an ) . In practical implementation , an early stopping trick is adopted in training the model ensemble . To be more specific , when training each individual model , a hold-out dataset will be created , and the training will early stop if the loss evaluated on the hold-out data does not decrease . Every time the ensemble models are trained , they are then used to generate k-length imaginary rollouts , which begin from states sampled from real data and follow the current policy . Each time the policy interacts with the target-MDP , it will be trained for G gradient updates using both real data and imaginary data in a certain ratio . To make it clear , in each batch for policy training , we denote the proportion of real data to the whole batch as β and the proportion of imaginary data as 1 − β . We call β the real ratio and call G the policy training iteration throughout the paper . The policy optimization algorithm is SAC ( Soft Actor-Critic ) [ 11 ] , and the loss functions for the actor πω and the critic Qφ are Lπ ( ω ) = Est∼D [ Eat∼πω [ α log ( πω ( at|st ) ) −Qφ ( st , at ) ] ] , ( 2 ) LQ ( φ ) = E ( st , at ) ∼D [ 1 2 ( Qφ ( st , at ) − ( r ( st , at ) ( 3 ) + γEst+1 , at+1 [ Qφ̂ ( st+1 , at+1 ) − α log πω ( at+1|st+1 ) ] ) ) 2 ] , where Qφ̂ is the target Q-function .

Summary: The paper is about learning a hyper-parameter setter for some of the hyper-parameters used to train a model-based reinforcement learning algorithm. Theoretical results are provided to show that there exists an optimal ratio of real-data to model-generated data during training. Experiments show that setting hyper-parameters with a parametric policy improves performance over manual hyper-parameter selection or hand-crafted schedules.