Finite-Time Convergence of Distributionally Robust Q-Learning with Linear Function Approximation

Technical Summary: Finite-Time Convergence of Distributionally Robust Q-Learning with Linear Function Approximation

Problem Statement
The paper addresses the challenge of learning optimal policies in Distributionally Robust Reinforcement Learning (DRRL) when the deployment environment differs from the training (nominal) environment. Specifically, the authors focus on model-free, discounted robust Q-learning under an $(s, a)$-rectangular chi-square uncertainty set.

The setting involves a learner operating on a single Markovian trajectory generated by an unknown nominal transition kernel $P_0$ and a behavior policy $\pi_b$. The learner lacks generative access (cannot sample arbitrary state-action pairs) and does not know $P_0$. The goal is to learn the optimal robust Q-function, $Q^\star_\chi$, which maximizes the worst-case discounted value over the uncertainty set.

The problem presents two coupled difficulties:

1. Function Approximation Difficulty: While the exact robust Bellman operator is a contraction in $\ell_\infty$, the projected robust Bellman operator induced by linear function approximation is not known to be a contraction for general discount factors $\gamma \in (0, 1)$.
2. Robust Estimation Difficulty: The chi-square robust Bellman update involves an inner minimization over transition distributions. Its dual gradient depends on conditional moments (specifically, the positive part of the Bellman gap and its square) that must be estimated from a single trajectory. Standard analyses for tabular robust Q-learning or linear Q-learning do not directly extend to this setting.

Methodology
The authors propose Algorithm 1, a multi-stage stochastic approximation algorithm that combines a target-network outer loop with a dual function-approximation scheme. The algorithm operates in outer blocks where the target Q-function is frozen. Within each block, three stages occur:

1. Stage 1: Dual Optimization with Moment Tracking.
   The algorithm solves the inner minimization of the chi-square robust Bellman update via its scalar dual form. Instead of solving a separate dual problem for every state-action pair, the dual variable $\eta$ is approximated linearly ($\eta_\nu(i) = \psi(i)^\top \nu$).

   • Two-Time-Scale Scheme: The dual variable $\nu$ is updated on a slow time scale, while two "moment-tracking critics" ($\omega_1, \omega_2$) estimate the required conditional moments ($Z_1$ and $Z_2$) on a faster time scale.
   • Smoothing: To handle the lack of global smoothness in the chi-square dual objective (which causes issues with Markovian noise analysis), the authors optimize a smoothed surrogate objective using a tunable parameter $\tau > 0$. This ensures the dual gradient is Lipschitz continuous.
   • Suffix Averaging: The dual iterates are suffix-averaged to obtain a stable estimate $\bar{\nu}_t$.

2. Stage 2: Fresh Evaluation.
   Before updating the Q-function, the algorithm performs a "fresh evaluation" of the variance-like moment ($Z_1$) using the averaged dual parameter $\bar{\nu}_t$. This stage isolates the estimation of the moment from the optimization trajectory, preventing bias accumulation that would occur if the Stage 1 critics (learned while $\nu$ was moving) were reused directly.

3. Stage 3: Q-Update.
   The algorithm performs a slow temporal difference update for the robust Q-function parameter $\theta$ against the fixed, evaluated target derived from Stages 1 and 2.

Key Contributions

1. Finite-Time Convergence Guarantee: The paper provides the first finite-time convergence bound for discounted model-free robust Q-learning with linear function approximation from single-trajectory data, without imposing small-discount-factor restrictions.
2. Target-Network Analysis for Robust RL: The authors adapt the target-network technique (common in non-robust deep RL) to the robust setting. By freezing the target Q-function within outer blocks, they bypass the non-contractivity of the projected robust Bellman operator, allowing the analysis to rely on the contraction of the unprojected operator across blocks.
3. Scalable Dual Estimation: They introduce a method to approximate the dual variable itself using a linear function class, coupled with moment-tracking critics. This avoids the scalability issues of solving independent dual problems for large state spaces.
4. Smoothing and Bias Control: The paper introduces a smoothed surrogate for the chi-square dual objective. This resolves the non-Lipschitz gradient issue inherent in the exact chi-square formulation, enabling rigorous finite-time stochastic approximation bounds under Markovian noise. The authors explicitly note that previous proofs for tabular robust Q-learning with chi-square uncertainty relied on global smoothness assumptions that do not hold; their proof corrects this gap.

Results
The main theoretical result (Theorem 1) establishes an end-to-end convergence bound of the form:
$$\mathbb{E}[\| \hat{Q}_T - Q^\star_\chi \|_\infty] \leq \gamma^T \| \hat{Q}_0 - Q^\star_\chi \|_\infty + \frac{E_{stat}(N, M, K, \tau) + E_{app}(\tau)}{1 - \gamma}$$
where:

• The first term represents the target-network contraction error.
• $E_{stat}$ captures finite-sample errors from the Q-update, dual optimization, smoothing bias, and fresh evaluation.
• $E_{app}$ captures approximation errors due to function approximation and smoothing.

Corollary 2 derives the sample complexity, showing that to achieve an error $\epsilon$, the total trajectory length scales as $\tilde{O}(\epsilon^{-5/(1-a)})$ (where $a$ is a step-size parameter), up to logarithmic factors.

Experiments on the FrozenLake MDP validate the theory, demonstrating convergence toward the robust reference value and illustrating the trade-off introduced by the smoothing parameter $\tau$ (smaller $\tau$ reduces bias but increases variance/optimization error).

Significance and Claims
The paper claims to close a significant gap between the empirical success of robust RL algorithms and the theoretical guarantees available for their non-robust counterparts. Prior work on model-free robust RL was limited to tabular settings, required generative access, assumed active exploration, or imposed restrictive conditions on the discount factor.

The authors state that their work provides the first finite-time convergence guarantee for discounted model-free robust Q-learning with linear function approximation from single-trajectory data without small-discount-factor assumptions. Furthermore, they assert that their proof for the chi-square case (even in the tabular limit) appears to be the first correct proof of convergence, as it addresses the lack of global smoothness in the dual objective which previous analyses had overlooked or assumed away. The framework is presented as an analyzable template that can be extended to large-scale neural network implementations.