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

arXiv: https://arxiv.org/abs/2510.03181 · HTML: https://mlsanity.com/paper/2510.03181/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 22 of its 52 references that are in the map.
- First posted on arXiv 3 Oct 2025.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 22 of 22)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Minimax Regret Bounds for Reinforcement Learning](https://mlsanity.com/paper/1703.05449/) (2017)
- [Is Q-learning Provably Efficient?](https://mlsanity.com/paper/1807.03765/) (2018)
- [#Exploration: A Study of Count-Based Exploration for Deep Reinforcement Learning](https://mlsanity.com/paper/1611.04717/) (2016)
- [Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes](https://mlsanity.com/paper/2012.08507/) (2020)
- [Reinforcement Learning for Non-Stationary Markov Decision Processes: The Blessing of (More) Optimism](https://mlsanity.com/paper/2006.14389/) (2020)
- [Discretizing Continuous Action Space for On-Policy Optimization](https://mlsanity.com/paper/1901.10500/) (2019)
- [Algorithms for multi-armed bandit problems](https://mlsanity.com/paper/1402.6028/) (2014)
- [Q-learning with Logarithmic Regret](https://mlsanity.com/paper/2006.09118/) (2020)
- [Q-learning with UCB Exploration is Sample Efficient for Infinite-Horizon MDP](https://mlsanity.com/paper/1901.09311/) (2019)
- [UCB Momentum Q-learning: Correcting the bias without forgetting](https://mlsanity.com/paper/2103.01312/) (2021)
- [Variance-Dependent Regret Bounds for Linear Bandits and Reinforcement Learning: Adaptivity and Computational Efficiency](https://mlsanity.com/paper/2302.10371/) (2023)
- [A Sliding-Window Algorithm for Markov Decision Processes with Arbitrarily Changing Rewards and Transitions](https://mlsanity.com/paper/1805.10066/) (2018)
- [Tightening Exploration in Upper Confidence Reinforcement Learning](https://mlsanity.com/paper/2004.09656/) (2020)
- [The Optimal Approximation Factor in Density Estimation](https://mlsanity.com/paper/1902.05876/) (2019)
- [Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts](https://mlsanity.com/paper/2302.06495/) (2023)
- [Combinatorial Multi-armed Bandits for Resource Allocation](https://mlsanity.com/paper/2105.04373/) (2021)
- [Online Calibrated and Conformal Prediction Improves Bayesian Optimization](https://mlsanity.com/paper/2112.04620/) (2021)
- [Near-Optimal Density Estimation in Near-Linear Time Using Variable-Width Histograms](https://mlsanity.com/paper/1411.0169/) (2014)
- [Density-Regression: Efficient and Distance-Aware Deep Regressor for Uncertainty Estimation under Distribution Shifts](https://mlsanity.com/paper/2403.05600/) (2024)
- [Variance-Aware Linear UCB with Deep Representation for Neural Contextual Bandits](https://mlsanity.com/paper/2411.05979/) (2024)

## Built on by (most cited first, 1 of 1)

- [Efficient and Uncertainty-Aware Diffusion Framework for Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2605.30776/) (2026)

## Abstract (arXiv)

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.