UTILITY: Utilizing Explainable Reinforcement Learning to Improve Reinforcement Learning
Authors: Shicheng Liu
Minghui Zhu
Venue: ICLR 2025 Poster (2025)
Keywords: explainable reinforcement learning
Primary area: reinforcement learning

Abstract: Reinforcement learning (RL) faces two challenges: (1) The RL agent lacks explainability. (2) The trained RL agent is, in many cases, non-optimal and even far from optimal. To address the first challenge, explainable reinforcement learning (XRL) is proposed to explain the decision-making of the RL agent. In this paper, we demonstrate that XRL can also be used to address the second challenge, i.e., improve RL performance. Our method has two parts. The first part provides a two-level explanation for why the RL agent is not optimal by identifying the mistakes made by the RL agent. Since this explanation includes the mistakes of the RL agent, it has the potential to help correct the mistakes and thus improve RL performance. The second part formulates a constrained bi-level optimization problem to learn how to best utilize the two-level explanation to improve RL performance. In specific, the upper level learns how to use the high-level explanation to shape the reward so that the corresponding policy can maximize the cumulative ground truth reward, and the lower level learns the corresponding policy by solving a constrained RL problem formulated using the low-level explanation. We propose a novel algorithm to solve this constrained bi-level optimization problem, and theoretically guarantee that the algorithm attains global optimality. We use MuJoCo experiments to show that our method outperforms state-of-the-art baselines.

Published as a conference paper at ICLR 2025
UTILITY: UTILIZING EXPLAINABLE REINFORCE- MENT LEARNING TO IMPROVE REMNNTTTEMENT LEARNING
Shicheng Liu Minghui Zhu Department of Electrical Engineering Pennsylvania State University University Park, PA 16802, USA {s f1f15539 9, muz1 1 6 ppsusu edu
ABSTRACT
Reinforcement learning (RL) faces two challenges: (1) The RI. agent lacks ex- plainability. (2) The trained RL agent is, in many cases, non-optimal and even far from optimal. To address the first challenge, explainable reinforcement learning (XRL) proposed to explain the decision-making of the RL agent. In this paper, demonstrate that XRL can also be used address the second challenge, i.e., improve RL performance. Our method has two parts. The first part provides two- level explanation for why the RL agent is optimal by identifying the mistakes made by the RL agent. Since this explanation includes the mistakes of the RL agent, has the potential to help correct the mistakes and thus improve RL perfor- mance. The second part formulates constrained bi-level optimization problem learn how to best utilize the two-level explanation improve RI. performance. In specific, the upper level learns how to use the high-level explanation to shape the reward so that the corresponding policy can maximize the cumulative ground truth reward, and the lower level learns the corresponding policy by solving con- strained RL problem formulated using the low-level explanation. We propose a novel algorithm to solve this constrained bi-level optimization problem, and the- oretically guarantee that the algorithm attains global optimality. We use MuJoCo experiments show that our method outperforms state-of-the-arr baselines.
INTRODUCTION
While reinforcement learning (RL) has been implemented in a wide range applications, it faces two significant challenges: (1) The RL agent lacks transparency due to its black-box nature. (2) has been widely observed in the RL community (Haarnoja et al. 20188] Henderson et al.) 2018
Dulac-Arnold al., 22199; Cheng al. (0244) that, in many cases, the trained RI. agent does achieve maximum cumulative reward (i.c., non-optimal and even far from optimal). These two challenges motivate the need improve the transparency and performance of the RL agent.
To address the first challenge, explainable reinforcement learning (XRL) methods are proposed to explain the decision-making of the RL agents, including learning an interpretable policy (Bastani et al. 2018: Bewley & Lawry 2021), pinpointing regions in the observations that are critical for choosing certain actions (Atrey al., 2019 Puri et al. 2019). learning the reward function that is actually maximized (Xie et al. 2022). and identifying the critical states influential to the cumulative reward (Guo et al., 2021 Cheng et al.) 2023). These XRI methods generate various explanations that improve the transparency of the RL agent and help people build trust in the RL agent.
This paper demonstrates that XRL can also be used to address the second challenge of RL, i.e. RL improvement. Given non-optimal RL agent, we use XRL to explain why this RL agent is not optimal by finding the mistakes made by the RL agent. Since our explanation provides insights into the RL agent's mistakes. it has the potential to help correct the mistakes and thus improve performance. Some recent works (Guo et al., [021]] [heng et al., 2023] 2024) also use XRL to improve the RI performance. In specific, they propose to first identify the critical states that

---

Published as a conference paper at ICLR 2025
most influential the cumulative reward explanation, and then perturb the actions (Guo al, 2021] or fine-tune the policy (Cheng et al. 2023] (022) at those critical states. However, they do not explain why the RL agent does not maximize the cumulative reward. This paper proposes novel framework that first explains why the RI. agent is not optimal, and then learns how to utilize the generated explanations to improve RI. performance. We summarize contributions follows:
Contribution statement. This paper proposes an optimization- based framework that aims to learn how to best utilize XRL improve RL. We refer to this framework "utilizing explainable RL to improve reinforcement learning efficacy" (UTILITY), Our contributions are threefold:
First, we provide two-level explanation for why the RL agent is optimal. The high-level planation learns reward function to which the RL agent is actually optimal, and then explains why the RL agent not optimal by comparing this learned reward function the ground truth reward function. The low-level explanation identifies the state-action pairs that lead the RL agent to be non- optimal. We refer to these state-action pairs as "misleading" state-action pairs, and rigorously derive mathematical metric to identify the "misleading" state-action pairs the low-level explanation.
Second, we formalize the problem of utilizing the two-level explanations improve the performance as constrained bi-level optimization problem. The upper-level problem aims to learn how to use the high-level explanation (i.e., the learned reward function) to shape the ground truth reward function to help the corresponding policy maximize the cumulative (ground truth) reward. The lower-level problem learns the corresponding policy by solving a constrained RL problem where the objective to maximize the cumulative shaping reward and the constraint is discourage from visiting the low-level explanation (i.e., the "misleading" state-action pairs). A key insight is that the formulated optimization problem facilitates policy improvement. Current constrained bi-level optimization can only deal with the case where the lower-level problem is strongly convex. However, in our case, both the objective function and constraint in the lower-level problem are highly non-convex. Therefore, novel theoretical framework is desired to solve this constrained bi-level optimization problem.
Third, we develop novel theoretical framework and thereby an algorithm to solve the constrained bi-level optimization problem. In specific, we first use a dual method to transform the constrained bi-level optimization problem to an equivalent unconstrained bi-level optimization problem, and then propose approximation-ban based triple-loop algorithm to solve this unconstrained bi-level op- timization problem. We quantify the approximation error each loop and prove that the algorithm attains global optimality. Experiments show that UTILITY outperforms state-of-the-ar baselines.
(b)
Figure 1: (a) An RL task where drone starts from the lower-left corner and navigates to the orange goal at the uiper-right corner. (b) failing trajectory (blue) generated by the RL agent and heat map that visualizes the ground truth reward. (c) The two-level explanation of why the blue trajectory fails to reach the goal. The high-level explanation is the learned reward (visualized as the heat map) to which the RL agent is actually optimal. The low-level explanation is the misleading state-action pairs (red circles and linked red arrows). (d) The trajectory after improvement reaches the goal.
Illustrative example. Figure Tsses example to illustrate proposed framework. Suppose use RL to navigate drone to the orange goal in Figure Ia] The state is the 2-D coordinate and the action is the moving direction. The ground truth reward is the goal states and zero otherwise. Figure Tb uses a heat map to visualize the ground truth reward and the blue trajectory is generated by the learned policy. This learned policy is not optimal because it fails reach the goal.
Figure Ic visualizes our two-level explanation why the learned policy not optimal. At the high level, the heat map visualize the learned reward function to which the RL agent's
2

---

Published as a conference paper at ICLR 2025
trajectory (policy) is actually optimal. The high-level explanation is that the RL agent's policy is actually optimal to the learned reward function (visualized as the heat map in Figure Ic]. and this learned reward function is very different from the ground truth reward function (visualized as the heat map in Figure lbj. Note that we normalize all the reward functions in Figures Ib Id to [0, 1] for better comparison. For the low-level explanation, we identify the top five "misleading" state-action pairs (i.e., the red circles and arrows) in the blue trajectory where the red circles are the states and the linked red arrows are the corresponding actions chosen by the non-optimal RL agent. These state-action pairs are "misleading" since the correct actions should point to the goal.
Figure shows the improvement where the heat map visualizes the learned shaping reward function and the blue trajectory is generated by the learned policy after improvement. We can see that the learned policy after improvement successfully reaches the goal.
2 RELATED WORKS
Due to the space limit, we only include the related works improving RL performance here, and we include more related works in Appendix E
Reward shaping. Reward shaping can improve the RL performance by shaping the ground truth reward function. Current works on reward shaping has two main categories. The first category (Ng al.] 1999 |uu et al.) 2020 Devlin & Kudenko 2012 Gupta et al. 2022) requires source, such as human expert, to provide domain knowledge as an ingredient to shape the ground truth reward function. However, when the tasks become complicated, it could be difficult and even infeasible for humans to provide domain knowledge. The second category does not need domain knowledge, including reward shaping based on exploration bonus (Bellemare et al.. 2016] Ostrovski et al. 2017). learning an intrinsic reward (Zheng et al., 2018, Memarian et al. (222)) I), and combining exploration bonus and intrinsic reward (Devidze et al. (022). The first category usually has better performance, while the second category does not require human-domain knowledge. Our method enjoys the benefits of both categories because we use domain knowledge but the domain knowledge is not given by human expert but learned by XRL.
external
Other methods that can improve RL performance. Lazy-MDP (Jacq et al., 2022) shows perfor- mance improvement with the help of a provided default policy. It uses the "lazy-gau to determine whether to choose greedy action or follow a default policy on each state s. Self-imitation learning (Oh al. 2018)) aims encourage deep exploration by reproducing previous good decisions. Pa- pers (Wang& & Taylorr [017]] Taylor] 20188] TTaylor et al.] 2023) aim to improve the RI performance by utilizing external assistance, such as the assistance of pre-trained RL agent (Wang & Taylor Taylor 2018 Taylor et al. 2023), which may not be accessible in some scenarios.
2017) or a human
TWO-LEVEL EXPLANATION OF WHY THE RL AGENT IS NON-OPTIMAL
This section provides a two-level explanation to explain why the RL agent is not optimal. The RL agent's decision making is based a Markov decision process (MDP) (S,A, 7, Po, P,r) which consists of a state set S, an action set AA, a discount factor y E (0, an initial state distribu- tion Po(-), state transition function -(-,.-), and the ground truth reward function r(.,)) The RL agent's learned policy is denoted by TA and the cumulative reward is defined J,( A Ex IE 7"((s,, [z)] where the initial state is drawn from Po- When we say that the RL agent not optimal, means that F o arg max, (T)).
1 The black-box assumption. To ensure practicability, following (Bewley & Lawry 2021 Guo
al. 202T, Cheng et al., 2023) Guidotti et al., 2019). we only treat the RL agent as black box with no access to its internal structure. In specific, we do not assume the access to the learned value/Q-function nor the learned policy TA of the RL agent. We can only observe a set of m trajectories D & } demonstrated by the RL agent (using the non-optimal policy TA) where each trajectory Ç = Su, ag..... ... is a state-action sequence.
The high-level explanation. At a high level, since the RL agent is not optimal the ground truth reward function r, we can learn a reward function f to which the RL agent's s policy na is actually optimal, and use this learned reward function f generate explanations. A recent work (Xie et al.)
3

---

