Introduction
I'm studying Reinforcement Learning, and in order to increase my understanding I've been challenging myself by trying to write proofs that show that the right hand side of the Bellman equations and Bellman optimality equations for state-value and action-value functions are contraction operators. The fact that the Bellman operators are contractions is of great importance, since this proves that the Bellman equations have unique fixed points that can be solved for iteratively, due to the Banach fixed-point theorem.
Below are my attempts at the proofs, my questions are:
- Are the proofs correct?
- Any feedback on what can be improved (I'm very inexperienced in writing proofs and would like to get better)?
I will assume some knowledge about finite Markov Decision Processes (MDPs), please let me know if you want me to elaborate.
Bellman equations
Let $\mathcal{S}$ denote the finite set of all states, $\mathcal{R} \subset \mathbb{R}$ the finite set of all rewards, and $\mathcal{A}(s)$ the finite set of all available actions $a$ in state $s$. Let $r_\pi(s)$ denote the expected immediate reward and $p_\pi(s^\prime | s)$ the probability of transitioning to state $s^\prime$ when following policy $\pi$ from state $s$. Let $\gamma \in (0, 1)$ denote the discount rate for future rewards. The Bellman equation for the state-value function for policy $\pi$ can be defined as follows:
$$ \begin{aligned} v_{\pi}(s) &= \mathbb{E}_\pi \big[R_{t+1} + \gamma v_{\pi}(S_{t+1}) | S_t = s \big] \\ &= \sum_{a \in \mathcal{A}(s)} \pi(a | s) \bigg[\sum_{r \in \mathcal{R}} p(r | s, a) r + \gamma \sum_{s^\prime \in \mathcal{S}} p(s^\prime | s, a) v_{\pi}(s^\prime) \bigg] \\ &= r_\pi(s) + \gamma \sum_{s^\prime \in \mathcal{S}} p_\pi(s^\prime | s) v_{\pi}(s^\prime) \end{aligned} \tag 1 $$
for all $s \in \mathcal{S}$. Let $n = |\mathcal{S}|$, we can write the equation in matrix form:
$$ \begin{bmatrix} v_\pi(1) \\ \vdots \\ v_\pi(n) \end{bmatrix}= \begin{bmatrix} r_\pi(1) \\ \vdots \\ r_\pi(n) \end{bmatrix} +\gamma \begin{bmatrix} p_\pi(1 | 1) & \dots & p_\pi(n | 1) \\ \vdots & \ddots & \vdots\\ p_\pi(1 | n) & \dots & p_\pi(n | n) \end{bmatrix} \begin{bmatrix} v_\pi(1) \\ \vdots \\ v_\pi(n) \end{bmatrix} \tag 2 $$
More compactly:
$$ v_\pi = r_\pi + \gamma P_\pi v_\pi \tag 3 $$
We can define an expected Bellman operator $\mathcal{T}^\pi : \mathbb{R}^n \to \mathbb{R}^n$ as:
$$ \mathcal{T}^\pi(v) = r_\pi + \gamma P_\pi v \tag 4 $$
for any $v \in \mathbb{R}^n$.
Let $||\cdot||$ be a norm in $\mathbb{R}^n$. If there exists a $\gamma \in (0, 1)$ such that $||\mathcal{T}^\pi(v_1) - \mathcal{T}^\pi(v_2)|| \leq \gamma ||v_1 - v_2||$ for all $v_1, v_2 \in \mathbb{R}^n$, then $\mathcal{T}^\pi$ is a contraction mapping. In all proofs in this post $|\cdot|$ and $\leq$ are elementwise, and the norm used is the max norm $||x||_\infty = \max(|x|) = \max(|x_1|, \dots, |x_n|)$, where $\max(\cdot) : \mathbb{R}^n \to \mathbb{R}$ chooses the largest element in a vector.
$$ \begin{aligned} ||\mathcal{T}^\pi(v_1) - \mathcal{T}^\pi(v_2)||_\infty &= \max \big(|r_\pi + \gamma P_\pi v_1 - (r_\pi + \gamma P_\pi v_2)| \big) \\ &= \gamma \max \big(|P_\pi(v_1 - v_2)| \big) \\ &\leq \gamma \max \big(P_\pi|v_1 - v_2| \big) \\ &\leq \gamma \max \big(|v_1 - v_2| \big) \\ &= \gamma ||v_1 - v_2||_\infty \end{aligned} \tag 5 $$
Thus $\mathcal{T}^\pi$ is a contraction. The last inequality is due to the rows of $P_\pi$ containing only non-negative elements that sum to 1.
Let $r(s, a)$ denote the expected immediate reward when selecting action $a$ in state $s$, and $p_\pi(s^\prime, a^\prime | s, a)$ the probability of transitioning to state $s^\prime$ and selecting action $a^\prime$ when selecting action $a$ in state $s$ and following policy $\pi$ after. The Bellman equation for the action-value function for policy $\pi$ can be defined as follows:
$$ \begin{aligned} q_{\pi}(s, a) &= \mathbb{E}_\pi \big[R_{t+1} + \gamma q_{\pi}(S_{t+1}, A_{t+1}) | S_t = s, A_t = a \big] \\ &= \sum_{r \in \mathcal{R}} p(r | s, a) r + \gamma \sum_{s^\prime \in \mathcal{S}} p(s^\prime | s, a) \sum_{a^\prime \in \mathcal{A}(s^\prime)} \pi(a^\prime | s^\prime) q_{\pi}(s^\prime, a^\prime) \\ &= r(s, a) + \gamma \sum_{s^\prime \in \mathcal{S}} \sum_{a^\prime \in \mathcal{A}(s^\prime)} p_\pi(s^\prime, a^\prime | s, a) q_{\pi}(s^\prime, a^\prime) \end{aligned} \tag 6 $$
for all $s \in \mathcal{S}$, $a \in \mathcal{A}(s)$. Let $n_s = |\mathcal{A}(s)|$, we can write the equation in matrix form:
$$ \begin{bmatrix} q_\pi(1, 1) \\ q_\pi(1, 2) \\ \vdots \\ q_\pi(n, n_n) \end{bmatrix}= \begin{bmatrix} r_\pi(1, 1) \\ r_\pi(1, 2) \\ \vdots \\ r_\pi(n, n_n) \end{bmatrix} +\gamma \begin{bmatrix} p_\pi(1, 1 | 1, 1) & p_\pi(1, 2 | 1, 1) & \dots & p_\pi(n, n_n | 1, 1) \\ p_\pi(1, 1 | 1, 2) & p_\pi(1, 2 | 1, 2) & \dots & p_\pi(n, n_n | 1, 2) \\ \vdots & \vdots & \ddots & \vdots \\ p_\pi(1, 1 | n, n_n) & p_\pi(1, 2 | n, n_n) & \dots & p_\pi(n, n_n | n, n_n) \end{bmatrix} \begin{bmatrix} q_\pi(1, 1) \\ q_\pi(1, 2) \\ \vdots \\ q_\pi(n, n_n) \end{bmatrix} \tag 7 $$
More compactly:
$$ q_\pi = r_\pi + \gamma P_\pi q_\pi \tag 8 $$
The only difference compared to the equation for the state-value function is that the vectors and matrices are larger. Thus, we can define an expected Bellman operator in identical fashion (with $n$ denoting the number of state-action pairs rather than the number of states) and the proof will be identical (the rows of $P_\pi$ still contain only non-negative elements that sum to 1).
Bellman optimality equations
For brevity I will go directly into the compact matrix form of the Bellman optimality equation for the optimal state-value function:
$$ v_* = \max_\pi(r_\pi + \gamma P_\pi v_*) \tag 9 $$
We can define an expected Bellman optimality operator $\mathcal{T}^* : \mathbb{R}^n \to \mathbb{R}^n$ as:
$$ \mathcal{T}^*(v) = \max_\pi(r_\pi + \gamma P_\pi v) \tag {11} $$
The rest of the post will focus on trying to prove that this is a contraction operator, because once again I believe that an identical proof works for the action-value function (please correct me if this is wrong).
Consider any two vectors $v_1, v_2 \in \mathbb{R}^n$, and let $\pi_1^* = \text{argmax}_\pi(r_\pi + \gamma P_\pi v_1)$ and $\pi_2^* = \text{argmax}_\pi(r_\pi + \gamma P_\pi v_2)$. Then we have:
$$ \mathcal{T}^*(v_1) = \max_\pi(r_\pi + \gamma P_\pi v_1) = r_{\pi_1^*} + \gamma P_{\pi_1^*} v_1 \geq r_{\pi_2^*} + \gamma P_{\pi_2^*} v_1 \\ \mathcal{T}^*(v_2) = \max_\pi(r_\pi + \gamma P_\pi v_2) = r_{\pi_2^*} + \gamma P_{\pi_2^*} v_2 \geq r_{\pi_1^*} + \gamma P_{\pi_1^*} v_2 \tag{12} $$
$$ \begin{aligned} \mathcal{T}^*(v_1) - \mathcal{T}^*(v_2) &= r_{\pi_1^*} + \gamma P_{\pi_1^*} v_1 - (r_{\pi_2^*} + \gamma P_{\pi_2^*} v_2) \\ &\leq r_{\pi_1^*} + \gamma P_{\pi_1^*} v_1 - (r_{\pi_1^*} + \gamma P_{\pi_1^*} v_2) \\ &= \gamma P_{\pi_1^*} (v_1 - v_2) \end{aligned} $$
Similarly we have $\mathcal{T}^*(v_2) - \mathcal{T}^*(v_1) \leq \gamma P_{\pi_2^*} (v_2 - v_1)$, which implies that $\mathcal{T}^*(v_1) - \mathcal{T}^*(v_2) \geq \gamma P_{\pi_2^*} (v_1 - v_2)$, and thus we have:
$$ \gamma P_{\pi_2^*} (v_1 - v_2) \leq \mathcal{T}^*(v_1) - \mathcal{T}^*(v_2) \leq \gamma P_{\pi_1^*} (v_1 - v_2) \tag{13} $$
Let $\max\{\cdot, \cdot \} : \mathbb{R}^n \times \mathbb{R}^n \to \mathbb{R}^n$ choose the largest values between two vectors elementwise, we have:
$$ \begin{aligned} |\mathcal{T}^*(v_1) - \mathcal{T}^*(v_2)| &\leq \max \{|\gamma P_{\pi_2^*} (v_1 - v_2)|, |\gamma P_{\pi_1^*} (v_1 - v_2)| \} \\ &\leq \gamma \max \{P_{\pi_2^*} |v_1 - v_2|, P_{\pi_1^*} |v_1 - v_2| \} \end{aligned} $$Let $\max(\cdot) : \mathbb{R}^n \to \mathbb{R}$ choose the largest element in a vector (as previously defined), we have:
$$ \begin{aligned} ||\mathcal{T}^*(v_1) - \mathcal{T}^*(v_2)||_\infty &= \max \big(|\mathcal{T}^*(v_1) - \mathcal{T}^*(v_2)| \big) \\ &\leq \gamma \max \big(\max \{P_{\pi_2^*} |v_1 - v_2|, P_{\pi_1^*} |v_1 - v_2| \} \big) \\ &\leq \gamma \max \big(|v_1 - v_2| \big) \\ &= \gamma ||v_1 - v_2||_\infty \end{aligned} $$Thus $\mathcal{T}^*$ is a contraction. Note that the last inequality is once again due to the rows of $P_\pi$ containing only non-negative elements that sum to 1.
Thank you for reading this far! If you have the time, please let me know if the proofs are correct, if the proofs for the state-value and action-value functions are identical as I suggest, and any other feedback.