How to Compute a Value Function? Policy Evaluation

What is Policy Evaluation?

The policy evaluation is the process of calculating the value function of a given (fixed) policy $\pi$.

Tip

Given a policy, policy evaluation asks: if I keep following this policy, how good is each state?

Assumption

The assumptions in this post are:

• the model is known,
• the environment is fully observable. (i.e., we know the current state $s$)
• the state space is finite and discrete,
• the action space is finite and discrete,
• the state transition model $T$ is known,
• the reward model $R$ is known,
• the policy $\pi$ is fixed,
• the discount factor satisfies $0\leq\gamma1$.

I will use the reward convention from the book:

It means the expected immediate reward after taking action $a$ in state $s$.

Notation

I’ll follow the notation convention in ADM:

• $U$ refers to utility, or value
• $R$ refers to reward function
• $T$ refers to state-transition probability function
• $\gamma$ refers to the discount factor

Remark

In this post, I mainly use a deterministic policy, so $\pi(s)$ directly returns an action. If $\pi$ is stochastic, i.e.$\pi(a \mid s)$, then we need to average over actions. For ease of reading, we assume $\pi$ is deterministic.

Solution 1: Iterative Approach

If the policy is executed for only one step, the utility is

This is easy to understand. If the agent moves only 1 step, we only care about the immediate reward. Given the current state $s$, the policy chooses action $\pi(s)$, and the reward is $R(s,\pi(s)).$

For more steps, we can use the recursive update:

^c76ebd

This equation is called the lookahead. Its equivalent julia code is:

function lookahead(𝒫::MDP, U, s, a)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    return R(s,a) + γ*sum(T(s,a,s′)*U(s′) for s′ in 𝒮)
end
function lookahead(𝒫::MDP, U::Vector, s, a)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    return R(s,a) + γ*sum(T(s,a,s′)*U[i] for (i,s′) in enumerate(𝒮))
end

Remark

Treating U as a function is closer to the mathematical definition. Treating U as a vector is easier to compute in discrete, finite environments like GridWorld.

In short,

• $U$-as-function is math-friendly.
• $U$-as-vector is implementation-friendly.

After enough iteration, $U^\pi_k$ converges to the value function $U^\pi$ due to contraction mapping(guaranteed to converge). At convergence, the old value and the new value are the same:

This equation is called Bellman expectation equation, and its equivalent code is

function iterative_policy_evaluation(𝒫::MDP, π, k_max)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    U = [0.0 for s in 𝒮]
    for k in 1:k_max
        U = [lookahead(𝒫, U, s, π(s)) for s in 𝒮]
    end
    return U
end

Remark

In many literature, the Bellman equation is loosely called. Depending on the context, it might be referred to

• Bellman expectation equation: evaluate a fixed policy $\pi$. This is for policy evaluation.
• Bellman optimality equation: choose best action by taking a max over action. This is for control / optimization. ^457806

Remark

I want to add some notes on contraction mapping. A function $f$ is a contraction mapping if $d(f(x), f(y)) \le k \cdot d(x, y)$ where $0\leq k<1$.

This means that the distance between the 2 outputs $f(x)$ and $f(y)$ is smaller than the distance between the 2 inputs $x$ and $y$.

Intuitively, if we repeatedly apply a contraction mapping $f(f(f(x)))\dots$, the result is guaranteed to move toward a fixed point where $f(x)=x.$

In the code U = [lookahead(𝒫, U, s, π(s)) for s in 𝒮], the old value U goes into the Bellman update, and the new value U comes out. Repeating this process moves U closer to the true value function.

Solution 2: System of Equations

ADM points out that policy evaluation can also be done without iteration by solving the Bellman expectation equation directly as system of equations. The matrix form is:

First, the vector $\mathbf{U}^{\pi}$ contains the value of every state, we have $|\mathcal{S}|$ states. If the states are $s_1,\dots,s_n$, then:

Remark

The $|\mathcal{S}|$ denotes the cardinal number of the state set. In this case, $|\mathcal{S}|=n$.

The vector $\mathbf{R}^{\pi}$ is the reward vector under policy $\pi$:

Remark

Because the policy is fixed, each $\pi(s_i)$ gives one specific action.

The matrix $\mathbf{T}^\pi$ is an $n \times n$ square matrix. Its entry in row $i$ and column $j$ is:

This is the probability of moving from state $s_i$​ to state $s_j$ when following policy $\pi$.

Remark

Note that the $\mathbf{T}^\pi$ is a stochastic matrix. Each row is a probability distribution, so each row sums to 1.

Remark