Published as a conference paper at ICLR 2025
2022 uses the state-action pairs (s.a) with the highest f(s. a) explanation, however, these state-action pairs cannot explain why the RL agent is not optimal. Therefore, we extend (Xie et al. 2022) by comparing the learned reward function f to the ground truth reward function r to explain why T is not optimal to the ground truth r. Figures [ Tc provide an example of our high-level explanation: the policy æ is actually optimal to the learned reward function f (visualized in Figure Ic). and this learned reward function f is very different from the ground truth reward function r (visualized in Figure lb].
Inverse reinforcement learning (IRL) (Abbeel & Ng 2004 Ziebart et al. 2008 Arora & Doshi 202I] can learn the reward function and an associated policy TA from the demonstration set D such that the behaviors of policy Fa demonstrated in DD are optimal to the reward function & learned by IRI, and the learned policy 7 can imitate the policy πΛ- We use maximum likelihood IRI. Zeng et al. 2022 2023) to learn the reward function f and policy R..
While the learned reward function f can be used for the high-level explanation, is only interpretable to humans in low dimension, e.g., we can use heat maps to plot reward functions (as in Figure D. When the state and action become high dimensional, the learned reward function f is hard for humans to understand and thus it is difficult to straightforwardy compare ř to (as we did in Figures 1b Icp. Therefore, we need the low-level explanation which is still interpretable in high dimension.
The low-level explanation, At low level, the RL agent is not optimal meaning that it visits some critical points that lead to the non-optimality. Recent works (Guo et al. 202I: Cheng et al.] [023] Amir & Amir 2018 Jacq et al. 2022) identify the states that are most influential to the cumulaa lative reward as critical points. In order to explain why the RL agent is not optimal, we extend their idea by redefining the critical points as the state-action pairs that lead T be non-optimal. We refer to these critical points as "misleading" state-action pairs and we aim to identify the top K "misleading" state-action pairs in the demonstration set as the low-level explanation, Note that we use infinite time-horizon MDP and it is not possible to identify the top K "misleading" state-pairs if trajectory has infinitely many different state-action pairs. However, in practice, the trajectory length is usually finite and we can unify the notions of finite time horizon and infinite time horizon by introducing "absorbing state" (Sutton & Barto 2018). In specific, we can treat the terminal state of finite-time- state keeping transitioning only to itself and generating zero reward.
horizon trajectory
The key challenge to identify the top K "misleading" state-action pairs in the demonstration set D is to propose proper criterion or metric to define what a "misleading" state-action pair is. A straight- forward way is to identify the state-action pairs that an optimal policy will not visit as "misleading", and thus use π A (als) T (als) as the metric where x² is optimal policy. However, this metric is infeasible because the optimal policy n is not accessible.
In contrast, derive feasible metric in Definition I that uses Q Q-function to find misleading state- action pairs in the demonstration set D. The Q-function under the policy π and reward function r is (7, & E Eis 7'r(st, at) [sp = s, ag = a].
Definition 1. A state-action pair (s, a) E DD is misleading state-action pair (ss > 0 where a(),a) A max.' Q (s,a') - QF (s, a) is referred as "misleading level". The larger the mis- leading level 1(8, a) is, the more misleading the state-action pair (8, a) is.
We include the derivation of how we come up with this metric I and the proof of why (s, a) D is misleading if ((s, a) > 0 in Appendix B.I In brief, we prove in Appendix B. that the policy TA will be an optimal policy if ((s, a) = 0 for all (s, a) such that a E A where (s)ss) is the set of actions that the policy T has nonzero probability to choose at the state s. Therefore, any state-action pair (s. E D such that l(s, a) > can be regarded as a "misleading" state-action pair that leads the policy TA to be non-optimal. The larger the misleading level 1(x, a) is, the more "misleading" the state-action pair (s, a) is, because the Q value of the chosen action a has a larger gap from the maximum Q value at the state s. We denote the set of the identified top K "misleading" state-action pairs by C, which serves as the low-level explanation.
While we cannot access the policy A. we have already learned the policy # using IRL and the policy A imitates the policy AA Therefore, we can use Q, to substitute for Q Given that we can access A and r, we can simply learn a by sampling the environment to collect enough data and doing regression. Due to the space limit, we include the method of learning Q in Appendix We are aware that in RL, it is usually sample inefficient and computationally expensive if we
B.2
4

---

Published as a conference paper at ICLR 2025
want to sample enough data to learn precise Q-functions corresponding to all the learning policy in the learning procedure. However, our case is different because we only need to learn one precise Q-function, which corresponds to the specific policy A.-
4 UTILIZING THE TWO-LEVEL EXPLANATION TO IMPROVE RL
This section provides theoretical framework that utilizes the two-level explanation in Section 3 to improve the RL performance. In specific, Subsection 4.1 formulates the problem as a constrained bi-level optimization problem. Subsection 4.2 proposes novel theoretical framework and thereby an algorithm to solve the constrained bi-level optimization problem.
4.1 PROBLEM FORMULATION
We aim to utilize the two-level explanation improve the RL performance. For the high-level explanation f, we use it to formulate a domain knowledge and learn how to use this domain knowl- edge to shape the ground truth reward such that the learned shaping reward can lead the pol- icy to maximize the cumulative ground truth reward (T). In specific, we use the comparison r - f between the ground truth reward r and the high-level explanation & as the domain knowl- edge. Note that in practice, we need to first scale f to the same scale with for better compari- son. Since this comparison quantifies the RL agent's misunderstanding of the ground truth reward function it has the potential to help patch the error. Towards this end, we propose to learn shaping function -(..)). For a given state-action pair (s, a). the corresponding shaping reward is To a) =sss a) + 0(s, (r(ss, f(s, a)).
For the low-level explanation C, we discourage the RL agent from visiting the "misleading" state- action pairs in C. Towards this end, we design cost function c(.,)) such that c(s. a) (0, Cmaxl when (s, a) E C, and c(s, a) = 0 otherwise, where Cmax is a positive constant. We discourage from visiting C by constraining the cumulative cost JT(TT) - E [Er= 7'((#e 0z)] under a budget b.
To utilize the two-level explanation, formulate constrained bi-level optimization problem:
max (Tro), where Tra = arg maxx(Jrr (T) + (x), s.t. Je(x)= S b), 8
where (T) E Ei (rro(r(sr ae), rss,, as) - f(s,, a:))] is the cumulative shaping reward, and the causal entropy H(T)) F E -7" log T (ax) [s:]] is to encourage exploration and is widely used in soft Q-learning (Haarnoja et al., 2017) and soft actor-eritic (Haarnoja et al. 2018).
In the problem L the upper level aims to learn shaping reward function such that the cor- responding policy Tre can achieve maximum cumulative ground truth reward J, (TTg)) Given the current learned shaping reward function re, the lower-level problem in (T) is to compute the corre- sponding policy Fr, by solving a constrained RL problem. The constrained RL problem encourages Tre to maximize the entropy-regularizgg ages Tre from visiting C by controlling the cumulative cost Je(TT) under the budget b. Note that we choose 0, it means that the policy Try should totally avoid the set C.
cumulative shaping reward (Jr, (T) + H H(x)) and discour-
Remark on why the problem (1) can improve AA Recall that TA is an non-optimal policy obtained by running RL algorithm until convergence, i.e., TA stops improving. The lower level in (I) improves A by facilitating a policy improvement step. In specific, according to monotonic policy improvement theorem, TA will improve if it chooses greedy actions according its Q-function Q at all the states. Given that A stops improving, it means that RA must choose some nongreedy actions at some states, i.e., the misleading state-action pairs. Constraining these misleading state- action pairs means that we constrain the nongreedy actions T originally chooses at the states. This constraint can help A choose greedy actions because it eliminates some nongreedy actions and thus A only needs to find greedy actions from smaller action sets. Since this constraint can help A find greedy actions, it can help improve T.. The upper level calibrates the reward function to improve performance. In specific, we can treat the RL agent as a biased black box such that if the input is the ground truth reward function r, the output is a policy TA that is optimal to another reward function f. We want to change the imput of this biased black box such that the output is policy that is optimal to the ground truth reward function r. More formally, we aim to learn a shaping reward ro as the input such that the output is optimal to r. A straightforward way is directly learn neural reward
5

---

Published as a conference paper at ICLR 2025
ro S x A + R. However, in practice, this is difficult to learn because the neural reward needs to search over a very large function space. We find empirically that if we use the shaping reward in the form of ro (s. a) = r(s, + ((s, ()(r(s, f(s, a)), the learning results will be much better. facilitates policy improvement, we can further improve the final learned policy of problem ( by generating its two-level explanation and use the explanation formulate problem (T) again. We can iteratively formulate and solve problem (I) to keep improving the learned policy. In this paper, we only discuss the scenario where we formulate and solve the problem [ once.
As the problem
Before solving the problem T we need to first make sure that the problem ①I is well-defined. In specific, since the lower-level problem in (I) is non-convex, it may have more than one optimal solution, i.e., Tre is not unique. Therefore, given a reward parameter 8, the corresponding upper- level objective function value (Tr,) may not be unique as Tr, is not unique. This will make the problem (T) ill-defined. The following theorem guarantees that the problem [ is well-defined.
Theorem 1. Given reward re, the optimal solution Try of the lower-level problem in 7 is unique.
4.2 THEORETICAL FRAMEWORK
While the current state-of-the-arts (Xu & Zhu 2023a Khanduri et al. 2023) constrained bi-level optimization can only deal with strongly convex lower-level problems, both the objective function and the constraint of the lower-level problem in (T) are non-convex. Therefore, novel theoret- ical framework is desired to solve the problem This subsection proposes a novel theoretical
framework solve the problem (1)
The proposed theoretical framework has three parts. (i) The first part transforms the original con- strained bi-level optimization problem ( to an equivalent unconstrained bi-level optimization prob- lem. The benefit of this transformation is that the equivalent unconstrained bi-level optimization problem has an unconstrained and convex lower-level problem, which is more tractable and easier to solve. (ii) The second part proposes novel algorithm to solve the problem (I) by solving the equivalent unconstrained bi-level optimization problem. (iii The third part theoretically guarantees that the proposed algorithm attains global optimality.
4.2.1 PROBLEM TRANSFORMATION
The lower-level problem of the problem is non-convex. To deal with the non-convexity issue, we
introduce the dual function of the lower-level problem in [ G(x:; 0) & max, Jr, (T) + (
X(Je(T) b) where À is the dual variable. The dual function G(A; 6) is convex in À since it is the point-wise maximum over a set of affine functions of X (Boyd & Vandenberghe 2004).
Theorem 2. The optimal solution of the lower-level problem in ! is uniquely the constrained soft policy (e)e where (0) is the unique optimal solution of the dual problem MİnA G();O).
We include the analytical expression of the constrained soft policy T (0)0 (Liu & Zhu 2022 in Appendix C Theorem
2 indicates that T (0);8 is the unique optimal solution of the Tower-level
problem in (i.e., (0)(0) = Ty,), and (0) = arg miny G();;) 0). Therefore, we can replace Try with x- (9):0 and replace the lower-level problem in ① with its dual problem, and thereby
transform the constrained bi-level optimization problem I to the following unconstrained bi-level
optimization problem:
max J, (T))))0))))) where X X²(0) arg min G();6).
Compared to the original problem (T). the lower-level problem of the problem 2 is unconstrained and convex. However, there are still two challenges to solve the new problem
Challenge (i): Evaluating the dual function G(A; needs to obtain the constrained soft policy TA;0- arg max,
However, current RL algorithms can only approach TA at a certain rate and only obtain the exact тA29 when iteration number goes to infinity. In practice, we can only an algorithm for finite iterations and thus we cannot obtain the exact x This will cause errors when we evaluate the dual function G.
Challenge (ii): Even if solution A (0) of the lower-level problem in finite time. This makes it difficult to evaluate and solve the upper-level problem in (2) since the ueper-level problem in (2) requires X² (0).
obtain the exact T we cannot guarantee to get the exact optimal

---

Published as a conference paper at ICLR 2025
Algorithm 1 Utilizing explainable reinforcement learning to improve reinforcement learning Input: Demonstration set 2, initial shaping reward parameter d,, dual parameter λo, and policy To Output: Shaping reward Fow and the policy after improvement म "((N)):N
1: Generate the two-level explanation (f, C)
2: for = 0, 1,........ NN 1 do
3: for ń = 0, , do
4: for ñ = 0, ................................ ,. - do
5: Compute the constrained soft Q function Vxa:9.
6:
Update the policy (a)s(a)s) x exp(Q An; (s, a)) for any (s,a) E E x A
7:
end for
8:
Set Aa, = ① No and use A to compute the approximated gradient 9x.:8.. Update AR+1 = An angAn:0n
9:
10:
end for
11: Set A(O,,)
N En=O A and compute # :. via (N s)-step soft policy iteration
12:
Use î (00 to compute the approximated gradient ge.. and update px+1 = 0n + Snge.
13:
end for
4.2.2 THE PROPOSED ALGORITHM
This part proposes novel algorithm that solves problem (1) by solving problem (2). The proposed algorithm is triple-loop where the inner loop approximates the constrained soft policy тт and tackles Challenge (i), the middle loop approximates the optimal solution X (0) of the lower-level problem in (2) and tackles Challenge (ii), and the outer loop solves the upper-level problem in (2). We use n, n. and ñ to respectively denote the iteration indices of outer, middle, and inner loop.
Algorithm D first generates the two-level explanation (line D and then uses three loops to utilize the generated two-level explanation. specific, the inner loop (lines approximates the constrained soft policy A: With the approximated policy :x:o (line 8 the middle loop solves the lower-level problem in
2 via (N - 1)-step gradient descent (line to approximate the optimal solution X (0)))-
With the approximated parameter A(0) (line
the outer loop solves the upper-level problem in ② via (N 1)-step gradient ascent (line 12). In the following, we elaborate each loop respectively.
The inner loop. Given the parameter (x, 0), the inner loop aims approximate the constrained soft policy TA via -tep soft policy iteration (Haarooj al. 2017). and Nn. Rt + 1. Soft policy iteration has two steps: policy evaluation and policy improvement. Policy evaluation computes the constrained soft Q-function Q%, corresponding to the current policy πa, dual parameter X, and reward parameter O. We include the expression of the constrained soft Q-function in Appendix C Policy improvement aims to update the policy according to T (als( xx exp(QX;0 (s, a)) for any a) E S x A. The output of the inner loop is the approximated policy âx, π Ng In practical implementations, we can update the policy Ry via the policy update in soft Q-learning (Haarnoja et al. (2017) or actor update in soft actor-critic (Haarnoja et al. (018)). While soft Q-learning and soft actor-critic are designed for unconstrained RI, we show in Appendix Chat we can revise them to approximate the constrained soft policy.
The middle loop. We aim solve the lever-level problem in 2 via (N-1)-step gradient descent. Lemma 1. The gradient of the dual function G is VxG(X:0)
The gradient Vx f) requires the exact constrained soft policy πA which is inaccessible. There- fore, we use the approximated policy fx:0 obtained from the inner loop to approximate the gradient VxG(A; 0) via the gradient approximation gx.0 b (f(fx.)) and solve the lower-level problem
via (N-1)-step gradient descent An+1 = A 0..Aa.:0 The output is A(0) = + N-1 An-
The outer loop. We solve the upper-level problem in (2) via (N 1)-step gradient ascent. Towards this end, we generalize the Q/value function (Sutton & Barto 2018). In specific, we define the Q- E Ei ('c((r 04) = 8, = a) and value function of cost as (s) A E 1 Dr rec(st, 0x) s). We define the fuunctioo of reward gradient Varo Qoro (8, a) A F De ? Voro a)) 8,40 a) and value function of
function of cost e under policy ा as (s, oo
7

---

Published as a conference paper at ICLR 2025
reward gradient as V E V s]. We define state-action visitation frequency as (s, EJ (0 =, ar = a}] where 11:} is the indicator function.
= E(sa))mp>x100 [10 V= (s) Coy- (Qo (e (8,a) - Ve Vo (s)) Qr -->'0000 (s, a)] where Cn is constant vector we fix policy n, and we include the expression of + in Appendix D.3)
Lemma 2. The upper-level gradient is dJ, (T)(Tx-(0;=))
(9)0 (s,
Since the gradient (d)() requires the exact optimal solution (0) and the exact constrained soft policy T (8))8) we can only use the policy *x(0);9 to approximate (d)(())()))))) via go = Excoi (s) - (Qe 2xter (s, a) Ve *xcere (s))) fxcmn (8, a)].
db
d0
E (s.a) s0000 [(Q
Exroni (s, a) Vera
We then solve the upper-level problem in (2) via (N 1)-step gradient ascent on+1 = On + Bag9n.
4.2.3 THEORETICAL ANALYSIS
This part quantifies the optimality of the policy after improvement [x(ww The main difficulty is that the inner loop and middle loop can only approximate the policy Tx:0 and the optimal solution X (0), and the approximation error may accumulate and ruin the convergence of the outer loop. In the following context, we sequentially quantify the convergence from the inner loop the outer loop.
Lemma (convergence of the inner loop). Given the parameter (2,0), the output tx:0 of the inner loop satisfies log (als) - log [als) S (n7"5)) ) for any (s,0) E S x A.
Lemma shows that inner loop converges linearly to the exact constrained soft policy xx:0-
Assumption 1. (i) It holds that (ra(;)) ") < C₁ for any 0 where C₁ a positive constant. (ii) It holds that || ([ III S C2 and Vo ( -1 < Cs. where C2 and C3 are some positive constants.
Assumption assumes that rg is bounded, Lipschitz continuous, and smooth to which a com- mon assumption in RL (Zheng et al.
2024b Lan et al. 2024a] Zheng et al. [024a Zhang et al.) 2025 Zheng et al., [024cc Lan et al. 2024b). We next quantify the convergence of the middle loop:
Lemma 4 (convergence of the middle loop). Suppose Assumption (ii) holds and let aa 1/(n +
1)" where ij E (1/2, 1), the outputs (A(0), * ((0xa) ofe middle loop satisfy that (i) [(0) x²) s O(1/N1-); (ti) log मे (000:
log (o); (a [c]s]] S O(1/N1- for any (s, Sx A.
Lemma 4 shows that if the iteration Ñ of middle loop is sufficiently large, the approximation error of X* (0) and (0) can arbitrarily small. We next quantify the convergence of the outer loop: Theorem (convergence of the outer loop). Suppose Assumption 7 and the conditions in Lemma hold and let 3n = 1/(n + 1)7 where 1) E (1/2, 1), then it holds EN-1 [TJ Il O(1/N' + 1/ 20 + 22.).
Theorem shows that Algorithm converges to stationarity when the iteration numbers and N go to infinity. When the state-action space is finite, we have the following stronger result:
Theorem 4 (optimality of the outer loop). Suppose the conditions in Lemma f hold and the state-action space is finite. Let the step size 3,, < min((1
/2/ L/L,, then it holds that lim - J (x(*x(0x)0x)) JØ = 0 where J is the maximum value of (T),,, and L is the smoothness constant of J, (0)=)))))) whose expression is in Lemma 9 in Appendix.
Theorem 4 shows that when the state-action space is finite, Algorithm can find an optimal policy asymptotically when the iteration numbers N and N go infinity.
EXPERIMENT
This section provides experiment results for the proposed framework. In specific, aim to answer the question: How does the proposed framework (UTILITY) compare to other RL improvement methods in terms of improving the RI. performance. Towards this end, we introduce three RI.
8

---

Published as a conference paper at ICLR 2025
Table Experiment results.
SAC
UIILITY
RICE
SIL
LIR
Delayed HalfCheetah
383.45 + 45.50
715.96 + 42.78
456.14 = 36.32
510.34 1 39.28
548.28 +: 47.94
Delayed Hopper
192.90 1 27.18
317.99 F 19.62
232.55
16.96
263.46
20.72
247.27 F 31
Delayed Walker2d
134.91 + 20.80
242.63 + 14.11
177.45
20.14
172.28
24.57
204.72 + 25.99
Delayed Ant
68.11 = 12.52
105.80 + 14.38
77.01 10.89
81.05 13.43
78.23
13.11
SAC
UNLITY
RICE
SIL
LI
300
300
200
350
154
202
160
40
60
(427)
(a) Delayed HalfCheetah
(b) Delayed Hopper
(c) Delayed Walker
(d) Delayed Ant
Figure 2: Improvement curve.
improvement methods for comparisons. (i) Fine-tune policy on initial states and critical states (RICE) (Cheng et al.] (022)): This method fine-tunes the policy starting at the original initial states and the states that are most influential to the cumulative reward. Note that (Guo et al.) 2021| Cheng al. 2023) also use the most influential states to improve performance and (Cheng al. 2024)
shows performance superiority over (Guo et al. 202 Cheng al.) 2023), thus we pick (Cheng et al. 2024) to compare. (ii) Self imitation learning (SIL) (Oh et al. 2018): This method reproduces previous good decisions in order to encourage deep exploration. (ii) Learning intrinsic reward (LIR): This method aims to learn an intrinsic reward i to formulate the shaping reward r + (Zheng et al. 2018). We choose these three methods to compare because they respectively belong to three different categories: XRL method (RICE), reward shaping method (LIR), and other methods that can improve RL (SIL). We use soft actor-critic (SAC) (Haarnoja et al. 2018 as the baseline RL algorithm that all the above RL improvement methods use and improve from. We aim to show the improvement of the above RL improvement methods compared SAC. We first run SAC until convergence to get the non-optimal policy πa and use nA to generate the demonstration set D for UTILITY to generate the two-level explanation and thus improve.
We test the algorithms on delayed MuJoCo environments (Zheng et al. 2018 Memarian et al. 2021 Oh al., 2018) where the rewards are accumulated for 20 time steps and provided only at the end of these periods. Note that this makes the reward become sparse, and we include the additional experiment results on dense reward in Appendix [ We use four different delayed MuJoCo tasks: delayed HalfCheetah, delayed Hopper, delayed Walker2d, and delayed Ant. Following (Finn al. 2017). each episode has the length of 100 in our experiments.
Figure 2 shows the learning curves of the algorithms where the z-axis is the interaction steps with the MDP environment and the y-axis is the cumulative reward. We plot both the mean (i.e., the solid line) and standard deviation (i.e, the shadow area) of the algorithms. The mean and standard deviation are computed using five random seeds. From the figures, we can observe that UTILITY improves the baseline RL algorithm SAC by large margin. While the other three methods (i.e.. RICE, SIL, and LIR) can also improve SAC to some extent, UTILITY achieves the highest cumula- tive reward. This due to the fact that UTILITY (i) learns a shaping reward that makes it easier to learn good policy and (ii) discourages the policy from making the mistakes (i.c., the "mislcading" state-action pairs) made by SAC. Note that the learned shaping reward dense while the ground truth reward in the delayed MuloCo environments is sparse, so that UTILITY can guide the learned policy to achieve higher reward. In contrast, RICE and SIL still suffer from the sparse reward. While LIR can also learn a dense shaping reward, UTILITY has the domain knowledge mulated by the high-level explanation to help better shape the reward. Moreover, UTILITY has the "misleading" ftat-acction pairs avoid. Note that the learned shaping reward not only helps in the sparse reward scenario, we include additional results in Appendix F.I to show that UTILITY can still largely improve SAC when the ground truth reward is dense.
f for-

---

Published as a conference paper at ICLR 2025
Table shows the final performance of the algorithms. We can observe that UTILITY has the highest cumulative reward among all the algorithms.
The ablation study. Since UTILITY uses both the high-level explanation (i.e., the learned reward) to shape the ground truth reward and the low-level explanation (i.c., misleading state-action pairs) to formulate constraint to improve SAC, we include ablation study to separately study the effect of the shaping reward and the constraint. In specific, we consider two methods: "shaping only" and "constraint only". The "shaping only" method only uses the high-level explanation learn shaping reward but does not uses the low-level explanation to formulate a constraint. The "constraint only" method only uses the low-level explanation to formulate the constraint but does not shape the original reward. We include the results for the delayed environments in Table 2 and the results for the dense environments both the "shaping only" and "constraint only" methods can improve SAC. Moreover, the "shaping only" method has a larger impact to improve the performance. This is because the shaping reward improves the policy globally as it changes the reward value for all (s, a), while the constraint may only improve the policy locally around the misleading state-action pairs.
Appendix
The results in Table 2 and Appendix F22 show that
Table 2: Ablation study.
SAC
UITLITY
shaping only
constraint only
Delayed HalfCheetah
383.45 + 45.50
715.96 = 42.78
695.63 + 33. 66
422.15 + 22.86
Delayed Hopper
192.90 + 27.18
317.99 = 19.62
289.10 + 18.41
210.12 + 15.77
Delayed Walker2d
134.91 + 20.80
242.63 14.11
211.37 + 18.64
175.66 + 15.27
Delayed Ant
68.11 + 12.52
105.80 = 14.38
88.18 + 8,66
75.16 + 6.58
Evaluation of the generated two-level explanation. Following (Guo et al. 2021 Cheng et al., 2023), we use fidelity as the metric to respectively evaluate the high-level and level explanations. The fidelity means the correctness of the two-level explanation. Since the two-level explanation is to explain why the RL agent is not optimal, one way to validate the fidelity of the explanation is to see whether the cumulative reward increases after we improve from the explanations.
we can see that both the high-level and low-level explana- tions are correct explanations because both the shaping only method and the constraint only method improve the performance. Moreover, the shaping only method (the fourth column in Table (2) has higher cumulative reward than LIR (the last column in Table TD. and the constraint only method (the last column in Table 2 has higher cumulative reward than RICE (the fourth column in Table D. This shows the high fidelity of our two-level explanation.
From the last two columns in Table 2
To compare the fidelity of our explanation with other methods, we fix the improvement method and change the explanation to compare. For the low-level explanation, we compare with RICE. In specific, we still use the constraint only method but now the constraint is to discourage from visiting the critical states identified by RICE. We refer to this method as "RICE+constainn.. Note that it is expected that "RICE+constainnt has low fidelity in our case because RICE does not aim to explain why the RI. agent is not optimal. For the high-level explanation, since there is no existing XRL method to compare, we use our shaping-only method without the domain knowledge to compare. We refer to this method as "shaping without r - f" We include the results for sparse reward and dense reward in Appendix [3 The results show that both the high-level and low-level explanations of UTILITY have high fidelity.
[F3
CONCLUSION
This paper proposes theoretical and systematic framework that utilizes XRL to improve RL. We first provide an explanation for why the RL agent is not optimal, and then formulate the problem of utilizing the explanation to improve RL as constrained bi-level optimization problem. We propose a novel theoretical framework to solve this problem, and use experiments to validate that the proposed framework can improve the RL performance. Despite the benefit, one limitation of the proposed algorithm is that it requires to interact with the environment. Therefore, one future work to extend our method to the offline RL. setting.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGEMENTS
This work is partially supported by the National Science Foundation through grants ECCS 1846706 and ECCS 2140175. We thank the reviewers for their insightful and constructive suggestions.
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning, pp. 1 1-8, 2004.
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International Conference Machine Learning, pp. 22-31, 2017.
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1-76, 2021.
Dan Amir and Ofra Amir. Highlights: Summarizing agent behavior to people. In International Conference on Autonomous Agents and MultiAgent Systems, pp. 1168-1176, 2018.
Saurabh Arora and Prashant Doshi. A survey of inverse reinforcement learning: Challenges, meth- ods and progress. Artificial Intelligence, 297:103500, 2021.
Akanksha Atrey, Kaleigh Clary, and David Jensen. Exploratory not explanatory: Counterfactual analysis of saliency maps for deep reinforcement leinfin. In International Conference on Learn ing Representations, 2019.
Osbert Bastani, Yewen Pu, and Armando Solar-Lezama. Verifiable reinorrrmet learning via policy extraction. Advances in Neural Information Processing Systems, 31: 31:2499-255 2509, 2018.
Amir Beck. First-order methods in optimization. Society for Industrial and Applied Mathematics, 2017.
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying bount-based exploration and intrinsic motivation. Advances in Neural Information Pro- cessing Systems, pp. 1479-1487, 2016.
Tom Bewley and Jonathan Lawry. Tripletree: A versatile interpretable representation of black box agents and their environments. In AAAI Conference on Artificial Intelligence, volume 35, pp- 11415-11433 11422, 2021.
Stephen P Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge university press, 2004.
Zelei Cheng, Xian Wu, Jiahao Yu. Wenhai Sun, Wenbo Guo, and Xinyu Xing. StateMask: Ex- plaining deep reinforcement learning through state mask. In Advances Neural Information Processing Systems, 2023.
Zelei Cheng, Xian Wu, Jiahao Yu, Sabrina Yang, Gang Wang, and Xinyu Xing. Rice: Breaking through the training bottlenecks of reinforcement learning with explanation. In International Conference Machine Learning, 2024.
Yinlam Chow, Ofir Nachum, Edgar Quenez-Guzman and Mohammad Ghavamzadeh. A Iyapunov- based approach to safe reinforcement learning. Advances in Neural Information Processing Systems, pp. 8103-8112, 2018.
Rati Devidze, Parameswaran Kamalaruban, and Adish Singla. Exploration-guided reward shaping for reinforcement learning under sparse rewards. Advances in Neural Information Processing Systems, pp. 5829-5842, 2022.
Sam Michael Devlin and Daniel Kudenko. Dynamic potential-based reward shaping. In Interna- tional Conference on Autonomous Agents and Multiagent Systems, pp. 433-440, 2012.
Gabriel Dulac-Amold, Daniel Mankowitz, and Todd Hester. Challenges of worl-world reinforcement learning. arXiv preprint arXiv: /290/, 2019.

---

Published as a conference paper at ICLR 2025
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learing for fast adaptation of deep networks. In International Conference on Machine Learning, pp. 1126- 1135, 2017.
Riccardo Guidotti, Anna Monreale, Fosca Giannotti, Dino Pedreschi, Salvatore Ruggieri, and Franco Turini. Factual and counterfactual explanations for black box decision making. IEEE Intelligent Systems, 34(6):14-2) 2019.
Wenbo Guo, Xian Wu, Usmann Khan, and Xinyu Xing. Edge: Explaining deep reinforcement learning policies. Advances in Neural Information Processing Systems, 34: 14:12222-12236, 2021.
Abhishek Gupta, Aldo Pacchiano, Yuexiang Zhai, Sham Kakade, and Sergey Levine. Unpacking re- ward shaping: Understanding the benefits of reward engineering on sample complexity. Advances in Neural Information Processing Systems, 35:15281-15299 2022.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning, pp. 1352-1361, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Confer- ence on Machine Learning, pp. 1861-1870, 2018
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. AAAI Conference on Artificial Intelligence, vol- ume 32, 2018
Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: A new approach of reward shaping. Advances in Neural Information Processing Systems, 33:15931 18:15931-15944 1-59911,,, 2020.
Sandy H Huang, Kush Bhatia, Picter Abbeel, and Anca D Dragan. Establishing appropriate trust via critical states. In IEEE International Conference Intelligent Robots and Systems, pp. 3929- 3936, 2018.
Alexis Jacq, Johan Ferret, Olivier Pietquin, and Matthieu Geist. Lazy-MDPs: Towards interpretable by learning when act. International Conference on Autonomous Agents and Multiagent Systems, pp. 669-677, 2022.
Zong Ke, Jingyu Xu, Zizhou Zhang, Yu Cheng, and Wenjun Wu. A consolidated volatility prediction with back propagation neural network and genetic algorithm. arXiv preprint arXiv:2412.07223, 2024.
Prashant Khanduri, Ioannis Tsaknakis, Yihua Zhang, Jia Liu, Sijia Liu, Jiawei Zhang, and Mingyi Hong. Linearly constrained bilevel optimization: A smoothed implicit gradient approach. In International Conference on Machine Learning, pp. 16291-16325, 2023.
Guangchen Lan, Dong-Jun Han, Abolfazl Hashemi, Vaneet Aggarwal, and Christopher G Brinton. Asynchronous federated reinforcement learning with policy gradient updates: Algorithm design and convergence analysis. arXiv preprint (4f::2404008033. 2024a.
Guangehen Lan, Han Wang, James Anderson, Christopher Brinton, and Vaneet Aggarwal. Improved communication efficiency in federated natural policy gradient via admm-based gradient updates. Advances in Neural Information Processing Systems, 36, 2024b.
Shicheng Liu and Minghui Zhu. Distributed inverse constrained reinforcement learning for multi- agent systems. Advances in Neural Information Processing Systems, 35:33444-33456 2022.
Shicheng Liu and Minghui Zhu. Meta inverse constrained reinforcement learning: Convergence guarantee and generalization analysis. In International Conference on Learning Representations, 2023.
Shicheng Liu and Minghui Zhu. Learning aalti- behaviors from distributed and streaming demonstrations. Advances in Neural Information Processing Systems, 36, 2024a.
12

---

Published as a conference paper at ICLR 2025
Shicheng Liu and Minghui Zhu. In-trajectory inverse reinforcement learning: Learn incrementally before an ongoing trajectory terminates. arXiv preprint arXiv:2410.15612, 2024b.
Yongshuai Liu, Jiaxin Ding, and Xin Liu. IPO: Interior-point policy optimization under constraints. In AAAI Conference on Artificial Intelligence, volume 34, pp. 4940-4947, 2020.
Farzan Memarian, Wonjoon Goo, Rudolf Lioutikov, Scott Nickum, and Ufuk Topcu. Self-supervised online reward shaping sparse-rewart environments. In IEEE International Conference on In- telligent Robots and Systems, pp. 2369-2375 2021.
Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning, pp- 278-287 1999.
Junhyuk Oh, Yijie Guo, Satinder Singh, and Honglak Lee. Self-imitation learning. In International conference on machine learning, pp. 3878-3887, 2018.
Georg Ostrovski, Marc G Bellemare, Aäron Oord, and Rémi Munos. Count-based exploration with neural density models. In International Conference Machine Learning, pp. 2721-2212222 7721-2730, 2017.
Nikaash Puri, Sukriti Verma, Piyush Gupta, Dhruv Kayastha, Shripad Deshmukh, Balaji Krishna- murthy, and Sameer Singh. Explain your move: Understanding agent actions using specific and relevant feature attribution. In International Conference on Learning Representations, 2019.
Guanren Qiao, Guiliang Liu, Pascal Poupart, and Zhiqiang Xu. Multi-modal inverse constrained reinforcement learning from a mixture of demonstrations. Advances Neural Information Pro- cessing Systems, 36:60384-60399 2023.
Guanren Qiao, Guorui Quan, Rongxiao Qu, and Guiliang Liu. Modelling competitive behaviors in autonomous driving under generative world model. In European Conference on Computer Vision, pp. 19-36, 2024.
Adam Stooke, Joshua Achiam, and Pieter Abbeel. Responsive safety reinforcement learning by pid lagrangian methods. International Conference on Machine Learning, pp. 9133-9143, 2020.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Matthew E Taylor. Improving reinforcement learning with human input. In International Joint
Conference on Artificial Intelligence, pp. 9724-5728 2018.
Matthew E Taylor. Nicholas Nissen, Yuan Wang. and Neda Navidi. Improving reinforcement learn- ing with human assistance: argument for human subject studies with hippo gym. Neural Computing and Applications, 35(32))23499-2339, 2023.
Chen Tessler, Daniel J Mankowitz, and Shie Mannor. Reward constrained policy optimization. In International Conference on Learning Representations, 2018.
Zhaodong Wang and Matthew E Taylor. Improving reinforcement learning with confidence-based demonstrations. In International Joint Conference on Artificial Intelligence, pp. 3027-3033, 2017.
Yuansheng Xie, Soroush Vosoughi, and Saeed Hassanpour. Towards interpretable deep reinforce- ment learning models via inverse reinforcement learning. In International Conference Pattern Recognition, pp. 5067-5074, 2022.
Siyuan Xu and Minghui Zhu. Meta value learning for fast policy-centric optimal motion planning. In Robotics science and systems, 2022.
Siyuan Xu and Minghui Zhu. Efficient gradient approximation method for constrained bilevel opti- mization. In AAAI Conference on Artificial Intelligence, volume 37, pp. 12509-12517. 2023a.
Siyuan Xu and Minghui Zhu. Online constrained meta- leta-learning: Provable guarantees for general- ization. Advances in Neural Information Processing Systems, 15531-115554, 2023b.
13

---

Published as a conference paper at ICLR 2025
Siyuan Xu and Minghui Zhu. Metf-reiinorreeent learning with universal policy adaptation: Prov- able near-optimality under all-task optimum comparator. arXiv preprint arXiv:2410.09728, 2024.
Tengyu Xu, Yinghin Liang, and Guanghui Lan. CRPO: A approach for safe reinforcement learning with convergence guarantee. In International Conference on Machine Learning, pp- 11480-1 1491, 2021.
Q Yu, 7. Ke, G Xiong, Cheng, and X Guo. Identifying money laundering risks digital asset transactions based Al algorithms. 2025.
Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Maximum-likelihood inverse re- inforcement learning with finite-time guarantees. Advances in Neural Information Processing Systems, 15:10122-10135, 2022.
Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. When demonstrations meet genera- tive world models: A maximum likelihood framework for offline inverse reinforcement learning. Advances in Neural Information Processing Systems, 36:65531 655661, 2023.
Haochen Zhang, Zhong 7heng, and Lingzhou Xue. Gap-dependent bounds for federated q-learning. arXiv preprint arXiv:2502.02859, 2025.
Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. Advances in Neural Information Processing Systems, pp. 4649-4659, 2018.
Zhong Zheng, Fengyu Gao, Lingzhou Xue, and Jing Yang. Federated Q-leaming: Linear regret speedup with low communication cost. International Conference on Learning Representations, 2024a.
Zhong Zheng, Haochen Zhang, and Lingzhou Xue. Federated q-learning with reference-advantage decomposition: Almost optimal regret and logarithmic communication cost. arXiv preprint arXiv:2405.1879) 2024b.
Zhong Zheng, Haochen Zhang, and Lingzhou Xue. Gap-dependent bounds for q-learning using reference-advantagee decomposition. arXiv preprint urXiv:2410.075544 2024c.
Zhengyuan Zhou, Michael Bloem, and Nicholas Bambos. Infinite time horizon maximum causal entropy inverse reinforcement learning. IEEE Transactions on Automatic Control, 63(9):2787- 2802, 2017.
Brian D Ziebart, Andrew L Maas, Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, pp. 1433-1438, 2008.
A APPENDIX
This appendix has three sections. B additional content of the two-level explanation. Section C provides notions and notations that serve as the building blocks of the appendix. Section D provides the proof of all the lemmas and theorems in the paper. Section F provides experiment details.
B TWO-LEVEL EXPLANATION
This section has two subsections. Subsection B. provides the derivation steps of how we come up with the mathematical metric / to find "misleading" state-action pairs in Section 3 and proves that a atate-action pair (s, a) E D is misleading if I(s, a) > 0. Subsection B.2 provides an algorithm to learn the Q-function aAaaA
14

---

Published as a conference paper at ICLR 2025
B.1 JUSTIFICATION OF THE MATHEMATICAL METRIC TO FIND THE "MISLEADING" PAIRS
Since the mislcading state-action pairs lead the policy TA to be non-optimal, we can say that the pol- icy TA will be an optimal policy if it does not visit misleading state-action pairs. Therefore, in order to identify "misleading" state-action pairs using Q-functions, we need to first build a connection between Q-functions and policy:
Definition 2. (i) When we say that a Q-function O7 indicates a policy n, it means that T(s)
arg m Q7 (s. a) for any s E S where ((s) is the set of actions that the policy a will choose at the state s.
(ūi) We use O to denote the set of all the Q-functions that indicate the policy n, and thus we can say that 2 indicates the policy и. Moreover, when we say that 2 indicates an action a at a state it means that a E arg m
o Q7 (s,aa)) a') where Q is an arbitrary Q-function in 27.
(iti) When we say that Q7 disagrees with OF on (s,a), a). it means that O,7 indicates the action a at the state but & arg max,' (s,
Definition establishes connection between Q-functions and policy. With this connection, the following theorem provides a way to define a mathematical metric to find "misleading" state-action pairs using Q-functions:
Theorem 5. The policy TA will be an optimal policy if n never visits any state acte-nction pair (s, a), on which Q disugrees with A
Proof. Suppose there is no state-action pair (s, E S x A, on which Q disagrees with 27A. Then for any (s.a) S A, if a arg max Q (s, a') where Q E 2,a, it holds that arg max, QFA Q (s,a'). Since Tx(s) = arg max, Q (s, a) where Q E O it holds that nT(s)) S arg max, Q (s. a) for any 8 E S. Recall that V and Q are respec- tively the value function and Q-function of the policy holds that (7*(s,a) a) = r(ss a) +
under the reward function r, then it and (s) = Ea'nma(t)ai [Q a')) Since (s)s) S arg max, Qs, a) for any s S. then (s) = En (11)))/444 (F, max, QTA (s, a') for any s S. Therefore, the Q-function O satisfies the Bellman optimality equation QT* a) = r(ss() a) +7E maxx,' Q=A(s a')] for any (s, a) E S x A. and thus QF* is the optimal Q-function because the Bellman optimality equation is uniquely satisfied by the optimal Q-function. Since (s) S arg max, Q (s, a) for any E S, the policy TA should be an optimal policy.
[V.
Since the misleading state-action pairs lead the policy πA to be non-optimal, we can say that the policy TA will be an optimal policy if it has zero probability to visit misleading state-action pairs. Therefore, Theorem 5 shows that the "misleading" state-action pairs can be mathematically de- fined as the state-action pairs, on which Q disagrees with 2.A. Therefore, we can develop mathematical metric to find the top "misleading" state-action pairs in the demonstration set D using Q-functions. Since the demonstration set D is generated by the policy TA and Q.AAA indicates the policy TA Q,aa will indicate the action a at the state s for any (s, a) E D. In order to find the "misleading" state-actions, we need to find (s,a) E D such that a & arg max,' QFA (s, d') or in other words, (=A(s,a) < max. Q77 (s, d'). Therefore, max,' QF* (s, a''-a'''- Q= a) as the metric to identify "misleading" state-action pairs in the demon- stration set DD The larger the loss I(s,a) a) is, the more "misleading" the state-action pair (s, a) is, because the Q value of the chosen action a has a larger gap from the maximum Q value at the state under the Q-function Q,AAA
can use the function I(s, a) &
B.2 METHOD
In this subsection, we use standard regression method to learn the Q-function Q+AAA In specific, we roll out the policy A to generate a set D of many (s, a) samples. For each (s, a) E D. we use f to generate many trajectories starting from (s, a) and use these trajectories estimate (s, Since we can generate many trajectories, we can estimate the Q value QF (s, a) for each (s, a) E D quite accurately. Then we use neural network Qu parameterized by solve the following
15

---

Published as a conference paper at ICLR 2025
regression problem:
min
Σ
(s) -
(s,a)ED
C NOTIONS AND NOTATIONS
The shaping reward function ro(r(s, a(s, i function of r(s, a) and r(s, ((s.a)-fin oo-rfss, a), and r(s, a) and r(s, a) - f(ss, a) are both functions of (s.a). Therefore, the shaping reward function is also function of (s, a). For simple notations, we use (s to denote the shaping reward of (s, a). Given the policy x and the parameters (A, 0)s, the corresponding constrained soft Q-function and constrained soft value function are:
(x a) r)(s, ad(s,a) +7 /
S'ES
(x) & F" (ro(ss), - )o(se, - log (axl)))) 80 =
Moreover, it can be shown ZZnng et al. exp(V (s) SoEa exp( (Q Qx,a (s, a) da.
2022 Haarnoja et al, 2017
Liu & Zhu
2024a b) that
The constrained soft policy (Liu & Zhu
2022
2023) is
Tx:o(als) exp(Qx:0(s, a))
(3)
exp(Vx:0(s))
Qx:((s,) = rr((s, a) >c(s, a) +7 / P(s'(s, (V)
(4)
(s) log /
exp(Qx:((s,
(5)
JaEA
We can obtain the constrained soft policy via soft Q-learning (Haarnjja et al. critic (Haarnoja et al. 2018) by treating 7ę Ac as the new reward function. mulative cost under the policy π starting from (s, a) as (s, 2 E" De 7'c(st, de)lso
2017 soft actor-
We define the cu-
og = a] and the cumulative cost starting from 8 (s)
E"
Dr oo
7 /a((s,, ac)
s]. We define the cumulative reward under the policy क starting from (s,a)
(s, A
E IE t frr(s,, ae)lso s, ag = a] and the cumulative reward starting from s as ViT(s) &
E Er ("r((1,, [s) = s). We define the cumulative reward gradient under the policy æ starting
from (8,a) as 27 (s,a) a E E' 7' Varoo (e) = 8, 0g = a) and the cumulative reward
Vara (s) & E E r Vororl 4t) = s]. We define the state visitation fre- quency under policy x as 1hat (s) Ea E 1 {st = s)] and state-action visitation frequency
starting from
V
as sin (s,a) a) A E ET 7 "1 s, a, = aHl where 1 } is the indicator function.
D PROOF
This section provides the proof of all the lemmas and theorems the paper.
Lemma 5. The gradients of the constrained soft policy are respectively VA log (als)) =
V - Q- (s, and V log (xls(a)s) = Q (s, a) V (s).
Proof. Recall from B.. we know that VA Tolo VAQxi(s, VA Recall from we know that
V vAQxo(s, = -c(ss +7 /
JES
( -((s + P(8'(8,)) exp(Vxx((s))
/
Vxexp(Qxols) d'')dd' as'
rEs
dEA
16

---

Published as a conference paper at ICLR 2025
/
a( exp(Qxi(s'd)) Vx d'
-c( +2
exp((xia(s')
JES A'EA
(6) (s,aa) +7
J'EA P( VAQui d''dd'
--((s,)) JuES I P(6(Is.))Txo((ls) -c(s' + S"ES
where (a) follows
(b) follows
Keep the expansion, we can see that
V2Qx0(s,a)= a
TA /'(s,, ar]lo s, ao = a] -Qa (s,
(6)
Similarly, we can get that:
VAVv:o(s) = E
y /c(s), a0)l0 = = -
(7)
1=0
voQxi((,a) = A
3,00 = al = Quors (s,a),
(8)
(oVA:(())= E Σ
VVrr((:; = = =
Voro (s);
Therefore, compute the gradients
V log (als) VaQxpo(s,a) - VAVANO()) = VoTAio (s) -Qaa (s, (s,a),
V log (als) voQA:o(s,a) VoVA,((s) QTin (s,a) Viaa (s).
Lemma 6. The constrained soft operator Txp
/ (ol(Q(s''aa))dd ds',
(Tx, (s, a) Ac(s, + P(s' (is'ss,0)oog
W'EA exp ro(s.a) (o(((rr))-((sss ac(s, + P(5')s,a)V(s')ds da
dES
()) e log
cxp
aEA
yES
is a contraction map with constant ?/.
Proof. It has been proved that Tx,l Q is contraction map with constant (Appendix A.2 in Haarnoja al. 2017) if we replace with re Ac). Here we show that T V is contrac- ton map with constant y. Define a norm of V ||V V/ll supses (V( V2(s)) and suppose [[V 1/21 c. Then we have that
log / exp (o((s,)) (c(s,,)) +2 / P(s'\s, da
Tx (i(s)
AEA
YES
ro( - Ac(s, + / P(8')s,a)(Va(()) + clds' da
< log
exp
DEA
PES
- Ac(s, +7 /
log
exp ro(s,a)
)ds + 70 da
.ES
log
exp(1c) exp ro(s,a) Ac(s, +
J P(5')s,a)V2(5))0) )da
JAEA
Txdl (2(s) 7
V2)) yc. Therefore, || V2 Vall I 7C Similarly, get V)))
>
Vall.
17

---

Published as a conference paper at ICLR 2025
Lemma 7. It holds that 'x 2 Txxl (Qxx)) a) and Vx00 2 Tx,0 (s)) for any (s,a).
Proof.
@ (s, ro((s,)) )e(s, +7 / JEE
10x00' (s', - log n 1 (a'ls' (a'ls"))ds'
(6)
2 ro(s,a) -)((s. +7 /
P(s' s, 0, (1 1xx a²) ) log mn+1 ss)]
JOES
= =ro(s, a) - ac(s, + P(s'ls, a) log
exp(Q%.0 da' ds'
dES
A'EA
-xx00 (s.)
where (a) follow equations (2)-(3) in (Harrnoja et al. 2018 and (b) follows policy improvement theorem (Theorem 4 in (Haarnoja et al. 2017)). Similarly, we can get that
A; (6)
Q2,0 (5,a) log T&+ (a|s)], 1Qx (s, log (al1))
E Ba
log
exp((x, (s. (s,
JEEA exp ro(s (ocos)-)((saaa a) +7 f a) dd jees
log
7xi8
(VTB))
D.1 PROOF OF LEMMA I
It has been proved in Theorem in (Haarnoja al. 2017) that the constrained soft policy TA arg max, Jro (T) AJe(T) + (()) where we treat Tg Ac as the new reward function. Recall that the dual function is G)):)) max, Jr, (T) X(J((T) - b) + HT),, therefore, we know that G(); Jrn (===)) ((Jo(Tx;0) b) + H Since Tx.0 is the optimal solution and the G is differentiable to the policy, we know that (a
aG(A:9) for every (8,a) a) E S Note that
aG();0) = 1 OTla)(al))
o (s, ro(s,0) - Ac(s,a) log (aa((ass) =0
(10)
VxG(A:0) V.
Jn (Tx:))) {dc(TA:)) + (TT.OO)
(Jo(TTA:)) b),
- JOEA
Vx { (opxi((s,0)) ro(s,a)= xc((s llog logmrolals) dadds (Je(T(::)) - b),
VATs;o(als).. dTxxo(als) (5,00 rr((saa) Ac(s, log mxp((ass)) Jdlads
JSES JaEA (Je(TA;)) b).
(4) Jc(Tx:))
where (a) follows
10
D.2 PROOF OF THEOREM 1 AND THEOREM 2
We first prove Theorem 2 and then prove Theorem 1. The fundamental logic is that we first prove that the dual problem has unique optimal solution X² (0) and then we prove that the primal problem
18

---

Published as a conference paper at ICLR 2025
(L.e. the lower-level problem in (2)) and the dual problem have the same set of optimal solutions. Since the optimal solution of the dual problem is unique, then the optimal solution of the primal problem is also unique.
To show the optimal solution of the dual problem unique, we prove that the dual function is strictly convex by showing that the Hessian of the dual function to À is positive definite.
From Lemma
!
we know that VxG(A; = Je(T)))), therefore, we have that
GAG((A;))
Po(so) / TX:((ao(s0)
c(su, ao) +> / P(s (so, (s1)ds)
-
dagdsos
DDEA
SLES
/ VATAiO(a0lso))
((800 ao) +7 / Mes
Po(so)
P(si|so, ao) (si)dsi
SOES
ADEA
77 JIES
P(si|sa, ViQE Vo (x (s1)dss) dagdso:
= -- SOES
I VATrio(aolso)
Po(so)
(5:00(0000)))
ADEA
} Jdroddo
+ Tx,0 80)
P(silsn, ao)
(100(a111)) th "0 (81,
Keep the expansion, we can get that
J m/ (A) Q² (s, (s)a) JEEA I6TAID (S)
VRAG(A;0)
1 aEd
VA log (al)) Q as,a)s
SES
(aio(als)) Q Ve vetui(s) (s) Qe (s, a)dads,
(11)
SES
aEA
where (a) follows Lemma variable Xsa to denote Q
5
Note that VeTA: (s) Ea 1Q 10 (s, a)]. we use the random
(s, a). then its expectation is Ea (X = V. (s). We know
that the variance V Var(a Var(X,a)
E[XealXea E(Xea)]]- Therefore, we can see that the equation
is actually variance:
G)= = SoES *xa (s)Var(X
From the expression we know that Tx is always stochastic. Therefore, the variance and thus G is strictly convex. Therefore, Var(Xsa) O. Then know that GAA; the optimal solution X*(() is unique.
Since G(A; 8) attains its minimum at (#)), the gradient of G (0) should be zero, i.c. Je(Tx-(0):o) - = O. Let p' and d' be the optimal value of the primal problem and the dual problem. Since (AA;; f)) = max, x Jro (T) + H(T)) ) - A(J((T)) b), we know that G(); 0) 2 (T) + H(r)
for any (A, 0), which means that d' 2 p" Therefore, we have that:
(b) Jr. (T))))))) + ((TX-(0))) Sp:
where (b) follows the fact that (=()))-0)))) - = O. Therefore, ( is an optimal solution of the primal problem. Suppose the primal problem has another optimal solution n', then it holds that r' arg max, ()) However, it has been proved in Lemma in (Zouu et al.) 2017) that given an arbitrary X. the optimal policy of max, Jr, (TT) + H(T) ) X(Ja . b) is unique. Therefore, xx (0):0 the unique optimal solution of the primal problem (i.e., the lower-level problem in (2)).
D.3 PROOF OF LEMMA 2
Since (0) = arg min G(>; 0), we know that VxG(>*(0); 0) = 0. Therefore, we have that
dV XV (0): = 0,
d0
19

---

Published as a conference paper at ICLR 2025
G) + (0 = 0, VA AA
(12)
Now take look the term ViA From Lemma 1 we know that VxG(A;0) b Je(Tx;o), therefore, we have that
G(A:0)
-VoJe(Tx.))
-Vo
Po(su)
(a)0(a0)50) Q: (so, an aodddod0
ADES
aaEd
Po(so)
Twrolau|so)
Vo log (w: (0)) Q0 ou) + V ao) daodso:
TMIn
aDEA
- Po(3))
Tu: (o)
Vo log Tw: (00)
20*00(000 an)
Jeees
aoEA
Volc(su, an) + /
P(s so, ao) ao) ( (s1) dandso,
9IES
V log (aolls)))) QC I Twro(aolso)
Po(so)
(50. ao)
JDES
ADEA
P(si/s0, Vo
( TAp((alls1QE= a. (daadid dagdoo.
JIES
aaEA
Keep the expansion, we can get that
S.S./. AAEA
3T (s, (o (s, (s.a)dads.
* (8, a)
2 (0 -
Vonro (s) Qe a)dads,
(13)
JE6 JEA
where (a) follows Lemma 5
Now, we take the full gradient log Tx- (x)(0) to::
d dog T (appo(a))s) = V log ((0) V log ((()))((a))))) (X*
d8
(
Tarcorio V
(s, a)
V (s) - (Q 10) (s, - VATA (
Voro
Q (s, a)
Vera
V (8) (Qe =x=00000 (0)0 (s,a) Ve ((TM)))()))
Sees SaEA (a
Q (s,a) Vigro
T>r(one(s) (s)
Q: x- (s, adddd
SSES SaEA (TTAcoo(sss
Qe Quaiion (s, a)
VeTA VoTA (s)
Qe Q-2100000 (s, (s,a)dads
- (s,
-Verr (-TA(0):00() Cax-cenn (Q TT- (s,a) a) (03)0 (s)).
(14)
Sees Seen v (hM
Tai (v.a)
TAi
Q. (
where Craip
(b) follows Lemma 5 and
Juss Sucia gi (s,a)
9% (s,a) (s)
Q" 200 (s.a)dads
(c) follows
and 13) Note that we can equivalently reformulate C as:
F(s.a)) 11Q Vora (s,
Vorr
(s))Q= A. (x,a)
Cax.
(15)
B(s.a)=aa" A (.)) [QQ RAI" (s,a) Ve
AI (s, AM (s))Q²
Therefore, we can compute the hyper-gradient as:
d.J, ( (d)
dloo (0)
O TA (0)0 (s, (x
((
d0
d0
20

---

Published as a conference paper at ICLR 2025
cc) (o
(
(
Vore V (s)
Cax-cora (Qe (0 (s, (=
Q Q 10xe rone (s.
where (d) follows the standard result of policy gradient (Sutonn & Barto 2018). and (e) follows
14
D.4 PROOF OF LEMMA 3
Since (als)
& exp(QX;0 a)), from Appendix C
we can see that (ali(a)s) =
exp(Q (s.e))
exp(VA. (x))
Mogma+i (als) log o QI, (s - Vit (s) Qx a) + V2:0(s),
(0) Qx:0(s,a) Q2i(( a) + Vs:((s) - (s),,,,
Qxi(ss a) Tx000 T200 (0 (s, a) + (s Txx00 ( (V (s)))
10 (Qx;0) Tx (2x (s, a) =Tx (Vx: (Vx:))s)) -Tx,8 -T (V (s);
(d) > (Qro(s.a) 2x (s. 0) + Vx Vx (s)
(s. S. ( (= (s) V (s)
where (a) follows policy improvement theorem (Theorem in (Haarnoja et al.
2017 (017)) (note that
Qxi and V are the optimal Q/value functions under (A, 0)), (b) Tollows Lemma 7 (c) follows the
fact that the optimal Q/value functions are the fixed points of the contraction operator Txo and (d)
follows Lemma 6
Lemma 8. For any 0 and ñ 2 0, it holds that V2 G(Az; 0) ž Tal where TO is positive constant.
Proof. It has been proved in Subsection D.2 that V G G(;;)) > 0. To prove that V G(a 0) x ToI, first prove that V2 G(A; 0) is continuous À and then prove that the trajectory of Xn bounded within a compact set for any ñ 2 0.
From 11 we know that
ViAG(>;)) = ATIA Σ
xt Q QF V (s, VC
Qt QE (s.
Since TA;0, Q.Ai (s, a), and VOTAI (s. a) differentiable X, we know that V G(); is contin- uous to X. Now, we show that the trajectory of Xa is bounded within a compact set.
Da411 - X²)0] = | angAn - (2)
|| - x²(0)11² + alllA
an (g Aa - (²())),
-I1Aa - x=(0) + az :
an (VG 0), Az (²))
Onn (9Am.: VVG(() 0), An >²(0)),
S ||Aa - (0)1 + a 10: - 0) (())))
OR (9A.:0 - V,G(A>; 0), Aa - -x*(0)):
(16)
|| X 20)1 + a (0 12 anignn VAG(An; 0), An - (0))),
A (0)1 a 0m 1 coallgxad
VAG(AR: 0)1 1|Aa (1))).l
|A - 2² + a 19 00
+ ann C |In X²(0)1l.
--1 AA - 20 + + Cn x²
asga.ell
21

---

Published as a conference paper at ICLR 2025
R-1
$ || - (0)1 + aollssn: -XX0))
+ an Cy
1=0
(K) 1 || - X00) + a Cmax )² + tan09"1100 Xx + +ax On
a-1
Cmax
ag +
k=0
(17)
where (a) follows 18 and (b) follows that |ggxel =111- Jx(o) S bb Sumx 1-7 that Ex=1 On7 Zi= 0G is bounded. Since or x 1/19, we know that Ei= 0₁ O(nt-1).
Now we show
Therefore, we know that 0; Ei= a, = O(n 27) I C' where C' is a positive constant. Therefore,
2m=10
a a R-1
xc 0i S C <CE.= 7at is bounded. Now, we sum the both sides of 17 from ñ
to N 1:
N
Da+t1 - (0)1)
N-1
<
|m (0)11 + (b + Cmax ja + +anC7*11> X-(0)]l + a.,
Σ
ag(b (b+ +
Cmax
1
0
1=0
1 (0)11
P-1
â-1
1 - x01111++
+
Σ
Cmax )2 + Can - -11(0)1 + Cn ah + 1-7
Σ
ar(b +
Cmax
n=0
1=0
n-1
S 1A - X-(0)1
a + Cmax )2 + a2 X2 0.7 Cy"
Σ az +
Gmax
1-7¹
7=
Note that an x [x+1) and 7 E ( 1). it is obvious that |Ax - (0) is bounded. Therefore, the trajectory of Ay is bounded for any ń 2 0. Therefore, we can always find positive constant TO such that VRA G(A; 0) 2 Tgl.
D.5 PROOF OF LEMMA 4
We first quantify the gradient approximation error VVx(();)) gxxe| and then show the convergence of the middle loop.
O)) gx:0l = |Jo(T(x0) - Jo(fx:e)l
= -V [ (s,a) us (3) (s)a))c(s, dddds
JES SOEd
I Cmax
JOES JEA
(S a) -b a) dads.
(0)
Cmax Co
Q2i((s 2 (s, a) dads,
JOES JAEA
Crnxx CaCSA max (sa)ESxA
(10x: a) Q tx. (s,a))) }}
(6) Cmax CaCsA
Nu
max (x.a)eSxA
(1Qx: (x
Ny
(18)
where (a) follows step (iv) of equation (64) (Zng et al. 2022) and C2 is a positive constant, CSa
can be any positive constant that is larger that the area of S x A, (b) follows the proof in Subsection D.4 and = Ca SA max (IQ a) - (x, a)]3.
Now, we quantify the convergence of the middle loop. From the expression 202
11 of V G(A: 8), we
know that ||V 0)1l <
(I-7). From (16), we know that:
an -(x 0) ||Aa X (0)]] ||Aa+t X² (
22

---

Published as a conference paper at ICLR 2025
+ a 0h a (9 VxG(AA;)), An - (8))
||An X2 X0(0]11 || X" X0)0) + o 112 on I1gAn VAG(An; 0)1 ||A (O)I::
(0)
W Ι - X- X0)0112 || - (0)0 + 19ll 11 a
In |A - X X0(0)112 - || - (0)) + a (b + Cmax
1- /p + O
(19)
where (c) follows 18 and the fact that || x^)0)] bounded (proved in Lemma 8). Ö is
positive constant, and (d) follows that |1x00l| = | | ( S b Smax
1-y" Telescoping (19) from
to N
R-1
Σ R=0
an (G - -((())))
N-1 $ |1:0 (0)011 - 1 (0)1
N-1 a
Σ ab + Cmax )2 + Σ カ==
n=0
Since 0g 1b+ and 17 E (2 1). there positive constant D Dinx such that Σ 0 +
ja + カ= a C < D maax. Therefore, we have that
R Σ No
N-1
G((A::)) 0) G(A -
Σ
an -(()))))))((())).)))))))))))))))
A=0
< |1 (0)11 - 11xx - (0)11 Dimax:
N-1
1
G(An: G(X(()):)]] NI-1
1100 (-(0)112 - 11 (0)0) + Dmax
(20)
Therefore, we have that
( 2 1 N-1
(I²(0) x2(0)1l (e) 2
- 2 (c({()):)) - G(X(()): - 2 Σ
G(4; 0) -((())))
G():());
TG
To n=0
0 N1
(21)
where (e) follows the fact that G(A;0) is Ta-strongly convex (Lemma 8 (f) follows Jensen's s N-1 inequality (note that x(0) = */4 En Xa), and (g) follows 20).
Now, we take look at the term
log x (a)) - log मे x (a|s)),
< logg ( log क 100) (a|s) + log T (a)s) - log A (a
(N) 2cmax
[\(0) (0) + log (als)) log (a)s))
1 γ
$
0
N1-6 )T
1
(22)
5 such that V log Tx:0l S 20max Lemma 9, The upper-level loss function ( (8):)) is L-Lipschitz and L-smooth where and L are positive constants. Moreover, it holds that |lgoll < L and [ogol) $ L.
where (h) follows Lemma
1-y and (i) follows 21 and Lemma 3
Proof. This suffices show that the norms [TX and [[V2J, are upper bounded by and L. From Subsection D.3 we know that
d.J, (0)))
d0
23

---

Published as a conference paper at ICLR 2025
- (
Voroo V (s) Cix (Qe (00000000ss Ve Vo (s))
eo (8.0))
where C C2)(0):0 = (X
Since [[23),(( S 1 Qa (o (s,a) Voro (s))Q==*(s, (s,a)]]] S
(Lemma and G(; = ((s.a))= TAI (.)
2Cr0mux Therefore, it holds that 2C we know that ||O || I (1-7)=TG
(1-(())((000))) do
1
[(
(222 + 202cmax 20max
202
C1 1-7
5 1-7
= L.
(23)
(1-7) -T= 1-
Similarly, we can see that ||9o|| I L.
Now, we take a look at the Hessian term V (Tx- We define h²" (8, a) Forr (s,, az)
Veoro (st) (x) (st, a1) (s)) at), H (s, - E E nhh² (st, su 8, 0,00
a). and (s) & E [E= 7th² (st, as) = s). We know that H (s. a)ll I L and || H (8) S
L. Therefore, have that
V (0)))))))))))= -V J
S JODES
/
Po(oo)
JADEA
(50, ao dagdso:
= =ff Po(so) /
(2) (50
DES + (=)10))((005s)) VH* (so, ao dandso
AAEAA
=f Pa(so)
[ aaEA (x. 80) V log x (0); - H= on)
( ao) +7 /
+ (0=00)((aolso))
P(si|su, a(s VH daydsu)
ME
Keep the expansion, we know that
( (8, a) Vh V + V log x-(o))(aas)) Ha (s, (s,aa dads. dads.
(24)
Since V Vlog (())) and H bound Vh a)ll. We aim to bound || (8, by bounding each term =
both bounded, the only thing left is to
in || "> (s. Note that VC = [[ no ((()))) - G(A(())))
TA-(010 Therefore, suffices to show that dl Q (s,a)l and d Q x (0:0 (s, a) are bounded.
d Txrcnn (8,a) = d
20 E
d0 Q
Vero(s)lar)lo s, a0 = a),
d0
= d0 1
Po(oo) / S
@=rrcone no) (Tx=(0))(u0lso))
+ (x)
VQ (s0:
(so. ao) daodso.
//30 Po(so) /
Tx-(0):0(a0(5))) V log TA (0)o0 Quarror (so, ao)
d0
S0ES
SSDEA
(Végro(so. ao) +0) of MES
P(si|so, an) VV Vara (s1)ds) ) dandoo.
+ =T0(0)00(00l0))
24

---

Published as a conference paper at ICLR 2025
Keep the expansion, get that
Qvornin (s,a),
dn
- (((a)EBXA ( Vooro(su, ao) + V log T (A) so) Qvora (so, au) dads.
Therefore, we can see that ll Pwrr (s, a)ll S
Cs 1-7 + 20max (1-2)
202
C2
re(1->)
(1-7)
(1-7)e-
Similarly, we can also see that d constant L such that (Tx)) - L. With the same procedure, L.
de Qe ** 10)0 (s, a) is also bounded. Therefore, we can find positive
see that V Vagall
S
Lemma 10. The hyper-gradient approximation error is upper bounded, i.e., 11 JT gell S 067 + x.-a).
Proof.
d
Jr (TT((10)):)) - goll.
d0
[
y Rh (1) ()) mar (1) (s, (s. dads
(0)
$ (1 (1-7) L)
I (uns vs (0)00 (s, a)lldads,
(S.A)ESXA
$ $(1 511-7)L
(100000 - 310 a)ll + +| "1000 a) - adldds,, dads,
(SSAEESXA
(1 CSA
(ax1102000000.0)
100000000000000.0 a))) + 1ax (x-
(
3 811-7)C.CAA
7 maxi112(0)
2 (8,)0)
+ maxx11220)000..0) [1)100))))
(r (1-7)CaCCA
1mxx112100,000.00 - 2 N0oy0 (s, as) + Cmax (A) A(OrI]
2 (1 C)Csss
(c)
NI-
where (a) follows (3 (b) follows step (iv) of equation (64) in Zeng et al. 2022). (c) follows
18
(d) follows the fact that ([V 'Qa (s,a))
Cuaas and (e) follows Lemma
D.6 PROOF OF THEOREM 3
We define a function f(0) such that V Vf(0) = and f() = Vege, therefore, we have that
_
I
f(0n+1) < ((0)) (V.( 0.+1 - 0,)
110n+1 - -0al13
2
= f((w) (. L [Vf(0))11
L
+ 3n Vf(0.)) < f(0.) f(0n+1) +
(25)
2
Telescoping 25
from n = to 1, we have that
N-1
L
f < f(0)) -((0)) + Σ
V
2
0=
25

---

Published as a conference paper at ICLR 2025
VS(0,)) -
N-1
Σ
B Val - -(00)- f(0x)
(0)
-
+
Σ
L3
N
2
R=
n=
n=0
N-1
(1(6)) (((0)) ([00))[(ON)) + Σ
L
( 0
[VS(0)) < NI-7
(26)
2
Ni-n
m=0
where (a) follows Lemma 9 and (6) follows the fact that _n S is bounded as B, = (m+1y1
and
7 E (2) 1). Therefore, we have that
N-1
N-1
+ +/V - V..((x)0000))?)
m=0
(c)
1
0
N1-9 + +
N2-24
where (c) follows 26 and Lemma 10
D.7 PROOF OF THEOREM 4
We first define F(0) J. (+)-) Theorem 10 in (Agarwal et al. 2021) shows that the policy gradient method can achieve global optimality asymptotically under softmax policy parameteriza- tion. The constrained soft policy ( lim *A(0)=0 can be regarded
softmax policy parameterized by Qx-(07e but the decision variable in our case is 0 instead of Qx-(9):9- How- ever, we can still build on Theorem 10 in (Agarwal et al., 2021) by building connections between and @x)00): In specific, in order to use the result Theorem 10 in (Agarwal et al. 202I), we need to prove (i) F(O,)) is monotonically increasing, i.e., F( 2 F(B,) for any 2 0; (ii) dF(O)
a4xx.0;;0 if dF 0 where Qx (0)0 is vector with the length [S] |A| whose components
(s.a)SS..A Once proving these two, given that lim lim dF(0x) are {Qx-10).0 (s. a)
de
from Theorem 3 we can use Theorem 10 in (Agarwal et al.) 2021) to prove Theorem
Now, we first show that F(,.) is monotonically increasing. This is a straightforward result of Theorem 10.15 in (Beck 2017 if we choose Bn Σ 1.- Note that Theorem 10 in (Agarwal et al,
2021 requires B, <
so that we can choose 3n S minf 1 (1-y)² 8 }. We next show that
dF(P) (F()) if
dQ
49
We know that dF
dF
dQ dQ so that it suffices show that
4Q 4Qx-(0).(5.0) d8
# for any
(s, a) and any 0. Therefore, we have that
aQx- (0)
V.QA Tx V
d8
(s, a)
@ (1, (x))),
(27)
where (a) follows 1 and 8 Now, we first prove that the term Q (s,
1 Ee= 7 Vora(s ar) is nonzero. Define 1(0; s,a) - E [E ('ro( at)l0 = 8, 4. a). and therefore Vol(O; s, A E |E 7'ro(st, at)sso ag a). We use vi" (s' a) and d al to denote the state and state-action visitation frequency when the initial acte-nction is a). Now, we take look at the Hessian term s.):
s,,,())A,saa).
JaJE -Vioress,a + / P(sl)s,0) / [xelayls) ddad
JaiES
Vooro(s, + J,Es P(sils,a) / VOtro(allsi) Vol(0; 81,
aTEA
+ + Vool(0; 1,0;; 1100;,,810000 01) daydsi.
26

---

Published as a conference paper at ICLR 2025
Q a)) we can get that Keep the expansion and note that Vol(0: s,
SEES l. I 1 () Txo Q (s', a d')dd' da'd'', ds',
(iol(0;
(s
Qvoro (s', a ddd ds',
= JES O'EA
/ (
I (a.o(a'lss)
@ (s',a') Voro (5') Q (oon0( (5', (5',a') ya'd (28)
JES
AAEA
Note that Vorr (8") = Eey F 1Q TAi (8', a')] If we use the ran- dom variable Yo to denote Quer (5', a,),, then its expectation is Bo'nTs (Yorr) Vara (s").
where (b) follows Lemma
We know that the variance Var (Yr(aa)))) = Li
B(Ysaa)) Therefore, we can see that the
equation 28 is actually a variance:
Vool(0; (0,((),,,))==- JEES
u (
= 0.
Therefore, the function 1(0; X, 8, a) is convex in 0 for any (A, If Vol(O; A, means that [(;; X, 8, achieves its optimum. However, 1(0;), 8, a) does not have an optimum, i.e., I(;; X, s, a) can be infinitely large or infinitely small because re(s, a) can be any arbitrarily large value. This is contradiction, therefore, 1(0;),s,a) & for any (A, a). Then, ((0) (0), Q
s,a) O, this
(om (s, a) #0.
Recall from (7 that dQ2 (Q = Q (s, - Qe (0110 (s, (. VX) For any (s,a) & C,
d9
Q (5 (s.a) a) because the policy (x= satisfies the constraint of the lower-level problem
in (1), i.e., avoiding the set C. Therefore, ( = Q (s,a) * O. For any (s,a) E C, we know that Qe (0 (s, a) c(s, a) because the policy (0)(0); avoids the C unless its starting acte-nction pair is in C. Therefore, can also design c(s, a) such that Q (s)
du
ess a) (0) + 0 for (s, a) E C. Therefore, we can ensure that 4Qa- dF(O) 0 if dF(0)
+ for any (8,a)
de
A. Therefore, dQx=<6).
do O.
E RELATED WORKS
XRL methods that has the potential to be used to improve RL performance. There are some XRI. methods that have the potential to improve the RI. performance even if they do not mention that they can improve the RL performance. Value-max (Amir & Amir 2018: Huang et al. 2018) use the value function (ss) to identify the states with highest value as critical points. We can perturb the actions on these critical states to improve the RL performance.
Constrained reinforcement learning (CRL). The lower-level problem in is CRL problem. The current works on CRL have two major categories: primal-dual approach and primal approach. [ The primal-dual approach (Achiam al. 2017 Tessler et al. 2018 Stooke et al. 2020) converts the
CRL problem into an unconstrained optimization problem by using the dual method. Our approach be categorized as a primal-dual approach. The primal approach (Liu al.) 2020) Chow al. 2018 Xu et al. 2021) enforce constraints via various designs of the objective function or the update process without an introduction of dual variables. However, these previous methods on CRL may not be suitable to the context of constrained bi-level optimization because they cannot guarantee that the upper-level problem in (I) is smooth after the lower-level problem solved. The non- non-smoothness of the upper-level problem can make the constrained bi-level optimization problem difficult to solve. In contrast, our approach ensures the smoothness of the upper-level problem in (T) because we derive an analytical solution of the constrained soft policy and this policy is smooth w.r.t. 0.
F EXPERIMENT DETAILS
The code was running on a laptop whose CPU is Intel Core 19 12900k and GPU is NVIDIA RTX 3080. The operating system is Windows 10. We use neural network to parameterize the learned
27

---

Published as a conference paper at ICLR 2025
reward function. The neural network has two hidden layers where each hidden layer has 64 neurons. The activation functions are respectively ReLU and Tanh.
The delayed MuJoCo environments. The delayed Mujoco environments are widely used in RL improvement literature (Zheng et al.] 2018 Memarian et al., 202T Oh et al. 2018) where the reward is accumulated by 20 time steps and only provided at the end. For example, for episode with length 100 (i.e., 0. ,9), only the time steps 19, 39, 59, 79, 99 receive nonzero reward while all the other steps receive zero reward. The time step 19 receives the reward that is accumulated from time O to time 19. We can see that the delayed MuJoCo tasks have sparse reward.
We also conduct experiments on the original MuloCo environments, which widely used in RL literature (Xu & Zhu 2023b
2024) We will explore some other tasks like motion planning (Xu & Qiao al. 2023 2024) and finance (Yu et al. 2025 Ke al. 2024).
Zhu 2022
F.1
EXPERIMENT RESULTS ON DENSE REWARD
To show that our method can also improve the performance on dense reward scenarios, we include the experiment results on the original MuJoCo environment (where the reward is dense) below:
Table 3: Experiment results (original MuJoCo environment with dense reward).
SAC
UITLITY
RICE
LIR
HaliChectah
686.40 = 41.24
824.42 + 42.18
701.44 + 45.83
716,62 = 39.17
718.25 + 42.10
Hopper
238.14 29.94
348.16 26.32
242.99 + 19.10
264.86 27.11
277.58 1 22.69
Walker2d
182.21 24.14
269.14 + 25.08
189.19
39.11
197.63 = 21.39
214.16 + 21.24
Ant
299.79 = 26.53
421.63 + 25.16
322.26 + 34.15
340.14 = 26.53
351.14 + 28.92
Table shows the final results on the original MuJoCo environment (dense reward). We can observe that UTILITY achieves the highest reward and largely improves SAC.
F.2 ABLATION STUDY
Since our method has two components to improve the performance: the shaping reward and the constrained formulated by the "misleading" state-action pairs. Here, we separately study the effect of the learned shaping reward and the constraint. In specific, we test the performance of the shaping only method and the constraint only method, and provide the results below:
Table 4: Ablation study for dense reward.
SAC
UITLITY
shaping only
constraint only
HalfCheetah
686.40 + 51.24
824.42 1 42.18
764.25 + 48.11
701.19 + 47.43
Hopper
238.14 + 29.94
348.16 1 26.32
311.78
34.24
268.15 1 42.66
Walker2d
182.21 + 24.14
269.14 = 25.08
242.18 = 29.62
196.77 + 22.19
Ant
299.79 + 26.53
421.63 + 25.16
396.05 + 29.18
317.12 + 34.59
that both the learned shaping reward and the constraint can improve the performance, and the shaping reward has larger impact. Moreover, even if we only use the shaping reward, the performance is better than LIR. This is because our shaping reward formulated by the high-level explanation. Even if we only use the constraint, the performance is better than RICE. The reason that the "misleading" state-action pairs we find are the points that lead to the failure, and thus avoiding these state-action pairs can improve performance. In contrast, RICE finds the states that are most influential cumulative reward, however, these states may not be the states that lead the RI. agent to be non-optimal.
the domain knowledge
F.3 FIDELITY OF THE GENERATED TWO-LEVEL EXPLANATION
The fidelity means the correctness of the two-level explanation (Guo et al. 2021 Cheng et al.) 2023). Since the two-level explanation is explain why the RL. agent (i.e., SAC) not optimal, one way to
28

---

Published as a conference paper at ICLR 2025
validate the fidelity of the explanation is to whether the performance improves after we improve from the explanations. From the last two columms in Table we can see that both the high-level and the correct explanations because both the shaping only method and the
low-level explanations constraint only method improve the performance. Moreover, the shaping only method (the fourth column in Table 4 has higher cumulative reward than LJR (the last column in Table 3 and the constraint only method (the last column in Table ( has higher cumulative reward than RICE (the fourth column in Table 3) This shows the high fidelity of our two-level explanation.
Table 5: Fidelity comparison for sparse reward.
shuping only (ours)
shaping without
constraint only (ours)
RICE#constrint
Delayed HalfCheetah
383.15 15.50
695.63 33.66
611.08 = 39.41
122.15
22.80
369.11
19.10
Delayed Hopper
192.90 + 27.18
280.10 18.41
255.18 + 16.57
210.12 15.77
181.45 + 12.11
Delayed Walker2d
134.91 1 20.80
211.37
18.64
191,15 4: 11.26
175.06 15,27
140.26 11.33
Delayed
68.11 12.32
88.18 1 8.08
77.11 = 5.52
75.16 6.58
63.11
Table 6: Fidelity comparison for dense reward.
SAC
shaping only (ours)
shaping without r
constraint only (ours)
KICE+constrain
HalfCheetah
686.40 :: 51.24
764.25 + 48.11
710.99 = 35.16
701.19 47.43 43
672.15 1: 25.16
Hopper
238.14 29.94
311.78 + 34.24
279.12 = 27.44
208.15 42.66
211.03 + 21.20
Walker2d
182.21 24.14
242.18 1: 29.62
213.15
18.94
196.77 22.19
177.19 1 15.33
299.79 26.33
396.05 29.18
339.14 22.85
317.12 34.59
288.19 18.02
Tables 5 and [ show that both the high-level and low-level explanations of our method have higher fidelity. The method "RICE+constraimt" has even worse performance than SAC because the critical states influential to the cumulative reward may be the states that lead to high cumulative reward, and thus constraining them may even make the performance worse. However, even if we not constrain these states but the fine-tune method as in (Cheng et al. (223)) instead, our constraint- only method (the last column in Table 4 still outperforms RICE (the fourth column in Table 3). For the high-level explanation, can see that our shaping only method achieves higher cumulative reward than the method "shaping without r - pr This shows the high fidelity of our high-level explanation.
F.4 HOW TO ACCELERATE THE TRIPLE-LOOP ALGORITIIM
The total iterations of UTILITY is N x N x Ñ where N is the iteration number the outer loop, Ñ is the iteration number of the middle loop, and is the iteration number of the inner loop. While the triple-loop structure looks computationally expensive, in practice, we can significantly accelerate the algorithm using warm start in the inner loop and middle loop. Take the inner loop as an example, given the shaping parameter 0 and the current dual parameter Aa. need to compute the corresponding constrained soft Bellman policy Txa;0 in the inner loop. Instead of starting from a random policy initialization,
use the policy a 1:0 learned in last inner loop as the initialization where Xa 1:8 is an approximation of TAa 1:9- The intuition behind this is that since Xa and AR-1 are close (only different by one-step gradient descent), close. Therefore, using # as the initialization makes it easier to approach TA Therefore,
expected that A and Tx are
the warm start trick enables us use fewer iterations for the inner loop, i.e., the iteration number Ñ reduces. We use the similar warm start trick for the middle loop to reduce the iteration number N.
G POTENTIAL SOCIETAL IMPACT
This paper has positive impact which is to improve the performance of RL agents. However, the paper also has potential negative impacts. Since the two-level explanation identifies the mistakes made by the RL agents. A malicious entity may use these weaknesses or mistakes to launch attack to the RL agents. To alleviate this issue, one solution is to keep the demonstration of the RL agent private, so that the malicious entity cannot get access to the demonstration and thus cannot find the weakness.
29

---

Published as a conference paper at ICLR 2025
H LIMITATIONS
One limitation of the method is that requires to interact with the environment. Therefore, one future work is extend this method to the offline RL setting.
I EVOLUTION FRO FIGURE Ic TO FIGURE
Iteration=0
Iteration=50
(c) Iteration=
Iteration=1550
Iteration=200
Figure 3: Evolution from Figure 1c Figure 1
For this scenario. we run UTILITY for 200 iterations. Figure 3 show the evolution from Figure Ic (iteration=0) to Figure Id (iteration=200). From the evolution, we can see that the learned reward becomes closer and closer the ground truth reward and the learned policy becomes more and more optimal.
30


Reviews:
[{"review_id": "3t5xv9PKvB", "reviewer": "Reviewer_euxG", "summary": "The authors propose a new approach of using explanations for reinforcement learning developed in the XRL community to improve RL algorithms. The design a two part method – the first part generates the explanations (two levels – learning an implied reward function and identifying ‘misleading’ states) and the second part utilizes this explanation to improve an RL agent’s performance by formulating a constrained bi-level optimization problem. Specifically, these two levels involve learning a reward shaping function and then a policy to optimize for this shaped reward. In addition to providing theoretical guarantees, the authors empirically demonstrate their method on four MuJoCo environments.", "questions": "1. The key idea in this paper of learning an implied reward function of a learned RL agent and also identifying the state-action pairs that cause the agent to be non-optimal to be used in learning a new policy seems novel and interesting. However, wouldn’t this require re-learning? Wouldn’t generating the explanations require either access to the environment for generating additional data or knowledge of the environment parameters such as transition function, maximum attainable reward, goal state etc.? Can the authors add a discussion on this topic detailing what extra data and computation are needed by their method?\n2. Can the authors also compare their approach with other RL algorithms from literature where the other algorithms have access to the same total quantity of data used by the authors’ approach? The current comparison in Fig 2, does have total environment interactions, can the authors confirm if these interactions also include the additional data mentioned in Point 1 above?\n3. It is known from the literature on inverse RL that learning a reward function from data is a problem that can have multiple solutions. Can the authors describe how they address the non-uniqueness problem and how do they choose a single reward function?\n4. Can the authors clarify if different plausible high-level explanations (learned reward functions) can lead to different low-level explanations (misleading states)? If so, how do the authors propose to select the best high-level explanation? What is the authors’ opinion on using feedback from the RL agent’s performance post using their approach as input to improve their explanation generation. i.e., can their approach benefit by making the process of explanation generation and policy improvement iterative? \n5. Can the authors confirm if misleading states are states where the agent does not act optimally as per its own Q-function? Is this definition valid even if the Q-function learned by the RL agent is not accurate, which is often the case in several practical RL problems?\n6. The authors state in Line 222 that: “However, our case is different because we only need to learn one precise Q-function, which corresponds to the specific policy $\\hat{\\pi}_A$”. But won’t the authors need this Q-function to be accurate for also state-action pairs that are not visited or rarely visited under this policy? This is because in defining the $l(s,a)$ metric, the authors use a max over all $Q$ values of a given state $s$. How do the authors propose to learn these values accurately or adjust the definition of their proposed metric to ensure that this does not impact their approach?\n7. Can the authors clarify if the shaping reward function requires access to the ground-truth reward function? If so, how do the authors propose to obtain this?\n8. Can the authors specify the loss function used for learning the proposed shaping reward function used in equation (1), i.e., can the authors provide a discussion on the gradients proposed in Lemma 2 and its relation to this shaping reward function?\n9. When the authors claim uniqueness of solution in Theorems 1 and 2, it would be helpful if they can clarify whether only $J$ value is unique or is $\\pi$ also unique?\n10. Minor: There is a typo in “RICE+constaint” which occurs twice in the paper (Lines 513 and 514).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The authors address an important problem in RL – which is using explainability to improve performance. Their framework “UTILITY” is a novel approach to using XRL to optimize the RL agent’s performance. \n2. The paper is well written. The authors provide a good background literature survey, explain the problem and their motivation well and also clearly describe their solution.\n3. The authors provide a detailed theoretical analysis and proofs for their claims. While I am unable to verify these completely, I have not found any error. However, I think it will still be helpful if the authors provide a proof sketch or outline in the main paper itself.", "weaknesses": "1. Can the authors add a discussion on stochastic environments and the impact of the same on their approach? I think the current approach has only been demonstrated in deterministic environments. Specifically, can the authors discuss how their theoretical guarantees and empirical results might change in stochastic environments, and whether any modifications to their approach would be needed?\n2. The current method seems to need a lot of environment interaction. Can the authors detail the additional interaction needed for their approach?\n3. Can the authors add a discussion on the scalability of their proposed approach, in terms of state/action space size, or computational complexity as the problem size increases?\n4. The impact of the error in explanation on approach and final policy has not been covered, i.e., errors in reward learning and Q value error and behavior policy learning error. Can the authors address these with a discussion in the paper? Specifically, can the authors provide a sensitivity analysis showing how errors in each component (reward learning, Q-value estimation, behavior policy learning) impact the final performance with some discussion on potential mitigation strategies for these errors?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZSZ4nFXZB6", "reviewer": "Reviewer_sY65", "summary": "This paper proposes an explainable reinforcement learning (XRL) based policy update algorithm. The algorithm utilizes notion of \"misleading state-action pairs\" in the agent's current behavior that lead the agent to perform poorly compared to the optimal policy. These misleading $(s, a)$ pairs are then to be avoided by the agent to be able to act optimally. The criterion for defining misleading $(s, a)$ is whether the advantage of the policy is negative. For these misleading experiences, a positive cost is assigned. Parallely, from the agent's demonstrations with application of inverse RL on top, a reward function $\\hat{r}$ is recovered. This reward function is then used to explain agent's current actions, as the reward function explains why a particular action is taken in the demonstration by assigning higher values. The misleading $(s,a)$ constitute the low-level explanations and the IRL reward function $\\hat{r}$ the high-level explanation. A bi-level algorithm is then set up, where the agent has to learn a policy such that its low-level cost is bounded, i.e., not many disadvantageous actions are taken while high-level **ground-truth returns** are maximized. The first level includes both the cost $c(s, a)$ and the reward shaping using ground truth reward $r$ and IRL reward $\\hat{r}$.\n\nThe paper then goes into the details of optimization where a dual approximation of bilevel algorithm is proposed that substitutes reward function with dual variables. Further, theoretical guarantees are provided about the convergence of a practical approach that uses low-level $c$ and high-level $\\hat{r}$ and outputs an optimal policy. \n\n\nAs for the experiments, this paper uses sparser variants of the MuJoCo environments, and tries to answer how much effective improvement their algorithm UTILITY provides against other XRL techniques when combined with standard SAC algorithm. Paper provides ablations for different considerations in the modeling of the algorithm. By removing low-level and high-level explanations, paper shows why both are needed to get superior performance.", "questions": "Suggestion -- I would suggest the authors to use RL-consistent notation and substitute -A(s, a) instead of l(s,a). This will also simplifies elucidation of the technique mentioned in the paper.\n\nQ. Why did you take the dual into consideration? How does that help the algorithm? I mean, please provide description about how the dual forms an upper or a lower bound on the true objective against which this approximation is being taken. Further adding details about how this approximation is still valid with its loosening of the original objective would be helpful too. I find it hard to think that dual formulation for making the lower level optimization unconstrained and convex, and then using the optimal dual solution as reward function is still a valid way.\n\nSuggestion -- To help illustrate UTILITY's explanations during the training iterations, it would be helpful if the authors provided evolution of Figure 1-c to Figure 1-d tracked across policy updates. This would support the utility of the UTILITY.\n\nQ. I tried finding how the cost function is defined for the lower-level optimization. However, I could not find a technical section on definition of $c(s, a)$. Can authors please shed some light on this? Also, how do you decide the top-K misleading state-action pairs? This also was not mentioned anywhere, or I might have missed it.\n\nQ. Can authors please demonstrate UTILITY's execution on a few more gridworld settings with reward variations: sparse, dense, multiple terminal states, etc.?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. I find a close connection between the proposed algorithm and algorithms in generalized policy iteration class, because (1) avoiding disadvantageous behavior under current policy is analogous to improving the policy (imagine giving a leeway when taking a greedy action and instead use budget over choosing worse actions), (2) learning a Q-function to understand the policy's performance is analogous to policy evaluation. So, as long as these two steps are implemented properly, the proposed algorithm is bound to improve the policy, and eventually help reach the optimum in state values. Again, I might have skipped the specific details about the convergence constants, learning rates, etc., but the algorithm seems about right.\n2. The idea of finding K misleading $(s, a)$ to inform the policy update is very neat. This in itself is valuable as many times a post-hoc policy analysis might be necessary for identification of its limitations. The paper does a great job at also integrating the misleading with policy's current beliefs and using ground truth rewards to rectify those beliefs.\n3. The paper is well written, the ideas are well presented, figures are illustrative. I like how authors wrote \"challenges\" bullet points for every proposed method and then provided solutions for them. The theorems and lemmas are invoked appropriately. Some parts were not clear to me, which I have asked in the following sections. Nevertheless, I enjoyed reading the paper for its theory-grounded explainable RL approach to policy optimization.", "weaknesses": "There are a few weaknesses of the paper: \n\n1. The notion of l(s, a): The paper assumes access to environment interaction and hence I am considering that there is enough knowledge about the actions per given state. If that is so, the way l(s, a) is defined by the paper in definition 1, is simply negative of advantage function for policy $\\pi_A$. \n2. The choice of MuJoCo environment and plausible issues: MuJoCo environments allow the agent to complete an episode for full length while other environments might terminate early if certain states are reached. Subsequently, MuJoCo environments provide correlated \"second-chances\" to the agent to make similar mistakes and the misleading effect is easy to capture. The algorithm's performance depends highly on the coverage of the demonstrations given by the earlier policy. MuJoCo would end up allowing coverage over optimal state-action visitation distribution and hence the algorithm proposed in the paper would work with single low-level and high-level explanation. However, this will not hold in the other complex environments. Also, in the case of sparse environments, where the agent has not received enough feedback from the environment to capture the \"misleading\" state-actions, the algorithm might misattribute notion of misleading.\n3. Issue with complexity: To me, the bilevel algorithm seems an overkill if explanations are to be used for RL policy improvement. First of all, the idea of low-level optimization is simply to keep the policy improving by ensuring only a budgeted amount of disadvantageous actions are taken. Secondly, solving an IRL problem simply to gather a reward function to compare with ground truth reward function and reshaping it for policy update is again analogous to policy evaluation, done in an unorthodox manner. Why should we not simply assign meanings from XRL literature to the policy iteration or rather generalized policy iteration (GPI) procedural steps and keep training the RL agent. \nI don't see in any way the explanations during the training are helpful for the humans. The complexity of bilevel approach makes it even harder for the human to understand how various non-linear function approximators might be playing their role. Moreover, I am afraid that allowing SAC the same amount of compute as UTITLITY for per-step improvement, i.e., giving access to online rollouts and updating the policy under UTILITY's computational budget, we might end up with performance as good as that of UTILITY, if we are comparing only the extent of improvement to underlying RL policy updates.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J55xNch1GW", "reviewer": "Reviewer_KPMH", "summary": "The paper proposes to solve two problems: 1) the incomprehensibility of reinforcement learning agents and 2) the suboptimal performance of the agents. It introduces a two-level explanation, consisting of a high-level explanation that illustrates the reward function the agent has learned internally. The low-level explanation shows the state-action pairs where an agent chooses suboptimal actions with respect to an optimal policy. The suboptimal state-action pairs are called misleading. They use the low-level explanations to explain when it is infeasible to utilize the high-level explanations, for example, environments with high dimensional features. The misleading state actions pairs and the high-level explanations together are used to improve policy performance via a new loss function. A new algorithm is proposed to solve the loss function introduced to improve policy performance. The proposed method is experimentally tested in several MuJoCo environments.\n\nNote: my expertise lies in explainability, thus my opinion on the second contribution (model performance) should weigh less.", "questions": "- Definition 1: Why is the Q-function taken with respect to $\\pi_A$. Isn't $\\pi_A$ the suboptimal policy?\n- Line 245: What is the functional form of the reward shaping function?\n- Line 249: Why do we not want to visit misleading states? I thought it is where the agent takes a suboptimal action. Can we not fix the problem by just taking the optimal action when we encounter a misleading state?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper proposes a new way to utilize misleading state action pairs to improve policy performance. Furthermore, they explain why the policy did not attain optimal performance in the first place.\n- The paper has done a great job on covering related work and contrasted to them.\n- The method proposed shows better performance in terms of return than all the baseline methods. And has good ablation tests to see which part of the method that contributes most to the results.\n- The authors have included code with the paper.", "weaknesses": "- It is unclear to me why the policy would learn the wrong reward function internally in the first place. I might have misunderstood something.\n- It is unclear why theorem 1 is true. I do not see reference to any proof.\n- I am unsure how the performance would translate to environments with high-dimensional features. For example, Atari environments. Is there a reason why only MuJoCo environments are chosen for experiments?\n- It is unknown how useful these explanations are to human users, and who is the target audience of these explanations.\n- The main novelty is on improving model performance using explanations. The explanations themselves are not new to this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gjoFhyi6ZJ", "reviewer": "Reviewer_iSK2", "summary": "The paper proposes to improve an RL policy by using IRL to learn a reward for which the policy is optimal, and then use this IRL reward to shape the original reward by adding two terms: (1) an additional reward proportional to the difference between IRL reward and true reward, and (2) a constraint that the policy does not visit significantly suboptimal states under the Q function of IRL reward. This results in a constrained RL problem, which the authors propose an iterative solution for. The algorithm is shown to improve over SAC and several other baselines that shape SAC on continuous control tasks.", "questions": "see above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Strengths:\n1. The idea of using IRL as a signal for RL is interesting.\n2. Figure 1 is very clear and helps understand the method.\n3. The experiments and ablation studies seem rather extensive, although I am not familiar enough with the relevant baselines to say if anything important is missing.\n4. The theoretical analysis of convergence for the iterative solution of the optimization problem in Section 4.2.3 is nice.", "weaknesses": "Post rebuttal: after an interesting discussion with the authors, there is a way to formally establish that Eq. (1) is a sound idea for a kind of \"policy iteration\" step that is guaranteed to improve the initial policy. This makes the paper an interesting contribution.\nThe authors promised to make significant chanegs in the writing that will improve the paper's presentation and suitability for the general RL community.\nBased on this, I recommend to accept the paper and changed my score.\n\nWeaknesses:\n1. In the problem formulation, the authors formulate the constrained optimization problem in Eq. (1) as the problem that the agent needs tosolve. My biggest concern is that while the authors spend considerable efforts on showing that their solution of Eq. (1) is sound, the derivations leading to Eq. (1) seem very ad hoc to me, and I cannot understand why this method would work in general RL problems (or, for what problems should it work?)\n2. The writing and presentation could be improved. For example, the paper talks a lot about explainability, but I did not find explainability to be relevant here so much. At some point, the “explainable RL” idea is replaced by standard IRL. I think the authors should tone down the XRL parts and focus instead on using IRL as a signal in RL.\n\nDetails:\n\n1. The derivations leading to Eq. (1), the main objective of the method, are not clear to me. I collect here several unclear passages.\n    1. Line 185: “At a low level, the RL agent is not optimal meaning that it visits some critical points that lead to the non-optimality.” - I don’t understand this. It could be that the agent is slightly suboptimal *on all* states. This statement seems very ad hoc. Why should it be suboptimal only on some critical points?\n    2. Line 199: \\pi* and \\pi_A are distributions. What does \\pi_A - \\pi* mean exactly? Why is it a metric?\n    3. Line 209: isn’t this a well known result? I.e., a policy is optimal iff it satisfies the Bellman equation.\n    4. Line 215: I’m not sure why Top K “misleading” states make sense. Because changing an action in a state affects the future states the agent will visit, I don’t see how ordering by misleading level makes any sense. For example, consider an agent that at the first step can walk either right or left. If it chooses right, it reaches the goal after some time and obtains reward. If it chooses left, after some time it can either choose to reach the goal, or die and get a huge negative reward. Consider a policy that both takes left in the beginning, and then chooses to die and get the negative reward. Because of discounting, the most misleading state here would be the last state with the negative reward. But if the agent chose right in the beginning, it would not have even reached this state at all, and that would be the state that we want to fix.\n    5. Line 240: I don’t understand what the difference r - \\hat{r}  really means. Is there any formal result that ties between  \\hat{r}  that is the max likelihood IRL solution to the true reward in the system? Is there any formal result that makes this ad hoc choice make sense?\n    6. Line 247-257: This whole paragraph seems very ad hoc. Why should we discourage the misleading state-actions? Maybe under a different policy, they are actually necessary for optimal behavior? E.g., consider an agent that *must* cross some “bridge” to get to the goal. Now, because the agent is not optimal, state actions on the bridge are misleading. But if we restrict the agent from visiting them, it will *never* get to the goal! The causal entropy term is also not well motivated - why is exploration crucial here? \n\nSummary (and explanation for score): I think this paper really needs more work in explaining *why* the approach makes sense. A theoretical result that ties the max-likelihood IRL that the authors build on to Eq. (1) and some optimal policy in the MDP, for example, would go a long way in establishing the soundness of this work, and what MDPs it should work for (the examples I mentioned above hint that for some MDPs we should not expect this approach to work well).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)