Paper deep dive
Loop-Free Inverse Reinforcement Learning via Sequential Value Recovery with Q-Score Matching
Yang chen, Yitan Zhang, Michael Witbrock, Shuyue Hu
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Inverse Reinforcement Learning (IRL) aims to recover a reward function that explains expert demonstrations. Existing IRL methods typically rely on a bi-level optimization procedure that alternates between reward learning and policy optimization, leading to substantial computational burden and training instability. In this work, we introduce a different route that eliminates policy optimization entirely by leveraging diffusion policies. Our key insight is that a diffusion policy encodes the action-gradient structure of the optimal soft Q function, enabling reward learning to be cast as a sequence of value recovery problems, thereby allowing us to bypass reward-policy loops inherent in prior IRL methods. Specifically, our method proceeds in three stages: (I) recovering the optimal soft Q function via action-gradient matching and estimating the corresponding soft value function (LogSumExp of Q values) in a way inspired by Gumbel regression; (II) calibrating these soft values by inferring a state-dependent offset; (III) extracting the reward by enforcing Bellman consistency. This leads to Loop-Free Inverse Reinforcement Learning (LFIRL), a fully offline algorithm that operates in a simple, loop-free, and sequential manner. LFIRL is simple to implement and significantly improves training efficiency while maintaining strong reward recovery performance. Empirically, across Maze, Franka Kitchen, Adroit Hand Pen, and Push-T benchmarks, LFIRL achieves 2-3x speedup over the fastest baselines, while matching or surpassing state-of-the-art methods in reward recovery quality.
Tags
Links
- Source: https://arxiv.org/abs/2609.38955v1
- Canonical: https://arxiv.org/abs/2609.38955v1
Trouble viewing inline? Open PDF directly →
Full Text
83,123 characters extracted from source content.
Expand or collapse full text
Loop-Free Inverse Reinforcement Learning via Sequential Value Recovery with Q-Score Matching Yang Chen †thanks: Equal contribution. Affiliation: Shanghai Artificial Intelligence Laboratory Email: chenyang4@pjlab.org.cn Yitan Zhang11footnotemark: 1 Affiliation: University of Auckland Email: yzhb332@aucklanduni.ac.nz Michael J. Witbrock Affiliation: University of Auckland Email: m.witbrock@auckland.ac.nz Shuyue Hu †thanks: Corresponding author. Affiliation: Shanghai Artificial Intelligence Laboratory Email: hushuyue@pjlab.org.cn Abstract Inverse Reinforcement Learning (IRL) aims to recover a reward function that explains expert demonstrations. Existing IRL methods typically rely on a bi-level optimization procedure that alternates between reward learning and policy optimization, leading to substantial computational burden and training instability. In this work, we introduce a different route that eliminates policy optimization entirely by leveraging diffusion policies. Our key insight is that a diffusion policy encodes the action-gradient structure of the optimal soft Q function, enabling reward learning to be cast as a sequence of value recovery problems, thereby allowing us to bypass reward-policy loops inherent in prior IRL methods. Specifically, our method proceeds in three stages: (I) recovering the optimal soft Q function via action-gradient matching and estimating the corresponding soft value function (LogSumExp of Q values) in a way inspired by Gumbel regression; (I) calibrating these soft values by inferring a state-dependent offset; (I) extracting the reward by enforcing Bellman consistency. This leads to Loop-Free Inverse Reinforcement Learning (LFIRL), a fully offline algorithm that operates in a simple, loop-free, and sequential manner. LFIRL is simple to implement and significantly improves training efficiency while maintaining strong reward recovery performance. Empirically, across Maze, Franka Kitchen, Adroit Hand Pen, and Push-T benchmarks, LFIRL achieves 2-3x speedup over the fastest baselines, while matching or surpassing state-of-the-art methods in reward recovery quality. 11 1 The implementation is available at https://github.com/NothingThrough/LFIRL. 1 Introduction Inverse reinforcement learning (IRL) [27, 1] studies how to recover a reward function from expert demonstrations, with the goal of explaining expert behavior through an underlying optimization principle. It has become an important paradigm for reward design, behavior understanding, and transferable decision making [21, 34]. In practice, however, a major bottleneck in IRL is time efficiency due to the inherent loop structure underlying many classic and modern approaches [1, 51, 9, 11, 47]: the alternating update of the reward function and the policy. This loop arises because most IRL methods generally cannot extract reward-supervision signals directly from demonstrations. Instead, the reward is usually learned through differentiating expert behavior from the learner behavior. This mutual dependence naturally leads to a classic bi-level optimization process, where a reward function is updated by contrasting expert demonstrations with learner behavior, the learner policy is then updated under the updated reward, and this process is repeated until convergence. Such a looped training paradigm is computationally expensive, especially in high-dimensional tasks, and also makes the overall training process harder to stabilize. This raises our central question: can we design an IRL method without the reward-policy loop while still recovering a useful reward function? Figure 1: Traditional inverse reinforcement learning alternates between reward learning and policy optimization, while our loop-free version trains each component once without a reward-policy loop. If we could extract reward-supervision signals directly from demonstrations, then reward recovery would no longer need to be mediated by a reward-policy loop. Intuitively, demonstrations already contain local information about which nearby actions are more or less preferable around expert behavior. Diffusion policy methods [5, 30, 33] provide a natural way to represent this information. Built on the recently burgeoning diffusion models [17, 36], a diffusion policy represents expert action distribution by transforming noise into actions through an iterative refinement process: it learns the gradient of the action distribution’s score function, ∇logπ∗(|) _a π^*(a|s), that describes the direction in which actions should move towards expert ones, and iteratively optimizes wrt this gradient field via a series of stochastic steps. This score-based signal is exactly what we need to break the reward-policy loop. Specifically, in the commonly used maximum-entropy reinforcement learning [14] setting, this score function encodes information of the gradient of the expert’s optimal soft Q function wrt actions, i.e., ∇logπ∗(|)=∇Q∗(,) _a π^*(a|s)= _aQ^*(a,s). Therefore, by matching its action gradient, we can estimate an optimal soft Q function directly from demonstrations, without relying on the differential information from an iteratively updated learner policy. Based on this observation, we propose a new IRL algorithm that avoids the conventional reward-policy loops and instead recovers the reward in a fully sequential manner (see Fig. 1). Concretely, it proceeds in three sequential stages: (I) This stage first recovers an optimal soft Q function by matching its action gradients to the score induced by a diffusion policy, where we use a Sobolev-style training mechanism [6] to favor expert actions over nearby perturbations; then estimates the corresponding soft value function (V == LogSumExp Q) in a way inspired by Gumbel regression [13] (Sec. 4.1). (I) For deriving calibrated value functions, the second stage infers a state-dependent offset function, appending which to the already estimated Q and V yields their more accurate estimates (Sec. 4.2). (I) The final stage recovers the reward function from the calibrated Q and V functions by enforcing Bellman consistency (Sec. 4.3). The entire pipeline is fully offline, using only a static set of trajectories. We empirically demonstrate that our resulting algorithm matches or surpasses state-of-the-art IRL baselines in reward recovery while significantly reducing training time (Sec. 5). Our contributions are summarized as follows: 1. We introduce diffusion policy to IRL and propose the Loop-Free Inverse Reinforcement Learning (LFIRL) algorithm. It is a fully offline IRL algorithm that avoids the computationally expensive looped reward-policy updates in conventional IRL methods by leveraging the diffusion policy to realize a sequential value recovery procedure. 2. We conduct empirical studies on tasks including Maze, Franka Kitchen, Adroit Hand Pen, and Push-T, showing that our method matches or surpasses strong online and offline baselines in reward recovery while improving training time efficiency by 2-3x in most settings. 2 Related Work Inverse reinforcement learning. IRL has been studied since the early formulations of reward recovery and apprenticeship learning [1, 27, 31, 51, 50], and was later extended to deep-learning variants such as Guided Cost Learning [9]. More recent adversarial formulations remain a dominant route for IRL and imitation learning, but they typically rely on a nested loop between reward/discriminator learning and policy optimization, including GAIL [16], AIRL [11], DAC [22], and observation-only adversarial variants such as GAIfO [39]. A broader unifying view is provided by the moment-matching perspective of imitation learning [37], while empirical studies such as [29] further highlight the sensitivity of adversarial pipelines to design choices. Several recent methods reduce or modify this classical reward-policy loop: Maximum-Likelihood IRL [47] proposes a single-loop update, ValueDICE [23] avoids a separate RL optimization procedure through off-policy distribution matching, IQ-Learn [12] and LS-IQ [2] learn implicit Q-function formulations, PIRO [4] stabilizes non-adversarial reward-policy learning through trust-region reward optimization. Other approaches modify the imitation or policy-search process in different ways, including Coherent Soft Imitation Learning [41], successor-feature matching [20], FILTER [38], and Hybrid IRL [32]. In contrast to both nested-loop methods and single-loop approaches that tightly entangle reward and policy learning, our method adopts a fully staged, sequential IRL pipeline: it uses diffusion-derived local action-score information to recover soft values, and then recovers rewards. Diffusion-Based Imitation Learning. Diffusion models have rapidly become strong policy classes for imitation learning, especially in robotics and multimodal control. Diffusion Policy [5] models actions through a conditional denoising process and has become a standard baseline for visuomotor behavior cloning, while subsequent work extends diffusion-based imitation learning through goal-conditioned policies [33], diffusion-augmented behavioral cloning [3], improved visual or self-supervised representations [25, 46], and architectural scaling studies [7, 49]. A smaller line of work uses diffusion models more directly for reward learning, for example by extracting reward-like quantities from diffusion models [28], learning rewards from expert videos [19], leveraging score information for policy optimization [30], or integrating diffusion into adversarial imitation pipelines as the discriminator such as DiffAIL [40], DRAIL [24], and DIFO [18]. Recent score-matching-based imitation methods such as [42] likewise emphasize the value of diffusion-derived gradient information. Our method differs from these directions in that we use a diffusion policy only as a score estimator for local supervision of soft Q-recovery, and then recover V and r sequentially, rather than using diffusion as the final policy class, extracting relative trajectory-level rewards from multiple diffusion models, or embedding diffusion inside an adversarial reward-policy loop. 3 Preliminaries A Markov decision process (MDP) is defined by the tuple (,⊥,,P,γ,r)(S,S_ ,A,P,γ,r), where S and A denote the state and action spaces, ⊥⊆S_ is a set of absorbing states, P(⋅|,)P(·|s,a) is the transition kernel, γ∈(0,1)γ∈(0,1) is the discount factor, and r:×→ℝr:S×A is the reward function. For any absorbing state ⊥∈⊥s_ _ , we have P(⊥|⊥,)=1P(s_ |s_ ,a)=1 and r(⊥,)=0r(s_ ,a)=0 for all ∈a , which describes can-never-escape states, e.g., completions or failures of a game. Let π(|)π(a|s) be a stochastic policy, and let ρπ(,)≔11−γ∑t=0∞Pr(t=,t=|π,P)ρ^π(s,a) 11-γ _t=0^∞ (s_t=s,a_t=a|π,P) denote its induced state-action density. 3.1 Maximum-Entropy RL and IRL The goal of classic RL is to find a policy to maximize the expected discounted long-term rewards (,)∼ρπ[r(,)]E_(s,a) ρ^π[r(s,a)]. We consider a generalized version of Maximum-entropy (MaxEnt) RL that augments the reward objective with the relative entropy ℋμ(π)≔(,)∼ρπ[−logπ(|)μ(|)]H_μ(π) _(s,a) ρ^π[- π(a|s)μ(a|s)] between π and a reference policy μ: (,)∼ρπ[r(,)]+εℋμ(π)E_(s,a) ρ^π[r(s,a)]+ \,H_μ(π), where ε>0 >0 is the temperature parameter. It recovers the standard MaxEnt RL [14] objective up to a constant when μ is uniform. Under this formulation of MaxEnt RL, the optimal policy admits the energy-based form π∗(|)=μ(|)exp(1ε(Q∗(,)−V∗())),π^*(a|s)=μ(a|s) ( 1 (Q^*(s,a)-V^*(s) ) ), (1) where the optimal soft value function V∗()V^*(s) and optimal soft action value function Q∗(,)Q^*(s,a) satisfy V∗()=εlog∫exp(1εQ∗(,))μ(|),Q∗(,)=r(,)+γ′∼P[V∗(′)].V^*(s)= _A ( 1 Q^*(s,a) )dμ(a|s), Q^*(s,a)=r(s,a)+ _s P [V^*(s ) ]. (2) Note that the MaxEnt optimality condition in Eq. (1) implies the following relation, π∗(|)/μ(|)π^*(a|s)/μ(a|s) and Q∗(|)Q^*(a|s) have the same gradient wrt actions, which we will leverage to build our method upon: ∇logπ∗(|)μ(|)=1ε∇Q∗(,). _a π^*(a|s)μ(a|s)= 1 _aQ^*(s,a). (3) In MaxEnt inverse RL (MaxEnt IRL), the reward function is unknown, while a set of expert demonstrations E=τii=1ND_E=\ _i\_i=1^N is observed, where each trajectory τ=(0,0,1,1,…)τ=(s_0,a_0,s_1,a_1,…) is generated by an expert policy πE _E. The goal is to recover a reward function under which πE _E is optimal under the above MaxEnt RL framework, which can be interpreted by the following optimization problem: minrmaxπ(,)∼ρπ[r(,)]+εℋμ(π)−(,)∼ρπE[r(,)]. _r _πE_(s,a) ρ^π[r(s,a)]+ \,H_μ(π)-E_(s,a) ρ _E[r(s,a)]. (4) The min-max structure of Eq. (4) reveals an inherent bi-level optimization procedure of many classic and modern IRL methods [16, 8, 11, 48]: the upper level updates the reward by contrasting expert behavior with the learner behavior induced by the current reward, while the lower level updates the learner’s imitation policy. Classic IRL methods typically implement this bi-level procedure as a computationally expensive nested reward-policy loop because the lower level is a full RL procedure for computing an optimal policy. Some recent methods [47, 48] improve time efficiency by introducing a relatively cheap single reward-policy loop, where the lower level conducts one or several steps of policy improvement. Although this bi-level optimization structure is well theoretically grounded [16, 47] and has been shown to be effective in applications [43, 10], its induced reward-policy loops introduce inevitable computational burden and instability to the optimization process. 3.2 Diffusion Models and Diffusion Policy Diffusion models [17, 36] generate samples by modeling a forward process that gradually adds noise to data and a reverse process that iteratively removes this noise to recover data samples. Diffusion policy methods use this idea for imitating expert behavior, where a policy π(⋅|t)π(·|s_t) is modeled as an action generation process: given an MDP time step t, let tkk=0K\a_t^k\_k=0^K denote the latent action sequence indexed by the diffusion step k, where t0a_t^0 corresponds to the initial action (samples from expert demonstrations) and larger k corresponds to noisier ones. The forward process uses a variance schedule βkk=1K\ _k\_k=1^K with αk≔1−βk _k 1- _k, α¯k≔∏j=1kαj α_k _j=1^k _j, and Gaussian noise ϵ∼(0,I) ε (0,I); the noisy action at step k can be written in closed form as tk=α¯kt0+1−α¯kϵ.a_t^k= α_k\,a_t^0+ 1- α_k\, ε. (5) The reverse process trains a noise predictor ϵϕ ε_ φ by minimizing the following MSE loss: ℒdiff(ϕ)≔(t,t0)∼ρπE,k∼Unif1,…,K,ϵ∼(0,I)‖ϵ−ϵϕ(tk,t,k)‖22.L_diff( φ) _ subarrayc(s_t,a_t^0) ρ _E,\\ k \1,…,K\,\, ε (0,I) subarray \| ε- ε_ φ (a_t^k,s_t,k ) \|_2^2. (6) Training the noise predictor ϵϕ ε_φ can be interpreted as learning ∇logπE(|t) _a _E(a|s_t), the so-called score of a diffusion policy, which is a vector field indicating the direction to adjust a noisy action to increase its likelihood under the expert action distribution. In the idealized limit of exact denoising, it can be shown that ∇logπE(|t)|=tk=−11−α¯kϵϕ(tk,t,k)≕gϕ(tk,t,k) _a _E(a|s_t)|_a=a_t^k=- 1 1- α_k ε_ φ(a_t^k,s_t,k) g_ φ(a_t^k,s_t,k). Combining this with Eq. (3), under the condition that the expert policy is soft-optimal, i.e., πE=π∗ _E=π^*, we have gϕ(tk,t,k)=∇logπE(|t)=∇logμ(|t)|=tk+1ε∇Qr∗(t,)|=tk.g_φ(a_t^k,s_t,k)= _a _E(a|s_t)= _a μ(a|s_t)|_a=a_t^k+ 1 _aQ_r^*(s_t,a)|_a=a_t^k. (7) Thus, a pretrained diffusion policy can provide the action-gradient supervision signal to recover the optimal soft Q function. We will utilize this property to design our method. To stay self-contained, detailed derivations of the identity relationship in Eq. (7) are given in Appendix A.1. 4 Methods In this section, we present Loop-Free Inverse Reinforcement Learning (LFIRL), a fully offline IRL algorithm that avoids reward-policy loops in conventional IRL approaches. Instead, it proceeds in three sequential stages (see Fig. 1): (I) first recover a soft Q function from diffusion-policy using Sobolev training and recover the value function V from the learned Q in a way inspired by Gumbel regression; (I) then infer the state-dependent offset and calibrate the Q and V; (I) finally recover the reward from the learned Q and V by enforcing Bellman consistency. Figure 2: Illustration of matching policy’s score on Push-T task, where the goal is to push the T-shaped block to the target pose. The red circle marks the agent’s current position (initially in the Start panel), and the dark dashed curves indicate two optimal routes to reach the correct position for pushing the block. The arrows depict the action-gradient field supervised by the diffusion policy’s score; the blue dots are perturbed actions whose Q values are constrained below that of the expert action by a margin, thereby providing anchoring information. 4.1 Stage I: Recover the Soft Q Function and Value function In this stage, we recover the optimal soft Q function through Sobolev Training [6, 35], which trains neural networks by jointly matching function values and their derivatives. In our formulation, the score function of the diffusion policy provides direct supervision on action-gradients. Specifically, we train a Q-network Q^(,) Q_ ω(s,a) such that its action-gradient matches the diffusion policy’s score after correcting for the reference policy. Concretely, we align ∇Q^(,)|=k _a Q_ ω(s,a)|_a=a^k with gϕ(k,,k)−∇logμ(|)|=kg_ φ(a^k,s,k)- _a μ(a|s)|_a=a^k, where ka^k is obtained by adding noise to an expert action 0a^0 via Eq. (5). However, the derivation of value supervision is not straightforward, as the optimal Q values themselves are indeed what we seek for. To address this, we introduce an implicit form of value supervision by enforcing a small margin between the Q values of expert actions and those of their locally perturbed counterparts. We do not impose this margin globally, in order to preserve action diversity, as multiple actions may be optimal under the same state. Intuitively, this margin-based constraint effectively regularizes the Q function, using expert actions as “anchor points” to shape a more precise function surface. These settings, illustrated with Push-T task in Fig. 2, lead to the following loss function for training Q^(,) Q_ ω(s,a): ℒQ^():= _ Q( ω):= ∈,∈,k∼Unif1,…,K[‖gϕ(k,,k)−∇logμ(|t)|=tk−1ε∇Q^(,)|=k‖22]⏟Gradient Matching _ subarraycs ,a ,\\ k \1,…,K\ subarray [ \|g_φ(a^k,s,k)- _a μ(a|s_t)|_a=a_t^k- 1 _a Q_ ω(s,a)|_a=a^k \|_2^2 ]_Gradient Matching (8) +λ(,)∼E,~∼η(⋅∣)[ReLU(ξ−(Q^(,)−Q^(,~)))]⏟Value Anchoring, +λ\,E_ subarrayc(s,a) _E,\\ a η(· ) subarray [ReLU (ξ- ( Q_ ω(s,a)- Q_ ω(s, a) ) ) ]_Value Anchoring, where ξ>0ξ>0 is an intended margin, η(⋅|)η(·|a) denotes the neighborhood of action a, ReLU(x)≔max0,xReLU(x) \0,x\ denotes the Rectified Linear Unit function, and λ≥0λ≥ 0 is a weighting coefficient. In practice, we can instantiate the reference policy μ(|)μ(a|s) as a simple distribution (e.g., uniform over bounded actions or a Gaussian prior), whose score ∇logμ(|) _a μ(a|s) is tractable. After obtaining Q Q_ ω, we estimate the corresponding soft value function. However, direct computation of V^()=εlog∫exp(Q^(,)/ε)μ(|) V(s)= _A ( Q_ ω(s,a)/ )\,dμ(a|s) is generally intractable in continuous action spaces. To address this, we reformulate the problem via an equivalent moment condition, drawing inspirations from Gumbel regression in RL [13], which models action selection by perturbing Q values with Gumbel distribution noise [15] and leverages Gumbel-Max Trick [15] to infer the value function as the expected stochastic maximum, yielding a LogSumExp form. Specifically, the soft value V^() V(s) is the unique solution to ∼μ(⋅|)[exp(Q^(,)−V()ε)]=1.E_a μ(·|s)[ ( Q_ ω(s,a)-V(s) )]=1. This characterization allows us to estimate V^() V(s) without explicitly computing the LogSumExp. We then construct a convex surrogate whose first-order condition recovers the above moment constraint: ℓV(v;)≔∼μ(⋅|)[exp(Q^(,)−vε)−Q^(,)−vε−1]. _V(v;s) _a μ(·|s) [ ( Q_ ω(s,a)-v )- Q_ ω(s,a)-v -1 ]. (9) This objective is strictly convex in v and admits a unique minimizer, which coincides with the soft value V^() V(s). The derivation is provided in Appendix A.2. In practice, we freeze Q Q_ ω and train a value network V V_ χ on a set of trajectories SD_S sampled from μ through the following empirical objective ℒV^()≔(,)∼DS[exp(z)−z−1],z=Q^(,)−V^()ε.L_ V( χ) _(s,a) D_S [ (z)-z-1 ], z= Q_ ω(s,a)- V_ χ(s) . (10) Notably, although this objective coincides with the Gumbel regression loss [13], our derivation does not rely on a stochastic model for Q values, as Q Q_ ω is fixed during training V V_ χ; instead, it follows directly from the deterministic characterization of the LogSumExp under the reference policy μ. 4.2 Stage I: Infer the State-Dependent Offset and Calibrate Soft Values The soft values recovered in Stage I are uncalibrated, as the action-gradient matching determines only how soft values change wrt actions, but does not reflect their state-dependent factors. This fact can be formally explained by the following proposition, whose justification is given in Appendix A.3. Proposition 1. Assume πE=π∗ _E=π^* and that the learned diffusion policy’s score is exact in the sense of Eq. (7). For each state s, if the action-gradient matching term in Eq. (8) is minimized exactly, then there exists a state-dependent function b∗:→ℝb^*:S such that Q∗(,)=Q^(,)+b∗()Q^*(s,a)= Q_ ω(s,a)+b^*(s). This reflects the ambiguity caused by shift invariance: adding any state-dependent offset b()b(s) leaves action-gradient unchanged, i.e., ∇Q^(,)=∇(Q^(,)+b()) _a Q_ ω(s,a)= _a ( Q_ ω(s,a)+b(s) ), making Q Q_ ω identifiable only up to b()b(s). We therefore introduce an offset network b:→ℝb_ ψ:S and define the calibrated soft Q function as Q,(,)≔Q^(,)+b()Q_ ω, ψ(s,a) Q_ ω(s,a)+b_ ψ(s) and write the calibrated soft value function from Stage I as V^()+b() V_ χ(s)+b_ ψ(s). A key challenge is that the offset b_ ψ is not identifiable from action-gradient information. While penalizing b2b_ ψ^2 can fix the shift ambiguity by selecting the minimum-norm solution among all equivalent offsets, it ignores the transition structure and may yield solutions that violate Bellman consistency. Therefore, we regularize the Bellman-implied reward, δ,,(t,t,t+1)≔Q^(t,t)+b(t)−γ(V^(t+1)+b(t+1)), _ ω, χ, ψ(s_t,a_t,s_t+1) Q_ ω(s_t,a_t)+b_ ψ(s_t)-γ ( V_ χ(s_t+1)+b_ ψ(s_t+1) ), (11) which couples b_ ψ across consecutive states and enforces consistency with the underlying dynamics. Intuitively, minimizing the squared norm of δ,, _ ω, χ, ψ penalizes inconsistent offset value shifts across transitions, encouraging the offset to vary coherently along trajectories. In addition, absorbing states provide direct supervision: for any ⊥∈⊥s_ _ , we have Q∗(⊥,)=0Q^*(s_ ,a)=0 for all a. 22 2 This zero value is caused by the absorbing state itself, not by a particular action, so it is used to calibrate the state-dependent offset rather than to provide the action-level value supervision used in Stage I. Thus, the resulting objective is ℒb()≔(t,t,t+1)∼E∪S[δ,,2+λbb2(t)⏟Offset regularization]+ _b( ψ) _ subarrayc(s_t,a_t,s_t+1)\\ _E _S subarray[ δ^2_ ω, χ, ψ+ _b\,b_ ψ^2(s_t)_Offset regularization]+ ⊥∈⊥,[(Q^(⊥,)+b(⊥))2⏟Absorbing states supervision], _s_ _ ,a[ ( Q_ ω(s_ ,a)+b_ ψ(s_ ) )^2_Absorbing states supervision], (12) where λb≥0 _b≥ 0 is a regularization coefficient, Q Q_ ω and V V_ χ are fixed, and only b_ ψ is updated. After calibrating the soft Q function, we re-estimate the value function using the calibrated Q,Q_ ω, ψ. Specifically, we freeze Q,Q_ ω, ψ and optimize the Gumbel regression objective again: ℒV(′)≔,∼μ(⋅|)[exp(z¯)−z¯−1],z¯≔Q,(t,t)−V′(t)ε.L_V( χ ) _s,a μ(·|s) [ ( z)- z-1 ], z Q_ ω, ψ(s_t,a_t)-V_ χ (s_t) . (13) This second value-fitting step is needed because the soft value function must be consistent with the offset-calibrated soft Q function. 4.3 Stage I: Recover the Reward Function Given the calibrated soft Q function Q,Q_ ω, ψ and value function V′V_ χ , reward recovery reduces to enforcing Bellman consistency. For a transition (t,t,t+1)(s_t,a_t,s_t+1), we write the Bellman-implied reward target r=Q,(t,t)−γV′(t+1).r=Q_ ω, ψ(s_t,a_t)-γ V_ χ (s_t+1). We parameterize the reward by a network r_ θ and learn it by minimizing Bellman residual errors. To stabilize training, we apply clipping and normalization: r~=clip(r,−cr,cr)−mrσr+ζr,r~=r(t,t)−mrσr+ζr, r= clip(r,-c_r,c_r)-m_r _r+ _r, r_ θ= r_ θ(s_t,a_t)-m_r _r+ _r, (14) where cr>0c_r>0 is the clipping threshold, (mr,σr)(m_r, _r) are the running mean and standard deviation of the clipped reward targets, and ζr>0 _r>0 is a small numerical constant. The reward network is trained via ℒr()=(t,t,t+1)∼E∪S[(r~−r~)2].L_r( θ)=E_(s_t,a_t,s_t+1) _E _S [( r_ θ- r)^2 ]. (15) Algorithm 1 Loop-Free Inverse Reinforcement Learning (LFIRL) 1: Input: Expert demonstrations ED_E; initialized networks Q Q_ ω, b_ ψ, V_ χ, and r_ θ; reference policy μ. 2: Train a (or load a pre-trained) diffusion policy on ED_E. Stage I 3: Recover Q^(,) Q_ ω(s,a) on ED_E via matching diffusion policy’s score. ⊳ Eq. (8) 4: Recover V()V_ χ(s) on SD_S using Gumbel-regression-style value fitting. ⊳ Eq. (10) 5: Recover state-dependent offset b()b_ ψ(s) on ED_E and SD_S. Stage I ⊳ Eq. (12) 6: Calibrate V^′() V_ χ (s) by the calibrated Q values Q^+b() Q_ ω+b_ ψ(s). ⊳ Eq. (13) 7: Recover r_ θ on ED_E and SD_S by minimizing Bellman residual errors. Stage I ⊳ Eq. (15) 8: Output: recovered reward function r_ θ. 4.4 Algorithm Summary The complete procedure of the proposed Loop-Free IRL framework is summarized in Algorithm 1. In contrast to traditional IRL pipelines, our method eliminates the reward-policy loop entirely: each stage is executed only once in a sequential manner, and every component is trained solely on the frozen outputs produced by preceding stages. Importantly, the algorithm is fully offline, as SD_S is an offline dataset collected from the reference policy μ. Consequently, training requires neither environment interactions nor rollouts from a learned dynamics model during training. Overall, this loop-free offline design substantially improves training efficiency, as will be demonstrated experimentally. 5 Experiments We evaluate LFIRL from two perspectives: (1) reward recovery quality and (2) training time efficiency. 5.1 Experimental Setup We consider four different kinds of tasks (see Fig. 3): PointMaze with the UMaze, Medium, and Large settings; Franka Kitchen consists of four subtasks; Adroit Hand Pen; and the Push-T manipulation [5]. We compare against four classes of baselines: (1) Non-Adversarial IRL baselines: IQ-Learn [12] and ML-IRL [47]; (2) the Adversarial IRL baseline: AIRL [11]; (3) the Offline IRL baseline: Offline ML-IRL [48], CLARE [45], and ValueDICE [23]; and (4) Diffusion-based GAIL baselines: DRAIL [24] and DIFO [18]. All reported results are averaged over five random runs. All methods are trained with the same environment-step budget, and within each environment we use the same number of expert demonstrations for all methods. Full implementation details, including hyperparameters and network architectures, are in Appendix E.1. To evaluate reward recovery on PointMaze and Franka Kitchen, we train an SAC [14] policy on the recovered reward, and then evaluate this policy in the original environment. We report success rate on PointMaze and the mean number of completed subtasks on Franka Kitchen. Since our method is designed to recover a reward function rather than directly output a policy, this setting provides a task-level assessment of whether the recovered reward induces the intended behavior. For the higher-dimensional Adroit Hand Pen and Push-T environments, we instead evaluate reward quality through discrimination: we construct balanced sets of high-quality and low-quality trajectories, score them with each learned reward model, and report the resulting classification accuracy. We additionally compare the recovered and ground-truth rewards on the same state-action samples using Pearson’s and Spearman’s correlation coefficients; the results are presented in Sec. 5.4. To evaluate efficiency, we record the training time required by each method to consume the same number of trajectory throughput, with each method run independently and without parallelized competition for compute resources. In addition to the main comparison, we also report an offset-ablation study and a demonstration-reduction study; the corresponding results are provided in Sec. 5.5 and Sec. 5.6. 5.2 Reward Recovery The main results are shown in Fig. 4. Overall, LFIRL recovers highly effective rewards and exhibits strong cross-task robustness. Across various tasks, the recovered reward consistently supports competitive downstream control and achieves the best or near-best final performance in most settings. Even in cases where it does not attain the single best score, its performance remains highly competitive, indicating that the proposed sequential recovery pipeline yields stable and transferable reward signals. UMaze Maze(M) Maze(L) Kitchen Pen Push-T Figure 3: Illustrations of the evaluation tasks used in our experiments. Figure 4: Reward-recovery results on PointMaze and Franka Kitchen. We report success rate on PointMaze and average completed subtasks on Franka Kitchen. Since LFIRL recovers the reward only at the final stage, its result is shown as a fixed final score rather than a policy-learning curve. For online IRL methods, “steps” denotes environment interaction steps. For offline IRL methods, “steps” denotes the number of transitions sampled from the offline dataset during training. For the higher-dimensional Push-T and Adroit Hand Pen, we instead evaluate whether the recovered reward can distinguish high-quality trajectories from low-quality ones. The results are reported in Tab. 1. These results further support a consistent overall picture: the reward recovered by LFIRL captures meaningful behavioral differences and generalizes well across qualitatively different evaluation protocols, including both downstream control and trajectory discrimination. Table 1: Trajectory-quality classification accuracy (%) on Push-T and Adroit Hand Pen using reward networks recovered by different IRL algorithms. Results are reported as mean ± standard deviation over five runs. Task LFIRL AIRL IQ-Learn DIFO DRAIL ML-IRL Offline ML-IRL CLARE ValueDICE Push-T 74.17± 1.96 59.67± 2.08 39.67± 0.58 79.00± 1.00 49.67± 1.15 71.33± 2.33 42.67± 1.20 71.00± 3.67 78.67± 2.52 Pen 97.17± 2.75 90.00± 2.45 74.33± 7.59 91.67± 2.05 98.33± 0.94 71.00± 4.32 66.00± 12.43 94.67± 0.58 95.67± 1.15 Table 2: Average running time (in hours). “Load D.P.” denotes LFIRL with a pre-trained diffusion policy, while “Train D.P.” denotes LFIRL including diffusion-policy training time. “T.R.” is short for “time reduction” and reports the relative time saved compared with the fastest baseline (marked with an underline). Task Online methods Offline methods LFIRL AIRL IQ-Learn DIFO DRAIL ML-IRL ML-IRL CLARE ValueDICE Load D.P. T.R. (SpeedUp) Train D.P. T.R. (SpeedUp) UMaze 1.38 6.88 2.47 1.70 9.76 3.09 5.86 4.66 0.28 79.7% (5x) 0.43 68.8% (3x) Maze(M) 3.67 10.10 7.05 2.65 13.80 4.99 5.98 4.87 0.77 70.9% (3x) 1.07 59.6% (2x) Maze(L) 6.88 10.88 7.73 3.05 13.33 5.52 6.67 5.17 1.42 53.4% (2x) 1.80 41.0% (2x) Kitchen 9.55 13.72 7.40 3.33 18.42 6.83 7.80 5.91 1.52 54.4% (2x) 1.87 43.8% (2x) Push-T 1.57 8.55 4.05 1.95 9.50 2.10 3.20 5.30 0.35 77.7% (4x) 0.62 60.5% (3x) Pen 4.83 14.88 10.83 7.50 18.52 6.53 3.79 4.71 1.18 68.9% (3x) 1.71 54.9% (2x) 5.3 Average Running Time Tab. 2 reports the average running time under the same step budget. Overall, LFIRL substantially reduces the time cost of training. Across tasks, it is consistently the fastest method, typically achieving about a 2–3x reduction in runtime compared with the fastest baseline in each setting, and a roughly 2–5x reduction when loading a pre-trained diffusion policy. These results confirm that removing the reward-policy optimization loop brings not only pipeline simplicity, but also a substantial practical gain in time efficiency. 5.4 Correlation with Ground-Truth Rewards The preceding evaluations assess whether a recovered reward supports the intended behavior or distinguishes trajectories of different quality. We further evaluate reward recovery directly by comparing each learned reward with the environment’s ground-truth reward on the same set of state-action samples. We report Pearson’s correlation coefficient (PCC), which measures linear association between the recovered and ground-truth reward values, and Spearman’s rank correlation coefficient (SCC), which measures agreement between their rankings. A higher PCC indicates that the recovered reward more closely follows the ground-truth reward’s linear variation across samples; a higher SCC indicates that it better preserves which samples receive higher or lower reward. These measures assess agreement in structure and ordering, respectively, without requiring the recovered reward to have exactly the same scale or offset as the ground-truth reward. Table 3: Direct comparison of recovered and ground-truth rewards on the same state-action samples. Each entry reports PCC / SCC; higher values indicate stronger linear association / rank agreement. Bold indicates the highest value for each metric in each task. Method UMaze Medium Large Kitchen Push-T Pen Average PCC / SCC PCC / SCC PCC / SCC PCC / SCC PCC / SCC PCC / SCC PCC / SCC LFIRL (Ours) 0.83 / 0.91 0.76 / 0.82 0.84 / 0.87 0.67 / 0.81 0.72 / 0.81 0.79 / 0.85 0.77 / 0.85 AIRL 0.65 / 0.86 0.63 / 0.74 0.77 / 0.75 0.44 / 0.63 0.69 / 0.74 0.75 / 0.71 0.66 / 0.74 IQ-Learn 0.62 / 0.74 0.59 / 0.62 0.65 / 0.70 0.34 / 0.51 0.48 / 0.65 0.47 / 0.62 0.53 / 0.64 DIFO 0.69 / 0.79 0.53 / 0.67 0.68 / 0.81 0.61 / 0.77 0.83 / 0.79 0.71 / 0.72 0.68 / 0.76 DRAIL 0.67 / 0.71 0.74 / 0.84 0.69 / 0.78 0.56 / 0.76 0.51 / 0.56 0.63 / 0.62 0.63 / 0.71 ML-IRL 0.78 / 0.75 0.67 / 0.72 0.64 / 0.71 0.63 / 0.69 0.66 / 0.73 0.78 / 0.69 0.69 / 0.72 Offline ML-IRL 0.49 / 0.52 0.47 / 0.58 0.42 / 0.62 0.45 / 0.54 0.45 / 0.53 0.53 / 0.43 0.47 / 0.54 CLARE 0.72 / 0.80 0.68 / 0.74 0.49 / 0.61 0.24 / 0.46 0.51 / 0.79 0.58 / 0.55 0.54 / 0.66 ValueDICE 0.72 / 0.81 0.61 / 0.79 0.74 / 0.71 0.53 / 0.56 0.65 / 0.76 0.57 / 0.63 0.64 / 0.71 As shown in Tab. 3, LFIRL obtains the highest PCC on five of the six tasks and the highest SCC on five of the six tasks. Thus, LFIRL’s consistently strong results indicate that the sequential recovery procedure preserves both the linear variation and relative ordering of the ground-truth reward across a range of environments. 5.5 Ablation Study Figure 5: Ablation results on PointMaze, Kitchen, Pen, and Push-T. We compare LFIRL with its variant that removes the state-dependent offset b()b(s). We conduct an ablation study to evaluate the contribution of the state-dependent offset b()b(s) in LFIRL. The results are shown in Fig. 5. Overall, removing b()b(s) consistently degrades performance across the evaluated tasks, showing that inferring the offset and calibrating soft values are important for recovering a well-aligned value structure. This result supports our analysis in Sec. 4.2: calibrating this offset is necessary for obtaining a reliable downstream reward. 5.6 Reduced-Data Experiments We further study how LFIRL behaves when the number of expert demonstrations is reduced. The results are shown in Fig. 6, where we use the full dataset, one-half of the demonstrations, and one-quarter of the demonstrations. Overall, the performance of LFIRL decreases as the number of demonstrations becomes smaller. This trend is consistent with the design of our method: the first stage recovers the Q function from local supervision around expert actions, and the number of available expert state-action anchors directly affects how accurately this action-value structure can be reconstructed. When the demonstration set is reduced, anchor coverage becomes sparser, which weakens the quality of Q recovery and subsequently affects downstream value and reward recovery. Nevertheless, LFIRL remains competitive across the reduced-data settings. Figure 6: Reduced-data experiments with the full, one-half, and one-quarter of the demonstrations. We report success rate on PointMaze, average completed subtasks on Franka Kitchen, and reward on Push-T and Pen. 6 Discussion Source of the time efficiency gain. The efficiency of LFIRL mainly comes from its linear training pipeline, not merely from using offline training framework. Offline IRL methods can still be slow if they keep a reward-policy loop, since reward updates repeatedly depend on policy or value optimization under the current reward, no matter if the policy interacts with the environment. Even methods that reduce this cost, such as IQ-Learn [12], may avoid training an explicit policy in discrete action spaces but still rely on value or policy quantities induced by the current reward-related objective. LFIRL removes this repeated coupling with each stage using the frozen output of previous stages, and this is why LFIRL remains much faster than offline IRL methods in Tab. 2. Choice of generative models. In fact, LFIRL uses a conditional Denoising Diffusion Probabilistic Model [17] (DDPM)-based diffusion policy to provide the derivative signal for Q function recovery, following common practice in diffusion-based imitation learning methods. In principle, any state-conditioned generative model that can provide a reliable score, or an action-gradient signal, can be used as the source of derivative supervision in Stage I. In Appendix B.1, we instantiate a version based on Flow Matching [26], another popular generative model paradigm. However, Appendix B.2 shows that, under the same number of training epochs, this variant gives weaker reward recovery. One possible reason is that the two pretraining objectives may have different convergence behavior under the same epoch budget. Conditional DDPM training directly optimizes denoising over multiple noise levels, which provides a strong and stable supervision signal for action refinement around expert demonstrations. Flow Matching instead learns a continuous transport vector field, whose action-generation quality can be more sensitive to integration accuracy, schedule design, and training budget. As a result, the pretrained diffusion policy is less accurate for giving the score signal used for Q function recovery, which then propagates to the final reward. 7 Conclusion We propose LFIRL, a novel staged IRL framework that removes the conventional reward-policy optimization loop by turning reward recovery into a sequence of value-recovery problems. Concretely, LFIRL first recovers an optimal Q function from diffusion-policy score signals and estimates the soft value function via Gumbel-regression-style value fitting. It then infers a state-dependent offset to calibrate the recovered soft values and re-estimates the value function using the calibrated Q function. Finally, it recovers the reward function by enforcing Bellman consistency. Experiments on PointMaze, Franka Kitchen, Adroit Hand Pen, and Push-T show that LFIRL matches or surpasses strong baselines in reward recovery while substantially reducing training time, achieving 2-3x efficiency gains in most settings. Limitations. The quality of the recovered Q function depends on sufficient local anchors around expert actions. When anchor coverage is sparse, the method may require more demonstrations to reconstruct the action-value structure accurately. Future work may explore broader classes of generative models beyond DDPM and Flow Matching for providing action-gradient supervision, and develop more sample-efficient ways to recover reliable value and reward information from limited expert demonstrations. References [1] P. Abbeel and A. Y. Ng (2004) Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, p. 1. Cited by: §1, §2. [2] F. Al-Hafez, D. Tateo, O. Arenz, G. Zhao, and J. Peters (2023) LS-iq: implicit reward regularization for inverse reinforcement learning. In 11th International Conference on Learning Representations, ICLR 2023, Cited by: §2. [3] S. Chen, H. Wang, M. Hsu, C. Lai, and S. Sun (2024) Diffusion model-augmented behavioral cloning. In Forty-first International Conference on Machine Learning, Cited by: §2. [4] Y. Chen, M. Zou, J. Zhang, Y. Zhang, J. Yang, G. Gendron, L. Zhang, J. Liu, and M. J. Witbrock (2025) Trust region reward optimization and proximal inverse reward optimization algorithm. In Advances in Neural Information Processing Systems, Cited by: §2. [5] C. Chi, Z. Xu, S. Feng, E. Cousineau, Y. Du, B. Burchfiel, R. Tedrake, and S. Song (2025) Diffusion policy: visuomotor policy learning via action diffusion. The International Journal of Robotics Research 44 (10-11), p. 1684–1704. Cited by: §E.2, §1, §2, §5.1. [6] W. M. Czarnecki, S. Osindero, M. Jaderberg, G. Swirszcz, and R. Pascanu (2017) Sobolev training for neural networks. Advances in neural information processing systems 30. Cited by: §1, §4.1. [7] S. Dasari, O. Mees, S. Zhao, M. K. Srirama, and S. Levine (2025) The ingredients for robotic diffusion transformers. In 2025 IEEE International Conference on Robotics and Automation (ICRA), p. 15617–15625. Cited by: §2. [8] C. Finn, P. Christiano, P. Abbeel, and S. Levine (2016) A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models. arXiv preprint arXiv:1611.03852. Cited by: §3.1. [9] C. Finn, S. Levine, and P. Abbeel (2016) Guided cost learning: deep inverse optimal control via policy optimization. In International conference on machine learning, p. 49–58. Cited by: §1, §2. [10] J. Fu, A. Korattikara, S. Levine, and S. Guadarrama (2023) From language to goals: inverse reinforcement learning for vision-based instruction following. In International Conference on Learning Representations, Cited by: §3.1. [11] J. Fu, K. Luo, and S. Levine (2018) Learning robust rewards with adversarial inverse reinforcement learning. In International Conference on Learning Representations, Cited by: §1, §2, §3.1, §5.1. [12] D. Garg, S. Chakraborty, C. Cundy, J. Song, and S. Ermon (2021) Iq-learn: inverse soft-q learning for imitation. Advances in Neural Information Processing Systems 34, p. 4028–4039. Cited by: §2, §5.1, §6. [13] D. Garg, J. Hejna, M. Geist, and S. Ermon (2023) Extreme q-learning: maxent rl without entropy. In The Eleventh International Conference on Learning Representations, Cited by: §1, §4.1, §4.1. [14] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine (2018) Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, p. 1861–1870. Cited by: §1, §3.1, §5.1. [15] T. Hazan and T. Jaakkola (2012) On the partition function and random maximum a-posteriori perturbations. In Proceedings of the 29th International Conference on Machine Learning, p. 1667–1674. Cited by: §4.1. [16] J. Ho and S. Ermon (2016) Generative adversarial imitation learning. Advances in neural information processing systems 29. Cited by: §2, §3.1. [17] J. Ho, A. Jain, and P. Abbeel (2020) Denoising diffusion probabilistic models. Advances in neural information processing systems 33, p. 6840–6851. Cited by: §1, §3.2, §6. [18] B. Huang, C. Yang, C. Lai, D. Wu, and S. Sun (2024) Diffusion imitation from observation. Advances in Neural Information Processing Systems 37, p. 137190–137217. Cited by: §2, §5.1. [19] T. Huang, G. Jiang, Y. Ze, and H. Xu (2024) Diffusion reward: learning rewards via conditional video diffusion. European Conference on Computer Vision (ECCV). Cited by: §2. [20] A. K. Jain, H. Wiltzer, J. Farebrother, I. Rish, G. Berseth, and S. Choudhury (2025) Non-adversarial inverse reinforcement learning via successor feature matching. In The Thirteenth International Conference on Learning Representations, Cited by: §2. [21] M. Kalakrishnan, P. Pastor, L. Righetti, and S. Schaal (2013) Learning objective functions for manipulation. In 2013 IEEE International Conference on Robotics and Automation, p. 1331–1336. Cited by: §1. [22] I. Kostrikov, K. K. Agrawal, D. Dwibedi, S. Levine, and J. Tompson (2019) Discriminator-actor-critic: addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations, Cited by: §2. [23] I. Kostrikov, O. Nachum, and J. Tompson (2019) Imitation learning via off-policy distribution matching. arXiv preprint arXiv:1912.05032. Cited by: §2, §5.1. [24] C. Lai, H. Wang, P. Hsieh, Y. F. Wang, M. Chen, and S. Sun (2024) Diffusion-reward adversarial imitation learning. Advances in Neural Information Processing Systems 37, p. 95456–95487. Cited by: §2, §5.1. [25] X. Li, V. Belagali, J. Shang, and M. S. Ryoo (2024) Crossway diffusion: improving diffusion-based visuomotor policy via self-supervised learning. In 2024 IEEE International Conference on Robotics and Automation (ICRA), p. 16841–16849. Cited by: §2. [26] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le (2023) Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, Cited by: §B.1, §6. [27] A. Y. Ng S. Russell et al. (2000) Algorithms for inverse reinforcement learning.. In Icml, Vol. 1, p. 2. Cited by: §1, §2. [28] F. Nuti, T. Franzmeyer, and J. F. Henriques (2023) Extracting reward functions from diffusion models. Advances in Neural Information Processing Systems 36, p. 50196–50220. Cited by: §2. [29] M. Orsini, A. Raichuk, L. Hussenot, D. Vincent, R. Dadashi, S. Girgin, M. Geist, O. Bachem, O. Pietquin, and M. Andrychowicz (2021) What matters for adversarial imitation learning?. Advances in Neural Information Processing Systems 34, p. 14656–14668. Cited by: §2. [30] M. Psenka, A. Escontrela, P. Abbeel, and Y. Ma (2024) Learning a diffusion model policy from rewards via q-score matching. In Forty-first International Conference on Machine Learning, Cited by: §1, §2. [31] N. D. Ratliff, J. A. Bagnell, and M. A. Zinkevich (2006) Maximum margin planning. In Proceedings of the 23rd international conference on Machine learning, p. 729–736. Cited by: §2. [32] J. Ren, G. Swamy, Z. S. Wu, J. A. Bagnell, and S. Choudhury (2024) Hybrid inverse reinforcement learning. arXiv preprint arXiv:2402.08848. Cited by: §2. [33] M. Reuss, M. Li, X. Jia, and R. Lioutikov (2023) Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532. Cited by: §1, §2. [34] S. Sharifzadeh, I. Chiotellis, R. Triebel, and D. Cremers (2016) Learning to drive using inverse reinforcement learning and deep q-networks. arXiv preprint arXiv:1612.03653. Cited by: §1. [35] H. Son, J. W. Jang, W. J. Han, and H. J. Hwang (2021) Sobolev training for physics informed neural networks. arXiv preprint arXiv:2101.08932. Cited by: §4.1. [36] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2021) Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, Cited by: §1, §3.2. [37] G. Swamy, S. Choudhury, J. A. Bagnell, and S. Wu (2021) Of moments and matching: a game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning, p. 10022–10032. Cited by: §2. [38] G. Swamy, D. Wu, S. Choudhury, D. Bagnell, and S. Wu (2023) Inverse reinforcement learning without reinforcement learning. In International Conference on Machine Learning, p. 33299–33318. Cited by: §2. [39] F. Torabi, G. Warnell, and P. Stone (2018) Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158. Cited by: §2. [40] B. Wang, G. Wu, T. Pang, Y. Zhang, and Y. Yin (2024) Diffail: diffusion adversarial imitation learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 38, p. 15447–15455. Cited by: §2. [41] J. Watson, S. Huang, and N. Heess (2023) Coherent soft imitation learning. Advances in Neural Information Processing Systems 36, p. 14540–14583. Cited by: §2. [42] R. Wu, Y. Chen, G. Swamy, K. Brantley, and W. Sun (2025) Diffusing states and matching scores: a new framework for imitation learning. In The Thirteenth International Conference on Learning Representations, Cited by: §2. [43] Z. Wu, L. Sun, W. Zhan, C. Yang, and M. Tomizuka (2020) Efficient sampling-based maximum entropy inverse reinforcement learning with application to autonomous driving. IEEE Robotics and Automation Letters 5 (4), p. 5355–5362. Cited by: §3.1. [44] O. G. Younis, R. Perez-Vicente, J. U. Balis, W. Dudley, A. Davey, and J. K. Terry (2024) Minari. Zenodo. External Links: Document, Link Cited by: §E.2. [45] S. Yue, G. Wang, W. Shao, Z. Zhang, S. Lin, J. Ren, and J. Zhang (2023) CLARE: conservative model-based reward learning for offline inverse reinforcement learning. In 11th International Conference on Learning Representations, ICLR 2023, Cited by: §5.1. [46] Y. Ze, G. Zhang, K. Zhang, C. Hu, M. Wang, and H. Xu (2024) 3D diffusion policy: generalizable visuomotor policy learning via simple 3d representations. In Proceedings of Robotics: Science and Systems (RSS), Cited by: §2. [47] S. Zeng, C. Li, A. Garcia, and M. Hong (2022) Maximum-likelihood inverse reinforcement learning with finite-time guarantees. Advances in Neural Information Processing Systems 35, p. 10122–10135. Cited by: §1, §2, §3.1, §5.1. [48] S. Zeng, C. Li, A. Garcia, and M. Hong (2023) When demonstrations meet generative world models: a maximum likelihood framework for offline inverse reinforcement learning. Advances in Neural Information Processing Systems 36, p. 65531–65565. Cited by: §3.1, §5.1. [49] M. Zhu, Y. Zhu, J. Li, J. Wen, Z. Xu, N. Liu, R. Cheng, C. Shen, Y. Peng, F. Feng, et al. (2025) Scaling diffusion policy in transformer to 1 billion parameters for robotic manipulation. In 2025 IEEE International Conference on Robotics and Automation (ICRA), p. 10838–10845. Cited by: §2. [50] B. D. Ziebart, J. A. Bagnell, and A. K. Dey (2010) Modeling interaction via the principle of maximum causal entropy. In Proceedings of the 27th International Conference on Machine Learning, p. 1255–1262. Cited by: §2. [51] B. D. Ziebart, A. L. Maas, J. A. Bagnell, A. K. Dey, et al. (2008) Maximum entropy inverse reinforcement learning.. In Aaai, Vol. 8, p. 1433–1438. Cited by: §1, §2. Appendix A Detailed Derivations A.1 Derivation of Diffusion Scores and Their Relation to Soft Q-Gradients In this appendix, we derive the identities underlying Section 3. We use t for the MDP time index and k for the diffusion step. For a fixed state ts_t, let t0∼πE(⋅∣t)a_t^0 _E(· _t) denote a clean expert action, and let tka_t^k denote its noisy version at diffusion step k. The DDPM forward process is defined by q(tk∣tk−1)=(tk,1−βktk−1,βkI),αk≔1−βk,α¯k≔∏j=1kαj.q(a_t^k _t^k-1)=N\! (a_t^k; 1- _k\,a_t^k-1, _kI ), _k 1- _k, α_k _j=1^k _j. This gives the closed-form marginal q(tk∣t0)=(tk,α¯kt0,(1−α¯k)I),q(a_t^k _t^0)=N\! (a_t^k; α_k\,a_t^0,(1- α_k)I ), or equivalently, tk=α¯kt0+1−α¯kϵ,ϵ∼(0,I).a_t^k= α_k\,a_t^0+ 1- α_k\, ε, ε (0,I). (16) This is the same forward noising relation used in Eq. (5). For each fixed (t,k)(s_t,k), let pk(∣t)p_k(a _t) denote the conditional density of noisy actions at diffusion step k, induced by sampling t0∼πE(⋅∣t)a_t^0 _E(· _t) and then applying Eq. (16). The score of this noisy action density is ∇logpk(∣t). _a p_k(a _t). The diffusion policy is trained by the noise-prediction objective ℒdiff(ϕ)=(t,t0)∼ρπEk∼Unif1,…,K,ϵ∼(0,I)‖ϵ−ϵϕ(tk,t,k)‖22.L_diff( φ)=E_ subarrayc(s_t,a_t^0) ρ _E\\ k \1,…,K\,\, ε (0,I) subarray \| ε- ε_ φ (a_t^k,s_t,k ) \|_2^2. For a fixed (tk,t,k)(a_t^k,s_t,k), the pointwise minimizer of this MSE objective is ϵϕ∗(tk,t,k)=[ϵ∣tk,t,k]. ε_ φ^*(a_t^k,s_t,k)=E\! [ ε _t^k,s_t,k ]. Using the Gaussian corruption in Eq. (16), the noisy-action score satisfies ∇logpk(∣t)|=tk=−11−α¯k[ϵ∣tk,t,k]. _a p_k(a _t) |_a=a_t^k=- 1 1- α_k\,E\! [ ε _t^k,s_t,k ]. Therefore, the learned noise predictor induces the score estimator gϕ(tk,t,k)≔−11−α¯kϵϕ(tk,t,k)≈∇logpk(∣t)|=tk.g_ φ(a_t^k,s_t,k) - 1 1- α_k ε_ φ(a_t^k,s_t,k)≈ _a p_k(a _t) |_a=a_t^k. (17) In the idealized limit of exact denoising, the approximation in Eq. (17) becomes exact for the noisy conditional density pk(∣t)p_k(a _t). For small diffusion noise, tka_t^k is close to the clean action t0a_t^0, and the noisy-action score approximates the clean expert-policy score in a local neighborhood of the expert action: gϕ(tk,t,k)≈∇logπE(∣t)|=tk.g_ φ(a_t^k,s_t,k)≈ _a _E(a _t) |_a=a_t^k. If the expert policy is soft-optimal, πE=π∗ _E=π^*. Under the reference-policy MaxEnt formulation in Eq. (1), we have π∗(∣)=μ(∣)exp(1ε(Q∗(,)−V∗())).π^*(a )=μ(a ) \! ( 1 (Q^*(s,a)-V^*(s) ) ). Taking the action gradient gives ∇logπ∗(∣)=∇logμ(∣)+1ε∇Q∗(,), _a π^*(a )= _a μ(a )+ 1 _aQ^*(s,a), because V∗()V^*(s) does not depend on a. Equivalently, ∇logπ∗(∣)μ(∣)=1ε∇Q∗(,), _a π^*(a )μ(a )= 1 _aQ^*(s,a), which is Eq. (3). Combining the low-noise score approximation with Eq. (3), we obtain gϕ(tk,t,k)−∇logμ(∣t)|=tk≈1ε∇Q∗(t,)|=tk.g_ φ(a_t^k,s_t,k)- _a μ(a _t) |_a=a_t^k≈ 1 _aQ^*(s_t,a) |_a=a_t^k. (18) This is the identity used in Eq. (8): the diffusion-policy score, after subtracting the reference-policy score, provides the action-gradient supervision for recovering the soft Q-function. A.2 Derivations for Stage I: Q-Recovery and Initial Value Fitting In this subsection, we justify the first-stage recovery of the uncalibrated soft Q-function and the initial value function. The key point is that matching action derivatives recovers the action-dependent geometry of Q∗Q^*, but cannot determine its state-dependent offset. We next justify the Gumbel-regression objective used to obtain the initial value estimate from the frozen Q Q_ ω. For a fixed state s and a fixed Q function Q~ Q, define ℓV(v;):=∼μ(⋅∣)[exp(Q~(,)−vε)−Q~(,)−vε−1]. _V(v;s):=E_a μ(· ) [ \! ( Q(s,a)-v )- Q(s,a)-v -1 ]. Differentiating with respect to v gives ∂vℓV(v;)=−1ε∼μ(⋅∣)[exp(Q~(,)−vε)]+1ε. ∂ v _V(v;s)=- 1 E_a μ(· ) [ \! ( Q(s,a)-v ) ]+ 1 . The first-order condition is therefore ∼μ(⋅∣)[exp(Q~(,)−vε)]=1.E_a μ(· ) [ \! ( Q(s,a)-v ) ]=1. Rearranging gives v=εlog∼μ(⋅∣)[exp(1εQ~(,))].v= _a μ(· ) [ \! ( 1 Q(s,a) ) ]. Moreover, ∂2∂v2ℓV(v;)=1ε2∼μ(⋅∣)[exp(Q~(,)−vε)]>0, ∂^2∂ v^2 _V(v;s)= 1 ^2E_a μ(· ) [ \! ( Q(s,a)-v ) ]>0, so the minimizer is unique. In Stage I, we instantiate Q~ Q as the frozen action value function Q Q_ ω. Replacing the expectation over μ(⋅∣)μ(· ) by actions gives ℒV()=,∼μ(⋅|)[exp(z)−z−1],z=Q^(t,t)−V^(t)ε.L_V( χ)=E_s,a μ(·|s) [ (z)-z-1 ], z= Q_ ω(s_t,a_t)- V_ χ(s_t) . This is Eq. (10). Since Q Q_ ω is fixed in this step, gradients are taken only with respect to χ. The resulting value network V V_ χ is the initial value estimate used in Stage I for offset calibration. A.3 Derivations for Stage I: Offset Calibration and Value Re-estimation Proof of Proposition 1. Fix any state s and a connected action neighborhood on which the low-noise diffusion-score supervision is reliable. If the derivative-matching term in Eq. (8) is minimized exactly, then for every initial action point a in this neighborhood, 1ε∇Q^(,)=gϕ(k,,k), 1 _a Q_ ω(s,a)=g_φ(a^k,s,k), where ka^k denotes the low-noise corrupted version of a. By Eq. (7), under the idealized assumptions that πE=π∗ _E=π^* and the diffusion score is exact, gϕ(k,,k)=1ε∇Q∗(,).g_φ(a^k,s,k)= 1 _aQ^*(s,a). Thus, ∇Q^(,)=∇Q∗(,). _a Q_ ω(s,a)= _aQ^*(s,a). Since the gradients match only with respect to the action variable, the two functions can differ by a term that depends on the state but not on the action. Therefore, there exists a state-dependent function b∗:→ℝb^*:S such that Q∗(,)=Q^(,)+b∗()Q^*(s,a)= Q_ ω(s,a)+b^*(s) on the sampled action neighborhood. This proves Proposition 1. ∎ Appendix B Flow-Matching Version of Diffusion Policy B.1 Flow-Matching Policy and Score Estimation The main method uses a DDPM-based diffusion policy to provide the score estimator gϕ(tk,t,k)g_φ(a_t^k,s_t,k). To keep the presentation concise, this subsection follows the simplified case that the reference policy μ is uniform. In this case, ∇logμ(∣)=0 _a μ(a )=0 inside the support, so the generative-model score can be directly used as the action-gradient signal for Stage I. This is not the only possible generative pretraining choice. A Flow-Matching-based model can also be used as a conditional action generator, and under a simple Gaussian interpolation path it can provide a score estimator that plays the same role in Stage I. This section explains this replacement. Flow Matching trains a time-dependent vector field that transports samples from a simple noise distribution to the data distribution [26]. In the policy setting, the data distribution is the expert action distribution conditioned on the current state. Let u∈[0,1]u∈[0,1] denote the flow time, where u=0u=0 corresponds to Gaussian noise and u=1u=1 corresponds to expert actions. Given an expert pair (t,t)∼E(s_t,a_t) _E and Gaussian noise ϵ∼(0,I) ε (0,I), the interpolation is defined as: tu=(1−u)ϵ+ut.a_t^u=(1-u) ε+ua_t. (19) Here, tua_t^u is an intermediate noisy action along the path from noise to the expert action. The derivative of this path with respect to u is dutu=t−ϵ. ddua_t^u=a_t- ε. A Flow-Matching policy therefore learns a state-conditioned velocity field vϕ(tu,t,u)v_φ(a_t^u,s_t,u) by the regression objective ℒFM(ϕ):=(t,t)∼Eu∼Unif(0,1),ϵ∼(0,I)[‖vϕ(tu,t,u)−(t−ϵ)‖22].L_FM(φ):=E_ subarrayc(s_t,a_t) _E\\ u (0,1),\, ε (0,I) subarray [ \|v_φ(a_t^u,s_t,u)-(a_t- ε) \|_2^2 ]. (20) After training, action generation can be performed by starting from Gaussian noise and integrating duu=vϕ(u,t,u) ddua^u=v_φ(a^u,s_t,u) from u=0u=0 to u=1u=1. The resulting endpoint is used as an action sampled from the learned policy. For LFIRL, we need not only a sampling procedure, but also a score signal for action-gradient matching. Let pu(∣t)p_u(a _t) denote the conditional density of the intermediate action tua_t^u at flow time u. Its score is ∇logpu(∣t), _a p_u(a _t), which indicates how an intermediate action should be adjusted to move toward higher-probability expert actions under state ts_t. Under the interpolation defined in Eq. (19), this score can be estimated from the learned velocity field. If the learned velocity is exact, then at a fixed (tu,t,u)(a_t^u,s_t,u) it estimates the conditional mean of the path derivative, vϕ(tu,t,u)≈[t−ϵ∣tu,t,u].v_φ(a_t^u,s_t,u) [a_t- ε _t^u,s_t,u]. Combining this relation with the interpolation equation gives [ϵ∣tu,t,u]≈tu−uvϕ(tu,t,u).E[ ε _t^u,s_t,u] _t^u-u\,v_φ(a_t^u,s_t,u). For the same Gaussian path, the score of pu(∣t)p_u(a _t) can be written in terms of this conditional noise estimate. Thus, for u<1u<1, we obtain the Flow-Matching score estimator gϕFM(tu,t,u):=−tu−uvϕ(tu,t,u)1−u≈∇logpu(∣t)|=tu.g_φ^FM(a_t^u,s_t,u):=- a_t^u-u\,v_φ(a_t^u,s_t,u)1-u≈ _a p_u(a _t) |_a=a_t^u. (21) The expression is used for u<1u<1; at u=1u=1, the path reaches the expert action endpoint and the denominator vanishes. In practice, the analogue of the low-noise DDPM regime is to use flow times close to, but not exactly equal to, 11. This gives the connection needed by LFIRL. When u is close to 11, the intermediate action tua_t^u is close to the initial expert action. Therefore, gϕFM(tu,t,u)g_φ^FM(a_t^u,s_t,u) estimates the local score of the expert action distribution around that action. Thus, in the low-noise regime, Eq. (3) gives gϕFM(tu,t,u)≈1ε∇Q∗(t,)|=tu.g_φ^FM(a_t^u,s_t,u)≈ 1 _aQ^*(s_t,a) |_a=a_t^u. (22) Therefore, a Flow-Matching-based diffusion policy can be used in Stage I by replacing the DDPM score estimator gϕ(tk,t,k)g_φ(a_t^k,s_t,k) in Eq. (8) with gϕFM(tu,t,u)g_φ^FM(a_t^u,s_t,u). Figure 7: Additional ablation comparing DDPM-based and Flow-Matching-based pretraining in LFIRL. The two variants use the same subsequent sequential reward-recovery pipeline, but differ in how the pretrained generative policy is obtained. Higher is better. B.2 Additional Results: DDPM vs. Flow Matching Pretraining In the main experiments, LFIRL uses a DDPM-based diffusion policy to provide the score signal for recovering the soft Q-function. We also evaluate an alternative implementation in which the diffusion-policy pretraining stage is replaced by a Flow Matching version, while keeping the remaining LFIRL pipeline unchanged. This comparison isolates the effect of the pretrained generative policy used to provide local action-gradient supervision. As shown in Fig. 7, under the same number of pretraining epochs, the Flow-Matching-based variant generally performs worse than the DDPM-based variant. This suggests that, in our setting, DDPM pretraining provides a more effective score signal for the subsequent Q-recovery stage. Since LFIRL relies on the pretrained generative policy to supply local action-gradient information, a weaker pretraining stage can directly reduce the quality of the recovered Q-function and subsequently degrade the final recovered reward. Appendix C Sensitivity to the Reference Policy The reference policy μ specifies the behavioral baseline relative to which the expert policy is interpreted. Under the reference-policy MaxEnt formulation, ∇logπE(∣)=∇logμ(∣)+1ε∇Q∗(,). _a _E(a )= _a μ(a )+ 1 _aQ^*(s,a). (23) We examine how the choice of μ affects LFIRL by comparing a uniform reference policy with a state-conditioned Gaussian reference policy. The two variants use the same expert demonstrations and LFIRL training configuration. The uniform variant uses a uniform distribution over the bounded action space, for which ∇logμ(∣)=0 _a μ(a )=0 in the interior of its support. For the Gaussian variant, we initialize an SAC actor before LFIRL training and freeze it throughout the reward-recovery pipeline. The actor outputs the mean and standard deviation of a state-conditioned action distribution and is not fitted to the expert demonstrations. At each noisy action tka_t^k, we evaluate the score of this frozen reference policy and use gϕ(tk,t,k)−∇logμ(∣t)|=tkg_ φ(a_t^k,s_t,k)- _a μ(a _t) |_a=a_t^k as the target for ε−1∇Q^(t,tk) ^-1 _a Q_ ω(s_t,a_t^k). Actions used for soft-value estimation are also sampled from the same frozen reference policy. We evaluate these variants on MuJoCo tasks, where the expert demonstrations can be generated by stochastic sampling from pretrained SAC experts. This allows the reference-policy comparison to be conducted with a specified source of expert actions; the generating policy distributions of the robotic-manipulation datasets cannot be verified in the same way. Table 4 reports the resulting episodic returns. Table 4: Sensitivity to the reference policy on MuJoCo tasks. Results are episodic returns, reported as mean ± standard deviation. The expert return is included for context. Environment LFIRL-Uniform LFIRL-Gaussian Expert Swimmer-v5 284.56± 27.64 208.42± 7.25 315.55± 1.38 Walker2d-v5 4933.41± 52.79 4722.40± 230.42 5861.18± 73.99 Hopper-v5 3413.50± 210.23 3868.04± 79.58 4098.17± 247.70 The choice of reference policy affects performance, and the direction of the effect depends on the task. The uniform variant performs better on Swimmer and Walker2d, whereas the Gaussian variant performs better on Hopper. LFIRL therefore is not invariant to μ: changing the behavioral baseline changes the reward interpretation and can alter the recovered reward. Nevertheless, the Gaussian-reference variant remains effective, showing that the recovery procedure is not restricted to the uniform-reference simplification. This experiment assumes that the reference policy is specified. It does not infer μ from expert demonstrations. In general, demonstrations alone do not uniquely separate an unknown reference policy from the reward without additional assumptions; jointly inferring the two is beyond the scope of this work. Appendix D Sensitivity to Stage-I Hyperparameters LFIRL has several sequential recovery stages, but their hyperparameters need not have the same effect on performance. The value-recovery stage uses a Gumbel-regression-style objective, while the final reward-recovery stage enforces Bellman consistency. Here, we focus on three choices that directly affect the local Q-recovery signal in Stage I: the magnitude of the diffusion noise, the value-anchoring coefficient λ, and the intended margin ξ in value anchoring. We vary one parameter at a time while holding the other two at their default settings: the low-noise schedule, λ=1λ=1, and ξ=1ξ=1. The noise settings scale the coefficient 1−α¯k 1- α_k in the forward noising process by 1×1×, 4×4×, and 8×8×, respectively. Table 5 reports PointMaze success rates under these settings. Table 5: Sensitivity of PointMaze success rate to Stage-I hyperparameters. Each group varies one parameter while the remaining parameters are held at their default values. Bold indicates the highest result within each parameter group for each environment. Diffusion-noise magnitude Value-anchoring coefficient Intended margin Environment Low Medium High λ=0.5λ=0.5 λ=1λ=1 λ=2λ=2 ξ=0.5ξ=0.5 ξ=1ξ=1 ξ=2ξ=2 UMaze 0.98 0.87 0.68 0.92 0.98 0.97 0.57 0.98 0.61 Medium 0.39 0.35 0.25 0.38 0.39 0.35 0.24 0.39 0.26 Large 0.31 0.29 0.25 0.26 0.31 0.34 0.18 0.31 0.24 Diffusion-noise magnitude. Performance decreases as the noise coefficient increases. Increasing 1−α¯k 1- α_k moves noisy actions farther from the clean expert-action distribution, where the local approximation between the noisy-action score and the expert-policy score becomes less reliable. The decline is modest from low to medium noise on Medium and Large, but larger at high noise across all three environments. This supports using the low-noise schedule as the default for local Q-gradient supervision. Value-anchoring coefficient λ. Performance is comparatively stable between λ=1λ=1 and λ=2λ=2. Increasing λ to 22 slightly improves Large while causing a small decrease on UMaze and Medium. Reducing it to 0.50.5 weakens the anchoring signal and lowers performance, particularly on Large. Thus, the method is less sensitive to this coefficient once anchoring is sufficiently strong, with λ=1λ=1 providing a stable choice across the three environments. Intended margin ξ. The intended margin has the strongest effect among the three tested parameters. Both ξ=0.5ξ=0.5 and ξ=2ξ=2 underperform ξ=1ξ=1 on every environment. A small margin provides insufficient separation between expert actions and their perturbed neighbors. An excessively large margin demands larger local Q differences and may conflict with gradient matching or suppress near-optimal actions. In these experiments, ξ=1ξ=1 provides the most reliable balance across all three environments. Overall, these results identify the diffusion-noise magnitude and the anchoring margin as the more consequential Stage-I settings. The value-anchoring coefficient is less sensitive once it is large enough to provide effective anchoring. Appendix E Additional Experimental Details E.1 Implementation Details and Hyperparameter Setup The overall training procedure of LFIRL is given in Alg. 1. Key network architecture and hyperparameter settings for each environment are summarized in Tab. 6 and Tab. 7. Table 6: Network architecture and hyperparameter setup for the PointMaze environments. PointMaze_UMaze-v3 PointMaze_Medium-v3 PointMaze_Large-v3 Expert demonstrations 2000 2000 2000 Policy pretraining epochs 60 40 40 Q/b/V/r learning rate 3e-5 3e-5 3e-5 Stage passes (Q,b,V,r)(Q,b,V,r) 8/8/20/20 8/8/20/20 8/8/20/20 Q, b, V, r hidden layers 256, 256, 256, 256 256, 256, 256, 256 256, 256, 256, 256 IRL batch size 256 256 256 Discount factor γ 0.99 0.99 0.99 EMA coefficient τ 0.005 0.005 0.005 Table 7: Network architecture and hyperparameter setup for the robotic manipulation environments. FrankaKitchen-v1 gym_pusht/PushT-v0 AdroitHandPen-v1 Expert demonstrations 19 200 200 Policy pretraining epochs 20 200 40 Q/b/V/r learning rate 3e-5 3e-4 3e-5 Stage passes (Q,b,V,r)(Q,b,V,r) 8/8/20/20 40/40/100/100 40/40/100/100 Q, b, V, r hidden layers 256, 256, 256, 256 256, 256, 256, 256 256, 256, 256, 256 IRL batch size 256 256 256 Discount factor γ 0.99 0.99 0.99 EMA coefficient τ 0.005 0.005 0.005 E.2 Expert Demonstrations Sources The sources of offline trajectory datasets for experts are provided in Tab. 8. In Maze and Kitchen tasks, we directly use the expert trajectories from the Minari Offline Reinforcement Learning datasets [44]. For the Push-T task, we use the expert trajectories from the dataset in Diffusion Policy [5]. Table 8: The sources of expert policies or demonstrations. Task Source UMaze https://minari.farama.org/datasets/D4RL/pointmaze/umaze-v2/ Medium https://minari.farama.org/datasets/D4RL/pointmaze/medium-v2/ Large https://minari.farama.org/datasets/D4RL/pointmaze/large-v2/ Franka Kitchen https://minari.farama.org/datasets/D4RL/kitchen/complete-v2/ Push-T https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip Pen https://minari.farama.org/datasets/D4RL/pen/expert-v2/ E.3 Hardware Information Hardware specifications are provided in Tab. 9. Table 9: Hardware configuration used in experiments. Hardware Specifications CPU AMD EPYC 7713 64-Core Processor @ 2 GHz GPU NVIDIA A100-SXM4-80GB @ 1215 MHz Memory 2 TB