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 deals with the problem of hyper-parameters optimization in model-based reinforcement learning (MBRL) Specifically, the paper tries to answer the question of how to schedule hyperparameters such as real data ratio, model training frequency, policy training iteration, and rollout length to achieve better performance. Supplementals materials are provided.   There are 2 contributions in the paper: - The paper provides some analysis to show that real data ratios affect the performance of MBRL. Specifically, they show that increasing the real data ratio improves the performance of MBRL. - They propose a new MBRL algorithm that learns to schedule hyper-parameters on the fly.