Imitation Learning as Return Distribution Matching
Authors: Filippo Lazzati
Alberto Maria Metelli
Venue: ICLR 2026 Poster (2026)
Keywords: Imitation Learning
Behavioral Cloning
Risk
Theory
Primary area: reinforcement learning
TL;DR: We introduce and study the problem of finding a policy that induces a return distribution close to that of the expert.

Abstract: We study the problem of training a risk-sensitive reinforcement learning (RL) agent through imitation learning (IL). Unlike standard IL, our goal is not only to train an agent that matches the expert’s expected return (i.e., its *average performance*) but also its *risk attitude* (i.e., other features of the return distribution, such as variance). We propose a general formulation of the risk-sensitive IL problem in which the
objective is to match the expert’s return distribution in Wasserstein distance. We focus on the tabular setting and assume the expert’s reward is *known*. After demonstrating the limited expressivity of Markovian policies for this task, we introduce an efficient and sufficiently expressive subclass of non-Markovian policies tailored to it. Building on this subclass, we develop two provably efficient algorithms—RS-BC and RS-KT —for solving the problem when the transition model is unknown and known, respectively. We show that RS-KT achieves substantially lower sample complexity than RS-BC by exploiting dynamics information. We further demonstrate the sample efficiency of return distribution matching in the setting where the expert’s reward is *unknown* by designing an oracle-based variant of RS-KT. Finally, we complement our theoretical analysis of RS-KT and RS-BC with numerical simulations, highlighting both their sample efficiency and the advantages of non-Markovian policies over standard sample-efficient IL algorithms.

