Offline Model-based Adaptable Policy Learning

1 Introduction . Recent studies have shown that reinforcement learning ( RL ) is a promising approach for real-world applications , e.g. , sequential recommendation systems [ 1 , 2 , 3 , 4 ] and robotic locomotion skill learning [ 5 , 6 ] . However , the trial-and-error of RL in the real world [ 7 ] obstructs further applications in cost-sensitive scenarios [ 8 ] . Offline ( batch ) RL learns a policy within a static dataset collected by a behavior policy without additional interactions with the environment [ 8 , 9 , 10 , 11 ] . Since it avoids costly trial-and-error in real-world environments , offline RL is a promising way to handle the challenge in cost-sensitive applications . A significant challenge of offline RL is in answering counterfactual queries , which asks about how the performance ( e.g. , Q value ) would have been if the agent were to execute an unseen action sequence , then learning to make optimal decisions based on the performance [ 8 ] . Fujimoto et al . [ 10 ] have shown that the distributional shift of states and actions , which comes from the discrepancy between evaluated policies and behavior policies , often leads to large extrapolation error in value function estimation . In traditional model-free algorithms , the extrapolation error in value function estimation hurts the generalization performance of the learned policies . Since the additional samples , which can correct value estimation errors , are unavailable in the offline setting , the performance of learned policies based on value function is unstable [ 10 ] . On the other hand , model-based RL techniques , which learn dynamics models from collected datasets and learn the value function and policies based on the dynamics models , do not need to estimate ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the value functions rely on the collected datasets . However , similar challenges occur in dynamics model approximation . The dynamics model might overfit the limited dataset and suffer extrapolation errors in regions that behavior policies have not visited , which causes instability of the learned policy when deployment [ 12 ] . Here we call it out-of-support regions . Moreover , in model inference , the compounding error , that is , the accumulated prediction errors between simulation trajectories and reality , would be large even if the one-step prediction error is small [ 13 , 14 ] . Recent studies in offline model-based RL [ 12 , 15 ] have made significant progress in MuJoCo tasks [ 16 ] . These methods constrain policy sampling in dynamics models for robust policy learning . By using large penalty [ 12 ] or trajectory truncation [ 15 , 12 ] in the regions with large prediction uncertainty ( uncertainty is a designed metric to evaluate the confidence of prediction correctness ) or compounding error , policy exploration is constrained in the regions of dynamics models where the predictions are corrected with high confidence , so as to avoid exploiting regions with risks of large extrapolation error . However , the constraints on dynamics models lead to a conservative policy learning process , which limits the potential of leveraging dynamics models : The visits on states and actions in out-of-support regions are more likely to be inhibited by the constraints , making the learned policy restrict the agent to be in similar regions as the behavior policy . From the perspectives of counterfactual queries , we consider that model-based RL is promising to handle offline RL — ideal reconstructed dynamics models can simulate the transition dataset without the distributional-shift problem given any policy , and the value function can be estimated via the “ simulated ” transition dataset directly . The bottleneck of offline model-based RL comes from the policy learning in the approximated dynamics model with extrapolation error . In this paper , instead of learning by tightly constraining policy exploration in in-support regions , we investigate decision-making in out-of-support regions directly . Finally , we propose a new offline policy learning framework , offline Model-based Adaptable Policy LEarning ( MAPLE ) , to address the aforementioned issues . Ideally , MAPLE tries to model all possible transition dynamics in the out-of-support regions . Then an Adaptable policy is learned to be aware of each case to adapt its behavior to reach optimal performance . In the practical version of MAPLE , we use an ensemble technique to construct ensemble dynamics models . To be aware of each case of the transition dynamics and learn an adaptable policy , we use a meta-learning technique that introduces an extra environment-context extractor structure to represent dynamics patterns , and the policy adjusts itself according to the environment contexts . We conduct experiments on the MuJoCo tasks . The results show that the sampling regions for robust offline policy learning can be extended via constructing transition patterns in out-of-support regions to cover the real case . The output adaptable policy yields better performance than SOTA algorithms when deployed . MAPLE gives a new direction to handle the offline policy learning problem in the dynamics models : Besides constraining on sampling and training dynamics models with better generalization , we can also model out-of-distribution regions by constructing all possible transition patterns . 2 Related Work . Reinforcement learning ( RL ) has shown to be a promising approach to complex real-world decisionmaking problems [ 1 , 2 , 3 , 4 ] . However , unconstrained online trial-and-error in the training of RL agents prevents further applications of RL in safety-critical scenarios since it might result in large economic losses [ 8 , 17 , 18 , 19 ] . Many studies propose to overcome the problem by offline ( batch ) RL algorithms [ 20 ] . Prior works on offline RL are based on model-free algorithms . To overcome the extrapolation error , which is introduced by the discrepancy between the offline dataset and true state-action distribution of learned target policies [ 10 ] , these methods are designed to constrain target policies to be close to the behavior policies [ 10 , 11 , 21 ] , to apply ensemble methods for robust value function estimation [ 22 ] , or to re-weight samples in datasets with importance sampling [ 23 ] . Most recent studies have shown that policy learning with an approximated dynamics model has good potential to take robust actions outside the action distribution of behavior policies [ 15 , 12 ] . The challenge comes from the extrapolation error of the dynamics models in out-of-support regions . To address the issues , these methods learn policies from dynamics models with uncertainty constraints . Uncertainty is a measure of prediction confidence on next states . The uncertainty is often computed by the inconsistency in the ensemble dynamics model predictions for each state-action pair . Kidambi et al . [ 15 ] construct terminating states based on a hard threshold on uncertainty , while Yu et al . [ 12 ] use a soft reward penalty to incorporate uncertainty . The penalty constrains policy exploration and optimization to the regions with high consistency for better worst-case performance in the deployment environment [ 15 , 12 ] . The difference between the aforementioned model-based methods and MAPLE is shown in Figure 1 . Compared with previous methods [ 15 , 12 ] learned by constraining ( in Figure 1 ( b ) ) , we learn to adapt to all possible dynamics transitions in the states in out-of-support regions ( in Figure 1 ( a ) ) . 3 Background and Notation . In the standard RL framework , an agent interacts with an environment governed by a Markov Decision Process ( MDP ) [ 24 ] . The agent learns a policy π ( at|st ) , which chooses an action at ∈ A given a particular state st ∈ S , at each time-step t ∈ { 0 , 1 , ... , T } , where T is the trajectory length . S and A denote the state and action spaces , respectively . The reward function rt = r ( st , at ) ∈ R evaluates the immediate performance of the action at given the state st . The goal of RL is to find an optimal policy π∗ which maximizes the multi-step cumulative discounted reward ( i.e. , long-term performance ) . The objective of RL is maxπ Jρ ( π ) : = Eτ∼p ( τ |π , ρ ) [ ∑T k=0 γ krk ] , where γ is the discount factor , and p ( τ |π , ρ ) is the probability of generating a trajectory τ : = [ s0 , a0 , ... , aT−1 , sT ] under the policy π and a dynamics model ρ ( st+1|st , at ) . In particular , p ( τ | π ) : = d0 ( s0 ) ∏T−1 t=0 ρ ( st+1 | st , at ) π ( at|st ) , where d0 ( s0 ) is the initial state distribution . A common way to find an optimal policy π∗ is to optimize the policy with gradient ascent along∇Jρ ( π ) [ 24 , 25 ] . In the offline RL setting , we are given only a static dataset D = { ( si , ai , ri , si+1 ) } collected by some unknown policy . The goal is to obtain a policy that maximizes Jρ by only using the static dataset . 4 Method . We argue that in current offline model-based methods , the constraints of sampling to in-support regions of the dynamics model lead to a conservative policy learning process , limiting the potential of leveraging dynamics models . In this paper , to relax the constraints on the dynamics model , we investigate decision-making in out-of-support regions directly . In this section , we first give a motivating example to show our ideal solution to out-of-support region decision-making ( in Section 4.1 ) . Then , we introduce a practical algorithm of the proposed solution for complex tasks , based on meta-learning techniques ( in Section 4.2 and Section 4.3 ) . 4.1 Decision-Making in Out-of-Support Regions . By rethinking the scheme of offline model-based RL , without loss of generality , we first formulate the problem as decision-making with a partially known dynamics model ( Pak-DM ) in a surrogate objective . In this problem , we have two dynamics models : a target dynamics model ρ and a partially known dynamics model ρ′ , where ρ is the deployment environment in the offline RL setting , and ρ′ is used to approximate ρ . Due to the bias of data sampling in the offline setting and the limitation on the capacity of the function approximator , only in part of state-action space , we have ρ′ ( s′|s , a ) = ρ ( s′|s , a ) , while the transitions in other parts of space are uncertain . We call the satisfied space “ accessible space ” ( a.k.a. , in-support regions ) and its complement “ inaccessible space ” ( a.k.a. , out-of-support regions ) . In this problem , we assume the two subspaces have been predefined in some ways as in the offline setting ( e.g. , we can define the space in which an uncertainty quantification is larger than a threshold as the inaccessible space ) and then discuss the decision-making problem in this setting . Formally , given an accessible space Xa and an inaccessible space Xi , the partially known dynamics model is defined as : ρ′ ( s′|s , a ) = { ρ ( s′|s , a ) [ s , a ] ∈ Xa Unknown [ s , a ] ∈ Xi , where X denotes the state-action concatenated space for brevity and [ s , a ] denote a vector concatenating s and a . Our objective is to find a policy π∗ to maximize Jρ by only querying the partially known dynamics model ρ′ . If Xi = ∅ , that is ρ′ ( s′|s , a ) = ρ ( s′|s , a ) , ∀s ∈ S , a ∈ A , the problem is reduced to a vanilla model-based policy learning problem with an oracle dynamics model . For simplification , we assume the oracle reward function r is given , but it can also be formulated as a partially known reward function in a similar way . Now we give an example of a Pak-DM in Figure 2 . In this setting , the model-based policy learning algorithms with constraints ( [ 12 , 15 ] ) can be summarized as : finding the optimal policy without reaching inaccessible space . It might output a conservative policy since it avoids making decisions that might lead agents to out-of-support regions . Taking Figure 2 as an example , the output policy would be run in the loop of A → B → α → A . It will avoid taking β on state B . If the real transition is ρ ( B , β ) = C2 , the policy can avoid the penalty −20 since C2 will not be reached . On the other hand , while ρ ( B , β ) = C1 , the policy would miss the large bonus 100 in C1 . Our question is , how should we make good decisions in out-of-support regions directly so that we can make better use of the approximated ρ′ for better performance ? We give a new paradigm to solve the problem which is the ideal implementation version of MAPLE : 1 . ( Training ) Construct a dynamics model set { ρ̂i } via modeling all possible transitions in Xi . ( It is impractical to do that with infinite state space . We will give a practical solution in Section 4.3 ) ; 2 . ( Training ) Learn the optimal policy from each model ρ̂i to form the optimal policy set { π∗ρ̂i } ; 3 . ( Deployment ) Initialize a state s0 from the deployment environment ρ ; 4 . ( Deployment ) Probe the environment ρ by selecting an action a such that [ s , a ] ∈ Xi . After getting the next state s′ = ρ ( s′|s , a ) , store the tuple ( s , a , s′ ) in a memory ( e.g. , a replay buffer ) D. If there is no action a that allows [ s , a ] ∈ Xi , randomly select a policy from the policy set to take an action . 5 . ( Deployment ) Reduce the policy set by only keeping the policies whose corresponding transition model ρ̂i can explain the experiences in the memory : { ρ̂i } ← { ρ | ρ ( s′|s , a ) = s′ , ∀ ( s , a , s′ ) ∈ D , ∀ρ ∈ { ρ̂i } } and { π∗ρ̂i } ← { π ∗ ρ|ρ ∈ { ρ̂i } } ; 6 . ( Deployment ) Repeat Step 4 and 5 until the policy set is reduced to a single policy . In this paradigm , we solve the decision-making problem in out-of-support regions by probing the uncertainty part of the deployment environment and adapting the policy for the environment . In Figure 2 , the ideal MAPLE solution would construct two dynamics models : ρ̂1 where ρ̂1 ( B , β ) = C1 , and ρ̂2 where ρ̂2 ( B , β ) = C2 . Then we learn two optimal policies { π∗ρ1 , π ∗ ρ2 } for each dynamics model . At deployment , we first randomly select a policy from the policy set to make decisions . Upon reaching B for the first time , where the transition on action β is uncertain , we take action β and get the next state . If the next state is C1 , the policy will reduce to π∗ρ̂1 , otherwise to π ∗ ρ̂2 . Therefore , if ρ ( B , β ) = C2 , the policy would initially run A → B → β → C2 → D → A and then run in the loop of A → B → α → A because the latter yields higher rewards . If ρ ( B , β ) = C1 , the policy would always run in the loop of A→ B → β → C1 → D → A . We then dive into the performance difference of the two paradigms for decision-making . Formally , we give Theorem 1 to describe it . The full proof can be found in Appendix A. Theorem 1 Given a target dynamics model ρ , a policy πa learned by adapting , a policy πc learned by constraints , and the maximum step Nm taken by πa to probe and reduce the policy set to a single policy , the performance gap between πa and πc satisfies : Jρ ( πa ) − Jρ ( πc ) ≥ ∆c −∆p − γNm+1Jρ∆ ( π∗ ) , where ∆c denotes the performance gap of the optimal policy π∗ and πc , while ∆p denotes the performance degradation of MAPLE compared with π∗ because of the phase of probing . Jρ∆ ( π ∗ ) denotes the performance degradation of π∗ on the dynamics model ρ caused by different initial state distribution : Jρ∆ ( π ∗ ) = Edπ ? Nm+1 ( s ) [ V ? ( s ) ] − EdπaNm+1 ( s ) [ V ? ( s ) ] , where dπ ? Nm+1 ( s ) and dπaNm+1 ( s ) denote the state distribution induced by π ? and πa at theNm+1 step and V ∗ ( s ) denotes the expected long-term rewards of π∗ at state s. We can see that the performance gain of πa is that it can automatically converge to the optimal policy after the loop of probing and reducing ( i.e. , ∆c ) , while the cost of πa comes from additional probing on inaccessible space , including less reward getting when probing ( i.e. , ∆p ) and a worse initial state distribution after probing ( i.e. , Jρ∆ ( π ∗ ) ) . Based on Theorem 1 , we give the principles for choosing between the paradigms : Firstly , with a larger performance gap of ∆c , πa can reach a better performance than πc . On the other hand , the tasks with large penalties on undesired behavior might make ∆p larger , which reduces the overall performance of πa . Besides , the tasks where sub-optimal behaviors easily lead agents to states with low value , e.g. , unsafe states which are prone to terminate the trajectory , might make Jρ∆ ( π ∗ ) large , which also reduces the overall performance of πa .

Summary: This paper presents Model-based Adaptable Policy LEearning (MAPLE), a model-based offline RL algorithm that learns a policy that can quickly adapt to an unknown environment during execution. Existing offline RL algorithms constrain the policy to visit only in-support regions, but MAPLE aims to model all possible dynamics models for out-of-support regions and try to learn a dynamic-conditioned (context-conditioned) policy. Policies for various uncertain dynamics models are prepared first, and gradually filter out policies that do not correspond well to transition observations when deployed. However, solving it exactly is intractable, and a practical approximation is presented, which uses an ensemble dynamics model set and an RNN-based environment-context extractor. Experimental results show that MAPLE outperforms baselines in D4RL benchmarks, showing better generalization ability.