IQ-Learn: Inverse soft-Q Learning for Imitation
Year: 2022

Abstract:
In many sequential decision-making problems (e.g., robotics control, game playing, sequential prediction), human or expert data is available containing useful information about the task. However, imitation learning (IL) from a small amount of expert data can be challenging in high-dimensional environments with complex dynamics. Behavioral cloning is a simple method that is widely used due to its simplicity of implementation and stable convergence but doesn’t utilize any information involving the environment’s dynamics. Many existing methods that exploit dynamics information are difficult to train in practice due to an adversarial optimization process over reward and policy approximators or biased, high variance gradient estimators. We introduce a method for dynamics-aware IL which avoids adversarial training by learning a single Q-function, implicitly representing both reward and policy. On standard benchmarks, the implicitly learned rewards show a high positive correlation with the ground-truth rewards, illustrating our method can also be used for inverse reinforcement learning (IRL). Our method, Inverse soft-Q learning (IQ-Learn) obtains state-of-the-art results in offline and online imitation learning settings, significantly outperforming existing methods both in the number of required environment interactions and scalability in high-dimensional spaces, often by more than 3x.

heading: 1 Introduction
text: 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.

heading: 2 Background
text: 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 .

heading: 3 Inverse soft Q-learning (IQ-Learn) Framework
text: 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.

heading: 4 Approach
text: In this section, we develop our inverse soft-Q learning (IQ-Learn) algorithm, such that it recovers the optimal soft Q-function for a MDP from a given expert distribution. We start by learning energy-based models for the policy similar to soft Q-learning and later learn an explicit policy similar to actor-critic methods.

heading: 4.1 General Inverse RL Objective
text: For designing a practical algorithm using regularizers of the form g (from Eq. 6), we define g using
a concave function : R ! R, such that g(x) = ⇢
x (x) if x 2 R +1 otherwise
with the rewards constrained in R .
For this choice of , the Inverse RL objective L(⇡, r) takes the form of Eq. 4 with a distance measure:
d (⇢, ⇢E) = max r2R E⇢E [ (r(s, a))] E⇢[r(s, a)], (8)
This forms a general learning objective that allows the use of a wide-range of statistical distances including Integral Probability Metrics (IPMs) and f-divergences (see Appendix B). 3

heading: 4.2 Choice of Statistical Distances
text: While choosing a practical regularizer, it can be useful to obtain certain properties on the reward functions we recover. Some (natural) nice properties are: having rewards bounded in a range, learning smooth functions or enforcing a norm-penalty.
In fact, we find these properties correspond to the Total Variation distance, the Wasserstein-1 distance and the 2-divergence respectively. The regularizers and the induced statistical distances are summarized in Table 2:
Table 2: Enforced reward property, corresponding regularizer and statistical distance (Rmax,K,↵ 2 R+ )
Reward Property d Bound range = 0 if |r|  Rmax and +1 otherwise 2Rmax · TV(⇢, ⇢E) Smoothness = 0 if krkLip  K and +1 otherwise K ·W1(⇢, ⇢E) L2 Penalization (r) = ↵r2 14↵ ·
2(⇢, ⇢E) 3We recover IPMs when using identity and restricted reward family R
We find that these choice of regularizers4 work very well in our experiments. In Appendix B, we further give a table for the well known f -divergences, the corresponding and the learnt reward estimators, along with a result ablation on using different divergences. Compared to 2, we find other f -divergences like Jensen-Shannon result in similar performances but are not as readily interpretable.

heading: 4.3 Inverse soft-Q update (Discrete control)
text: Optimization along the optimal policy manifold gives the concave objective (Prop 3.5): max Q2⌦ J ⇤(Q) = E⇢E [ (Q(s, a) Es0⇠P(·|s,a)V ⇤(s0))] (1 )E⇢0 [V ⇤(s0)], (9) with V ⇤(s) = log P
a exp Q(s, a).
For each Q, we get a corresponding reward r(s, a) = Q(s, a) Es0⇠P(·|s,a)[log P a0 exp Q (s 0 , a
0)]. This correspondence is unique (Lemma A.1 in Appendix), and every update step can be seen as finding a better reward for IRL.
Note that estimating V ⇤(s) exactly is only possible in discrete action spaces. Our objective forms a variant of soft-Q learning: to learn the optimal Q-function given an expert distribution.

