Optimal Multi-Reward Reinforcement Learning

Zijun Chen; Zihan Zhang

cs.LGSep 29, 2026

Optimal Multi-Reward Reinforcement Learning

Authors: Zijun Chen, Zihan Zhang

Organizations: Department of Computer Science and Engineering Hong Kong University of Science and Technology

Abstract

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 $ε$ -optimal policy for every reward using online episodic interaction only. Performance is measured by the policy error $V_{0}^{*, m} - V_{0}^{\widehatπ^{m}, m}$ where $m\in [M]$ represents the reward function and $V_{0}^{*, m}=\mathbb{E}_{s_1\sim μ}[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}{ε^2}\log M \mathrm{polylog}\left(\frac{SAH\log M}{\min\left\{ε, 1\right\}δ}\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\{ε, 1\right\}δ))$ . 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.

Figures & tables

Reference	Sample complexity	Comment
[ 10 ]	$\widetilde{O}\!\left(\frac{S^{2}AH^{5}}{\epsilon^{2}}+\frac{S^{4}AH^{7}}{\epsilon}\right)$	Reward-free exploration.
[ 25 ]	$\widetilde{O}\!\left(\frac{SAH^{5}\log M}{\epsilon^{2}}\right)$	–
[ 15 ]	$\widetilde{O}\!\left(\frac{S^{2}AH^{3}}{\epsilon^{2}}\par\right)$	Reward-free exploration.
[ 12 ]	$\widetilde{O}\!\left(\frac{SAH^{3}\log M}{\epsilon^{2}}+\frac{S^{4}A^{4}H^{6}\log M}{\epsilon}\right)$	–
[ 18 ]	$\widetilde{O}\!\left(\frac{SAH^{3}\log M}{\epsilon^{2}}+\frac{S^{2}AH^{4}\log M}{\epsilon}\right)$	–
This work	$\widetilde{O}\!\left(\frac{SAH^{3}\log M}{\epsilon^{2}}\right)$	No burn-in terms.

Table 1: Comparison of online episodic sample-complexity bounds for learning near-optimal policies for

M

reward functions.

p_{\ell+1}(m)\leftarrow\frac{p_{\ell}(m)\exp\bigl((\hat{v}_{\ell,m}^{+}-\hat{v}_{\ell,m}^{-})/H\bigr)}{\sum_{j=1}^{M}p_{\ell}(j)\exp\bigl((\hat{v}_{\ell,j}^{+}-\hat{v}_{\ell,j}^{-})/H\bigr)}.

Algorithm 1 Online-Multi-Reward-Learning-with-Exponential-Weight

(\mathcal{S},\mathcal{A},H,K,L,\{r^{m}\}_{m\in[M]},\epsilon,\delta)

	$\displaystyle b_{h}^{k,m}(s,a)$	$\displaystyle\leftarrow c_{1}\sqrt{\frac{\mathbb{V}(\widehat{P}_{h,s,a},\overline{V}_{h+1}^{k,m})\log\frac{1}{\delta^{\prime}}}{\max\{N_{h}(s,a),1\}}}+c_{2}\frac{H\log\frac{1}{\delta^{\prime}}}{\max\{N_{h}(s,a),1\}},$		(3)
	$\displaystyle\overline{Q}_{h}^{k,m}(s,a)$	$\displaystyle\leftarrow\min\big\{r_{h}^{m}(s,a)+\langle\widehat{P}_{h,s,a},\overline{V}_{h+1}^{k,m}\rangle+b_{h}^{k,m}(s,a),H\big\}.$		(4)

Algorithm 2

\texttt{Multi-Reward-Switching-MVP}(p)

Panel of Figure 1 — Figure 1: Empirical episode efficiency in multi-reward reinforcement learning. (a) Mean worst-reward policy error versus total episode budget at $M=16$ , averaged over 12 independent runs. (b) Estimated episode budget required to reach a mean worst-reward error of $0.1$ as $M$ varies, using 8 independent runs. Shaded regions indicate pointwise 95% confidence intervals.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

$\mathcal{S},\;S=\|\mathcal{S}\|$	State space and its size
$\mathcal{A},\;A=\|\mathcal{A}\|$	Action space and its size
$H$	Horizon length
$K$	Number of learning episodes per phase
$L$	Number of phases
$M$	Number of reward functions
$s,s^{\prime}\in\mathcal{S}$	States

Appendix

Table 2: Parameters of MDP

\underline{Q}_{h}(s,a)\leftarrow\begin{cases}\max\left\{{r_{h}(s,a)+\dfrac{N_{h}(s,a)}{N_{h}(s,a)+1}\langle P_{h,s,a},\underline{V}_{h+1}\rangle,0}\right\},&\begin{subarray}{c}(h,s,a)\in\mathcal{R}\end{subarray},\\ r_{h}(s,a),&\text{otherwise}.\end{cases}

Appendix

Algorithm 4

\texttt{Conservative-Evaluation}(P,N,\mathcal{R},\pi,r)

