PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal Multi-Reward Reinforcement Learning

Zijun Chen, Zihan Zhang

arXiv:2609.36486v1cs.LGstat.ML

TL;DR

The paper studies how to learn near-optimal policies for multiple known rewards in one shared environment, rather than using separate or reward-agnostic exploration. It uses optimistic and conservative value estimates with gap-based multiplicative weights, and reports a minimax sample-complexity bound with no additional burn-in term. The upper bound matches the lower bound up to polylog factors when log M ≤ S.

  • Problem

    Existing results typically focus on a single reward, while reward-agnostic exploration does not use rewards that are known in advance and may collect more data than the given reward class requires.

    Source: §1 · pp. 1–2

  • Method

    The approach estimates reward-specific surpluses from optimistic and conservative value gaps and uses multiplicative weights to turn gap-weighted progress into simultaneous guarantees.

    Source: §4 · p. 7

  • Results

    The algorithm has minimax-optimal sample complexity with no additional burn-in term, and when log M ≤ S its upper bound matches the lower bound up to a polylog factor.

    Source: §1.1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Policies for multiple rewards can be learned jointly with at most an additional log M factor over the single-reward minimax sample complexity.

    Source: §1.1 · p. 2

    The stated setting assumes M ≥ 2 known rewards sharing the same transition kernel; the single-reward case is covered by Zhang et al..

    Source: §2 · pp. 3–4

Abstract

from arXiv · show

We study an unknown-transition finite-horizon Markov decision process (MDP) with a finite collection of known reward functions $\{r^1, r^2, \ldots, r^M\}$. The goal is to output an $\epsilon$-optimal policy for every reward using online episodic interaction only. Performance is measured by the policy error $V_{0}^{*, m} - V_{0}^{\widehat\pi^{m}, m}$ where $m\in [M]$ represents the reward function and $V_{0}^{*, m}=\mathbb{E}_{s_1\sim \mu}[V_{1}^{*, m}(s_1)]$. Under this setting, we design a provably efficient algorithm to establish a minimax sample complexity bound of $$ O\left(\frac{SAH^3}{\epsilon^2}\log M \mathrm{polylog}\left(\frac{SAH\log M}{\min\left\{\epsilon, 1\right\}\delta}\right)\right)$$ episodes, with no additional burn-in cost. This matches the information-theoretic lower bound up to a factor of $ \mathrm{polylog}(SAH\log M/(\min\left\{\epsilon, 1\right\}\delta))$. Our method combines three technical ingredients. First, we adapt MVP to reward-switching learning to construct optimistic value estimates. Second, we use fresh replay samples to conservatively evaluate the candidate policies. Third, gap-based multiplicative weights updates adjust the reward-sampling distribution using the differences between these estimates, converting weighted learning progress into simultaneous guarantees for all rewards.

1 Introduction

The paper addresses simultaneous policy learning for multiple known rewards in a shared environment, between single-reward learning and reward-agnostic exploration. It gives a jointly efficient approach and establishes a near-minimax sample-complexity result under a stated condition on M.

1 Introduction

  • A naive separate run for each reward costs a factor proportional to M, while reward-agnostic exploration can require more data than the given reward class needs.The setting instead allows the M reward functions to be known during exploration and uses their shared transition dynamics.Source: §1 · p. 2

1.1 Our Contribution

  • The algorithm learns policies for all rewards jointly with at most an additional log M factor over the single-reward minimax sample complexity.The contribution also adapts a reward-agnostic lower-bound construction to simultaneous policy learning.Source: §1.1 · p. 2
  • When log M ≤ S, the upper bound matches the lower bound up to a multiplicative polylog factor.The cited single-reward result establishes minimax-optimal sample complexity without an additional burn-in term.Source: §1.1 · p. 2, §1 · p. 2, §1.2 · p. 2

1.2 Related Works

  • Prior reward-agnostic methods explore without access to downstream rewards, whereas this setting uses a fixed reward class available during exploration.The related work also situates the question against single-reward minimax learning.Source: §1 · p. 2, §1.2 · p. 3