heading: 4.4 Inverse soft actor-critic update (Continuous control)
text: In continuous action spaces, it might not be possible to exactly obtain the optimal policy ⇡Q, which forms an energy-based model of the Q-function, and we use an explicit policy ⇡ to approximate ⇡Q.
For any policy ⇡, we have a objective (from Eq. 5): J (⇡, Q) = E⇢E [ (Q Es0⇠P(·|s,a)V ⇡(s0))] (1 )E⇢0 [V ⇡(s0)] (10)
For a fixed Q, soft actor-critic (SAC) update: min ⇡ Es⇠D,a⇠⇡(·|s)[Q(s, a) log ⇡(a|s)], brings ⇡ closer to ⇡Q while always minimizing Eq. 10 (Lemma A.4 in Appendix). Here D is the distribution of previously sampled states, or a replay buffer.
Thus, we obtain the modified actor-critic update rule to learn Q-functions from the expert distribution:
1. For a fixed ⇡, optimize Q by maximizing J (⇡, Q). 2. For a fixed Q, apply SAC update to optimize ⇡ towards ⇡Q.
This differs from ValueDICE [22], where the actor is updated adverserially and the objective may not always converge (Appendix C).

heading: 5 Practical Algorithm
text: Pseudocode in Algorithm 1, shows our Q-learning and actor-critic variants, with differences with conventional RL algorithms in red (we optimize -J to use gradient descent). We can implement our algorithm IQ-Learn in 15 lines of code on top of standard implementations of (soft) DQN [14] for discrete control or soft actor-critic (SAC) [13] for continuous control, with a change on the objective for the Q-function. Default hyperparameters from [14, 13] work well, except for tuning the entropy regularization. Target networks were helpful for continuous control. We elaborate details in Appendix D.

heading: 5.1 Training methodology
text: Corollary 2.1 in Appendix A states E(s,a)⇠µ[V ⇡(s) Es0⇠P(·|s,a)V ⇡(s0)] = (1 )Es⇠p0 [V ⇡(s)], where µ is any policy’s occupancy. We use this to stabilize training instead of using Eq. 9 directly.
Online: Instead of directly estimating Ep0 [V ⇡(s0)] in our algorithm, we can sample (s, a, s0) from a replay buffer and get a single-sample estimate E(s,a,s0)⇠replay[V ⇡(s) V ⇡(s0)]. This removes the issue where we are only optimizing Q in the inital states resulting in overfitting of V ⇡(s0), and improves the stability for convergence in our experiments. We find sampling half from the policy buffer and half from the expert distribution gives the best performances. Note that this is makes our learning online, requiring environment interactions.
4The additional scalar terms scale the entropy regularization strength and can be ignored in practice
Algorithm 1 Inverse soft Q-Learning (both variants) 1: Initialize Q-function Q✓ , and optionally a policy ⇡ 2: for step t in {1...N} do 3: Train Q-function using objective from Equation 9:
✓t+1 ✓t ↵Qr✓[ J (✓)] (Use V ⇤ for Q-learning and V ⇡ for actor-critic)
4: (only with actor-critic) Improve policy ⇡ with SAC style actor update: t+1 t ↵⇡r Es⇠D,a⇠⇡ (·|s)[Q(s, a) log ⇡ (a|s)] 5: end for
Algorithm 2 Recover policy and reward 1: Given trained Q-function Q✓ , and optionally a trained policy ⇡ 2: Recover policy ⇡:
(Q-learning) ⇡ := 1Z expQ✓ (actor-critic) ⇡ := ⇡
3: For state s, action a and s0 ⇠ P(·|s,a) 4: Recover reward r(s,a, s0) = Q✓(s,a) V ⇡ (s0)
Offline: Although Ep0 [V ⇡(s0)] can be estimated offline we still observe an overfitting issue. Instead of requiring policy samples we use only expert samples to estimate E(s,a,s0)⇠expert[V ⇡(s) V
⇡(s0)] to sufficiently approximate the term. This methodology gives us state-of-art results for offline IL.

heading: 5.2 Recovering rewards
text: Instead of the conventional reward function r(s, a) on state and action pairs, our algorithm allows recovering rewards for each transition (s, a, s0) using the learnt Q-values as follows: r(s, a, s0) = Q(s, a) V ⇡ (s0) (11) Now, Es0⇠P(·|s,a)[Q(s, a) V ⇡ (s0)] = Q(s, a) Es0⇠P(·|s,a)[V ⇡ (s0)] = T ⇡Q(s, a). This is just the reward function r(s, a) we want. So by marginalizing over next-states, our expression correctly recovers the reward over state-actions. Thus, Eq. 11 gives the reward over transitions.
Our rewards require s0 which can be sampled from the environment, or by using a dynamics model.