Panel of Figure 2 — Figure 2: The lower-bound construction. (a) A waiting state and a deterministic binary tree lead to shared absorbing states. At stage $h_{t}=t+1+\ell$ , action $a$ at leaf $s_{x}$ has transition distribution $q_{a,J_{a}^{(t,x)}}$ . (b) Each time–leaf pair $(t,x)$ has its own hidden tuple $J^{(t,x)}\in[L]^{A}$ ; different times share the same physical states. (c) The publicly known reward $r_{t,x,\nu}$ localizes the policy at the designated pair before evaluating the absorbing-state reward $\nu$ . The entry reward at stage $t$ applies only to actions $a\neq 1$ , whereas action $1$ stays at $w$ .

Markov Decision Processes

Markov decision processes model sequential choices under uncertainty, linking actions, state transitions, and rewards to planning and reinforcement learning.

See all 153 papers

Feb 18, 2026

Improved Bounds for Reward-Agnostic and Reward-Free Exploration

We study reward-free and reward-agnostic exploration in episodic finite-horizon Markov decision processes (MDPs), where an agent explores an unknown environment without observing external rewards. Reward-free exploration aims to enable

ε

-optimal policies for any reward revealed after exploration, while reward-agnostic exploration targets

ε

-optimality for rewards drawn from a small finite class. In the reward-agnostic setting, Li, Yan, Chen, and Fan achieve minimax sample complexity, but only for restrictively small accuracy parameter

ε

. We propose a new algorithm that significantly relaxes the requirement on

ε

. Our approach is novel and of technical interest by itself. Our algorithm employs an online learning procedure with carefully designed rewards to construct an exploration policy, which is used to gather data sufficient for accurate dynamics estimation and subsequent computation of an

ε

-optimal policy once the reward is revealed. Finally, we establish a tight lower bound for reward-free exploration, closing the gap between known upper and lower bounds.

Oran Ridel, Alon Cohen

Department of Engineering, Tel Aviv University, Tel Aviv, Israel. · Google Research, Tel Aviv, Israel.

RL Exploration Reward-Free RL Sample-Efficient RL

Ask PDF

May 19, 2026

Minimax Optimal Variance-Aware Regret Bounds for Multinomial Logistic MDPs

We study reinforcement learning for episodic Markov Decision Processes (MDPs) whose transitions are modelled by a multinomial logistic (MNL) model. Existing algorithms for MNL mixture MDPs yield a regret of

\smash{\tilde{O}(dH^2\sqrt{T})}

(Li et al., 2024), where

d

is the feature dimension,

H

the episode length, and

T

the number of episodes. Inspired by the logistic bandit literature (Abeille et al., 2021; Faury et al., 2022; Boudart et al., 2026), we introduce a problem-dependent constant

\barσ\_T \leq 1/2

, measuring the normalised average variance of the optimal downstream value function along the learner's trajectory. We propose an algorithm achieving a regret of

\smash{\tilde{O}(dH^2\barσ\_T\sqrt{T})}

, which recovers the existing bound in the worst case and improves upon it for structured MDPs. For instance, for KL-constrained robust MDPs,

\barσ\_T = O(H^{-1})

, reducing the horizon dependence by a factor

H

. We further establish a matching

\smash{Ω(dH^2\barσ\_T\sqrt{T})}

lower bound, proving minimax optimality (up to logarithmic factors) and fully characterising the regret complexity of MNL mixture MDPs for the first time.

Pierre Boudart, Pierre Gaillard, Alessandro Rudi

SIERRA · INRIA, École Normale Supérieure CNRS, PSL Research University Paris, France · Thoth +3

Regret Minimization in RL Markov Decision Processes Reinforcement Learning Minimax Regret

Ask PDF

Mar 24, 2026

End-to-End Efficient RL for Linear Bellman Complete MDPs with Deterministic Transitions

We study reinforcement learning (RL) with linear function approximation in Markov Decision Processes (MDPs) satisfying linear Bellman completeness -- a fundamental setting where the Bellman backup of any linear value function remains linear. While statistically tractable, prior computationally efficient algorithms are either limited to small action spaces or require strong oracle assumptions over the feature space. We provide a computationally efficient algorithm for linear Bellman complete MDPs with deterministic transitions, stochastic initial states, and stochastic rewards. For finite action spaces, our algorithm is end-to-end efficient; for large or infinite action spaces, we require only a standard argmax oracle over actions. Our algorithm learns an

\varepsilon

-optimal policy with sample and computational complexity polynomial in the horizon, feature dimension, and

1/\varepsilon

Zakaria Mhammedi, Alexander Rakhlin, Nneka Okolo

Google Research, NYC · MIT

Reinforcement Learning Linear Markov Decision Processes Policy Learning Sample-Efficient RL Function Approximation in RL

Ask PDF

Abstract

Markov Decision Processes

Explore similar work

Improved Bounds for Reward-Agnostic and Reward-Free Exploration

Minimax Optimal Variance-Aware Regret Bounds for Multinomial Logistic MDPs

End-to-End Efficient RL for Linear Bellman Complete MDPs with Deterministic Transitions