IQ-Learn: Inverse soft-Q Learning for Imitation

1 Introduction . Imitation of an expert has long been recognized as a powerful approach for sequential decisionmaking [ 29 , 1 ] , with applications as diverse as healthcare [ 39 ] , autonomous driving [ 41 ] , and playing complex strategic games [ 8 ] . In the imitation learning ( IL ) setting , we are given a set of expert trajectories , with the goal of learning a policy which induces behavior similar to the expert ’ s . The learner has no access to the reward , and no explicit knowledge of the dynamics . The simple behavioural cloning [ 34 ] approach simply maximizes the probability of the expert ’ s actions under the learned policy , approaching the IL problem as a supervised learning problem . While this can work well in simple environments and with large quantities of data , it ignores the sequential nature of the decision-making problem , and small errors can quickly compound when the learned policy departs from the states observed under the expert . A natural way of introducing the environment dynamics is by framing the IL problem as an Inverse RL ( IRL ) problem , aiming to learn a reward function under which the expert ’ s trajectory is optimal , and from which the learned imitation policy can be trained [ 1 ] . This framing has inspired several approaches which use rewards either explicitly or implicitly to incorporate dynamics while learning an imitation policy [ 17 , 10 , 33 , 22 ] . However , these dynamics-aware methods are typically hard to put into practice due to unstable learning which can be sensitive to hyperparameter choice or minor implementation details [ 21 ] . In this work , we introduce a dynamics-aware imitation learning method which has stable , nonadversarial training , allowing us to achieve state-of-the-art performance on imitation learning bench- 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . marks . Our key insight is that much of the difficulty with previous IL methods arises from the IRL-motivated representation of the IL problem as a min-max problem over reward and policy [ 17 , 1 ] . This introduces a requirement to separately model the reward and policy , and train these two functions jointly , often in an adversarial fashion . Drawing on connections between RL and energy-based models [ 13 , 14 ] , we propose learning a single model for the Q-value . The Q-value then implicitly defines both a reward and policy function . This turns a difficult min-max problem over policy and reward functions into a simpler minimization problem over a single function , the Q-value . Since our problem has a one-to-one correspondence with the min-max problem studied in adversarial IL [ 17 ] , we maintain the generality and guarantees of these previous approaches , resulting in a meaningful reward that may be used for inverse reinforcement learning . Furthermore , our method may be used to minimize a variety of statistical divergences between the expert and learned policy . We show that we recover several previously-described approaches as special cases of particular divergences , such as the regularized behavioural cloning of [ 30 ] , and the conservative Q-learning of [ 23 ] . In our experiments , we find that our method is performant even with very sparse data - surpassing prior methods using one expert demonstration in the completely offline setting - and can scale to complex image-based tasks like Atari reaching expert performance . Moreover , our learnt rewards are highly predictive of the original environment rewards . Concretely , our contributions are as follows : • We present a modified Q-learning update rule for imitation learning that can be implemented on top of soft-Q learning or soft actor-critic ( SAC ) algorithms in fewer than 15 lines of code . • We introduce a simple framework to minimize a wide range of statistical distances : Integral Probability Metrics ( IPMs ) and f-divergences , between the expert and learned distributions . • We empirically show state-of-art results in a variety of imitation learning settings : online and offline IL . On the complex Atari suite , we outperform prior methods by 3-7x while requiring 3x less environment steps . • We characterize our learnt rewards and show a high positive correlation with the ground-truth rewards , justifying the use of our method for Inverse Reinforcement Learning . 2 Background . Preliminaries We consider environments represented as a Markov decision process ( MDP ) , which is defined by a tuple ( S , A , p0 , P , r , ) . S , A represent state and action spaces , p0 and P ( s0|s , a ) represent the initial state distribution and the dynamics , r ( s , a ) represents the reward function , and 2 ( 0 , 1 ) represents the discount factor . RS⇥A = { x : S ⇥ A ! R } will denote the set of all functions in the state-action space and R will denote the extended real numbers R [ { 1 } . Section 3 and 4 will work with finite state and action spaces S and A , but our algorithms and experiments later in the paper use continuous environments . ⇧ is the set of all stationary stochastic policies that take actions in A given states in S. We work in the -discounted infinite horizon setting , and we will use an expectation with respect to a policy ⇡ 2 ⇧ to denote an expectation with respect to the trajectory it generates : E⇡ [ r ( s , a ) ] , E [ P1 t=0 t r ( st , at ) ] , where s0 ⇠ p0 , at ⇠ ⇡ ( ·|st ) , and st+1 ⇠ P ( ·|st , at ) for t 0 . For a policy ⇡ 2 ⇧ , we define its occupancy measure ⇢⇡ : S ⇥ A ! R as ⇢⇡ ( s , a ) = ⇡ ( a|s ) P1 t=0 t P ( st = s|⇡ ) . We refer to the expert policy as ⇡E and its occupancy measure as ⇢E . In practice , ⇡E is unknown and we have access to a sampled dataset of demonstrations . For brevity , we refer to ⇢⇡ as ⇢ for a learnt policy in the paper . Soft Q-functions For a reward r 2 RS⇥A and ⇡ 2 ⇧ , the soft Bellman operator B⇡ : RS⇥A ! RS⇥A defined as ( B⇡Q ) ( s , a ) = r ( s , a ) + Es0⇠P ( s , a ) V ⇡ ( s0 ) with V ⇡ ( s ) = Ea⇠⇡ ( ·|s ) [ Q ( s , a ) log ⇡ ( a|s ) ] . The soft Bellman operator is contractive [ 13 ] and defines a unique soft Q-function for r , given as Q = B⇡Q . Max Entropy Reinforcement Learning For a given reward function r 2 RS⇥A , maximum entropy RL [ 14 , 5 ] aims to learn a policy that maximizes the expected cumulative discounted reward along with the entropy in each state : max⇡2⇧ E⇡ [ r ( s , a ) ] +H ( ⇡ ) . Where H ( ⇡ ) , E⇡ [ log ⇡ ( a|s ) ] is the discounted causal entropy of the policy ⇡ . The optimal policy satisfies [ 42 , 5 ] : ⇡ ⇤ ( a|s ) = 1 Zs exp ( Q ( s , a ) ) , ( 1 ) where Q is the soft Q-function and Zs is the normalization factor given as P a0 exp ( Q ( s , a 0 ) ) . Q satisfies the soft-Bellman equation : Q ( s , a ) = r ( s , a ) + Es0⇠P ( ·|s , a ) h log X a0 exp ( Q ( s0 , a0 ) ) i ( 2 ) In continuous action spaces , Zs becomes intractable and soft actor-critic methods like SAC [ 13 ] can be used to learn an explicit policy . Max Entropy Inverse Reinforcement Learning Given demonstrations sampled using the policy ⇡E , maximum entropy Inverse RL aims to recover the reward function in a family of functions R that rationalizes the expert behavior by solving the optimization problem : maxr2R min⇡2⇧ E⇡E [ r ( s , a ) ] ( E⇡ [ r ( s , a ) ] + H ( ⇡ ) ) , where the expected reward of ⇡E is empirically approximated . It looks for a reward function that assigns high reward to the expert policy and a low reward to other policies , while searching for the best policy for the reward function in an inner loop . The Inverse RL objective can be reformulated in terms of its occupancy measure , and with a convex reward regularizer : RS⇥A ! R [ 17 ] max r2R min ⇡2⇧ L ( ⇡ , r ) = E⇢E [ r ( s , a ) ] E⇢ [ r ( s , a ) ] H ( ⇡ ) ( r ) ( 3 ) In general , we can exchange the max-min resulting in an objective that minimizes the statistical distance parameterized by , between the expert and the policy [ 17 ] min ⇡2⇧ max r2R L ( ⇡ , r ) = min ⇡2⇧ d ( ⇢ , ⇢E ) H ( ⇡ ) , ( 4 ) with d , ⇤ ( ⇢E ⇢ ) , where ⇤ is the convex conjugate of . 3 Inverse soft Q-learning ( IQ-Learn ) Framework . A naive solution to the IRL problem in ( Eq . 3 ) involves ( 1 ) an outer loop learning rewards and ( 2 ) executing RL in an inner loop to find an optimal policy for them . However , we know that this optimal policy can be obtained analytically in terms of soft Q-functions ( Eq . 1 ) . Interestingly , as we will show later , the rewards can also be represented in terms of Q ( Eq . 2 ) . Together , these observations suggest it might be possible to directly solve the IRL problem by optimizing only over the Q-function . To motivate the search of an imitation learning algorithm that depends only on the Q-function , we characterize the space of Q-functions and policies obtained using Inverse RL . We will study ⇡ 2 ⇧ , r 2 R and Q-functions Q 2 ⌦ where R = ⌦ = RS⇥A . We assume ⇧ is convex , compact and that ⇡E 2 ⇧1 . We define V ⇡ ( s ) = Ea⇠⇡ ( ·|s ) [ Q ( s , a ) log ⇡ ( a|s ) ] . We start with analysis developed in [ 17 ] : The regularized IRL objective L ( ⇡ , r ) given by Eq . 3 , is concave in the policy and convex in rewards . And has a unique saddle point where it is optimized . To characterize the Q-functions it is useful to transform the optimization problem over rewards to a problem over Q-functions . We can get a one-to-one correspondence between r and Q : Define the inverse soft bellman operator T ⇡ : RS⇥A ! RS⇥A such that ( T ⇡Q ) ( s , a ) = Q ( s , a ) Es0⇠P ( s , a ) V ⇡ ( s0 ) , Lemma 3.1 . The inverse soft bellman operator T ⇡ is bijective , and for any r , ( T ⇡ ) 1r is the unique contraction of B⇡ . The proof of this lemma is in Appendix A.1 . For a policy ⇡ , we are thus justified in changing between rewards and their corresponding soft-Q functions . We can freely transform functions from the reward-policy space : ⇧ ⇥ R to the Q-policy space : ⇧ ⇥ ⌦ , giving us the lemma : Lemma 3.2 . If L ( ⇡ , r ) = E⇢E [ r ( s , a ) ] E⇢ [ r ( s , a ) ] H ( ⇡ ) ( r ) and J ( ⇡ , Q ) = E⇢E [ ( T ⇡Q ) ( s , a ) ] E⇢ [ ( T ⇡Q ) ( s , a ) ] H ( ⇡ ) ( T ⇡Q ) , then for all policies ⇡ 2 ⇧ , L ( ⇡ , r ) = J ( ⇡ , ( T ⇡ ) 1r ) for all r 2 R , and J ( ⇡ , Q ) = L ( ⇡ , T ⇡Q ) , for all Q 2 ⌦ . Lemma 3.1 and 3.2 allow us to adapt the Inverse RL objective L ( ⇡ , r ) to learning Q through J ( ⇡ , Q ) . Simplifying our new objective ( using Lemma A.3 in Appendix ) : J ( ⇡ , Q ) = Es , a⇠⇢E [ Q Es0⇠P ( ·|s , a ) V ⇡ ( s0 ) ] ( 1 ) Es0⇠p0 [ V ⇡ ( s0 ) ] ( T ⇡Q ) , ( 5 ) We are now ready to study J ( ⇡ , Q ) , the Inverse RL optimization problem in the Q-policy space . As the regularizer depends on both Q and ⇡ , a general analysis over all functions in RS⇥A becomes too difficult . We restrict ourselves to regularizers induced by a convex function g : R ! R such that g ( r ) = E⇢E [ g ( r ( s , a ) ) ] ( 6 ) This allows us to simplify our analysis to the set of all real functions while retaining generality2 . We further motivate this choice in Section 4 . Proposition 3.3 . In the Q-policy space , there exists a unique saddle point ( ⇡⇤ , Q⇤ ) that optimizes J . i.e . Q⇤ = argmaxQ2⌦ min⇡2⇧ J ( ⇡ , Q ) and ⇡⇤ = argmin⇡2⇧ maxQ2⌦ J ( ⇡ , Q ) . Furthermore , ⇡ ⇤ and r⇤ = T ⇡⇤Q⇤ are the solution to the Inverse RL objective L ( ⇡ , r ) . Thus we have , maxQ2⌦ min⇡2⇧ J ( ⇡ , Q ) = maxr2R min⇡2⇧ L ( ⇡ , r ) . This tells us , even after transforming to Q-functions we have retained the saddle point property of the original IRL objective and optimizing J ( ⇡ , Q ) recovers this saddle point . In the Q-policy space , we can get an additional property : Proposition 3.4 . For a fixed Q , argmin⇡2⇧ J ( ⇡ , Q ) is the solution to max entropy RL with rewards r = T ⇡Q . Thus , this forms a manifold in the Q-policy space , that satisfies ⇡Q ( a|s ) = 1 Zs exp ( Q ( s , a ) ) , with normalization factor Zs = P a exp Q ( s , a ) and ⇡Q defined as the ⇡ corresponding to Q . Proposition 3.3 and 3.4 are telling us that if we know Q , then the inner optimization problem in terms of policy is trivial , and obtained in a closed form ! Thus , we can recover an objective that only requires learning Q : max Q2⌦ min ⇡2⇧ J ( ⇡ , Q ) = max Q2⌦ J ( ⇡Q , Q ) ( 7 ) Furthermore , we have : 1The full policy class satisfies all these assumptions 2Averaging over the expert occupancy allows to adjust to arbitrary experts and accommodate multimodality Proposition 3.5 . Let J ⇤ ( Q ) = J ( ⇡Q , Q ) . Then J ⇤ is concave in Q . Thus , this new optimization objective is well-behaved and is maximized only at the saddle point . < latexit sha1_base64= '' ig4yF5qJxp0Cc2lIlU5zGovCIUI= '' > AAAB9HicbVBNS8NAEN34WetX1aOXxSJUkJJIRY9FLx48VLAf0ISy2WzapZts3J0USujv8OJBEa/+GG/+G7dtDtr6YODx3gwz8/xEcA22/W2trK6tb2wWtorbO7t7+6WDw5aWqaKsSaWQquMTzQSPWRM4CNZJFCORL1jbH95O/faIKc1l/AjjhHkR6cc85JSAkbz7ipvwc+zSQMJZr1S2q/YMeJk4OSmjHI1e6csNJE0jFgMVROuuYyfgZUQBp4JNim6qWULokPRZ19CYREx72ezoCT41SoBDqUzFgGfq74mMRFqPI990RgQGetGbiv953RTCay/jcZICi+l8UZgKDBJPE8ABV4yCGBtCqOLmVkwHRBEKJqeiCcFZfHmZtC6qTq16+VAr12/yOAroGJ2gCnLQFaqjO9RATUTRE3pGr+jNGlkv1rv1MW9dsfKZI/QH1ucPTQ+RKA== < /latexit > < latexit sha1_base64= '' t0faVr2Ta73XcO9Ogw+Hsj1Rzps= '' > AAAB8nicbVBNSwMxEM36WetX1aOXYBEqSNmVih6LXjx4qGA/YLuUbDbbhmaTJZkVSunP8OJBEa/+Gm/+G9N2D9r6YODx3gwz88JUcAOu++2srK6tb2wWtorbO7t7+6WDw5ZRmaasSZVQuhMSwwSXrAkcBOukmpEkFKwdDm+nfvuJacOVfIRRyoKE9CWPOSVgJf++0qWRgnOsz3qlslt1Z8DLxMtJGeVo9Epf3UjRLGESqCDG+J6bQjAmGjgVbFLsZoalhA5Jn/mWSpIwE4xnJ0/wqVUiHCttSwKeqb8nxiQxZpSEtjMhMDCL3lT8z/MziK+DMZdpBkzS+aI4ExgUnv6PI64ZBTGyhFDN7a2YDogmFGxKRRuCt/jyMmldVL1a9fKhVq7f5HEU0DE6QRXkoStUR3eogZqIIoWe0St6c8B5cd6dj3nripPPHKE/cD5/ANUPkFE= < /latexit > < latexit sha1_base64= '' BumFioJ7GF/Lp9Isv4bReWWK9pA= '' > AAAB/nicbVDLSgMxFM3UV62vUXHlJliEClJmpKLLohtx1YJ9QGcomUzahmaSIckIZSj4K25cKOLW73Dn35hpZ6GtBwKHc+7lnpwgZlRpx/m2Ciura+sbxc3S1vbO7p69f9BWIpGYtLBgQnYDpAijnLQ01Yx0Y0lQFDDSCca3md95JFJRwR/0JCZ+hIacDihG2kh9+8iLkB5hxNL7acXDodDnsHnWt8tO1ZkBLhM3J2WQo9G3v7xQ4CQiXGOGlOq5Tqz9FElNMSPTkpcoEiM8RkPSM5SjiCg/ncWfwlOjhHAgpHlcw5n6eyNFkVKTKDCTWVi16GXif14v0YNrP6U8TjTheH5okDCoBcy6gCGVBGs2MQRhSU1WiEdIIqxNYyVTgrv45WXSvqi6tepls1au3+R1FMExOAEV4IIrUAd3oAFaAIMUPINX8GY9WS/Wu/UxHy1Y+c4h+APr8wcgWpTx < /latexit > < latexit sha1_base64= '' SB7eN8/tjH0kd8P8DZ1Oq7+cTn8= '' > AAACAHicbVDLSsNAFJ3UV62vqAsXbgaLUEFKIhVdFt2Iqwr2AU0ok8mkHTqZCTMToYRs/BU3LhRx62e482+ctllo9cCFwzn3cu89QcKo0o7zZZWWlldW18rrlY3Nre0de3evo0QqMWljwYTsBUgRRjlpa6oZ6SWSoDhgpBuMr6d+94FIRQW/15OE+DEachpRjLSRBvaBFyM9wohlt3nNS+gp9HAo9MnArjp1Zwb4l7gFqYICrYH96YUCpzHhGjOkVN91Eu1nSGqKGckrXqpIgvAYDUnfUI5iovxs9kAOj40SwkhIU1zDmfpzIkOxUpM4MJ3Tc9WiNxX/8/qpji79jPIk1YTj+aIoZVALOE0DhlQSrNnEEIQlNbdCPEISYW0yq5gQ3MWX/5LOWd1t1M/vGtXmVRFHGRyCI1ADLrgATXADWqANMMjBE3gBr9aj9Wy9We/z1pJVzOyDX7A+vgHTLpXp < /latexit > < latexit sha1_base64= '' 94BIgKv/3TFlXWxL95WraRqiLwY= '' > AAAB83icbVBNS8NAEJ3Ur1q/qh69LBahFimJVPRY9OKxgv2AJi2b7aZdukmW3Y1QQv+GFw+KePXPePPfuG1z0NYHA4/3ZpiZ5wvOlLbtbyu3tr6xuZXfLuzs7u0fFA+PWipOJKFNEvNYdnysKGcRbWqmOe0ISXHoc9r2x3czv/1EpWJx9KgngnohHkYsYARrI7llV7Be5QLJXuW8XyzZVXsOtEqcjJQgQ6Nf/HIHMUlCGmnCsVJdxxbaS7HUjHA6LbiJogKTMR7SrqERDqny0vnNU3RmlAEKYmkq0miu/p5IcajUJPRNZ4j1SC17M/E/r5vo4MZLWSQSTSOyWBQkHOkYzQJAAyYp0XxiCCaSmVsRGWGJiTYxFUwIzvLLq6R1WXVq1auHWql+m8WRhxM4hTI4cA11uIcGNIGAgGd4hTcrsV6sd+tj0Zqzsplj+APr8wfog5BO < /latexit > < latexit sha1_base64= '' 9P+Yq3viFrFovsJoEM/kjy4XXic= '' > AAAB83icbVBNS8NAEJ34WetX1aOXxSLUIiWRih6LXjy2YD+gSctmu2mXbpJldyOU0L/hxYMiXv0z3vw3btsctPXBwOO9GWbm+YIzpW3721pb39jc2s7t5Hf39g8OC0fHLRUnktAmiXksOz5WlLOINjXTnHaEpDj0OW374/uZ336iUrE4etQTQb0QDyMWMIK1kdySK1ivfIkavfJFv1C0K/YcaJU4GSlChnq/8OUOYpKENNKEY6W6ji20l2KpGeF0mncTRQUmYzykXUMjHFLlpfObp+jcKAMUxNJUpNFc/T2R4lCpSeibzhDrkVr2ZuJ/XjfRwa2XskgkmkZksShIONIxmgWABkxSovnEEEwkM7ciMsISE21iypsQnOWXV0nrquJUK9eNarF2l8WRg1M4gxI4cAM1eIA6NIGAgGd4hTcrsV6sd+tj0bpmZTMn8AfW5w+2HJAt < /latexit > < latexit sha1_base64= '' ihRPGf7o9AxxBdveWwqYHm29dhI= '' > AAACBHicbVDLSsNAFJ3UV62vqMtuBotQNyWRii4UCm5ctmAf2IQwmU7aoZNJmJmIJWThxl9x40IRt36EO//GaZuFth64cDjnXu69x48Zlcqyvo3Cyura+kZxs7S1vbO7Z+4fdGSUCEzaOGKR6PlIEkY5aSuqGOnFgqDQZ6Trj6+nfveeCEkjfqsmMXFDNOQ0oBgpLXlm2YkpvIJOIBBO7Sy982QGHfIQV1snnlmxatYMcJnYOamAHE3P/HIGEU5CwhVmSMq+bcXKTZFQFDOSlZxEkhjhMRqSvqYchUS66eyJDB5rZQCDSOjiCs7U3xMpCqWchL7uDJEayUVvKv7n9RMVXLgp5XGiCMfzRUHCoIrgNBE4oIJgxSaaICyovhXiEdJ5KJ1bSYdgL768TDqnNbteO2vVK43LPI4iKIMjUAU2OAcNcAOaoA0weATP4BW8GU/Gi/FufMxbC0Y+cwj+wPj8ASBBlx0= < /latexit > < latexit sha1_base64= '' B/F7GSCY381gFwbuqoCBPCMgSBA= '' > AAAB+XicbVDLSgMxFL1TX7W+Rl26CRahuigzUtFl0Y24asE+oB1LJk3b0ExmSDKFMvRP3LhQxK1/4s6/MdPOQlsPBA7n3Ms9OX7EmdKO823l1tY3Nrfy24Wd3b39A/vwqKnCWBLaICEPZdvHinImaEMzzWk7khQHPqctf3yX+q0JlYqF4lFPI+oFeCjYgBGsjdSz7W6A9YhgnjzMni5K9fOeXXTKzhxolbgZKUKGWs/+6vZDEgdUaMKxUh3XibSXYKkZ4XRW6MaKRpiM8ZB2DBU4oMpL5sln6MwofTQIpXlCo7n6eyPBgVLTwDeTaU617KXif14n1oMbL2EiijUVZHFoEHOkQ5TWgPpMUqL51BBMJDNZERlhiYk2ZRVMCe7yl1dJ87LsVspX9UqxepvVkYcTOIUSuHANVbiHGjSAwASe4RXerMR6sd6tj8Vozsp2juEPrM8fjheS9Q== < /latexit > In Appendix C , we expand on our analysis and characterize the behavior for different choices of regularizer , while giving proofs of all our propositions . Figure 1 summarizes the properties for the IRL objective : there exists a optimal policy manifold depending on Q , allowing optimization along it ( using J ⇤ ) to converge to the saddle point . We further present analysis of IL methods that learn Q-functions like SQIL [ 33 ] and ValueDICE [ 22 ] and find subtle fallacies affecting their learning . Note that although the same analysis holds in the reward-policy space , the optimal policy manifold depends on Q , which isn ’ t trivially known unlike when in the Q-policy space .

Summary: The paper presents an IL/IRL method that avoids estimating the reward during optimization by alternating between updates of the policy and Q-function instead, which builds on recent insights from ValueDice. Indeed, the optimization problem for learning the Q-function is very similar to the one used by ValueDice, when optimizing the reverse KL. Notable differences compared to ValueDice are the generalization to a large class of divergences (including f-divergences), non-adversarial formulation, and a procedure to recover rewards. In particular, I think that the modification for recovering a state-only reward function is very interesting. The algorithm only requires mild modifications of SAC (mainly by using a different loss for the Q-function) and shows promising results in experiment evaluations.