If the policy is stochastic, then the same matrix form still works. We only need to define: $R_i^\pi​=\sum_a \pi(a\mid s_i​)R(s_i​,a),$ and $T_{ij}^\pi = \sum_a \pi(a \mid s_i) T(s_j \mid s_i, a).$

Now all notation is ready. Let’s unpack the matrix equation:

Recall the scalar Bellman expectation equation:

The matrix-vector multiplication $\mathbf{T}^{\pi} \mathbf{U}^{\pi}$ is exactly the vector form of:

Suppose the diagram uses 3 states, i.e., $|\mathcal{S}|=3$:

Then:

where $U_i = U^\pi(s_i)$.

The transition matrix $\mathbf{T}^\pi$ is a $3\times3$ square matrix:

Here, $T_{s_is_j}=T(s_j\mid s_i,\pi(s_i)).$

The symbol $\displaystyle \sum_{s'}$ means we enumerates all possible next states. This is why the matrix-vector product uses a $3$-vector. For the row corresponding to current state $s_1$, we get

Therefore

Using Strang’s favorite “column picture”, we can have the following diagrams help us understand the matrix-vector product better.

So the matrix equation

is simply the vectorized form of

Now, how do we solve it?

Start from:

Rearrange to isolate $\mathbf{U}^{\pi}$:

$\mathbf{U}^{\pi} - \gamma \mathbf{T}^{\pi} \mathbf{U}^{\pi} = \mathbf{R}^{\pi}.$

Factor out $\mathbf{U}^{\pi}$ ($\mathbf{I}$ is the $n \times n$ identity matrix):

$(\mathbf{I} - \gamma \mathbf{T}^{\pi}) \mathbf{U}^{\pi} = \mathbf{R}^{\pi}.$

Finally, multiply both sides by the inverse:

$\mathbf{U}^{\pi} = (\mathbf{I} - \gamma \mathbf{T}^{\pi})^{-1} \mathbf{R}^{\pi}.$

In code, however, we usually should not compute the inverse directly. It is better to solve the linear system:

function policy_evaluation(𝒫::MDP, π)
	𝒮, R, T, γ = 𝒫.𝒮, 𝒫.R, 𝒫.T, 𝒫.γ
	R′ = [R(s, π(s)) for s in 𝒮]
	T′ = [T(s, π(s), s′) for s in 𝒮, s′ in 𝒮]
	return (I - γ*T′)\R′
end

Remark

The backslash operator \ solves the linear system directly.

Supplement Question

Now we know the “what”. But let’s ask “why”.

Why we can write:

At first, this may feel confusing because $\mathbf{U}^{\pi}$ appears on both sides. But in the iterative code, we have:

function iterative_policy_evaluation(𝒫::MDP, π, k_max)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    U = [0.0 for s in 𝒮]
    for k in 1:k_max
        U = [lookahead(𝒫, U, s, π(s)) for s in 𝒮]
    end
    return U
end

So it looks like there are 2 different ideas:

1. iterative update: old $U$ goes in, new $U$ comes out
2. matrix equation: the same $U^\pi$ appears on both sides

How should we understand this?🤔

The key idea is that the matrix equation describes the fixed point of the iterative process.

The iterative method is

Define the Bellman expectation operator $\mathcal{B}_{\pi}$ as:

At convergence, the value no longer changes:

So the fixed-point equation becomes

which is exactly

Tip

$\mathcal{B}_{\pi}$ is a contraction mapping in the $\ell_\infty$ norm with Lipschitz constant $\gamma < 1$.

More explicitly:

$\|\mathcal{B}_{\pi}(\mathbf{U}) - \mathcal{B}_{\pi}(\mathbf{V})\|_\infty \;\leq\; \gamma \|\mathbf{U} - \mathbf{V}\|_\infty.$

Why?

The inequality holds because each row of $\mathbf{T}^\pi$ is a probability distribution.

Remark

By the Banach fixed-point theorem, the fixed point exists, is unique, and the iterative methods converges to it.

Now let’s connect this with the matrix solution.

Because $\mathbf{T}^\pi$ is stochastic matrix, its eigenvalues have magnitude at most 1. Since $0\leq\gamma<1$, the eigenvalues of $\gamma\mathbf{T}^\pi$ have magnitude strictly less than 1.

Therefore, the Neumann series converges:

So:

This also explains the iterative method.

If we start with ($\mathbf{U}_0=\mathbf{0}$), then:

After many iterations, we get:

That is the same infinite series as the matrix inverse solution!

So the 2 methods are not contradictory:

• iterative policy evaluation approaches the fixed point step by step
• the linear-system method solves the fixed point directly

To conclude,

• Iterative policy evaluation is fixed-point iteration for the Bellman expectation equation.
• Exact policy evaluation solves the same fixed point as a linear system.