TD(0) & Q-Learning, as Robbins-Monro

Motivation

When I started studying TD prediction and Q-learning, I was surprised by how much they reminded me of dynamic programming. More concretely, TD(0) prediction looked very close to the Bellman policy operator used for policy evaluation, and Q-learning looked very close to the Bellman optimality operator for state-action value functions. The similarity is hard to miss, and for anyone with basic familiarity with SGD and stochastic approximation, Robbins-Monro is probably the first intuitive guess for the missing connection. I do not think this connection is especially difficult to make, but Sutton and Barto's book, which I was studying at the time, does not go into much detail about it. So I thought it would be useful to write a post covering it. After all, the best way to learn something is to explain it.

I worked through most of these derivations with pen and paper. After that, I took photos of the pages and used ChatGPT Pro/Codex as a kind of proof-review partner and for helping me build the HTML version of my derivations.

What We Want to Show

We want to show that TD(0) prediction and Q-learning are essentially Robbins-Monro procedures that try to find roots of two corresponding Bellman residuals: the Bellman policy residual and the Bellman optimality residual. In the TD(0) prediction setting, this means finding the root of $T_\pi v - v = 0$, whose solution is the policy value function $v_\pi$. In Q-learning, the analogous root-finding problem is $T_*Q - Q = 0$, whose solution is the optimal action-value function $Q_*$.

Terminology and Notation

In this post, we consider a finite discounted Markov decision process with bounded rewards and discount factor $0 \le \gamma < 1$. I will also make two Markov assumptions explicit. First, the environment satisfies the Markov property:

$$ p(S_{t+1}, R_{t+1}\mid S_t,A_t,S_{t-1},A_{t-1},\ldots) = p(S_{t+1}, R_{t+1}\mid S_t,A_t). $$

Second, the policy is Markov:

$$ \Pr(A_t=a\mid S_t,S_{t-1},A_{t-1},\ldots) = \pi(a\mid S_t). $$

We will also use $\mathcal F_t$ to denote the filtration generated by the history up to time $t$. One fact we will frequently refer to is that any random variable that is measurable with respect to filtration at time $t$ has the property $\mathbb E[X\mid\mathcal F_t]=X$. The easiest intuitive way of thinking about this is that being measurable with respect to $\mathcal F_t$ means that the quantity is already knowable from the history available at time $t$. There is no future randomness left in it, which is why such terms can be treated as fixed and pulled out of conditional expectations conditioned on $\mathcal F_t$.

Bellman Policy Operator

For a fixed Markov policy $\pi(a\mid s)$, the policy induces a Markov reward process over states. The value function of this policy is

$$ v_\pi(s) = \mathbb E_\pi[G_t\mid S_t=s], \qquad G_t=\sum_{k=0}^{\infty}\gamma^k R_{t+k+1}. $$

The Bellman policy operator maps a candidate value function $V$ to a new value function:

