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 considers the challenge in offline meta-RL problems, particularly, the difficulty with balancing a learned policy between exploring OOD state-actions by following the meta-policy and exploiting the offline dataset. To solve the problem, it starts with proximal meta-RL approach in the online setting, points out the problem of a degraded meta-policy in the offline setting due to lack of feedback from online interactions and proposes to regularize the task policy with additional penalty from deviating from the behaviour policy. It also integrates the model-based offline RL method for single tasks (COMBO) and offline meta-learning for task specific dynamics into its final form (MerPO). This paper provides theoretical analysis on its advantage over behaviour policy and the meta policy. Extensive experiments with ablation show the empirical performance meets its expected behaviour.