DeepAveragers: Offline Reinforcement Learning By Solving Derived Non-Parametric MDPs

1 INTRODUCTION . Research in automated planning and control has produced powerful algorithms to solve for optimal , or near-optimal , decisions given accurate environment models . Examples include the classic valueand policy-iteration algorithms for tabular representations or more sophisticated symbolic variants for graphical model representations ( e.g . Boutilier et al . ( 2000 ) ; Raghavan et al . ( 2012 ) ) . In concept , these planners address many of the traditional challenges in reinforcement learning ( RL ) . They can perform “ zero-shot transfer ” to new goals and changes to the environment model , accurately account for sparse reward or low-probability events , and solve for different optimization objectives ( e.g . robustness ) . Effectively leveraging these planners , however , requires an accurate model grounded in observations and expressed in the planner ’ s representation . On the other hand , model-based reinforcement learning ( MBRL ) aims to learn grounded models to improve RL ’ s data efficiency . Despite developing grounded environment models , the vast majority of current MBRL approaches do not leverage near-optimal planners to help address the above challenges . Rather , the models are used as black-box simulators for experience augmentation and/or Monte-Carlo search . Alternatively , model learning is sometimes treated as purely an auxiliary task to support representation learning . The high-level goal of this paper is to move toward MBRL approaches that can effectively leverage near-optimal planners for improved data efficiency and flexibility in complex environments . However , there are at least two significant challenges . First , there is a mismatch between the deep model representations typically learned in MBRL ( e.g . continuous state mappings ) and the representations assumed by many planners ( e.g . discrete tables or graphical models ) . Second , near-optimal planners are well-known for exploiting model inaccuracies in ways that hurt performance in the real environment , e.g . ( Atkeson , 1998 ) . This second challenge is particularly significant for offline RL , where the training experience for model learning is fixed and limited . We address the first challenge above by focusing on tabular representations , which are perhaps the simplest , but most universal representation for optimal planning . Our main contribution is an offline MBRL approach based on optimally solving a new model called the Deep Averagers with Costs MDP ( DAC-MDP ) . A DAC-MDP is a non-parametric model derived from an experience dataset and a corresponding ( possibly learned ) latent state representation . While the DAC-MDP is defined over the entire continuous latent state space , its full optimal policy can be computed by solving a standard ( finite ) tabular MDP derived from the dataset . This supports optimal planning via any tabular MDP solver , e.g . value iteration . To scale this approach to typical offline RL problems , we develop a simple GPU implementation of value iteration that scales to millions of states . As an additional engineering contribution , this implementation will be made public . To address the second challenge of model inaccuracy due to limited data , DAC-MDPs follow the pessimism in the face of uncertainty principle , which has been shown effective in a number of prior contexts ( e.g . ( Fonteneau et al. , 2013 ) ) . In particular , DAC-MDPs extend Gordon ’ s Averagers framework ( Gordon , 1995 ) with additional costs for exploiting transitions that are under-represented in the data . Our second contribution is to give a theoretical analysis of this model , which provides conditions under which a DAC-MDP solution will perform near optimally in the real environment . Our final contribution is to empirically investigate the DAC-MDP approach using simple latent representations derived from random projections and those learned by Q-iteration algorithms . Among other results , we demonstrate the ability to scale to Atari-scale problems , which is the first demonstration of optimal planning being effectively applied across multiple Atari games . In addition , we provide case studies in 3D first-person navigation that demonstrate the flexibility and adaptability afforded by integrating optimal planning into offline MBRL . These results show the promise of our approach for marrying advances in representation learning with optimal planning . 2 FORMAL PRELIMINARIES . A Markov Decision Process ( MDP ) is a tuple 〈S , A , T , R〉 ( Puterman , 1994 ) , with state set S , action set A , transition function T ( s , a , s′ ) , and reward function R ( s , a ) . A policy π maps states to actions and has Q-function Qπ ( s , a ) giving the expected infinite-horizon β-discounted reward of following π after taking action a in s. The optimal policy π∗ maximizes the Q-function over all policies and state-action pairs . Q∗ corresponds to the optimal Q-function that satisfies π∗ ( s ) = arg maxaQ ∗ ( s , a ) . Q∗ can be computed given the MDP by repeated application of the Bellman Backup Operator B , which for any Q-function Q , returns a new Q-function given by , B [ Q ] ( s , a ) = R ( s , a ) + γEs′∼T ( s , a , · ) [ max a Q ( s′ , a ) ] . ( 1 ) The objective of RL is to find a near-optimal policy without prior knowledge of the MDP . In the online RL setting , this is done by actively exploring actions in the environment . Rather , in the offline RL ( Levine et al. , 2020 ) , which is the focus of this paper , learning is based on a static dataset D = { ( si , ai , ri , s′i ) } , where each tuple gives the reward ri and next state s′i observed after taking action ai in state si . In strict offline RL setting , the final policy selection must be done using only the dataset , without direct access to the environment . This includes all hyperparameter tuning and the choice of when to stop learning . Evaluations of offline RL , however , often blur this distinction , for example , reporting performance of the best policy obtained across various hyperparameter settings as evaluated via new online experiences ( Gulcehre et al. , 2020 ) . Here we consider an evaluation protocol that makes the amount of online access to the environment explicit . In particular , the offline RL algorithm is allowed to use the environment to evaluate Ne policies ( e.g . an average over repeated trials for each policy ) , which , for example , may derive from different hyperparameter choices . The best of the evaluated policies can then be selected . Note that Ne = 1 corresponds to pure offline RL . 3 DEEP AVERAGERS WITH COSTS MDPS ( DAC-MDPS ) . From a practical perspective our approach carries out the following steps as illustrated in Figure 1 . ( 1 ) We start with a static experience dataset , where the states are assumed to come from a continuous latent state space . For example , states encoded via random or learned deep representations . ( 2 ) Next we compile the dataset into a tabular MDP over the “ core states ” of the DAC-MDP ( those in the dataset ) . This compilation uses k-nearest neighbor ( kNN ) queries to define the reward and transition functions ( Equation 2 ) functions of the core states . ( 3 ) Next we use a GPU implementation of value iteration to solve for the tabular MDP ’ s optimal Q-function . ( 4 ) Finally , this tabular Q-function is used to define the Q-function over the entire DAC-MDP ( Equation 3 ) . Previously unseen states at test time are assigned Q-values and in turn a policy action via kNN queries over the core states . Conceptually , our DAC-MDP model is inspired by Gordon ’ s ( 1995 ) early work that showed the convergence of ( offline ) approximate value iteration for a class of function approximators ; averagers , which includes methods such as k nearest neighbor ( kNN ) regression , among others . It was also observed that approximate value iteration using an averager was equivalent to solving an MDP derived from the offline data . That observation , however , was not investigated experimentally and has yet to be integrated with deep representation learning . Here we develop and evaluate such an integration . The quality of an averagers MDP , and model-learning in general , depends on the size and distribution of the dataset . In particular , an optimal planner can exploit inaccuracies in the underrepresented parts of the state-action space , which can lead to poor performance in the real environment . The DACMDPs aim to avoid this by augmenting the derived MDPs with costs/penalties on under-represented transitions . This turns out to be essential to achieving good performance on challenging benchmarks . 3.1 DAC-MDP DEFINITION . A DAC-MDP is defined in terms of an experience dataset D = { ( si , ai , ri , s′i ) } from the true MDP M with continuous latent state space S and finite action space A . The DAC-MDP M̃ = ( S , A , R̃ , T̃ ) shares the same state and action spaces as M , but defines the reward and transition functions in terms of empirical averages over the k nearest neighbors of ( s , a ) in D. The distance metric d ( s , a , s′ , a′ ) gives the distance between pairs ( s , a ) and ( s′ , a′ ) . This metric considers ( s , a ) pairs with different actions to be infinitely distant . Otherwise , the distance between pairs involving the same action is the euclidean distance between their states . In particular , the distance between ( s , a ) and a data tuple ( si , ai , ri , s′i ) is given by d ( s , a , si , ai ) . Also , we let kNN ( s , a ) denote the set of indices of the k nearest neighbors to ( s , a ) in D , noting that the dependence on D and d is left implicit . Given hyperparameters k ( smoothing factor ) and C ( cost factor ) we can now specify the DAC-MDP reward and transition function . R̃ ( s , a ) = 1 k ∑ i∈kNN ( s , a ) ri − C · d ( s , a , si , ai ) , T̃ ( s , a , s′ ) = 1 k ∑ i∈kNN ( s , a ) I [ s′ = s′i ] ( 2 ) The reward for ( s , a ) is simply the average reward of the nearest neighbors with a penalty for each neighbor that grows linearly with the distance to a neighbor . Thus , the farther ( s , a ) is to its nearest neighbor set , the less desirable its immediate reward will be . The transition function is simply the empirical distribution over destination states of the nearest neighbor set . Importantly , even though a DAC-MDP has an infinite continuous state space , it has a special finite structure . Since the transition function T̃ only allows transitions to states appearing as destination states in D. We can view M̃ as having a finite core set of states SD = { s′i | ( si , ai , ri , s′i ) ∈ D } . States in this core do not transition to non-core states and each non-core state immediately transitions to the core for any action . Hence , the value of core states is not influenced by non-core states . Further , once the core values are known , we can compute the values of any non-core state via onestep look ahead using T̃ . Thus , we can optimally solve a DAC-MDP by solving just its finite core . Specifically , let Q̃ be the optimal Q-function of M̃ . We can compute Q̃ for the core states by solving the finite MDP M̃D = ( SD , A , R̃ , T̃ ) . We can then compute Q̃ for any non-core state on demand via the following one-step look-ahead expression.1 This allows us to compute the optimal policy of M̃ , denoted π̃ , using any solver of finite MDPs . Q̃ ( s , a ) = 1 k ∑ i∈kNN ( s , a ) ri + γmax a Q̃ ( s′i , a ) − C · d ( s , a , si , ai ) ( 3 )

Summary: The authors propose to learn a non-parametric MDP model from batch data, which can be solved efficiently using discrete value iteration (by solving for the “core” states which are all the end-states in observed transitions) and which provides a Q-value defined over the full continuous space through a kNN lookup. There is an interesting penalty term when the estimate relies on far-away support data. The value of the optimal policy in the approximated MDP can be bounded, under some smoothness assumptions, in the original MDP.