Inverse Reinforcement Learning with the Average Reward Criterion
Authors: Feiyang Wu
Jingyang Ke
Anqi Wu
Venue: NeurIPS 2023 poster (2023)
Keywords: Machine Learning
Reinforcement Learning
Inverse Reinforcement Learning
Markov Decision Process
stochastic optimization
complexity analysis
TL;DR: We develop first-order methods for solving average-reward RL and inverse RL with theoretical guarantees.

Abstract: We study the problem of Inverse Reinforcement Learning (IRL) with an average-reward criterion. The goal is to recover an unknown policy and a reward function when the agent only has samples of states and actions from an experienced agent. Previous IRL methods assume that the expert is trained in a discounted environment, and the discount factor is known. This work alleviates this assumption by proposing an average-reward framework with efficient learning algorithms. We develop novel stochastic first-order methods to solve the IRL problem under the average-reward setting, which requires solving an Average-reward Markov Decision Process (AMDP) as a subproblem. To solve the subproblem, we develop a Stochastic Policy Mirror Descent (SPMD) method under general state and action spaces that needs $\mathcal{O}(1/\varepsilon)$ steps of gradient computation. Equipped with SPMD, we propose the Inverse Policy Mirror Descent (IPMD) method for solving the IRL problem with a $\mathcal{O}(1/\varepsilon^2)$ complexity. To the best of our knowledge, the aforementioned complexity results are new in IRL with the average reward criterion. Finally, we corroborate our analysis with numerical experiments using the MuJoCo benchmark and additional control tasks.