2 Preliminaries

The preliminaries specify a finite-horizon, time-inhomogeneous MDP and define the simultaneous learning objective across known rewards. They also introduce notation for variance under a transition distribution.

2 Preliminaries

  • The setting is a time-inhomogeneous finite-horizon MDP with finite state and action spaces, stage-dependent transitions, and rewards in [0, 1].Policies map states to action distributions, and episode-level mixtures are also allowed.Source: §2 · p. 3
  • The M known reward functions share one transition kernel, and the stated results focus on M ≥ 2.The single-reward case is covered by Zhang et al..Source: §2 · pp. 3–4
  • The PAC objective is to return a policy for each reward whose value is ε-optimal, simultaneously with high probability.Each reward function induces its own optimal policy and value function.Source: §2 · p. 4

Notations.

  • V(P, X) denotes the variance of X under distribution P, computed as ⟨P, X²⟩ − ⟨P, X⟩².Here X² is the element-wise square of X.Source: Notations · p. 4

3 Algorithm

The algorithm combines shared-model multi-reward learning, fresh sealed replay, conservative candidate-policy evaluation, and gap-based multiplicative weights. It returns a uniform mixture of counterfactual policies for each reward.

3 Algorithm

  • Each phase generates reward-specific candidate policies and optimistic estimates, collects fresh replay samples, evaluates candidates conservatively, and updates the reward-sampling distribution.The procedure adapts the MVP framework to the multi-reward setting using shared-model learning, sealed replay, and multiplicative weights.Source: §3 · p. 4
  • Multi-Reward-Switching-MVP computes optimistic reward-specific values and counterfactual policies, then samples a reward according to pℓ to choose the rollout policy.The empirical transition model and doubling-schedule statistics carry across phases, allowing later learning to use earlier transitions.Source: §3 · pp. 4–5
  • Sealed replay builds a fresh empirical transition model from replay samples retained for sufficiently visited state-action-stage tuples; those samples do not update the learning model.Conservative evaluation shrinks continuation values according to visit counts rather than subtracting an explicit pessimistic bonus, controlling upward fluctuations when counts are small.Source: §3 · p. 5

Gap-based multiplicative weights update.

  • Rewards with larger optimistic–conservative certificate gaps receive larger multiplicative weight increases, and each reward’s output policy is a uniform mixture over its phase-and-episode counterfactual policies.Each phase uses K learning episodes and 8K replay episodes, for 9KL total interaction episodes across L phases.Source: Gap-based multiplicative weights update · pp. 5–6

4 Sample Complexity Analysis

The analysis extends MVP’s burn-in-free guarantee to multiple rewards, with optimal logarithmic dependence on the number of rewards up to additional logarithmic factors. It combines optimistic learning, conservative evaluation using fresh replay samples, and gap-based multiplicative weights to control error across rewards.

4 Sample Complexity Analysis

  • Theorem 1 gives a minimax-optimal multi-reward sample-complexity guarantee, with logarithmic dependence on M up to additional logarithmic factors.The theorem’s phase construction returns an epsilon-optimal policy for every reward with probability at least 1 − delta.Source: §4 · p. 6

4.3 Gap-based Multiplicative Weights Update

  • Gap-based multiplicative weights convert reward-weighted learning progress into simultaneous error control without requiring every phase’s evaluation gap to be nonnegative.Conservative evaluation provides conditional-mean and exponential-moment control, limiting the additional overhead to O(H log(M/delta)/L).Source: §4 · pp. 7–8

4.1 Multi-Reward Learning Guarantee.

  • Optimistic value estimates bound policy suboptimality, while the learning analysis transfers selected surplus to its weighted average.The optimism and weighted-surplus guarantees hold with high probability.Source: §4.1 · p. 8

