💻 computer science

Online KL-Regularized Reinforcement Learning with Function Approximation under Misspecification

This paper introduces KL-regularized formulations for contextual bandits and episodic reinforcement learning under general function approximation with model misspecification, establishing high-probability regret guarantees for regression-based algorithms that explicitly account for approximation errors.

Original authors: Haoyang Hong, Zichen Wang, Quanquan Gu, Huazheng Wang

Published 2026-06-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haoyang Hong, Zichen Wang, Quanquan Gu, Huazheng Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to play a complex video game. The goal is for the robot to learn the best moves to win. In the world of artificial intelligence, this is called Reinforcement Learning (RL).

Usually, scientists assume the robot has a "perfect map" of the game world. They assume the robot can learn a model that exactly matches reality. But in the real world, this assumption often fails. The game might be too complex, or the robot's "brain" (its mathematical model) might be too simple to capture every nuance. This is called Model Misspecification. It's like trying to describe a 3D landscape using only a 2D drawing; you'll always miss some details, no matter how hard you try.

This paper tackles a specific, modern version of this problem: teaching robots to learn while being "gentle" with their existing knowledge.

The "Gentle Nudge" (KL-Regularization)

In modern AI (like the systems that power chatbots), we don't just want the robot to learn new things; we want it to learn without forgetting its original personality or going off the rails. To do this, we use a "gentle nudge" called KL-Regularization.

Think of it like a student learning a new subject.

The Reference Policy: This is the student's original, safe way of thinking.
The New Policy: This is the student's new, optimized way of thinking after studying.
The KL Penalty: This is a rule that says, "You can learn new things, but don't stray too far from your original, safe way of thinking." If the student changes too drastically, they get "fined" (penalized). This keeps the learning stable and prevents the robot from making wild, dangerous guesses.

The Problem: The "Rough Map"

The authors ask: What happens if the robot's map is fundamentally flawed (misspecified) AND we are trying to keep it on a gentle path?

Previous theories said: "If your map is wrong, the robot will fail to learn efficiently."
This paper says: "Not necessarily. We can still prove the robot will learn well, even with a rough map, as long as we account for how rough the map is."

The Solution: The "Safety Margin"

The authors designed new algorithms (MR-KL-UCB and MR-KL-LSVI) that act like a cautious explorer with a safety margin.

The Explorer's Strategy: The robot tries to guess the best move. But because it knows its map might be slightly wrong, it adds a "safety margin" (a bonus) to its guesses.
The "Misspecification" Term: The key innovation is that this safety margin explicitly includes a term for the "roughness" of the map.
- Analogy: Imagine walking in the fog. If you know the fog is thick (high misspecification), you take smaller steps and stay closer to the path. If the fog is thin, you can walk faster. The algorithm automatically adjusts its "cautiousness" based on how bad the map is.
The Gibbs Policy: Instead of just picking the single "best" move (which might be a fluke), the robot picks moves based on a probability distribution (a "Gibbs policy"). It's like rolling a weighted die where the best moves have a higher chance of being picked, but the robot still explores other options. This randomness helps it avoid getting stuck in bad habits caused by a bad map.

The Results: "Good Enough" is Proven

The paper provides mathematical proof (regret bounds) showing that:

Even if the robot's model is imperfect, it will still learn to play the game well.
The "cost" of the imperfect model is clearly visible in the math. It shows exactly how much slower the robot learns because of the bad map.
If the map were perfect (the old, ideal scenario), the math simplifies to the standard, known results. This proves the new method is a true upgrade that covers both the perfect and imperfect worlds.

In a Nutshell

This paper is about building AI that is robust. It acknowledges that AI models are often imperfect approximations of reality. Instead of pretending the models are perfect, the authors built a system that admits, "My map is a bit blurry," and adjusts its learning strategy accordingly. It ensures that even with a blurry map and a rule to stay "gentle," the AI will still learn effectively and safely.

Key Takeaway: You don't need a perfect map to navigate; you just need a strategy that knows how to handle the fog. This paper provides that strategy for AI.

Technical Summary: Online KL-Regularized Reinforcement Learning with Function Approximation under Misspecification

1. Problem Formulation

This work addresses the problem of KL-regularized contextual bandits and episodic reinforcement learning (RL) under general function approximation when the chosen model class is misspecified.

In standard settings, theoretical guarantees often rely on the realizability assumption, which posits that the ground-truth reward or value functions lie exactly within the learner's function class. However, in practical applications such as Reinforcement Learning from Human Feedback (RLHF) and Large Language Model (LLM) post-training, learners frequently employ restricted function classes (e.g., linear or low-rank models) that cannot perfectly capture complex, nonlinear reward structures. This leads to model misspecification, where even with infinite data, the best predictor in the chosen class retains a non-zero approximation bias.

The specific objectives are:

Contextual Bandits: Minimize the cumulative KL-regret against a KL-optimal Gibbs policy, defined by a reference policy $\pi_{\text{ref}}$ and a regularization parameter $\eta$ .
Episodic RL: Minimize cumulative KL-regret in an episodic MDP setting, where policy updates are regularized by a KL penalty relative to stage-wise reference policies.

The core challenge is that existing misspecification analyses for standard reward-regret formulations do not directly transfer to KL-regularized settings due to structural differences in the optimization geometry and Bellman targets.

