Source-linked AI summary
Fast Rates for Offline Q-Learning with KL Regularization
Yize Hao, Guangxin Jiang, Shuang Qiu
TL;DR
Offline Q-learning’s ability to achieve fast rates with KL regularization was unresolved. The paper proposes KLQ and pessimistic KLPQ, reporting eO(1/ϵ) sample complexity and removal of KLPQ’s explicit exponential dependence on regularization and horizon.
Problem
Whether incremental offline Q-learning can achieve fast rates for a KL-regularized objective remains unclear, despite fast-rate results in other settings.
Source: §1 · p. 2Method
The paper develops KLQ, which uses a regularized Bellman update, and KLPQ, which incorporates pessimism into KLQ.
Source: §1 · pp. 2–3, §3 · p. 6Results
KLQ and KLPQ achieve eO(1/ϵ) sample complexity, with KLPQ removing explicit exponential dependence on the regularization parameter and horizon.
Source: §6 · p. 11Takeaways & Limitations
These results establish fast-rate guarantees for offline Q-learning in tabular finite-horizon KL-regularized MDPs and show a sample-complexity benefit from pessimism.
Source: §6 · p. 11KLQ’s exponential factor can make its guarantee less informative as the horizon grows; choosing η to control it also changes the learning objective.
Source: §4 · p. 10
Abstract
from arXiv · showhide
This paper investigates offline Q-learning for finite-horizon Markov decision processes with Kullback-Leibler (KL) regularization. Recent studies have shown that KL regularization can yield fast rates in offline decision-making problems. However, it remains unclear whether Q-learning, which adopts incremental stochastic approximation updates, can achieve such fast rates in the offline setting. We first propose KL-regularized Q-Learning (KLQ), a simple variant of classical Q-learning. By exploiting KL regularization, we prove that KLQ achieves a sample complexity of $\widetilde{O}(1/\epsilon)$ without explicit pessimism. This guarantee, however, has an exponential dependence on the KL regularization parameter and the horizon. To address this limitation, we further propose KL-regularized Pessimistic Q-Learning (KLPQ), which incorporates pessimism into KLQ. We prove that KLPQ retains the same \(\widetilde O(1/\epsilon)\) rate while removing the explicit exponential dependence, thereby further improving the sample complexity. To the best of our knowledge, these results establish the first fast-rate guarantees for offline Q-learning in tabular finite-horizon KL-regularized MDPs.
1 Introduction
The paper asks whether incremental offline Q-learning can achieve fast rates under KL regularization, and introduces KLQ and pessimistic KLPQ to address that gap. It reports a first eO(1/ϵ) guarantee without pessimism and a KLPQ variant that removes explicit exponential dependence on the regularization parameter and horizon.
1 Introduction
- Prior KL-regularized fast-rate analyses use backward induction, leaving analogous guarantees for incremental offline Q-learning unresolved.Source: §1 · p. 2
- The cited offline Q-learning guarantees concern the unregularized objective and rely on pessimism.Source: §1 · p. 2
- KLQ establishes the first eO(1/ϵ) sample-complexity guarantee for offline finite-horizon tabular Q-learning without pessimism.Source: §1 · p. 2
- KLPQ adds pessimism to KLQ, removes its explicit exponential dependence on regularization and horizon, and retains the fast rate.Source: §1 · p. 3
- KL regularization supports the analysis of incremental Q-learning through a squared Q-estimation-error recursion, which pessimism further sharpens.Source: §1 · p. 3
1.2 Notation
- The paper uses eO(f) to suppress polylogarithmic factors in f, problem parameters, and inverse failure probability.Source: §1.2 · p. 4
2 Problem Setup
The setup is a finite-horizon episodic MDP with a KL-regularized policy objective relative to a reference policy. Learning uses fixed offline trajectories and assumes coverage of the KL-regularized optimal policy by the behavior policy.
2.1 Finite-Horizon KL-Regularized MDPs
- The model is a finite-horizon episodic MDP with finite state and action sets, horizon H, and rewards in [0, 1].Source: §2.1 · p. 4
- The policy value is regularized by KL divergence from a reference policy, with η determining the regularization strength.Source: §2.1 · p. 4
2.2 Offline Learning and Data Coverage
- Because the learner cannot interact further with the environment, finite-sample guarantees require data coverage of state-action pairs relevant to the target policy.Source: §2.2 · p. 5
- Single-policy concentrability requires the behavior policy to cover the occupancy of the KL-regularized optimal policy, without requiring it to equal the reference policy.Source: §2.2 · p. 6
- The learning objective is to maximize expected KL-regularized return from the initial-state distribution, and performance is measured by policy suboptimality.Source: §2.2 · p. 6
3 Pessimism-Free KL-Regularized Q-Learning
KLQ keeps Q-learning’s incremental update form while replacing the hard maximum with KL-regularized value and policy maps. The section gives a finite-sample guarantee, while noting that its bound has an exponential factor in the regularization parameter and horizon.
3 Pessimism-Free KL-Regularized Q-Learning
- KLQ replaces the hard maximum with a regularized Bellman value operator and derives the associated policy from the Q-estimate.It retains the stochastic-approximation form of Q-learning and updates only the visited state-action pair on each offline trajectory.Source: §3 · p. 6
- KLQ uses the classical Q-learning update without an explicit pessimistic bonus or monotone max update.The value and policy are instead determined by the KL-regularized Bellman equations.Source: §3 · p. 7
- Theorem 3.1 gives a finite-sample guarantee and an ϵ-optimal policy sample-complexity bound for KLQ.The supplied theorem passage does not include the bound’s full expression.Source: §3 · p. 7
- The KLQ bound contains e^(ηH(H+1)), which can become large as the horizon grows for fixed η > 0.The paper attributes this dependence to likelihood-ratio comparisons that contribute to stagewise error propagation and converting Q-estimation error to policy suboptimality.Source: §3 · p. 7
4 KL-Regularized Pessimistic Q-Learning
KLPQ adds pessimism to KLQ’s regularized updates to improve its dependence on the regularization parameter and horizon. Its guarantee removes the explicit exponential factor and improves the polynomial dependence on the horizon by one factor.
4 KL-Regularized Pessimistic Q-Learning
- KLPQ retains KLQ’s regularized value and policy updates while adding a pessimistic bonus and a monotone Q-estimate update.The auxiliary iterate uses the bonus-adjusted Q-learning update; the monotone estimate is used to compute the regularized value.Source: §4 · p. 8
- Theorem 4.1 gives a finite-sample guarantee for the policy returned by KLPQ and an ϵ-optimal policy sample-complexity bound.The supplied theorem passage does not include the bound’s full expression.Source: §4 · p. 9
- KLPQ removes the explicit e^(ηH(H+1)) factor and improves the polynomial dependence on H by one factor, for fixed η > 0.The improvement does not change the KL-regularized objective or the concentrability coefficient C⋆_η.Source: §4 · p. 9
- KLPQ’s KL-regularized guarantee is not directly comparable with the unregularized offline Q-learning bound, whose leading accuracy dependence is ϵ^-2.The paper notes that the objectives and corresponding coverage requirements differ; it reports KLPQ’s accuracy dependence as ϵ^-1 for fixed η.Source: §4 · p. 9
5 Theoretical Analysis
The analysis traces how statistical error propagates in KLQ and shows how KLPQ’s pessimistic construction changes that recursion. It also gives KLPQ a sharper horizon dependence through direct averaging.
5 Theoretical Analysis
- The proof sketches suppress polynomial dependencies on H, S, A, and C*_η and logarithmic factors, while retaining explicit dependence on K and η.The full parameter dependencies are given in the main theorems and detailed proofs.Source: §5 · p. 10
5.1 Analysis for KLQ
- KLQ’s recursion combines propagated next-stage error with accumulated statistical error, while converting the cumulative bound to final-iterate error requires an additional argument.The analysis attributes KLQ’s exponential factor to stagewise propagation and conversion to policy suboptimality.Source: §5.1 · p. 10
5.2 Analysis for KLPQ
- KLPQ’s pointwise pessimism avoids an extra factor in the squared-error recursion and permits direct control of policy suboptimality without comparing occupancy measures.The analysis contrasts this with KLQ’s likelihood-ratio and occupancy-measure comparison.Source: §5.2 · p. 11
- Backward iteration uses the polylogarithmic statistical-error dependence on K; KLPQ’s pessimism and monotone update support its error bound.The KLPQ argument also uses direct averaging, saving one factor of H relative to the final-iterate argument.Source: §5.2 · p. 11
6 Conclusion
The paper studies offline Q-learning in finite-horizon KL-regularized MDPs and proposes KLQ and KLPQ. Both obtain an eO(1/ϵ) sample-complexity rate, with KLPQ removing the explicit exponential dependence on the regularization parameter and horizon.
6 Conclusion
- KLQ achieves eO(1/ϵ) sample complexity without explicit pessimism for offline Q-learning in finite-horizon KL-regularized MDPs.The conclusion identifies this as the paper’s initial algorithmic guarantee.Source: §6 · p. 11
- KLPQ retains the eO(1/ϵ) rate while removing the explicit exponential dependence on the regularization parameter and horizon.The conclusion says this improves sample complexity relative to KLQ.Source: §6 · p. 11
- The results establish fast-rate guarantees for offline Q-learning with KL regularization and demonstrate a benefit of pessimism in improving sample complexity.Source: §6 · p. 11
A Proof of Proposition 2.1
The proof defines value, Q, and policy quantities recursively from the end of the horizon, then uses the Gibbs variational identity to verify the induced policy’s value.
A Proof of Proposition 2.1
- The proof sets the terminal value to zero and recursively defines Q-values, values through F, and policies through G.Source: A Proof of Proposition 2.1 · p. 13
- The induction uses the Gibbs variational identity, with the policy defined by G as the maximizer in that identity.Source: A Proof of Proposition 2.1 · p. 13
- F_h,s(q) lies between the minimum and maximum coordinates of q, which the proof uses to establish the value range.Source: A Proof of Proposition 2.1 · p. 13
B Supporting Lemmas
The supporting lemmas provide concentration and visitation-count tools used in the statistical analysis, including bounds for predictable value functions and martingale differences.
B Supporting Lemmas
- Freedman’s martingale inequality supplies a concentration bound for martingale differences under a bounded conditional-variance assumption.Source: B Supporting Lemmas · p. 14
- A concentration lemma bounds errors for predictable, bounded value functions simultaneously across stages, state-action pairs, and visit counts.Source: B Supporting Lemmas · p. 14
- The visitation-count lemmas use binomial counts and Chernoff bounds to control when behavior-occupied state-action pairs are observed.They bound the number of episodes with zero count for a pair with positive behavior occupancy.Source: B Supporting Lemmas · pp. 15–16
C Proofs for Algorithm 1
- The appendix absorbs constant factors in logarithmic confidence terms into universal constants.Source: C Proofs for Algorithm 1 · p. 16
C.1 Lemmas for Theorem 3.1
These lemmas establish bounded iterates and error controls used to prove Theorem 3.1, then connect Gibbs-policy bounds to output-policy suboptimality.
C.1 Lemmas for Theorem 3.1
- The iterates satisfy 0 ≤ Q^k_h(s,a) ≤ L_h and 0 ≤ V^k_h(s) ≤ L_h, an invariant maintained by convex-combination updates.Source: §C.1 · p. 17
- The proof converts a pointwise KL bound between Gibbs policies into output-policy suboptimality using Gibbs and Bellman identities, with stage terms telescoping.Source: §C.1 · p. 18
- A pointwise squared Q-estimation-error recursion is derived to control both cumulative squared error and the final iterate’s error.Source: §C.1 · p. 19
- Lemma C.4’s simultaneous bounds are established under Assumption 2.2 with probability at least 1−δ.Source: §C.1 · p. 19
C.2 Proof of Theorem 3.1
- Theorem 3.1’s proof combines Lemmas C.4 and C.5 on an event of probability at least 1−δ to bound output-policy suboptimality.Source: §C.2 · p. 24
D Proofs for Algorithm 2
This appendix fixes the confidence level and notation used for the proof of Algorithm 2.
D Proofs for Algorithm 2
- The proof fixes the confidence level δ used in Algorithm 2.Source: D Proofs for Algorithm 2 · p. 24
- Iterate-dependent quantities refer to the single KLPQ run at that fixed confidence level.Source: D Proofs for Algorithm 2 · p. 24
- The appendix retains the definitions of the estimation error, filtration, and transition-probability shorthand from earlier sections.Source: D Proofs for Algorithm 2 · p. 24
D.1 Lemmas for Theorem 4.1
The lemmas for Theorem 4.1 establish bounded, ordered iterates and use them to control estimation error and the output policy.
D.1 Lemmas for Theorem 4.1
- KLPQ’s Q- and value iterates remain within their deterministic stagewise ranges, with coordinatewise monotonicity preserved across updates.Source: §D.1 · pp. 24–25
- On event E1, Lemma D.2 establishes the stated value ordering simultaneously over episodes, stages, and state-action pairs.Source: §D.1 · pp. 25–26
- Lemma D.3 bounds value error and the output policy using the established ordering of the Q-iterates.Source: §D.1 · pp. 26–28
- Lemma D.4 derives a squared Q-error recursion by subtracting the iterated update from the optimal Bellman equation and applying Jensen’s inequality.Source: §D.1 · pp. 28–29
D.2 Proof of Theorem 4.1
Theorem 4.1 follows from the KLPQ error bounds; the proof’s propagated error term avoids the factor present in the earlier argument.
D.2 Proof of Theorem 4.1
- Theorem 4.1 follows by repeating the earlier bound calculation with KLPQ’s propagated error term, which has no factor γ_h+1.Source: §D.2 · p. 30
- The proof uses backward iteration without products of γ_j to obtain the accumulated squared Q-estimation-error bound.Source: §D.2 · p. 30