$$ (T_\pi V)(s) = \sum_a \pi(a\mid s) \sum_{r,s'}p(r,s'\mid s,a) \left[r+\gamma V(s')\right]. $$

In dynamic-programming policy evaluation, the Bellman policy operator defines the iteration $v_{k+1} = T_\pi v_k$. We will not prove this here, but the fixed point of the Bellman policy operator, meaning $(T_\pi V)(s) - V(s) = 0$ for every state $s$, is the value function of the policy.

Bellman Optimality Operator

For state-action value functions, the Bellman optimality operator maps a candidate action-value function $Q$ to a new action-value function:

$$ (T_*Q)(s,a) = \sum_{r,s'}p(r,s'\mid s,a) \left[ r+\gamma \max_{a'} Q(s',a') \right]. $$

Again, we will not prove this here, but the fixed point of the Bellman optimality operator, meaning $(T_*Q)(s,a) - Q(s,a) = 0$ for every state-action pair $(s,a)$, is the optimal state-action value function.

Robbins-Monro Stochastic Approximation

Robbins-Monro is a general way to solve a root-finding problem when the true function cannot be observed directly. Suppose we want to find $x^\star$ such that

$$ h(x^\star)=0. $$

Instead of observing $h(x_t)$ directly, we observe a random quantity whose conditional expectation is the root function:

$$ \mathbb E[H(x_t,\xi_{t+1})\mid\mathcal F_t]=h(x_t). $$

The update is

$$ x_{t+1}=x_t+\alpha_t H(x_t,\xi_{t+1}). $$

In this post, the root functions will be Bellman residuals:

$$ h(V)=T_\pi V - V, \qquad h(Q)=T_*Q - Q. $$

TD(0) Prediction as Robbins-Monro, and Why It Converges to the Policy Value Function

We start from a quick refresher on TD(0). For a fixed policy $\pi$, it maintains a tabular representation of the state value function (again, we are only working with the tabular case), which in essence is just a vector of floats with dimensionality equal to the size of the state space. After observing a transition $(S_t,A_t,R_{t+1},S_{t+1})$, it forms the TD error

$$ \delta_t = R_{t+1}+\gamma V_t(S_{t+1})-V_t(S_t). $$

which is then used to update the state that was visited:

$$ V_{t+1}(S_t)=V_t(S_t)+\alpha_t\delta_t, $$

while leaving all other states unchanged:

$$ V_{t+1}(s)=V_t(s), \qquad s\ne S_t. $$

This update is applied online as a trajectory is generated, so TD(0) may update the same state's estimate many times. Intuitively, after the transition we get a new one-step estimate for how valuable $S_t$ was: the reward we just received plus the discounted value of the state we transitioned to $S_{t+1}$. The TD error is the difference between this new one-step estimate and the old estimate we had for $S_t$ before seeing the transition.

What we want to prove next, in order to make the case of TD(0) as Robbins-Monro, is that the TD error is an unbiased estimate of the Bellman residual.

Derivation

Because $S_t$ and $V_t$ are already known under $\mathcal F_t$, they are treated as fixed inside the conditional expectation. In particular, $V_t(S_t)$ can be pulled out:

$$ \begin{aligned} \mathbb E_\pi[\delta_t\mid\mathcal F_t] &= \mathbb E_\pi[ R_{t+1}+\gamma V_t(S_{t+1})-V_t(S_t) \mid\mathcal F_t ]\\ &= \mathbb E_\pi[ R_{t+1}+\gamma V_t(S_{t+1}) \mid\mathcal F_t ] - V_t(S_t). \end{aligned} $$

The remaining randomness comes from the action selected by $\pi$ and then from the environment transition. Since $S_t$ is already part of $\mathcal F_t$, conditioning on both $\mathcal F_t$ and $S_t$ would not add new information: $\sigma(\mathcal F_t,S_t)=\mathcal F_t$. Still, from an educational perspective, it is useful to write $S_t$ explicitly so we can see the connection to the environment Markov property defined earlier, following Sutton and Barto's notation. We can condition on the next action and sum over actions:

$$ \begin{aligned} &\mathbb E_\pi[ R_{t+1}+\gamma V_t(S_{t+1}) \mid\mathcal F_t ]\\ &= \sum_a \pi(a\mid S_t) \mathbb E[ R_{t+1}+\gamma V_t(S_{t+1}) \mid \mathcal F_t, S_t, A_t=a ]. \end{aligned} $$

Now use the Markov property of the MDP. Once we know $S_t$ and $A_t=a$, the conditional distribution of $(R_{t+1},S_{t+1})$ is $p(r,s'\mid S_t,a)$. The current value estimate $V_t$ is still fixed because it is $\mathcal F_t$-measurable, so

$$ \begin{aligned} &\mathbb E[ R_{t+1}+\gamma V_t(S_{t+1}) \mid \mathcal F_t, S_t, A_t=a ]\\ &= \sum_{r,s'}p(r,s'\mid S_t,a) \left[r+\gamma V_t(s')\right]. \end{aligned} $$

Putting it all together gives

$$ \begin{aligned} \mathbb E_\pi[\delta_t\mid\mathcal F_t] &= \sum_a \pi(a\mid S_t) \sum_{r,s'}p(r,s'\mid S_t,a) \left[r+\gamma V_t(s')\right] - V_t(S_t)\\ &= (T_\pi V_t)(S_t)-V_t(S_t). \end{aligned} $$

If we define the Bellman residual function as

$$ h(V)=T_\pi V - V, $$

then the previous equation becomes

$$ \mathbb E_\pi[\delta_t\mid\mathcal F_t] = h(V_t)(S_t). $$

Result

The TD(0) error is a noisy unbiased observation of the Bellman residual at the sampled coordinate $S_t$.

Therefore, TD(0) prediction can be viewed as a Robbins-Monro stochastic approximation procedure for finding the root of the Bellman residual.^*

Robbins-Monro	TD(0) prediction
Find a root $h(x^\star)=0$.	Find a value function where $h(V)=T_\pi V - V = 0$.
Observe a noisy sample $H(x_t,\xi_{t+1})$ with conditional expectation $h(x_t)$.	Observe the TD error $\delta_t$, with $\mathbb E_\pi[\delta_t\mid\mathcal F_t]=h(V_t)(S_t)$.
Update by $x_{t+1}=x_t+\alpha_t H(x_t,\xi_{t+1})$.	Update the sampled state by $V_{t+1}(S_t)=V_t(S_t)+\alpha_t\delta_t$.

^* This convergence reading depends on the usual tabular on-policy TD assumptions: bounded rewards, sufficient state visitation, Markov noise, and learning rates satisfying $\sum_t \alpha_t=\infty$ and $\sum_t \alpha_t^2<\infty$.

Result

From the Bellman equation, we know that the unique root of the Bellman residual is the policy value function $v_\pi$. Therefore, under the usual tabular on-policy TD assumptions, TD(0) converges to the value function of the policy.

Q-learning as Robbins-Monro, and Why It Converges to the Optimal State-Action Value Function

Let's briefly revise how tabular Q-learning works. Q-learning maintains an action-value estimate $Q_t$. After observing a transition $(S_t,A_t,R_{t+1},S_{t+1})$, it forms the Q-learning TD error

$$ \delta_t^Q = R_{t+1} +\gamma \max_{a'} Q_t(S_{t+1},a') - Q_t(S_t,A_t). $$

Then it updates the sampled state-action pair:

$$ Q_{t+1}(S_t,A_t) = Q_t(S_t,A_t)+\alpha_t\delta_t^Q, $$

while leaving all other state-action values unchanged:

$$ Q_{t+1}(s,a)=Q_t(s,a), \qquad (s,a)\ne(S_t,A_t). $$

The key claim is analogous to the TD(0) case: once we condition on the current state-action pair, the Q-learning TD error is an unbiased estimate of the Bellman optimality residual.

Derivation

For the Q-learning update, we condition on the post-action information $\sigma(\mathcal F_t,A_t)$. Since $S_t$ is already $\mathcal F_t$-measurable, this is the same as $\sigma(\mathcal F_t,S_t,A_t)$. Under this conditioning, $Q_t$, $S_t$, and $A_t$ are fixed, so $Q_t(S_t,A_t)$ can be pulled out:

$$ \begin{aligned} \mathbb E[\delta_t^Q\mid\mathcal F_t,S_t,A_t] &= \mathbb E[ R_{t+1} +\gamma \max_{a'} Q_t(S_{t+1},a') -Q_t(S_t,A_t) \mid\mathcal F_t,S_t,A_t ]\\ &= \mathbb E[ R_{t+1} +\gamma \max_{a'} Q_t(S_{t+1},a') \mid\mathcal F_t,S_t,A_t ] - Q_t(S_t,A_t). \end{aligned} $$

Now use the environment Markov property. Once we know $(S_t,A_t)$, the conditional distribution of $(R_{t+1},S_{t+1})$ is $p(r,s'\mid S_t,A_t)$. The current estimate $Q_t$ is fixed because it is measurable with respect to the information available before the transition, so

$$ \begin{aligned} &\mathbb E[ R_{t+1} +\gamma \max_{a'} Q_t(S_{t+1},a') \mid\mathcal F_t,S_t,A_t ]\\ &= \sum_{r,s'}p(r,s'\mid S_t,A_t) \left[ r+\gamma \max_{a'}Q_t(s',a') \right]. \end{aligned} $$

Putting it all together gives

$$ \begin{aligned} \mathbb E[\delta_t^Q\mid\mathcal F_t,S_t,A_t] &= \sum_{r,s'}p(r,s'\mid S_t,A_t) \left[ r+\gamma \max_{a'}Q_t(s',a') \right] - Q_t(S_t,A_t)\\ &= (T_*Q_t)(S_t,A_t)-Q_t(S_t,A_t). \end{aligned} $$

If we define the Bellman optimality residual as

$$ h(Q)=T_*Q-Q, $$

then the previous equation becomes

$$ \mathbb E[\delta_t^Q\mid\mathcal F_t,S_t,A_t] = h(Q_t)(S_t,A_t). $$

Result

The Q-learning TD error is a noisy unbiased observation of the Bellman optimality residual at the sampled state-action coordinate $(S_t,A_t)$.

Therefore, Q-learning can be viewed as an asynchronous Robbins-Monro stochastic approximation procedure for finding the root of the Bellman optimality residual.

Robbins-Monro	Q-learning
Find a root $h(x^\star)=0$.	Find an action-value function where $h(Q)=T_*Q-Q=0$.
Observe a noisy sample $H(x_t,\xi_{t+1})$ with conditional expectation $h(x_t)$.	Observe $\delta_t^Q$, with $\mathbb E[\delta_t^Q\mid\mathcal F_t,S_t,A_t]=h(Q_t)(S_t,A_t)$.
Update by $x_{t+1}=x_t+\alpha_t H(x_t,\xi_{t+1})$.	Update the sampled coordinate by $Q_{t+1}(S_t,A_t)=Q_t(S_t,A_t)+\alpha_t\delta_t^Q$.

Result

The unique root of the Bellman optimality residual is the optimal state-action value function $Q_*$. Therefore, under the usual tabular Q-learning assumptions, Q-learning converges to $Q_*$.