Published as a conference paper at ICLR 2026
IMITATION LEARNING AS RETURN DISTRIBUTION MATCHING
Filippo Lazzati
Alberto Maria Metelli
Politecnico di Milano Milan, Italy
Politecnico di Milano
Milan, Italy a bertoma r www.mmteeeee...... imi
filippo. laz zat
ABSTRACT
We study the problem training a risk-sensitive reinforcement learning (RL) agent through imitation learning (IL). Unlike standard IL, goal is not only to train an agent that matches the expert's expected return (i.e., its average performance) but also its risk attitude (i.e., other features of the return distribution, such variance). We propose a general formulation of the sissssenst IL problem in which the objective is to match the expert's is return distribution in Wasserstein distance. We focus on the tabular setting and assume the expert's 's reward is known. After demon- strating the limited expressivity of Markovian policies for this task, we introduce an efficient and sufficiently expressive subclass of non non-Markovian policies tailored to it. Building on this subclass, we develop two provably efficient algorithms RR--BC and RS-KT for solving the problem when the transition model is unknown and known, respectively. We show that RS-KT achieves substantially lower sample complexity than RS-BC by exploiting dynamics information. We further demon- strate the sample efficiency of return distribution matching in the setting where the expert's reward is unknown by designing oracle-basee variant of RS-KT. Finally, we complement our theoretical analysis of RS-KT and RS-BC with numer- ical simulations, highlighting both their sample efficiency and the advantages of non-Mark Mon-Markovian policies over standard sample-efficient IL. algorithms.
INTRODUCTION
Imitation Learning (IL) (Abbeel & Ng 2004 Osa et al. 2018) is the problem of training an agent to
behave by mimicking demonstrations from an expert. By removing the need for designing a reward function for the bask-which is often a difficult challenge (Hadfield-Mend) et al., 2 has
been successfully applied in diverse domains, including robotics (Argall | 2009 autonomous
driving (Le Mero et al. 2022). finance (Goluža et al. 2023), and LLMs (Zhao et al. 2005).
Most existing algorithms- including BC (Behavioral Cloning) (Pomerleau 1988), GAIL (Ho Ermon] 2016), and others (Ziebart) 2010, et al. 2020; Garg et al.] 2022 focus finding Markovian policy that best matches the expert's 's occupancy measure. This focus motivated by two observations. First, matching occupancy measures guarantees that the expecied return of our policy close to the expert's, regardless of the expert's unknown reward (Abbeel & Ng 2004). Second, Markovian policies sufficiently expressive. Indeed, for any arbitrary policy, there exists Markovian policy with the same occupancy measure (Puterman 1994).
By focusing solely on the occupancy measure wwiih captures the expected value of the return Instribution-standdee IL algorithms are inherently risk- risk-neutral, ignoring other characteristics of the return distribution such as the variance (Mannor & Tsitsiklis, 2011). However, expert demonstrations often come from humans who, in domains like finance (Föllmer & Schied, (016) or autonomous driving (Bernhard al. 2019), exhibit risk-sensitive behavior under stochasticity. these settings, the key aspect of the demonstrations is the expert's 's risk attitude, encoded in the shape of the return distribution (Bellemare al. 2023), but overlooked by standard IL methods.
$ *Coreesonding author.

---

Published as a conference paper at ICLR 2026
To address this, Santara et al. (2018) ; Lacotte et al. (2019) proposed extending occupancy measure matching to -iskiiensttvve settings by additionally matching the Conditional Value at Risk (CVaR) (Rockafellar & Uryasev 2000) at a chosen level of the expert's 's return distribution. Intuitively, by seeking the Markovian policy that best matches both the expectation and the CVaR at level o, these algorithms attempt to imitate not only the expert's average performance but also its tail behavior.
However, this approach faces two main limitations: (t) matching only the expectation and the CVaR level a captures narrow slice of the expert's full return distribution and thus provides a weak imitation of risk attitude, and (fi) Markovian policies are not expressive enough to capture all relevant risk- sensitive behaviors (Bellemare et al., 2023), leading to misspecification error. To overcome these challenges, we reformulate risk-sensitiv IL as matching the expert's 's entire return distribution, and we design algorithms that perform policy search efficiently in the space of non-Markoviag policies.
Contributions. Our main contributions are as follows:
We formulate II. as the problem of matching the expert's return distribution in Wasserstein distance. We motivate this setting and demonstrate the importance of non-Markovian policies (Section (
We introduce an efficient and expressive subclass of non-Markovian policies for the tabular setting with known expert reward, and use it to develop two provably efficient algorithms, RS-BC and RS-KT, for the cases where the transition model is unknown and known, respectively (Section 4
We show that the tabular setting with an unknown expert reward but a known transition model, sample efficiency can still be achieved by devising an oracle-based variant of RS-KT (Section 5)
Finally, we conduct numerical simulations to empirically evaluate RS-BC and RS-KT, comparing them in particular against standard provably efficient IL algorithms (Section 6
All proofs are provided in Appendix B D and additional related work is discussed in Appendix
2 PRELIMINARIES
Notation. Given natural number n. E N, we define [n] {1, ni.. Given a real number I E I. we let [x] == maxme2:mss m be the floor function. Given two sets X, y, we denote by A and AJ. respectively. the set of probability measures on X' and the set of functions from y to A Given probabilities p, E A on finite support X' = R. with cumulative distributions Fp and Fq. the (1-)Wasserstein distance is W(p,g) q) := SR 12
F2(x) |dx (Villani 2008). and the total variation distance is TV(p, q) == 1 Erex (p(x) (x))). The CVaR at level (0. 1) of is CV) So F (w)dd, where F inf seR: (2) Rockafellar & Uryasev 2000
Markoy Decision Processes (MDPs). A tabular finit-hhoizzon episodic Markov Decision Process
without reward (MDP) (Puterman 1994 Abbeel & Ng. 2004) is a tuple M == (S, A, H, so: p). where S is the finite state space (S = [S]. A is the finite action space (A = [A). H E N is the horizon, 80 E S is the initial state, and p E A
enriched with a reward r S x A x [#] SxAx [H] is the transition model. An MDP R M can be denote the of state-action trajectories length
[0, 1]. to obtain an MDP M, == (S, A, H, so, P.r.. We
1 as S4 (8 t for all E [H + I]. and define Ω = Unatn Sr, For any trajectory w = (s1, a1...., S/: an) and reward we let G(wrr) En Th' (sh; an)) denote the sum of rewards of
policy prescribes actions in states. We denote by IINM = Adxs the set of non-Markovian (i.e., history-dependent) policies and by IIM- = AMMxs the set of Markovian policies. Note that I = NM Playing policy E NM in an MDP M (or an MDP) induces probability distribution over trajectories 1dT A82+ The occupancy measure dt of in M is the marginal of pa over state-action pairs at given stage: di (s, a)= 1pt (sn. 8, ag a). Given reward r, the random sum of rewards Σ 1ehhrrr(spp (r an)) induced by the execution of π is the return, and we denote its distribution, called the return distribution
(Bellemare et al. 2023), as 1/. (g) == 1at (I (8h. on) = g) for all g [0, H). Lastly, we denote by JH the expectation of n).
Neglecting past rewards in NM without loss of generality since we consider deterministic rewards.
2

---

Published as a conference paper at ICLR 2026
Imitation Learning (IL). In IL, are given dataset D fwa [ of N trajectories wi QH+1, collected by (potentially (on-Markovian) expert policy n e I7 and the goal is to find policy it with expected return close to that of T8 under the expert's 's reward , (Abbeel & Ng 2004):
arg min meTINM
Since F is usually unknown, the problem is reformulated in robust terms as finding a policy that performs comparably to the expert for any possible reward:
â E arg min RETTI xxXxxx [0.1]
max
J
Ho & Ermon (006) showed that Eq. essentially reduces to finding policy â whose occupancy measure d" is close to the expert's dm Thus, Eq. (2) (and D can be addressed with Markovian policies ITM, which are known to be expressive enough for occupancy measure matching problems (e.g., see sereLar & Tachet Des Combes (2023)). Based on
Interestingly, Abbeel & Ng 2004
these insights, recent theoretical work (Rajaraman et al. | Foster et al. 2024) demonstrated that IL be solved provably efficiently.
Risk-sensitiv IL. Optimizing Egs. 4 guarantees imitation of the expert's average perfor- mance, i.e., its expected return Je but does not guarantee imitation of its risk attitude, encoded
in the shape of its return distribution n7 For this reason, Santara et al. 2018); Lacotte et al. (2019) proposed strengthening the standard IL formulation by also matching the CVaR at some level
E (0, 1) of 77 in addition to JE Formally, in the unknown r setting, they propose
(J. J.) + (CCVaRa(n; (V
arg min max REIIN rXSxAx [
(3)
where p(x) if & and +oo otherwise. This extension, however, makes the problem substantially harder than standard II, since the optimal solution to Eq. (3) cannot, in general. found among the Markovian policies IIM (even if r was known) Nevertheless, Santara et al. (2018); Lacotte et al. (201)) ignored this aspect and proposed algorithms that output Markovian policies, which, however, may not be suited for general non-Markovian experts like humans (Mandlekar et al. 2022
RETURN DISTRIBUTION MATCHING
Our goal is to train agents that match not only the expert's 's expected return but also its risk attitude. Standard IL is unsuitable since it ignores risk, while existing risk-sensitive IL methods only capture a limited aspect of the expert's return distribution, i.e., the CVaR at fixed level. We therefore propose an alternative formulation, called return distribution matching (RDM), which requires matching the entire expert return distribution 7, in Wasserstein distance:
R E arg min W T 11,5
TEIINM
This objective extends Eq. 4 and assumes knowledge of the expert reward E Our focus will primarily be this known-reward setting (see Section 4 both because is of independent interest (similar to inverse constrained RL (Malik et al. (02TT) and utility learning (Lazzati & Metelli (023]), challenging unknown-reward case, which next (3)
and because it provides foundation for the ( we propose a robust version of RDM: formalize. When f is unknown, following Eqs.
arg min
max MM r:Sx [M] [0.1]]
W
which requires matching the expert's return distribution for all possible rewards (see Section 5}. We now provide three key observations about RDM. First, matching return distributions in Wasserstein
Thh formulation of Santara optimizing the CVaR; however, the legh-level idea and the issues with non-Markovian policies remain same. Indeed, since the optimal policy to CVaR optimization problem is, in general, non-Markovian (Bäuerle & On 201 1), the solution Eq. 3 also helongs to NT due the hard constraint enforced by p.
2018 is slightly different, they require matching the expectation while
3

---

Published as a conference paper at ICLR 2026
distance is strictly more general than matching expected return or CVaR. Indeed, if W( M 7 then J - & and C. (1,²2 -- (1).s)) & e/a for any E (0, 1) (see Appendix B.1] Second, Wasserstein distance is essential for favorable sample complexity. Stronger metries, such as total variation. require exponentially many expert trajectories in some instances, even if the MDP M, fully known (see Appendix B.2 for the proof):
Theorem 3.1. There exist MDP M, with S, A, H > 2 and expert policy N II such trajectories, any algorithm 2 satisfies
that, even with = ( 1) - 1-
DE TV TV(1))7)11 (7, MTE
1
2c
where f is the output of given in input M, and DE
Finally, we remark that Markovian policies are not expressive enough for RDM, since they fail to capture the behavior of non- Markovaon experts even for the simpler risk-sensitive problem Eq.
3 Note that the gap can be significant even with very short horizons (proof in Appendix (B22):
Proposition 3.2. There exist an MDP M, with horizon H 3 and an expert policy F II
such that any Markovian policy π IIM satisfies
W W(1 2 1 77,5
3 0.5.
Therefore, new algorithms that output non- Markkoaoo policies are needed tackle RDM effectively.
KNOWN-REWARD SETTING
In this section, assume the expert's 's reward y is known and present our main contributions. In Section [.11, we introduce an efficient and sufficiently expressive subset of non-Markovian policies for RDM. Building on this, in Sections 4.2 and 4.3 we develop two provably efficient algorithms, RS-BC and RS-KT, for the cases where the transition model is unknown and known, respectively.
4.1 AN EFFCIENT CLASS OF MON-MARKOVIAN POLICIES
Proposition 3.2 shows that Markovian policies M are not expressive enough for RDM. At the same time, optimizing Eq. (4) over the entire set of non-Markovian policies IM is intractable due to the curse dimensionality. In this section, we introduce a subclass of non Markovavnian policies II(r,"), lying between NI and IIa, that allows us to address RDM accurately without sacrificing efficiency, The trade-off between accuracy and efficiency is controlled by parameter 0 > O. To this end, we first establish some notation. For any reward r, define II(r) S IINM as the set of policies whose choice of action depends only the current state stage h, and the cumulative reward so far G(w;; r):
(T()) = (T E NM
3Φ E A ())ls,w) = (alss, G(w;r)) Vs E S E A, h [M] w E Rn)
where G, (g E - B] Ω G(w:r) 9) denotes the set of possible cumulative reward values attainable at any stage except the last c Observe that each π E I[(r) can be interpreted as a Markovian policy in the MDP obtained by augmenting the state space with the cumulative reward G,. Next, for any reward and expert policy T ITNM, define whose probability of taking an action in state s with history w E Ur coincides with the "average" probability with which T selects in after accumulating G(w;r) r) reward:
e I[r)) as the policy
P (sh = s, an = =a hh-1 Th ax) = G(w;r))
TH (als,
(6)
p (sh = 8, Eh-1 (w aw) G(w;r))
If the denominator is zero, we (a(sss w) 1/A. With these two ingredients, II()) and m., can state the following important intermediate result (see Appendix (ppendix/C.I.Iffor the proof):
Lemma 4.1. Let M, be any MDP and let F IM be any policy. Then, the policy n, E [I
R satisfies 7 (g) THE (g) for all E [0, H].
Note that is always finite in tabular MDPs with deterministic rewards.
4

---

Published as a conference paper at ICLR 2026
Algorithm 1: RS-BC (Risk-Sensitive Behavior Cloning)
Input:Datasee D = [(i] ai...., db., kENN, reward r parameter 0 // Count the state-act ieleiiooccumppp ivi reward eccurrences
M,(s.g, EMIN 1ssh s, E (sk,, -=)a=))))) g} W
// Retrieve the policy #(als,w)
Mr (s. E
if E. (a wh [H]] 8 a
My (a. (s.((()))))))
otherwise
Return f
In words, Lemma 4.1 guarantees that IT always contains at least one policy with exactly the
same return distribution the expert, namely one that minimizes Eq. (4). Moreover, provides an
analytical expression for such policy, 7, (see Eq. 6). Unfortunately, [[(r_) ) is not always desirable. As discussed in Appendix C.1.1 there exist reward functions E for which I[rr IT and, thus, T.e may be an arbitrary Markoaakn policy, inefficient to store. Intuitively, this occurs when each trajectory yields a different return value, leading to the exponential dependence |G.e = (SA) #-1 To overcome this limitation, inspired by prior work (Bastani et al.) [022 Lazzati & Metelli] 2025), we adopt a discretization approach. Given a parameter E (0, 1], we define a d-covering of the interval [0, h 1] as 2%h = {0,6 0, 20, ..... 1h- 1/010) for all h E [H + 1]. and set ) := ) H+1 Then, for any reward r, we construct the discretized reward ro as (breaking ties arbitrarily):
(8, == arg min |r (s)
(ss,a,h)eSx Ax [H].
xE>
Crucially, note that Gra S )d for any reward since summing discretized rewards always yields discretized values. Because yd has "small" (polynomial) size. | O(H/0), the policy set II(rff) is also small, [I(rr")| < [IINM], and every policy 11(rf'), including T can be efficiently stored (with ((SAH)>*) memory). The following lemma shows that the approximation error introduced by using policies in IT(r)) instead of (Tr for RDM can be tightly controlled by (see Appendix for the proof):
Lemma 4.2. Let 0 (0, 1]. Let M,E be any MDP and π E ITM any policy. Then, the policy H8.
T E [[(FF) satisfies (M, 7 -T,E <H8.
In short, Lemma 4.2 shows that efficient and accurate solutions to the RDM problem can be sought within (l(r ). In particular, 11())) contains n,5. whose error can be reduced by decreasing 0, the cost of increased memory requirements for storing the policy, which scale as O(1/8). In the next two sections, show how to build efficient RDM algorithms based on II(r and T,5-
4.2 NO-INTERACTION SETTING
In this section, present RS-BC (Risk-Sensitive Behavior Cloning, Algorithm efficient algorithm for RDM in the no-interaction (offline) setting, where we neither know nor have access to the transition model of the environment M, and are instead given only dataset D E of N expert trajectories together with the expert's ' reward F The idea of RS-BC is simple: directly use E estimate policy R, whose return distribution is guaranteed by Lemma 4.2 be close to that of the expert. RS-BC estimates n (als, w) as the fraction of the times in D that the expert
provably
took action a in state after collecting G(w;r discretized cumulative reward (see Line 2) This "empirical" estimator follows closely the definition of T in Eq. (6), with probability terms replaced by counts M (computed at Line The next result shows that RS-BC sample-efficient by providing a worst-case upper bound on its sample complexity (proof in Appendix C.2.
Theorem 4.3. Let c (0, H and 5 (0. 1). Let M, be any MDP and let E E TM be any policy. Then, choosing f = c/(4H, with probability at least 8, the policy â output by Algorithm
'The "behavior cloning" part in RS-BC comes from the intuition that RS-BC can be seen as performing BC after augmenting the state space with the (discretized) cumulative reward.
5

---

Published as a conference paper at ICLR 2026
Algorithm 2: RS-KT (Risk-Sensitive imitation with Known Transition)
Input: Dataset D aisi,.i,..... din.)i) reward r", parameter 0, transition model p
// Estimate the return distribution of the expert V/, ((g) & iij 1 af) 9) Vg E y
// Computo the policy in I((ff) closest to f via Eq. arg ( W(r
10
Return f
satisfies W(n ( & with number of samples:
SH In 1 (A+ + (A+In ln)
N & θ
In words, Theorem
4.3 shows that RS-BC requires a polynomial (in the quantities of interest S, A, H, 1/e. In( (n(1))) number of samples to output a good imitation policy for RDM with high probability Compared to the best existing upper bound for standard IL. O(SAH In( (n(1 (Corollary 3.1 of Foster et al. (2024)), we observe a gap (H /eaee///nnn///// (n(1/5)). This reasonable, as RDM appears more complex than occupancy measure matching (e.g.. it requires non-Mart non-Markovian poli- cies). We conjecture that the gap unimprovable, while the (H²) dependence may be large but comparable to the ((11²) rate in the related setting of IL from observation alone (Theorem 3.3 of Sun et al. (2019). See Appendix C.2.2 for further discussion. Finally, RS-BC is also computationally efficient, since both Lines]
2 require only O(SAH)*) iterations and memory.
4.3 KNOWN-TRANSING
SETTING
In this section, we present RS-KT (Risk-Sensitive imitation with Known Transition, Algorithm 2 a provably efficient algorithm for RDM in the nnown-transition setting, where we have access to the transition model p of the enviromment, in addition to the expert's dataset D and reward r By leveraging knowledge of p, RS-KT achieves a drastic reduction sample complexity compared to RS-BC. The idea behind RS-KT is straightforward. First, the expert dataset DD to compute estimate f of the expert's 's return distribution T7, (Line 1 Then, exploit knowledge of r and p to identify the policy in 11(r5) whose return distribution is closest to 7 (Line) We now show that RS-KT is sample efficient (see Appendix C.3.1 for proof):
Theorem 4.4. Let (0, H] and E (0, 1). Let M,E be any MDP and T 11 any policy. Assume that the optimization problem in Line solved exactly. Then, choosing = e/(7H), with probability 1 - 8, the policy t output by Algorithm 2 satisfies W(n .Taa) & with:
2
NEO
(9)
Interestingly, Theorem 4.4 shows that RS-KT has sample complexity independent of S and A, making it substantially more sample efficient than any algorithm for standard IL in large MDPs, where ((S) samples are required even with knowledge of p (Theorem 5.1 of Rajaraman et al. (2020). We believe the (11/²) rate is tight, it matches the lower bound for estimating distribution in Wasserstein distance (Theorem 3.1 of Bobkov & Ledoux (2019)) (2019)) Compared to RS-BC, the reduction in sample complexity is drastic, (SSAH*)e). Extending RS-KT to settings where p is unknown but can be estimated from online interaction with the environment is an interesting direction for future work (see Xu al. (2023) for the standard IL case). If Line is solved with error <ppx. then Theorem 4.4
guarantees W(nn", )x < 6 Eapx- In distributional RL terms (Bellemare et al. 2023), RS-KT adopts a fixed-size categorical representation at Line | Finally, we refer the reader to Appendices C.4 and C.5 for extensions of RS-BC and RT to the cases where , is unknown but either belongs to a given finite set or is linear in a given feature map, and to Appendix C.6 for greater generalization.
notation omits logarithmic in S, A, H, 1/c, In(16))

---

Published as a conference paper at ICLR 2026
We now turn to the computational complexity of RS-KT. Crucially, the optimization problem in Line can be formulated as a linear program (LP) with a polynomial number (SSAH)))) of variables and constraints, and therefore can be solved efficiently. To see this, let M = (S, A, Su. p) be the MDP R with augmented state space S = S y initial state Su (su. 0) E S, and transition model p defined as Ph, (s.g) g| o a) = Pa o g r (s', (5', a')] i.e., the probability of reaching (s, e S by playing in (s',g') g¹) at stage h coincides with pa (s'1 Is, a) if the reward received in 6', a' is 9 - g'. and 0 otherwise. Since [ coincides with the set of Markovian policies in M (Bäuerle & Rieder 2014, Lazzati & Metelli 2025), we can rewrite Line 2 as a variant of the occupancy measure matching problem in M:
min deK.,ne A33
W(n. 7)
s.t. ((g))
Σ
du(s,9 -
E yo,
(10)
(s.0)eSxA
where the constraint enforces that 7 is the return distribution induced by the occupancy measure and K denotes the set of feasible occupancy measures in M (Puterman 1994)
{de ==
S 1
Σ (50, a) = 1 v& 5, 1 2: Σ (3) = Σ dn d d')Pp'' (3]3 dd) d)} 3'a"'
In words, Eq. 10p searches for an occupancy measure d IC that induces the return distribution closest to 77. From such solution, policy î E 11(rz)) with occupancy da (and thus return distribution u. v)) can recovered via
da(s, G(w;
(a)s,w)) E. (s, Glw; (.)
Ah [H] SE S, A A, WE Sr.
when the denominator is nonzero, and ((a w) 1/A otherwise (Syed et al. 2008). Observe that all the constraints in Eq. (10) are linear in d and 17, and that the objective W(n. 7)) can also be written linearly (see Peyré Cuturi 2019) and Appendix C.3.2).
STATISTICAL INSIGHTS ON THE UNKNOWN- SETTING
In this section, we assume that the expert's reward F is unknown, and provide some statistical insights on RDM in the "robust" form Eq. (5). Specifically, we show that, perhaps surprisingly, this complex problem requires only polynomial number of expert demonstrations be accurately solved when the transition model is known, even in the worst case. To establish this, we first prove that a polynomial number of expert demonstrations suffices to accurately estimate the expert's 's return distribution under any reward (proof in Appendix
Theorem 5.1. Let E (0, H] and 8 e (0, 1). Let M,, be any MDP and n II any policy. Then, choosing 0 = E/ e2HH), a number of samples
N θ SAH3
))
In
88)
(11)
suffices guarantee that, with probability N EwEDS Il (G((;r))) = g) Vg, T, we have:
least 8, for the estimator nr(g)
W 1r mr) E
max r P.SAA M [0.1]
In brief, an expert dataset DE of size in Eq. 11 suffices to accurately estimate the expert's 's return distribution nf under any reward r via the estimator îr. As consequence, any policy fे that induces return distributions close to î for all possible rewards accurately solves the robust RDM problem:
Theorem 5.2. Under the conditions of Theoremm.. assume access to computational oracle that takes as input the dataset DE and the transition model p. and outputs a solution to
Å arg min
max
W(T TT mn)) (r))
(12)
TEllor r: Ax
[0.1]
Then, with probability at least
8, using the number of samples in Eq. it holds that
W N1,"", ,n²²) & 2c.
max r: -x8xAx[[1--(0.]]
7

---

Published as a conference paper at ICLR 2026
20
80
N 300
N 1000
N 10000
RS-BC
0.081=0.039
0.038:0.016
0.022=0.013
0.012=0.005
0.005=0.002
RS-KT
0.095=0.036
0.049=0.017
0.03:0.013
0.01940.007
0.011=0.006
BC
0.099=0.056
0.076+0.054
0.072+0.056
0.069#0.058
0.068+0.055
MT MIMIC-MD
0.127=0.062
0.086#0.055
0.07440.056
0.07=0.057
0.068
N = 20
= = 80
N = 300
N 1000
N = 10000
RS-BC
0.087=0.04
0.051=0.022
0.035+0.015
0.027=0.016
0.022+0.016
RS-KT
0.144+0.053
0.119#0.039
0.109=0.04
0. 0.108+0.038
0.106+0. (06+0.039
BC
0.103=0.057
0.08+0.053
0.072=0.056
0.069+0.058
0.068+0.056 058
MIMIC-MD
0.13240.065
0.09#0.055
0.076+0.055
0.071=0.057
0.06840.00 058
N = 20
N = 80
N = 300
N = 1000
N = 10000
RS-BC
0.10240.031
0.052+0.015
0.026#0.008
0.015+0.005
0.00440.001
RS-KT
0.11820.036
0.05920.017
0.031=0.0099
0.021+0.007
0.01:0.004
BC
0.085=0.035
0.04120.016
0.021=0.008
0.012+0.005
0.003=0.002
MIMIC-MD
0.132+0.052
0.06#0.022
0.03+0.01
0.016:0.006
0.005=0.002
N = 20
N = 80
N = 300
N = 1000
N = 10000
RS-BC
0.16940.00 0.79
1.76200079
0. 1.11
0.165+0.081
0. 0.16010.081
BC
0.168+0.077
0. 0.16600.778
0. 0..0000
0.177=0.091
0.174=0.093
?
0. 0.16940.049
0.08=0.018
0.043=0.01
0.02410.006
0.008=0.002
Table I: Results of the simulations described in Section The best results in each column are highlighted in bold. (Top) Simulation with A, H (2, 2, 5) for Q1. (Upper middle) Simulation with / = 0.5 for Q2. (Lower middle) Simulation with Markovian expert for Q3. (Bottom) Simulation with S, A, (300, 5,5) 5) for Q4.
In words, Theorem 5.2 establishes that the robust RDM problem is sample efficient for any algorithm that accurately solves Eq. (12.. Intuitively, such an algorithm can be viewed as an extension of
RS-KT the unknown-reward setting. Note that the minimization in Eq. (12) is over NTNM, since a satisfactory policy may no longer exist in the policy class IT(r)), ). which is also unknown due to the unknown reward. We leave future work the interesting question of whether an (approximate) solution to Eq. (12) can be computed efficiently, and note that restricting the optimization to a subset ll NI reduce computation time the cost some misspecification error.
NUMERICAL SIMULATIONS
In this section, we study RS-BC and RS-KT from a practical perspective by conducting simulations aimed at answering the following four questions
1. What is the performance improvement of RS-BC and RS-KT on the RDM problem compared standard II. algorithms?
2. How are the results affected by the choice of 0?
3. What happens if the expert's policy is Markovian?
4. Does RS-KT truly reduce sample complexity compared to RS-BC, as predicted by theory?
We select BC (Foster et al. 2024) and MIMI C-MD Rajjaaaan et al. 2020 2021) as baseline IL algorithms for, respectively, the no-interaction and woro--tranitionn settings, and address all questions by conducting various simulations following the next three-step process. (i) First, we randomly generate 50 MDPs and expert policies with fixed size S, A, H (details in Appendix F. (fi) Second, for each number of expert trajectories N E {20, 80, 300, 1000, 10000), we collect three different expert datasets of N trajectories for each MDP, provide them as input the four algorithms RS-BC, RS-KT, BC, and MIMIC-MD and record the average (over the three seeds) Wasserstein distance
Thhe code for running simulations can be found at risk
www.fww......m.... http/fillppppoo i
8

---

Published as a conference paper at ICLR 2026
between the expert's return distribution and the return distribution induced by each algorithm's output policy. (éi)) Finally, we average these distances across all 50 MDPs, obtaining, for each algorithm and value of N, number representing the expected error of that algorithm when given access to N expert trajectories in the considered setting.
Below, describe the specific simulations conducted and the results obtained for each question. The numerical results are reported in Table
Question 1 (Q1). To address Q1, we conducted three simulations with different problem sizes S, H e ((2,2, 5), (50, 5, 5), (2, 2, 20) , all with non-Markovian policies and 0 = 0.05. The results for (2, 2) are reported in Table (top), while the other two are shown in Tables 3 in Appendix E.6 Crucially, Table (top) reveals that BC and MIMIC CDD, by relying Markovian policies, are biased and cannot match the expert's ' return distribution satisfactorily. even with a large dataset of N = 10000 trajectories. In contrast, by leveraging our efficient non- Mar-Maakovinn policy class, RS-BC and RS-KT continue to reduce the error as the number of trajectories N increases, up to a limit determined by choice of f. A similar pattern is observed for larger S, A in Table and especially for larger horizons H in Table where the limited expressivity of Markovian policies becomes even more pronounced, We also note that larger S, A, H! slows down RS-KT significantly due to solving an LP with poly(5, A, H) variables and constraints, and increases the approximation error due to 0, as discussed in Appendix OopendixEE.6jar E.6 and in Q2.
Question 2 (Q2). To address Q2, we first conducted simulation with increased 0 = 0.5, keeping S, A, H = (2 5) and a non-Markovian expert. As shown in Table (upper middle), a larger consistently increases the approximation error, causing RS-BC and RS-KT to perform worse than with 0 = 0.05 (see Table (top). Nevertheless, while RS-BC remains fairly robust and continues outperform BC and MIMIC MD intuitively because it corresponds policy class -RS-KT tends to reach the worst-case approximation error predicted by Lemma H0/2 s 0.2, as also discussed in Appendix E.6 Additionally, we conducted three further simulations with values of 6 small enough to eliminate approximation error, observing a consistent increase in performance, particularly for RS-KT in settings with larger S, A, H (see Appendix FE.7.
with more expressive
Question 3 (Q3). We simulation with a Markovian expert and S, A, H (2, 2,),, 5), with results reported in Table [ (lower middle). Interestingly, BC outperforms all other algorithms in terms of both sample and computational efficiency. Intuitively, this occurs because operates on much smaller hypothesis space (i.e., 11 than RS-BC and RS-KT, and it is unbiased, since the expert is Markovian. The policy output by BC aims to match the expert's trajectory distribution (Foster et al. 2024), and thus also its return distribution for any reward, which is not the case for MIMIC-MD, explaining its comparatively worse performance.
Question 4 (Q4). According our theoretical results (Theorems and 4.4) RS-KT should outperform RS-BC in terms of sample complexity for large S, A. as its performance does not depend on them. To verify this, expert, and 0 = 0.05. To speed up computation (particularly the IP in RS-KT), we avoided running RS-KT directly and instead compared the expert's return distribution with the estimate î computed at Line [ By the triangle inequality, this guarantees that the error between the output of RS-KT and the expert most twice the error of f). Results are reported in Table I (bottom), showing dramatic improvement in sample complexity. In particular, both RS-BC and BC struggle even with N = 10000, while RS-KT achieves significant performance with as few as = 300 or 1000 trajectories.
simulation with large A, H (300, 5, 5),, non-Markovian
In summary, the key takeaways of this section are:
RS-BC and RS-KT generally outperform BC and MIMIC-MD due to the use of more expressive non-Markovian policies and reward information, particularly for large H.
RS-BC is faster than RS-KT because it does not require solving linear program and is more robust large 0, while RS-KT much more sample efficient for large MDPs.
BC performs well when the expert is Markovian, even without access reward information.

---

Published as a conference paper at ICLR 2026
7 CONCLUSION
In this paper, we introduced and analyzed RDM, a general formulation of (risk-sensitive) IL as the problem of matching the expert's return distribution. Remarkably, we showed that both the known- and unknown-reward settings are statistically tractable. For the known-reward case, we proposed two algorithms, RS-BC and RS-KT, which not only come with strong theoretical guarantees but also empirically outperform standard IL methods at accurately matching the expert's 's return distribution.
Limitations and future directions. This work has several limitations. We focus only on the tabular setting, and our theoretical analysis lacks lower bounds, so it remains unclear whether the upper bounds in Theorems 4.3 4.4 and are tight. Furthermore, the unknown-reward setting lacks a practical algorithm, and our empirical study does not include real-world data. Future work could extend our results to state-only feedback settings (Sun et al. 2019), develop practical and scalable versions of RS-BC and RS-KT for large continuous environments, and design algorithms for the Mkknown-reward setting.
ETHICS STATEMENT
This work aims to develop learning machines that replicate human behavior not only in task per- formance but also in risk preferences. While Al research has extensively explored improving performance, aligning machines with human risk attitudes remains underexplored and raises im- portant ethical considerations. We envision future where autonomous agents assist humans in everyday tasks, from driving to cooking. Because uncertainty is inherent to human life, agents must account for risk to behave safely and effectively. Embedding risk- sensisensttivtty is therefore essential for supporting human well-being, but it must be approached with care to minimize potential harm. Being foundational research, this work does pose direct societal risks. Nevertheless, recognize that ethical considerations will be crucial as future applications of risk-sensitive agents move closer to real-world deployment.
REPRODUCIBILITY STATEMENT
All theoretical results in this work accompanied by formal proofs in the appendix. In the main text, immediately before or after each theorem, lemma, or proposition, we provide a reference to the specific appendix section containing the formal proof, as well additional discussion on the implications of the result when present. For the numerical simulations, Appendix E contains all supplementary details not included in the main text, along with additional simulation results. We include with this submission zip file containing all code required to replicate the experiments, as well as .ny files with the exact simulation results. Following publication, we will make a publicly accessible repository containing all code and data, and we will reference it in the main text.
LLM USAGE
We used LLMs in paper writing only to fix potential grammar mistakes and improve clarity.
ACKNOWLEDGMENTTS
This publication was funded with the contribution of Ministero dell' Università e della ricerca pursuant to D.D. 7206 of 17 April 2025 BANDO FIS 2. Project "Unified Learning from Diverse Human Feedback" (HUml.m). CUP: D53C25000710001.
(Starting Grant), title:
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Prashanth L. A. and Michael Fu. Risk-sensitive reinorrmmen learning via policy gradient search, 2022. BRL/attss///axii wwtp:iww.orgaab///iiiee 09126
Pieter Abbeel and Andrew Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference Machine Learning 21 (ICML), 2004.
Brenna D. Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robolics and Autonomous Systems, 57:469-483, 2009.
Osbert Bastani, Jason Yecheng Ma, Estelle Shen, and Wanqiao Xu. Regret bounds for risk-sensitive reinforcement learning. In Advances Neural Information Processing Systems 35 (NeurIPS), pp. (6259-36266) 2022.
Mare G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcmmett Learning. MIT Press, 2023.
Julian Bernhard, Stefan Pollok, and Alois Knoll. Addressing inherent uncertainty: Risk-sensitive behavior generation for automated driving using distributional reinforcement learning. IEEE Intelligent Vehicles Symposium (IV), pp. 2148 2155, 2019.
Adam Block, Daniel Pfrommer, and Max Simchowitz. On the imitation of mon-markovian demon- strations: From low-level stability to high-level planning. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023. URL
Sergey G. Bobkov and Michel Ledoux. One-dimensional empirical measures, order statistics, and
kantorovich transport distances. Memoirs of the American Mathematical Sociery, 2019.
Kiante Brantley, Wen Sun, and Mikael Henaff. Disagreement-reguue reguulaazze imitation learning. In International Conference Learning Representations 8 (ICLR), 2020.
Nicole Bäuerle and Jonathan Ott. Markov Decision Processes with Average-Valuee Value-at- Risk criteria. Mathematical Methods of Operations Research, 74:361-379, 2011.
Nicole Bäuerle and Ulrich Rieder. More risk-sensitive markov decision processes. Mathematics of Operations Research, 39(1):105-120, 2014.
Yu Chen, Xiangcheng Zhang, Siwei Wang, and Longbo Huang. Provable risk-sensitive distributional reinforcement learning with general function approximation, 2024.
Yinlam Chow, Mohammad Ghavamzadeh, Lucas Janson, and Marco Pavone. Risk-constrained reinforcement learning with percentile risk criteria. Journal of Machine Learning Research (JMLR), 18:1-51, 2018.
Robert Dadashi, Leonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. In International Conference Learning Representations 9 (ICLR), 2021.
Yingjie Fei, Zhuoran Yang, Yudong Chen, Zhaoran Wang, and Qiaomin Xie. Risk-sensitive reinforee- ment learning: Near-optimal risk-sample tradeoff in regret. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 22384-22395, 2020.
Yingjie Fei, Zhuoran Yang. Yudong Chen, and Zhaoran Wang. Exponential bellman equation and improved regret bounds for risk-sensitive reinforcement learning. In Advances in Neural Information Processing Systems 34 (NeuriPS), pp. 20436-20446, 2021.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International Conference on Machine Learning 33 (ICML., volume 48, pp. 49-58, 2016.
Dylan J. Foster, Adam Block, and Dipendra Misra. Is behavior cloning all you need? understanding horizon in imitation learning. In Advances in Neural Information Processing Systems 37 (NeurIPS), pp. 100002-106666 2024.
11

---

Published as a conference paper at ICLR 2026
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforce- ment learning. In International Conference on Learning Representations 5 (ICLR), 2017.
Hans Föllmer and Alexander Schied. Stochastic Finance: An Introduction in Discrete Time. De Gruyter, 2016.
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. In Advances in Neural Information Processing Systems 34 (NeurlPS), pp. 4028 4039, 2021.
Sven Goluža, Tessa Bauman, Tomislav Kovačević, and Zvonko Kostanjčar, Imitation learning for financial applications. In MIPRO ICT and Electronics Convention 46 (MIPRO), pp. 1130-1135, 2023.
Dylan Hadfield-Menell, Smitha Milli, Pieter Abbeel, Stuart Russell, and Anca Dragan. Inverse reward design. In Advances in Neural Information Processing Systems 30 (NeurIPS), 2017.
William B. Haskell and Rahul Jain. A convex analytic approach risk-aware markov decision processes. SIAM Journal Control and Optimization, 53:1569-1598, 2015.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances Neural Information Processing Systems 29 (NeurIPS), 2016.
Ronald A. Howard and James Matheson. Risk-sensitive markov decision processes. Management
Science, 18(7):356-369 1972.
Emilie Kaufmann, Pierre Menard, Omar Darwiche Domingues, Anders Jonsson, Edouard Leurent, and Michal Valko. Adaptive reward-free cxploration. In International Conference Algorithmic Learning Theory 32 (ALT 2021), pp. 865-891, 2021.
J. Kiefer and J. Wolfowitz. Asymptotic minimax character of the sample distribution function for vector chance variables. The Annals of Mathematical Statistics, 30:463-489, 1959.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Diseriminator-actorrs Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference Learning Representations 7 (ICLR), 2019.
Jonathan Lacotte, Mohammad Ghavamzadeh, Yinlam Chow, and Marco Pavone. Risk-sensitive generative adversarial imitation learning. In International Conference on Artificial Intelligence and Statistics 22 (AISTATS), volume 89, pp. 2154-2163, 2019.
Romain Laroche and Remi Tachet Des Combes. On the occupancy measure of non-Markovian policies in continuous MDPs. In International Conference on Machine Learning 40 (ICML), volume 202, pp. 18548 18562, 2023.
Filippo Lazzati and Alberto Maria Metelli. Learning utilities from demonstrations in markov decision processes. International Conference Machine Learning 42 (ICML), 2025. URL
https:://ppenreviee net wtt ottp://eeeeeeeeeeeeooooooooooaaa
Filippo Lazzati, Mirco Mutti, and Alberto Maria Metelli. Offline inverse New solution concepts and provably efficient algorithms. In International Conference on Machine Learning 41 (ICML), 2024.
Luc Le Mero, Dewei Yi, Mehrdad Dianati, and Alexandros Mouzakitis. A survey on imitation learning techniques for end-to-end autonomous vehicles. IEEE Transactions on Intelligent Transportation Systems, 23:14128 14144, 2022.
Hao Liang and Zhi--Qan Luo. Bridging distributional and risk-sensitive reinforcement learning with provable regret bounds. Journal of Machine Learning Research, 25:1-56, 2024.
Anirudha Majumdar, Sumeet Singh, Ajay Mandlekar, and Marco Pavone. Risk-sensitive inverse reinforcement learning via coherent risk models. In Robotics: Science and Systems 13 (RSS), 2017.
12

---

Published as a conference paper at ICLR 2026
Shehryar Malik, Usman Anwar, Alireza Aghasi, and Ali Ahmed. Inverse constrained reinforcement learning. In International Conference on Machine Learning 38 (ICML), volume 139, pp. 7390- 7399, 2021.
Ajay Mandlekar, Fabio Ramos, Byron Boots, Silvio Savarese, Li Fei-Fei, Animesh Garg, and Dieter Fox. Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data. In IEEE International Conference on Robotics and Automation 37 (ICRA), pp. 1414-4420, 2020.
Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei- Fei, Silvio Savarese, Yuke Zhu, and Roberto Martin-Martin What matters in learning from offline human demonstrations for robot manipulation. In Conference on Robot Learning 5 (CoRL), volume 164, pp. 1678-1690, 2022.
Shie Mannor and John N. Tsitsiklis. Mean-variance optimization in markov decision processes. In International Conference on Machine Learning 28 (ICML), pp. 177-184, 2011.
P. Massart. The tight constant in the 18:1 1869-19,,,,, 1990.
inequality. The Annals of Probability,
Alberto Maria Metelli, Giorgia Ramponi, Alessandro Concetti, and Marcello Restelli. Provably efficient learning of transferable rewards. In International Conference on Machine Learning 38 (ICML), volume 139, pp. 7665-7676, 2021.
Alberto Maria Metelli, Filippo Lazzati, and Marcello Restelli. Towards theoretical understanding of inverse reinforcement learning. In International Conference on Machine Learning 40 (ICML), pp- 24555-24599 2023.
Aneri Muni, Esther Derman, Vincent Taboga, Pierre-Luc Bacon, and Erick Delage. What matters when modeling human behavior using imitation learning? In 2nd Workshop on Models of Human Feedback for Al Alignment, 2025. URL pottpss:/openreeeeee
net forum ? 2id=
9t@NFc99Th
Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning 17 (ICML 2000). pp. 663-670, 2000.
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J. Andrew Bagnell, Pieter Abbeel, and Jan Peters. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7:1-179, 2018.
Victor M. Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual Review of Statistics and Its Application, 6(1):405-433 431, 2019.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning. 11:355-607, 2019.
Dean A. Pomerleau. Alvinn: An autonomous land vehicle in neural network. In Advances in Neural Information Processing Systems (NeurIPS), 1988.
Martin Lee Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John
Wiley & Sons, Inc., 1994.
Aoyang Qin, Feng Gao, Qing Li, Song-Chun Zhu, and Sirui Xie. Learning non-markovian decision- making from state-only sequences. In Advances in Neural Information Processing Systems 36 (NeurIPS), pp. 6596-6618, 2023.
Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp- 2914-2924, 2020.
Nived Rajaraman, Yanjun Han, Lin F. Yang. Kannan Ramchandran, and Jiantao Jiao. Provably breaking the quadratic enror compounding barrier in imitation learning, optimally, 2021. URL
http://rrw...oor//a/////000..11 wwtp/a//////oooo 12948
13

---

Published as a conference paper at ICLR 2026
Lillian J. Ratliff and Eric Mazumdar. Inverse risk-sensitive reinforcement learning. IEEE Transactions on Automatic Control,
2020.
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. SQIL:
imitation learning via reinforcement learning with sparse rewards. International Conference on Learning Representations (ICLR), 2020.
Allen Z. Ren, Justin Lidard, Lars Lien Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. In International Conference Learning Representations 13 (ICLR), 2025.
R. Tyrrell Rockafellar and Stanislav Uryasev. Optimization of conditional value-at risk. Journal of Risk, 3:21-41, 2000.
Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. In International Conference on Artificial Intelligence and Statistics 13 (AISTATS), volume 9, pp. 661-668, 2010.
Anirban Santara, Abhishek Naik, Balaraman Ravindran, Dipankar Das, Dheevatsa Mudigere, Sasikanth Avancha, and Bharat Kaul. Rail: Risk-averse imitation learning. In International Conference on Autonomous Agents and MultiAgent Systems /7 (AAMAS), pp- 2062-2063, 2018
Wen Sun, Anirudh Vemula, Byron Boots, and Drew Bagnell. Provably efficient imitation learning from observation alone. In International Conference Machine Learning 36 (ICML), volume 97, pp. 6036-6045, 2019.
Umar Syed, Michael Bowling, and Robert E. Schapire. Apprenticeship learning using linear pro- gramming. In International Conference on Machine Learning 25 (ICML), pp. 1032-1001 1029,039,000 2008.
Marcel Torne, Andy Tang, Yuejiang Liu, and Chelsea Finn. Learning long-context diffusion policies via past-token prediction, 2025.
09561
Luca Viano, Stratis Skoulakis, and Volkan Cevher. Imitation learning in discounted linear mdps without exploration assumptions, 2024. URL httt:s:://axxl..org/bbbbba 2405 02181
Cédric Villani. Optimal Transport: Old and New. Springer Berlin, Heidelberg, 2008.
Kaiwen Wang, Nathan Kallus, and Wen Sun. Sear-minimax-optinn optimal risk-sensitive reinforcement learning with CVaR. In International Conference on Machine Learning 40 (ICML), volume 202, pp. 35864-35907, 2023.
Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy Finetuning: Bridging sample-efficient offline and online reinforcement learning. In Advances in Neural Information Processing Systems 34 (NeurIPS 2021), pp.
2021.
Tian Xu, Ziniu Li, Yang Yu, and Zhi-Quan Luo. Provably efficient adversarial imitation learning with unknown transitions. In Conference on Uncertainty in Artificial Intelligence 39 (UAI), volume 216, pp. 2367-2378, 2023.
Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual ma- nipulation with low-cost hardware, 2023. URL
org. wttpg//b///////////////// 13705
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang. Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang. Zikang Liu, Peiyu Liu, Jian- Yun Nie, and Ji-Rong Wen. A survey of large language models, 2025. URL
2303. 18223
Brian D. Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy, 2010.
14

---

Published as a conference paper at ICLR 2026
A ADDITIONAL RELATED WORK
Standard II... There are two main approaches to address the standard occupancy measure matching formulation of II: Behavior Cloning (BC), that treats II. as supervised learning problem, directly
learning a mapping from states to actions (Pomerleau 1988 Ross & Bagnell 2010), and Inverse Reinforcement Learning (IRL) that infers reward function that reflects the expert's 's preferences. and then derives a policy via planning (Ng & Russell, 2000 Abbeel & Ng 2004] Ziebart 2010; Finn
et al. 2016 Fu et al. 2017). Other popular occupancy measure matching methods include Ho & Ermon (2016); (2016) Kostrikov et al. (2019): Reedy et al. (2020); (rantley et al. (2020); Garg et al. (2021);
Dadashi et al. (2021). Recently, there have been various efforts into providing theoretical guarantees on BC (Rajaraman et al. [2000] Foster et al., (024) and IRL (Metelli et al.) 2021: 2023), and also for
related IL algorithms (Sun al. [ Viano al. 2024). In this paper we provide a theoretical
analysis analogous to that in Rajaraman et al. (2020) Foster et al. (2024), but for the novel return
distribution matching setting.
Risk-sensitive RL. The first paper to address risk sensitivity in Markov Decision Processes (MDPs) is Howard & Matheson (1922). Since then, various researchers have explored this problem (see the survey A. & Fu (2022) and the recent book [Bellemare et al. (2023)). Notably, Mannor & Tsitsiklis (2011) emphasizes the importance of non-Markovian policies in mean- variance optimization. Additionally. ly, Bäuerle & Onl ( 2011) and [aäeele & Rieder (2014) show that optimal behavior in CVaR and expected utility planning problems can generally be achieved using non-Markovian policies that base actions on both the current state and the cumulative reward. This idea is exploited in Section to construct our algorithms. Some risk-sensitive RL algorithms employing such Markovinv policies include Haskell & Jain (2015); Chow et al. (2018). In addition, we mention Fei et al. (2020) 2021),
which study the risk risk-sensitive RL problem from a theoretical perspective, by providing provably efficient risk-sensitive RL algorithms under the entropic risk measure. [Bastani et al.] (2022); Wang
al. (2023) provide analogous regret analysis for the CVaR objective. Lastly, recent Chen et al.
(2024) Liang & Luo (2024) consider a theoretical distributional RL analysis for risk-sensitive RL.
Risk-Sensitive II. Some works extend occupancy measure matching with additional objectives to capture the expert's risk attitude. [Santara et af. (2018) and Lacotte et al. (2019) are most similar to ours, as they aim to match both expected return and CVaR. However, as noted in the introduction, they restrict to Markovian policies, which limits their ability to fully capture the expert's risk preferences. |Rathff& & Mazumdar (2020) propose a risk-sensitive IRL algorithm focused learning risk parameters. also assuming a Markovian expert. MMjumdaa et al. (2017) study risk-sensitive IL with form of Mor-korko policy, but in environments far simpler than tabular MDPs. Lazzati & Metelli (2025) consider non-Markovian experts and imitation, but primarily aim to recover the expert's 's utility, and their II. approach struggles when demonstrations come from a single environment. Nevertheless, our policy class in Section 4 draws inspiration their discretization approach. Muni cct al. (2025) also mentions the importance of non-Markovian policies for lisk-sensitive IL, but do not present any algorithm.
Non-Markovian II. The importance of adopting non-Markovian policies for IL has already been recognized by some IL works, especially in the field of robotics. Mandlekar et al. (2020) 2022) identify one cause of non- Non-Markovianity in human demonstrations in partial observability. and consider variants of BC with recurrent neural networks for imitation. Zhao et al.) (2023) and subsequent literature (e-g., see Torne et al. (025);; Ren et al. (2025)) p) learn Mor-Markovian policies implicitly through action chunking, an open-loop control technique in which at each state our policy outputs a "chunk" (i.c., a sequence) of actions. We mention Qin et al. (2023), which learn non- Markovian policies as energy-based priors from state-only sequences, and Block et al. (2023), which study imitation in non- con-linear systems. Importantly, none of these works address Mon-Markovian policies arising from risk sensitivity, which is novel aspect of ours. Moreover, note that fitting general non-Mann non-Markovian policies may require an amount of data exponential in the horizon in the worst-case (see Appendix B.3
15

---

Published as a conference paper at ICLR 2026
B ADDITIONAL RESULTS AND PROOFS FOR SECTION 3
In this appendix, we first show that matching the return distribution in Wasserstein distance is strietly more expressive than matching the expectation the CVaR a given level (Appendix [B.Ip. Next, we provide the proof of other results in Section (Appendix [.2).. Finally, we prove that matching the trajectory distribution of arbitray non-Markovian policies sample inefficient (Appendix B.3
B. ADDITIONAL INSIGHTS RDM
We show here that matching the expert's s return distribution in Wasserstein distance implies eloseness terms of expected return, the variance the return, and the CVaR any level.
First, observe that, for any MDP M, and policies Th 9:
- Wnn 11,5)
This follows after having realized that the identity function (denoted below as I(-) is 1ee11-Lpppn -ipischhiz and using the dual form of the Wasserstein distance (see Eq. 6.3 of Villani 2008)).
J
J
E x
[ By - [Y]
X
E
[(X)] - Ey B [T(Y]
sup
Ex Ex [f(X]] - By [f(Y)]
W(n 17.))
Second, for any E (0, 1), we have:
(CVaRa CVaR. (Te W( n.
a This follows using an alternative expression for the Wasserstein distance. Formally, for any pair of distributions p. g with support on [0, H]:
1
(VVaRa CVaR. (1)
FØ (x) (x) dx
J0
1 a
F (x) - (x) dr
J0
Fx)
Fx)(2)) dr
0
W(p.g),
where (1) use that the 1-Wasserstein distance can be written this way (Panaretos & Zemel 2019), where recall that (x) == inf (z)=2 z.
Lastly, observe that closedness in Wasserstein distance also implies closedness between the variance of returns:
\var(n var(n,e)] & 4H1 W (1,5 nAE)
To see it, for any pair of distributions p, with support on [0, HJ. we can write:
|var(p) var(g))
Xx= (Ex=p Er [Y3] (EEr~g (Y])²
<
Bx [X2]
Ey Y2] + (Ex~p (x]] (By (F)))
(1)
2H1 sup
EY [f(x)]
(Ex=plX)) [X] Ey Bx]) (Ex=p(X) [X] Ry-g(D])
[
16

---

Published as a conference paper at ICLR 2026
21 W(p, q) + 211 [xxp[X] By ~ |
& 4H W),,)),
where at (1) we use that the function f x - x²2 2H Lipschitz on [0, H] (since (((w))/(M) supx+y ž- s (I 2) 2H for æ, y [0, H)), and we rescaled obtain functions that are Lipschitz, (2) use the dual form of the Wasserstein distance (see Eq. 6.3 of Villani (2008) after having noticed that the absolute value can be removed as set {f I 1 symmetric, and we upper bounded E [x-][X]] + By-ggi " 2H. Finally, at (3)
|
supxAy |r-y]
we use the result proved earlier.
B.2 PROOFS
Theorem 3.1. There exist an MDP M,, with A, H 3 2 and an expert policy n IIAM such that, even with N (S 11 trajectories, any algorithm 2 satisfies
TV x n E DE TV
3
2c²
where āे is the oulput of a given in input M, and DE
Proof. To prove this result, we simply provide a reward function r S x A x [H] ~ [0, 1] for which, in the "hard" MDP used in the proof of Theorem B.I have:
TV
for any pair of policies. Then, the claim in the theorem follows directly from the result Theorem B.I
First, let us associate a different integer 0, I,..., S to each state {s1...., ,83,. and different integer 0, 1.....,A. A - 1 to each action (a1,.... a.)). and denote v, S - N and A N these mappings. Then, the reward that we provide is the following:
E (s. a) == 10 Th
Simply, observe that every triple 8, U, / is associated a reward value belonging to a different power of 10, thus when sum them to compute the return of trajectory obtain different return value for every possible trajectory:
TV The
11,5)
= 2
1
Σ
Nre (g)
n. (0)
2 Σ
Σ
( (w) Pa (w)
w:Glwir m
Σ
(w) p (w)
pa
pal
This concludes the proof.
Proposition 3.2. There exist an MDP M,. with horizon H 3 and an expert policy z. IIM such that any Markovian policy π E II satisfies
W( E W 7 2 0.5.
Proof. Consider the MDP M,E (S, p. r) in Fig. where S {sinit, 81, 82, 83).
A {a), 02), Hl 3, 30 = Sinit, the transition model p such that:
P1 (81 Sinit; a) (82 Sinit; a) 1/2 Va A,
17

---

Published as a conference paper at ICLR 2026
$1
01, 02,02, P
0
1/2
01 02.
start
Sinit
83
1/2
02.7 2
$2
Figure 1: MDP for the proof of Proposition 3.2
P2 $1
=1 Va A.
Ps( (3 (353333,,0 1.
and the reward function r is defined as:
T (Sinis a) -0 Va E A,
12 (81, a) = Va E A, E
1 (82, a) = Va E A,
E 1 (83,a1) O,
r3 (83, 02) 1.
Let R be the deterministic non-Markovian policy that plays always action a every time it is not in 83, and then in 83 plays a1 previously we passed through 81, otherwise play 02. Formally, for any E Ω and S:
क (ails,w) = {
if / 83 (s otherwise
83 A (Smit: 01 81, 01))
It is easy to note that the return distribution of π in M,E is:
0h 81,
where notation & denotes the Dirac delta on æ. With 81 we always get return
Now, let Th be any Markovian policy parameterized by [0, as (we do not specify values in other states and stages, because they are not relevant for the return distribution):
Tg (sa)
The return distribution of n M is:
2 50
82.
namely, irrespective the policy x", the return distribution is 1 w.p. 1/2. Let us compute the Wasserstein distance between these return distributions:
F W
11e
=
F E (g) T
F
(g) dg
0
1 +
+ 1
2
irrespective of a. This concludes the proof.
18

---

Published as a conference paper at ICLR 2026
81
81
81
82
82
start
KK
135
$5
85
89
Figure 2: The MDP R M used the proof of Theorem B.1
B.3 STATISTICAL INEFFICIENCY OF GENERAL NON-MARKOVIAN POLICIES
In this appendix, we provide an explicit proof that the problem of matching the distribution over trajectories of an arbitrary non-Markovian expert's 's policies may require an exponential (in the horizon) amount of data. Then, the proof of Theorem [3. can be simply obtained through a reduction to this problem. The lower bound is provided in the next theorem, which makes use of the family of hard instances in Fig. 2
Theorem B.1. There exist MDP.RM (S,A, H0, p) with S > 2, A 2 2, H 2 and non-Markovian expert's 's policy T E 11 such that any learning algorithm 21 taking in input dataset of N = (S- 1)"-1- expert's trajectories DE [wi]ie[n] satisfies:
E E
Tp
where रे is the policy outputted by A when taking in input both DE and M.
Proof. The proof draws inspiration from that of Theorem 5.1 in Rajaraman et al. (2020) for the known-transition setting.
We begin by describing the class of hard instances that will be considered for proving this lower bound. Note that we are considering here the "II." problem of computing policy with trajectory distribution close to that of the expert. Thus, problem instance is pair MDP, R -expert's policy, Let ? = {(M, क.. x be the family of problem instances where the MDP R M = (S, A, H, s0 p) represented in Fig. 2 and the expert's policy n. any deterministic policy that for < always plays 01, while at the last stage play an arbitrary action. Formally, S 2 {81,85) (i.e., it has > 2 at least 81, 8g, and potentially other states 82, A [a1,02), > 2, 80 81, and transition model p is defined as a function of scalar < E [0, 3 that will choose later:
S pn(s'(s,a) = - (S (S- 1)ç
if # sg s' #88 if + ss A 5g otherwise
for any action a A. Then, the expert's policy nh is defined as any policy in Ihardd defined as:
Ihhard_
T NM th [H 1], vs S, P4 <a)lss, = 1).
Then, B is formally defined as:
$ = (M' M = M [
We set Ç (N+1)= and observe that, to guarantee that M exists, we need to enforce that
(S - (S-1)Ç, i.e., the total probability assigned to reaching state * sg. is smaller than 1. So:
S
(S-1)
N (S- 1)-11
1+
19

---

Published as a conference paper at ICLR 2026
To proceed, we choose distribution P over problem instances in Z. We select Pe A as the uniform probability distribution over the family of problems R, ie.:
x(M E : P(M',
P3
Then, we can show that:
(S 1)"
Ei (²) Dll pat 1p= - 1e
e(N + 1)
then we can conclude that there is at least problem instance (M E & such that:
(8-1) H-1 e(N + 1)
EDE-P+ 1TT
1pē|
from which if we insert the choice N (S - 1)11 1, we obtain the claim of the theorem. To this aim, we can write:
E -
p pa
Σ
Ixx (w) 1p (w)
wellH
ED ((M
pa'
(w) (w)(
pa (w)
wE
En (M's
Σ
p (s',aa, s ""
sll atl
- -s2 a
# )
ED (M (M
Σ
Σ
pxt (s1, (s,.,,,,,......,,,,,,
aEA
- p=18,,,01,1 82 A1,
a)
ED (M))
Σ
Σ
p(s2) ....) H)
m aed
<'(as11,01, 01, 01
pat (81, 01, 22 00 .......,, 8",a))
EM
Σ
Σ
p(02)-
p(s)) p(s
DEA
<'(als1 0
p(s²)
p(s
#(c|s1,
a 81, a1, 82 a1 st))
he[H-1]
Σ
p(s2) (()))p(s²) p(5)) .... M) Σ
"1
aEA
(alss,, 8 ai.................... " - a1, 2 al
s " .K
EEiin
Σ mls?) p(s²)
p(s B(M ((D))
Σ
2,,
aEd
n' a1 2 - f[(als al, a al.....,,
" K
EDEng Σ
pls²) p(s)) (1 a1 s² a " 8 DE}
B(M(A))-((D))
Σ
(alls an, 52 a,,,,nn...................... "a.....,,
and
20

---

Published as a conference paper at ICLR 2026
- ",,
a,,
x|
(
D u
Σ
p(89) p(8)). (( 01, s²,a a "......,.
)) $ DE)
62,5 pm!!!!!!!!!!
Σ
1aa= = bj - - 01 8 ........ , K
A
AEA
asA
E EDang
Σ
pls²) p(s))
pls ( ) [[()), 01 s² ,.. DC)
$2,09 p
7(b/811,1111 al, 8² an, ,., K)
A
A
Σ
A(als), 3 ,a,,....... " K
+
dEA((b)
-DE=_g q
Σ
pls3) p(s)) ... p(s") an, 2t ",.....,s")) DC}
s H
K
1 + A
a1 .1.......... s")
PaA
(E)
R(b|s1, (1b a1, 8², a1 .a......s") ......,")))
EDE-q Σ
pls²) pls²) ... an, st,,,................... ",.....,s") DC)
(1
K
Σ
(1- (37(b)81, a1, x² an,
A
BEA
EDE-g
Σ
pls3) p(s))) ... 1(
a1 22 a.. $ DD)
(1 + K 2K/A)
ED q
Σ
pls2) p(s))(().... ((... a1, 2² s
, ) $
DE}
H
p(s²) p(35)
p(s (( a1, 3 . $ D
Σ
p(s²) p(s))) .........$$. p(st)(( (1
p(s)) (1-1(3))
p(03)
p(att))
p(s?) p(s2) ... p(s (1 1 - pl(3) -(s3) ..... x(s)))
(2e..e
1 1 N
N N+1 1
Σ
N+1
(10) 1
Σ
N+1
te (H.... H)
(S - 1)H-1
N+1
where at (1) we define q E A Q 11 as a distribution that generates datasets D as if we sampled first (M", x²) P. and then we collected N trajectories from x"" in M", and we define (DD as the
21

---

Published as a conference paper at ICLR 2026
uniform distribution over the set of deterministic policies:
ITmimii (DC)= == TT IIhard E [H - 11. Vw E (M(D)) r(w) = T TDe (w)).
where 22 (DF) denotes the set of trajectories (#1, (11; d1........... 8/-1: (, 8h) visited in some trajectory
in DE, and T (w) denotes the corresponding action present DE, and T(w) denotes the determin- istic action taken by π. At (2) we lower bound by considering the error only for the trajectories in which at all h < H the action played is d1, and we note that in .M' = M, the initial state always $1- At (3) we define distribution p A over states as p(s) < Ç * ss, and p(ss) = 1- (S 1)Ç. 1)ç. Note that since any policy x" E IDbard plays action di for < H. then we set this probability to 1. At (4) we rewrite also pa using P and chain rule of conditional probabilities. At (5) we bring the p terms outside and define K = [[ 7(a1 $1, Q1. s², $².00 ".a,..... s") for brevity. At (6) we use that P'(DF) gives always M' .M, and that it gives, when the trajectory is not observed (i.e., is not in DR,, with equal weight, policies that play any action E A given trajectory $1, a1, s², .i......" s". Thus, since there is no dependence on the actions assigned by such policies given different trajec- tories, then the expectation simplifies to just E A (...). Then, note that, using this notation, a'(a|s1, 01 x2, 01.--.. ........ 5H)) 1aa- b). At (7) use that K is product of probabilities, thus it lies in [0, 1], and so, since for A 2 we have 1 - 2/A = (A - 2)/A 3 O, then the quantity is lower bounded by using K O. At (8) we bring the expectation inside and use the fact that the expectation of the indicator is the probability. At (9) we realize that the probability of a trajectory does not depend on the specific policy (inside I that generated it. Thus, we can replace q using p. and the complementary of an event of a Binomial distribution. At (10) we proceed as in Lemma A.21 of Rajaraman al. (2020) by lower bounding the exponential term by 1/e.
C ADDITIONAL RESULTS AND PROOFS FOR SECTION
In Appendix C we collect additional results and proofs for Section while in Appendices
Lastly, in Appendices C.4 C.5 we RS-BC and RS-KT result the settings in which, respectively, r is known to belong to a given finite set R and it is known to be linear in known feature map ø. Finally, in Appendix C.6 we sketch how to extend RS-BC arbitrary IL problems.
C.2
we collect those for, respectively, Sections 4.2 4.3
extend our results
C.1 ON THE CLASS OF POLICIES IN SECTION 4.1
In Appendix C.I.I we report the proof of Lemmas and 4.2 while in Appendix C.1.2 we discuss the size of the policy class II)
C.1.1 PROOFS
Lemma 4.1. Let M, be any MDP and let T E MM be any policy. Then, the policy nT E II(rF))
satisfies 7. (g) n (g) for all g e [o, H].
Proof. Thanks to Lemma C.2 know that:
P (GH g SH = p (GH-gΛSH
Vge [0, H es,
where GH = - F. (SN, ax) denotes the random return at stage H. Then, for any g' E H).
can write:
H_R (g') 1
TT. (G (H(s an) gg)
Σ
Σ
P (GH = 9H s, s.r aH) = =g' -g)
gE9,r, R.
(2) Σ
Σ
Σ
FT. (GH = SH = S) " (a (als.g)
ge9.E F S68 aE (1.
Σ
Σ
Σ
ps (Gn g, = 3) (a|s,9)
g89,8,M . S ae -3
22

---

Published as a conference paper at ICLR 2026
Σ
Σ
p (G#r 9:80=11110 =8))
Σ
(a)())))
g E, M SES
e (s.0)==y'-
P (G = 9. = 5,00 0)
Σ
Σ
(GH 9.8H =
Σ
P (G 9,8H
959,E N SES
o (=,0)==''9
pot E (GH 9. = s0s8 a) a) = gg)
ge9, "1 #ES AEA
Σ
E (sw aw) = oo)
(g'),
where at (1) define symbol Gr.H == (g [0, 1] E ΩH G(w;r g). (2) recognize that, by definition, R,E takes actions only depending on the current state, stage and past rewards, and we denote with brevity this fact with л, (alaa,, g), at (3) we use the aforementioned result from Lemma C.2 at (4) we use the definition of π. (a) 8, g) where the denominator is not O. noting that, in that case, the entire expression evaluates
Lemma 4.2. Let 0 (0, Let M, be any MDP and n E TT any policy. Then, the policy R. aB) Tre E IT(r8) satisfies W(1, Mpp & HO.
Proof. We can write:
J W
(1)
< W
7,5
The
TH + W
mr Wig
THE =²)
+ W
1h5 n,
(2)
- -rgllo + ww a
<
2/|
mrr
F x.) < HO W
n1 N.
2 HO,
where (1) we apply twice the triangle's 's inequality, at (2) we apply twice Lemma C at (3) we realize that, by definition of r it holds that r 0 cc & 0/2, and finally, at (4), we apply Lemma 4.1 with reward r and expert's 'seetts policy F
Lemma C.1. Let M be any MDP R and let n I be any policy, Leir Ax [0, be any pair of reward functions. functioss Then, it holds that:
W(n 1p2e)) & H'r 1²2 32100-
Proof. We can write:
(Tom Sdi
S
W(n w
sup Mlllus
Sdutfa)
J0.m
sup
Σ
f(g)n
Σ f(e)n)(())
flfles
gesupp(n;1)
gEsupp(1)2)
(3) sup
Σ
f(0)
Σ
(T(W)
(esupp(n;,)
w
G(wrr )=
Σ
f(g)
Σ
p (w)
geaupp(n)a)
w
G(wr?)=
Σ
((G(w;"'))PT
sup
welt
23

---

Published as a conference paper at ICLR 2026
Σ
we
Pp (w) (PT(W) [f(G(w))))) [ f(G((:r
Σ
sup
f:lf wellir
Σ sup /:1flue wEL
PT
(w:))) - G(w;r
w
Σ
PT
(w)
Σ
(riss,, an) - (sh an)
WERr++
kE[H]
Σ
P (w)
Σ
r an)
(sh; an)
WE
WE [H]
Σ
P
(w)
Σ
max (s.a)ESxA
a) -1)(s,))
+
NE[H]
(1(s,a) - -r)(s,a)-
Σ (n)(w)
max (s.0)ESxA
he[l]]
wEBn++
Σ
()(3sa) - r))
max (s.a)ES x.A
he he[H]
11
max [5xx.))
(l(s, - r)
where at (1) we apply the duality formula for the Wasserstein distance (see Eq. 6.3 in Section 6 of Villani 2008) ). Note that interpret the return distributions 77, and 17,5, probability measures, and that a function f : R - R is L-Lipschitz (i.e., |f L) if, for all x. E R, we have [f(x) f(y] < D|r yl- At (2) realize that 17,1 and 7,22 have finite supports, since in a tabular MDP with deterministic reward the number of trajectories is finite, and, thus, also the number of corresponding returns must be finite. We denote by supp(.) the support of distribution. At (3) use the definition of return distributions, at (4) we use the definition of Lipschitz functions, at (5) we denote by ($1, Oi...., SH, ar) the state-action trajectory w e QH and we use the definition of operators G(::r" r(::")) and G)):r").
Lemma C.2. Let My be any MDP and let IT be any expert's policy. Let n. II(()) be the policy defined as in Eq. (6) for expert's 's policy æ. Then, for all h [H]. s and g e [0, 11, holds that:
P
(M aw) = sn = =) =P P
Thrss, = sh =)
W=1
Proof. For simplicity, let us denote by G5 the random return at time step h under reward r. Formally:
G₁ = Σ (w(ss,ak) Wh e [[]]
In this way, the claim of the theorem be easily rewritten as:
psr (Gss = g 5p = =) = 1aT GG. - g 8h =s)
We prove the result by induction. Let us begin with the base case: h = 1. For all se S and (0). we have:
Td. (G = -9 =1ig 0}1{s su) pm (c: -gs1=s))
24

---

Published as a conference paper at ICLR 2026
where noticed that, for h 1, no action is taken yet. Now, let consider any stage h (2.3..... H. and let us make the induction hypothesis that, for all h [ - 1II for all e S and ge [0, h' 11, it holds that:
Tpa (Gw = 9 Si = 5) PM (Gw A 8w = 8).
Then, for any 8 E S and g' e 10, h- 17, we can write:
1p (Gn = g 3h
(
E
WA - 8h= = =6)
WERt w
((
Σ
Σ
Σ
(T - 84-1 04 a 8h = 5))
gEgr.h-1 wEfh-ii
(aA))
G(wir) -rr, 165.a)-9'-9
Σ
Σ
gEgr, w (x.
G(wir)= (s-1(=;))= g
-p (an-1 alw, I
(5h s,)
Σ
Σ
Σ
Ps - 5h-1=
gEGr,A-1 wEf 1: (s,a)ESxA
G(w:r)= 1h-1(s.a)-g-
1p (mr(ah- (ah-1
alw, 8) 1 (8,0)
Σ
Σ
Σ
("r(wh-1 = W 8A-11=
welt 1: (x.))SSAA: GE.r.
G(wir)= 9 Th-1 (5,)(s;)=-'-
(alw, s) (8' (s(() a)
(6)
Σ
Σ ppr (pr(w-- = -
gEGr.
(s,a) AA: wein-1: The ((s,)) -g G(wrr)-g
(alw,, 1)
Σ (s.a)ESxA:
(r(alg, s)
(wh-1= = 1
gE
weil-1:
1h-1(11a-9'-9
G(wir)=g
Σ
ge
(s,a)ESxA:
Ph-1(5a))gg-g
(8)
Σ
Σ
(10 8h =
(s.4)ESxA:
Σ
Pst
Pa Pa g 86-1=8)
g-
(s,4)ESxA:
1r-1(sa)--g-g
Ph
Σ
Σ
PT 8h ah-1 ())p-1(5(1s,a)
geGr,h
(s.0)ESxA:
(n-1(sa)=g'-g
Σ
Σ
1
gEgr.h
(s.a)ESXA:
144,))-'-
Σ
1PT (G = - (h-1(9h-1100 =g' -2--- 8h 5')
gEG
25

---

Published as a conference paper at ICLR 2026
=P =G = 8/ s'),
where at (1) we use symbol wh" to denote the random trajectory long h" stages, i..., whose realizations belong to Qn,, for any h" E [H]] At (2) we define symbols Gr,b == [g E [0, h - 1] 3w Ω2 G(w;r) = 9} for any E [H]. denoting the set of possible values of cumulative reward obtainable at stage h, and we sum over all such values (note that they are finite in tabular MDPs with deterministic rewards), at (3) we use the chain rule of conditional probabilities, (4) we use the Markovianity of the environment, at (5) we note that 1PT (ah-1 alw, s) actually is TT- [aw,,, s), at (6) we exchange the two summations, at (7) we recognize that, by definition, म, m.-a|a: (w, 8) takes on the same value for all the trajectories with the same value of return, and thus we can bring this quantity outside the summation over the w. We use symbol गr (a|g. s) to denote this fact for brevity. We do the same also for Ph-1 (8'l [8, a). At (8) we the induction hypothesis, at (9) we replace (rlag 8) with its definition when Pat (Gh-1 9 A 8h-1 8) > 0 as in the opposite case the entire formula takes on value zero
C.1.2 II IS TOO LARGE FOR SOME y
Consider any reward n. : S A x H - [0, 1] that assigns, every possible trajectory w Ω of this kind exist as the finite, while the set of possible return values that we can assign to each 8, a, h is infinite (continuous) [0, 1j. An example of a reward of this kind was provided in the proof of Theorem 3.1
a different return value G(w;r E [0. H - 17. Observe that rewards p
number of possible trajectories [2]] O((SA)=-1)
(s, a) 10
V(s, h) E S x A [W],
where v. S N, Fa A - N are arbitrary injective functions mapping the set of states and actions to the set of natural numbers. Intuitively, given the return G(w;; R of any trajectory:
(s1, a1 82, ag........ Sh-1; On-1),
we can easily reconstruct which s, a, actually belong to w by checking which decimal numbers in G((ur are "flagged".
Given a reward E of this kind, have that [G.e] in the same order as |Ω] = (1()) and so the set of policies E) defined in Section 1.1 reduces to the whole set of non-Markovian policies IINM, with all its disadvantages.
C.2 No INTERACTION SETTING
In Appendix C.2.1 on Theorem
we prove Theorem 4.3 while in Appendix C.2.2 we provide additional discussion
C.222
C.2.1 PROOF OF THEOREM 4.3
Theorem 4.3. Lete (0, H) and e (0. 1). Let M,, be any MDP and let T E IM be any policy. Then, choosing f = c4HH, with probability at least 1 - 8, the policy āे output by Algorithm satisfies W(n n & c, with number of samples:
SH In
(A (++ In 1))
N δ
(8)
Proof. We can write:
= + w(
ThE The
w W F J nHe nTE
+ W
77
J.f + w
J *
+ W
nT 175
2HO W
+
2HO + H 11.E
2H0 + HI Σ Σ
Σ
THE (Gr g gg =3) - - (-1.91)
hE[H]
9 B. S6S
26

---

Published as a conference paper at ICLR 2026
2HO + H
Σ Σ
Σ
p (Gh = g A.s. sn = s. (1(19,))
he[H]] gEg E eES
2H0 Hé'
where at (1) we apply triangle's inequality, at (2) we apply Lemma 4.2 and Lemma C.1 twice, at (3) use Particular Case 6.13 of Villani (2008), which tells us that can upper bound the Wasserstein distance between two distributions supported on set X' by the diameter of X (max,, x - x') times the one norm between the two distributions. Since X' [o. H in case, get the expression written above. At (4) we apply Lemma C.3 with the notation defined in that lemma. observing that policies 7. and म satisfy the hypothesis for reward r At (5) we use Lemma C.2 observing that the random return GA is in terms of reward r and recalling the definition of T. Lastly, at (6) apply Lemma C.4 with accuracy e'.
The result follows by imposing that & 277 and 2H10 & 5. which can be achieved by taking = 27F and 0 4f7 and by observing that:
Σ
19
H
Σ
[
he[#]
Σ 110,0,20 20, .(.. 1)/0101
h=[#]
Σ (1+ (h- 1)/0)
hE[H]
H(1 1/0)
Σ h
he[H]
H(H + 1)
H(1 1/0) +
20
H
(1) 4H3
where (10) used the formula for arithmetic sums, and at (11) used the previous choice
0 = 3F
Replacing into the number of samples in Lemma C.4 (and also é'
c
we get the result:
2H
N 3 In 8SH²
8SH
+ (A 1) 2048eSH In 8SH &
In
dc
By using Ô notation to hide logarithmic terms in A, H, 1 In 3. we get the result.
Lemma C.3 (Error Propagation). Let Mr be any MDP. For any pair of policies n, n' II such that, all h E [H] a E A, sE S and w, Or with G(w;r) = G(wi;r):
x(alws s) x(alw',s)
<'(lw.s) ='(alu',
holds that:
Σ P (G,, = gg 5) Th (T(-1s.g) Th(-1s.9))
1r
H gEg..A
where Gr.h := {g E [0, 1]
Ωn G(w;r) - g) for all h E [ + 1B. Gh. := Em=1 Tw(Sw, aw) denotes the random return stage h. and Tu (n g) and n (-1s,, g) denote the unique probability with which the policies ा and x' prescribe actions at under any trajectory Stp with G(w;r) g.
27

---

Published as a conference paper at ICLR 2026
Proof. To prove the result, first demonstrate by induction that, for all h E [2, H]] it holds that:
Σ
Σ
ps
(Gn = 8n = pm (Gn = g 10 =
gEgr,N
E Σ
Σ
g A S/ -
We[h-1 11 gE9,,, SE8
We remark that, by hypothesis, both π and n' prescribe the same actions when faced with trajectories w. w' with the same return G(w;r) = G(w'; r)) that are followed by the same state s.
We begin with the base case
We can write:
Σ
Σ
PT (G2 = g 82 = 5') - pa (G2 - 82 = 8')
gE9r,2 d'ES
Σ
Σ
pst (ri((11, = g 82 8) p (11($1,01)= = ga 82 5))
gea.,2 s'eS
Σ
Σ
Σ
(P* = = = as
geg,,a s'eS
(5,a)e5xA: ri(s,a)-g
(82
8118, (al als) (01 8))
(
Σ Σ
Σ
(s=
0E9r, SEES
(5.0)ESXA: ri(s,a)=g
- p(s' a s0)
Σ
Σ
Σ
p(s"|so, 2) [(a|so) - x [alss))
gEg..2
AEA:
EE
Σ
Σ
(also))- - (also)
gEGr2
A
(Ho.))g
Σ
m(c|so) - x (also)
aEA
T-130) - n (1100)))
Σ Σ Σ (G= sw = -)][TT((sss (1s.))
[ DES $68
where (1) realize that in MM the initial state is always so, and that the transition model is Markovian and independent of the policy, at (2) we apply triangle's 's inequality and keep only 8o because of the indicator, and at (3) we have simply rewritten the expression in more convenient way for proving the result (note that [1] = {1} and Grr,1 = {0) and ['T 0 ).)).
81 =) 1
Now, let us consider any stage h E [3, H]. Let us make the inductive hypothesis that:
Σ
Σ SES
ps
(G 9 8h-1 s) - p (G 9 A 8h-1 = =s)
gE9r.A-1
Σ
Σ
p (GH gA SM = 5). (- -1,.(-1999)
Meih-21
geG, se5
Then, we can write (we use symbol w/ to denote the random trajectory (81, 01........5 8b, 0z) up to stage h):
Σ
Σ
PT = sh = = FF (Gr g' ^ =s){
G'E9. d'ES
28

---

Published as a conference paper at ICLR 2026
Σ 3ΣΣ Σ (
g' S'ES geSr,,-- 1
Iam (Ch-1 = 9 A 1 = g'-----
g 8h =
- (G
Th (8h- gh
g' g ^ 8h =
Σ Σ
Σ
Σ
Σ
g'ea..h s'ES gEfr,
wef F: (ea)ESXA ((
Pa = w -
- 8h 5')
10 5))] =)))
44) Σ Σ
Σ
Σ
g'Egn, SEE n- w (Ro)ESxA:
-(uir)-9 9 1h-1(s,))
1d (wh-2 w 8h-1 3) (01-1= =0 a A 8h = s)
pa (wh-2 2 - = s) (an-1 (an-1 a St s' s/(w,s))
(5) Σ Σ
Σ
g'Eg..A h SES gest,n-1 w (a.))SSA
((
Pa (T(wn w 8h-1 = S) s) (Ph-
pad (wh w w = ()) ())
8) s)pn
Σ g'EGr, S'E8 gEgr,h-1 weih--1: (x.)ESSA:
Σ
Ph- Ph-1 (5)
G(wrr)-gr 1(5.0)-6'-g
1s (wp_ = w A s - s) s)
pr' (wh-2 2 = 8h - = ss)
+ PTt (wh-
S) A 8/ 1))-1=)) (a((()))
E Σ g'Eg..n SEE gEg G(war) 9 (ea0
Σ
Ph-1(8'1s,a)
well (s.jeSxA:
Ppt (whh2 w A 8h-1
8) (a)w,,)) - n' (a(w,,))
Σ Σ
Σ
Σ
()
g'eS.. S'eS gE wefl-1: (x.))SSXA
=
A (
Σ Σ Σ
Σ
g'Eg..n g- win (x.))ESAA:
G(wir)=g y (h-1 -g
p pT = w A 8h-1 8) a() - <'(aw,, 8)
Σ Σ Σ
Σ
(5-1(5'|s,a)
Σ
<(alw, s)
I g'Er..h S'eS ge
(xA)ESXA:
Fh 18.00)
G(wrr)-g
(PTT = 1 {)
= $h-1 (
-
29

---

Published as a conference paper at ICLR 2026
8
Σ
Σ
Th Th-1(als,g) - (
g- (s.a)5xxA
Σ
P"
wERR-11 1:
G(war)-g
Σ
E
Σ
Th xh-1(a|s.g)
Σ
g' gefrch
(s.))
WERR 1 G(wir]=g
rh- (4,0)=9'
(wh-2 w 8h-1 8h-1= -x)]
= Σ
Σ
Th Th-1(als,g) - 1 (GniI= = = 1 8)
gEgr,M-1 (s.a)ESxA
Σ
Σ
g'Eg,,h gES
(a.))SXA:
(h-1(s_a)=g'pp
-)) PP (G
-p
= = AA 8h-1=
9 Σ
Σ
-Gh-1= TT
-/11((18,
gEgr,h-1 SES
Σ
p. (Th g ^ 8h-1 = 5) -p (G = g 56-1
gean, 1 SES
(10)
Σ Σ
1pt (Gh-1 = = 8h-1 = 8)
- क 1-1,0-1,,0)
gE SES
Σ
P"N = g sN s)
We[h 2] G SES
(-1s, 9) - (1.
Σ Σ
Σ P' (G= = 8) (1), =/o(11100)
|e[[-- GEG,, SES
where at (4) we use the chain rule of conditional probabilities, at (5) we do it again, and we recognize the policies π and n', and also that the transition model Markovian, at (6) we use triangle's inequality to split the summations and bring the absolute value inside, at (7), in the first term, we note that p (s, a) is the only term that depends on and that it sums to 1, while in the second term we exchange the order of two summations and apply triangle's inequality to bring one inside, at (8), in the first term, we first remove the summation on g' along with the indicator function that forces us consider subset of state-action pairs, and then exchange two other summations and note that the policies do not depend by hypothesis on the entire past trajectory, but just on the return so far. Instead, in the second term, use that Eves Ph-1 (5' [s, a) 1, and also that, by hypothesis, T' does not depend on the entire past trajectory, but just on At (9), i.a., we use that Engeg... (ra-1(s, g' g) 1 and that EaeA T/ (a|s, g) 1. Finally, at (10), apply the inductive hypothesis.
Thanks to this result, we can finally prove the claim in the lemma, using passages analogous to those above, with the difference that we do not have the summation over the states the current stage (i.e.. H + 1):
1F n 1/1
Σ
(T) n. (o)
geg 1+1
Σ
pst (GH+1 = g) -P
Pp (GH+1= = 9)
M
33 Σ
Σ
(rn((s',a) = -g- --gggee G(,(w) 9') (
gEG gEG
S'ES g'e.. welly
30

---

Published as a conference paper at ICLR 2026
P = an aA=H= = =
pd (WH-1 w SH = an
- =d)]
Σ Σ Σ
1(r((ss, =9 g', G(wrr) -
M s'ES d'EAA g'Eg,, welh
P" ("(WH-1 w A 8H = (
p (WH-1= = SH 5 (d)
(wH--==== sp = =d)
(
Σ
Σ
PT (G g 50
THH (1s
g M SES
Σ
Σ
P (G = gg = =) pp (GH g s)
g59r, RES
(12)
Σ
E
Σ
P
s 5) (-1,,g) - (-1s
M][H] GEG,,N F.M! NES
where (11) we made the same passages above, all in one, and at (12) applied the result proved earlier by induction.
This concludes the proof. Note that, from high-level perspective, our proof approach resembles the "reduction to supervised learning" approach made in Rajaraman et al. (2020) for stochastic policies, with the difference that we work in augmented state-space MDP.
Lemma C.4 (Concentration). Let e (0, H] and ò E (0, 1). Let M, be any MDP, xt E TTNM be any expert's policy, and पt be the output of Algorithm 7 Then, with probability 1 5, we have that (we use the notation in Lemma C.3
Σ KE[H] BEST EE.S.E.... SES
1pT
(Gn =gA= 8) F. (.).1s. E..
with number of samples:
1935 HG 280
2S9
128 In
250 8
N
In
+ (A - ln
2
8
2²
where G = Enerm 99
Proof. We can write:
Σ
Σ
1x
(GA = g 8h = T. [ - 7 (-((1s.9))
H] gEg F
J. SES
Σ
Σ
= pa E (G,, = g s4 =)) Σ
T (aggs)
H
GEG
SES
aEd
Mn(s.g.a)
Mn(s,, -0 of 111 Ma(s.g,a)) = o))
Mn(s.g,a))
Σ
EPT =g
he H DEg SES
In 280 (A - 1) In(e(1 + Ent (s.g.a)"))
2/2
A-1
Ev Mi(s,g,a)
2SG
+ (A - 1) In (c(1+
)
N
2V2 In
8
31

---

Published as a conference paper at ICLR 2026
ΣΣX
Σ
p (Gn = g gn sn = Bn E. Mn (n(s,g,
he|H| gEG, .F.h SES
( 2v In
2sg
(A - In
1+ + A 1)
N
8
8 25Q pa
(GA g 8h
Σ
N
RE[H] gEG, E. SES
In 2SQ
2SC
+ (A- 1nn-)in(((1+
A A-1 )
In
N
Σ
Σ
1px (Gr = g sh = 8)
ke H] gEg F
(4)
In 25g
2SG
+ (A (A- (A-1)ln(e)) in
A M1 1
N
N
Σ k=]
SIG B A
Σ
Σ
pat (Gn
8h 8)
yEG, F.a NES
81
S In 257
A
In
2S7
+ (A (A-1))mm) 1)ln(e(1
N ))
Σ hE[H]
V G.g.
N
8
(5)
25 SHŌ
In
2SG
(A (A-11)mn 1) 1 1)
8N
5
In
N
8
where at (1) we use that, if E. MA (s..9..0 d') = 0, then:
Σ
Mn(s.g.a)
1 Σ Ma(s,g,a') 0) + 1
A 1
Ma(s,g,d)= =0))
T
E. Mn(8,0,a)
AsA
Σ A
m.g (rggn( 1
2,
as we the total variation distance between two probability distributions cannot exceed 1. Instead, E.d Ma(s,g, g, d') > O, conditioning on E.e Ma(s. g, a'), at all 8, where pa (G,, A 8h = 8) > note that Ms a)/ E., Mn(s, g, a') is the empirical vector of probabilities of n. x, ( (a)g.s)
(recall its definition from Eq. 6). thus we can apply Lemma 8 of Kaufmann al. (2021) to get that, for any å E (0, 1):
Ma(s.g.))
E 5 [KL
Eg (s (=
In 1 (A 1) In(((11 + En Ma (a.9.a))) A-1
> 1-d. d.
Enp Ma(s,g,a')
Combining this result with the Pinsker's inequality, that tells us that |r 1r- $ / 2K yy) and with union bound over all h [Н]] E S, g E 9.5 As we get the passage (1) w.p. 1 6/2. Note that we add an additional 2 for the case Σ., MM, g,a O, and we define ÿ := Σme[mn G9 nl- At (2) we bound Σ., Mn g, a') & N. and bring that quantity outside, (3) we apply Lemma A.l of Xie al. (021). after having noticed that Σ., Mr(s, g, a') Bin Bin(N, [N, (G,, = g ^ sh. = s)). and make it hold for all 8, g, w.p. 1 5/2. At (4) and (5) apply the Cauchy-Sehwarz' inequality.
32

---

Published as a conference paper at ICLR 2026
Now, we impose that this quantity is smaller than
SHŞ hn 259
2S9
+ (A 1) 1) ()
x-1))
N
E
5
64SH In2
280
64SH9 (A
In
250
e.V
In
1-1 + +
2
2²
Thanks to Lemma J.3 of Lazzati et al. (2024). we know that this inequality is satisfied with:
N 128SHGhi In 285
250
128e5HG1H 28
In
A
2²
Rearranging and applying a final union bound concludes the proof.
C.2.2 ADDITIONAL DISCUSSION ON RS-BC
First, observe that the sample complexity bound in Theorem 4.3 cannot be improved if we extend our analysis with that of Foster et al. (2024). Indeed, Foster et al. (2024) also provides an 1/e² dependence our proof, that combined with the additional 1/6 due discretization, would give the same 1/c² rate. Moreover, Foster et al. (024) would not allow to improve even the H dependence in the horizon in our proof, as their Corollary 3.1 combined with simple variation of our proof that considers an MDP with an augmented state space, would still provide an Н dependence that should be combined with the 2 arising from bounding the Wasserstein with H times the total variation, and taking the square. Note that these considerations on Foster et al. (2024) implicitly assumed that the proof Foster et al. (2024) can be extended to nor-Markovian expert's policies with the same rate, which has to be demonstrated as well.
Second, we mention that Rajaraman al. 2020) provides for IL a 1/c dependence instead of 1/2.. However, note remark that the result is in expectation, and not with high probability, remarked also by Foster et al. (2024) in their footnote 21. Indeed, this explains why the 1/c rate of Rajaraman et al. 2020) seems to overcome the 1/62 in the lower bound of Theorem G.1 of Foster et al. (2024).
Lastly, we mention that RS-BC (and also the theoretical guarantees Theorem 4.3) can easily extended to the setting in which r is unknown but observed, namely, in which expert's 's trajectories are state-action-rewan computation of the returns ((w; BB does not require knowledge of F in state-action pairs never observed. This different from RS-KT, in which we require knowledge of n. everywhere.
trajectories (31, 01, ......). Indeed, looking at Algorithm note that the
C.3 KNOWN-TRANSITION SETTING
In Appendix C.3.1
we prove Theorem while in Appendix C33.2 we write down explicitly the
LP in Eq. 10
C.3.1 PROOF OF THEOREM 4.4
Theorem 4.4. Let E (0, H) and (0, 1). Let M,E be any MDP and F NM any policy. Assume that the optimization problem in Line solved exactly. Then, choosing = e/(7H), with E probability 1 8, the policy ã output by Algorithm satisfies Wyr n,x)) & e, with:
H 2 1
N < 0
(9)
Proof. We begin by showing that the estimate of return distribution computed by RS-KT Line is close to the expert's return distribution with high probability. To this aim, we write:
W
E The
(1)
< W
ww
(2)
<
HO/2
W
33

---

Published as a conference paper at ICLR 2026
H +
c (x) - FG(x)) dx
< HO/2
sup
F 0/ "0
(xx) Fagla") dic
='E[0.H]
11
(xx)) F3(att) |s
H sup
F n
do
x'E(0,H]
0
HH/2 + H sup F (x') F3((1))
x'E(U.H]
nig
H0/2 + Hd',
where at (1) we use triangle's inequality, at (2) we apply Lemma C.1 at (3) we use symbol Fg for the distribution function of any probability measure y, at (4) we apply the DKW inequality (Kiefer & Wol-
fowitz 1959.
Massart 1990), as Fa(()) = En'ge Max') En'sw N EEENN 1 (E 2ē (s), %j))
N LEENN " ( To. (5h, a/) & x') corresponds to the empirical distribution function of Specifically, in our setting, the DKW inequality tells us that, for any é' > 0, it holds that:
p sup "'
F
(x)) Fa((z') < ) 11-2e 2e- 2N2()
(13)
By imposing the term the right hand side to be 1 8, and solving w.r.t N, we get that:
In 2
N E 2(6/)2
5"
Now, building on this result, write:
W w( E R (5)
$ W n 1) + W i7
7 1r
+ W
(1ig ( mp
(6) HH/2 + He' + w( (ī. -
mrg ) + w (1.g:1,0)
< H0+ + Hd + w n7 n
HO + He + min
W ((7. 1/g0
=
(9)
& H0 HH + +w(
Tif
n.g
(10)
$ HH Hé + +w( ((:
W ;,,,,
3
+ W
11,11
F.B x
(11)
< 2HO + 2H + +W(1 1 dR
W
3. :T.gg
:TIE")
7
(125
HO + 2Hd + +WW &
1T M,E
(12) HO + 2HC'
&
where at (5) we apply triangle's inequality, at (6) use the result above, (7) we apply Lemma C.1 at (8) we use the definition of âे and the hypothesis of solving the mininization problem exactly, at (9) upper bound with specific choice of policy, i.e., म, (recall Eq. 6). at (10) we apply triangle's inequality again, at (11) we apply the result above again, (12) we use again Lemma C.T and finally, at (13), we apply Lemma 4.2
If we now choose 0 e/(7H), and e' = e/(AH), we get that, with probability 8- it holds that:
W The : &
34

---

Published as a conference paper at ICLR 2026
with:
2H N &
2 8
C.3.2
EXPLICIT FORMULATION OF THE LP
The optimization problem in Eq. 10
can be written more explicitly follows:
min MERDD
Σ t(g)
SANID' deR deR
JER 1D/8
gE3"
s.t.
Σ
(1s.9.a)=1 =1
(14)
(xxxg)
Σ
di(s0, 0,a)=1
(15)
aEA
Σ dn(s, 9)
-E
dn-1(s', d'''s'''ddd')nn_ille" p 1') -g- - -9-
8'
(ss 3 (2, ........ H} (16)
n(9)= = Σ
dh(s,g - (0,
(17)
(s,a)ExxA
x(g) = Σ
(1()g) - (('')
Vge
(18)
g'Ey": g6
-t(g) < (()) (())
Ag e y
(19)
where Eqs. 14 15
denote the flow constraints. i.e., define the set of feasible occupancy enforces that is the return distribution in M corresponding to occupancy permit to rewrite the Wasserstein distance in a linear manner.
16
measures AC, Eq.
measure d, while Eqs.
18
19
Observe that the number of optimization variables SAH + 1y& ((SHH) while the number constraints 2 (H 1)8 3 O(SAH))
C.4 WHEN E BELONGS TO FINITE SET d REWARDS
In this appendix, we consider variant of the known-reward setting, in which 5 is unknown, but we have knowledge of set (r- (r'...., rd) containing > 1 reward functions, and we also know that E R. We consider the following robust variant of RDM for this setting:
a E arg min max W 17 WEllNM rER
(20)
To tackle this problem, we will proceed to Section 4 by considering the cartesian product of all the rewards in set Specifically, we first present class of non-Markovian policies sufficiently expressive for addressing this task, and then we present two variants of RS-BC and RS-KT. Crucially, we will have exponential dependencies the number rewards for both the computational and sample complexities.
Let us begin by presenting IT(R), a generalization of IT()) to multiple rewards:
II(R) { T NM
f E A A 1 xSx0, xx..xx
(a]ss,w = Φn (a|s, G(w; ((..,) ))) Vs E S, A, [H] S2
Intuitively, (RR) contains policies that depend on the amount of rewards collected so far for every possible reward p² in R.
Now, let us define मæ. E [[()), an analogous of policy म, (Eq. 6). For any set of rewards R and expert policy nf NTNM define कर IT(R) the policy whose probability of taking action in
35

---

Published as a conference paper at ICLR 2026
Algorithm 3: Variant of RS-BC for set of rewards
D = (ii,,i, .i...., sh. aiz) MENI, set of rewards R = fr',,..,,"" r") parameter 0
// Count the state-acti aeate-action-cmmee aeccccto--cuulll lative reward eccurrences
EMENT 1sa, s, a. Ei (s = Ma(s.g ."...""...... .a) g...... - nd
(sh (ar(s))aa))) ald) = = vh E [H] s E S. y E Vh. a
11 Retrieve the policy w())sww)
My Male:GGorr)a..
if denominator
(a .... Gtwigg)
Vh [M]] S.w As A
otherwise
Return 7
state with history E On coincides with the "average" probability with which n selects in after accumulating G(w;; r') reward for each reward r" E R:
Ipr (sh- =, an a, G G(w;rt), ... C = G(w;ra)) p (sh = s, G = G(w;r²),
TR (R(a|ss
(21)
GA = G(w;rd))
where we defined the random return at stage h under reward r² R as Gj,, == 2 (sy ax). the denominator is zero, we set ,,, (a|s,, w) 1/A.
We have the following result replicating Lemma
Lemma C.5. Let M be any MDP R fr'....,r") r any set of 2 rewards containing the unknown r", and π E II be any policy. Then, the policy गR E II( is minimizer of Eq. (20). and satisfies 11
R (g) n,-- (gr all g.
We prove it in Appendix C.4.1 However, II) IINM of course for some rewards. Thus, we
(rg..... rg), i..., by discretizing each reward inside R. Then, it should be clear that the memory required for storing policy in II(R²) scales ((SAN)))))) (because we do the cartesian product). Analogously to Lemma 4.2 we can bound the approximation error (proof in Appendix (C..I]):
can discretize, by defining: R
Lemma C.6. Let 0 E (0, 1.. Let M be any MDP R = [r2 (rt,...,,
.......)) any set of d 3 1 rewards containing the unknown F and _ I1 be any policy. Then, the policy TR= E II(RR) satisfies ( T .T,E H0.
Now, we address this problem with a variant of RS-BC for the no-interaction setting, and a variant of RS-KT for the known-transition setting.
We begin with a variant of RS-BC, reported in Algorithm 3 Simply, we count the (discretized) occurrences for every reward in R. We have the following result (proof in Appendix C.4.2p:
Theorem C.7. Let (0, H and 8 (0. 1). Let M be any MDP R, R (r' ..... rd any of
> rewards containing the unknown F. and x F e NI be any policy. Then, choosing e/ e/(4H),
with probability at least 1 - 6, the policy îे output by Algorithm 3 satisfies W(n m with number of samples:
N & δ
(
SHH+2d&& In
(A+m²))
(22)
2+d
Observe that, for 1, we retrieve the known-reward setting, and the number of samples in Eq. 22 matches that of RS-BC (Eq. ().
Now, we do the same for RS-KT. See Algorithm for variant of the algorithm. We have (proof in Appendix C.4.3):
Theorem C.8. Lere (0, H] and d (0, 1). Let M be any MDP.R. R = (r',....,") any of d > 1 rewards containing the unknown F. and Fh E NM be any policy. Assume that the optimization problem Line 2 is solved exactly. Then, choosing 0 e/(4H), with probability 1 8.
36

---

Published as a conference paper at ICLR 2026
Algorithm 4: Variant of RS-KT for set of rewards
Input: Dataset D = [(si, ai, ai...., a)., (i)) [ej]) set of rewards R. = (r' ,...,,"),, parameter 0
transition model
// Estimate the return distribution of the expert nf." for any TE R rom a = Age V", Vr R
(x())) & EieNN
// Compute the policy in I(RR) closest to Rrr E R via Eq.
24
E arg (ET((R) max,eR W(r.,.
Return 7
= the policy ă output by Algorithm 4 satisfies W(n,s nT,5))
with:
H 2 In
N E
0
(23)
Interestingly, the bound here is still polynomial. Now, we show an extension of the LP formulation in Eq- 10 for addressing Line Of Algorithm Specifically, we just construct augmented MDP that keeps track of the past rewards for every possible reward in R, and note that II(R)) describes the of Markovian policies in this MDP. So, want match sort of augmented return distribution
for this problem:
min deK.n AEA
max W(T', TT-))
(24)
P'RR
s.t. (1(g)
Σ
dn(s, "a..t...gg..... a)) .
(25)
8,0.8'
Vg E Vie [d].
(26)
Intuitively, the constraints above enforce that ry' is the return distribution induced by the occupancy measure w.r.t. the reward rj. for all the rewards r' R. and K. denotes the set of feasible occupancy measures in this augmented MDP (Puterman, 1994):
K= de A AxA H
Σ di (30, a) 1 V3 2: Σ (3,30) = Σ dd ad)
8,.
where the state space = S y .... yd d times, 30 = (s0, 0, O,.......... model is:
,.. 0) and the transition
gd) In words, Eq. 24p searches for an occupancy measure E K. that induces the return distribution ry' closest to 7 From such a solution, policy î E II(R)) with occupancy d" = d (and thus return distribution nf for all i) can be recovered via:
(a(s', 2d 1s.g I- a)) pn a) (rx(s, - - ) (r2((s
(s, G(w; r.)......... ................................ G(u, (((ssww) E. (s, r() G(w; r a')
Yh E [H], S S, a A, w E Sn.
when the denominator is nonzero, and @(a|s, w) = 1/A otherwise (Syed et al. 2008). We remark that, being the set R. finite, then the minmax above can be formulated LP minimization problem.
C.4.1 TECHNICAL RESULTS AND PROOFS FOR THE POLICY CLASS
Lemma C.5. Let M be MDP,R, R = fr fr'....,r") any set of d > rewards containing the unknown E. and T e NM be any policy. Then, the policy कR e [[R) is minimizer of Eq. (20).
and satisfies v x (g) ni. (g) for all g.
Proof. To prove this result, we show that, for any o E R, the return distribution RR coincides with the expert's return distribution n
37

---

Published as a conference paper at ICLR 2026
For any 1² R and g' H], we can write (we use C == - (sw, and for all j e [d]]:
my (g') = P (G + (H aH) = g
1 Σ
Σ
P (G =
- g9)
9 A H SES
Σ
Σ
Σ
g 10
MEG " #ES
paR
(G)
g
Gd = g s sh
s, s. g g')
= Σ
Σ
Σ
Σ
8'E0,1,M g' #ES aE (3
PaR (G
g
Gd
g
88 - an 8 aa)
Σ
Σ Σ
Σ
g g'E9,1,H g'EG,d, #e8 AE (sa))))) (sa))-''-"
paR (G) # g"
C
= +H= s)
Σ
Σ
Σ
Σ
g g-E0,1,M A0, H
#ES aE (s,a)
p (G
G 9 ,8
8) ) 2....
Σ
Σ
Σ
P (G) = gg...., G 9"h 8H = 8)
0'E0,1,8 H
G RES
Σ
TR(als g")
aEAiri -gg0)
(4)
Σ
Σ
p (G) = g² " 5 =
SH
g' SES
(SH =s, ag = G h g .... G gd)
Σ
p 1s s, G = g¹ G - =
aEAirt (5:0)-9' -g
Σ
Σ
ΣΣ
= ()) g
g E
StS aeA
pa E (sH = 5, On = a, G = gو
G = gd)
Σ
Σ (rrn(s, g g | p (sH s, ay a, G = g')
ge .i eES aEA
1/p (g'),
- where at (I) we define symbol G.H == (ge H1 - 1w E ΩH : G(w;r) = g] for any r, at (2) we recognize that, by definition, RR takes actions only depending past rewards for any r², and we denote with brevity this fact with Ty (als. gl..................... ."),, at (3) we use Lemma C.9 at (4) we use the definition of [a] g) (Eq. 21 where the denominator is not O. noting that, in that case, the entire expression evaluates to
the current state, stage and
Lemma C.6. Let 0 E (0, 1]. Let M be any MDP'R, R = .r.,..., , any set of d 3 1 rewards containing the unknown E and Nt E NM be any policy. Then, the policy g E II(?)) satisfies
W(T nz mTT)) < НО.
Proof. For any reward od R, we can write:
W( W(1, ZR0 11e 11/ & W
FF)
x=20) + w
W(1 u .1xh)) + W W(155.07"
11 -1
38

---

Published as a conference paper at ICLR 2026
211
r' ról00 1
H0 W
Txe0 =e 1r :
(4) H0,
where at (1) we apply twice the triangle's 's inequality, at (2) we apply twice Lemma C.T at (3) we realize that, by definition of ry. it holds that ||| rglloo & 0/2, and finally, at (4), we apply Lemma
C.5 with set R and expert's policy n².
The proof is concluded after having observed that F E R by hypothesis, and so these passages hold also for P.
Lemma C.9. Let M be any MDP R, R = fr' fr"....,,"" rdj any set of d 2 rewards containing the unknown F and π E IINM be any policy. Let TR E (TR)) be the policy defined as in Eq. (1 for
expert's policy Then, for all h E [H]] se S and g',............. .....9a E 10, 11, it holds that:
PIT (ch= == A..... Gh o² ^ sh A ss) Pm (h) g A... C =g" A
where we used Gj, == Σ (sw ax)
Proof. We prove the result by induction. Let us begin with the base case: h 1. For all E S and [00,, we have:
g g.....'e
Pa G = g² A....A a GH = od s = 8) = -110* = 0 ^ o = 0) fo}
= -P (G) = g¹ a G = o² ^ = 8).
where we noticed that, for h = 1, no action is taken yet. Now, let us consider any stage h (2.3.
and let us make the induction hypothesis that, for all h [h - II. for all s E S and g'...... [0, - 11d, it holds that:
H). .....
g"
aad (Gl- = gg A... A...AG.-- = 9" ^ sw = = 1p (Gh, g' A a... Gy = =g² A sw = s).
Then, for any 8' and Ft . E [0,h- - 1]d, we can write:
pax (G), = g
G g A sh = =s)
Σ
PpR = A.W.WAA
=
wEBK- (xx0)xxxx
Olwer + +ri (5,a)- Vi
Σ
Σ
Σ
Σ
g'E9,1,N11
geEg R-1
we 12
(s.a)ESxA:
G(war') )-g" Vi (s.a)==
pare (wih
w
8h-1
a = a
8h = 88)
Σ
Σ
Σ weR
Σ
pax (wh-1 AA
9'E9,2,1
gde R-1
G(war") )=g" (s.a)ESxA:
1
.PF (ap-1 = alw, s) s) = 8'1 [W,,,, a)
Σ
Σ wi
Σ
gdEs pl h-1
(s.a)ESXA:
G(wer") ) Vi (9,a)=
pa p = alw, (8)
Σ
Σ
Σ
Σ
g' 1-
gdes pd k-1 -
(s,a)ESxA:
G(wer") )-g" Viri (3.)==""
39

---

Published as a conference paper at ICLR 2026
TR(ak,, 1
Σ
Σ
Σ
P (wh-1 -
(x
wEfn-1
( vi Gluir))="" W
TR(aaw, 18ph)-
Σ
Σ
9'E9,1, n-
(a,a)ESxA:
-
TR(aag" "...,g" 1))n,,,,))pnn(1(111
Σ
PER = 5h-11
weRn-1:
G(arr)= Vi
Σ
Σ
Σ
TR(alg'... ....,9" s s))pn-(6'|ss,)
g1Eg,1, h-1
geeU yd 1-1
(x,
(t-1(=,)== Vi
P G h- = gt
Gh Σ
8h-1
Σ
TR(alg' !....9",
g'E0,1, h-1
(s,a)ESxA:
16,,0)='-99" Vi
(TT(G) = g
Gh-1= = 9" sh-1=s) Σ
g'E0,11K-1
9dEg 1 A-1 (s;a)ESxA:
" 1(,,0)=''gh)
1 (G)_1 gd A...A Gy gd A 8 0h-111111-
P PT 9¹ A... G Gh-1 9d 8h-1
Ph P1-1(''1s, (G = A... A Gh = =g¹ ^ 3h-11 8) Σ
Σ
Σ
g'E0,1, h-1
(s,a)E5xA:
The (5,0)
Pdt PTT(C)
=g
Gh-1= = 9d
Σ
Σ
gLE0,2,N-1
gMEO, p h-1
(a,a)ESxA:
(5-1(5
Ps PT1(C// gt A... -h-1= 9² -A
=
P (G) A... GH -d ^ S/ s')
where at (1) we use symbol w to denote the random trajectory long h²" stages, i.e., whose realizations
belong to 2x,, for any h" e [H]. At (2) we define symbols Gr, = {g E [0, 1 1]
On G(w;r) = g) for any r, at (3) we use the chain rule of conditional probabilities, at (4) we use the Markovianity of the environment, at (5) we note that pa (an-1 is TR(alw s), at (6) we exchange the two summations, at (7) we recognize that, by definition, TR (R(aa (alw,,s) s) takes on the same value for all the trajectories with the same value of return for all rewards r² R, and thus we can bring this quantity outside the summation over the w. We use .....-""
alw,s) actually
to denote this fact for brevity. We do the same also for Ps'1(s'|s, a). At (8) we use the induction hypothesis, at (9) we replace (R(alg] ....,. ~², s) with its definition when pa (G) = g² Gd. G 9d 5h-1 5) in the opposite case the entire formula takes on value zero.
symbol (alag" ...."" -
C.4.2 PROOF OF THEOREM
Theorem C.7. Ler E (0. H] and E (0, 1). Let M be any MDP R. R= "e..,..,,"" any set of d > 1 rewards containing the unknown E., and nt NT be any policy. Then, choosing c/(4H),
40

---

Published as a conference paper at ICLR 2026
with probability at least - 6, the policy ât output by Algorithm satisfies W(n nT,5) with number of samples:
SH In
(A + In (nn 2))
N & δ
(22)
24d
Proof. We can write:
W The (re)
W W(1 1,, ,7,E + W
Twe)
7nE
Tza Tza
- + W
Txx" .T.e) (1,
THE
+ *w n n7-2))
$ 2HO + W (11 R -
! (155)
2HO + H1
7
n
2HO + H Σ Σ
Σ
Σ
P = A a Gh = s
he[H] g² 1j
g .g.n
$ES
( ................................ 9d)
801
2HO + H Σ Σ
Σ
Σ
P 2 (G) g A... Gh = =
sh ~
he [1 g² s'
0 og
sES
R (1s. ...."
8-) tx (1s. 9²) |
& 2HO + Hé',
where at (1) we apply triangle's inequality, at (2) we apply Lemma C.6 and Lemma C.1 twice, at (3) we use Particular Case 6.13 of Villani (2008), which tells us that we can upper bound the Wasserstein distance between two distributions supported on set X by the diameter of X (max,, Jx - æ') times the one norm between the two distributions. Since X [0. HJ in our case, we get the expression written above. At (4) we apply Lemma Lemm with the notation defined that lemma with set Rg and policies TR and î. Moreover, this holds recalling that r e R and so r E R (5) we use Lemma C.9 Lastly, at (6) we apply Lemma C.11 with accuracy e'.
The result follows by imposing that é' & 2H and 2HO $. which can be achieved by taking é' 2n and a and by observing that:
=
Σ
II
[E[H] [é[e]
Σ
D
h=[H]
Σ DA
Σ (1+ (h 1)/02² [=[l]
o (H(H100²²
where at (11) we used the previous choice 9
4H
Replacing into the number of samples in Lemma C.11 11
(and also e'
get the result:
2H
SH
d² In
1
A + In 1 1
N E
2+d
By using δ notation to hide logarithmic terms in A, H, 1 In 1 d, we get the result.
41

---

Published as a conference paper at ICLR 2026
Lemma C.10 (Error Propagation). Let M be any MDP, Rand R {rl ...... rd any set of d 1
rewards. For any pair of policies т. I' IINM such that, for all E [H], a E A, s E S and w,w' Uh with G(w,r") = G(w'; ('') Vi [d]:
x(aww,)) =(alw's))
=) x'(aw'
it holds that, for any r¹ R:
n 7p
Σ
Σ
pst (G), gt A... a C g ^ sh =)
gdeg,d AES
TA(ls, .......g)) n (.1., ..... ....,,9)]
1] Ba S2 G(w;r) = for any reward r, Gj = En (sp ard denotes the random return stage h under reward r², and n g'......... .-.., g and x (-1s, ,..., go) denote the unique probability with which the policies a and n' prescribe actions in s at h under any trajectory w Sn with G(urr") = g Vi [d]]
where Gr, ==
Proof. To prove the result, we first demonstrate by induction that, for all h E [2, HI. it holds that:
Σ
Σ
Σ
pa (G) gt A... Gh 9d A sh s)
g¹ g
Pr (G), =
g
A
A G = g ^ 8h = 8)
Σ
Σ
Σ Σ
P. (G = A A...
s s) C g A sw
HeTh- 11 1
grE SES
(11,9
gd) ( (".,,,....
We begin with the base case h = 2. We can write:
Σ
Σ
Σ
P (T] g6 A A ~ 82
g 0'
SES
1p
(G) g A A.... G ~2
Σ
Σ
Σ
(s,a)eSxA: ((())==00
((818,8)111
80) 50) (als)
Σ
Σ
Σ
Σ
p(o'ssus a))T())))) 2'((/=))
8'
s'es
ri (s0,a) VVa
Σ
Σ
Σ
- -
8 20
gdEG,d
(1(00.))-9
Σ T(a|s))) - <aals0))
|s((150) - r'(1150)||
Σ Σ
Σ
> p. (G = = A... G = a sw =)
Wel]
gtE gdE SES
[ ...0..".''"" "²t)
- ,.... 0z)
where at (1) we realize that in M the initial state is always sp, and that the transition model is Markovian and independent of the policy, (2) we apply triangle's inequality and keep only so
42

---

Published as a conference paper at ICLR 2026
because of the indicator, and at (3) we have simply rewritten the expression in more convenient way for proving the result.
Now, let consider any stage h
[3, H]. Let make the inductive hypothesis that:
Σ
Σ
Σ
P (Gh-1= =g' A...A a Gh-11 = 9" A sh-1 =s) 5)
g' g'
g guEl h-1 SES
- pr (Gl-1 = g' A... A Gh-1 = g" A sh-1 =s)
Σ
Σ
Σ Σ
SPT(C g¹ n... Gh = 9 A s
- 1',
gMEg,A,N SES
[
the (-1,,
Then, we can write (we use symbol w/ to denote the random trajectory (81, ..,...,,........................ 0b) up to stage h):
Σ
Σ
Σ
PT (G), = =g¹ A... a Ch = 9d A sh = s)
g g'
9 gM SES
Pa (G), g A... G o² ^ Sh = s)
Σ
Σ
Σ
Σ
Σ
Σ
11
3 888 g'
WESt
(8.0)ESxA
G(wer') ) VVV (1.0)) -r
Pe ("((1-2 8h-1
8)TP (an
8h = 8 s'(w,))
p (wh-2 W
S)PP
(ah-1
a -
8' |W 8)
Σ
Σ
Σ
Σ
g¹ 1 b
odE d'ES g'EG,1,n-1
WESt (MA)ESXA:
G(weri) Vir =0.a)=="'-""
PPT (P 8h-1 ) (p
pp' (wh-2 = W 8h-1
5) (alw,, 8) s) 1 (6/1s,0))
Σ
gdE d'ES 8'89,1,n-1 3ME
wEfL
(s.)ESxA:
G(weri)=3 vink
(
Ph- (s |s, a) p.t (wh-2 8h-1 = ()T(alw,s) s)
p (Wh-2 = W 8h-1 = sa/( 8)
1 +1P = W 8h-1 (aw,, 8)
Σ
Σ ΣΣ Σ
Σ
Σ
Σ
gME0,2, J'ES g'
J'eS h-1
w
(s.a)ESXA:
G(wr')=
((".))-''-"' -''Vi
Ph- Ph-1(s'|s,o) a) PT (Wnn22 = W A sh-1 8)
T(alss - sa)lw
Σ
Σ
Σ
Σ
Σ
s'eS g'
Σ
Ph- Ph- (0)
well (s,0)ESxA:
G(wi)) "= (s -7'v
P (Wh-2 = Wn 8h-1 8) - ppt
(Wh-2 = w A -))1=)))
43

---

Published as a conference paper at ICLR 2026
Σ
E.......... Σ
Σ
Σ
Σ
gdE0,2 1 2-
-
h-1 G(wir") (5
wefl
(c.0)eSxA:
(" = A 8h-1
5) T(aws, 9) - x (
Σ
Σ
Σ
Σ
Σ
Σ
g' 0'E9,1
g'E s'ES s'E33'E,,
g'
(A.))ESAA:
Ph- (5(1,
Σ
s)
((wrr") -ig'v"
(PT = 1 10t (wh-2 = w A 8h-1 = =))]
T ......,............... g4) - T/ "al, g4)
J'E9,1 1-
g g"ES, A a-1 (A.))S5A
Σ ("(wh22= = = A
WER 1 G(w:"))-''
E
E
g gLEG,IM
T'E ri 1(:.0)=g' -'g'i
(e.a)ESXA:
T/ (alss g..
Σ
wEft-11: G(wer 3=v'vi
(PFTWN-2 w 8h-1
8) P (wh-2 w 8h-1 = =s)
Σ
Σ
T Th-1(als,7 ..... , 1 77))
3 F'E (a,))SSA
Pp ("
h-
g
G
A 8h =)
Σ
Σ
Σ
Σ
Σ
+
g 1
g"Eg,A,
9'E9,1
gdea,d
(s.))SSA
Th (als,9) ..... .h)
1TT (G)-1 = gg A A..., A A = =J 1 8h-1=s)
FF' (G) = A g" ^ 8h-1=10 8))
Σ
Gh-1 gt A...,AGH_ g² 8h-1 = 8)
34 1-11
g'EE - SES
Th (-18, ....,.........................
Σ
Σ Σ
7EE0,IM-11
3dE n-1 #ES
1sT (G/h11 = go ^ A..., A Gdd
h-1 = -
FF' (G) = g² A A..., A A = 9d ^ A =s)
(10)
Σ
Σ Σ
FT ( = = ^ A..., 1 1 g ^ sh-1= s)
g g'Eg 1-1 b-
44

---

Published as a conference paper at ICLR 2026
Th (--1(-1s,5 ."....""))
Th 1h-1(-1s, .")
Σ
Σ
We 21 g' gd SES
P" P' = A...
Gk 9"t 8/ =
)d)
t ......9)]
Σ Σ
Σ
Σ
Melhh
g N SES
(T = g6 A G g Sh s)
Th (18,9 ..... l² (1ss, m
where at (4) we use the chain rule of conditional probabilities, (5) we do it again, and we recognize the policies π and n², and also that the transition model is Markovian, at (6) we use triangle's inequality to split the summations and bring the absolute value inside, (7), in the first term, we note that Ph-1(s'|ss [s, a) is the only term that depends on s' and that it sums to 1, while in the second term we exchange the order of two summations and apply triangle's 's inequality to bring one inside, at (8), in the first term, we first remove the summation on y² along with the indicator function that forces us to consider a subset of state-action pairs, and then we exchange two other summations and note that the policies do not depend by hypothesis on the entire past trajectory, but just on the return so far for any r' E R. Instead, in the second term, we use that Eies Ph-1 (s' (s, a) = 1, and also that, by hypothesis, r' does not depend on the entire past trajectory, but just on g' ...... .g" At (9), i..., we use that Lgies, E ... Le 1(r a) g" viv and that E.EA T (a g ...... ge)) 1. Finally, at (10), we apply the inductive hypothesis.
Thanks to this result, we finally prove the claim in the lemma, using passages analogous those above, with the difference that we do not have the summation over the states at the current stage (i.c., H 1). For any r E R:
Σ
M (g)
n (
gEa, 1212
Σ
pst
(G)
(G/++1 9)
pr (G
=g)
ge
Σ
Σ
Σ
E
gEg, H+1
8'E9,1.8
g gdEg,
welly:
(s.0)ESxA
G(u;T') (lu:r')-g'ver vi (5,0==9---)
PTT (WH
Sh ())
- p (WH-1
5H 5) (aloss)
(IP (WH-I= - s) s))
(I)) Σ
Σ
1sT (G), g² ...., A =1a 8n 8)
g'eg,1,.
g"E9, SES
TH (H(+ss.g .....g"))
.'.................. gd))) TH (((-ss.g)) ..... 9r)
Σ Σ
Σ Σ
1q[H-1]g'EG.1.
gdEg,d SES
1.T (G = g¹ A A...A A G gd ^ 8k = s)
.e...gg....."......... اوو (1s,
g8) |
where at (11) we made the same passages as above from (6) on, all in one, with the only differences that we sum only over g Grs H1 instead of doing so for any reward of R. and also that we do
45

---

Published as a conference paper at ICLR 2026
not have the sum over the next state s'. The result follows by summing the two terms in the last expression written.
This concludes the proof.
Lemma C.11 (Concentration). Let e E (0, H] and E E (0,1). Let M be any MDP.R R and 'R (r)..... rd) any set of d & rewards, T e IINM be any expert's 's policy, and ft be the output of Then, with probability 1 - 5, we have that (we use the notation in Lemma C.10
Algorithm 3
Σ
Σ
Σ SES
p (G), g A..... Gh = 96 Ah
(- ."..,..) t ..... الل||
......")]] sc,
with number of samples:
N 193SHG1
280
2Sg
128&SHGI In 28
+ (A- In
22
8
6²
where G Ene(in Il hee 19ig.n
Proof. We can write:
Σ Σ
Σ
Σ
(G), = gو A ..... Gh = Am
he 1 gl.9.gm
G"E.gg.n SES
TRR n (-1s,
Σ
Σ
Σ
Σ prp (G) = g¹ A A...A
A Gh = 9d 1h =
NE| s'E0,j.
grer M
SES
Mn(s,gt g"
Σ
T n (als, g ...... 94)
En Mn ..... a')
aeA
> 0} +
1 11 Σ Ma
-,a)
A
Σ
Σ
Σ
Σ
p E (G) = g ^
G = g =s)
hE[H] g' g'Eg,g.A
gdeg .g
SES
In 280
(A- - 1) In(eee11 + E. Mn(agg of
2
A-1
Eng. Mn(8,g go.. aa)
(2)
2SG
+ (A- 1)
N
& 2v 2
A-1))
Σ Σ Σ
Σ
p
(G), g¹ A... Gh = g² 8h =
H=] g
gee d
SES
V E. Mn(8.g' .a. "g.a""
(3)
2SG
+ (AAA 1) (e(1 1 -
N
< 2
In
5
Σ
Σ
Σ
Σ
H
gieg,
SES
46

---

Published as a conference paper at ICLR 2026
8 In
2SQ p (G) g ^ ... a
s =
N
28G
25g
+ (A (A- 1) In (c(1+ 1)
81
In
N
8
Σ
Σ
Σ
he[Hi [1 gLEg,).
J.
9d yg.. SES
P (G), g¹
C% 8h s
In 250
2Sg
(A - 1) ( l) +
- 1) VS3
N
81
In 8
A-
Σ
Σ
Σ
Σ
p (G), g¹ A...A
a Gh gD ^ sh =
M| g g'eG gi
90 gg.h
SES
Sğ In 259
25g
+ (A - 1) [e(((((+ A ))1))²²
N
81
Σ
N
5
hE[H1
81
28G SHG In In
8
25g
+ (A- 1)
1) in ( (1+ A- 1
N
N
8
where at (I) we use that, if E. Mn (n,g ""..,........."...... ".. = 0. then:
M,
Σ AeA
T ('ls
"...s,g.......f. ,²,aa) 112 Mn ....g.,......... > 0}
En Ma(s.g .....
Mn(s,g' .......',, ".aaa)) a') 01)
1
T (als, .'...)))
A
< 2,
we the total variation distance between probability distributions cannot exceed 1. Instead, if
End MM(8, ......" Ih, a') > 0, conditioning on En. Mr(8,9' ....g""............... a'), at all 8, "........ 9d where P (Gn 9 A 8h 8) > O, we note that Ma (s, g gr..g.......... .....,. g a)/ E., Mn. Mn(s, ...., gd a') is the empirical vector of probabilities of TTg (als,
sss g" ........ gg)) (recall its definition from Eq. 21 thus we
can apply Lemma 8 of Kaufmann et al. (2021) to get that, for any å E (0, 1):
Mr(s, ,..., Σ. Mn ......................
Ipr (KL
TRR STR",0--3,.8 ...., 9n)
....., ndh a²)
E. Ma "
In (A-1) Egg (M
A-1
3 1 - 5.
:.... (...".......... g".
Combining this result with the Pinsker's inequality, that tells us that la - y/1 & 22 [. (y), and with a union bound over all E [H]] 8 S, g² E Grg.n...., .g E Grg.ho h we get the passage in (I)
w.p. 1 5/2. Note that add an additional 2 for the case ENr Mu(s, "t..sgg..,.,00 ...... "., ,.aaa
0, and
define := EmeTHn [Tae nr At (2) we bound E,, Mn g ..... a') & N, and bring that quantity outside, at (3) we apply Lemma A.I of Xie al. (202I). after having noticed that Ene Ma(s.g'.. ......" gth a²)
Bin NN (G), g² A... G g Sh. s)) and make it hold for all 8, 3....""'"" g h w.p. 5/2. At (4) and (5) we apply the Cauchy-Schward inequality.
47

---

Published as a conference paper at ICLR 2026
Now, we impose that this quantity is smaller than
SHGh In 2SC
In ô ()
2Sg
(A-
N 1-1))=
E
64SHGhi In
250
64SH
2507
eNV
A A-1 +e).
e²
e²
Thanks to Lemma J.3 of Lazzati et al.
(2024).
we know that this inequality is satisfied with:
128SH In
280 192S
1)
280 In
128eSHG In
280
N <
8
+ A 1-1. 1.
6
22
62
Rearranging and applying final union bound concludes the proof.
C.4.3 PROOF OF THEOREM C.8
Theorem C.8. E (0, H] and (0,1). Let M be any MDP,R, R = (r-, ,.....a] any set of d 3 1 rewards containing the unknown r and x E IINM be any policy. Assume that the optimization problem in Line is solved exactly. Then, choosing A = c/(4H), with probability 1 8. the policy से output by Algorithm satisfies W(1 g n & e, with:
H
N E
In
(23)
2
Proof. Observe that, in the same way as in the proof of Theorem 4.4 or4or 5.1
it is simple to see that:
max W
nr .
H0/2 + Hé,
PE
with probability by using:
2d
N &
In
2(e))2
8"
data, where we made a union bound over the rewards in RD. Then, conditioning on this event and proceeding as in the proof of Theorem 5.2 we have:
2 w n 7
max W PER Th m.,
())
max W
Th 1/ .n))
+ max W (Trs niae rER
max W FER
(n TTo
(2)
He' HO
max
W (Trs nhe) rER
(3)
Hé + HO + min
max W Prr
mEll(RO)
reR
& H + H0 + max VER
W Arr m
< 2H + 211,
where at (1) we use triangle's ' inequality, at (2) we use that event above holds and Lemma C.I at (3) the definition of म. (4) upper bound the minimum with specific choice of reward in NTNM i.e., E and finally, at (5), we apply again that event holds and Lemma C.1
If we choose 0 c/(4H), and é' c/(4H), get that, with probability 1 5, the claim of the theorem holds with data:
8H2
2d
N
In
2
48

---

Published as a conference paper at ICLR 2026
C.5 WHEN 7 IS LINEAR IN KNOWN FEATURE MAP
In this appendix, we consider a variant of the known-reward setting, in which r is unknown, but have knowledge of dd-dimnnsional feature map d S A [-1,+1]e such that the expert reward ,.d can be written as:
A (s, a) (Ø(s, a). wh) Σ a.(s,a)una, a)wr,, ie[d]
for some unknown vectors Sh E e[-1,+1j4 Note that the linear reward setting is common in the IL literature (Abbeel Ng. 2004). We consider the following robust variant of RDM for this setting:
a arg min TEll w -1]
max
W (Tow:
Tew
(27)
where notation Φw denotes the reward obtained through the dot product between & and w.
We now sketch how this setting can be easily addressed through the technique presented in Appendix
First, consider each known feature oi S x A - [-1,+1] reward function, and define the set of "rewards" 4 ==
}. Then, consider the set of policies II(@) using definition in Appendix C.4 and let πT. be the policy defined as in Eq. 2T. Then, from Lemma C.5 we have the guarantee
that:
H
H
Σ
(shi an) =9) = =p=
-))=
(Σ
Φ (sh, an) = [-, Vie [d].
As a consequence, we have that, for any wh E [-1, + -11 it holds that:
10 V (g)
4g.
since the cumulative feature map collected is the same. Simply put, this means that set IT (and also policy Th) suffice for this "linear" variant of the RDM problem. However, as the feature map arbitrary, (1( might be too large. Thus, we may want discretize.
Extend the discretization approach of Section 4.1 to "rewards" in [-1, +1] and define set Φe :=
(0) ... dg). Then, observe that the policy Tœe satisfies a variant of Lemma C.6
W
T Now Tow)
max 111111
(w(12n" To0 To00 W (loo ( + W W Nooww vobu) w (onoow now
max
w:: (=[M)-[-1,+]]
2dH20 +
max -11
W
Foa Now Noww
2dH
where at (1) we use triangle's inequality and denote du the feature map which discretized each dimension, at (2) apply Lemma Lem twick and observe that, using Holder's inequality and the
definition of discretization: max low @gw| | & max& 10llell1 11 & d max, cello & dH0/2,
and at (3) use the aforementioned property (i.c., Lemmaa C.5)
Therefore, the policy TRo, and so the set II(*)), suffice for this new robust RDM problem. It is
immediate to extend RS-BC and RS-KT to this setting by simply extending Algorithms 3 and 4 using
input rewards taking values in 1-1,+1], and we are done. Then, by adjusting Theorems C.7 and
C.8 to keep track of this small variation, we can have also theoretical guarantees. Specifically, for the number of samples in Eqs. (22) and (23) we can guarantee that the policy output by the newly
constructed algorithms has a return distribution close to that of <099 for any w, which in turn has a return distribution close to that of the expert for any w as shown above. To do this for the variant of RS-BC, note that just need to extend the proof of Lemma C.10 The crucial insight in doing so is that both Tge and our estimate play actions with same probability at all trajectories with the same cumulative discretized features. Regarding the variant of RS-KT, the proof is immediate as we just need a union bound over all the features.
Modulo some slight difference as rewards are in [0, 1] but features are in [-1, +1]. We upper bound with additional factor of to keep into account that rewards are in [-1,+1]
49

---

Published as a conference paper at ICLR 2026
C.6 GENERALIZATION TO ARBITRARY PROBLEMS
In this appendix, we sketch how to extend RS-BC and its analysis to arbitrary IL problems of the following kind, in which we aim find a policy á that minimizes:
arg min TelMMM
Σ
Σ
I (wh E
pa p= x))
KE|H]
SEAA
where X' = 1h is any partition of the set of trajectories satisfying a certain property, specifically:
U x = P4+1 zEAr,
th [H],
x x n' 1 Va, e Nr, thE [H],
Væ E V(s,a) 6 S A, 3r's E Xn+1, Vw Eee w.s- vh [H 1].
where w a denotes the trajectory obtained by concatenating w with s, a. First, define:
1 (8h Oh a, Wh-1 xn-1(w))
x x (alsswa) w)
==
(28)
px (Sn = 8, Wh-1 (nw)(w))
when the denominator is not O, and 1/A otherwise, and (.r()) denotes the set of trajectories which belongs. Then, this policy satisfies:
Theorem C.12. holds that:
pax (sh Wh-1 E
p (sh wh-1 E x))
wh [H se, Xn-1-
Proof. The proof follows that of Lemma C.2 and C.9 Specifically, prove it by induction. At 1, note that trivially holds as X, Ø. Now, make the induction hypothesis that at any < h. have:
PEx (sw = = Wh-1 E x) = 1 (sw = s ww-1 Ex)
S S. Aw-
Then, at h, for any s E S and x' E X4-1, we can write:
pax (sh s' Wh-1E21 Σ s.w.s.aEz))
5,a,a
$ w)
Σ ΣΣ I(u.ssaaE
zEX, WEZ S.
i Σ /e
8,0
Σ (sh-1 =
Σ
Σ Selessssee//////esoaaaaa.a..
x 8,0
pax ( =
P
Σ Sebessaeeppppp//.ee0x
zEA
pa = -
Σ aa
XE 8.G
(8h-1 0h-1 -
50

---

Published as a conference paper at ICLR 2026
Ip= (sh wh-1 E x²),
where at (1) use notation x to denote that all the trajectories in when combined to a given 8, a give birth to the same x' by hypothesis, at (2) we use the induction hypothesis, and at (3) the definition of xxx
Then, RS-BC can be easily extended by counting the number of occurrences in each set of trajectories, and also the theoretical guarantees can be easily extended this setting.
D ADDITIONAL RESULTS AND PROOFS FOR SECTION
Theorem 5.1. Let E (0, H and ò e (0, 1). Let M,, be any MDP and π e IIMM any policy. Then, choosing (/H2), a number of samples
N δ SAH
2 --nnn 3)
(11)
suffices to guarantee that, with probability at least 6, for the estimator ()) = DWEDE ([G(W;))) 9) Vg,r, r, have:
W W(n 1r (r sc
E
max
Ax (0.00
Proof. For any reward r S x A x
H - [0. 1]. we can write:
w w(n " "
7r
(1)
< W
Th
-ra W -(TT VTo :7r)
@ 110
+ W Tro (Tr)
( H0
+ Hc',
where (1) we use triangle's inequality, at (2) we apply Lemma C.I and at (3) we use the same derivation as in the proof of Theorem 4.4] after having noticed that the estimate û, is the same estimate used in [f RS-KT. So, by the DKW inequality, we have that the last passage holds with probability 1 using number of samples:
& 2(6")2 8
This holds for a single reward r S x x [H] [0. 1]. To make this hold for any possible reward, observe that it suffices to guarantec that it holds for all the rewards in the set R, defined the set of all reward functions taking on discretized values:
R = S S Ax H }
Indeed, R represents an / 2-coerrng of the set of all the real-valued reward functions. Therefore, the result follows through the application of a union bound over all the rewards in R. Since they [D SAH then we obtain a number of samples:
2SAH
2
SAH's 2²
N
In
Ö
In
)
to guarantee that W or 1 + 1/0 1 H/c.
(7r) & c for all r, where we set é' c/(2H), = H/H, and D1
L
Theorem 5.2. Under the conditions of Theorem 5.1 assume access computational oracle that takes as imput the dataset DE and the transition model p, and outpuls solution to
â arg min aEllar Ax -f0.1]
max
W nF 9r
(12)
Then, with probability at least
8, using the number of samples in Eq. it holds that
W
(177 1r .n²² & 2c.
max
51

---

Published as a conference paper at ICLR 2026
Proof. Define the good event & as:
W (1F ( {
max
Ax
Then, under ε, it holds that:
max x
W Tr 1/
1
max
W( TT .
max
W Arr Dz
r: xxxxAx[H]-
(2)
W( Trr ni
max
[0.1]
(3)
W (frr W(rri
E min
max
#EIINM r:S AX
(4)
66+ max
W
Trr 177
r:SxAx [M] [0.1]
(5)
2e,
where at (1) we use triangle's inequality, at (2) we use that event & holds, (3) we use the definition of f. at (4) we upper bound the minimum with specific choice of reward in II i.c., π and finally, at (5), we apply again that event holds. The proof is concluded by applying Theorem 5.1 which shows that, with the samples in Eq. 111. the event & holds w.p. 8.
E ADDITIONAL DETAILS ON NUMERICAL SIMULATIONS
We describe how we sample an MDP and an expert's policy in the experiments in Appendix E. In Appendix E.2 we discuss the (immpossibility of implementing a variant of RS-KT that works with Markovian policies, in Appendix E.3 we provide implementation details on MIMIC-MD, in Appendix|E.4 E.4 we provide a graphical example of expert's return distribution and its estimates from the four algorithms considered in Section 6 and in Appendix E..3 we provide a simulation comparing RS-BC with the W-RS-GA IT. (Lacotte et al., 2019 Finally, in Appendices [.66 and E.7 we provide additional details and results on the simulations conducted to address, respectively, questions and 2. We mention that all simulations took place in some hours on AMD Ryzen 5500U processor.
E. DETAILS ON SAMPLING THE MDPs AND THE EXPERT'S POLICIES
All the MDPs are sampled follows. The initial state so S is sampled uniformly random from S. The reward function F. : S A x [H] - [0, 1] is sampled, in each s, a, h, uniformly at random from set (0, p. 2p,... -.. [1/ple), whose values are controlled by parameter p (0, 1] (intuitively, the difference - gives insights into the approximation error). The transition model p is obtained in two steps. First, we sample it uniformly at random each s, a, h from the simplex A S (by sampling from the Dirichlet distribution), but then we make the transition of each 8, a, h deterministic with probability 0.7 increase the variety of the MDP.
Regarding the expert's policy, when we sample Markovian expert's policies, we simply sample them uniformly from the simplex A A in every S x [[]]. Instead, when we mention Morn-Markovinn" expert's 's policies, then, of course, we cannot sample them uniformly at random from 11 due to the curse of dimensionality. Instead, what we do is sample randomly from parameteric subset of NTNM defined as follows. We map states {81,..... 85} = S to (O.....S- 13. Then, each policy projects each past history (#1, a1, 82, a2, 5h, an) to R
S-13, 1},
and actions (a)...., aa A 8b,
to {O..... after having mapped it to integers and padded with zeros to reach size 2H, by using projection matrix of size (2H, 16) randomly generated. Then, we obtain the probabilities of playing each action by multiplying such 16 (6-dimensional vector by matrix of weigths of size (16, A) corresponding to the current state 8h. (we randomly generated one of these weight matrices for every state).
E.2 ON ADDRESSING RETURN DISTRIBUTION MATCHING WITH MARKOVIAN POLICIES
We mention that finding the best Markovian policy n' e II with return distribution closest given array n, ie.. addressing min W( 7), seems a problem that cannot be formulated as
52

---

Published as a conference paper at ICLR 2026
an LP and not constraints. so, generating random MDPs and (non-Markovian) expert policies to understand how much Markovian policies outperformed by Mon-Marovvian policies for return distribution matching may be difficult. We helieve that understanding more in-depth this fact will be interesting for future works.
general convex optimization problem, it seems to require bilinear
E.3 IMPLEMENTATION DETAILS OF MIMIC-MD
Algorithm MIMIC MD has been devised by Rajaraman et al. (0200) but an efficient formulation
is provided in Rajaraman al. (0221 (see its proof of Theorem 2). Although it concerns only deterministic expert's 's policies, we extend to stochastic policies as the full algorithm as follows.
Simply, in addition to the constraints for being valid occupancy measure, replace constraint (s.(s.)) Na(a) Σn dr (8, a')
(s(8, a) {a o (s)} (valid for deterministic policies) with d, a)
where (s) * 0 to extend to stochatic experts. Note that it is linear in d.
E.4 AN EXAMPLE OF EXPERT'S RETURN DISTRIBUTION AND ITS ESTIMATES
In Figure we plot the (discretized) return distribution of the expert's policy 7f, (in green, ) obtained after having sampled at random an MDP with size S, A, H (3, 2, 8) and the policy as well. We discretized with bins at distance of 0.5 for computation and plotting.
Then, we have sampled a dataset D of N = 300 trajectories from nh and given it (and potentially the transition model) in input the four algorithms considered in Section|6) i.e RS-BC, RS-KT, BC and MIMIC-MD, obtaining the return distributions 17 7R85-BC, .1 RR .7th and M/MMMIC.MD.
Observe that the return distributions of both RS-BC and RS-KT are close to 17. Instead, BC and MIMIC-MD. MD, by working with Markovian policies, are biased, and their return distributions are not that close to 17
0.30
0.25
MIMICND
Probablity
0.20
0.15
0.10
0.05
0.00
Return values
Figure 3: An example of expert's return distribution and various estimates computed with RS-BC, RS-KT, BC and MIMIC-MD.
E.5 COMPARISON WITH LACOTTE ET AL.
2019
The algorithm W-RS-GAIL. (Lacotte et al. 2019) aims to match both the expectation and the CVaR at some given level a of the return distribution. Crucially, addresses this problem using Markovian and stationary policies, and function approximation. To make the comparison with RS-BC fairer, implement W-RS-GAII IL without function approximation, by using SA "parameters". In addition, note that RS-BC does not dynamics information, i.e., is purely offline, while W-RS-GAI GAIL interact with the environment. At the same time, however, RS-BC has access to the reward of the expert, while W-RS-GATH does not.
53

---

Published as a conference paper at ICLR 2026
N 100
= 1000
RS-BC
0.045#0.022
0.025 + 0.017
W W-RS-GAT GRRS-GAIL L, 0.3
0.226 0.143
0.22 = 0.147
W 1-RS-GATI L,, a = 0.7
0.202 0.122
0.197 # 0.123
Table 2: Results of simulation in Appendix
E.5
N = 20
N 80
N = 300
N = 1000
N = 10000
RS-BC
0.101=0.041
0.059:0.017
0.032:0.011
0.016#0.006
0.005=0.002
RS-KT
0.164#0.04
0.084x0.021
0.052#0.013
0.033+0.01
0.02=0.006
BC
0.10420,041
0.05820.018
0.03520.012
0.02440.01
0.018=0.009
MIMIC MD
0.13920.058
0.079:0.028
0.045#0.016
0.029#0.012
0.019=0.009
Table 3: Results of simulation with S, A, H 550,5,5) for Q1.
N 20
N 80
N = 300
= 1000
= 10000
RS-BC
0.19320.086
0.087=0.035
0.04620.019
0.027=0.01
0.012#0.005
RS-KT
0.223=0.066
0.115=00033
0.07220.023
0.05320.017
0.041=0.018
BC
0.208=0.08
0.162#0.083
0.15620.086
0.151#0.086
0.15100.085
MIMIC-MD
0.265+0.106
0.18+0.086
0.159=0.084
0.153+0.087
0.15=0.085
Table 4: Results of simulation with A, 22, 20) for QI.
Specifically, model both the cost function and the policy to learn æ by using a parameter w,,a and 0,aa for every possible (s, a) E S x A. For the policy, we compute the probability using a softmax: (a(a)s) e" En " so as to avoid the need for normalization after every gradient step. Note that
the gradient in 22 of the log policy is 10
d er.a
dTo(als) d6za
En e'
11s 8) (11a=a) - a}
d03.a
En ee
For the sake of simplicity (and for avoiding extending results from the discounted mniiiit-horizonn setting to ours), we drop the causal entropy term. Moreover, we implement the updates of both the cost function and the policy parameters through gradient ascent/descent, avoiding the usage of Adam and KLLconstrained gradient descent step w.r.t. a linear approximation of the objective.
We generate at random 50 MDPs with S A 2, H 5 and 50 Markovaoo stochastic policies. For each environment, for each number of trajectories in N (100, 1000). we generate 2 datasets of N expert's trajectories. We then run RS-BC and two different versions of W- the datasets; estimate the return distributions, compute the Wasserstein distance and finally compute the average distance among the environments and the datasets. The results are reported Table 2
on
The two versions of W-RS--O- G considered differ for the level of the CVaR desired, where one uses a = 0.3, and the other a = 0.7. Both share the values of other hyperparameters: we set À (to balance the CVaR and the mean objectives), learning rate 0.0005 for both the cost and the policy parameters, 3000 number of iterations, and 500 trajectories collected at each training loop.
By looking at Table since by increasing the number of expert trajectories from N 100 to 1000 the error of W--S--AAIL does not decrease, we observe that of the policy class adopted by this algorithm (Markovian and stationary policies), as well as the fact that it does not aim to match the whole return distribution, but just the CVaR specific level This limited expressivity is expected (e.g., see Proposition 3.2)
realize that RS-BC outpeforms both versions of W-RS
Moreover,
a matter of sample efficiency, but of limited expressivity
E.6 ADDITIONAL DETAILS ON Q1
For simulations, we set p 0.03 < 6, meaning that there is some approximation error.
This is used for the update of the policy parameters.
54

---

Published as a conference paper at ICLR 2026
N
80
N 300
1000
10000
RS-BC
0.081=0.036
0.035:0.016
0.019:0.011
0.011=0.005
0.004=0.002
RS-KT
0.095=0.042
0.043#0.019
0.024#0.012
0.013=0.005
0.004#0.002
BC
0.104=0.053
0.07640.048
0.068+0.048
0.066+0.049
0.065=0.049
MIMIC- C-MD
0.1320.063
0.085+0.046
0.071#0.046
0.068+0.049
0.065=0.049
Table 5: Results of simulation with A, 2,,
and 0 p for Q2.
N = 20
N= = 80
= 300
= = 1000
10000
RS-BC
0.088=0.026
0.048+0.022
0.025+0.011
0.012+0.005
0.00420.002
RS-KT
0.135=0.037
0.068+0.02
0.034#0.01
0.019#0.005
0.00620.002
BC
0.092:0.032
0.05440.029
0.038+0.023
0.031#0.022
0.028=0.024
MIMIC C-
0.118=0.042
0.06640.028
0.04420.023
0.03340.021
0.029=0.024
Table 6: Results of simulation with S, A, H
320,3,5) for Q2.
= = 20
N = 80
N 300
N 1000
N = 10000
RS-BC
0.177=0.067
0.091=0.038
0.047:0.018
0.023+0.008
0.008=0.003
RS-KT
0.224=0.083
0.109
0.057=0.018
0.031#0.01
0.011#0.004
BC
0.196#0. 104
0.15940.100 102
0.148+0.100
0.145+0. 104
0.14420.
MTMTC-M
0.246#0.115
0.17440. 103
0.151#0.103
0.145+0. 104
0.144 106
Table 7: Results of simulation with S, A, H (2, 2, 20) for Q2.
Regarding Table 3] would like to discuss some points. First, mention that the increase in size of S and not sufficiently big to permit to RS-KT to outperform the sample complexity of RS-BC, discussed in Q4. Second, increasing S, A, H makes RS-KT much more time-consuming. requires solving an LP with much more variables and constraints. Third, RS-KT performs comparably BC and MIMIC-MID MDD but this is due to increment of approximation error due the increase of S, A, as clear from Table 6 where in absence of approximation error RS-KT outperforms BC and MIMIC-MD. Note that this is not the fact for RS-BC, which seems more robust to approximation error for this problem size (intuitively, the reason is that it is strictly more expressive than BC for any choice of 8).
Regarding Table 4 we mention that larger H implies a larger approximation error in particular for RS-KT, as clear from Table where absence of approximation error RS-KT outperforms BC and MIMIC-MD.
E.7 ADDITIONAL DETAILS ON Q2
The three additional simulations have all a non-Markovian expert and 0 (to enforce no approx- imation error) with parameters S, A, H, ((2, 2.5. 5e 2), (20, 5e 2), (2, 20, 1e 1)),
and the results are reported respectively in Tables 5 and
By comparing these tables respectively with Table [ (top), and 4 where there is approximation emror due to 0 = 0.05 > 0.03 = we realize that the approximation error mostly concerns RS-KT and with larger horizons (as expected from Lemma 4.2 and from knowledge of how RS-BC works).
55


Reviews:
[{"review_id": "cBFPRobGHO", "reviewer": "Reviewer_P3jS", "summary": "The paper studies training a risk-sensitive reinforcement learning agent through imitation learning. The classical imitation learning methods only focus on matching the expert's expected return (average performance), which may not capture risk attitudes. To overcome this, the authors propose matching the entire expert return distribution, such as variance and tail risks. The approach formulates a return distribution matching (RDM) problem using the Wasserstein distance between distributions.\n\nThe authors propose two provably efficient algorithms, RS-BC (Risk-Sensitive Behavior Cloning) for unknown transition models and RS-KT for known transition models that learn a class of non-Markovian policies that efficiently express risk-sensitive behaviors. These algorithms improve over standard imitation learning by matching the full return distribution, accommodating risk attitudes. They provide theoretical guarantees on sample complexity and policy approximation, and then empirically validated that RS-BC and RS-KT outperform baseline methods.", "questions": "Please address the weaknesses of the paper. I am open to changing my score based on the authors' responses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "**The following are the strengths of the paper:**\n1. This paper considers training a risk-sensitive reinforcement learning agent through imitation learning while matching the entire expert return distribution to capture the risk attitudes. \n\n2. The authors propose two provably efficient algorithms, RS-BC and RS-KT, that learn a class of non-Markovian policies that efficiently express risk-sensitive behaviors, while having theoretical guarantees on sample complexity and policy approximation.\n\n3. Finally, the authors empirically validated that RS-BC and RS-KT outperform baseline methods.", "weaknesses": "**The following are the weaknesses of the paper:**\n1. The proposed algorithms are restricted to finite tabular RL environments, limiting their scalability to RL problems with large or continuous state spaces. \n\n2. It is unclear how expensive Step 2 of RS-KT is, especially for large LP. How practical it is to assume the expert's reward function (also the transition model for RS-KT) is known.", "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": "HhMohf28N9", "reviewer": "Reviewer_sUz6", "summary": "The paper studies the problem of risk-sensitive imitation learning. The goal is to train an risk-sensitive agent that matches the expert’s return rather than merely the expected return. They show that Markovian policies are in general insufficient for this task,  and introduce an efficiently storable subclass of non-Markovian policies based on discretized cumulative reward. They propose  two provably efficient algorithms, RS-BC and RS-KT, when the transition model is unknown and known, respectively for the tabular setting. They prove sample complexity guarantees for both algorithms, show a polynomial-sample result for the unknown-reward robust variant. They also provide simulation evidence that RS-BC and RS-KT better match expert return distributions than standard baselines in a range of synthetic MDPs.", "questions": "questions/comments\n-  Could the authors provide an experiment comparing RDM to a strong CVaR-matching IL method across a range of \\alpha levels, to show where full distribution matching yields material improvement?  \n- how to choose \\theta in practice is still not very clear\n- Lower bounds is missing (the authors have noted this as a limitation)\n- There are some small typos and places where exposition could be tightened \n- Missing related works, to name a few,  sample complexity/regret guarantees on risk-sensitive RL [1][2], CVaR RL [3][4], risk-sensitive distributional RL [5][6]\n\n[1]Fei, Yingjie, et al. \"Risk-sensitive reinforcement learning: Near-optimal risk-sample tradeoff in regret.\" Advances in Neural Information Processing Systems 33 (2020): 22384-22395.\n\n[2]Fei, Yingjie, et al. \"Exponential bellman equation and improved regret bounds for risk-sensitive reinforcement learning.\" Advances in neural information processing systems 34 (2021): 20436-20446.\n\n[3]Bastani, Osbert, et al. \"Regret bounds for risk-sensitive reinforcement learning.\" Advances in Neural Information Processing Systems 35 (2022): 36259-36269.\n\n[4]Wang, Kaiwen, Nathan Kallus, and Wen Sun. \"Near-minimax-optimal risk-sensitive reinforcement learning with cvar.\" International Conference on Machine Learning. PMLR, 2023\n\n[5]Liang, Hao, and Zhi-Quan Luo. \"Bridging distributional and risk-sensitive reinforcement learning with provable regret bounds.\" Journal of Machine Learning Research 25.221 (2024): 1-56.\n\n[6]Chen, Yu, et al. \"Provable risk-sensitive distributional reinforcement learning with general function approximation.\" arXiv preprint arXiv:2402.18159 (2024).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem reformulation of risk-sensitive IL is clear and well-motivated. Since return contains for relavant distribution  for risk-sensitive decision-making, matching return distributions can generalize expectation/CVaR matching to general risk-attitude imitation.\n- The paper is well organized and the main ideas are presented clearly. The transitions are fluent and follow a logical flow.\n- All three algorithms are conceptually simple and intuitive. Both RS-BC and RS-KT are computationally efficient and  come with  sample complexity guarantees. For the known reward setting, the RS-KT algorithm attains a strong sample bound independent of S,A. These sample complexity results with different dependence regimes, e.g., with and without transition knowledge are novel and interesting.\n- The paper provides empirical validations across many random MDPs, which shows the proposed methods reduce Wasserstein error relative to standard baselines.", "weaknesses": "- One major limitation is the computational cost of RS-KT. It requires solving an LP with size scaling as $O(S A H |Y_\\theta|)$. For moderate/large tabular MDPs this can become expensive.\n- Unknown-reward setting lacks practical algorithm. Theoretical sample-efficiency for the unknown-reward case is established by using an oracle-based argument, but a concrete practical algorithm for this setting is deferred. This is an important practical case since expert reward unknown.", "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": "TCSEBPHugF", "reviewer": "Reviewer_9FFJ", "summary": "The paper presents the risk-sensitive imitation learning problem as matching the expert's return distribution in Wasserstein distance, and designs two types of algorithms, RS-BC (without model, based on behavior cloning non-Markov strategy class) and RS-KT (known transition, constructing LP search occupancy measure to match distribution), for the cases with/without known transition models. It provides theoretical sample complexity bounds (Theorem 4.3/4.4/5.1/5.2) and conducts a large number of simulation experiments.", "questions": "1. The RS-KT algorithm needs to solve an LP with a scale of $\\mathcal{O}(SAH|\\mathcal{Y}^{\\theta}|)$, which seems difficult to scale in practice. Have the authors considered extending it to an approximate method in the continuous domain (for example, using function approximation to replace tabular occupancy measures)?\n2. Can the algorithm be extended to an online interactive version, which can learn unknown transitions and unknown expert rewards from the environment. Maybe the author can add more discussion on this approach and the comparison with risk-sensitive RL.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Directly promoting risk-sensitive imitation learning to return distribution matching is more comprehensive in depicting \"risk attitude\" than merely matching expected value or several quantiles, which has both theoretical and intuitive appeal.\n2. This paper proposed a  subclass of non-Markovian policies $\\Pi(r_\\theta)$ that can be stored and computed. The sample complexity bounds for RS-BC (no interaction) and RS-KT (known transition) are provided, with clear analysis and comparison. The theorem gives upper bounds for various scenarios.\n3. The experimental design is clear, aiming to answer four specific research questions (Q1-Q4).", "weaknesses": "1. This paper proposes the main algorithms (especially RS-BC and RS-KT), which all assume that the expert's reward function $r^E(s,a)$ is known, or at least the reward for each step can be accurately calculated from the expert's trajectory. Section 5 of the paper explicitly points out that this is an important limitation, and it specifically discusses the \"unknown $r^E$\" case, but this part only stays at the theoretical analysis level (oracle-based) and does not propose an actual runnable algorithm.\n2. The theoretical guarantees depend on strong assumptions (e.g., known transitions, exact LP solution), yet the experiments use approximations without quantifying the induced errors, leaving the practical scalability and validity of the theoretical bounds unclear.", "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": "I5MeluvTa7", "reviewer": "Reviewer_XjG5", "summary": "The paper addresses the problem of risk-sensitive imitation learning (IL) in tabular finite-horizon Markov Decision Processes (MDPs). Traditional IL focuses on matching the expert's occupancy measure, which only imitates the expected return (average performance) and is inherently risk-neutral, ignoring aspects like variance or tail behavior in the return distribution. The authors reformulate IL as \"Return Distribution Matching\" (RDM), where the goal is to train an agent that matches the expert's entire return distribution (encoding risk attitude) in Wasserstein distance, rather than just the expectation or limited risk measures like Conditional Value at Risk (CVaR). They consider settings where the expert's reward is known (primary focus) or unknown, and the transition model is either unknown (no-interaction/offline) or known. The aim is to develop provably sample-efficient algorithms that handle non-Markovian experts, who may exhibit complex risk-sensitive behaviors not capturable by Markovian policies.", "questions": "1.It is not clear what about Algorithm 1 is ‘risk-sensitive’. It seems to just be learning a policy, without a clear link to the return distribution.\n2.The paper could benefit from a broader discussion as to how the proposed methods connect to the distributional RL literature. For example, how does the proposed method relate to a categorical or quantile representation.\n3.Since we know the expert's rewards, wouldn't it be better to solve the forward RL with the known reward to completely restore the expert's policy? What are the benefits of the proposed methods over the forward RL?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The analysis carried out is sound and well-motivated. \n2. The paper is polished and easy to read. \n3. The problem setting is interesting. \n4. The empirical results are encouraging.", "weaknesses": "1. Primary focus assumes known r^E, which may not hold in real IL (e.g., human demos); unknown-reward variant relies on oracles, potentially impractical. If the expert’s rewards were known, we could have solved the forward RL, which allows us for the full access to the rich literature: function approximation, generalization, etc. Assumptions: Requires known rewards (or oracle for unknown); no online interaction for learning transitions.\n2. There is no compelling empirical evidence presented in the main body that validates whether the correct distribution is being learned (the results presented in the main body only show the overall performance). \n3. There is a figure in the appendix that seems to show the learned distribution, however this figure is hard to understand, and can likely be significantly improved. I strongly encourage the authors to improve this figure and include it in the main body.", "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)