Stage-Based Q-Learning: Ref-Advantage Decomposition

• The paper presents a stage-based Q-learning algorithm that decomposes Q-values into reference and advantage components to achieve significant variance reduction and gap-dependent regret bounds.
• It employs dual estimates—with standard UCB and reference-advantage updates—to stabilize learning and enable nearly linear speedup in federated and multi-agent settings.
• The method offers rigorous guarantees including improved policy switching cost bounds and better dependence on problem-specific gaps compared to classical Q-learning approaches.

Stage-based Q-Learning with Reference-Advantage Decomposition is a family of model-free reinforcement learning algorithms designed for finite-horizon episodic Markov Decision Processes (MDPs) that exploit structured error decomposition and variance reduction. These approaches, notably UCB-Advantage and Q-EarlySettled-Advantage, partition state-action-horizon visits into exponentially growing stages and utilize reference-advantage decomposed confidence intervals to obtain near-optimal and gap-dependent regret bounds. Extension to federated and multi-agent settings has further demonstrated nearly linear speedup and log-scale communication with analogous theoretical guarantees.

1. Mathematical Foundation and Problem Setting

In these algorithms, the environment is a finite-horizon, episodic MDP $$M = (\mathcal{S}, \mathcal{A}, H, \{\mathbb{P}_h\}_{h=1}^H, \{r_h\}_{h=1}^H)$$, with $\mathcal{S}$ and $\mathcal{A}$ as state and action sets of cardinalities $S$ and $A$, and horizon $H$. Transitions are governed by $\mathbb{P}_h(\cdot|s,a)$, and rewards $r_h(s,a) \in [0,1]$ may be deterministic. For each episode $k$ and stage $h$, the policy $\mathcal{S}$0 yields $\mathcal{S}$1 (expected cumulative future reward) and $\mathcal{S}$2.

Optimal values satisfy Bellman equations:

$\mathcal{S}$3

Suboptimality gaps are denoted $\mathcal{S}$4, and the minimum nonzero gap is $\mathcal{S}$5. The maximum conditional variance of value propagation is $\mathcal{S}$6. The cumulative regret over $\mathcal{S}$7 episodes (total steps $\mathcal{S}$8) is $\mathcal{S}$9. Additionally, policy switching cost (the number of per-step policy changes) is

$\mathcal{A}$0

(Zheng et al., 2024).

2. Reference-Advantage Decomposition and Algorithmic Structure

The reference-advantage methodology decomposes $\mathcal{A}$1 as $\mathcal{A}$2, with $\mathcal{A}$3 a stage-settled baseline and $\mathcal{A}$4 the local advantage. By fixing $\mathcal{A}$5 within a stage, variance in $\mathcal{A}$6 estimates is minimized, and estimation focuses on the residual $\mathcal{A}$7.

Stage-based updates divide visits to each $\mathcal{A}$8 into stages of length $\mathcal{A}$9. Two Q-estimates are maintained:

• Standard UCB estimate: $S$0,
• Reference-Advantage estimate: $S$1,

where $S$2 and $S$3 are running means of $S$4 and $S$5, and $S$6 is a variance-based bonus:

$S$7

with $S$8 (Zheng et al., 2024).

Each stage update replaces $S$9 with $A$0. Reference values $A$1 are "settled" once total visits to $A$2 surpass a threshold $A$3, ensuring the reference is close to optimal. Q-EarlySettled-Advantage uses similar updates but fixes $A$4 once upper and lower confidence bounds converge within $A$5.

3. Error Decomposition and Variance Reduction

Regret analysis relies on decomposing the suboptimality for each $A$6:

$A$7

• The first term is controlled by the standard UCB bonus.
• The reference-advantage term is bounded by the variance of the advantage component and can be made logarithmic in $A$8 and inversely proportional to $A$9.
• The difference between Q-estimates is non-negative due to the min-update and does not increase regret.

Reference-advantage decomposition is crucial to excellent variance properties. The reference part, once settled, enjoys variance decay $H$0, while the advantage has variance at most $H$1, enabling much tighter confidence intervals than Hoeffding-only bonus schemes (Zheng et al., 2024).