4.2 Conservative Evaluation with Sealed Replay

  • Fresh replay samples independently evaluate candidate policies conservatively, with replay counts matching learning-phase visits on retained tuples.The replay procedure executes each recorded rollout policy eight times; the count guarantee holds with high probability.Source: §4.2 · pp. 8–9

5 Numerical Experiments

Experiments compare the algorithm with four reward-agnostic baselines on random nonstationary finite-horizon MDPs. Across the tested budgets and reward counts, the method generally reaches low worst-reward error with modest budget growth as M increases.

5 Numerical Experiments

  • The experiments used random nonstationary finite-horizon MDPs with S = 10, A = 5, and H = 6, measuring average worst-reward policy error.The first experiment fixed M = 16; the second estimated the budget needed to reach mean worst-reward error 0.1 as M varied.Source: §5 · p. 9
  • At M = 16, the algorithm reached mean worst-reward error 0.1 at 2,048 episodes, earlier than the four baselines on the tested grid.It had the lowest mean worst-reward error at 7 of 9 budgets; Jin et al. [10] had lower point estimates at 8,192 and 32,768 episodes.Source: Figure 1 · p. 10, §5 · p. 10
  • At 10 of 11 tested reward counts, the method required the lowest estimated budget; Ridel and Cohen [18] performed better at M = 96.As M increased from 4 to 128, its estimated budget rose from approximately 1,240 to 3,153 episodes (2.54×).Source: Figure 1 · p. 10, §5 · p. 10

6 Discussion

The paper presents a minimax-optimal sample-complexity result for learning policies across multiple rewards without an additional burn-in term. It identifies reward-agnostic burn-in removal and adaptation to reward-class complexity as future directions.

6 Discussion

  • The proposed algorithm has sample complexity eO(SAH^3 log(M)/epsilon^2), is minimax optimal, and has no additional burn-in term.The stated result concerns reinforcement learning with M different rewards.Source: §6 · p. 10
  • One open direction is whether reward-agnostic reinforcement learning can achieve minimax-optimal sample complexity without burn-in.The paper frames this as future work.Source: §6 · p. 10
  • Another open direction is adapting sample complexity to the intrinsic complexity of the reward class rather than its cardinality alone.The paper asks whether shared reward structure can yield sharper bounds.Source: §6 · p. 10

Appendices

The appendices provide MDP notation, technical lemmas, and auxiliary algorithms supporting the analysis. They include concentration results and procedures for sealed replay and conservative evaluation.

A.1 Notations

A.2 Technical Lemmas

  • The technical lemmas include bounds for bounded i.i.d. variables and martingale concentration.The appendix includes Bennett’s inequality, an empirical-mean setup, and Freedman’s inequality.Source: §A.2 · p. 15

B Auxiliary Algorithms

  • Sealed-Replay takes retained tuples, fixed sample counts, and rollout policies, then returns replay estimates and counts.It runs eight replay rounds and collects fresh transitions for retained tuples until their target counts are met.Source: B Auxiliary Algorithms · p. 16
  • The auxiliary conservative-evaluation algorithm evaluates a policy under a reward using an evaluation kernel and count function.Source: B Auxiliary Algorithms · p. 16

C Missing Proofs in Section 4

The proof combines optimistic learning, conservative evaluation, and reward-weight updates to guarantee simultaneous reward-specific performance. It accounts for learning and replay interaction in the episode total.

C Missing Proofs in Section 4

  • When Algorithm 1 runs for L phases with K learning episodes per phase, it returns episode-level mixture policies {πm}m∈[M] with probability at least 1−δ, using 9KL total episodes, including learning and replay.Source: C Missing Proofs in Section 4 · pp. 16–17

C.2 Conservative Evaluation and Multiplicative Weight Update Parts for Multi-Reward MDPs

  • The conservative-evaluation and update analysis combines optimistic error with evaluation bias to control reward-specific cumulative certificates across phases.The proof outline assigns the weighted optimistic error and conditional mean evaluation bias to separate supporting analyses before combining them.Source: C Missing Proofs in Section 4 · pp. 17, 19

