Q-Learning with Shift-Aware Upper Confidence Bound in Non-Stationary Reinforcement Learning

View PDF HTML (experimental)

Abstract:We study the Non-Stationary Reinforcement Learning (RL) under distribution shifts in both finite-horizon episodic and infinite-horizon discounted Markov Decision Processes (MDPs). In the finite-horizon case, the transition functions may suddenly change at a particular episode. In the infinite-horizon setting, such changes can occur at an arbitrary time step during the agent's interaction with the environment. While the Q-learning Upper Confidence Bound algorithm (QUCB) can discover a proper policy during learning, due to the distribution shifts, this policy can exploit sub-optimal rewards after the shift happens. To address this issue, we propose Density-QUCB (DQUCB), a shift-aware Q-learning UCB algorithm, which uses a transition density function to detect distribution shifts, then leverages its likelihood to enhance the uncertainty estimation quality of Q-learning UCB, resulting in a balance between exploration and exploitation. Theoretically, we prove that our oracle DQUCB achieves a better regret guarantee than QUCB. Empirically, our DQUCB enjoys the computational efficiency of model-free RL and outperforms QUCB baselines by having a lower regret across RL tasks, as well as a COVID-19 patient hospital allocation task using a Deep-Q-learning architecture.

Comments: International Conference on Artificial Intelligence and Statistics, 2026

Subjects:

Machine Learning (cs.LG)

Cite as: arXiv:2510.03181 [cs.LG]

(or arXiv:2510.03181v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2510.03181

arXiv-issued DOI via DataCite

Submission history

From: Ha Manh Bui [view email] [v1] Fri, 3 Oct 2025 16:56:47 UTC (1,107 KB) [v2] Sun, 29 Mar 2026 19:46:06 UTC (1,113 KB)