4. Main Theoretical Guarantees

Gap-dependent bounds for stage-based Q-learning with reference-advantage decomposition substantially improve upon prior results:

• UCB-Advantage: For $H$2, $H$3, and $H$4,

$H$5

• Q-EarlySettled-Advantage:

$H$6

• UCB-Advantage policy-switching cost: With high probability,

$H$7

(Zheng et al., 2024).

Comparison with Hoeffding-based methods (e.g., Q-Hoeffding) shows improvement from $H$8 to $H$9 in representative regimes. For deterministic or low-variance MDPs, careful $\mathbb{P}_h(\cdot|s,a)$0 tuning can further improve dependence to $\mathbb{P}_h(\cdot|s,a)$1.

FedQ-Advantage, a federated extension, provides $\mathbb{P}_h(\cdot|s,a)$2 regret with $\mathbb{P}_h(\cdot|s,a)$3 total communication cost, reflecting nearly linear speedup in the number of agents $\mathbb{P}_h(\cdot|s,a)$4 and maintaining variance-reduction advantages (Zheng et al., 2024).

5. Multi-Agent and Zero-Sum Markov Game Extensions

Reference-advantage decomposition generalizes to multi-agent and zero-sum Markov games. In two-player zero-sum Markov games, model-free stage-based Q-learning with min-gap reference-advantage (tracking optimistic/pessimistic value estimates and setting references at the minimal observed gap) recovers the model-based $\mathbb{P}_h(\cdot|s,a)$5 sample complexity previously unattainable with purely model-free methods (Feng et al., 2023).

Key innovations for Markov games include:

• Maintaining optimistic/pessimistic $\mathbb{P}_h(\cdot|s,a)$6 and $\mathbb{P}_h(\cdot|s,a)$7 estimates.
• Updating reference pairs when current value-gap is minimal.
• Applying coarse correlated equilibrium (CCE) policies to break single-agent monotonicity, with statistical control restored via min-gap reference updates.

Variance-reduced bonuses and reference-advantage separation ensure optimal horizon dependence and sublinear regret scaling in multi-agent contexts, with certified $\mathbb{P}_h(\cdot|s,a)$8-optimal Nash equilibria.

6. Memory Efficiency and Practical Considerations

Stage-based reference-advantage methods (e.g., Q-EarlySettled-Advantage) achieve regret-optimality ($\mathbb{P}_h(\cdot|s,a)$9 up to logs) with space complexity $r_h(s,a) \in [0,1]$0 and much smaller sample-size requirements than previous memory-efficient algorithms. The early-settle rule for references allows aggressive stage-based freezing of $r_h(s,a) \in [0,1]$1 once the confidence interval converges, controlling drift and maintaining statistical guarantees.

Key properties:

• Monotonicity and optimism: Q-estimates are non-increasing and preserve UCB/LCB sandwiching.
• Reference closeness: Early-settle logic ensures $r_h(s,a) \in [0,1]$2 after burning in.
• Policy switching and rare-switches: Bounds are provided for policy changes, critical in practical deployment (Li et al., 2021).

7. Significance and Comparison to Prior Work

Stage-based Q-learning with reference-advantage decomposition represents the first fully gap-dependent logarithmic-$r_h(s,a) \in [0,1]$3 regret analysis for variance-reduced Q-learning and provides the first gap-dependent bounds on policy switching cost. It generalizes efficiently to federated settings, zero-sum Markov games, and large-scale multi-agent scenarios, offering memory efficiency, favorable communication cost, and nearly information-theoretic sample complexity.

A plausible implication is that reference-advantage decomposition will remain central to future advances in model-free RL, particularly in scalable and distributed environments where low-variance estimation and fast stage-wise convergence are essential. The separation into reference and advantage estimation is the primary driver of improved gap-dependent rates and enables adaptive exploration while maintaining statistical control (Zheng et al., 2024, Zheng et al., 2024, Feng et al., 2023, Li et al., 2021).