Towards Formalizing Reinforcement Learning Theory (2511.03618v1)

Abstract: In this paper, we formalize the almost sure convergence of $Q$-learning and linear temporal difference (TD) learning with Markovian samples using the Lean 4 theorem prover based on the Mathlib library. $Q$-learning and linear TD are among the earliest and most influential reinforcement learning (RL) algorithms. The investigation of their convergence properties is not only a major research topic during the early development of the RL field but also receives increasing attention nowadays. This paper formally verifies their almost sure convergence in a unified framework based on the Robbins-Siegmund theorem. The framework developed in this work can be easily extended to convergence rates and other modes of convergence. This work thus makes an important step towards fully formalizing convergent RL results. The code is available at https://github.com/ShangtongZhang/rl-theory-in-lean.

• The paper introduces a formal Lean 4 framework that rigorously proves almost sure convergence for Q-learning and linear TD algorithms.
• It employs modern techniques including the Robbins-Siegmund theorem and Lyapunov functions to overcome limitations of traditional ODE-based proofs.
• The framework rigorously constructs probability spaces for Markovian trajectories and is extensible to other reinforcement learning methods.

Formalizing the Almost Sure Convergence of $Q$-Learning and Linear TD in Reinforcement Learning

Introduction

This paper presents a formalization of the almost sure convergence of two foundational reinforcement learning (RL) algorithms—$Q$-learning and linear temporal difference (TD) learning—under Markovian sampling, using the Lean 4 theorem prover and the Mathlib library. The work addresses a longstanding gap in RL theory: while convergence proofs for these algorithms are well-established in the literature, they often rely on intricate, error-prone arguments, particularly those based on the ODE method, and typically omit full measure-theoretic rigor. This formalization provides a unified, extensible framework for verifying convergence properties, leveraging modern proof techniques centered on the Robbins-Siegmund theorem and Lyapunov functions, and sets a new standard for the robustness of RL theory.

Motivation and Context

The convergence of $Q$-learning and linear TD has been a central topic in RL theory, with numerous works addressing both finite and infinite state-action spaces. However, traditional proofs are delicate for two primary reasons:

1. ODE-Based Proof Fragility: The canonical approach uses stochastic approximation via ODE methods, which are detail-heavy and susceptible to subtle errors, as evidenced by errata in peer-reviewed works and even textbooks.
2. Measure-Theoretic Gaps: RL theory is typically developed in the MDP framework, requiring rigorous construction of probability spaces for infinite trajectories (via the Ionescu-Tulcea theorem), and careful handling of measurability and integrability. These aspects are often glossed over, especially for infinite state spaces.

Prior formalizations in proof assistants (Coq, Isabelle/HOL) have focused on dynamic programming or stochastic approximation for outdated algorithm variants, without addressing the modern forms of $Q$-learning and linear TD with Markovian noise. This work fills that gap, providing the first formal, machine-checked convergence proofs for these algorithms in their standard forms.

Formalization Framework

Mathematical Setting

The formalization considers an infinite-horizon MDP with finite state and action spaces, reward and transition functions, and a discount factor $\gamma \in [0,1)$. The two algorithms are:

• Linear TD: For policy evaluation, with updates

$$w_{t+1} = w_t + \alpha_t (R_{t+1} + \gamma x(S_{t+1})^\top w_t - x(S_t)^\top w_t) x(S_t)$$

where $x(\cdot)$ is a feature map and $w_t$ the parameter vector.

• $Q$-Learning: For control, with updates

$$q_{t+1}(s, a) = q_t(s, a) + \alpha_t (R_{t+1} + \gamma \max_{a'} q_t(S_{t+1}, a') - q_t(S_t, A_t)) \mathbb{I}_{(s, a) = (S_t, A_t)}$$

The goal is to formally prove that, under standard assumptions (irreducibility, aperiodicity, appropriate step sizes), the iterates converge almost surely to the correct fixed points ($w_*$ for TD, $q_*$ for $Q$-learning).

Probability Space Construction

A key technical contribution is the explicit construction of the probability space for sample paths using the Ionescu-Tulcea theorem, as formalized in Mathlib. This enables rigorous measure-theoretic treatment of Markovian trajectories, conditional expectations, and filtrations, which are essential for the convergence analysis.

Unified Proof Strategy

The formalization eschews the ODE-based approach in favor of a modern method based on the Robbins-Siegmund theorem, Lyapunov functions, and the skeleton iterates technique. The main steps are:

1. Rewriting Updates: Both algorithms are cast in the form

