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: This paper suggests an alternative approach, equivalent to adversarial IRL (that learns both the policy and the reward in an adversarial fashion), but in a direct way (learning the Q-function that encompasses both the reward and the policy). As in adversarial methods (GAIL, AIRL etc) they use the maximum entropy assumption and more explicitly, they assume that the optimal policy is greedy wrt an optimal soft-Q-function.   Using this relation between the optimal policy and the Q-function, they restrict the set of solutions to a manifold in which the objective only depends on the Q-function. They just have to minimize this objective (ensuring on the side that the policy stays in the manifold, i.e. it stays greedy) to obtain an optimal policy and Q-function. Their solution, in theory equivalent to the ones from an adversarial setting, is expected to outperform them as they avoid the learning difficulties of adversarial learning. Optionally, they can recover a reward function from the Q-function by simply inverting the Bellman operator.