heading: 5.3 Implementation of Statistical Distances
text: Implementing TV and W1 distances is fairly trivial and we give details in Appendix B. For the
2-divergence, we note that it corresponds to (x) = x 14↵x 2. On substituting in Eq. 9, we get
max Q2⌦
E⇢E [(Q(s, a) Es0⇠P(·|s,a)V ⇤(s0))] (1 )Ep0 [V ⇤(s0)] 14↵E⇢E [(Q(s, a) Es0⇠P(·|s,a)V ⇤(s0))2]
In a fully offline setting, this can be further simplified as (using the offline methodology in Sec 5.1):
min Q2⌦
E⇢E [(Q(s, a) V ⇤(s))] + 1
4↵ E⇢E [(Q(s, a) Es0⇠P(·|s,a)V ⇤(s0))2] (12)
This is interestingly the same as the Q-learning objective in CQL [23], an state-of-art method for offline RL (using 0 rewards), and shares similarities with regularized behavior cloning [33] 5.

heading: 5.4 Learning state-only reward functions
text: Previous works like AIRL [10] propose learning rewards that are only function of the state, and claim that these form of reward functions generalize between different MDPs. We find our method can predict state-only rewards by using the policy and expert state-marginals with a modification to Eq. 9:
max Q2⌦
J ⇤(Q) = Es⇠⇢E(s)[Ea⇠⇡(·|s)[ (Q(s, a) Es0⇠P(·|s,a)V ⇤(s0))]] (1 )Ep0 [V ⇤(s0)]
Interestingly, our objective no longer depends on the the expert actions ⇡E and can be used for IL using only observations. For the sake of brevity, we expand on this in Section 1 in Appendix A.

heading: 6 Related Work
text: Classical IL: Imitation learning has a long history, with early works using supervised learning to match a policy’s actions to those of the expert [15, 35]. A significant advance was made with the formulation of IL as the composition of RL and IRL [29, 1, 43], recovering the expert’s policy
5The simplification to get Eq. (12) is not applicable in the online IL setting where our method differs
by inferring the expert’s reward function, then finding the policy which maximizes reward under this reward function. These early approaches required a hand-designed featurization of the MDP, limiting their applicability to complex MDPs. In this setting, early approaches [9, 31] noted a formal equivalence between IRL and IL using an inverse Bellman operator similar to our own.
Online IL: More recent work aims to leverage the power of modern machine learning approaches to learn good featurizations and extend IL to complex settings. Recent work generally falls into one of two settings: online or offline. In the online setting, the IL algorithm is able to interact with the environment to obtain dynamics information. GAIL [17] takes the nested RL/IRL formulation of earlier work , optimizing over all reward functions with a convex regularizer. This results in the objective in Eq. (3), with a max-min adversarial problem similar to a GAN [11]. A variety of further work has built on this adversarial approach [21, 10, 3]. A separate line of work aims to simplify the problem in Eq. (3) by using a fixed r or ⇡. In SQIL [33], r is chosen to be the 1-0 indicator on the expert demonstrations, while ASAF [4] takes the GAN approach and uses a discriminator (with role similar to r) of fixed form, consisting of a ratio of expert and learner densities. AdRIL [38] is a recent extension of SQIL, additionally assigning decaying negative reward to previous policy rollouts.
Offline IL: In the offline setting, the learner has no access to the environment. The simple behavioural cloning (BC) [34] approach is offline, but doesn’t use any dynamics information. ValueDICE [22] is a dynamics-aware offline approach with an objective somewhat similar to ours, motivated from minimization of a variational representation of the KL-divergence between expert and learner policies. ValueDICE requires adversarial optimization to learn the policy and Q-functions, with a biased gradient estimator for training. We show a way to recover a unbiased gradient estimate for the KL-divergence in Appendix C. The O-NAIL algorithm [2] builds on ValueDICE and combines with an SAC update to obtain a method that is similar to our algorithm described in section 4.4, with the specific choice of reverse KL-divergence as the relevant statistical distance. The EDM method [19] incorporates dynamics via learning an explicit energy based model for the expert state occupancy, although some theoretical details have been called into question (see [37] for details). The recent AVRIL approach [6] uses a variational method to solve a probabilistic formulation of IL, finding a posterior distribution over r and ⇡. Illustrating the potential benefits of alternative distances for IL, the PWIL [7] algorithm gives a non-adversarial procedure to minimize the Wasserstein distance between expert and learned occupancies. The approach is specific to the primal form of the W1-distance, while our method (when used with the Wasserstein distance) targets the dual form.

heading: 7 Experiments

