Model-Based Offline Meta-Reinforcement Learning with Regularization

1 INTRODUCTION . Offline reinforcement learning ( a.k.a. , batch RL ) has recently attracted extensive attention by learning from offline datasets previously collected via some behavior policy ( Kumar et al. , 2020 ) . However , the performance of existing offline RL methods could degrade significantly due to the following issues : 1 ) the possibly poor quality of offline datasets ( Levine et al. , 2020 ) and 2 ) the inability to generalize to different environments ( Li et al. , 2020b ) . To tackle these challenges , offline Meta-RL ( Li et al. , 2020a ; Dorfman & Tamar , 2020 ; Mitchell et al. , 2020 ; Li et al. , 2020b ) has emerged very recently by leveraging the knowledge of similar offline RL tasks ( Yu et al. , 2021a ) . The main aim of these studies is to enable quick policy adaptation for new offline tasks , by learning a meta-policy with robust task structure inference that captures the structural properties across training tasks . Because tasks are trained on offline datasets , value overestimation ( Fujimoto et al. , 2019 ) inevitably occurs in dynamic programming based offline Meta-RL , resulted from the distribution shift between the behavior policy and the learnt task-specific policy . To guarantee the learning performance on new offline tasks , a right balance has to be carefully calibrated between “ exploring ” the out-of-distribution state-actions by following the meta-policy , and “ exploiting ” the offline dataset by staying close to the behavior policy . However , such a unique “ exploration-exploitation ” tradeoff has not been considered in existing offline Meta-RL approaches , which would likely limit their ability to handle diverse offline datasets particularly towards those with good behavior policies . To illustrate this issue more concretely , we compare the performance between a state-of-the-art offline Meta-RL algorithm FOCAL ( Li et al. , 2020b ) and an offline single-task RL method COMBO ( Yu et al. , 2021b ) in two new offline tasks . As illustrated in Figure 1 , while FOCAL performs better than COMBO on the task with a bad-quality dataset ( left plot in Figure 1 ) , it is outperformed by COMBO on the task with a good-quality dataset ( right plot in Figure 1 ) . Clearly , existing offline Meta-RL fails in several standard environments ( see Figure 1 and Figure 11 ) to generalize universally well over datasets with varied quality . In order to fill such a substantial gap , we seek to answer the following key question in offline Meta-RL : How to design an efficient offline Meta-RL algorithm to strike the right balance between exploring with the meta-policy and exploiting the offline dataset ? To this end , we propose MerPO , a model-based offline Meta-RL approach with regularized Policy Optimization , which learns a meta-model for efficient task structure inference and an informative meta-policy for safe exploration of out-of-distribution state-actions . Compared to existing approaches , MerPO achieves : ( 1 ) safe policy improvement : performance improvement can be guaranteed for offline tasks regardless of the quality of the dataset , by strike the right balance between exploring with the meta-policy and exploiting the offline dataset ; and ( 2 ) better generalization capability : through a conservative utilization of the learnt model to generate synthetic data , MerPO aligns well with a recently emerging trend in supervised meta-learning to improve the generalization ability by augmenting the tasks with “ more data ” ( Rajendran et al. , 2020 ; Yao et al. , 2021 ) . Our main contributions can be summarized as follows : ( 1 ) Learnt dynamics models not only serve as a natural remedy for task structure inference in offline Meta-RL , but also facilitate better exploration of out-of-distribution state-actions by generating synthetic rollouts . With this insight , we develop a model-based approach , where an offline metamodel is learnt to enable efficient task model learning for each offline task . More importantly , we propose a meta-regularized model-based actor-critic method ( RAC ) for within-task policy optimization , where a novel regularized policy improvement module is devised to calibrate the unique “ exploration-exploitation ” tradeoff by using an interpolation between two regularizers , one based on the behavior policy and the other on the meta-policy . Intuitively , RAC generalizes COMBO to the multi-task setting , with introduction of a novel regularized policy improvement module to strike a right balance between the impacts of the meta-policy and the behavior policy . ( 2 ) We theoretically show that under mild conditions , the learnt task-specific policy based on MerPO offers safe performance improvement over both the behavior policy and the meta-policy with high probability . Our results also provide a guidance for the algorithm design in terms of how to appropriately select the weights in the interpolation , such that the performance improvement can be guaranteed for new offline RL tasks . ( 3 ) We conduct extensive experiments to evaluate the performance of MerPO . More specifically , the experiments clearly show the safe policy improvement offered in MerPO , corroborating our theoretical results . Further , the superior performance of MerPO over existing offline Meta-RL methods suggests that model-based approaches can be more beneficial in offline Meta-RL . 2 RELATED WORK . Offline single-task RL . Many existing model-free offline RL methods regularize the learnt policy to be close to the behavior policy by , e.g. , distributional matching ( Fujimoto et al. , 2019 ) , support matching ( Kumar et al. , 2019 ) , importance sampling ( Nachum et al. , 2019 ; Liu et al. , 2020 ) , learning lower bounds of true Q-values ( Kumar et al. , 2020 ) . Along a different avenue , model-based algorithms learn policies by leveraging a dynamics model obtained with the offline dataset . ( Matsushima et al. , 2020 ) directly constrains the learnt policy to the behavior policy as in model-free algorithms . To penalize the policy for visiting states where the learnt model is likely to be incorrect , MOPO ( Yu et al. , 2020 ) and MoREL ( Kidambi et al. , 2020 ) modify the learnt dynamics such that the value estimates are conservative when the model uncertainty is above a threshold . To remove the need of uncertainty quantification , COMBO ( Yu et al. , 2021b ) is proposed by combining model-based policy optimization ( Janner et al. , 2019 ) and conservative policy evaluation ( Kumar et al. , 2020 ) . Offline Meta-RL . A few very recent studies have explored the offline Meta-RL . Particularly , ( Li et al. , 2020a ) considers a special scenario where the task identity is spuriously inferred due to biased datasets , and applies the triplet loss to robustify the task inference with reward relabelling . ( Dorfman & Tamar , 2020 ) extends an online Meta-RL method VariBAD ( Zintgraf et al. , 2019 ) to the offline setup , and assumes known reward functions and shared dynamics across tasks . Based on MAML ( Finn et al. , 2017 ) , ( Mitchell et al. , 2020 ) proposes an offline Meta-RL algorithm with advantage weighting loss , and learns initializations for both the value function and the policy , where they consider the offline dataset in the format of full trajectories in order to evaluate the advantage . Based on the off-policy Meta-RL method PEARL ( Rakelly et al. , 2019 ) , ( Li et al. , 2020b ) combines the idea of deterministic context encoder and behavior regularization , under the assumption of deterministic MDP . Different from the above works , we study a more general offline Meta-RL problem . More importantly , MerPO strikes a right balance between exploring with the meta-policy and exploiting the offline dataset , which guarantees safe performance improvement for new offline tasks . 3 PRELIMINARIES . Consider a Markov decision process ( MDP ) M = ( S , A , T , r , µ0 , γ ) with state space S , action space A , the environment dynamics T ( s′|s , a ) , reward function r ( s , a ) , initial state distribution µ0 , and γ ∈ ( 0 , 1 ) is the discount factor . Without loss of generality , we assume that |r ( s , a ) | ≤ Rmax . Given a policy π , let dπM ( s ) : = ( 1 − γ ) ∑∞ t=0 γ tPM ( st = s|π ) denote the discounted marginal state distribution , where PM ( st = s|π ) denotes the probability of being in state s at time t by rolling out π in M. Accordingly , let dπM ( s , a ) : = dπM ( s ) π ( a|s ) denote the discounted marginal state-action distribution , and J ( M , π ) : = 11−γE ( s , a ) ∼dπM ( s , a ) [ r ( s , a ) ] denote the expected discounted return . The goal of RL is to find the optimal policy that maximizes J ( M , π ) . In offline RL , no interactions with the environment are allowed , and we only have access to a fixed dataset D = { ( s , a , r , s′ ) } collected by some unknown behavior policy πβ . Let d πβ M ( s ) be the discounted marginal state distribution of πβ . The datasetD is indeed sampled from d πβ M ( s , a ) = d πβ M ( s ) πβ ( a|s ) . DenoteM as the empirical MDP induced by D and d ( s , a ) as a sample-based version of dπβM ( s , a ) . In offline Meta-RL , consider a distribution of RL tasks p ( M ) as in standard Meta-RL ( Finn et al. , 2017 ; Rakelly et al. , 2019 ) , where each taskMn is an MDP , i.e. , Mn = ( S , A , Tn , rn , µ0 , n , γ ) , with task-shared state and action spaces , and unknown task-specific dynamics and reward function.For each task Mn , no interactions with the environment are allowed and we only have access to an offline dataset Dn , collected by some unknown behavior policy πβ , n . The main objective is to learn a meta-policy based on a set of offline training tasks { Mn } Nn=1 . Conservative Offline Model-Based Policy Optimization ( COMBO ) . Recent model-based offline RL algorithms , e.g. , COMBO ( Yu et al. , 2021b ) , have demonstrated promising performance on a single offline RL task by combining model-based policy optimization ( Janner et al. , 2019 ) and conservative policy evaluation ( CQL ( Kumar et al. , 2020 ) ) . Simply put , COMBO first trains a dynamics model T̂θ ( s′|s , a ) parameterized by θ , via supervised learning on the offline dataset D. The learnt MDP is constructed as M̂ : = ( S , A , T̂ , r , µ0 , γ ) . Then , the policy is learnt using D and model-generated rollouts . Specifically , define the action-value function ( Q-function ) as Qπ ( s , a ) : = E [ ∑∞ t=0 γ tr ( st , at ) |s0 = s , a0 = a ] , and the empirical Bellman operator as : B̂πQ ( s , a ) = r ( s , a ) + γE ( s , a , s′ ) ∼D [ Q ( s′ , a′ ) ] , for a′ ∼ π ( ·|s′ ) . To penalize the Q functions in out-of-distribution stateaction tuples , COMBO employs conservative policy evaluation based on CQL : Q̂k+1 ← arg min Q ( s , a ) β ( Es , a∼ρ [ Q ( s , a ) ] − Es , a∼D [ Q ( s , a ) ] ) + 1 2 Es , a , s′∼df [ ( Q ( s , a ) − B̂ πQ̂k ( s , a ) ) 2 ] ( 1 ) where ρ ( s , a ) : = dπ M̂ ( s ) π ( a|s ) is the discounted marginal distribution when rolling out π in M̂ , and df ( s , a ) = fd πβ M ( s , a ) + ( 1 − f ) ρ ( s , a ) for f ∈ [ 0 , 1 ] . The Bellman backup B̂π over df can be interpreted as an f -interpolation of the backup operators under the empirical MDP ( denoted by BπM ) and the learnt MDP ( denoted by B π M̂ ) . Given the Q-estimation Q̂π , the policy can be learnt by : π′ ← argmax π Es∼ρ ( s ) , a∼π ( ·|s ) [ Q̂π ( s , a ) ] . ( 2 )

Summary: This paper proposes a model-based meta-RL approach to offline learning over a distribution of MDPs (offline meta-RL). The proposed method is inspired by COMBO and in fact is a direct extension of it to the multi-task setting. The extension involves meta-learning a model-initialisation (or prior) that can be rapidly adapted to new offline RL tasks. They further add a proximal RL policy-improvement operator, where the prior policy is meta-learned over the task distribution. They establish theoretical guarantees for their method and conduct a careful empirical investigation to motivate its design, while establishing that it provides additional benefits compared to relevant baselines.