2. Methodology

2.1 Misspecification Formulations

The authors introduce specific misspecification conditions tailored to the KL-regularized objective:

Contextual Bandits: A pointwise misspecification formulation where the squared error between the true reward $R^\star$ and the best approximator in the class $\mathcal{R}$ is bounded by $\zeta_{\text{Bandit}}^2$ .
Episodic RL: A stagewise misspecification condition aligned with KL-regularized Bellman targets. This bounds the error between the true Bellman backup (involving the log-partition operator) and the best approximator in the stage-wise function class $\mathcal{R}_h$ , bounded by $\zeta_{\text{RL}}$ .

2.2 Algorithms

The paper proposes two regression-based algorithms that incorporate optimism to handle uncertainty and misspecification:

MR-KL-UCB (Contextual Bandits):
- Fits an Empirical Risk Minimizer (ERM) on observed data.
- Computes an uncertainty bonus based on the eluder dimension of the function class.
- Adds the known misspecification level $\zeta_{\text{Bandit}}$ to the bonus.
- Updates the policy using a Gibbs distribution derived from the optimistic reward estimate: $\pi_t(a|x) \propto \pi_{\text{ref}}(a|x) \exp(\eta \hat{R}_t(x,a))$ .
MR-KL-LSVI (Episodic RL):
- An extension of Least-Squares Value Iteration (LSVI) to the KL-regularized setting.
- Performs backward planning: at each stage $h$ , it recomputes regression labels using the continuation value from the next stage.
- Fits a stagewise ERM, computes an uncertainty bonus, and adds the misspecification level $\zeta_{\text{RL}}$ .
- Updates the policy via Gibbs policy improvement relative to the stage-wise reference policy.
- Uses a soft-value update (log-partition operator) to handle the KL regularization.

2.3 Theoretical Analysis Techniques

The analysis combines several advanced techniques:

Gibbs Variational Identity: Utilizes the exact relationship between the KL gap and the difference in scores to bound regret.
Quadratic Self-Bounding: Establishes that the KL gap between the learner's policy and the optimal policy is upper-bounded by the squared difference in their score functions (optimism gap).
Reduction to Bellman Residuals: Converts squared on-policy Q-gaps into squared Bellman residual terms, isolating KL-specific contributions.
Freedman's Inequality: Used to align predictable sums of squared bonuses with realized sums, ensuring high-probability guarantees.

3. Key Contributions

Novel Misspecification Formulations: The paper defines the first KL-aligned misspecification conditions for both contextual bandits (pointwise) and episodic RL (stagewise along Bellman paths).
Algorithmic Design: It introduces MR-KL-UCB and MR-KL-LSVI, which explicitly calibrate optimism bonuses to known misspecification levels ( $\zeta$ ) while maintaining the Gibbs policy structure required for KL-regularized objectives.
High-Probability Regret Guarantees: The authors establish high-probability KL-regret bounds that explicitly separate statistical complexity (via eluder dimension) from the approximation error due to misspecification.
- For bandits: $O(\eta \sqrt{\log(N_R T/\delta) \cdot d(\mathcal{R}, \lambda, T)} + T\zeta_{\text{Bandit}}^2)$ .
- For RL: $\tilde{O}(\eta H^2 \sqrt{d_{\text{RL}}(\lambda, T)} + \eta H^3 T \zeta_{\text{RL}}^2)$ .
Recovery of Realizable Cases: The analysis demonstrates that when the misspecification level is zero ( $\zeta = 0$ ), the bounds recover the state-of-the-art guarantees for realizable KL-regularized settings.

4. Results

The main theoretical results (Theorem 5.1 and Corollary 5.4) provide the following guarantees:

Contextual Bandits: Under the pointwise misspecification model, the cumulative KL-regret scales with the square root of the eluder dimension of the function class and linearly with the total time horizon $T$ multiplied by the squared misspecification error $\zeta_{\text{Bandit}}^2$ .
Episodic RL: Under the stagewise misspecification condition, the regret scales with the horizon $H$ , the regularization parameter $\eta$ , the aggregate eluder dimension, and the squared misspecification error $\zeta_{\text{RL}}^2$ . The bound includes an explicit term for the sum of squared bonuses, which is controlled by the misspecification level.

These results hold with high probability ( $1-\delta$ ) and assume conditionally sub-Gaussian noise.

5. Significance and Claims

The paper claims to fill a critical gap in the theoretical understanding of RLHF and preference-based learning, where model misspecification is common but often ignored in theoretical analyses.

Bridging Theory and Practice: By moving beyond the realizability assumption, the work provides a more realistic theoretical framework for modern RL applications where simple function classes (like linear models) are used to approximate complex reward landscapes.
Structural Insight: The analysis highlights that KL-regularized objectives introduce specific geometric challenges (e.g., log-partition operators and Gibbs policies) that require distinct misspecification conditions compared to standard reward-regret formulations.
Robustness: The proposed algorithms are shown to be robust to approximation errors, degrading gracefully as the misspecification level increases, rather than failing completely as some realizability-dependent bounds might suggest.

The authors note that their model assumes pointwise misspecification and light-tailed noise. They acknowledge that extending these guarantees to on-policy misspecification measures, heavy-tailed settings, and broader feedback models (like preference-based formulations beyond the current scope) remains an important direction for future work.