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 present a nearest neighbour method for learning a model offline from the statistics of the given data set. Representations are provided from other off-policy deep RL methods. Value iteration is used on top of the model to learn the final policy. The algorithm is tested in several Atari games over two data set sizes.