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 an offline RL algorithm named IQL, which generalizes between Bellman expectation equation and Bellman optimality equation with expectile regression. Expectile regression assigns low weights for low-performing samples and assigns high weights for high-performing samples to learn the expectile. By performing $\tau$-expectile regression on the randomness of action in the dataset (from the sampling process of behavior policy), IQL is able to obtain a value function that is $\tau$ optimal. IQL is generalizing SARSA and Q-learning in the sense that SARSA corresponds to IQL with $\tau=0.5$ and Q-learning corresponds to IQL with $tau=1.0$. It offers very stable learning of a near-optimal value since it is only using in-distribution samples and the effective dataset size can be controlled through $\tau$. In the experiments, the paper has shown that IQL can obtain state-of-the-art performance.