$w_{t+1} = w_t + \alpha_t (F(w_t, Y_{t+1}) - w_t)$

where $Y_{t+1}$ encodes the relevant Markovian noise.

1. Lyapunov Analysis: Existence of a suitable Lyapunov function $\phi$ is established, satisfying smoothness, norm equivalence, and contraction properties.
2. Noise Decomposition: The update is decomposed into a main term and two noise terms: a leading martingale difference sequence and a smaller bias term.
3. Robbins-Siegmund Application: By verifying that the Lyapunov function along the iterates forms an almost supermartingale, the Robbins-Siegmund theorem is invoked to conclude almost sure convergence.
4. Skeleton Iterates for Markovian Noise: The skeleton iterates technique is used to handle the dependence structure in Markovian samples, reducing the problem to one amenable to the martingale difference analysis.

Formal Theorem Statements

The main formalized results are:

• Linear TD (Markovian samples): Under irreducibility, aperiodicity, full column rank of features, and step size $\alpha_t = (t+2)^{-\nu}$ with $\nu \in (2/3, 1)$, the iterates converge almost surely to the TD fixed point.
• Linear TD (i.i.d. samples): Under the Robbins-Monro step size condition, convergence holds for i.i.d. samples.
• $Q$-Learning (Markovian samples): Under analogous conditions, the $Q$-learning iterates converge almost surely to the optimal action-value function.

The restriction $\nu \in (2/3, 1)$ arises from the proof technique; extending to $\nu \in (1/2, 2/3]$ would require formalizing more advanced ODE-based arguments.

Implementation in Lean 4

The formalization is implemented in approximately 10,000 lines of Lean 4 code, leveraging Mathlib for measure theory, probability, and linear algebra. Key implementation aspects include:

• Stochastic Matrices and Markov Chains: Classes for stochastic vectors, row-stochastic matrices, irreducibility, aperiodicity, and Doeblin minorization.
• Probability Measures on Trajectories: Construction of the sample path measure via the Ionescu-Tulcea theorem.
• Iterate Definitions: Encodings of the TD and $Q$-learning updates as recursive functions over sample paths.
• Conditional Expectation: Formalization of conditional expectations in the context of Markov chains, requiring substantial code to bridge the gap between abstract measure-theoretic definitions and concrete matrix computations.
• Supermartingale Arguments: Formalization of the Robbins-Siegmund theorem and its application to the Lyapunov function along the iterates.

The codebase is publicly available and can serve as a high-quality dataset for benchmarking LLMs on formal reasoning and code synthesis in mathematics and machine learning.

Numerical and Theoretical Implications

The formalization confirms, with machine-checked rigor, the almost sure convergence of $Q$-learning and linear TD under standard conditions. The framework is immediately extensible to:

• Convergence Rates: By telescoping the supermartingale inequality, $\mathcal{L}_2$ convergence rates can be obtained for i.i.d. samples; for Markovian samples, established techniques can be formalized.
• Other Modes of Convergence: The approach can be extended to high-probability concentration, $\mathcal{L}_p$ convergence, and exponential tail bounds, contingent on formalizing auxiliary results such as Hoeffding's lemma.
• Other Algorithms: The framework is adaptable to off-policy TD methods, gradient TD methods (without eligibility traces), and, with further development, to more complex algorithms involving time-inhomogeneous Markov chains (e.g., SARSA, policy gradient methods).

A notable technical challenge is the formalization of conditional expectations in the context of Markov chains, which required significant effort and code. This highlights the complexity of bridging intuitive probabilistic reasoning and formal, machine-checked proofs.

Implications for AI and Formal Methods

This work demonstrates the feasibility and value of fully formalizing nontrivial results in RL theory, setting a precedent for future efforts in the formal verification of machine learning algorithms. The formalization provides a robust foundation for further theoretical developments and can serve as a benchmark for evaluating the capabilities of LLMs and automated theorem provers in mathematical reasoning.

The project also provides empirical evidence regarding the current limitations of LLMs in formal mathematics: while LLMs (e.g., Gemini, ChatGPT) are valuable as tutors, search engines, and for small lemma synthesis, they are not yet capable of independently completing such a formalization. This underscores the need for continued research at the intersection of AI and formal methods.

Conclusion

The formalization of the almost sure convergence of $Q$-learning and linear TD in Lean 4 represents a significant advance in the rigor and reliability of RL theory. The developed framework is extensible to a broad class of RL algorithms and convergence properties, and the codebase provides a valuable resource for both the formal methods and machine learning communities. Future work will address more general step size regimes, more complex algorithms, and further integration with automated reasoning tools.