Density-QUCB: Shift-Aware Q-Learning

• DQUCB is a shift-aware algorithm that integrates Q-learning and UCB by using a transition density estimator to modulate the exploration bonus based on in- and out-of-distribution transitions.
• It employs a sliding window with KDE to continuously estimate transition likelihoods, ensuring low exploration during stable phases and increased exploration right after shifts.
• Theoretical and empirical results show that DQUCB mitigates the exploitation of stale policies, achieving lower cumulative regret in both tabular and deep RL settings compared to shift-unaware methods.

Density-QUCB (DQUCB) is a shift-aware Q-learning–UCB algorithm for Non-Stationary Reinforcement Learning under distribution shifts in both finite-horizon episodic and infinite-horizon discounted Markov Decision Processes. It augments Q-learning Upper Confidence Bound (QUCB) with a transition density function used as a likelihood in the optimism bonus, so that exploration is suppressed when observed transitions are in-distribution and increased when post-shift transitions appear out-of-distribution. In the formulation studied in "Q-Learning with Shift-Aware Upper Confidence Bound in Non-Stationary Reinforcement Learning," DQUCB targets the failure mode in which a policy learned before a change-point continues to exploit sub-optimal rewards after the shift; its central mechanism is to divide the UCB bonus by an estimated transition likelihood $p(s' \mid s,a;\theta)$, thereby making optimism explicitly shift-aware (Bui et al., 3 Oct 2025).

1. Problem setting and formal objectives

The paper studies two single-change-point non-stationary MDP models. In the finite-horizon episodic case, the MDP is

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$

with finite state space $\mathcal{S}$, action space $\mathcal{A}$, horizon $H \in \mathbb{N}$, step-dependent transition operator $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$, and deterministic rewards $r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$. Episodes are indexed by $k=1,\dots,K$, and at the start of each episode $k$, an adversary picks the initial state $s_1 \in \mathcal{S}$. The non-stationarity is a single unknown abrupt change-point at episode $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$0, where transitions switch from $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$1 to $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$2. Before the change, $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$3; after the change, $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$4. Rewards are fixed across time in the analysis.

For a policy $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$5, the paper defines the pre-shift and post-shift value and action-value functions as

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$6

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$7

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$8

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$9

The optimal quantities are denoted $\mathcal{S}$0, and episodic regret is

$\mathcal{S}$1

In the infinite-horizon discounted case, the MDP is

$\mathcal{S}$2

with discount factor $\mathcal{S}$3, time-homogeneous transition $\mathcal{S}$4 before shift, and a single change-point at time $\mathcal{S}$5 after which transitions follow $\mathcal{S}$6. Rewards are deterministic and satisfy $\mathcal{S}$7. Policies map histories to actions. The discounted regret over the first $\mathcal{S}$8 steps is

$\mathcal{S}$9

The discounted analysis also uses sub-optimality gaps

$\mathcal{A}$0

with $\mathcal{A}$1 and $\mathcal{A}$2 defined as the minimum strictly positive gaps.

The regularity conditions are deliberately close to standard QUCB analyses. Rewards are bounded in $\mathcal{A}$3, concentration relies on Azuma–Hoeffding martingale inequalities and standard learning-rate properties, and no explicit mixing, ergodicity, or concentrability assumptions are required beyond those implicit in Q-learning UCB analyses. In the discounted setting, strictly positive sub-optimality gaps are used to obtain logarithmic-in-$\mathcal{A}$4 dependence. The main theory addresses one change-point, while noting that multiple shifts can be handled by extending the regret decomposition across stationary segments.

2. Transition density and implicit shift detection

DQUCB equips QUCB with a transition density estimator $\mathcal{A}$5 and uses this likelihood directly inside the optimism bonus. The density is computed from a joint density model through Bayes’ rule,

$\mathcal{A}$6

The model is fit by maximum likelihood on a sliding window of the $\mathcal{A}$7 most recent transition tuples $\mathcal{A}$8. The paper instantiates $\mathcal{A}$9 with a nonparametric estimator, specifically Kernel Density Estimation (KDE), for computational simplicity. In tabular episodic settings with step-dependent dynamics, a separate parameter $H \in \mathbb{N}$0 is maintained for each step $H \in \mathbb{N}$1; in the discounted setting, a single global parameter $H \in \mathbb{N}$2 is used.

The paper interprets this quantity as an approximate density ratio,

$H \in \mathbb{N}$3

where $H \in \mathbb{N}$4 is the cumulative dataset used to fit $H \in \mathbb{N}$5 and $H \in \mathbb{N}$6 is the current sample. Before the shift, $H \in \mathbb{N}$7, so $H \in \mathbb{N}$8. At the change-point and shortly after, $H \in \mathbb{N}$9 while $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$0, so $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$1. After sufficient post-shift data are absorbed by the sliding window, the two distributions again match and $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$2.

A defining characteristic of DQUCB is that it does not perform explicit change-point testing. There is no GLR, no CUSUM, and no thresholded test statistic. Instead, shift handling is implicit: the algorithm continuously re-estimates $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$3 and scales the UCB bonus inversely by this likelihood. In-distribution samples therefore yield $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$4 close to $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$5, reducing exploration pressure, while out-of-distribution post-shift samples yield smaller $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$6 and inflate the bonus. False positives and false negatives are not explicitly controlled by a detector; robustness is attributed to the windowed density fit and to the smooth modulation of optimism rather than to hard reset logic.

This design also clarifies a common misconception. DQUCB is not a model-based RL method in the usual sense of explicitly estimating the full transition kernel. It uses density estimation only to modulate uncertainty in a model-free Q-learning–UCB update. A plausible implication is that the method occupies an intermediate position: it remains algorithmically model-free at the control level while importing a lightweight transition-likelihood signal for non-stationarity handling.

3. Q-learning updates and density-weighted optimism

The starting point is the standard tabular QUCB update. In the episodic setting, with visit counter $$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$7 and learning rate

$$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$8

QUCB uses the value backup

$$P_h : \mathcal{S}\times \mathcal{A}\to \Delta(\mathcal{S})$$9

and updates

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$0

with Hoeffding-style bonus

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$1

DQUCB changes only the bonus:

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$2

The episodic update becomes

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$3

When the environment is stable, $r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$4, so the bonus is close to that of ordinary QUCB or smaller in effective impact. Immediately after a shift, the paper argues that $r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$5 drops, approximately as $r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$6, which inflates the bonus and pushes the algorithm away from stale pre-shift Q-values.

In the discounted setting, the effective horizon for analysis is

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$7

with learning rate

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$8

The QUCB bonus is

$r_h : \mathcal{S}\times \mathcal{A}\to [0,1]$9

while the DQUCB bonus is

$k=1,\dots,K$0

The discounted update uses auxiliary optimistic estimates,

$k=1,\dots,K$1

and performs

$k=1,\dots,K$2

The procedural structure mirrors these formulas. In episodic DQUCB, the inputs are $k=1,\dots,K$3, window size $k=1,\dots,K$4, and a density model class such as KDE; initialization sets $k=1,\dots,K$5, $k=1,\dots,K$6, and a density model for each step. For each episode and step, the algorithm selects $k=1,\dots,K$7, observes $k=1,\dots,K$8, increments the visit count, computes the likelihood $k=1,\dots,K$9, sets the density-aware bonus, updates $k$0, updates $k$1, and then refits $k$2 on the most recent $k$3 transition tuples. In discounted DQUCB, the structure is analogous but uses a single density model updated online at every step.

The computational overhead comes entirely from density fitting. With KDE over a window of size $k$4, the episodic version requires $k$5 time per fit and overall $k$6 time, with memory $k$7 for Q-values plus $k$8 for recent buffers. In the discounted case, density refitting costs $k$9 per update and $s_1 \in \mathcal{S}$0 overall, with memory $s_1 \in \mathcal{S}$1 for Q-values plus $s_1 \in \mathcal{S}$2 for the buffer. Compared to QUCB, DQUCB therefore adds density buffers and model fits, but remains substantially lighter than explicit model-based transition estimation.

4. Regret guarantees and proof architecture

The main theoretical claim is that DQUCB achieves improved high-probability regret bounds over shift-unaware QUCB in both problem settings, with the gain arising from the elimination of the linear post-shift exploitation tail that can appear when stale pre-shift Q-values remain optimistic under the wrong transition law (Bui et al., 3 Oct 2025).

For the finite-horizon setting, Theorem 1 states that, for an absolute constant $s_1 \in \mathcal{S}$3, with probability at least $s_1 \in \mathcal{S}$4, choosing

$s_1 \in \mathcal{S}$5

yields

$s_1 \in \mathcal{S}$6

where $s_1 \in \mathcal{S}$7 is the estimator error of $s_1 \in \mathcal{S}$8. In the oracle limit $s_1 \in \mathcal{S}$9,

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$00

By contrast, the paper remarks that non-shift-aware QUCB with

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$01

can incur

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$02

which is asymptotically worse as $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$03 grows because of the linear $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$04 term.

For the discounted setting, Theorem 2 states that, for an absolute constant $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$05, with probability at least $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$06, choosing

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$07

gives

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$08

In the oracle limit $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$09, the same expression holds without the multiplicative $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$10. The corresponding remark for ordinary QUCB is

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$11

again asymptotically worse because of the linear post-shift term.

The proof strategy decomposes regret before and after the change-point and treats optimism relative to $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$12 and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$13 separately. Learning-rate lemmas for $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$14 control propagated error through visit mass, square-sum, and scale properties. Concentration is handled using Azuma–Hoeffding bounds on martingale difference sequences associated with empirical and true transitions. The density-aware bonus enters as

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$15

in the episodic case and

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$16

in the discounted case. The estimator error of $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$17 appears multiplicatively as the $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$18 factor.

A common misunderstanding would be to read these theorems as explicit guarantees for change-point detection accuracy. They are not. The guarantees are regret guarantees for a density-modulated optimistic learner under a single abrupt shift. The paper’s theoretical object is adaptation through bonus inflation, not certified detection delay or false-alarm control.

5. Empirical evaluation and observed behavior

The empirical study spans tabular benchmarks, a deep RL benchmark, and a real-world allocation problem, and reports that DQUCB consistently achieves lower cumulative regret than QUCB after shifts while preserving the computational efficiency of model-free RL (Bui et al., 3 Oct 2025).

In tabular experiments, the paper evaluates GridWorld and Frozen-Lake. The GridWorld environment is a $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$19 grid with a shift in transition noise $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$20 from $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$21 to $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$22 at episode $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$23, using $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$24 and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$25. Frozen-Lake is a $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$26 environment with slippery-level shifts $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$27 at $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$28, using $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$29 and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$30. The tabular baselines are QUCB, UCBMQ, and UCBVI. The reported pattern is that DQUCB consistently achieves lower cumulative regret than QUCB after shifts, confirming faster adaptation. Relative to UCBVI and UCBMQ, DQUCB trades slightly higher regret for significantly lower time and space complexity. For GridWorld, the complexity comparison states that model-based methods incur $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$31 time and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$32 space, whereas DQUCB uses $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$33 time and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$34 space.

In deep RL, the study uses CartPole-v0 with continuous state and two actions. The shift is the addition of Gaussian noise $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$35 to velocity features after episode $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$36, with $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$37. Density estimation still uses KDE with sliding window $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$38, while visit counts $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$39 are approximated with hashing in continuous state space. Baselines are random, DQN, and DQN-UCB. The reported result is that, despite instability of deep Q-learning, DQUCB shows the lowest cumulative regret, particularly post-shift. The accompanying calibration analysis reports that DQN-UCB becomes overconfident after shifts, whereas DQUCB remains better-calibrated because the bonus reacts to transition likelihood.

The real-world task is an infinite-horizon COVID-19 patient hospital allocation problem with $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$40 Texas hospitals and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$41 days. The shift occurs at day $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$42 through a change in the environment parameter $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$43, affecting occupancy dynamics. The state is $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$44, the action is $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$45 subject to $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$46, and the reward penalizes overflow. Additional baselines are UCB_RA and CUCB_RA. The reported outcome is that DQUCB exhibits lower regret, equivalently higher reward, than DQN-UCB and the bandit baselines, especially post-shift. The normalized likelihood from the density model drops at the change-point, and this drop correlates with an increased UCB bonus and with explorative reallocation that reduces overflow.

The ablation results sharpen the empirical characterization. Under varying shift intensity in Frozen-Lake, DQUCB remains superior to QUCB and UCBMQ across intensities, although extreme noise degrades all methods, including UCBVI. Under a map-change shift in GridWorld from one-room to two-rooms, DQUCB adapts faster than QUCB. Under perturbed transition estimates in a comparison to UCBVI, model-based UCBVI degrades more as transition estimation error increases, while DQUCB remains more robust because it does not explicitly fit a full transition matrix.

6. Practical considerations, limitations, and relation to prior work

The implementation described in the paper uses the same learning-rate schedule as QUCB: $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$47 in the episodic setting and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$48 in the discounted setting with effective horizon $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$49. The confidence terms are

$$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$50

and the bonus constants $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$51 and $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$52 are absolute constants, tuned in practice. The density model is KDE on approximately the last $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$53 samples, reflecting a trade-off between responsiveness and variance. In deep RL, hashing is used to approximate visit counts in continuous spaces (Bui et al., 3 Oct 2025).

The most direct limitation is density estimation accuracy. The regret bounds depend on the estimator error of $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$54, and the theory records this dependence through the multiplicative $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$55 factor. KDE may become costly or inaccurate in very high dimensions; the authors therefore suggest alternative density estimators with near-linear time and space. A second limitation is the shift model. DQUCB is tailored to abrupt, piecewise-stationary changes. Gradual drifts are still handled implicitly through the sliding window and the density-weighted bonus, but without explicit drift modeling. A third limitation concerns large state-action spaces: although DQUCB is more scalable than model-based RL, the per-step density models in the episodic setting and the global density fit in the discounted setting still add overhead relative to vanilla QUCB.

Its relation to prior work is sharply defined. Relative to QUCB and related UCB-based RL methods, DQUCB extends model-free Q-learning–UCB by making optimism shift-aware via a transition likelihood. Relative to model-based UCBVI or UCLR-style approaches, it avoids full transition estimation and the associated $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$56 costs. Relative to explicit change-point detection methods based on GLR or CUSUM, it replaces thresholding with continuous density-weighted optimism, which simplifies the mechanism but does not provide explicit false-alarm guarantees. Relative to momentum Q-learning methods such as UCBMQ, the paper positions DQUCB as simpler and more computationally efficient because UCBMQ requires bias-value function computation.

The paper also outlines several extensions. These include a model-based DQUCB that combines density-weighted bonuses with lightweight model learning such as linear mixture models; replacements for KDE in continuous state-action spaces, including normalizing flows or energy-based models and adaptation to actor-critic architectures; partial observability through densities over belief states or latent representations; and multiple change-points by decomposing the horizon into stationary segments and summing per-segment regret. These proposals are not developed into full algorithms in the paper, but they indicate the breadth of settings in which density-aware optimism could plausibly be transplanted.

At a conceptual level, DQUCB’s core contribution is narrow and specific: it modifies the UCB bonus in Q-learning by dividing by the estimated transition likelihood $$M := \langle \mathcal{S}, \mathcal{A}, H, P, r \rangle,$$57. This keeps exploration low when transitions match recent historical data and automatically increases exploration after an abrupt shift, mitigating exploitation of outdated policies. Within the scope studied, that modification is sufficient to produce improved oracle regret guarantees over shift-unaware QUCB and lower regret in tabular, deep RL, and real-world hospital-allocation experiments.