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: Conservatism, i.e. avoiding out-of-distribution actions (and states), is a major theme in offline RL because of the need to avoid erroneous extrapolation beyond the data. This paper takes an alternative approach and instead trains a policy that adapts when deployed beyond the data distribution. To do this, a context encoder RNN is trained to produce latent codes given the episode history, and the encoder and policy are jointly optimized to maximize average performance across a large ensemble of pretained dynamics models. The algorithm, offline Model-based Adaptable Policy LEarning (MAPLE), is shown to achieve strong results on the MuJoCo tasks from the D4RL benchmark.