Representation Balancing Offline Model-based Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has accomplished remarkable results in a wide range of domains , but its successes were mostly based on a large number of online interactions with the environment . However , in many real-world tasks , exploratory online interactions are either very expensive or dangerous ( e.g . robotics , autonomous driving , and healthcare ) , and applying a standard online RL would be impractical . Consequently , the ability to optimize RL agents reliably without online interactions has been considered as a key to practical deployment , which is the main goal of batch RL , also known as offline RL ( Fujimoto et al. , 2019 ; Levine et al. , 2020 ) . In an offline RL algorithm , accurate policy evaluation and reliable policy improvement are both crucial for the successful training of the agent . Evaluating policies in offline RL is essentially an off-policy evaluation ( OPE ) task , which aims to evaluate the target policy given the dataset collected from the behavior policy . The difference between the target and the behavior policies causes a distribution shift in the estimation , which needs to be adequately addressed for accurate policy evaluation . OPE itself is one of the long-standing hard problems in RL ( Sutton et al. , 1998 ; 2009 ; Thomas & Brunskill , 2016 ; Hallak & Mannor , 2017 ) . However , recent offline RL studies mainly focus on how to improve the policy conservatively while using a common policy evaluation technique without much considerations for the distribution shift , e.g . mean squared temporal difference error minimization or maximum-likelihood training of environment model ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Yu et al. , 2020 ) . While conservative policy improvement helps the policy evaluation by reducing the off-policyness , we hypothesize that addressing the distribution shift explicitly during the policy evaluation can further improve the overall performance , since it can provide a better foundation for policy improvement . To this end , we aim to explicitly address the distribution shift of the OPE estimator used in the offline RL algorithm . In particular , we focus on the model-based approach , where we train an environment model robust to the distribution shift . One of the notable prior works is Representation Balancing MDP ( RepBM ) ( Liu et al. , 2018b ) , which regularizes the representation learning of the model to be invariant between the distributions . However , despite the promising results by RepBM , its step-wise estimation of the distance between the distributions has a few drawbacks that limit the algorithm from being practical : not only it assumes a discrete-action task where the target policy is deterministic , but it also performs poorly in long-term tasks due to the curse of horizon of step-wise importance sampling ( IS ) estimators ( Liu et al. , 2018a ) . To address these limitations , we present the Representation Balancing with Stationary Distribution Estimation ( RepB-SDE ) framework , where we aim to learn a balanced representation by regularizing , in the representation space , the distance between the data distribution and the discounted stationary distribution induced by the target policy . Motivated by the recent advances in estimating stationary distribution corrections , we present a new representation balancing objective to train a model of the environment that no longer suffers from the curse of horizon . We empirically show that the model trained by the RepB-SDE objective is robust to the distribution shift for the OPE task , particularly when the difference between the target and the behavior is large . We also introduce a model-based offline RL algorithm based on the RepB-SDE framework and report its performance on the D4RL benchmark ( Fu et al. , 2020 ) , showing the state-of-the-art performance in a representative set of tasks . 2 RELATED WORK . Learning balanced representation Learning a representation invariant to specific aspects of data is an established method for overcoming distribution shift that arises in unsupervised domain adaptation ( Ben-David et al. , 2007 ; Zemel et al. , 2013 ) and in causal inference from observational data ( Shalit et al. , 2017 ; Johansson et al. , 2018 ) . They have shown that imposing a bound on the generalization error under the distribution shift leads to the objective that learns a balanced representation such that the training and the test distributions look similar . RepBM ( Liu et al. , 2018b ) can be seen as a direct extension to the sequential case , which encourages the representation to be invariant under the target and behavior policies in each timestep . Stationary distribution correction estimation ( DICE ) Step-wise importance sampling ( IS ) estimators ( Precup , 2000 ) compute importance weights by taking the product of per-step distribution ratios . Consequently , these methods suffer from exponentially high variance in the lengths of trajectories , which is a phenomenon called the curse of horizon ( Liu et al. , 2018a ) . Recently , techniques of computing a stationary DIstribution Correction Estimation ( DICE ) have made remarkable progress that effectively addresses the curse of horizon ( Liu et al. , 2018a ; Nachum et al. , 2019a ; Tang et al. , 2020 ; Zhang et al. , 2020 ; Mousavi et al. , 2020 ) . DICE has been also used to explicitly address the distribution shift in online model-free RL , by directly applying IS on the policy and action-value objectives ( Liu et al. , 2019 ; Gelada & Bellemare , 2019 ) . We adopt one of the estimation techniques , DualDICE ( Nachum et al. , 2019a ) , to measure the distance between the stationary distribution and the data distribution in the representation space . Offline reinforcement learning There are extensive studies on improving standard online modelfree RL algorithms ( Mnih et al. , 2015 ; Lillicrap et al. , 2016 ; Haarnoja et al. , 2018 ) for stable learning in the offline setting . The main idea behind them is to conservatively improve policy by ( 1 ) quantifying the uncertainty of value function estimate , e.g . using bootstrapped ensembles ( Kumar et al. , 2019 ; Agarwal et al. , 2020 ) , or/and ( 2 ) constraining the optimized target policy to be close to the behavior policy ( i.e . behavior regularization approaches ) ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Lee et al. , 2020 ) . A notable exception is AlgaeDICE ( Nachum et al. , 2019b ) , which implicitly uses DICE to regularize the discounted stationary distribution induced by the target policy to be kept inside of the data support , similar to this work . On the other hand , Yu et al . ( 2020 ) argued that the model-based approach can be advantageous due to its ability to generalize predictions on the states outside of the data support . They introduce MOPO ( Yu et al. , 2020 ) , which uses truncated rollouts and penalized rewards for conservative policy improvement . MOReL ( Kidambi et al. , 2020 ) trains a state-action novelty detector and use it to penalize rewards in the data-sparse region . Matsushima et al . ( 2020 ) , MOOSE ( Swazinna et al. , 2020 ) and MBOP ( Argenson & Dulac-Arnold , 2020 ) guide their policy optimization using the behavior policy , similar to the behavior regularization approaches . Note that these aforementioned offline RL methods build on the standard approximate dynamic programming algorithm for action-value estimation ( model-free ) or on a maximum-likelihood environment model ( model-based ) , without explicitly addressing the distribution shift in the estimator . In contrast , we augment the objective for model learning to obtain a robust model under the distribution shift , which is the first attempt for offline RL to the best of our knowledge . 3 PRELIMINARIES . A Markov Decision Process ( MDP ) is specified by a tuple M = 〈S , A , T , R , d0 , γ〉 , consisting of state space S , action space A , transition function T : S × A → ∆ ( S ) , reward function R : S × A → ∆ ( [ 0 , rmax ] ) , initial state distribution d0 , and discount rate γ . In this paper , we mainly focus on continuous state space S ⊆ Rds and conduct experiments on both discrete action spaces A = { a0 , ... ana } and continuous action spaces A ⊆ Rda . Given MDP M and policy π , which is a ( stochastic ) mapping from state to action , the trajectory can be generated in the form of s0 , a0 , r0 , s1 , a1 , r1 , ... , where s0 ∼ d0 and for each timestep t ≥ 0 , at ∼ π ( st ) , rt ∼ R ( st , at ) , and st+1 ∼ T ( st , at ) . The goal of RL is to optimize or evaluate a policy , based on the normalized expected discounted return : Rπ , ( 1− γ ) EM , π [ ∑∞ t=0 γ trt ] . A useful and important concept throughout the paper is the discounted stationary distribution , which represents the long-term occupancy of states : dπ ( s , a ) , ( 1− γ ) ∞∑ t=0 γt Pr ( st = s , at = a|M , π ) . From the definition , it can be observed that Rπ can be obtained by Rπ = E ( s , a ) ∼dπ [ r ( s , a ) ] . Offline RL and off-policy evaluation In this paper , we focus on the offline RL problem where the agent can only access a static dataset D = { ( si , ai , ri , s′i ) } Ni=1 for the maximization of Rπ . We consider a behavior-agnostic setting where we do not have any knowledge of the data collection process . We denote the empirical distribution of the dataset by dD . Before improving policy , we first aim to better evaluateRπ given a target policy π and a static dataset D , which corresponds to an off-policy evaluation ( OPE ) problem . We mainly focus on a modelbased approach where the algorithm first estimates the unknown dynamics T̂ , R̂ using the dataset D. This defines an approximate MDP M̂ = 〈S , A , T̂ , R̂ , d0 , γ〉 , with the approximate expected discounted return R̂π , ( 1−γ ) E M̂ , π [ ∑∞ t=0 γ trt ] obtained from M̂ . In this paper , we are interested in the MDP estimate M̂ that can effectively reduce the error in the evaluation of policy π , |Rπ−R̂π| . In order to do so , we need to learn a good representation of a model that results in a small OPE error . We assume a bijective representation function φ : S ×A → Z where Z ⊆ Rdz is the representation space . We define the transition and the reward models in terms of the representation function φ , i.e . T̂ = T̂z ◦ φ and R̂ = R̂z ◦ φ . In practice , where we use a neural network for T̂ and R̂ , z can be chosen to be the output of an intermediate hidden layer , making φ represented by lower layers and T̂z , R̂z by the remaining upper layers . We define dπφ ( z ) the discounted stationary distribution on Z induced by dπ ( s , a ) under the representation function z = φ ( s , a ) , and similarly for dDφ ( z ) .

Summary: In this paper, the authors propose a model-based approach with representation balancing (RepB-SDE)to cope with the distribution shift of offline reinforcement learning. RepB-SDE learns a robust representation for the model learning process, which regularizes the distance between the data distribution and the discount stationary distribution of the target policy in the representation space. RepB-SDE adopts the estimation techniques of DualDICE and a novel point is that RepB-SDE plugs this trick into the model-based representation learning and proposes an effective model-based offline RL algorithm.