Inverse Reinforcement Learning with the Average Reward Criterion
Feiyang Wu feiyangmu@@atech.cmm edu
Jingyang Ke j jingyang ke@gatech. edu
Anqi Wu
angisu@gatech. edu School of Computational Science and Engineering College of Computing Georgia Institute of Technology Atlanta, Georgia 30332
Abstract
We study the problem of Inverse Reinforcement Learning (IRL) with an average- reward criterion. The goal is to recover an unknown policy and a reward function when the agent only has samples of states and actions from experienced agent. Previous IRL methods assume that the expert is trained a discounted environment, and the discount factor is known. This work alleviates this assumption by proposing an average-reward framework with efficient learning algorithms. We develop novel stochastic first-order methods solve the IRL problem under the average- reward setting, which requires solving an Average-reward Markov Decision Process (AMDP) as a subproblem. To solve the subproblem, we develop a Stochastic Policy Mirror Descent (SPMD) method under general state and action spaces that needs (1/=)) steps of gradient computation. Equipped with SPMD, we propose the Inverse Policy Mirror Descent (IPMD) method for solving the IRI. problem with a (1/e²2 complexity. To the best of our knowledge, the aforementioned complexity results are new in IRI. literature with the average reward criterion. Finally, we corroborate our analysis with numerical experiments using the MuloCo benchmark and additional control tasks.
Introduction
Reinforcement Learning (RL) problems are frequently formulated Markov Decision Processes (MDPs). The agent learns a policy maximize the reward gained over time. However, in numerous engineering challenges, we are typically presented with set of state-action samples from experienced agents, or experts, without an explicit reward signal. Inverse Reinforcement Learning (IRL. aspires to recover the expert's policy and reward function from these collected samples.
reduces the IRL problem to predicting an expert's behavior without estimating the reward signal. Such formulation is undesirable in some scenarios when we either wish to continue training or analyze the structure of the agent's 's reward signal, e.g., in reward discovery for animal behavior study [3I] 34 26 [4]. Furthermore, as popular methodology in Imitation Learning. Generative Adversarial Network (GAN) 15 42 suffers from unstable training due mode collapse [2). The sampling-bast based Bayesian IRL (BIRL) approach 8 treats the reward function as a posterior distribution from the observation. This approach is conceptually straightffrward but suffers from slow convergence of sampling. Another line of research reformulates
Methods like Imitation Learning 29
15 10
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

the IRL as RL with general utility functions formulation does not learn reward functions.
36 39 12
While theoretically attractive, such
Furthermore, inferring rewards poses an additional challenge: without any additional assumption, the same expert behavior be explained by multiple reward signals 20 To resolve this ambiguity, the authors propose adding the principle of Maximum Entropy that favors a solution with higher likelihood of the trajectory. In the following work НТ the objective is replaced by maximizing the entropy of the policy, which better characterizes the nature of stochastic sequential decision-making body of work focuses on alleviating the computational burden of the nested optimization structure of under the Maximum Entropy framework. In [I] the authors skipped the reward function update stage by representing rewards using action-valle functions. f-IRL [23] minimize the divergence of state visitation distribution for finite-horizon MDPs. More recently, in [7]], the authors propose dual formulation of the original Maximum Entropy IRI. and a corresponding algorithm to solve IRI. in the discounted setting.
Despite the success in practice, theoretical understanding of RL/IRL methods is lacking. For RL, specifically for Discounted Markov Decision Processes (DMDPs), analysis of policy gradient methods for solving DMDPs is catching up only until recently 2 the Average-reward Markov Decision Processes (AMDPs) remains very limited, where the policy aims to maximize the long-run average reward gained. Li al. 20 propose stochastic first-order methods for solving AMDPs with rate of convergence (log(( ") for policy optimization, but is constrained to finite states and actions with policy explicitly represented in tabular forms. Moreover, for IRL, the only known finite time analysis are B5
16 5
19
[
Meanwhile, understanding of
2 for either solving DMDPs, or finite-horizon MDPs. There are no well-established convergence analyses for IRL with AMDPs. Moreover, prior works only focus on MDPs with finite state and action spaces, which do not fully address many real-life problems that inherently feature continuous elements (e.g., robotics [I7]). Such settings further require general function approximations, the usage of which lacks theoretical understanding in recent work as well.
37
38
In summary, all previous works either do not fundamentally address our problem, which necessitates the learning of reward functions, or they miss key elements such as general state and action spaces, general function approximation, and the analysis incorporating approximation errors. Each of these facets demands substantial effort. In this paper, we focus on convergence analyses for average-award MDPs. We extend the RL algorithm 20 for AMDPS to general state and action spaces with general function approximation and develop an algorithm to solve the IRL problem with the average-reward criterion. Analyses of both algorithms are new in the literature to the best of our knowledge.
1.1
Main contributions
Stochastic Policy Mirror Descent for solving AMDPs: Our rescarch introduces the Stochastic Policy Mirror Descent (SPMD) algorithm, a novel approach designed to tackle Average-reward Markov Decision Processes (AMDPs) that involve general state and action spaces. Leveraging performance difference lemma, we demonstrate that the SPMD algorithm converges within O(E -1 steps for specific nonlinear function approximation class, and in more than O(e steps for general function approximation.
Inverse Policy Mirror Descent for solving Maximum Entropy IRL:
We expound a dual form of the Maximum Entropy IRL problem under the average-reward framework with the principle of Maximum Entropy. This dual problem explicitly targets the minimization of discrepancies between the expert's s and agent's 's expected average reward. Consequently, we propose a first-order method named Inverse Policy Mirror Descent (IPMD) for effectively addressing the dual problem. This algorithm operates by partially resolving an Entropy Regularized RL problem at each iteration, which is systematically solved by employing SPMD. Drawing upon the two-timescale stochastic approximation analysis framework we present the convergence result and establish a O(E rate of convergence.
Our RL and IRL methodologies have been tested against the well-known roboties manipulation benchmark, MuloCo, as a means to substantiate our theoretical analysis. The results indicate that the proposed SPMD and IPMD algorithms generally outperform state-of-the-an algorithms. addition, we found that the IPMD algorithm notably reduces the error recovering the reward function.
Numerical experiments:

---

2 Background and problem setting
2.1 Average reward Markov decision processes
An Average-reward Markov decision process is described by tuple M (,, where S denotes the state space, A denotes the action space, P is the transition kernel, and e is the cost function. At each time step, the agent takes action a E A the current state policy π S + A. We use II as the space of all feasible policies. Then the agent moves to the next state S with probability P(s") s, o), while the agent receives an instantaneous cost c(s, a) (or reward r(s, a) = -(ss a)). The agent's ' goal is to determine policy that minimizes the long-temn cost
according its
(s(s) lim E
ET (c(se, (c(s, 0z) (s:)) =
(I)
where hdd is a closed convex function with respect to the policy i.e., there exists μl 2 0 such that (s) - hh (s) + (VA" (s, x(1s) -T' r-(+s)) 2 pr D(T(s)
(2)
where Vh denotes the subgradient of h at r' and ().,)) is the Bregman distance, i.e.,
D(a2, 01) == (a1) - (w(22) (w(a2), 01 - a2)l 2 1 || a2ll? for all a1, 02 E A.
(3)
Here w : A + R is a strongly convex function with the associated norm || || in the action space A and let denote ||. its dual norm. In this work also utilize span semi-norm [27] [[vll = max e min v. Vo E TR". hf = 0. Eq. reduces to the classical unregularized AMDP If 1e (s) = [- log -| =: H()(s)), i.e. the (differential) entropy, Eq. defines the average reward of the so-called entropy-regularized MDPs.
In this work, we consider the ergodic setting, for which we make the following assumption formally:
Assumption 2.1. For any feasible policy æ, the Markov chain induced by policy is ergodic, The Markov chain is Harris ergodic in general state and action spaces (see (221). The stationary distribution xT induced by any feasible policy exists and is unique. There is some constant number such that nxT (s) 2 T.
0 < T <
As result of Assumption 2.1 for any feasible policy æ,, the average-reward function does not depend on the initial state (see Section 8 of (27l). Given that, one can view n" as function of æ. For a given policy we also define the basic differential value function (also called bias function; see, e.g., (27
p= (*(() = = E (st, as) + h- (st) - -p so = s,0, ~ x(-(st), (o)lss).sss Se+1 ~ P(-[st, at)
(4)
and the basic differential action-value function (or basic differential Q-function) is defined as B at) + htT (s:) pz 50 0, ar T(-)s)) 8141 ~ P(-|se. az)
(5)
Moreover, we define the sets of differential value functions and differential action-value functions (differential Q- Q-functions) as the solution sets of the following Bellman equations, respectively,
(=(s) = [IQ"" a) T-|ss)]
(6)
Q" a) = c(s. + (s)
P(ds' (s,a)
(7)
Under Assumption 2.1 the solution of Eq. 6 (resp., Eq. 7 is unique up to an additive constant. Finally, our goal in solving an AMDP is to find an optimal policy x² that minimizes the average cost:
min pT s.t. *(188) IT, VsES.
(AMDP)
2.2 Inverse Reinforcement Learning
Suppose there is a near-optimal policy, and we are given its demonstrations Ç= (=:((s, (()))))))) IRL aims to recover a reward function such that the estimated reward best explains the demonstrations. We consider solving the IRL problem under the Maximum Entropy framework (MaxEnt-IR) -RL),, which
3

---

aims to find a reward representation that maximizes the entropy of the corresponding policy and incorporates feature matching as a constraint. Formally, the MaxEnt-IRe problem is described as
H(T)) == (x.))) log म (a|s)]
(MaxEnt-IRL)
s.t. (s.a)-dd [s)(s,a]] (s.a)~dd [(](s,a)]
where d and d' denote the state-action distribution induced by the expert and current policy n and respectively, and g(88 a) E IR" denotes the feature of a given (x, a) pair. If we assume that for given state-action pair, the cost a linear function to feature, i.e., c(8, a; 0) 0 "p(ss a) for some show that the parameter 0 is the
parameter 9 with the same dimension of the feature (s a), dual multiplier of the above optimization problem, as the Lagrangian function can be written as
[(TT, 0) H(Tx) + (a.a)raa (c(s,a;0) B(aaa)dd [c[s, az 0)]
Therefore, the dual problem is formulated
min (( = E (e (c(s,a;0) B(saa)dd [c(s, a(s,,;;))
(Dual IRL)
s.t. π = arg max E (s.a)=a'' c-sss a; 0)] + H(T').
Notice that to find we need to solve an Entropy Regularized Reinforcement Learning problem, i.e.,
solving AMDP with the regularizer set to be the negative entropy, h" -H(T). To this end, we first propose a Stochastic Policy Mirror Descent (SPMD) method for solving [AMDI] The solution methods are presented in section S Then we introduce an Inverse Policy Mirror Descent (IPMD) algorithm based SPMD for the inverse RL problem (Dual IRL introduced in section see that we only need to solve the subproblem AMDP partially
Dalving/AMDP
We will
Stochastic Policy Mirror Descent for AMDPs
This section proposes SPMD method to solve Regularized Reinforcement Learning problems for AMDPs. SPMD operates in an actor-critic fashion. In each step, the agent evaluates its current policy (critic step) and performs policy optimization step (actor step). In this work, assume there is a way to perform the critic step using standard methods, e.g., Temporal Difference learning using neural networks. Further discussion on implementation is included in section 3.1 We will focus designing a novel actor step and providing its complexity analysis.
Our policy optimization algorithm is motivated by the following performance difference lemma, which characterizes the difference in objective values of two policies म. n'.
Lemma 3.1. (Performance Difference) Assume that Assumption 2.1 holds. For any n. n' 11
p= pp" -p-= 0 - (s)
(ds))
(ss),
(9)
where
36 (s. x (s)) = (s) (T
+ -
(s).
(10)
Proof can be found in A. The above lemma shows the gradient of the objective function with respect to action a is Q" (s, (8)
+ hd (s), i.e.,
f
V Vap"
(11)
The existence of the above equation requires the differentiability of Q, h and locally Lipschitz continuity of x [see 24 25]. Note that z can be seen as some generalized advantage function. Lemma shows that the gradient of the objective relates to differential Q-functions. This inspires us to update the policy in the following mirror-descent style:
T (*)1(8) = arg min @ as) h*)) 1 ())
(12)
aeA
where тe is some predefined step size. The complexity analysis of this type of algorithm has been thoroughly studied in 20 in tabular forms or linear function approximation settings, i.c., lookup

---

table represents the policy, and the differential Q-functions are approximated as linear functions with respect to some feature space. practice, especially the Deep Reinforcement Learning community, both policy and the Q-functions are represented by neural networks. Thus, novel analysis required for general function approximation. Additionally, the exact value of Q" and Vw(x) in Eq. 19 can only be estimated through approximation. Note that Vw(TT) rises inside the Bregman distance (T a), referring to Eq. 13 We consider their stochastic estimators calculated from state-action sample pairs Ç. denoted as < a; o). Ve (T (s): Φ) respectively, where & denotes the parameters of the method of choice, for example, weights and biases in neural network. For the rest of the paper, we abbreviate Q. (s, a: Φ) as ((s, a; σ) when the context is clear. We described the Stochastic Policy Mirror Descent (SPMD) algorithm in Algorithm
Algorithm 1: The Stochastic Policy Mirror Descent (SPMD) algorithm for AMDPs
Input: Initialize random policy my and step size sequence (m)
2: for 00 1..... K do
3: Sample collection: & ((11, at, (t),12))
4: Critic step: Approximate Q (w)(+(()) with
a= R @uack (s, a;)); vw(Tx(s)) = (s(m)(s)
(13)
(14)
5: Actor step: Update the policy
T = arg min @ (8, (3,0)))) + 16) + Ti I. ((VW(T + +a)).
(15)
a
6: end for
The following paragraphs present the complexity analysis for Algorithm Note that we assume the approximated Q function might not be convex but weakly convex, i.e.,
Ö(s, a; Φ) + (oD(mo()),
(16)
convex w.r.t. a E A for some μo 2 0 and To is the initialized policy. This assumption is general any differential function with Lipschitz continuous gradients is weakly convex. Moreover, assume that (s), +)) and (s, -) are Lipschitz continuous with respect to the action.
Assumption 3.2. For all a, 02 E A there exist some constants My, MR. Ma such that
(s) - n22 (s) < Mallar a2|1
(17)
(0(s, ((s, a2; の) < Mallar - a2 ||
(18)
(s, - (s,(saa) < Malla - 02 ||
(19)
and ((s, a; σ) is -o we-weakly convex. i.c., Eq. 16 is convex.
Note that the strong convexity modulus of the objective function in Eq. 15 can be very large since Tkk
can be small, in which case the subproblem Eq. 15 is strongly convex, thus the solution of Eq. 15
is
unique due to strong convexity [88]
Our complexity analysis follows the style of convex optimization. We begin the analysis by decom- posing the approximation error in the following way
50 (s, a) = Q (s, a; (he) - T (s, a),
(20)
02 (s) = 1 v (T) ox) 1 Vw (s)),
E V (TA
1 Vw (m)) + 1 M (TT(S): E ( Eo (T (s): the
(21)
We refer to the first two and last two approximation error terms in 8 as deet and 8 respectively, as they represent& deterministic approximation error which we cannot reduce and a stochastic error related to the variance of the estimator. They represent the noise introduced in the critic step. To make
5

---

the noise tractable and the entire algorithm convergent, we must assume they are not arbitrarily large. But before making more assumptions about these error terms, we must distinguish the case when Q + h is convex and when it's not, i.e., when fr 2 0 or μ < O, where
PA Má for simplicity. We can obtain globally optimal solutions if μ 2 O. While in the other case, only stationary points can be obtained, e.g., both 11th (s, x' (s)) and D TT++) approach O. Note that v relates to the progress we make in each iteration. See A.6 for more detail. These two cases require different assumptions on the error terms, so they must be treated differently. We start with the optimistic case when μ 2 0, under which we make the following assumption about the approximation error 68 (8, a) and 0% (s): Assumption 3.3. When μ 2 O, the critic step has bounded errors, i.e., there exist some constants 5 ,,,5 ot such that
Ecr.s-n' (s, Th m + ( (+))) S S
(22)
= 1100 (+)] S 5 "Bee -110% 1110% [|3] S (at))
(23)
where Ç& is the samples we collect in each iteration.
With the above assumption, we present the complexity result of SPMD when μ 2 0.
Theorem 3.4. Suppose 2.1
3.2
3.3 hold and the step size T7 satisfies
S (
+ nhk - ) 2 1,
(24)
for some 3x 2 0. Then after running the SPMD algorithm for K iterations, we have
K-1 1k=IBB)- Bx)
- Bx p²) + 1 + nn 1
fo (TTO π*) > 3 (2M +M + Mn))/2++ (00))) < 12 Bx
K-1
5 + D
(25)
where D = maxaj.age E.A D(a1, 02).
Proof can be found in A.4 Equation 25 encapsulates a convergence overview without specifying explicit step size choices, To be specific, og,, 32 are both step sizes, and the central term > Ps: p (distance to the optimal value) plus (TTR,, x*) (the distance of the current policy to the optimal policy) will shrink as grows. The sum is bounded by some combination of ag and 33 plus an irreducible function approximation error s. The left-hand side is a weighted average of function value errors plus an average error from step 1 to K. The following result summarizes the convergence rate using particular step size choices. We obtain the first part of Corollary 3.3 by fixing the number of iteration K and optimizing the right-hand side of Eq. 25
We get the rest by straightforward computation.
Corollary 3.5. a) If μ = 0, and mh
D(al.=")
and B+ = 1, k = 1,2,3... K.
V K (2Mg+Ma++M)) M +(o
then the average reward will converge to the global optimal p²
I (p p0))/11 ( = OKK)1/2) 8 + 5 DA
(26)
If μ > O, and mhe # and Bk = L. then the average reward will converge the global optimal
1 1 ( p²)/(1 - TT) + MEE x ((K-))(K++))))))) +5 + 5h DA
(27)
If 1. > 0, 2 and Br = k + 1, then the average reward will converge to the global optimal
2k
K(K+I) Ck=o b+ (pTH - p²) DEELD(TKK T*]] + + 45 DA (28)
Remark The above result shows that if we carefully choose the function approximation Q and the regularizer h so that > O, we need O(e-1) number of SPMD iterations to obtain & precision solution, i.e., > 11 PT p² + E[D(TK, T T²)] I s. The complexity deteriorates to O(= if
J === Notice that ς and 5t5 are irreducible solely by policy optimization, as they arise due to function approximation and stochastic estimation errors.
When μ < global optimal solutions are unobtainable. The best we can do is to find a stationary point, that is when vo and D(Tk,, 1k+1) tend to O. In this case, we only need to assume that the estimation error terms on
guedet (s) and 0% (us.st are bounded, i.c..

---

Assumption 3.6. When μ < O, the critic step has bounded errors for any s E S, i..., there exist some constants 5t) Tt such that
|18g
(s)ll 5 and |10g (s)l S 8".
(29)
Theorem 3.7. Let the number of iterations K be fixed. Suppose that 2.1 and
3.6 hold. Also
assume that μ < O and that 1 min(lu//2 1 ...... K --1 1. Then for any E S, there exist iteration indices found by running K iterations of the SPMD method such that both the generalized advantage function and the distance between the iterates will converge to 0, i.e.
3 (8, (T(s))+(())) = O(K-)) O(K
+T Da/(1)((1
(30)
1 ((()) (s)) + (M + ( (S))
- +5
(31)
Remark This is the most general case, we don't assume any restrictions on the function approxi- mation class. As expected, the complexity bound becomes worse compared to previous results. To reach a e-precision stationary point, we require at least O iterations. Note that tà 0 indicates that Q (s, (s)) has virtually no difference than TTh (s). This has a similar implication of DTT++1 (s), (s)) approaching 0 which implies that reaching stationary point.
3.1 Practical algorithm for Entropy Regularized AMDP
The above SPMD algorithm computes (s) for every state. each iteration, we need to solve subproblem Eq. [1] for every state (or for any state encountered). In practice. we prefer performing each SPMD iteration in a mini-batch style, i.e., approximately solving Eq. [[] some trajectories collected. If the regularizer h and w are the negative entropy, we can approximately solve the SPMD)
actor step Eq. 15 by
Tk+1 arg min KL (T()) z(oo) exp( Ộ (s.a; *ell
(32)
where Z(b)) is some normalization constant and ÔT (s, a; (h)) =
T (s, az dh))
1++
T+7 log T& (s). In practice, apply the so-called reparameterization trick [ to obtain an unbiased gradient estimator. Let the policy be parameterized by
<alss π (fele; s)lss, S
(33)
where is input noise sampled from fixed distribution; represents the parameters of the policy. Then we can approximately perform the actor step by solving the following problem:
min (.,.)~<.. [oog( [og(=(
²T fe(c:s))
(34)
For large action spaces, the reparameterization trick efficiently approximates the solution of Eq, [5 As for the critic step, we can approximate Q by minimizing the Temporal Difference (TD) error, i.e., solving the following problem:
+ - 2 - (x)))))))))
min E (eaac
(35)
where 7nt is the estimated average cost, e.g., by taking p" = K Ee(s,a) 11(
Inverse Policy Mirror Descent for IRL
Equipped with efficient solver for AMDPs presented in section 3, can now solve the IRL problem in the form of Dual IRL). We suppose d parameterizes the reward function with arbitrary methods, e.g., neural networks. To solve the dual problem, we update the parameter by performing gradient descent step 0x+1 0₂ a& V (() where ag is some predefined step size. The gradient computation of the dual objective function is
V (0) (( [Vcc(s, a; 0)] - E(raa)ar [Vec(s, a; 0))
(36)
7

---

Algorithm 2: The Inverse Policy Mirror Descent (IPMD) algorithm
Input: Initialize random policy my and step size sequence {ax}
2: for 0, 1..... K do
3: Collect samples {(s,, ae) and compute c(s1, az; 0x)) based on current reward estimation.
4: Critic step: Implement a policy evaluation algorithm evaluation QT
(s,
QT (s. a).
(37)
5: Actor step: Update the policy by solving the following for every E S
(T()) 2 cxp( (xsp(2" (s,
T (8) arg m
(38)
6: Dual Update step:
Perform the stochastic update
0k+1 O2 -.
(39)
7: end for
Without accessing the transition kernel P, the true gradient is unavailable. So we use a stochastic ap- proximation instead, denoted Ib (0,; (f - g(0s ( where () = t En Ve(st, az; ox; 8) the stochastic estimator of the gradient of the average reward. We describe the proposed Inverse Policy Mirror Descent (IPMD) algorithm in Algorithm 2 In this section, provide our analysis that captures the algorithm behaviors across iterations since, in each iteration, Q is only an approximation of the true differential Q-function, which itself alters due the change of the reward estimation. The idea of the proof is based on the Lipschitz continuity of the iterates, as it controls the difference of these functions across iterations. We make the following formal assumptions for such a purpose.
Assumption 4.1. For any E S, A, the gradient of the reward function is bounded and Lipschitz continuous, i.e., there exist some constant real numbers T.r, Tug so that the following holds:
||Vo(s, a; 0) s L,, [|Vss a; 01) Vc(s, a; 02) I Lgl10 02112-
Note that Q is also function of 0. As shown in section ② can be parameterized but we only denote the estimator as (s o; 0x) since its parameters do not contribute to the analysis. Assumption 4.2. Suppose that least one of S,A A is continuous, we assume that
max 0
V Q (s, a; 0)]l < Lag,
(40)
where T.g is some positive constant. For convenience we denote Q" a; Ox)) also 2% (s,
We further assume that the estimation error from using O is bounded.
Assumption 4.3. For any & E S. E A. there exist some constants Sk. S, Vk, v such that the following inequalities hold:
|IoT - TTe
E5 10 Q ssp SS
Ese
S S
v
(41)
sp,.00
Notice that this error bound subsumes both the estimation errors and approximation errors. In general, the Bellman operator is nonexpansive in the average-reward setting, so analysis based on the contraction property of the Bellman operator with the infinity norm, as what is used in [ fails in this case. To this end, we assume that the operator is span contractive, i.e.,
Assumption 4.4. In the critic step, there's a way to construct a span contractive Bellman operator T such that there exists 0 < 7 < 1 for any differential Q-fanctions Q1, Q2.
||TQ1 TQ2 < - sp.o0.
(42)
In A.6 we provide conditions when the Bellman operator has J-step span contraction. Now we can discuss the convergence results the following theorems. The convergence of the differential Q-function is characterized in the following result.

---

Theorem 4.5. Suppose that assumptions 2.7 4.1 4.4 hold. If VK. the differential Q-function
will converge the optimal solution for given reward functiin parameterized by 0k. Le.,
E k=
k-1k
- Q0
NoN
sp..0 = (
1-71
(43)
where og 0 is some step size chosen, and y is some constant defined in Assumption
Proof can be found in A.7 The above theorem shows that the differential Q-function will approach the optimal differential Q-function Q% with respect to the current reward estimation 0x. If the reward estimation is accurate, is the differential f-function. Finally, in the next theorem, show that the policy converges to the optimal policy for a given reward function, and the reward function approximation converges to stationary point. Proof can be found in A.8
Theorem 4.6. Suppose that assumptions 2.1 4. 4.
hold. If = /K and run the proposed IPMD algorithm K iterations, the algorithm will produce near stationary solutions for the reward function and the optimal policy for such reward fanction. Specifically, the following holds,
k Ek= [E llog T - log To, 20 = ((K-))+0(()
1-7-
(44)
K ((I) O(K- ((K-)) 2.
(45)
where ao. v, Loss Ly Ca some constants and S |A| some measure of the state and action space.
Remark First, the above theorem shows that the policy will converge to the optimal policy given the reward parameterization, although each iteration's 's reward parameter differs. Second, the reward parameter will converge to a stationary point of (Dual I as the objective can be highly nonconvex under general function approximation. Third, for general AMDPs without 1-step span contraction, we can use a J-step Bellman operator (see Appendix A.7) for policy evaluation to maintain a O(s complexity for the entire algorithm.
5 Numerical experiments
In this section, we showcase the performance of the proposed SPMD and IPMD algorithms. Our code can be found at htttps://anonymooootet our implementation in the Appendix A.9
4open. scwence/rrIPMD--DDSSS See more detail about all
MuJoCo Robotics Manipulation Tasks for RL This experiment tests RL agent's performance
on robotics manipulation tasks. Our SPMD algorithm is based on the stable e- eetee--baeeessss [28). We compare the performance of our algorithm with Soft Acrot-Critic (SAC) [[]] implemented in 28 The policy network employs two fully connected hidden layers of dimension 256 each, taking actions as imput and outputting a distribution. Both the Q network and reward function share the same architecture, with ReLU activation used in hidden layers. A double Q-learning technique is used minimize overestimation 13.. During training, we found that setting the entropy coefficient term to 0.01 makes training stable and efficient. The learning rate is 34- Each step of the algorithm samples 512 state-action sample pairs. Table reports the numerical results of each model. Our proposed SPMD achieves on-par performance with SAC and exceptionally better performance in the Humanoid environment.
MuJoCo benchmark for IRL In this experiment, we compare the proposed IPMD method with IQ-Learn I1] and f-IRL 2 The authors of MI-IRL have not released its implementation the time we experimented. Nevertheless, the performance of MI-IRI. is comparable to f-IRI.. It is observed in the literature that imitation learning algorithms have inferior performance 37 Hence we omit these methods. Table reports the numerical results of each model. Note that we cannot record a competitive result for IQ-Learn with Humanoid as claimed in the original paper, which we highlight using * The result shows that IPMD outperforms in a variety of environments. One possible reason Ant is falling behind is that Ant has more ground contact since it has more legs. This will impact the mixing time of the MDP and our assumption on the 1-step contraction. Compared to Half-Cheetah, or Humanoid,
harder for the Ant to transition from an arbitrary state to another arbitrary state as it involves multiple legs working together. One remedy is that we construct J-step contractive operator for policy evaluation, as done in Z The success of humanoid possibly due to slightly different policy evaluation scheme compared to the discounted setting, where the entropy

---

Table 1. MuJoCo Results for RL. The average
Table 2. MuJoCo Results for IRL. The aver-
performance of RL agents across five runs,
age performance of IRL agents across five runs.
Task
SAC
SPMD (Ours)
IQL
J-IRL
IPMD (ours)
Expert
Hopper
3876 = 78
3619 = 5
1909
3083
3564
4115
Half-Cheetie
13300 46
13025 = 37
9132
11259
12634
15467
Walker
6115 + 41
4454 + 26
5155
5378
5423
5323
Ant
5118 + 57
5230 + 118
3486
5460
4053
5783
Humanoid
6923 + 11
10249 =
661
3580
7379
7137
Perdukmmer
10 IPMD
IPMO
Brro
067
Tphoodl
Episae
Figure Reward recovery performance IPMD and IQ-Learn. The lower the bar, better the performance.
term from the policy no longer plays a part, as the term c cancels out the additional regularization and entropy of the policy. We suspect this brings more stable training and thus higher performance.
Reward Recovery Finally, we compare the proposed IPMD method against IQ-Learn on recovering the expert's 's reward function in two environments: Pendulum and LunarLanderComtinut implemented in The result of the experiment is shown in Figure ! Using state-action pairs from 4 different episodes, reward the expert uses). The experts are trained with a discount factor 0.99. IQ-Learn's 's discount factor is 0.99. The result shows IPMD's superiority in reward recovery.
both
compare the span seminorm of the predicted reward and the true reward (the
Discussion
In this paper, we formulate the Maximum Entropy IRL, problem with the average-reward criterion and propose efficient algorithms to solve it. We devise the SPMD method for solving RL problems in general state and action spaces beyond the scope of linear function approximation with strong theoretical guarantees. We integrate this method develop the IPMD method for solving the entire IRL problem and provide convergence analysis. To reach a precision stationary point, IPMD requires (e--2) policy optimization steps.
We also notice there are possible improvements from this work. For example, we impose various assumptions both the function approximation classes and strong reliance the accuracy of the critic step. is natural to seek more robust methods when facing inaccurate estimations. Additionally, we impose continuity assumption in various Assumptions (e.g. Assump 4.2 Such assumptions violated neural networks are used as function approximations. In practice, is possible to alleviate the issue by "clipping" the gradient norm when performing gradient descent, which proves effective in related works 30). Furthermore. we only consider uniform ergodic chains. It is interesting to see how our method would hold under the unichain multi-chain setting. Lastly, the analysis based span contraction might be improved for general MDPs. This requires nontrivial work and thus we leave it for future development.
be
References
[1] Anas Barakat, Ilyas Fatkhullin, and Niao He. Reinforcement learning with general utilities: Simpler variance reduction and large state-action space. arXiv preprint arXiv:2306. 01854, 2023.
10

---

[2] Jalaj Bhandari and Daniel Russo. A note on the linear convergence of policy gradient methods. arXiv preprint arXiv:2007.11120, 2020.
[3] Vivek S Borkar and Vijaymohan Konda. The actor-critic algorithm as multi-time-scale stochastic approximation. Sadhana, 22:525-544 543, 1997.
[4] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
[5] Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research, 2021.
[6] Alex J Chan and Mihaela van der Schaar. Scalable bayesian inverse reinforcement learning. arXiv preprint arXiv: 102.:200.. 2021.
[7] Zaiwei Chen, Sajad Khodadadian, and Siva Theja Maguluri. Finite-sample analysis of off-policy natural actor-critic with linear function approximation. arXiv preprint arXiv:2105.125400 2021.
[8] Jaedeug Choi and Kee-Eung Kim. Nonparametric bayesian inverse reinforcement learning for multiple reward functions. Advances in neural information processing systems, 25, 2012.
(9] Jaedeug Choi and Kee-Eung Kim. Hierarchical bayesian inverse reinforcement learning. IEEE transactions on cybernetics, 45(():793-805, 2014.
[10] Yan Duan, Marein Andrychowicz. Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
[11] Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song. and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34: 4028 4039,4039,0000 2021.
[12] Matthicu Geist, Julien Pérolat, Mathicu Lauriere, Romuald Elie, Sarah Perrin, Olivier Bachem, Rémi Munos, and Olivier Pietquin. Concave utility reinforcement learning: the mean-field game viewpoint. arXiv preprint arXiv:2106.03787, 2021.
[13] Tuomas Haarnoja Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint 1rXh::1812.05905 2018.
[14] Tsubasa Hirakawa, Takayoshi Yamashita, Toru Tamaki, Hironobu Fujiyoshi, Yuta Umezu, Ichiro Takeuchi, Sakiko Matsumoto, and Ken Yoda. Can ai predict animal movements? filling gaps in animal trajectories using inverse reinforcement learning. Ecosphere, 911))e02447, 2018.
[15] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
[16] Sajad Khodadadian, Zaiwei Chen, and Siva Theja Maguluri. Finite-sample analysis of off- policy natural actor-critic algorithm. In International Conference Machine Learning, pages 5420-5431. PMLR, 2021.
[17] Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32(11):1238 2013.
[18] Guanghui Lan. Policy optimization over general state and action spaces. arXiv preprint arXiv:2211.16715 2022.
[19] Guanghui Lan. Policy mirror descent for reinforcement learningg Linear convergence, new sampling complexity, and generalized problem classes, Mathematical programming, 198(1): 1059-1106, 2023.
[20] Tianjiao Li, Feiyang Wu, and Guanghui Lan. Stochastic first-order methods for average-reward markov decision processes. arXiv preprint arXiv:2205.05800 0050000000 2022.
11

---

[21] Alberto Maria Metelli, Filippo Lazzati, and Marcello Restelli. Towards theoretical understanding of inverse reinforcement learning, arXiv preprint arXiv:2304. 12966,, 2023.
[22] Sean P Meyn and Richard L. Tweedie. Markov chains and stochustic stachllicsstaiility. Springer Science Business Media, 2012.
[23] Tianwei Ni, Harshit Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Ben Eysenbach. f-irl: Inverse reinforcement learning via state marginal matching. In Conference on Robot Learning, pages 529-551 PMLR, 2021.
[24] G Pflu.. On-line optimization of simulated markovian processes. Mathematics of Operations Research, 15(3):381-39) 1990.
[25] Georg Ch Pflug. Derivatives of probability measures and upplications to the opti- mization of stochastie systems. In Discrete Event Systems: Models and Applications: IIASA Conference Sopron, Hungary, August 3-7, 1987, pages 252-274. Springer, 2006.
[26] Robert Pinsler, Max Maag, Oleg Arenz, and Gerhard Neumann. Inverse reinforcement learning of bird flocking behavior. In ICRA Swarms Workshop, 2018.
[27] Martin I. Puterman, Markov decision processes: Discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[28] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-basclines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021.
[29] Stefan Schaal. Is imitation learning the route to humanoid robots? Trends in cognitive sciences, 3(6):233-244 1999.
[30] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 763477 2017.
[31] Can Eren Sezener, Eiji Uchibe, and Kenji Doya. Obtaining reward functions of rats using inverse reinforcement learning. 2014.
[32] Ziyu Wang, Josh S Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation of diverse behaviors. Advances in Neural Information Processing Systems, 30, 2017.
[33] Tengyu Xu, Zhe Wang, and Yingbin Liang. Improving sample complexity bounds for (natural) actor-critic algorithms. Advances in Neural Information Processing Systems, 33:4358 -3699, 2020.
[34] Shoichiro Yamaguchi, Honda Naoki, Muneki Ikeda, Yuki Tsukada, Shunji Nakano, Ikue Mori, and Shin Ishii. Identification of animal behavioral strategies by inverse reinforcement learning. PLoS computational biology, 14(5):e100612) 2018.
[35] Tom Zahavy, Alon Cohen, Haim Kaplan, and Yishay Mansour. Average reward reinforcement learning with unknown mixing times. In The Conference on Uncertainty in Artificial Intelligence (UAI), 2020.
[36] Tom Zahavy, Brendan O'Donoghue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex mdps. Advances in Neural Information Processing Systems, 34:25746-257555 2021.
[37] Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Maximum-likelihodd inverse reinforcement learning with finite-time guarantees. arXiv preprint 1rXiv:2210.01800 2022.
[38] Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Understanding expertise through demonstrations: maximum likelihood framework for offline inverse reinforcement learning. arXiv preprint arXiv:2302.07457 2023.
12

---

[39] Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities, Advances in Neural Information Processing Systems, 33:4572-45888 2020.
[40] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, 2008.
[41] Brian Ziebart, Andrew Bagnell, and Anind K Dey. Modeling interaction via the principle of maximum causal entropy. In ICML, 2010.
[42] Konrad Zolna, Scott Reed, Alexander Novikov. Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference Robot Learning, pages 247-263. PMLR, 2021.
13


Reviews:
[{"review_id": "aW79aa355o", "reviewer": "Reviewer_9xxW", "summary": "This paper aims to address the problem of inverse reinforcement learning (IRL) under the maximum entropy framework (MaxEnt-IRL) and an average reward criterion. The MaxEnt-IRL problem is formulated as a combination of an average reward Markov decision process (AMDP) and a dual IRL problem. The proposed algorithms, namely the SPMD algorithm and the IPMD algorithm, are utilized to solve the AMDP and dual IRL problem, respectively. The SPMD algorithm achieves a gradient computation step complexity of $\\mathcal{O}(1/\\epsilon)$ under general state and action spaces, while the IPMD algorithm has a complexity of $\\mathcal{O}(1/\\epsilon^2)$. The empirical studies conducted on the MuJoCo benchmark and various control tasks confirm the theoretical findings.\n", "questions": "Seen in the weaknesses section.", "limitations": "Seen in the weaknesses section.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The idea of formulating the IRL problem as a dual IRL problem, incorporating an AMDP as a subproblem, is interesting.\n2. The paper extends a previous RL algorithm [1] to general state and action spaces with a general function approximation class for AMDPs. The paper is clear and easy to follow.\n3. The proposed algorithms are well-founded and supported by theoretical analyses of average reward Markov decision processes. The convergence analysis of the algorithms is provided under certain assumptions on the approximation function classes.\n\nOverall, the paper is well-written and easy to follow. The proposed algorithms are well-motivated and supported by theoretical foundations. However, certain assumptions regarding function approximation classes and the critic step may be too restrictive, potentially resulting in the failure of convergence in scenarios involving deep neural networks. Despite these concerns, the theoretical contributions of the paper could be of interest to the community, and I recommend the acceptance of the paper.\n\n[1] Tianjiao Li, Feiyang Wu, and Guanghui Lan. Stochastic first-order methods for average-reward markov decision processes.arXiv preprint arXiv: 2205.05800, 2022.\n", "weaknesses": "1/ The advantages of employing the average reward criterion over discounted IRL are not adequately explained in the paper.\n2/ Assumption 3.2 assumes weak convexity and Lipschitz continuity for the approximated Q-functions. These assumptions may not hold in practice, particularly when neural networks are used for function approximation. As the global convergence result of SPMD relies on the convexity of the approximated Q-function and the regularizer h, it may not hold in practical scenarios.\n3/ Assumption 3.6 imposes a restriction that the errors of the critic step, typically represented by neural networks, are bounded. This assumption may be overly restrictive, as estimations can be highly inaccurate in many cases. The convergence result for general function approximation classes may also fail to hold in practice.\n4/ When the reward function could be parameterized by neural networks, which is normally the case in practice, Assumption 4.1 would be too strong.\n5/ In the reward recovery experiment, it would be beneficial to compare the proposed algorithms with discounted IRL algorithms since the expert is trained in a discounted environment (the author trained the expert agent using SAC with a discount factor of 0.99).\n6/ Including the running time of the proposed algorithms and the baselines would provide valuable insights.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your thoughtful review and the points you've raised. Your insights have been invaluable in refining our paper. Below, we address each of your comments:\n\n1. We apologize for any lack of clarity in explaining the advantages of employing the average reward criterion over discounted IRL. We appreciate your comprehensive understanding of this aspect. Discounted IRL has challenges posed by guessing the discount factor. Our average-reward criterion formulation addresses this issue effectively. First, considering demonstrations using an average-reward metric, our method excels in both reward and policy recovery. Furthermore, when the demonstration indeed uses a discount factor, but the discount factor is unknown, our approach yields improved reward estimates without the need for manual discount factor guesswork. We will enhance our paper to elucidate these advantages more explicitly.\n\n2. We thank you for highlighting the potential limitations of Assumption 3.2, particularly in the context of neural network-based function approximation. You're correct in noting that neural networks don't guarantee Lipschitz continuity. However, as we assume the state and action space are bounded sets, it is reasonable to assume there exists a Lipschitz constant, which can be fairly large. Regarding the convexity of the approximated Q-function, we have additional analysis when the approximated Q-function plus the regularizer h is not convex, which leads to SPMD converging to a stationary point. See Theorem 3.7. Please refer to our overall response for a detailed discussion.\n\n3. Your observation about the potentially restrictive nature of Assumption 3.6 is insightful. We recognize that bounded errors may not accurately capture all cases, especially when dealing with neural network-based estimations. Your suggestion to consider relaxing this assumption aligns with our intention to make our approach more applicable to a broader range of scenarios. We will explore ways to alleviate this assumption while ensuring the reliability of our convergence results. Please refer to our overall response for a detailed discussion.\n\n4. We appreciate your insight regarding the applicability of Assumption 4.1 when dealing with neural network parameterized reward functions. You're right that neural networks don't inherently guarantee Lipschitz continuity or bounded gradients. But as argued above, since the state space and the action space are compact (bounded and closed), the parameterized reward function is indeed Lipschitz continuous as long as its gradient is bounded, which is much less restrictive. It is even less restrictive when simpler models are preferred. For instance, the assumption aligns well with linear reward functions over feature spaces, commonly used in animal behavior studies, which brings room for further inference and interpretation. We will enhance the clarity of this rationale in our manuscript.\n\n5. Both IQL and f-IRL are trained with the same discount factor of 0.99 in all experiments, the same as the expert demonstrations.\n\n6. We thank you for highlighting the value of including running time information for our proposed algorithms and baselines. Our RL algorithm has the same computation efforts as similar algorithms in stable-baselines like PPO, SAC, etc. Training an agent in a single-thread environment with 5 million steps typically takes around 3.5 hours on the Apple M1 laptop. Our IRL algorithms, although carries additional reward estimation in each step, do not impose too much computational burden. Training an agent in a single-thread environment with 5 million steps takes around 3.5 to 4 hours. Please refer to our overall response for all running times. Notice that Humanoid is a larger instance and thus takes longer to train.\n\nOnce again, we extend our gratitude for your thorough review and constructive feedback, which have significantly contributed to improving our paper.\n"}, {"review_id": "LLoPFH8vNW", "reviewer": "Reviewer_cqFs", "summary": "This paper proposes an inverse reinforcement learning (IRL) algorithm for infinite horizon average reward Markov decision processes (AMDPs). At first, the authors show the stochastic policy mirror descent (SPMD) algorithm that achieves $\\mathcal{O}(\\varepsilon^{-1})$ rate of convergence. Then, the authors propose the inverse policy mirror descent (IPMD) algorithm that achieves $\\mathcal{O}(\\varepsilon^{-2})$ rate of convergence. The SPMD algorithm is compared with SAC on the MuJoCo benchmarks, and the SPMD achieves on-par performance with SAC. The IPMD algorithm is compared with f-IRL and IQL, and the experimental results show that IPMD is slightly better than f-IRL. Interestingly, the proposed algorithms perform much better than the baselines for the Humanoid environment that has many degrees of freedom. ", "questions": "Major comments\n- Is the SPDM algorithm for AMDPs (Algorithm 1) an on-policy method? The objective function for training $\\bar{\\mathcal{Q}}^{\\pi, \\zeta}$ (35) suggests an on-policy algorithm, but I did not find the explanation. In addition, (35) does not use a target network. Is it a benefit of the property of the on-policy algorithm? \n- I do not fully understand why $\\nabla \\omega (\\pi)$ is usually unavailable, as described in Lines 143-144, because $\\omega$ is selected by ourselves. Furthermore, it is unclear why its estimator is parameterized by $\\phi$ because $\\pi$ is parameterized by $\\xi$ (33). \n- Please show $g(\\theta; \\zeta)$ in detail because it is used to update $\\theta$. For example, is $g(\\theta_k; \\zeta^\\pi_k)$ the estimator of $\\mathbb{E}_{(s, a)~d^\\pi} [\\nabla_\\theta c(s, a; \\theta)]$? \n- The experimental results are promising, but one potential reason for successful results is that the reward function of the MuJoCo benchmark is well-shaped. For example, the reward of Walker2D is given by healthy_reward + forward_reward - control_cost: https://gymnasium.farama.org/environments/mujoco/walker2d/ . It may imply that the average reward formulation is more appropriate than the discount reward one in the MuJoCo benchmark. My interest is the performance of the proposed method when it is evaluated in sparse reward settings such as navigation tasks. \n\nMinor comments\n- The following paper proposes a framework based on stochastic mirror descent for AMDP. Please discuss the relationship between these papers and the proposed method. \n  - Y. Jin and A. Sidford. (2020). Efficiently Solving MDPs with Stochastic Mirror Descent. In Proc. of ICML. G. \n  - Neu et al. (2017). A unified view of entropy-regularized Markov decision processes. http://arxiv.org/abs/1705.07798 . \n  - C.-Y. Wei, et al. (2020). Model-free Reinforcement Learning in Infinite-horizon Average-reward Markov Decision Processes. In Proc. of ICML.\n- As mentioned in Strength, the following paper proposes a method to estimate the reward function under LMDP: K. Dvijotham and E. Torodov. (2010). Inverse Optimal Control with Linearly Solvable MDPs. In Proc. of ICML. It would be better to discuss their paper. \n", "limitations": "Minor comment: The authors briefly mention issues on potential societal impacts in A.10. \n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "- Originality: RL and IRL algorithms for AMDP with strong theoretical basis are novel, although Dvijotham and Todorov proposed IRL under the linearly solvable Markov decision processes (LMDP). \n- Quality: The paper provides some strong theoretical analysis on AMDP.\n- Clarity: The authors show detailed steps to prove the theorems in Appendix. \n- Significance: The standard IRL algorithms deal with infinite horizon discounted reward problems under the assumption that the discount factor is known in advance. It is problematic because the expert's discount factor is usually unknown. The proposed method is promising because it can avoid tuning the discount factor. \n", "weaknesses": "- Some recent studies regarding AMDPs are not mentioned in the manuscript. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time reviewing our paper. Your points are valid with respect to the arrangements of the paper. We appreciate your keen observation and suggestion regarding recent studies on Average-Reward Markov Decision Processes (AMDPs) that may contribute to our work. While we strive to provide a comprehensive review of relevant literature, we acknowledge that some recent studies might not have been explicitly referenced in the manuscript. We will conduct a thorough review of the literature to identify and appropriately reference any pertinent works that can enrich our discussion and provide context for our contributions.\n\n1. The SPDM algorithm for AMDPs, as outlined in Algorithm 1, is designed to accommodate both on-policy and off-policy training modes. In practice, our RL algorithm implementation is based on off-policy training schemes with a target network, similar to the approach employed in SAC. However, the adaptability of SPDM allows for on-policy training as well. \n\n2. We understand your query regarding the availability of $\\nabla \\omega(\\pi)$. By unavailable we meant it needs stochastic estimation. While we can indeed choose $\\omega$ ourselves, the precise value of $\\omega(\\pi)$ and its gradient $\\nabla \\omega(\\pi)$ needs estimation, just as the differential $Q$-function. Notice that even though $\\omega$ can be a general distance-generating function other than negative entropy, evaluating its value might require stochastic estimation from the sample we collected. Regarding the parameter $\\phi$, we note that in cases where the action value function is parameterized separately from the policy, $\\nabla \\omega(\\pi)$ might not share the same parameters as $\\xi$ in Equation (33). $\\omega$ may also be present in the form of a neural network or other models, which we denote its parameter $\\phi$, the same as the $Q$-function. We recognize that this aspect could be clearer in our explanation and will revise the manuscript accordingly.\n\n3. We appreciate your request for more detailed clarification regarding $g_k$. In practice, we employ a simple stochastic estimator using sample averages, denoted as $g(\\theta; \\zeta)= \\tfrac{1}{N} \\sum_{t=1}^N \\nabla c (s_t,a_t;\\theta)$, where $N=|\\zeta|$ represents the number of samples in the collection $\\zeta$. This estimator aids in the update of $\\theta$ as part of the algorithm. In this sense, $g(\\theta_k;\\zeta_k)$ is indeed the estimator of $E_{(s,a)\\sim d^{\\pi_k}}[\\nabla_{\\theta} c(s,a;\\theta_k)]$.\n\n4. Your observation about the well-shaped reward functions in the MuJoCo benchmark is insightful. Indeed, the presence of dense rewards in MuJoCo tasks contributes to the success of our proposed method. We acknowledge your interest in the performance of our method in sparse reward settings, such as navigation tasks. This is an area of interest for us as well, and we plan to explore and discuss the behavior of our algorithm in such environments in future work.\n\n5. Thank you for pointing out the relevance of [Jin and Sidford], [Neu et al. 2017], and [Chen et al., 2020] in the context of Mirror Descent-based methods. \n[Jin and Sidford] solve the average reward problem using linear programming-based methods, which is not comparable to policy gradient methods in practice due to scalability issues, although their analysis yields a novel complexity bound. [Neu et al. 2017] propose an entropy-regularized Reinforcement Learning framework, which covers important policy gradient methods and casts them as mirror descent or dual averaging. However, no convergence analysis is provided. [Chen et al., 2020] is more comparable to our RL analysis. However, the regret bound is known to be different from general convergence analysis. Nevertheless, all the above works consider MDPs with finite state and action spaces. Our work considers a more general setting where action and state can be continuous instead of categorical, for example, in robotics research and animal behavior studies. \n\n6. We appreciate your reference to Dvijotham and Todorov's work on IRL under linearly solvable MDPs (LMDP). While their approach offers an effective control-oriented perspective on IRL, we recognize that using Maximum Likelihood Estimation may face limitations when applied to larger and more complex problems, such as those involving continuous state and action spaces. Our choice to formulate the problem under a maximum entropy framework accounts for these challenges and allows for a more robust and flexible approach. We will certainly expand on the discussion regarding the relationship between our approach and the LMDP framework in our forthcoming revision.\n\nOnce again, we sincerely thank you for your thoughtful review, which has significantly contributed to the refinement of our paper. Please kindly refer to our global response for societal impacts."}, {"review_id": "ZP8aTRCSkM", "reviewer": "Reviewer_ifE7", "summary": "In this paper inverse reinforcement learning is studied when the teacher was using an average-reward criterion rather than discounted rewards with known discount factor. The paper proposes a stochastic first-order method starting from stochastic policy mirror descent for MDPs and continuing towards inverse policy mirror descent for solving the IRL problem. The paper contains also some numerical experiments based on MuJoCo benchmarks.\n", "questions": "Can you explain more about the improvements (as in Humanoid RL) and deficiencies (as in Ant IRL or walker RL)? What feature of the algorithm contributes to the success or suboptimality of the performance in each problem? Are the assumption not strictly satisfied, or is the sampling the reason for any suboptimality?\n\nWhich theorems have appeared in earlier work or have strong similarities to theorems in earlier work?\n\nI would assume that 1/(1-gamma) not much smaller than K, and the difference between them is related to the number of states and actions. In other words, is there a practical need for the O-symbols or is the complexity already implied by the known factor proportional to 1/(1-gamma)?\n\nWhy does the performance show variablity even for simple problems (see Fig. 1)?\n", "limitations": "The authors state in the supplementary material: \"We anticipate no potential negative societal impacts concerning this research.\" which could seem questionable in the context of IRL, because IRL has the potential to uncover hidden motives in legal human behavior without the consent of a person. The research seems fine, but some discussion, beyond the mere statement would be useful.", "rating": 3, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The paper is very clear and elegant in its use of average reward method in RL and IRL. There are many papers on duality of this type although not yet specifically on the average reward criterion, although in NAC papers similar considerations were made.\n\nThe supplementary material is strictly supplementary and very clear and comprehensive.\n\nIt's great the hints for practical use are included.", "weaknesses": "“To the best of our knowledge” is not a useful sentence in an abstract, in particular in an anonymous manuscript. It is more useful to discuss such opinions later in the paper in the context of related results (see e.g. previous NeurIPS conferences).\n\nDuality has been used in many similar contexts, but it could still be said more clearly why it is a good idea here. It is acceptable, if you want to avoid discussing natural AC, but then there is more to explain.\n\nExperiments are relatively few with just final results given, so no understanding why the proposed method is useful can be gained. The is no discussion of the experimental results, therefore, although one may guess, it is not clear why the proposed method struggles on Ant or whether there is any relation of the performances comparing tables 1 and 2. There is no attempt included to check whether the complexity results are tight or whether global optima are indeed found.\n\nCheck: “Acrot”\n\nAssumptions 4.2 and 4.3 are formally stated in the main text, but do not seem to be used. In the appendix only 4.2 is mentioned.\nRef. 28 is not mentioned in the text, perhaps a separate bibliography for the appendix would make sense.\n\nI find it a bit difficult to distinguish between all the Qs. The distinction between math-Q and cal-Q is probably necessary, For example, tilde-cal-Q (32) could as well be substituted by its definition. It occurs again only in (34) and (38), where at least a ref to the definition needs to be included. In “differential Q-function” (in italics in the main text vs. theorem environment) different fonts are used, which should be avoid this, ideally by using the math-Q font that seems to occur also in the theorem environment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and insightful comments. We appreciate your time spent on the review and detailed comments on the paper. First we want to point out that our paper focuses on the theoretical development of solving the Inverse Reinforcement Learning (IRL) problem under the average-reward setting. Note that almost all analyses on the discounted MDPs highly rely on (some form of) the contraction property given by the discount factor, and their work cannot be directly translated into the average-reward setting.\n\n1. Regarding \"To the best of our knowledge\" in the abstract, we acknowledge your point and will consider rephrasing it in subsequent revisions.\n\n2. We recognize the need for a clearer justification of the use of duality. In the context of Inverse Reinforcement Learning (IRL), we highlight that directly solving the primal problem presents challenges, including maintaining constraints and ensuring maximum entropy policies, which are not resolved effectively in the community. Exploiting duality has a history in IRL, allowing us to leverage a more structured dual problem, enhancing our ability to devise effective optimization techniques. Regarding using Natural Actor-Critic (NAC), it is true that one may solve the subproblem with NAC. But there's no existing theoretical analysis on the rate of convergence for NAC for solving the average-reward Markov Decision Processes (AMDPs). As mentioned in the paper, [Cen 2021] uses NAC to solve a discounted MDP with a finite state and action space. Whether its method and analysis can be adapted to solve our problem (average reward with general state and action space and general function approximation) needs nontrivial analysis. \n\nShicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence\n314 of natural policy gradient methods with entropy regularization. Operations Research, 2021.\n\n3. While we acknowledge the relatively limited number of experiments and the absence of an in-depth discussion of results, our primary focus was on the theoretical underpinnings of our proposed methods. The experiments validate our theoretical findings and illustrate competitive or superior performance, even without exhaustive exploration of hyperparameters. Notice that since the performance metric is the episodic return, methods that work with discounted MDPs are approximately solving the AMDP, and the performance is very sensitive to the discount factor. In IRL, guessing the discount factor is crucial to estimate the reward. Our method can naturally avoid these problems. \n\n4. We understand your point about complexity analysis and global optima. Given the complexities introduced by function approximation, proving tight complexity bounds or identifying global optima remains challenging. Notice that there are active efforts in the theoretical community hoping to derive a \"tight\" complexity bound for general Reinforcement Learning, yet no general results are obtained. Additionally, in the context of non-convex optimization, a global optimum is usually unattainable. Nevertheless, our approach offers insights into practical reward recovery, as demonstrated in the reward recovery experiment comparing our method to Inverse Q-learning (IQL), a discounted method. We show that even with the correct discount factor when using IQL, our method is 10x times more accurate.\n \n5. We appreciate your observation about Assumptions 4.2 and 4.3. Assumption 4.2 is crucial in our analysis framework, impacting Lemma A.4 (Equation 72), A.5 (Equation 78), and Thm 4.5 (Equation 88), while Assumption 4.3 is essential for Thm 4.5 (in Equation 93) and 4.6 (Equation 104-107).\n\n6. We understand your concern about distinguishing between different Q-functions. Our choice of $\\tilde{\\mathcal{Q}}$ aims for conciseness, though we will ensure better clarity in referencing definitions. We will consider using more distinct notations for denoting the differential $Q$-function and its stochastic estimator.\n\n7. We acknowledge the request for a deeper discussion of algorithmic improvements and deficiencies in various environments. For a detailed analysis of the algorithm, please refer to our overall reply.\n\n8. We clarify that our theoretical contributions leverage convex/non-convex optimization techniques. Our formulation of average-reward IRL is novel, bridging the gap between the current understanding of average-reward Markov Decision Processes (AMDPs) and IRL under such settings. While some theorems may appear similar to existing work, adapting them to this novel context requires nontrivial analysis, as evidenced by the cited lemmas.\n\n9. In our complexity analysis, the $O$ notation captures convergence rates, while constants such as $1/(1-\\gamma)$ play lesser importance as $K$ can grow sufficiently large. This notation is widely used to convey the speed of convergence. It provides a concise and informative way to describe the growth rate of a function relative to a specific parameter (usually the problem size or the number of iterations).\n\n10. The observed variability in performance, even for simple problems, can be attributed to the stochastic nature of the environments. When an environment is initialized, a random noise is added by the simulator. That contributes to the variability\n\n11. Your ethical concern is valid, and we share your perspective on the responsible use of IRL. Please refer to our overall response for a detailed discussion.\n\nOnce again, we thank you for your comprehensive review and look forward to addressing your valuable feedback in our revisions."}, {"review_id": "Rn0Vumh2Dy", "reviewer": "Reviewer_DtTS", "summary": "This paper addresses Inverse Reinforcement Learning (IRL) in the average-reward setting. The authors propose a Stochastic Policy Mirror Descent (SPMD) method to solve the Average Reward Markov Decision Process subproblem and use SPMD to propose the Inverse Policy Mirror Descent method to solve the IRL problem. The authors provide complexity results for these methods and experimentally validate them using the MuJoCo benchmark.", "questions": "- Does Assumption 3.2 hold in practice? Do the authors do anything to enforce these constraints?\n- What is $\\omega$ in Equation (3) and Algorithm 1 chosen to be in practice?\n- Can the authors provide architecture details for the policy, Q network, and learned reward function in the appendix?\n- Equation 24 in Theorem 3.4 makes an assumption about step size. Is this assumption implemented in practice (i.e. in Algorithm 1 what is $\\eta_k$ chosen to be?\n- It seems the authors assume access to a large number of expert trajectories (Appendix A.9 says expert demonstrations are collected for five million steps). I’m curious how IPMD would perform with a more limited number of expert trajectories (for instance, the f-IRL[1] paper has experiments where they just use 1 expert demonstration). Have the authors considered how number of expert trajectories affects performance of IPMD?\n\n[1] f-IRL: Inverse Reinforcement Learning via State Marginal Matching (Ni et al.)", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors develop a novel method and show its merits both theoretically and experimentally compared to existing methods.\n- The theoretical analysis is very thorough and the authors clearly outline the assumptions that are made.\n- The proposed IPMD method achieves good results on a variety of environments.", "weaknesses": "- The paper is a bit light on experiment details (e.g. hyperparameters). More detail here would be appreciated (e.g. see Questions).\n- It’d be nice if the authors could provide an intuitive explanation for the result of Theorem 3.4. Equation 25 is a bit difficult to understand.\n- The authors make a fair number of assumptions in their theoretical analysis. It'd be nice if the authors added some discussion on which assumptions hold in practice and/or how they enforce those assumptions, as well as how restrictive these assumptions are.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We extend our sincere appreciation for your insightful review and thoughtful comments, which greatly contribute to the refinement of our paper. We are pleased to address each of your points below:\n\n1. We recognize the importance of providing comprehensive experiment details, including hyperparameters, to enhance the transparency and reproducibility of our work. In response to your suggestion, we will furnish a more comprehensive account of hyperparameters, encompassing aspects such as learning rates, discount factors, and neural network architectures. Please see our response to later questions and our global reply.\n\n2. We value your input regarding Equation 25 and its comprehension. Equation 25 encapsulates a convergence overview without specifying explicit step size choices. Such formulations, common in optimization literature, highlight the significance of step size selection on convergence speed. To be specific, $\\alpha_k, \\beta_k$ are both step sizes and the central term $\\sum \\rho_k - \\rho^*$ (distance to the optimal value) plus $D(\\pi_k, \\pi^*)$ (the distance of the current policy to the optimal policy) will shrink as $k$ grows, which is be bounded by some combination of $\\alpha_k$and $\\beta_k$ plus an irreducible function approximation error $\\varsigma$. The left-hand side is a weighted average of function value error plus an average error from step $k=1$ to $k=K$. We acknowledge the less intuitive nature of Theorem 3.4 and will provide a more intuitive explanation of its implications, coupled with illustrative examples, to enhance clarity. For a specific step size choice and the condition on $\\mu$, we get different rates of convergence, which is presented in Corollary 3.5.\n\n3. Assumption 2.1 (uniform ergodicity) is considered restrictive but often necessary for analysis, albeit potential violations in practice. Assumption 3.2's practicality hinges on model choices; if neural networks are employed, gradient \"clipping\" may be utilized. This also applies to Assumptions 4.1 and 4.2. Assumptions 3.3, 3.6, and 4.3, concerning stochastic estimator proximity, are less restrictive given bounded state and action spaces.  Your insight into detecting issues in practice is well-taken, and we will incorporate a comprehensive discussion to highlight these nuances.\nIn practice, it is usually obvious when things go wrong, e.g., when the differential Q-function explodes to 10e6. This can happen due to rare occasions and the culprit is usually detectable, e.g., a large learning rate. Please refer to our overall response for additional discussion. \n\n4. As elucidated in Section 3.1, in practice, we select $h$ and $\\omega$ as negative entropy due to the subproblem's entropy-regularized nature. However, alternative distance-generating functions for Bregman distance beyond KL divergence are possible, and PMD can accommodate different choices. We will provide additional clarification on this choice and its implications. \n\n5. The policy network employs two fully connected hidden layers of dimension 256 each, taking actions as input and outputting a distribution. Both the Q network and reward function share the same architecture, with ReLU activation used in hidden layers. A double Q-learning technique minimizes overestimation. We will incorporate comprehensive architecture specifications in the appendix for reference. \n\n6. We appreciate your inquiry into the implementation of the assumption in Equation 24. Our comparative exploration of various step size schedules yielded minimal differences, possibly due to subproblem optimization via a few gradient descent steps. We will emphasize this aspect's relevance and practical implementation. \n\n7. Your query on IPMD's performance sensitivity to expert trajectory count is insightful. While unreported at submission, IPMD exhibits strong performance on locomotion tasks involving complex robots with a sole expert demonstration. For example, in one of our ongoing projects, we test IPMD on the MuJoCo Cassie robot (https://mujoco.readthedocs.io/en/stable/models.html), hope Cassie can walk on random terrain from demonstrations. IPMD can reach episodic reward 479.815 from one expert demonstration of episodic reward 447.1955, outperforming the expert demonstration. This suggests IPMD's efficiency in utilizing limited expert demonstrations. \n\nOnce again, we thank you for your comprehensive review and look forward to addressing your valuable feedback in our revisions."}]

Author rebuttal:
1. Regarding assumptions being too restrictive:
We acknowledge that some assumptions made in the paper are sometimes too restrictive. However, our particular problem setting has a special implication for the assumptions we made.  Assumption 2.1 (uniform ergodicity) is considered restrictive but often necessary for analysis, albeit potential violations in practice. Assumption 3.2's practicality hinges on model choices; if neural networks are employed, gradient "clipping" may be utilized. This also applies to Assumptions 4.1 and 4.2. Assumptions 3.3, 3.6, and 4.3, concerning stochastic estimator proximity, are less restrictive given bounded state and action spaces. 
There are possible remedies to encourage our algorithm to satisfy those assumptions. First, we can "clip" large gradients by a constant; or add penalty terms regarding the norm of the gradient. Regarding the span semi-norm contraction property, we can construct a $J-$step operator in the policy evaluation step, which is well explored in the literature (see [Puterman 1994], [Zhang 2021]).

Puterman, M. L. (2014). Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons.

Zhang, S., Zhang, Z., Maguluri, S. T. (2021). Finite Sample Analysis of Average-Reward TD Learning and $ Q $-Learning. Advances in Neural Information Processing Systems, 34, 1230-1242.

In practice, it is usually obvious when things go wrong, e.g., when the differential Q-function explodes to 10e6. This can happen due to rare occasions, and the culprit is usually detectable, e.g., a large learning rate. 

2. Hyperparameter settings:
The policy network employs two fully connected hidden layers of dimension 256 each, taking actions as input and outputting a distribution. Both the Q network and reward function share the same architecture, with ReLU activation used in hidden layers. A double Q-learning technique minimizes overestimation. 
During training, we found that setting the entropy coefficient term to 0.01 make training stable and efficient. The learning rate is set to be 3e-4. Each step of the algorithm samples 512 state-action sample pairs. Additionally, we include our running time of APMD and IPMD for all MuJoCo tasks in the pdf file attached.

3. Experiments analysis: 
We acknowledge the request for a deeper discussion of algorithmic improvements and deficiencies in various environments. While further performance analysis and tuning are indeed possible, our primary focus was on theoretical foundations. One possible reason Ant is falling behind is that Ant has more ground contact since it has more legs. This will impact the ergodicity of the MDP and our assumption on the 1-step contraction. It is harder to transition from an arbitrary state to another arbitrary state as it involves multiple legs working together. One remedy is that we construct a $J$-step contractive operator for policy evaluation, as done in [Chen 2021]. For the success of humanoid, it is possibly due to a slightly different policy evaluation scheme where the entropy term from the policy no longer plays a part, as the term $r-\rho$ cancels out the additional regularization and entropy of the policy. This is different from discounted setting. We suspect this brings more stable training and thus higher performance.

4. Ethical discussion: 
Our focus is on methodology, and we acknowledge that IRL, like many other machine learning techniques, has potential implications if misused. IRL can be used in violation of privacy (as the reviewer mentioned) by inferring an individual's intentions and preferences, potentially crafting convincing social engineering attacks or phishing attempts; IRL can also be used to model the behavior of specific demographics, which could result in biased algorithmic decision-making, leading to unfair treatment or discrimination against certain groups, etc. 

Meta-review: The reviews for this paper were mixed, with some reviewers accepting it and some reviewers rejecting it, giving the paper an average borderline score. The paper considers the intersection of average reward & inverse RL which is giving it a relatively narrow scope. After reading the paper and discussing related work with the authors, the novelty of the paper is now clear to me, but remains somewhat limited given the already narrow scope and existing related work. Nevertheless, the paper makes interesting contributions that can later on be generalised to other settings. I encourage the authors to include a discussion on how their work can be used more broadly, for example, for other problems that involve average reward or more general utilities. I also encourage them to include the detailed response and comparisons they have provided in the rebuttal regarding the related work. After further considerations, I am recommending to accept the paper.
Decision: Accept (poster)