Learning to Answer from Correct Demonstrations
Authors: Nirmit Joshi
Gene Li
Siddharth Bhandari
Shiva Kasiviswanathan
Cong Ma
Nathan Srebro
Venue: ICLR 2026 Poster (2026)
Keywords: Imitation Learning
Contextual Bandits
Likelihood Maximization
Primary area: learning theory
TL;DR: Taking a learning-theoretic view of SFT, we rethink existing modeling assumptions. Under relaxed assumptions on the reward model's capacity, we show MLE fails and design a new optimal learner.

Abstract: We study the problem of learning to generate an answer (or completion) to a question (or prompt), where there could be multiple correct answers, any one of which is acceptable at test time. Learning is based on demonstrations of some correct answer to each training question, as in Supervised Fine Tuning (SFT).  We formalize the problem as imitation learning (i.e., apprenticeship learning) in contextual bandits, with offline demonstrations from some expert (optimal, or very good) policy, without explicitly observed rewards. In contrast to prior work, which assumes the demonstrator belongs to a bounded-complexity policy class, we propose relying only on the underlying reward model (i.e., specifying which answers are correct) being in a bounded-complexity class, which we argue is a strictly weaker assumption. We show that likelihood-maximization methods can fail in this setting, and instead present an approach that learns to answer nearly as well as the demonstrator, with sample complexity logarithmic in the cardinality of the reward class. Our method is similar to Syed and Schapire 2007, when adapted to a contextual bandit (i.e., single step) setup, but is a simple one-pass online approach that enjoys an ``optimistic rate'' (i.e., $1/\varepsilon$ when the demonstrator is optimal, versus $1/\varepsilon^2$ in Syed and Schapire 2007, and works even with arbitrarily adaptive demonstrations.

Published as a conference paper at ICLR 2026
LEARNING TO ANSWER FROM CORRECT DEMON- STRATIONS
Nirmit Joshi; Gene Li Toyota Technological Institute at Chicago ni rmi t@ttic edu gene@ttic.cm edu
Siddharth Bhandari, Shiva Prasad Kasiviswanathan Amazon
s i dddhatthbbaddrr engg@ coom kasivi see com
Cong Ma University of Chicago Department of Statistics congm@ucht cago. edu
Nathan Srebro
Toyota Technological Institute at Chicago nsrebroêttic..
ABSTRACT
We study the problem of learning generate answer (or completion) ques- tion (or prompt), where there could be multiple correct answers, any one of which is acceptable test time. Learning is based on demonstrations of some correct swer to each training question, as in Supervised Fine Tuning (SFT). We formalize the problem imitation learning (i.e., apprenticeship learning) in contextual ban- dits, with offline demonstrations from some expert (optimal, or very good) policy, without explicitly observed rewards. In contrast prior work, which assumes the demonstrator belongs to bounded-complexity policy class, we propose relying only on the underlying reward model (i.e., specifying which answers are correct) being in a counded-complextny class, which we argue is a strictly weaker assump- methods can fail in this setting, and instead present approach that learns to answer nearly well the demonstra- tor, with sample complexity logarithmic in the cardinality of the reward class. Our method is similar to Syed & Schapire (2007), when adapted to contextual ban- dit (i.e., single step) setup, but is a simple one-pass online approach that enjoys "optimistic rate" (i.e., 1/s when the demonstrator is optimal, versus 1/8² in Syed & Schapire), and works even with arbitrarily adaptive demonstrations.
tion. show that seetihood-maximizann
Please see the full version arX:v:2510.154:
INTRODUCTION
Many real-world problems involve generating an answer to a question, where there may be many equally good responses. A math question have millions of equally valid but differently written solutions, coding task can admit many different perfectly working implementations, and a recom- mendation query can be satisfied by multiple items. The learner's challenge is not to reproduce all correct responses, but to generate a single good answer.'
This recurring structure can be formalized as contextual bandits. The context and actions are the questions x E X' and answers y E V, respectively. In the simplest case, consider an unknown binary reward function r. (z. y) (0, 1} indicating whether a response is correct. We get demonstrations that are always correct, i.e., from training set S = {x ): where ~ DD yi #- ri), where D is some context distribution and 7(- x)) is demonstration policy supported on the set of optimal or correct actions for the context given by²
T. (z) == arg max (I, fy y (x. y) = L} yEy
(1)
for 0-1 rewards with correct answers
"The work was done while NJ was interning at Amazon. I 'Throughout, use "good" and "correct" interchangeably. One should think of "correct" in this context as "full credit complete answer", not merely one that "logically true". II the special case of binary rewards with always correct demonstrations, every question has least
also implicitly assuming answer. This required in the more general treatment.

---

Published as a conference paper at ICLR 2026
Note that (x)] can be huge and there may be many correct answers.
More generally, we consider real-valued rewards r. r.(2,,) y) E [0, 1], and do not require the demon- strator स to be optimal (i.e., to always answer in o. (z)). Our goal is to compete with the value of the demonstrator, where the value of a policy x : X + ())) is defined as:
(T) =
(r, y)].
(2)
For binary rewards, Vr, (T) is exactly the accuracy of the predictor л. Our goal learn a predictor पे that produces good responses to unseen contexts sampled from DD as captured by Vr, (7). In particular, we would like to learn a predictor whose value (i.e., accuracy) nearly as good as that of the demonstrator:
Vr. (+) 2 (T)
This is apprenticeship learning (Abbeel & Ng. 2004; Syed & Schapire, 2007), aka imitation learn- ing, in contextual bandits from offline data. We argue that in modern large language model (LLM) training, this is exactly the problem addressed during the supervised fine-- fine-tuning (SFT) phase, where the model is trained on questions with expert-generated good (perhaps near-optimal) answers, and the goal is to be able to answer as well as the expert.
An important aspect of our objective in Eq. (2) is that it is entirely reward-driven, focusing solely on utility. We do not require matching the distribution of IT. In most question-answering or prompt- response systems, we argue that the true objective is often reward maximization rather than dis- tribution matching. For example, in tasks such as producing gold-medal-winning the solution set a- (question) = {all completely correct solution texts) (i.e., reward = 1), is unfath- omably large, due both to variation in solution approaches and details, and to differences in the solution text itself, down to word choices, spacing, and punctuation. Rather than matching the gen- eration style of particular expert, or some arbitrary distribution over experts, producing any single correct solution is sufficient to win a gold medal. The same principle underlies the deployment of modern LLMs, where success is judged by whether the model generates single high-utiliig an- swer to the user's question (OpenAI, 2023; Anthropic, 2024; DeepMind, 2023; DeepSeekAl, 2025; Touvron et al., 2023).
IMO solutions,
Demonstrator Class Assumption and Likelihood Maximization. A common approach when learning from demonstrations is to assume that the demonstrator से belongs to low-capacity policy class IT. This motivates maximum likelihood estimation (MLE), or equivalently log-loss minimiza- tion, as natural learning rule; see, e.g., Foster et al. (2024). Indeed, in SFT for LLMs, one fits model by minimizing log-loss on prompt-completion pairs. Foster et al. (2024)-relying on the fact
that MLE enjoys sharp O log [T]] certain notion of distance- also enjoys low reward suboptimality, with tight rates for the value sub- optimality in (2) (see Proposition in Section 3, based on Foster et al., 2024). That is, not only does MLE achieve small value suboptimality; it does so by cloning the demonstrator's action distribution.
convergence guarantees in distribution from m samples under a
m
However, for this guarantee to be meaningful, log [II must be small, and the demonstrator must lie in this class. This requires strong assumptions about the exact behavior of the demonstrator for example, modeling the generative process of specific students producing solutions to math problems, rather than just modeling what means for solution to be correct.
Reward Class Assumption and Our Results. Instead, we propose modeling the rewards them- selves rather than the demonstrator policy class, and rely only on the reward class having low cardi- nality (Section 2). Concretely, we assume the unknown reward function r, : X' x y -> [0, 1), spec- ifying the utility of particular answer for given question, comes from known low-cardinality reward class R C [0, x))
"Response generation by LLMs can be seen Markov Decision Process (MDP). where the response token sequence and each token is an action. In this MDP, transitions are deterministic, the 'state' includes the full prompt and all previously generated tokens, and rewards are sparse typically only given the end. Thus we find treating response generation as a contextual bandit with the entire response y as a single action, as more useful abstraction for our purposes. Indeed, treating as single action is also the view routinely taken in the study of question-answering for L.LMs (e.g. Rafailov et al., 2023; Huang al., 2025; Wu et al., 2025, etc.). See also Remark 3.
2

---

Published as a conference paper at ICLR 2026
Assump.
low-cardinality II
low-cardinality
Rule
A Algg [1|
log [1]]
MLE
May not learn
m
Proposition
Theorems and 2
log 11
A log [ [
log
Our Learner
Corollary 2
Theorem
Table 1: The case of optimal demonstrations provides a link between the low-cardinality assump- tions on IT and R. The black text compares the sample complexity of MIE and our learner under both assumptions when the demonstrator is optimal, establishing that low-cardinality assumption on R is strictly weaker than low-cardinality II. The slower decaying terms for both learners under their respective assumptions are shown in gray; these interpolate between the slow I/ and the fast m rates depending on the demonstrator suboptimality
For example, one think of the reward given by encoder-only transformer that takes the quences x and y as input and outputs a single real-valued reward. In this view, the class R corre- sponds all possible transformers obtained by varying the model's weights. To simplify analysis and presentation, we consider finite cardinality classes R. and rely on the log-cardinality log [R]. The log-cardinality should be thought of as corresponding to the number of parameters-eae if the transformer has N parameters, each 32-bit float, then log |R| < 32...
We show that the reward class assumption is, in a sense, weaker assumption than the demonstrator class assumption, and that it motivates looking beyond likelihood maximization/log-lon minimiza- tion and suggests different learning procedures:
In Section 2, we show that assuming the reward lies in low-cardinality class strictly weaker assumption, at least when the demonstrator 7 is optimal for the unknown reward r., compared to assuming the demonstrator lies in a low-cardinality policy class.
In Section 3, we show that MLE can fail to generalize for low-cardinality reward classes in simple situations (Theorems and 2).
In Section 4, we present learner that succeeds with optimal O(log [R|/m) samples in competing with an optimal demonstrator, and this sample complexity degrades grace- fully when the demonstrator is suboptimal via "optimistic rate" (Theorem 5). The guar- antee has no dependence on the size [)%, or the size the optimal action sets (o. (z)].
See Table for a comparison between MLE and our learner under low-cardinality assumptions on II and R. In Appendix A.3, we discuss the important distinction between distribution matching (which MLE attempts to achieve) and alternative approaches (like ours) that do not rely on dis- tribution matching to obtain good performance. In particular, distribution matching is impossible under the Reward Class Assumption; thus, any approach that ensures low subppiiubot must do so without relying distribution matching (unlike MLE). Finally, in Appendix F, also dis- cuss a @ extension our learner, which is minimax optimal for the pass@k metric when the demonstrations are optimal.
Relationship to Syed & Schapire (2007). While preparing this manuscript, we learned of this lated work in the more general setting of known, and (possibly) stochastic transitions. The method and analysis of Syed & Schapire (2007) can be adapted to contextual bandits, yielding method
similar ours, and already guaranteeing the reward suboptimality of spondence in Appendix A.2.
log [R] m2
We provide corre-
In this paper, emphasize the application to contextual bandits, or equivalently in the known model class of MDPs with deterministic transitions (see Remark 3), and contrast with policy-class clas--eeas based likelihood maximization methods. In our setting, method enjoys optimistic rate (as in Ta- ble 2), and allows us to argue strict dominance over MLE the case of an optimal demonstrator. Our result is derived from regret-minimization algorithm for the online setup with adversar-
3

---

Published as a conference paper at ICLR 2026
ial sequence of demonstrations (see Figure 1). Additionally, in the statistical setting, our learning algorithm is a one-pass online algorithm, as opposed to a multi-pas pailti-pass batch approach, with a self- contained and simpler analysis of fast rates; with the faster 1/E rate holding for arbitrarily adaptive but good demonstrations.
2 SETTING AND PROBLEM DEFINITION
Let X' and y respectively be the sets all possible contexts and actions. A policy (or predictor, the words interchangeably) is specified by mapping æ + x(- x) from contexts to distributions over actions. We overload notation and view a policy as a randomized mapping, so (a) denotes random variable. If π is deterministic, x (<x the action taken for the context x. Recall the definition of the value of a policy for a reward r from Eq. (2), given by V,(TT) = EgwplEyywrt-p [rr,, v)]. The value also depends the unknown population distribution DD context, which is implicit in the notation. We (2; y) D to denote joint distribution where x ~ 'D and y ~ TT.- r). Learning is based on an i.i.d. samples from 'D x म, for some demonstrator policy स. The goal to always compete with the value of the demonstrator V+, (T) w.r.t. the ground truth reward function r.. This is formalized in Appendix A.I.
Crucial our work is the distinction between types of hypothesis classes and corresponding realizability assumptions:
Reward Class Assumption. The unknown reward function r. a known hypothesis class of reward functions, E R. S [0, 1j The demonstrator policy स is arbitrary.
Demonstrator Class Assumption. The unknown demonstrator गे is in a known demonstrator pol- icy class से II (A())) The reward function r,, X' x y [0, is arbitrary.
WHY MLE IS GOOD FOR LOW-CARDINALITY II BUT FAILS FOR R.
Given policy class IT (A())) the (conditional) Maximum Likelihood Estimator (MLE) (or
the log-loss minimizer) is defined as
m
m
(LEn(() arg max Hlll
ग (yu x;)) = arg min #Ell
log x (m xi))
(MLE)
3.1 LEARNING UNDER THE DEMONSTRATOR CLASS ASSUMPTION
To contrast with work, first consider learning under the Demonstrator Class Assumption, where the demonstrator class II has small cardinality. In this setting, MLE is minimax optimal. The argument uses that MLE achieves an O log [ convergence rate in squared Hellinger distance to the demonstrator's distribution (see Proposition 2), where
(H Q
(VV(())
2EZ
for distributions P. Q over discrete domain 2. A simple change-of-meaaure argument then shows that, when the demonstrator is optimal, the value suboptimality is also bounded by (IP, Q). and hence scales Of log [T]] In any case, the value suboptimality can be bounded by total variation
distance, which decays as O( log IiI since (P, Q) S (P, Q). More generally, one can show the following proposition, with rates that interpolate between the two (in similar spirit to Foster al. (2024), and proved for completeness Appendix D.1):
Proposition For any finite II (A()))) where (X,Y) are countable, for any D, any के E II. with probability at least
over S ~ (D x F)m any Till E (LEn(S)) enjoys the following
guarantee:
(1) For any T. X x
[0, 1] for which मे is optimal, we have
6log(2
Vr,
Vr, (mll & (D x Aale. D x 7) <

---

Published as a conference paper at ICLR 2026
(2) For any r. : X x - [0, 11, we have
6 6log(2
Vr, (T) Vr. (Tmle) I DTv(Dx & mleer x R) =
(3) More generally, for any r. X x y + 10. 1], we have
Vr, (7t Vr, mle) 50 (V Ax log([
/og( log m
m
m
where A E (x.a))DDF Suupy r (2 (2, v)] is the value suboptimality of the demonstra-
tor.
The proof follows ideas from Foster al. (2024). We first use standard density-estimation guaran- tees to obtain convergence in squared Hellinger distance (D mlle, D x т) (see Proposition 2). A saange-of-measune argument then shows that, under optimality, the value suboptimality con- trolled by the squared Hellinger distance (and, in any case, by the total variation distance). This immediately implies the first two parts of the proposition. To obtain the tight interpolating rates, we rely on variance-dependee decomposition of the suboptimality in terms of the squared Hellinger distance established by Foster et al. (2024); see Appendix D. for details.
3.2 LEARNING UNDER THE REWARD CLASS ASSUMPTION.
We now ask what can be said about the performance of the Maximum Likelihood Estimator in terms of the log-cardinality of the reward class R. We show that even when R S. {0, 1)xxx)/xxx) is a binary reward class and the demonstrator is always correct, i.c., demonstrator with value 1, the MLE still fails to learn. We will specify a class R. such that the correct answers always exist, i.e., VrER. S.t. r(z, y) 1.
First, since our problem is only specified in terms of the reward class R, we need to specify which policy class II we want maximize the likelihood over. Since we are promised that the demonstrator is optimal, the most natural choice is to consider the class IT from Eq. (14) of all policies supported on the correct answers for the rewards R. Recall that Reward Class Assumption with R. along with optimality of the demonstrator, implies the Demonstrator Class Assumption with lR As discussed in Section 2, even if R is small, the class [l may be infinite (as long as there are even two correct actions for reward). This makes the guarantee in Proposition 1 vacuous. Indeed, we will see that MLE can fail to learn with any sample size, even when R = | =
Let us think of how Aill MLEE1e (S) would behave. For any observed in the training set, any distribution over correct answers is allowed, and only correct answers are observed, so Tm x) will match the empirical distribution s (a) of observed yı for 2; x. However, unseen contexts x, it may choose any distribution over actions E (x) for some r that is consistent with the data, i... in:
CONSR (S) = (r R V(gD)s r(x,8) = =1]
(4)
And the MLE predictor can be written as:
if a is observed; (m any distribution with support C (x(xx, otherwise.
(-
for some r E CONSR(S),
(5) Since the loss on observed contexts is zero, the loss of the MLE predictor Rile is exactly the same as just choosing any E CONS (S) consistent with the data, and predicting actions y E o, on any unseen context Maximizing the likelihood over policy class does not buy us anything for unseen contexts over just ensuring consistency. While relying on just consistency is fine, even optimal, in (realizable, finite cardinality) supervised learning, will see that it fails for our problem, and so does MLE which equivalent to it.
Consider only binary response space y {0, 1} and reward class R (ro, roi J. For every E X. Ta considers only as the correct answer (i.e., V ro(#, y) = 1 - y and thus only Tra (x) = {0)). However, Tui considers both (0, 1} the correct responses (i.e., Ya rzz 0) r(z, 1) 1). If the true hypothesis is r, = To, then all observed labels are from an optimal demonstrator. However, Yoi also remains consistent, and thus MLF (S) may output 1 at test time time-foiling to generalize
5

---

Published as a conference paper at ICLR 2026
and go beyond memorization. Thus, for imput distribution D such that the missing mass (i.e., unobserved contexts) is arbitrarily close to 1, we get the following failure.
Theorem (Failure MLE over I. There exists R S {0, 1}xxy with = =
and a choice of r R and गे such that for every sample size and (0,1), there exists marginal distribution DD such that (+) = but almost surely for ~ (D x T T some mll E
MLEug (S) has the following guarantee:
Vr. (flee Sc.
The failure of MLEng can be attributed the fact that the induced demonstrator class Ig is too large, with infinitely many hypothesis, even though R is small. We might instead consider using Maximum Likelihood Estimation, but over smaller policy class whose size matches that of R. A natural choice for such policy class is:
I = [ R} where (- x)) = (n))(
where for simplicity consider σr (1)), the set of optimal or correct actions, be finite here. The advantage that we now have π [R]. However, this policy class is misspecified in that the actual (and optimal) demonstrator need not be uniform all correct answers and so πT might not contain T. We again cannot rely on Proposition 1.
Remark (Overlap of MLF). Interestingly, despite the misspecification, for any т. E R and with any optimal demonstrator से (even if ग E IIR, i.e., under the Reward Class and Optimality assumptions), we can at least ensure that mil E |MLE (S) achieves a non-trivial overlap with the optimal (i.e., correct) answers o. (x). More specifically, Theorem 6 in Appendix D.3 ensures that for S (D 7 such that के acts optimally w.r.t. T, E R, any Mile MLE (S).
Ps POND 3gEor Amte
2 1 1--8) 2-88 for any
m 2 _ (log Rlg
(6)
where x = sup,, or (x)) are the maximum support size.
The guarantee in (6) says that the MLE generates a good response under r. with probability at least 1/x = 1/ supr.. (x). However, such an overlap is not our goal and it is not very helpful in general: The support sizes & can be huge (e.g., for essay answers with even minimal language variability, it is exponential in the length of the answer), and the error of Amle can be as high as 1/K. One can show this is indeed the case- witt large support sizes, the value can be arbitrarily close to zero, formalized below and proved in Appendix D.2:
Theorem 2 (Failure of MLE over 1IR). For any E (0, 1), there exists R S {0, 11 with R = 2, X 1,1 2[//e,, such that for some choice of 'D and T. E R such that V, ( 1, for every sample size m, and almost surely for S (D x 7)m, there is unique
Tile E MLET, (S) and it has the value:
Vr, (Tmlle Se.
Theorems I and 2 establish that natural choices of MLE fail learn simple binary reward classes R (cf. Definition 1) even from an always correct demonstrator, i.e., with value 1.
LEARNING ALGORITHMS
In this section, we present our main result: learning rule that can learn from demonstrations (as in Definition 1) with tight optimal sample complexities with optimistic rates. To achieve this, it is useful to focus on an adversarial online version learning from demonstrations. For conceptual simplicity, we first discuss the case of binary rewards with always correct demonstrator. We present a logarithmic Olog RR)) mistake online learner in Section 4.2. In Section 4.3 we general- ize this to an online regret bound for real-valued rewards with arbitrary (not necessarily optimal) demonstrations. The desired statistical learner is then obtained using an online-to allineto--btth conversion in Section 4.4.

---

Published as a conference paper at ICLR 2026
Online (Mistake-nnwwrry Learning Process for Contextual Bandits: At each round ..
1. The learner receives an instance Tt.
The learner outputs a response Jt (based on past instances and demonstrations)
In the case of 0-1 rewards with always correct demonstrator: the learner is said to make mistake iff r. (x, Vo) of 1. More generally, the learner is said to have received a reward r. (2r, Mr.. Importantly, the learner does not know whether a mistake occurred, generally, does not observe the obtained reward.
4. The learner receives demonstration Vr- In the case of 0-1 rewards with always correct demonstrator, we are promised that r. (Zt, 311 = 1. Generally, Y/ can arbitrary.
The online learning process studied. In contrast to the standard online contextual bandit setup, in our setup the learner does not know whether it made mistake (unobserved rewards), but does always receive some other (unrelated) demonstration. The demonstration is always promised be correct the special case considered in Sections 4.1 and 4.2. the most general case studied in Section 4.3, the demonstrations can be arbitrary.
Figure 1:
4.1 ONLINE SETTING
To make progress, we turn our attention to the even more challenging online version of the problem. In the online process, summarized in Figure I, the instances F2 X are received sequentially, and on each instance the learner first responds with Yi and then receives an arbitrary demonstration yt- We will return to the general reward case later in Section 4.3, but for now, let us focus on the case of binary rewards, with always correct demonstrations y/ s.t. (xe,, yo for some unknown reward function r, E R.
The goal of the learner is to output Te such that r )e) then the learner made a mistake at round L. Importantly, the learner does not know has made mistake- the learner receives no direct feedback on Vt. (unlike in traditional contextual bandit setup. e.g., Langford & Zhang, 2007; Dudík et al., 2011; Agarwal et al., 2014). only some (possibly unrelated) correct answer y/- Nevertheless, we want ensure that for any r. E R. and any sequence of instances I, and correct demonstrations with r, (Z Vt)) 1, the learner will only make few mistakes with respect to r+.
and say that if (2,, Uc) 1,
4.2 LOGARITHMIC MISTAKE ONLINE LEARNER WITH ALWAYS CORRECT DEMONSTRATIONS
We now present Algorithm 1, which we will prove makes at most log [R] mistakes. For now let us focus on binary rewards where the demonstrator is always correct, and setting the hyperparameter y 1 (see Section 4.3 for the general version). The algorithm maintains weights w(') (r), which are initialized to uniform over all r E R. Predictions are based on a weighted average of the reward. For binary rewards, this corresponds
outputting the response Tt which maximizes the sum of the weights of all rewards r under which it correct response to d (i.e., s.t. r r(r,, yt) = 1). This is similar to Maj. As with 7Maj. we zero out the weights of that are consistent with the demonstration yt, and thus completely remove them from consideration. The difference is that in addition, we also increase the weights of under which the predicted action Ot is not correct (i.e., r(ze, Fi) * 1, and so r's vote did not help elect Te). We do so despite the fact that we do not know whether De correct or not. turns out that this ue-weighting sufficient to exponentially reduce the number of mistakes the algorithm makes:
Theorem 3 (Online Guarantee). On any sequence (e),y that has only correct demonstra- tions with respect to some E (i.e., s.s. Ve r(z,, 31) 1), Algorithum with 7 makes at most log [R] mistakes with respect to r.
See Section E. for the proof. We can use an online-to-batch conversion, as in Algo- rithm 2, to obtain guarantee in the statistical setting of Definition with sample complexity
7

---

Published as a conference paper at ICLR 2026
Algorithm Online MUSTTAKF-UNNAAAEEEEEEII-----
Rule
Input: A finite reward R - [0, [) and a hyperparameter E [0, 1].
Initialize (r) for all E R.
In every round.
1. Receive Tt.
2. Output De == arg maxyey ErER (r) r(x..)).
So using the policy (x) == arg max ErER w!l (r) r(z, v), for the input aq.
3. Receive yr-
4. In the always correct demonstrator case (when = 1), update for each r E R:
if r y1) +
(r(+
(r)
if both r(ze, H1) r r 9c) = 1; if r 96) but r( De # 1.
(r)
More generally, update: w (r) + wh (r) (1 7)) - (1 ()))))))))-r( ((xe:))))
(8)
where r(x, = sup, r(,,
O(s-'(log [ + log )-)) However, rather than carrying this out formally at this point, we first generalize the online analysis beyond binary rewards with correct demonstrations, and then present a unified anline-to-bated conversion and statistical guarantee.
4.3 REGRET GUARANTEE FOR BOUNDED REWARDS AND ARBITRARY DEMONSTRATIONS
Still sticking to the online model, we now discuss the most general setting and provide a guarantee for Algorithm 1. This generalizes the earlier analysis in two ways: (i) it handles real-valued bounded rewards, and (ii) it accounts for suboptimal demonstrators. We refer to Figure for the general online learning process from demonstrations: after outputting for context Tz. the learner is said to have received reward r(te, Me), which not observed. The demonstrations y/ may arbitrary and are not guaranteed to be optimal. Accordingly, instead of counting "mistakes," reward. For a fixed reward r E R, consider the total reward collected by the algorithm and by the demonstrator, well as the optimal total reward achievable on the given sequence of contexts with respect to r, i.e., by playing an optimal action in each round:
track total accumulated
T (T( = Σ (.) J,(T)== = Σ ()), JT(T = Σ sup ()).))
T
(9)
Since the demonstrator may longer be optimal, cannot completely exclude from consideration reward models for which a given demonstration is not optimal. Instead of the hard update rule in (7), we use the softer variant in (8): we again decrease the weight of reward functions r for which the demonstrator is suboptimal, and increase the weight of reward functions for which our current prediction is suboptimal. However, the amount of increase or decrease now related to the degree of suboptimality (rather than being binary as before) and is controlled by the step-size hyperparameter y (instead of doubling or completely zeroing the weight).
With the update in Eq. (8), the total weight in the system, W,, remains monotone non-increasing. Moreover, by keeping track of the weights (r)).)). we can conclude the following (see Appendix E.2 for derivation details and the complete proof):
Theorem (Online Regret), For any finite R and any sequence ((xe, We))E,, after any number of
rounds T of Algorithm I with any y 1, we have:
- -J,)) log R
for all E (T - J. < (1 (1,T))
Theorem 3 can be seen edge case of Theorem 4. When the demonstrator is optimal, in Theorem we have (T) (T) 0 and we take y be arbitrary constant (e.g..
8

---

Published as a conference paper at ICLR 2026
Algorithm 2 Statistical Learner via Online-to-Batee Conversion
Input: Samples S {( y4) : [ml]], and an online learning algorithm Aonline-
Run once over an online learning algorithm Aonline- Record the policies { ke(m used at different rounds.
For every test context æ E X, pick a stopping time I Unif((1.... m)), and output
according to fy(- z). In other words, output a randomized mixed policy:
1 m
m
Tozo( x))
(e x)
(10)
= in the specialized analysis of Theorem 3). When the demonstrator is suboptimal, we need choose y to balance the two terms on the right-hand side, and doing so yields regret bound of of (V
(V::(T)-J.CT) -..()) log [2]
O( T R|). Since our final goal is statistical guarantee as in Definition 1, it will be easier to set after an online-to-batch conversion, which we do in the next subsection.
4.4 ANALYZING STATISTICAL ONLINE-TO-BATER ESTIMATOR
We now return to the statistical setting of Definition and consider learning (randomized) policy given training set
(D 7)m We do so using standard online-to-batch conversion of Algorithm 1, as specified in Algorithm 2: we run Algorithm 1 on the training set S (we do not need to randomize over the ordering), keeping track of the m policies मe used at each step t = 1,..., m. The output is then stochastic policy To2b. where for any context æ, R.2b | x)) is uniform over the m answers { (2) I(=1....m that the algorithm would have given for context x at different steps. Equivalently, given context x, Rozb picks random stopping time & fl..... .m}, runs Algorithm on the first examples, and then uses the policy म. reached after steps by outputting (x()..
To obtain "optimistic" rate that interpolates between 1/s sample complexity for optimal demon- strators and 1 1/2² in the worst case, we also rely on a bound A on the suboptimality of the demonstra- allow any demonstrator (since the worst demonstra-
tor, which we can always take A tor is I-suboptimal). An online-to-bated analysis (see details and proof in Section E) then ensures:
Theorem 5 (Statistical Guarantee via Online-to-Batee Conversion). Consider any finite reward class R containing bounded reward functions.
(i) For any sub-optimality bound E [o, 1]. consider rumning Algorithu with
log R + 6og (28
min
1 2:
6m
For any E (0, 1), sample size m, and demonstrator म. with probability at least 8 over (D 7m for all r E R s.t. tt has suboptimality V (+) S A, we have:
6 6 log R
+ 6 log ()))
log [ + 12log(
(+ - V2(Ro2b) 2
(11)
(ii) Optimal Demonstrationss Consider ranning Algorithm 2 with 7 1/2, For any 6 E (0. 1) and sample size m, with probability at least over F1,... Fm Midd D, and any y1 Your for any E R for which Vie(m) Ys E or (r(x), we have:
16 16 R + 12 log
V. - VV S
(12)
See Appendix E.3 for the proof. Theorem ensures that any finite reward class is learnable from demonstrations (Definition 1) with sample complexity O(g-² [R] + log 8-1)), by always choos- ing A as an upper bound on the demonstrator's suboptimality (in which case the guarantee in (11) holds for all R). Moreover, reward classes are learnable from optimal demonstrations with

---

Published as a conference paper at ICLR 2026
sample complexity O(E-' (log [R] + 1))) by choosing A 0 and thus 7 1/2. In fact, the theorem delivers more than what is required under Definition 1, since the optimal demonstrations y: may be chosen arbitrary optimal actions given the contexts 24, even adaptively, without following a fixed optimal policy. The guarantee holds as long as the contexts are drawn independently from D.
Remark We emphasize that Algorithm guarantees low value suboptimality with respect to by establishing guarantees uniformly over the entire reward class R (see Theorems to 5). The learner only has access to demonstrations. It achieves a small regret/value gap by ensuring its performance is good simultaneously for all reward functions in R, and in particular for whatever may be the true r, E R.
Remark 3. In this paper we focus contextual bandits. But our results apply more generally to MDPs with deterministic (possibly unknown) transitions. Consider an MDP model class of finite cardinality, where each model specifies deterministic transition function and a (possibly stochastic) reward. We can reduce this to a contextual bandit setting, where the entire MDP action sequence is treated single contextual bandit action. For auto-regressive language generation, viewing each token as an action and the entire context as the state, the transition dynamics are fixed and known and only the reward function is unknown; the entire generated sequence of tokens can be treated as a single action y,
5 SUMMARY AND OPEN ISSUES
We study the problem of learning to answer from correct demonstrations, i.e., imitation learning or apprenticeship learning in contextual bandits. We compare and contrast learning based on the low- cardinality assumptions on the Reward Class and the Demonstrator Class. We argue that the former is strictly weaker when the demonstrations are optimal. We show that natural approaches based on likelihood maximization fail under this assumption, necessitating going beyond them. We design a learner that achieves optimal sample complexity logarithmic in the reward class size, with tight rates. Our study motivates looking beyond likelihood maximization and distribution matching, and adopt new type of methods when the harder problem of distribution matching may not be achievable.
We point to Remark for some additional technical theoretical questions. We also lay out two broad directions for future research below.
Continuous Classes. Our analysis is for arbitrary abstract, but finite, reward classes R, and is stated in terms of the log-cardinality [R]. Many natural reward classes are better thought of as infinite classes with continuous parametrization, R = fr.. (X' x }) + R w E Rd. Relevant choices for R include generalized linear (or kernel) classes with rw (z.g) = g((w, Φ(z, y))) (for some specified feature map Φ(x, v) and link function g IR + RR, low-rank bilinear scores, or transformer models where Tur (2. y) E IR is the sealar output of transformer with weights w taking the sequences and y as imput. Can characterize learnability from demonstrations, and the sample complexity required, for such continuous classes as well?
Tractable Implementation and Supervised Fine Tuning of LLMs. Another limitation of our work is that, while the sample complexity is logarithmic in the reward class size, a naive implemen- tation of our methods requires runtime and memory linear in the cardinality, and thus exponential in the number of parameters for the continuous classes discussed above. We envision heuristic simplifi- cations of the algorithm that may be practically implementable. However, several research questions still remain open when it comes to the practical success of such an approach in the SFT of L.MSs.
In particular, what reward class R {ram> should be used in practice? After pretraining, we typ- ically only have an initialized policy лo- How important is this pretraining knowledge for down- stream success? How can also be used to construct different parameterized reward classes? To what extent, and in which situations, are methods based on reward classes more successful than policy-class- based likelihood maximization? How do they affect downstream performance in other phases of posttraining? Ultimately, the key question is whether this class of methods. which per- forms reward hedging via iterative planning and discrimination, can serve as a genuine alternative to behavior cloning, and if so, in which regimes and for what reasons.
10

---

Published as a conference paper at ICLR 2026
Acknowledgement. We grateful to Yishay Mansour for bringing the work of Syed & Schapire (2007) to our attention, and to Moulin et al. for pointing us to their recent related work.
REFERENCES
Picter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, pp. 1, 2004.
Alckh Agarwal, Daniel Hsu, Satyen Kale, John Langford, Lihong Li, and Robert Schapire. Taming the monster: A fast and simple algorithm for contextual bandits. In International conference on machine learning, pp. 1638-1646. PMLR, 2014.
Anthropic. Claude model family. https: www. www.aaathropcc.com/eeeeeoooe 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Pondé de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Sandhini Puri, Gretchen Krueger, Mikhail Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Michael Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Phil Tillet, Felipe Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William H. Guss, Alex Nichol, Igor Paino, Nikolas Tezak, Jie Tang. Igor Babuschkin, Suchir Balaji, Szymon Jain, William Saunders, Jesse Hsu, Ryan Crow- der, Anjali Srinivasan, Andrew Ho, Tom Gray, Nicholas Ryder, Dario Amodei, llya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374, 2021.
Uri Dalal, Meirav Segal, Zvika Ben-Haim, Dan Lahav, and Omer Nevo. Leveraging Ilm inconsis- tency to boost pass@k performance. arXiv preprint vrXiv:2505.12938, 2025.
Google DeepMind. https:/ //blog. Introducing gemini: Our largest and most capable ai model. gemini ai/, 2023.
DeepSeekAL. Deepseek models overview.
www.dddwpeeeeeeee com, 2025.
Miroslav Dudík, John Langford, and Lihong Li. Doubly robust policy evaluation and learning. arXiv
preprint arXiv:/ 102.0000344400 2011.
Dylan J. Foster and Alexander Rakhlin. Foundations of reinforcement learning and inter-
active decision making. Technical report, arXiv preprint arXiv:2312.16730, 2023. URL
https://arxit httppp///aarxu.rrrraa/////2222.2
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.1344 2021.
Dylan Foster, Adam Block, and Dipendra Misra. behavior cloning all you need? understanding horizon imitation learning. Advances in Neural Information Processing Systems, 37:120602- 120666, 2024.
Audrey Huang. Adam Block, Qinghua Liu, Nan Jiang, Akshay Krishnamurthy, and Dylan Foster. Is best-of-n the best of them? coverage, scaling, and optimality in inference-timg alignment. arXiv preprint arXiv:2503.211 218..1188, 2025.
Ziwei Ji, Nayeon Lee, Rita Frieske, Tiezheng Yu, Dan Su, Yan Xu, Eric Ishii, Yejin Bang, An- drea Madotto, and Pascale Fung. Survey hallucination in natural language generation. ACM Computing Surveys, 55(12):1-38, 2023.
Adam Tauman Kalai and Santosh S Vempala. Calibrated language models must hallucinate. In Proceedings of the 56th Annual ACM Symposium Theory of Computing, pp. 160-171, 2024.
Adam Tauman Kalai, Ofir Nachum, Santosh S. Vempala, and Edwin Zhang. Why language models hallucinate. Technical report, OpenAI, September 2025.

---

Published as a conference paper at ICLR 2026
Michael Kirchhof, Luca Füger,, Adam Goliñski, Eeshan Gunesh Dhekane, Amo Blaas, and Sinead Williamson, Self-reflective uncertainties: Do llms know their internal answer distribution? arXiv preprint (rXXh:22505.00995 2025.
John Langford and Tong Zhang. The epoch-greedy algorithm for multi-armed bandits with side information. Advances in neural information processing systems, 20, 2007.
Antoine Moulin, Gergely Neu, and Luca Viano. Inverse q-learning done right: Offline imitation learning in rearliiiall mdps. arXiv preprint arXiv:2505. 19946, 2025.
Andrew Y Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Ieml, vol- ume 1, 2000.
OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Gabriel Orlanski, Seonhye Yang, and Michael Healy. Evaluating how fine-tuning on bimodal data affects code generation. arXiv preprint arXiv:2211.08822, 2022.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in neural information processing systems, 36:53728-53744 2023.
Nived Rajaraman, Lin Yang. Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. Advances Neural Information Processing Systems, 33:2914 2924, 2020.
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems, 14211702-11716 2021.
Gokul Swamy, Sanjiban Choudhury, Andrew Bagnell, and Steven Wu. Of moments and match- ing: A game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning. pp. 10022-10032. PMLR. 2021.
Umar Syed and Robert E Schapire. A theorete approach to apprenticeship learning. Advances in neural information processing systems, 20, 2007.
Hugo Touvron et al. Llama: Open and efficient foundation language models. arXiv preprint 2321::3302.1397,, 2023.
Fang Wu, Weihao Xuan, Ximing Lu, Zaid Harchaoui, and Yejin Choi. The invisible leash: Why rivr may not escape its origin. arXiv preprint arXiv:2507.14843 2025.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
12

---

Published as a conference paper at ICLR 2026
APPENDIX TABLE OF CONTENTS
A Miscellaneous Deferred Discussions
13
A.I Learning Goals Under Two Assumptions and Comparison
13
A.2 Comparison with Syed & Schapire (2007)
15
A.3 Reward Maximization Distribution Matching
16
B Technical Preliminary Lemmas
18
Proofs of Simple Lemmas
19
C.I Proof of Lemma and Corollary
19
C.2 Impossibility of Cloning: Formalizing Observation
20
Proofs for Section
21
D.1 Proof of Proposition 1: MLE is Good for Low- Low-Cardinality 11
21
D.2 Simple MLE Failures for Now-Cardinality R
22
D.3 Overlap of MLE
24
Proofs from Section 4
24
E.1 Proof of Theorem 3
25
E.2 Proof of Theorem
25
E.3 Proof of Theorem 5: Online-to-Baie Analysis via Freedman's 's Inequality
26
pass@k Error Minimization
28
F.1
Online Learner for passukk Metric
29
F.2 Statistical Upper Bound
F3 Lower Bounds for Online and Statistical Settings for pass@k- Metric
32
A MISCELLANEOUS DEFERRED DISCUSSIONS
A.
LEARNING GOALS UNDER TWO ASSUMPTIONS AND A COMPARISON
In this, we formalize learning under each of the two assumptions on the Demonstrator Class and Reward Class, and argue that the latter is strictly weaker assumption. We start by discussing some terminology.
Demonstrator Optimality: Although our general results do not require the demonstrator to be opti- mal, it is instructive to pay particular attention the case where it optimal. As we shall see, can get faster convergence rates and lower sample complexity under demonstrator optimality, and this will also be important in relating between the two assumptions above in Appendix A.I. We say that the demonstrator स acts optimally with respect to iff
V, 2 V," sup Vr, (T).
(13)
This optimality equivalent to requiring that for almost surely over I ~ D, the policy 7(-2)) is supported the set o., (æ) of optimal actions, in Eq. (1) (i.e., "correct answers". in the special case of binary rewards where there is always correct answer).
Our focus is the Reward Class Assumption and learning in the following sense:
13

---

Published as a conference paper at ICLR 2026
Definition 1 (Learning from Demonstrations under Reward Class Assumption). We say that reward class R is learnalle from demonstrations by a learning rule A (X' x)" ) + (A()))* with sample complexity AR,A(e,, 6), if for any e. 8 E (0, 1), and any sample size m 2 A...AE, 6), for any D, स and RR we have that with probability at least 1 over the training set S (D x xm)
(A(S)) V. (+)
Furthermore, R is learnable from optimal demonstrations where the same definition applies but the learner is further promised that the demonstrator गे acts optimally w.r.l. T*-
We contrast this with learnability based Demonstrator Class Assumption:
Definition 2 (Learning from Demonstrations under Demonstrator Class Assumption). We say that a demonstrator policy class I1 is learnable from demonstrations by learning rule A (X x ) (A())) with sample complexity m 6), if for any e, 8 E (0, 1), and any sample size m 2 A 8), for any D, any E II, and any unknown reward Yw,, we have that with probability at least - å over the training set S ~ (D k)
(A(S)) 2 (T)
Furthermore, TI is learnable from optimal demonstrations where the same definition applies but the learner is further promised that the reward is such that II is optimal w.r.f r..
Let us compare the relationship between the Reward Class Assumption and the Demonstrator Class Assumption that we introduced. We will do so under demonstrator optimality, as this optimality assumption provides link between the demonstrator and the reward function. We will argue that under demonstrator optimality, the low-cardinality Reward Class Assumption is strictly weaker than the low-cardinality Demonstrator Class Assumption.
The Reward Class Assumption, combined with the optimality of म, implies that for some r E R, the demonstrator से is supported the set of optimal actions o, X + 2n given by:
(1) arg max (f,, yEY
(14)
Thus
T IR
U II, where II, (N X A()) V supp T(. (x)) PE
(15)
So the Reward Class Assumption with R implies the Demonstrator Class Assumption with T However, once [o,( 1, the class 1l is much larger than R. Even if R has finite cardinality, the class 1R is continuous and infinite. Even we discretize this probability space, 1l can still be introducing (x)](x)]- parameters per instance choose the distribution over correct answers, so overall a number of parameters that scales with the domain size and the number of correct answers. For example, even just restricting to policies that are uniform subset of o, (z)), there are (2),, - 1) optimal policies for each E R.
much larger, since
Under the Demonstrator Class Assumption and the optimality the demonstrator, the class of all reward functions under which the policies in Il are optimal can be huge as well. Nevertheless, it sufficient to consider an artificial reward class, where rewards essentially enforce matching the support of the policies in the demonstrator class II:
R (1-01 X (0,1)
MT) where
(2,9) =000)=1(0 1fy E supp T(: [=)]-
(16)
Although the actual reward r. may not be in R0-1 if some से E IT is optimal with respect r., then being (sub)optimal w.r.t. the corresponding _011 ensures also (sub)optimality w.r.. r.:
$a. introduced earlier in Eq. (I) is the same as σr. (i.c., evaluated at r r.). It is important to note that o-(z) may he empty: since y is infinite, the supremum supyey r(x, y) need not be attained. Thus, for an optimal policy to exist with respect to a given r. it must be the case that 0,(xx) is non-empty almost surely ~ D. Our general result does not rely on the existence of an optimal policy. One may instead consider (arbitrarily large hut finite) ) assume that the range of r(x, y) is finite (c.g., binary, taking only 100 possible values), either of which sufficient to guarantce the existence of an optimal policy.
14

---

Published as a conference paper at ICLR 2026
Lemma 1. For the demonstrator class assumption, and assuming optimality of the demonstrator,
Le,, т E II is optimal w.r.t some unknown r., we have: (1) π is optimal w.r.t r E R with
value V = 1; and for any policy æ. if is e-suboptimal to i w.r.t 7011 then a is e-suboptimal
w.r.l. also Tx- In other words, V,g1 (T) 2 implies V,, (T) 2 V
See Appendix C.I for a proof. This immediately leads to the following corollary. Corollary 1. For any demonstrator class if R is learnable from optimal demonstrations by an algorithm A with sample complexity m(c, (Definition 1), then II is learnable from optimal demonstrations by the same with the same sample complexity m(,,5).
Since 'Ri'l I [T], cardinality Reward Class Assumption, imply the same guarantees on learning from optimal demon- strations under the low-cardinality Demonstrator Class Assumption. A direct corollary of Theorem and Corollary is that our learner continues to succeed for cow-carinallity demonstrator classes, when applied to suitable reward class.
that guarantees on learning from optimal demonstrations under low-
Corollary (For low-cardinality Demonstrator Classes), Any finite demonstrator class TI is learn- able from optimal demonstrations with sample complexity Olog [I| + log 5-1)), by applying the method from Theorem 5 to the reward class Rf.1 from (16).
From theoretical learnability perspective, and at least under demonstrator optimality, the finite- ness of reward class assumption thus weaker, and so learning results under the Reward Class Assumptions are stronger and more general than under the Demonstrator Class Assumption.
A.2 COMPARISON WITH SYED & SCHAPIRE (2007)
We now discuss Syed & Schapire (2007) and instantiate their results for contextual bandits. Syed & Schapire study learning general stochastic Markov decision processes with linear reward R" (s) (w, Φss)), for -norm constrained reward classes. Their setup can be mapped to contex- tual bandits by considering the state space S = X' U (X' x i.e., the states consist all contexts as well as tuples of (context, action) pairs. While their paper requires the state and action spaces to be finite, we believe this restriction is not essential for their guarantees, since the sample complexity has no dependence on the sizes of these spaces.
For a finite class R, we can define feature map Φ S + R2 with d = | and express each reward as linear predictor with unit & norm. Consider |R] dimenniionll feature map whose coordinates are indexed by r E R as follows:
((r,))[r]= rrx,,) and ((x)(r]=
Then each r R can be realized by linear classifier w, that is the indicator vector in the th coordinate. Taking the discount factor y = 1, their expressions for the value and feature expectation (denoted by μu) (Syed& & Schapire, 2007, Section 2) under this embedding become:
V.(m) = E(ew))DDX
Σ (ur, ((n)) = E( (w., ((x v)) (x
= =B(ew)
p 5) = [-))(r]
(17)
Under this correspondence, we believe their Theorem 2 can be adapted to show that any finite class R is learnable with sample complexity O (e-2 [R] + log 6-1)) Their Algorithm I is similar to our strategy in that it maintains weights over the features, which correspond to weights over
"We note that their Theorem 2 is stated for a multi-step MDP with discounted rewards. One needs to adapt their proof to obtain the single-step guarantee. Namely, it introduces еæ the error term incurred by trajectory of finite length. But in contextual bandits, eн 0 even for length-one trajectory, and thus the condition required on H is degenerate. Also, dependence on 1/(1 - 7) in the sample complexity due to the scale of the cumulative reward, but even when taking 7 = in contextual bandits, the total reward is bounded in the single-step case. Thus, we can still obtain good estimate of the value of a policy with small er using ( samples. Finally, since can exactly embed R linear threshold, we have en 0, and assuming we can solve arg max (as do), even ep
15

---

Published as a conference paper at ICLR 2026
Our Learner
Syed & Schapire (2007)
Moulin et al. (2025)
A [R]([ RI/02
1og(1R1/8) m
log [R] log []]
+
m
172
Table 2: Comparison between rates of our learner and relevant prior work. For value suboptimality based on m samples, our learner has a decay rate that interpolates between 1/m (optimal demon- strator where A = 0) and 1/ //m (always), based on the suboptimality gap A.
rewards in our setting. The algorithm then (i) produces an optimal policy according to the reward given by the weighted combination of features (as in Line 2 of our Algorithm 1), followed by (ii) a multiplicative weight update over features (corresponding to the reward updates in Line 4 of our algorithm).
The main differences are as follows. To perform the mulippicative-weight update, their algorithm estimates the value of the current policy with respect all rewards in the class accuracy eF, and then updates based on the resulting value differences (equivalently, feature- eeature-eppcttttion differences via (17)). This leads pass online approach: we update the weights after processing each example, without revisiting the data.
multi-pass batch algorithm. Our algorithm, in contrast, simpler one-
We obtain faster optimistic rates, interpolating between ((1/E) when the demonstrator is optimal and ( in general: see Table 2. Our learner also enjoys regret minimization (Theorem 4) or a mistake bound (Theorem 3) for adversarial sequences of demonstrations in the online setup (Figure 1). Even in the statistical setting, our learner continues to enjoy these guarantees when the demonstrations are arbitrarily adaptive, long they good (Theorem 5).
Dual Updates. Another relevant work is Moulin et al. (2025), which can also be instantiated to admit guarantee for contextual bandits. They again study general MDPs under a certain type of realizability assumption on the Q-function. In the contextual bandit case, the reward class R and their Q-function class Q are equivalent, and their assumption simplifies assuming only a reward class, as in our setting. In contrast to our approach (and Syed & Schapire (2007)), which makes multiplicative weight updates over rewards, their algorithm makes iterative multiplicative updates over the space of policies, which can be viewed as dual updates. By adapting (Moulin et al., 2025, Theorem Algorithm 2), we believe it possible to show that their learner achieves e-value- suboptimality with O (E-4 log | log([R]/s)) samples. It is unclear whether the dependence on log | can be avoided for the dual update strategy. Note that log | may scale with the length of the responses the autoregressive generation setting, and thus, at least from purely statistical point of view, may not be desirable.
A.3 REWARD MAXIMIZATION VS DISTRIBUTION MATCHING
In this paper, we view reward maximization or utility as our goal when learning from demonstrations. That is, we measuring success not by how well we match the actions of the demonstrator policy, but only by how well we match its reward or value. This is in line with other work on imitation learning (e.g., Pomerleau, 1988; Ng et al., 2000; Syed & Schapire, 2007: Rajaraman et al., 2020; Rashidinejad et al., 2021; Foster et al., 2024, etc.), which also view the demonstrator as a resource, but the ultimate goal is to obtain good reward.
Cloning. One approach to obtaining good reward based on demonstrations is to match the ac- tion distribution, or clone, the demonstrator policy, as we have seen in Proposition 2. Indeed, the approach of Foster et al. ensures high reward by virtue of cloning the demonstrator, relying low-cardinality Demonstrator Class Assumption.
An important message we emphasize is that matching the distribution of the demonstrator is not necessary for matching the demonstrator's 's reward. In fact, matching the distribution (i.e. "cloning") might not even be possible, even if it is possible to match the reward.
Specifically, we show that, under a Reward Class Assumption, it is possible to match the reward of the demonstrator, but not through cloning maximum likelihood approach. To emphasize that
16

---

Published as a conference paper at ICLR 2026
cloning is not possible in this setting, it is sufficient to consider trivial reward classes where reward maximization is straightforward, but for which cloning the demonstrator is impossible.
Observation 1 (Informal, formalized in Proposition 3). Consider simple reward class R. = (r... which considers all answers to all questions as correct, with multiple answers [ and infinitely many questions X. Then reward maximization is possible without any samples, but matching the distribution of specific optimal demonstrator т (under TV, Hellinger, KI,, or reverse KI.) is impossible using any number of samples.
containing a single reward function r. (= y)
As such, any method (such ours) that is guaranteed to succeed under the Reward Class Assump- tion must do so not by ensuring cloning. Furthermore, departing from the Demonstrator Class Assumption and cloning also entails looking beyond maximum likelihood estimation (i.e., log-loss minimization). The likelihood maximization is inherently associated with, and typically analyzed through, cloning. Indeed, in Section 3 showed that, in our setting, maximum likelihood estima- tion fails at reward maximization- and therefore also at
new approaches.
Inverse Reinforcement Learning and Planning. Another approach to imitation learning in- verse reinforcement learning (IRL) (Ng et al., 2000; Ziebart et al., 2008), which aims to recover an underlying reward function that rationalizes the expert's behavior." ? If we have such a reward F. then one can plan a good policy for apprenticeship learning. Again this goal is impossible to achieve just based on demonstrations, and is a more difficult task than just learning a policy (i.e., apprenticeship learning).
Reward Hedging via Iterative Discriminating and Planning. Importantly, learning a policy from demonstrations still remains possible under the weaker Reward Class Assumption. To do so, all methods-including ours, Syed & Schapire (2007), and Moulin et al. (2025)reel on iterative reward hedging and planning, interleaving between: (i) finding reward that best discriminates the demonstrator from the current policy, and (ii) arriving at a new policy based on this reward. Cru- cially, we emphasize that this reward is not the true reward, nor even reward that rationalizes the demonstrator. We also note that this iterative hedging approach is also seen and referred to as ap- prenticeship learning via IRL in the seminal work of Abbeel & Ng (2004). We explicitly distinguish this from IRL, since the functional role of these rewards is to discriminate a currently maintained policy from the demonstrator policy, rather than to rationalize the demonstrations the original goal of IRL. (Ng et al., 2000).
Approaches to Apprenticeship Learning Cloning IRL
Possibility (?)
Some Relevant
for small Reward Classes
References
X
Foster et al. (2024)
X
Ng et al. (2000); Ziebart et al. (2008)
Reward Hedging via Iterative Discriminating & Planning
Abbeel & Ng (2004); Syed Schapire (2007);
Swamy et al. (2021): )); Moulin et al. (2025); Ours
Question-Answering with Language Models. For question-answering or prompt-completion with LLMs, the emphasis on reward maximization translates to a focus on producing good an- swers (i.e., correct or high-utility ones). We do not view matching the distribution of answers goal. This aligns with how LLM-based Al agents are used and evaluated on benchmarks (OpenAI, 2023; DeepSeekAI, 2025; Anthropic, 2024; DeepMind, 2023). Reward maximization is also the view taken in post-training of L.LMs. Even in limited cases where the answer distribution might matter-such as when there is no clear consensus-thh the model should arguably produce a more comprehensive response that reflects its uncertainty (c.g., Kirchhof et al., 2025), rather than attempt to match the distribution of some specific training "population".
Taking reward maximization as the goal for LLMs, the next important question is: what methods should we adopt when learning from demonstrations available during SFT? The predominant ap-
We use "rationalize" informally. While defining IRL IRL requires care and can be ill- ill-posed (e.g., see discussion in Ziebart et al., 2008), we believe one can adopt a notion under which the discussion remains meaningful. For example, one can aim to recover reward function that correctly classifies the demonstrator's 's action on future examples drawn from a distribution that, with probability 1/2 each, either shows the demonstrator's 's actual action (guaranteed to be good) or an unrelated action (not gond under the ground-truth reward).
17

---

Published as a conference paper at ICLR 2026
proach is log-loss minimization (i.e., Maximum Likelihood), which discussed earlier, succeeds through distribution matching under the small Demonstrator Class Assumption. Since supervised fine-tuned models not always satisfactory (e.g., Ji et al., 2023; OpenAI, 2023), it may be wise to question this assumption and approach. Perhaps in these SFT problems, distribution matching is impossible (with the policy class used model the demonstrator and the amount of data).
Likelihood maximization/log-bnn more difficult task. Kalai & Vempala (2024) and Kalai al. (2025) go even further and argue that distribution matching and reward maximization (i.e., accuracy) might even be in tension with each other, and that insisting on distribution matching could hurt the reward. As discussed, even in situations where matching the distribution could be important, the model should indicate its uncer- tainty. Thus the focus should perhaps be more directly and even exclusively on uillity or reward, rather than attempting to solve the more difficutt task of matching the distribution- sentiment well-captured in a quote by Vladimir Vapnik:
minimization implicitly encourages distribution matching which
"When solving problem of interest, do not solve a more general problem as an intermediate step."
might also be impossible to recover good reward model only from the demonstrations (like in the IRL goal), without other types of feedback. Despite this, learning good policy by adopting iterative reward hedging and planning type methods might be successful.
B TECHNICAL PRELIMINARY LEMMAS
We start by a technical lemma about the one-sided change of measure bound on an expectation of bounded function in terms of the Hellinger distance (e.g.. Lemma A.11 from Foster et al., 2021). We will use the exact variant from (Foster al., 2024, Lemma 3.11).
Lemma 2 (Change-of-measure bound via Hellinger distance, Foster al. (2024)). Let (2,F) be measurable space and let P., Q be probability measures on it. For every measurable function
h: 2- IR.
Ep 12
+ Eelh2
Er [h]
Rg[h]| $
(H(P,Q).
(18)
2
In particular for h: Z [0, R).
[E[M]] S 2 + (P.(Q.Q).
(19)
We now specify Freedman's inequality that allows for tight non-asymptotic control on the mar- tingale difference sequence. We will use Freedman's inequality for the analysis, the variant (Foster & Rakhlin, 2023, Lemma 35) in particular.
Lemma 3 (Freedman's 's inequality; Bernstein for martingales). Let (Xt)ssT be a real-valued mar- tingale difference sequence adapted to a filtration (Fi)IsT. [Xll S R almost surely, then for any E (0, 10.1/R) R), with probability at least - for all T T
log(5-1)
Σ
Ex.
n
E
[X3]
Maximum Likelihood Estimation for Distribution Learning. We now state guarantee for the maximum likelihood estimator (MLF) for density estimation, exactly similar to (Foster al., 2024, Section B.4). Given class of candidate densities G and i.i.d. samples 21...., 2m I+ (possibly not in 9), we define the empirical negative log-likelihooe (log-loss) of g E G as
Liog (g)
log (2x).
The maximum likelihood estimator is then
Imle E arg min Loog (g).
(20)
gEg
18

---

Published as a conference paper at ICLR 2026
Definition 3 (Log-loss covering number). For class G S A22), we say that subset G' S A(Z) an E-cover with respect to the log-loss if for all g E G there exists g E G' such that supzez log(g(z)/g(<z)) I E. We denote the size of the smallest such cover by Mog(G.E)
We have the following property of MLE's convergence in the squared Hellinger distance with high probability.
Proposition 2. With probabiliiy over i..d samples from any g- E G.
6 log (9, + 46 4s
Di Imlee)
inf 620
+ 2 inf (o + Dx² (9) + 2 Eppt-
9Eg
In particular, if 9 is finite and Eapt the maximum likelihood estimator satisfies
6log (218)))
(i Imlee $
+
inf log + Dxa (g+ l1
m
gEG
Note that the term inf log (1 xx² (g- l1 (g) corresponds to the misspecification error, and is zero if Eg
We note that the proof of (Foster al., 2024, Proposition B.1) contains couple of minor typo- graphical errors, noted here for posterity. Namely, in (20) therein, the authors aim to compare ğ and G. but ended up comparing 9 and g.- A similar mistake is repeated a couple of more times later in the proof without affecting the correctness of the argument.
C PROOFS OF SIMPLE LEMMAS
C.I PROOF OF LEMMA AND COROLLARY
Proof of Lemma 1. Recall that p-11 is the 0-1 reward that assigns reward 1 exactly when the action lies in the support of #(- æ), and otherwise. Hence, by construction, गे achieves reward almost surely under 0 which implies Viga (π) = L, which is maximum possible and thus, V 1. So
ग is optimal w.r.t. r E Rl--1
For the second part, suppose Ves 2 1 - E. By the definition of the 0-1 reward r 0-1 this means that under D and y T(- z), with probability least 1 - y supp(#(- z)). Since से is optimal w.r.t. т., the optimal action set supp(ī(- x)) S (x) almost surely x D. Summarizing this argument below where (I, ~ D n
D, have
((x)))Dx>(( E o..x)) = E[ e a. (z))) > E[1 fy E supp(#( x))}|- = Viga ( 2 1-
Thus,
Px (z)) E
(21)
This immediately implies the desired claim implies the desired claim according the following calculation, which (z, y) ~ D x n:
Vr, (r)) (aaa)=Dxx (x) > (x f (z))r.
1(v El (x)) sup r. (2))
((aa))DXF
v'ey
= EEND
sup r. (x, ( Eyy [1 (1 Ly'y
sup r, (2, v/) (1 PP Pynnt oo,(wi)) yyEy
= EynD
2 E E
sup r. (x. v/) L'Ey
EyND [Pa-s(ta) (y & (x)))
(since r. (xx1 $ 1)
> E
sup r, (x (x. v/) LWEY
-E
(By Eq. 21)
V"
19

---

Published as a conference paper at ICLR 2026
Proof of Corollary 1. This follows directly from Lemma 1. Suppose we can learn the reward class RH-1 with sample complexity m[e, 8) (ef. Definition 1). Then for any demonstrator मे E II and any reward function r, for which स is optimal, with probability at least 1 - 8 over the draw of m(=,6) samples, the output policy से satisfies
Vr. (T)
2 V, -E
1, and any policy that is e- suboptimal under 0-0-1 is also e-suboptimal under r.. Therefore, learning Rm with mmss,) samples immediately yields a learner for II from optimal demonstrations (cf. Definition 2) with the same sample complexity m(=,8).
Here we used Lemma 1-T is optimal for + E R1 with +
C.2 IMPOSSIBILITY OF CLONING: FORMALIZING OBSERVATION 1
We formalize Observation in the following proposition.
N and y = {0, 1). so that it is trivial to achieve Vr, ( with zero samples. However, for any sample size E N there exists an imput distribution DD over A(X) such that distribution matching of some optimal demonstrator (with value 1) is impossible in the following sense:
Proposition 3. There exists simple 0-1 reward class == (r+) with 1, X
inf sup (D= (Es=DDx))<EEwDDDD DH(F( x), #- [x:
#ell.,
2 inf sup Es-(DX;) DTV | I)) ..
x;; S))] 2 4'
# sell.
inf sup BS-(Dxajm (BwD DKL |T S;]) 2 4
Rell.,
and
inf sup (BA DKL (- S) || x))] = 00.
fell.,
Proof. Let X - f0, 1), and define VLex 0. (x) {0, 1}. That is, for every x E X, both outcomes 0 and 1 are correct. This implies V,, (T) can be achieved trivially with zero samples. We will now show the distribution matching to the optimal demonstration remains difficult.
Fix sample size m N. Let D be the uniform distribution on 2 2m.). Let IT C IT, be the
set of all deterministic policies supported on the optimal action sets a, (x) {0, 1} such that for each #(- æ) is either do d1 chosen independently for each Thus [1I| 2
By definition of Hellinger distance, (P,,, Q) > DIv(PP, Q). Therefore it suffices to prove the bound for Dtv. Given samples drawn from D, at most distinct x can appear. Thus
P PE~D(E unseen) 1
2
Now if we put uniform prior on the demonstrators in II, then for any unobserved x and any estimator policy
1
E U ( (.) [)) x;; S)) x unseen]
(PTv(
2
2
Combining these yields:
E Unr EgE (Tv(F(- (+)) 7(- S])]
>
2
7
Thus, by minimax principle:
inf sup EsE DTv(#(- [x), R. x; S)]]
#Elt 2 inf Ex EsLLE DDvV (DV #. [z;
20

---

Published as a conference paper at ICLR 2026
In order to get first direction of KL, note that we have for any unseen x, again have:
1
(ni [ =)| x;; S)) x unseen
2
Taking EgEDDD and applying minimax principle:
inf sup EsEx~D DKL ( >) x;; S))] #ell
2
For the opposite direction of KL, note that for any म., there exists सं E II such that 7. x)) assigns probability 0 to a label that R( 1; S) assigns probability non-zero on an unseen æ, hence:
DKL (f(- 2) x))
In fact all our lower bounds hold almost surely over sampling of S.
D PROOFS FOR SECTION 3
D.I PROOF OF PROPOSITION 1: MLE IS GOOD FOR LOW-CARDINALITY II
Proof of Proposition 1. Consider any unknown but fixed marginal distribution E A(X). First observe that for any S E (X x ))",, the joint law D x mile is the MLE among all joint distributions {D X π π E IJ, because of the shared marginal distribution. Using Proposition 2, for ~ (D T)
Ps D (D X mmle) I (21166)111 21-6.
(22)
We use the above guarantee to obtain the different parts of the proposition.
(1) The first part directly follows by applying Lemma 2 to achieve the change of measure via the squared Hellinger distance with h(z, y) = supy (I,z)) -r. (z, v). Clearly, h X X'yy + [0, 1]. Thus,
Exxama
(sup r. (x. - r(z, - 2 2 (sup r(x, y²) - r(z, v) + DH (D x n. x mlle yEY
V'EY
= D 7
(since में is optimal)
The LHS simply ED (supyeyy (x. v/) r(., v)) V V,, (mee)) by definition. Com- bined with (22), this gives the desired statement that with probability at least - 8,
Vr, (aile) - D (D 7, D mlle)
m
(2) The second part follows from the definition of TV and its relationship with the squared Hellinger distance: with probability least 1 8, for any r. : X
[0,
6 log (2gg())g 111/8)
Vr, (T) - V., (mmte) & Drv (D x म. D mile) W D (D x ग. D Rmle)) S
m
The first inequality follows from the variational definition of the TV distance, the second inequality follows from the standard relationship DTv(PP, Q) I D (IP, Q) and the final inequality follows
from (22) along with noticing that v- an increasing function.
(3) For the final part, to obtain the optimistic rate, we adapt the variance-dependen decomposition of the reward suboptimality in terms of (- derived by Foster al. (2024). Their Theorem 3.1
and Corollary 3.1, written for contextual bandits with rewards bounded in [0, 1]. give:
Vr. (T) Vr. (fiile) $O (V
log (111811) m
log [T [ log
m
Here,
: = :B(=EDDDDx
(u, - -rx(x,1)
< R
(supr=(x, v/) r- (20, v)
(as 2² S 2 for 2 E [0, 1)
= =V (T) = Az
21

---

Published as a conference paper at ICLR 2026
D.2 SIMPLE MLE FAILURES FOR CAW-CARDINALITY R
We first show that there is a simple instance of a support class, where some MLE over the entire class R = U, II, fails.
Proof of Theorem J. Fix any E (0, 1). Let y = (0.1), X = N, and consider the reward class R (ro, r defined by
ro(2,v)= =
and
roi (x. =1 v(z, E x
Thus, under to only the response rewarded, whereas under roi both {0, 1} receive reward
1. Choose the true reward
To. and let the demonstrator मे be optimal under т., meaning x 60(+) for all x; hence every observed label equals
Fix any sample size Let
= m
and let D be the uniform distribution on [g] = {1....,h 9). For sample S {1o,, v:)
(D 7) T)", define the of distinct unlabeled inputs
Sais := {z; iE [m]].
Now consider the predictor mile defined by
(0(+), a eE Sais
mm ( (1)(), x & Sais-
(23)
We show that Till E MT.. (S).
The log-likelihood any x E πT
Log (T;; S)
Σ log (0 x.))
Σ N₂ log 2))
xESa
where (S) counts occurrences of in S. This expression depends only π 7(0 x)) for æ E Sais and is maximized by setting 70 x)) for all x E Sais For x & Sais the likelihood imposes no constraint, and (23) provides a valid maximizer.
We now evaluate the value under the true reward. Since ro(æ, 8) = 1 iff y = 0, the population value of mile is
|Saial
Vr (f = PANDD (~)mm = 0] = Pa E Sais)
Because | S m and q 2 m/c,
Vro (mile I
q
This inequality holds deterministically for every sample S. Thus,
PS--DDa)== (Vr, (mmte) S E) 1.
Since Till is an MLE over II and achieves value most c, the theorem follows.
We show that the attempt to restrict the capacity of the class via another natural choice of considering II U,ER } also does not work, when the expert demonstrations में does not necessarily follow the distribution मr., while still showing optimal demonstrations.
Proof of Theorem 2. Fix E (0, 1) and set == [1/c]. Take X' = {x} and
y {8'} U {a....., a.-. U {b1,...,b.} b.).
so that
= = (s (s- + 2s = 1/], as required. We define two rewards T1, E {0, xxxx
as follows:
ri (x, = [1, E f a and r2(2,
y by..........b 10. otherwise.
{1.
0. otherwise,
22

---

Published as a conference paper at ICLR 2026
Thus, at the unique context z, the set of r1-optimal answers has size s, and the set of r2-optimal answers has size + 1, with (x, = ( (2(,,0) =1} = {v"}
and all other optimal answers disjoint. For each reward E R = F2 and x E A, define the corresponding "uniform" policy on optimal action set
(r)(),)
( x))
[-(())
otherwise.
We now specify the data-generating process. Let D be the point mass at x, and choose the true reward to be 72- Let the demonstrator be the deterministic policy
7 x) = 5y (+),
which always outputs y". Since (2, y²) this demonstrator is optimal, i.e..
Vr. (T) a~ D. ))(()1)) (r2 y) 11==
For any sample size m, every dataset S { y. (D x 8)m is almost surely equal to (1x, }²) )mm i.e., F2 and y/: yy for all i. The likelihoods of policies II are given by
m
Tr (y) x1) = (1=) (((11)
nr, (yi x4i) = ( (x)
Since
#+I we have
1
+1
so r,, is the unique maximizer of the likelihood over R.. Therefore, for every and almost surely for S~ (D hh कr, is the unique element MLF (S).
Finally, we evaluate the value of this MLE under the true reward r. = r2. The policy कr, puts mass 1/s on yy and mass 1/8 on each of a1,..., ax-1, and zero elsewhere. Among these, only y rewarded by r2.
Thus
1
Vr, (mil Vr (T,1) P DDD, Soun,, (1)) (2, 1) 1) n. (y"
[1/
Recall that the analysis holds almost surely over S (D म) and the theorem follows.
Proof of Theorem 2. Fix y E (0, 1) and let s :: [1/>]- Take X' = {a} and
{ U {aI...., a.-13 U {bi...., .,, b.).
so 11+1++(s 1) + 2s 2[1/>]. Define two reward class R S (0,1)xxx such that the rewards F1, F2 have correct answer sets given by a .a,,...a...]
(size s), so or, (x) rr (x) {x*} and they are otherwise disjoint.
or, or,(x) (x) {u" "a,,....
0r²(x) = (u' bi........................ b.} (size s-1),
I = (म : R},
क x)
[x)
0,
otherwise.
Set D to be the point mass at and choose the ground-truth support with data-generating conditional = dy- (always emit y²). For any m, every dataset (D x 7) equals (x,,"))?"
It is simple to see that Tyi MLLETT (S) is the unique maximum likelihood estimator. This is because
m
I (ys xx)) - II ( (
However, with r, = 72, the estimator Amle = Try has the value
1
Vr, ( = Vr, (Tr) = ((.D)))D (V (z)
< 7.
[1/7]
On the other other hand, V,, (T) = because y" E or, (r). All bounds hold almost surely over the sampling of training S, hence the desired theorem holds.
23

---

Published as a conference paper at ICLR 2026
D.3 OVERLAP OF MLE
We now show that MLE over the restricted capacity class RR attains a non-trivial overlap at the log- arithmic sample complexity, though its failure to directly optimize the objective of interest (cf. Sec- tion 3).
Theorem 6. Consider any finite class R S [0. 1|xxXy and an unknown joint realizable distribution DD x n. where ã acts optimally w.r.t. some unknown T* E R. For any estimator #ill E MMLETT (S),
have the following guarantee: for any sample size 2 e- (log [R] + log(1/5))
have
Ps
PawD 3gea.(u) (w) Fyle(y)a) 2
2 1 2 1-6,, for
Jorany any m 2 _ (log [R|+Hog(1) (1))),
where K. supr,a l0, (x)) are the maximum support size.
Proof of Theorem 6. The proof is simple. We first describe some notation: we let
or (x) arg max vIr
yey
We generalize the notion of consistency from Eq. (4) to all the reward functions that agree with the notion of optimality of the observed demonstrations S:
CONS& (S) = (r E R V(ayjes y E (.(()))
We start the proof by noting that the optimal policy supported on the actual ground-truth reward r. has non-zero likelihood. Therefore, any policy in the MLE (S) must have non-zero likelihood as well. Thus, for any r for which т, E MLE (S), we must have that r E CONS&(S). The conclusion of this argument is that MLF (S) S ( r CONSR(S)} Therefore, in order to establish
PS-(DXA)m (PoND ( (1/
1 K
2 1 -
2 21
it suffices to establish
Ps (Vr E CONSæ(() (PNNs (P 0) 2 - c) 2-1
(24)
This is equivalent to:
Ps (Vr E CONS&(S) CCC
(25)
Consider any bad r E R such that PP--D(a,()) O (x. (D 7), we have that becomes inconsistent with probability at least E, i.e., ((e...))((Dx)) (y/ & (,()))) Therefore, for any fixed over the draw of randomness of sample S ~ (D 7)
0) > e. With each draw (zi, yi) ~
Ps(r CONS& (5)) S (1 -1-) =M S 5m
Therefore, by standard union bound,
Ps bad CONSR(S) that survives)
R em S 2
The theorem follows by noting that the |R| 2 < 5 for any m > log [R]+Hg](1/8)
Remark 4. It may be possible to turn this into predictor that directly starts to produce good responses, depending on the overlap among hypotheses and other types of feedback available in post-training (e.g., whether generated response is good not). This overlap by parameter that reflects the need for repeated sampling and the number of feedback that must be queried, which in allows for more quantitative understanding of how many feedbacks are required to guarantee performance in terms of this parameter. However, we leave it open to formulate an interesting setup that enables study of both types of feedbacks together for our problem, and we do not attempt to investigate this any further.
be captured
E PROOFS FROM SECTION 4
We first start with proof of the mistake bound and regret analysis of our Algorithm for the online setup.
24

---

Published as a conference paper at ICLR 2026
R
(r 9x 1}
(r (1)
(r)+ O
Claarr L (0t44)) Goeenaao - (luorr
Figure 2: A visualization of the update rule of Algorithm during hth round. The weight of the hypotheses in the red, green, and blue regions are respectively set zero, unchanged, and doubled.
E.1 PROOF THEOREM
Proof of Theorem 3. Let W++1 = ErEr (r)+)((r) be the total weight in the system after com- pletion of i rounds. We first claim that the total weight W, non-increasing, i.e., W++1 I W,- The intuition that because D. is the weight-maximizer, the total weight of hypotheses such that raxe, Mt) but r(x,, 9c) + (i.e., the added weight) is most the total weight of hypotheses such that r(Z,, Øe) = 1 but (Tt,, yt) # I which we always zero out (see also Figure 2). Formally,
W
Σ
(r) +
Σ
(T)
Σ
(r)
Σ (r)
r(Fc8)==1 and #1)
r(and)=1 and $0
r(x).8))=1
<
Σ
w (r) +
Σ
(()
Σ
(r) - We (26)
r(x.3))== (+)
r(
where the equalities are rearrangements of terms inside the summations, and first inequality follows from the definition of ÿe in Algorithm which implies for any y E y (so Yt in particular):
w (r())) (=.,0)=1
Σ
Σ wh (r)r(())))) FE
S
Σ (r)r(r PER
Σ (r).
(27)
r==,000=1
Now, for any fixed E R. in the mistake round for meaning r(r,, 0c) # 1-the weight of is doubled. And after any number of rounds with M, mistakes w.r.t. some consistent R:
(r) M. s WT+1 S W1 [R], which implies M, S log R|
E.2 PROOF THEOREM
Proof. We will use the elementary inequalities
(1+>)" < 11+ yu and
(1-7)"<1110 -1- yu for u E [0.1], 7 [0, 1].
(28)
Easing the notation, let us denote
>((r.)) r(z,, x) r((,,)) sup r(z v/) r(xe v).
y'ey
25

---

Published as a conference paper at ICLR 2026
We first show that the total-weigh in the system Wi ErEr (1) is non-increasing.
W411 W.-
Σ w
())
(1-)) ))
Σ (r)
FER
FER
w (r)
D1) -de(r: - Σ (r) +1
rER
(by (28) & range(r) S [0,
Σ
n0 (r) Ac(r, 772 Σ w:
(r))e(r Dh) (e)
FE
FER
Σ
(r)((r)) (((e,)
Σ
w (T)Ar,, (o)de(r s 0,
FR
TER
where the last inequality holds because the first term is non-positive (the learner's prediction Ør maximizes the weighted reward) and the second term is non-positive anyway. Hence W++1 S W,, and by induction WT S W1 = |R].
Because } is non- non-increasing, and we started with W1 [R), we have that for any E R
(r) = (1 + (1-)) -000 = WT+1 S W1 = (R].
Taking logarithms and rearranging terms gives
log [R]
log I-
(T) - JT( -
log(1 +7) log +7) (J²(T) - J.()).
log
Finally, using [og[ S + and
S 1++y for 7 E (0, 1), we obtain the stated bound.
Tog(1 [7
We now prove Theorem 5 via an online-to-batch conversion.
E.3 PROOF OF THEOREM 5: ONLINE-TO TO-TACH ANALYSIS VIA FREEDMAN'S 'S INEQUALITY
We now analyze the online-to-batch conversion in Algorithm 2, deriving Theorem 5 with tight op- timistic rates. This is achieved by applying Freedman's inequality (Lemma 3) to an appropriate martingale difference sequence. Here we analyze Algorithm 2 only as applied to Algorithm 1, but believe the same analysis extends more general settings with different loss functions that admit a suitable regret guarantee, and may be of independent interest.
Proof of Theorem 5. We consider the natural filtration F, for the probability space of the first t samples, and denote by Er[-] expectation conditioned this filtration." 8 Fix any r E R. To analyze the value of the policy, observe that
(+ Ex-1(r( Vr(fte) = E [r]], and Vi" LEup y)]
Thus the following forms martingale difference sequence:
z (V. V,())) ((p,, -2 0:))
with | <
Since E, [2 0, the sequence (Z mm forms martingale difference sequence with respect to (F.). To control its conditional variance, note that
E, [2,)0 = Var.-1 = Var.-1 [Suur((ze. v) - r(xt, Dc)
E,
supr(xr,))) r(s,, De) V V2(
'We will always assume that the quantities of interest are measurable with respect to this filtration, without delving into the formal details.
26

---

Published as a conference paper at ICLR 2026
because V2 - (Ve is F1-1 and hence acts as constant. The only inequality above follows from the fact that sup, r(z y) - r( Do) E [0. 1], and for such random variables Var[-] s E[-].
Applying Freedman's inequality (Lemmaa 3) with R and confidence parameter 6/2, for any 7 E (0, 1) we have, with probability at least 5/2.
V - -V+(f() + log(25
Σ
Z S 1=1
Expanding Zi = (V²* V((())) (sup, r r(xe; y) - ((xt, and summing over 1,..., m gives
log(28-
Σ - (T+))
Σ (sup (supr(z., ( (.)) n - V.(
7
Dividing by m and noting that the final averaged policy Ro2b is defined so that
(TT(T)))
t=1
yields
(m) J.(m)
log(20-1)
- (V - (To2)) <
+
m
17m
Now setting 7 10.1 and using the elementary inequality we obtain
- S 1 2y for E 10.1.2).
V - (.o26)) _ (m) - (m)m
+ log(26-1
m(1 - 7)
7(1 7)m
I (1+ <11+2>) - - [m)m) 2 2log(20-1)
(29)
7/m2
We denote the suboptimality gap of the demonstrator by A Vi Vr(TT, and repeat the same analysis the martingale difference sequence"
2 = (supr(z v) - ((I).81) - (V, V, V2)))..
Applying Freedman's 's inequality (Lemma 3) with Rt and confidence parameter 8/2, for any 7. E (0,1) we obtain, with probability at least - 8/2,
m2
log(28-1)
log(28-1
Σ
E. [(Z)]]
< 7m A +
?
7
Substituting the expression for Zi gives
Σ (su,) v) r( vi)) AF S ym A + log(28-1
Rearranging after dividing both sides by yields
(m) - J,(m) < (1+7)
log(28-)
(30)
m
7/m
For this step of the proof, however, direct application of Bernstein's 's inequality also suffices, and is not necessary proceed via martingalc difference sequence.
27

---

Published as a conference paper at ICLR 2026
Now combine Eqs. (29) and (30) and apply union bound over the two events. Using the regret bound from Theorem 4, we obtain the following: with probability at least 6,
(1+ (m² - J,(m))
+ 2 2log(26-1)
V - Vin(F22))) <
(31)
m
7m
2 2log(2 (1+ + 27)² (m) - J, (m)) + (1+ + 27) log |R|
(32)
m
ym
? m
$ < + +7) - + (1 + 2y)²
22)2 (2g(22-1)
(1+ + 2y) log (R]
12g
+
m
m
7 171
(33)
log [R]
6log(28 ? m
$ 61+ A
where we used y E (0, (34)
mm
Subtracting A$ = Vi' - (T) from both sides gives, with probability at least 6,
(F) - Vi((oo2b) s <,> + 2log (R]
6 6log(25-1)
(35)
mm
mm
Finally, choosing
2 log RR| + 1log(28-1
y min
6m
yields for any E R. for which A$ S A, we have with probability at least -
6 6 log R 6lg
4log [R] + 12 log(288))
(*) - V, S 2
(36)
m
m
Finally, applying union bound over r E R and replacing 6 with 8/|R| yields the stated guarantee in Theorem 5 that holds for all r E R simultaneously with probability at least 1 - 6, completing the proof.
The second part of the theorem follows by substituting A = 0. In this case, the only randomness comes from the contexts .1...., Im miid D, and we assume the relevant quantities are measurable with respect to the filtration the probability space of the the contexts.
F pass@k k ERROR MINIMIZATION
As an extension of our setup, we also consider a passâk objective and show that, when the demon- strator optimal, the minimax sample complexity O(lgs++1 [R]). We note that when the demon- strator is suboptimal, it is unclear whether one can compete with the demonstrator with an improved sample complexity of Ologg+1 [R]). For a pass@k metric, the policy μ : - A()*) is allowed to output A responses (V)) .(k,,))))) and it receives the maximal reward among the candidate responses, i.e., max SiSk r. (xt, (1)).
maxiiicA
V,, Vr, EnwD (o)) max r- (x. ( ISisk
(37)
For binary rewards, this corresponds exactly to pass@k-accuracy, Such an objective is often used in benchmarks (e.g.. Chen al. (2021); Orlanski et al. (2022); Dalal et (2025)) and is natural in settings where it is sensible to output multiple completions and let the user choose (e.g., in writing assistance tasks), or when verifying an answer is easier than finding one (e.g., when the answer Lean proof, a Python program that be tested, or perhaps using another LLM verifier). The policy μ may output any joint distribution over the & actions, with arbitrary dependencies for instance, to encourage diversity coverage. This need not be a product distribution corresponding to repeated independent sampling from a single-response policy. Importantly, at training time, the leameer still receives only single correct answer y, for each training question Tz.
Our methods (Algorithm 1) and analysis can be extended to the pass@k k objective. In this setting, the additional flexibility available to the learner yields improved sample complexity when the
28

---

Published as a conference paper at ICLR 2026
demonstrator is optimal. The first modification is in Line 2, where sequentially and greedily as follows: each action ye is selected based on the weighted majority vote for whether (/ is optimal w.r.t. r among the yet "unsatisfied" rewards r (i.e., those for which none of the preceding inconsistent rewards r for which the observed y/ is not optimal, i.e., Yr & o, (xx) ) (as in Algorithm 1), and update
select the l actions
actions are good). Upon receiving Yt. we zero out the weights of
w (r) + (k + 1)+ (r),
(38)
more aggressively by factor of (A: + 1), for all rewards for which the optimal action set o (at)
contains y/ but none of (7/ (1) ..... Di (k)) Counting mistake when none of the generated answer by the algorithm were optimal, this new strategy leads to a mistake bound of logs+1 (R. in the online case (Theorem 8), which then translates to similar sample complexity in the statistical case, using the same auline-to-bateh conversion (Algorithm 2).
Theorem 7 (Learning guarantee for pasttik-metric). Consider any finite reward class R C [0, 1|xxy, For any k 1:
For any sequence (( ((xr,y)))sss I,),,, and any R. for which all the demonstrations are op- timal, ie., Veyr E (xt)),, Algorithm 3 makes at most loga+1 [R] mistake w.r.t. r (i.e.
aii) ViER Vi $ (x)))..
Letting Po2b be the policy given by Algorithm 2, for any ô (0, 1) and sample size m, with probability at least 1 ō, over .1,.... Im midd D, and y ........ Um S.I. Vee|ml Ur E o (at) for some T, E R, we have:
loga41 R + 420gg22
V,, Vr, (To2b))
S
We provide the proof and details of the algorithm in Appendices F.1 and F.2. We show a matching lower bound of - [R]) for both online and statistical settings (Theorem 9) in Appendix F.3.
F.1 ONLINE LEARNER FOR passuk METRIC
We first start by describing an online learner in Algorithm 3.
Algorithm 3 Online pass@k rule with greedy selection and mistake unaware updates Input: A finite model class R C [0, 1](xx)), and parameter i E N.
Initialize w: (r)))) = for all r E R.
In every round, receiving xę:
1. For each y 3) form the slice A = {r E R or (xt)).
2. (Greedy top-k selection). Let 3o- 0. For i 1,2, ........ set
(1) VI
w ( A
U
1
12 (o!
arg max JEY'!
xEV1-1
(Break ties arbitrarily.) Define U, : U
(1))
=(1) 3. Output the l labels Vr ...... (k)
4. (Weight update). Upon receiving label ye:
for all A
(++))))
for all r A m
n U2;
(As + 1) (r)),
for all r E A \ U.- y.
We now prove the guarantee for the online rule, which formalizes the first part of Theorem 7.
29

---

Published as a conference paper at ICLR 2026
Theorem 8 (Online pass@k guarantee). On any sequence VO),,, for any R such that V y E (ze),, Algorithm 3 makes at most loga IR) mistakes with respect to T (i.e., rounds with (()) Niely Di $ (1)) or equivalently (Fr ..... } (z) 0).
(1))
Dr
Proof. Define the potential W, == Erer (r), then we again have {Wi]e is non-increasing. is U2) (A', nU))
W41 = (k: + (0 \ Ut) + w (A2.)
= ( I _w \ A + (A \ Uv) + w: n U.)
kw (A \ Ut) + (A
-w, U A S Wee
where in the first inequality arises from the following key relation:
(U, \ A 2 k w (A \ Ut).
(39)
We defer its proof to the end of this section.
Now suppose the algorithm makes M, passåk mistakes (as defined in Theorem 8) by the end of round w.r.t. r. On each mistake round we must have a+ E A ue Us. its weight is multiplied by (k+ Therefore
(o.+ (k+1)M S W44 S W1 = [R].
which yields M S loga+1 RR).
Proof of the key inequality (39). To simplify the notation, we remove the time index I and write U, as U, and ye as Ay.- Define the uncovered mass in A after selecting first i labels greedily as:
a w An U A.) =E>,
that an = wo ' an t0 (Ay, \U),
and au 2 01 2
2 ag. Correspondingly, define
w ( Ay.))
U An).
$ a:-1 a--1-a a.
=EV--1
In addition, we define the uncovered weight for which of
got picked as
== w (Apo)
U A2)
SEVI-1
By the design of the greedy selections, we have
m: 0:-1
for all i [k].
(40)
With these definitions and relations in place, we are ready to prove the desired inequality. Basic set inclusion/ exclusion tells us that
w (U \ Ayn
w 10)
(A
\ Ap U A Aggu ... U Ant n)
21
U 1:)
w (A (A,m))
U Az) w((h w(t A
1=1
=EV1-1
xED.-1
(m, st),
where the last equality uses the definitions of m, and 8;. Using (40), we obtain
my 8; 2 a; (a3-1 0s)
ai
which allows us to further lower bound (t) (U \ A) as
w
(U (U Ay,))
Σ a k kkk-
Here use the fact that (a:); is non-increasing. This proves the claim.
30

---

Published as a conference paper at ICLR 2026
F.2 STATISTICAL UPPER BOUND
We do the same online-to-batch conversion from Algorithm 2, based on a random stopping time of the online nule in Algorithm 3. We now prove the statistical guarantee for this estimator, complete the proof of Theorem 7.
Proof of Theorem 7. We start by defining some notation, Let μ : X + A()*) be any policy for the pass@k metric. For any reward function X x y + [0, 1] and distribution over prompt D, define the 0-1 loss of the policy as follows.
(p) = BeaDDyy (=(()) a [ e [M),vl) & o,(())]
reduction of reward maximization to loss minimization for an optimal demonstrator, We now show that achieving small loss Lp,r suffices to ensure low value suboptimality w.r.t. By definition:
(i( BR (v.D (0)
max r(a, v()) Isish
ExNDEY
1 E y0 E (x)( max ((.y isish
EyND
sup r(r, y/) TE (v L'EY
[1 [k). 0 E (x)(1))]
sup r(x, 8/) ( (t E [b)) ( & oo(a)))] W
E
EEND
sup r(x, 3/) Ly'E)
BOnD ( (vi [k), & or(())]
(since r, (x. 3/) s 1)
=V-- (p).
This implies
V, - Vz(1 S LD,r(H),
(41)
and thus, it suffices to show that the learner achieves low loss, which we do now.
The estimator achieves low loss. Let Po2b the output policy and (ū+ === 1,.... m) be the policies used in different rounds by the online learner. Consider the natural filtration F, given by the first samples, and let Ex[-] B[- Fe] be the expectation conditioned on it. Let
f = 1 (vi [k] 0 (xx) $ (2,))
Because fú is deterministic function of Ser { v1) : < t}, have
E, E FFol = (Pte.
We define the martingale difference sequence
7 = (p+) f
where |Z|| < 1 almost surely.
Then E, [ZT] and
E.-1 = E, (LD. (7te) - [T) = Var,-1(() < E,-1(E)) = Lo,((().
Applying Freedman's inequality (Lemma 3) with R and confidence parameter 5, for any 1) E (0, 1) we have, with probability at least 1 - 8.
log
<
LD, ( +
"1
t=1
t=
31

---

Published as a conference paper at ICLR 2026
Substituting Zi (per f and rearranging gives
m Σ 4
1og(26-1)
L (T+)
7
By Theorem 8, we have that for any r E R. for which Vie(m) + Sr E (r()). we have Σt= f $ log++1 [R] In particular for r = rx. This yields with probability at least
log(28-1)
(1
Lp.r (Tt) S 7) logg+1 [
Substituting 7) = 1/2, and dividing both sides by m gives us
m
LD.r, (Tte) S logR+1 [R] 1og(225 (25-1)
172
m
Finally, noting that the final averaged policy Bo2b is defined so that
LD,r, (Ro2b)
Σ tm1
LD.r, (7x),
yields with probability least ò
V Vr, (Pa2h) S LD.r. (Pozh)) = Σ
1
Lr., (Pt) S loga+1 [R] + 1ogg (25-1)
t=1
here the first inequality follows from the equivalence derived earlier in Eq.(41).
F.3 LOWER BOUNDS FOR ONLINE AND STATISTICAL SETTINGS FOR passukk- METRIC
We next provide a lower bound that, information- thhrrreiillttorrtccallll. this dependence cannot be improved and we only gain factor of 1/ loo k in sample complexity well as mistake bound in the worst- case. The lower bound instance shown here holds for even the simple special case of binary rewards and multiclass classification,
single correct answer.
Theorem 9 (Online 2 [R]) pass@k mistake bound even for multiclass classification). Fix integers & > and 2 2. There exists problem instance R yx with [R] S d, [D]] + 1, 1 = [log&+1 d] such that for any deterministic online learning algorithm that outputs at most l labels, there exists sequence (It, Gt (E(x))))) realizable by some r. E RR such that it makes mistake on every round.
Note that our instance is an instance of multiclass classification problem r X + y. This is isomorphic to instance R C {0, 1) xxy, where r(æ, y) for only single y for every æ E X.
[loga+1 d] and take X {1...., m) and Consider the full product class R = y which has size (k+1)mm < d. First of all, observe that in (k) any round in which y, does not belong to the list of (7! (1) Dr the mistake is made because are in the multiclass classification setting.
Proof of Theorem's 9. Let
{1..... k+1).
For rounds t E [m], present a fresh coordinate 2; = t. Since | = k+1, there exists a label ys E y that the learner failed to output in the set; Y4 mistake on every round. Moreover, this sequence realizable since R = y contains all functions from X to y.
<lf) for all [B)]. Reveal this yt- This forces
We now prove lower bound for the statistical setting. For function X' > V, we write r(x) for the unique label it outputs on input æ. Let D x r denote the joint distribution where D distribution over contexts and y r(x) identically.
32

---

Published as a conference paper at ICLR 2026
Theorem 10 (Statistical lower bound of Sl(log [R()). Fix integers k 2 1 and 2 1. Let X' {1...., q). y {l..... 2k}. and take the hypothesis class R = =xx (all multiclass functions), so
its cardinality is d == [R] (2k)9. Let D be the uniform distribution on X. Then for any estimators T (X x vy + ())
inf sup (Dx(D PEWD P ((x) & [(x)] rER
( (1- 9
In particular, ensure expected error most E < 1 for all R, one needs ln(1/(2E)) m 2 - In(1 1/9)
2 In
23
Proof. Fix any (possibly randomized) estimator f.. Let 5 = { (x, y1 } be the training sample
drawn i...d from (D r) for ~ (R), and let Us (z, s m) C X' be the set of distinct inputs seen in S. Draw x ~ DD independently of S and then (ιx) pu( xr).
On any x $ Us, under the prior where uniform over R, for any (possibly randomized) k-list 7(x) ~ MV- x),
# of distinct labels in ((x)
Pr[r(x) E ((x) S, x $ Us] = E
S, x $ Us
D
2k
2
so (r(r()) ( S, Us! 2 (Allowing duplicates in (2x) cannot decrease this probabil- ity.)
I Us, the learner can always include the observed label and incur fore, for any estimator 7t,
that There-
1 2 Pla & Us}.
P (1 $ (xx) |s]
Taking expectation over S and using Unif(A) yields
Es |us/u 1 (1- m
Es P. F... (r(x) $ x())
2
since E|[Us|| = 4(1 (1 ) Finally, by minimax principle
inf sup Es P P..DDDPap P r(x) [(x)] 2 inf Ey-Unii(RR Es Pe-DPPa r(x) & [())
rER
5 (1 -
2
For the sample-complexity bound, solve 1 (1 1 1 < for and use - In(1 114) S 1/g.
Because d = (2k)),, we have = loge d, so the bound implies m S logg d) under D Un).
Remark 5. Both online (Theorem 9) and statistical lower bounds (Theorem 10) for pass@k essen- tially demonstrate that one cannot do better than memorization below 2ioog, d) barrier in the worst- case, even for the special case of the problem realizable multiclass classification R S (0, 1) x with single correct answer and always correct demonstrations.
33


Reviews:
[{"review_id": "0kyt2n8Fwm", "reviewer": "Reviewer_mbR1", "summary": "This paper studies the offline imitation learning problem in contextual bandits, under the model that only correct demonstrations are provided. This problem is of practical relevance for the supervised training of LLMs. Departing from prior work, the authors do not assume access to a small policy class (to model the conditional distribution of responses given prompts), but only the weaker assumption of access to a small \"reward\" class, assigning a binary reward of 1 to responses that are correct and 0 to responses that are incorrect. It is shown that standard algorithms (i.e., MLE) can fail under this weaker assumption, and new algorithms are devised that learn under this relaxed reward class assumption. Several extensions are provided.", "questions": "- Do you envision convenient parameterizations of the reward class S (e.g. linear classifiers) where efficient implementations are possible?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The shift from policy classes to reward classes is well-motivated and potentially practically interesting. The presentation is very clean. The technical results are interesting and check out (at a glance). The results establish that MLE fails and introduce new algorithmic ideas for imitation learning from reward classes, and these are tight in terms of their dependence on the log-cardinality class. The authors study several extensions of interest, such as non-optimal demonstrators and pass@k metrics.", "weaknesses": "I don't have many complaints. The main drawbacks include: 1. the analysis assumes a finite reward class, and since the algorithm is Majority/EXP-style this seems non-trivial to extend, 2. the algorithm requires enumeration over S and is thus not computationally implementable. It would be highly desirable to obtain a more computationally tractable algorithm.", "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": "xTNpvqVaRV", "reviewer": "Reviewer_FNKo", "summary": "This paper studies the problem of learning to generate answers (or completions) to questions (or prompts) when there can be multiple correct answers, based on demonstrations of correct answers. Compared to prior work, the main novelty of this paper is that it relies on assumptions that **the reward model lies in a low-cardinality class**, rather than the assumption that **the demonstrator belongs to a low-complexity policy class**.\n\nSpecifically, this paper formulates the problem in Section 2 and discusses why MLE can fail under this problem formulation in Section 3. Novel learning algorithms are proposed and analyzed in Section 4. Various possible extensions are discussed in Section 5. Concluding remarks and possible future work are discussed in Section 6.", "questions": "- Please add some experimental results if possible.\n\n- [Minor] I think the \"learning from arbitrary demonstrator\" setting (Section 5.2) is a more general and more interesting setting. It is not clear to me why the authors don't present that setting as the main problem formulation of this paper, and discuss the case when $\\pi_*$ is optimal as a special case. I am wondering why.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Overall, I think this is a thought-provoking and solid theoretical paper. Specifically,\n\n- This paper is motivated by very practical problems in LLMs, i.e., many QA problems in math/coding can have many equally valid but differently written solutions. It has also developed a rigorous contextual bandit framework to analyze them. More importantly, it has identified major limitations of MLE under this framework. As a researcher in the field of RL/bandit and LLM, I find this paper very thought-provoking.\n\n- The learning algorithms proposed in this paper are interesting.\n\n- The proposed learning algorithms, and some of their extensions, have been extensively analyzed. To the best of my knowledge, the analyses are rigorous and correct.\n\nOverall, I think this is a strong paper and recommend accepting it.", "weaknesses": "- One obvious limitation of this paper is that it has not provided any experimental results. Given that this paper is motivated by practical problems in LLMs, I recommend that the authors include some experimental results on LLMs in this paper. I think the experimental results will show the following:\n  - how well the proposed algorithm will perform in **practical** LLM settings, compared to MLE-based algorithms;\n  - the computational tractability/efficiency of the proposed algorithms in **practical** LLM settings.\nThese are hard to tell from mathematical analyses.", "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": "HgCrrdc3q6", "reviewer": "Reviewer_Cuid", "summary": "The paper studies the problem of learning from demonstrations when several correct outputs may exist for the same input. The authors argue that maximum likelihood estimation (MLE) fails in such cases because it imitates demonstrator distributions rather than learning correctness. They formalize this as an imitation learning problem with a low-cardinality reward class rather than a low-capacity policy class. Under this assumption, they show that MLE can overfit and fail to generalize, and propose an alternative “mistake-unaware” weighted-majority learner that provably converges to the correct reward mapping with logarithmic sample complexity. The theoretical results extend to non-optimal demonstrators and pass@k metrics", "questions": "1. Could the authors comment on whether variants of their learner, perhaps approximate or sampling-based, might make it computationally feasible in larger hypothesis spaces?\n\n2. Are there simple synthetic or illustrative experiments that could highlight the key behaviors (e.g. when MLE fails versus when the proposed learner succeeds)? Even a toy example might help ground the theory.\n\n3. The paper assumes a finite hypothesis class; how sensitive are the results to mild relaxations of this assumption?\n\n4. How might the “low-cardinality reward class” assumption be instantiated or approximated in realistic domains like LLMs?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper tackles a relevant and well-motivated problem. The gap it identifies in how MLE behaves when multiple valid answers exist is real and interesting.\n\n2. The proposed solution is conceptually neat: shifting from reasoning about small policy classes to small reward classes gives a different and insightful perspective on imitation learning.\n\n3. The theoretical development is rigorous and internally consistent, with clear proofs and sound use of classical online-learning tools such as weighted-majority and halving arguments.\n\n4. The exposition of the core idea, which distinguishes learning correctness from imitating distributions, is well articulated.\n\n5. The paper effectively demonstrates that MLE’s failure is not statistical noise but a structural limitation, and that an alternative learner can overcome it in theory.", "weaknesses": "1. The work is entirely theoretical, with no experiments or simulations to illustrate MLE’s failure or the proposed learner’s behavior. This limits how convincingly one can assess its practical value.\n\n2. The learner assumes a finite discrete hypothesis class and realizability (σ* ∈ S). While the authors acknowledge this, there is little discussion on what these assumptions imply for larger or continuous models, like LLMs.\n\n3. The computational demands of maintaining weights over all hypotheses are briefly noted but not discussed in depth; in realistic settings, this would be infeasible.\n\n4. While the motivation from LLMs is compelling, the theoretical framework (finite S, symbolic correctness sets) cannot directly model text generation. The connection to practical SFT or reward-learning pipelines remains purely conceptual.\n\n5. The introduction could be more focused: it introduces several formal terms early on, which could instead appear in the problem-definition section. Presenting the motivation, gap, and main contributions upfront would make it easier to follow.\n\n6. Although the proposed approach closes a theoretical gap, it remains unclear how it might inspire a usable method in practice. A short discussion of potential approximations or practical directions would help strengthen the paper.", "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": "Q7n2SVVVdP", "reviewer": "Reviewer_aCdu", "summary": "This paper addresses learning to generate answers to questions where multiple correct responses exist. The authors frame this as an offline imitation learning problem using training demonstrations, similar to supervised fine-tuning. Unlike previous approaches that assume demonstrators follow simple policies (justifying maximum likelihood methods), the authors make a weaker assumption: only the reward model (determining answer correctness) needs to be simple. They demonstrate that standard likelihood maximization can fail under this assumption. Their key contribution is a new learning algorithm with sample complexity that scales logarithmically with the number of possible reward functions. The method works robustly even with imperfect demonstrators and in relaxed evaluation settings (pass@k). The work suggests moving beyond likelihood maximization for learning from demonstrations may be beneficial.", "questions": "1. Why did you not carry out any empirical experiments?\n2. Are you able to actually implement your algorithm in code?\n3. How do you square your theoretical work predicting the failure of MLE, when in practice it appears to work reasonably well (if not perfectly)?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "1. The fundamental motivation, that a learner needs to produce a single correct answer and does not need to exactly mimic the demonstrator's correct answer, is a good one\n2. The mathematical analysis carried out in the paper is extensive and rigorous\n3. The paper is well-written and clear in its ideas, methods, and analysis", "weaknesses": "1. The reliance on reward model classes being finite and reward being binary is quite restrictive and limits the applicability of the theory of the method to tasks which are fuzzier and outside of domains such as maths. Even in programming one might want more nuanced reward signals (e.g. % of unit tests passed + some reward on how well formatted / structured the code is).\n2. For domains with large state spaces (e.g. long context language modelling, often found in agentic tasks), the methods may fail to be performant due to their dependence on the size of the state space\n3. No empirical experiments are carried out, not even for toy tasks.\n4. It's not clear algorithm 1 is feasible to implement, since it takes in as inputs distributions over binary reward functions, and involve arg-maxing over outputs with terms inside the arg-max summing over huge spaces of reward functions. Outside of problems with few contexts and responses (in the contextual bandit sense), the algorithm is likely computationally intractable. This does not support the claim that your learner is in any way \"efficient\".\n5. Critiques of the MLE paradigm are not supported by experiments or analysis of experimental data. Given the ability of LLMs to produce reasonable outputs (even if not perfect) after SFT training, this points to the mathematical framing of the paper being flawed, since some of its conclusions directly contradict the current state of affairs.\n6. The learning algorithm is claimed to be logarithmic on $|S|$, however, given $S \\subseteq (2^{\\mathcal{Y}})^{\\mathcal{X}}$, this makes the learner's convergence time scale linearly with the product of the context and output spaces in the worst case. Given the domain of interest is LLM finetuning, where these spaces are token sequences, the learner has an exponential dependence on the lengths of the input and output sequences. It is disingenuous to classify this as \"efficient\".", "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)