C.1 Optimistic Evaluation Part for Multi-Reward MDPs

  • The optimistic-evaluation analysis extends the MVP optimism argument to M reward functions, with concentration events covering the learning episodes.The stated proof structure follows and adds concentration events across rewards and learning episodes.Source: §C.1 · p. 19, §C.1.1 · pp. 19–20

C.1.1 Proof of Lemma 1

  • The Lemma 1 proof uses backward induction to establish optimism for the reward-specific value estimates under the shared empirical transition model.The shared transition model and its counts persist across phase boundaries.Source: §C.1.1 · pp. 19–20

C.1.2 Proof of Lemma 2

  • The Lemma 2 analysis bounds cumulative optimistic error over learning episodes and transfers the guarantee to each reward using concentration and initial-state handling.The proof notes that the initial-state correction is unnecessary at this step because the quantity already uses recorded initial states.Source: §C.1.2 · pp. 20–22

C.2.1 Retained row receives sufficiently many fresh replay samples.

This proof establishes that sealed replay supplies fresh samples for retained rows and supports conservative evaluation. Its arguments use conditional sampling bounds and a fresh empirical transition model.

C.2.1 Retained row receives sufficiently many fresh replay samples.

  • With probability at least 1−2SAHLδ′′, sealed replay supplies each retained row with as many samples as were collected during learning.The proof conditions on the pre-replay history, when policies, retained rows, and quotas are fixed, and uses conditional independence and a Chernoff bound.Source: §C.2.1 · pp. 24–25
  • On the learning-count event, each retained row has at least 8 log(1/δ′′) learning visits before its replay quota is set.This count condition is used to establish that the sealed replay quotas are filled.Source: §C.2.1 · p. 24

C.2.2 Proof of Lemma 4

  • For conservative evaluation, the proof constructs a fresh empirical transition model from replay samples and bounds its values within the finite-horizon range.Under the replay success event, the model uses Nℓ,h(s,a) fresh samples; the recursion gives values between zero and H−h+1.Source: §C.2.2 · pp. 25–26

C.3 Proof of Lemma 5

The Lemma 5 proof controls conditional mean evaluation error across retained and unretained state-action rows. The theorem proof then combines these controls with optimism and concentration to establish the algorithm’s guarantee.

C.3 Proof of Lemma 5

  • With probability at least 1−SAHLδ′′, the conditional mean evaluation-error bound holds for every phase, and the same event controls all rewards through weighted occupancy.Retained and unretained rows are handled separately; every retained tuple contributes at most 18 log(1/δ′′) to the retained part.Source: §C.3 · pp. 28–30

C.4 Proof of Theorem 1

  • On the optimism event, the theorem proof combines the cumulative optimistic bound with the conditional mean evaluation bound to control the cumulative certificates.The final proof intersects its confidence events and states that all returned policies are ε-optimal with probability at least 1−δ.Source: §C.4 · pp. 31–32

D Supplementary Proofs in Section 4

  • The supplementary proof distinguishes persistent learning statistics from replay data: replay samples do not update the learner’s empirical rows.Learning episodes are indexed consecutively across phases, while replay samples remain separate from the persistent learner.Source: D Supplementary Proofs in Section 4 · p. 32

D.1 Proof of Lemma 14

The supplementary proofs establish concentration controls for the persistent learner and its value estimates. The appendix also describes a numerical experiment’s setup and evaluation procedure, without giving its outcomes in these passages.

D.1 Proof of Lemma 14

  • The Lemma 14 proof bounds cumulative optimism surplus over learning episodes with probability at least 1−5δ′.The argument includes martingale sums and applies the bound to the entire run, including histories with replay outcomes and weight updates.Source: §D.1 · p. 33

D.2 Profile-Based Concentration Lemma for Multi Reward MDPs

  • The profile-based concentration lemma gives a uniform high-probability bound across deterministic profiles, batches, and candidate-label assignments.Its stated conditions include KL≥max{2,SAH}; the proof establishes a uniform event before substituting the realized profile.Source: §D.2 · p. 34

