Offline Reinforcement Learning with In-sample Q-Learning

Offline reinforcement learning requires reconciling two conflicting aims : learning a policy that improves over the behavior policy that collected the dataset , while at the same time minimizing the deviation from the behavior policy so as to avoid errors due to distributional shift . This trade-off is critical , because most current offline reinforcement learning methods need to query the value of unseen actions during training to improve the policy , and therefore need to either constrain these actions to be in-distribution , or else regularize their values . We propose a new offline RL method that never needs to evaluate actions outside of the dataset , but still enables the learned policy to improve substantially over the best behavior in the data through generalization . The main insight in our work is that , instead of evaluating unseen actions from the latest policy , we can approximate the policy improvement step implicitly by treating the state value function as a random variable , with randomness determined by the action ( while still integrating over the dynamics to avoid excessive optimism ) , and then taking a state conditional upper expectile of this random variable to estimate the value of the best actions in that state . This leverages the generalization capacity of the function approximator to estimate the value of the best available action at a given state without ever directly querying a Q-function with this unseen action . Our algorithm alternates between fitting this upper expectile value function and backing it up into a Q-function , without any explicit policy . Then , we extract the policy via advantage-weighted behavioral cloning , which also avoids querying out-of-sample actions . We dub our method in-sample Q-learning ( IQL ) . IQL is easy to implement , computationally efficient , and only requires fitting an additional critic with an asymmetric L2 loss . IQL demonstrates the state-of-the-art performance on D4RL , a standard benchmark for offline reinforcement learning . We also demonstrate that IQL achieves strong performance fine-tuning using online interaction after offline initialization . 1 INTRODUCTION . Offline reinforcement learning ( RL ) addresses the problem of learning effective policies entirely from previously collected data , without online interaction ( Fujimoto et al. , 2019 ; Lange et al. , 2012 ) . This is very appealing in a range of real-world domains , from robotics to logistics and operations research , where real-world exploration with untrained policies is costly or dangerous , but prior data is available . However , this also carries with it major challenges : improving the policy beyond the level of the behavior policy that collected the data requires estimating values for actions other than those that were seen in the dataset , and this , in turn , requires trading off policy improvement against distributional shift , since the values of actions that are too different from those in the data are unlikely to be estimated accurately . Prior methods generally address this by either constraining the policy to limit how far it deviates from the behavior policy ( Fujimoto et al. , 2019 ; Wu et al. , 2019 ; Fujimoto & Gu , 2021 ; Kumar et al. , 2019 ; Nair et al. , 2020 ; Wang et al. , 2020 ) , or by regularizing the learned value functions to assign low values to out-of-distribution actions ( Kumar et al. , 2020 ; Kostrikov et al. , 2021 ) . Nevertheless , this imposes a trade-off between how much the policy improves and how vulnerable it is to misestimation due to distributional shift . Can we devise an offline RL method that avoids this issue by never needing to directly query or estimate values for actions that were not seen in the data ? In this work , we start from an observation that in-distribution constraints widely used in prior work might not be sufficient to avoid value function extrapolation , and we ask whether it is possible to learn an optimal policy with in-sample learning , without ever querying the values of any unseen actions . The key idea in our method is to approximate an upper expectile of the distribution over values with respect to the distribution of dataset actions for each state . We alternate between fitting this value function with expectile regression , and then using it to compute Bellman backups for training the Q-function . We show that we can do this simply by modifying the loss function in a SARSA-style TD backup , without ever using out-of-sample actions in the target value . Once thisQfunction has converged , we extract the corresponding policy using advantage-weighted behavioral cloning . This approach does not require explicit constraints or explicit regularization of out-ofdistribution actions during value function training , though our policy extraction step does implicitly enforce a constraint , as discussed in prior work on advantage-weighted regression ( Peters & Schaal , 2007 ; Peng et al. , 2019 ; Nair et al. , 2020 ; Wang et al. , 2020 ) . Our main contribution is in-sample Q-learning ( IQL ) , a new offline RL algorithm that avoids ever querying values of unseen actions while still being able to perform multi-step dynamic programming updates . Our method is easy to implement by making a small change to the loss function in a simple SARSA-like TD update and is computationally very efficient . Furthermore , our approach demonstrates the state-of-the-art performance on D4RL , a popular benchmark for offline reinforcement learning . In particular , our approach significantly improves over the prior state-of-the-art on challenging Ant Maze tasks that require to “ stitch ” several sub-optimal trajectories . Finally , we demonstrate that our approach is suitable for finetuning ; after initialization from offline RL , IQL is capable of improving policy performance utilizing additional interactions . 2 RELATED WORK . A significant portion of recently proposed offline RL methods are based on either constrained or regularized approximate dynamic programming ( e.g. , Q-learning or actor-critic methods ) , with the constraint or regularizer serving to limit deviation from the behavior policy . We will refer to these methods as “ multi-step dynamic programming ” algorithms , since they perform true dynamic programming for multiple iterations , and therefore can in principle recover the optimal policy if provided with high-coverage data . The constraints can be implemented via an explicit density model ( Wu et al. , 2019 ; Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Ghasemipour et al. , 2021 ) , implicit divergence constraints ( Nair et al. , 2020 ; Wang et al. , 2020 ; Peters & Schaal , 2007 ; Peng et al. , 2019 ; Siegel et al. , 2020 ) , or by adding a supervised learning term to the policy improvement objective ( Fujimoto & Gu , 2021 ) . Several works have also proposed to directly regularize the Q-function to produce low values for out-of-distribution actions ( Kostrikov et al. , 2021 ; Kumar et al. , 2020 ; Fakoor et al. , 2021 ) . Our method is also a multi-step dynamic programming algorithm . However , in contrast to prior works , our method completely avoids directly querying the learned Q-function with unseen actions during training , removing the need for any constraint during this stage , though the subsequent policy extraction , which is based on advantage-weighted regression ( Peng et al. , 2019 ; Nair et al. , 2020 ) , does apply an implicit constraint . However , this policy does not actually influence value function training . In contrast to multi-step dynamic programming methods , several recent works have proposed methods that rely either on a single step of policy iteration , fitting the value function or Q-function of the behavior policy and then extracting the corresponding greedy policy ( Peng et al. , 2019 ; Brandfonbrener et al. , 2021 ; Gulcehre et al. , 2021 ) , or else avoid value functions completely and utilize behavioral cloning-style objectives ( Chen et al. , 2021 ) . We collectively refer to these as “ single-step ” approaches . These methods avoid needing to query unseen actions as well , since they either use no value function at all , or learn the value function of the behavior policy . Although these methods are simple to implement and effective on the MuJoCo locomotion tasks in D4RL , we show that such single-step methods perform very poorly on more complex datasets in D4RL , which require combining parts of suboptimal trajectories ( “ stitching ” ) . Prior multi-step dynamic programming methods perform much better in such settings , as does our method . We discuss this distinction in more detail in Section 5.1 . Our method also shares the simplicity and computational efficiency of single-step approaches , providing an appealing combination of the strengths of both types of methods . Our method is based on estimating the characteristics of a random variable . Several recent works involve approximating statistical quantities of the value function distribution . In particular , quan- tile regression ( Koenker & Hallock , 2001 ) has been previously used in reinforcement learning to estimate the quantile function of a state-action value function ( Dabney et al. , 2018b ; a ; Kuznetsov et al. , 2020 ) . Although our method is related , in that we perform expectile regression , our aim is not to estimate the distribution of values that results from stochastic transitions , but rather estimate expectiles of the state value function with respect to random actions . This is a very different statistic : our aim is not to determine how the Q-value can vary with different future outcomes , but how the Q-value can vary with different actions while averaging together future outcomes due to stochastic dynamics . While prior work on distributional RL can also be used for offline RL , it would suffer from the same action extrapolation issues as other methods , and would require similar constraints or regularization , while our method does not . 3 PRELIMINARIES . The RL problem is formulated in the context of a Markov decision process ( MDP ) ( S , A , p0 ( s ) , p ( s′|s , a ) , r ( s , a ) , γ ) , where S is a state space , A is an action space , p0 ( s ) is a distribution of initial states , p ( s′|s , a ) is the environment dynamics , r ( s , a ) is a reward function , and γ is a discount factor . The agent interacts with the MDP according to a policy π ( a|s ) . The goal is to obtain a policy that maximizes the cumulative discounted returns : π∗ = argmax π Eπ [ ∞∑ t=0 γtr ( st , at ) |s0 ∼ p0 ( · ) , at ∼ π ( ·|st ) , st+1 ∼ p ( ·|st , at ) ] . Off-policy RL methods based on approximate dynamic programming typically utilize a state-action value function ( Q-function ) , referred to as Q ( s , a ) , which corresponds to the discounted returns obtained by starting from the state s and action a , and then following the policy π. Offline reinforcement learning . In contrast to online ( on-policy or off-policy ) RL methods , offline RL uses previously collected data without any additional data collection . Like many recent offline RL methods , our work builds on approximate dynamic programming methods that minimize temporal difference error , according to the following loss : LTD ( θ ) = E ( s , a , s′ ) ∼D [ ( r ( s , a ) + γmax a′ Qθ̂ ( s ′ , a′ ) −Qθ ( s , a ) ) 2 ] , ( 1 ) where D is the dataset , Qθ ( s , a ) is a parameterized Q-function , Qθ̂ ( s , a ) is a target network ( e.g. , with soft parameters updates defined via Polyak averaging ) , and the policy is defined as π ( s ) = argmaxaQθ ( s , a ) . Most recent offline RL methods modify either the value function loss ( above ) to regularize the value function in a way that keeps the resulting policy close to the data , or constrain the argmax policy directly . This is important because out-of-distribution actions a′ can produce erroneous values for Qθ̂ ( s ′ , a′ ) in the above objective , often leading to overestimation as the policy is defined to maximize the ( estimated ) Q-value .

Summary: This paper proposes a new offline RL algorithm that uses in-sample policy evaluation and advantage-weighted regression during policy improvement. Particularly, it utilizes sarsa-like TD to update Q-function and avoids the query value of unseen actions during training. To evaluate their proposed method, they used the D4RL benchmark to compare with previous offline RL methods. In addition, they show results for fine-tuning of learned offline-policies during online deployment.