heading: 7.1 Experimental Setup
text: We compare IQ-Learn ("IQ") to prior work on a diverse collection of RL tasks and environments - ranging from low-dimensional control tasks: CartPole, Acrobot, LunarLander - to more challenging continuous control MuJoCo tasks: HalfCheetah, Hopper, Walker and Ant. Furthermore, we test on the visually challenging Atari Suite with high-dimensional image inputs. We compare on offline IL - with no access to the the environment while training, and online IL - with environment access. We show results on W1 and 2 as our statistical distances, as we found them more effective than TV distance. In all cases, we train until convergence and average over multiple seeds. Hyperparameter settings and training details are detailed in Appendix D.

heading: 7.2 Benchmarks
text: Offline IL We compare to the state-of-art IL methods EDM and AVRIL, following the same experimental setting as [6]. Furthermore, we compare with ValueDICE which also learns Q-functions, albeit with drawbacks such as adversarial optimization. We also experimented with SQIL, but found that it was not competitive in the offline setting. Finally, we utilize BC as an additional IL baseline.
Online IL We use MuJoCo and Atari environments and compare against state-of-art online IL methods: ValueDICE, SQIL and GAIL. We only show results on 2 as W1 was harder to stabilize on complex environments6. Using target updates stabilizes the Q-learning on MuJoCo. For brevity, further online IL results are shown in the Appendix D.
6 2 and W1 can be used together to still have a convex regularization and is more stable

heading: 7.3 Results
text: Offline IL We present results on the three offline control tasks in Figure 2. On all tasks, IQ strongly outperforms prior works we compare to in performance and sample efficiency. Using just one expert trajectory, we achieve expert performance on Acrobot and reach near expert on Cartpole.
Mujoco Control We present our results on the MuJoCo tasks using a single expert demo in Table 3. IQ achieves expert-level performance in all the tasks while outperforming prior methods like ValueDICE and GAIL. We did not find SQIL competitive in this setting, and skip it for brevity.
Atari We present our results on Atari using 20 expert demos in Figure 3. We reach expert performance on Space Invaders while being near expert on Pong and Breakout. Compared to prior methods like SQIL, IQ obtains 3-7x normalized score7 and converges in ⇠300k steps, being 3x faster compared to Q-learning based RL methods that take more than 1M steps to converge. Other popular methods like GAIL and ValueDICE perform near random even with 1M env steps.

heading: 7.4 Recovered Rewards
text: IQ has the added benefit of recovering rewards and can be used for IRL. On Hopper task, our learned rewards have a Pearson correlation of 0.99 with the true rewards. In Figure 4, we visualize our recovered rewards in a simple grid environment. We elaborate details in Appendix D.

heading: 8 Discussion and Outlook
text: We present a new principled framework for learning soft-Q functions for IL and recovering the optimal policy and the reward, building on past works in IRL [43]. Our algorithm IQ-Learn outperforms prior methods with very sparse expert data and scales to complex image-based environments. We also recover rewards highly correlated with actual rewards. It has applications in autonomous driving and complex decision-making, but proper considerations need to be taken into account to ensure safety
7normalizing rewards obtained from random behavior to 0 and expert to 1
and reduce uncertainty, before any deployment. Finally, human or expert data can have errors that can propagate. A limitation of our method is that our recovered rewards depend on the environment dynamics, preventing trivial use on reward transfer settings. One direction of future work could be to learn a reward model from the trained soft-Q model to make the rewards explicit.

heading: 9 Acknowledgements
text: We thank Kuno Kim and John Schulman for helpful discussions. We also thank Ian Goodfellow as some initial motivations for this work were developed under an internship with him.

heading: 10 Funding Transparency
text: This research was supported in part by NSF (#1651565, #1522054, #1733686), ONR (N00014-19-12145), AFOSR (FA9550-19-1-0024) and FLI.

Review summaries:
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.  

This work aims at improving the sample efficiency in terms of expert demonstrations in IL. In order to do so, it transforms the usual max-min objective of IRL into a single maximization objective by expressing both the reward and the policy in terms of the Q-function. The methods are supported with proofs and tested on both discrete and continuous action environments for the online setting and on simpler environments for the offline setting.

This paper proposes to get around the typical two-stage inverse RL pipeline of (1) learn a reward model and (2) optimize it, by encapsulating both rewards and dynamics in a single Q function, learned by imitation. This reframing of IRL preserves the standard properties of IRL but is simpler to optimize and avoids the need for adversarial training used by other methods. This results in a simple imitation learning method that is competitive with or outperforms reasonable baselines on a few standard benchmarks. The proposed algorithm has other nice features, including the ability to easily recover the learned reward and to be trained both offline and online, and it can be easily rephrased to obtain a method for imitation from observations.

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.