D.3 Statement and Proof of Lemma 20

  • The Lemma 20 proof applies profile-based concentration to auxiliary values and their squared, normalized values.Including zero among the candidate vectors allows the event to bound positive parts of model errors and their signed sum.Source: §D.3 · pp. 35–36

D.4 Statement and Proof of Lemma 21

  • The Lemma 21 proof reuses the preceding value-family concentration event and incurs no additional failure probability.It uses the variance bound from Lemma 20 and absorbs a linear term under the stated condition.Source: §D.4 · pp. 36–37

E Numerical Experiments

  • The numerical experiments use random nonstationary finite-horizon MDPs with 10 states, 5 actions, and horizon 6.Rewards vary by reward function, stage, and state, but are shared across actions at each state and stage.Source: E Numerical Experiments · p. 37

E.1 Experiment 1: Policy error versus episode budget

  • Experiment 1 compares mean worst-reward policy error across episode budgets using 12 paired environment–learner seed runs per budget.Its pointwise bootstrap intervals hold configurations fixed and do not account for configuration selection.Source: §E.1 · p. 37

E.2 Experiment 2: Required episode budget versus reward count

Experiment 2 estimates the episode budget needed to reach a mean error of at most 0.1 as the number of rewards varies. It uses repeated environment–learner seeds and bootstrap intervals to characterize the estimated crossing.

E.2 Experiment 2: Required episode budget versus reward count

  • The experiment varies M from 4 to 128, using eight independent environment–learner seed pairs and 128 reward functions per fixed transition kernel.For each M, it uses the first M reward functions and reruns each algorithm for every (B, M) setting.Source: §E.2 · p. 37
  • For each reward count M, the experiment estimates the episode budget B at which Error8(B, M) reaches at most 0.1.The crossing is selected where the mean error and mean errors at all larger tested budgets are at most 0.1.Source: §E.2 · pp. 37–38
  • The estimated crossing is interpolated on log error versus log budget, with budget grids refined near the crossing.Pointwise 95% intervals use 10,000 paired bootstrap resamples of the eight seed pairs; they quantify empirical variability under selected configurations, not the simultaneous PAC guarantee.Source: §E.2 · p. 38

F Lower Bound Construction

The lower-bound construction embeds independent hard transitions at Θ(SH) time–leaf pairs while reusing physical states. A fixed, publicly known reward class makes simultaneous ε-optimality imply exact recovery of the hidden indices, yielding a lower bound under the stated parameter conditions.

F Lower Bound Construction

  • The construction places independent hard transitions at Θ(SH) time–leaf pairs while sharing the same physical states.It uses a waiting state, a deterministic binary tree, and shared absorbing states; time indices add no physical states, and each episode has at most one informative transition.Source: F Lower Bound Construction · pp. 38–40
  • The construction uses balanced vectors and a decoder property to distinguish the hidden action indices from observations.For each action, the hidden tuple selects an observation distribution; the decoding argument states that distributions satisfying its condition determine the index exactly.Source: F Lower Bound Construction · pp. 38–39
  • A fixed, publicly known reward class is defined independently of the hidden transition tuples and may be revealed before exploration.The class contains rewards indexed by time–leaf pairs and packing vectors; its size is bounded by SHA^2L^2.Source: F Lower Bound Construction · pp. 38, 40
  • Simultaneous ε-optimality for the reward class permits exact recovery of the entire hidden-index array.The construction’s proposition and corollary give lower-bound claims for specified state, action, horizon, reward-count, and accuracy conditions.Source: F Lower Bound Construction · pp. 38, 41
  • The lower-bound proof compares stopped interaction transcripts under hidden kernels and a reference kernel using an adaptive KL chain rule.The cited passage bounds the divergence between each hard-transition distribution and the reference distribution by α^2.Source: F Lower Bound Construction · p. 41
Loading 2609.36486v1…