Value-based RL: From Value Iteration to DQN

Background

The Reinforcement Learning Problem

Reinforcement learning considers how an agent should act in an environment to achieve a certain goal. Much of the notation and structure of this post is derived from Peter Henderson’s COS 435. We define the environment as a Markov Decision Process (MDP): a tuple $(\mathcal{S}, \mathcal{A}, \mathcal{T}, R, \gamma)$ , where

$\mathcal{S}$ is the state space
$\mathcal{A}$ is the action space
$\mathcal{T}: \mathcal{S} \times \mathcal{A} \to \Delta(\mathcal{S})$ is the transition kernel, where $\mathcal{T}(s' \mid s, a) = P(s_{t+1} = s' \mid s_t = s, a_t = a)$
$R: \mathcal{S} \times \mathcal{A} \to \mathbb{R}$ is the expected immediate reward function, where $R(s, a) = \mathbb{E}[r_t \mid s_t = s, a_t = a]$
$\gamma \in [0, 1)$ is the discount factor

The agent’s goal is to find the optimal policy $\pi^*: \mathcal{S} \to \mathcal{A}$ that maximizes the expected discounted reward:

\pi^* = \arg \max_{\pi} \mathbb{E}_{\pi} \left[ \sum_{t=0}^\infty \gamma^t r_t \mid s_0 \right]

where $r_t$ is the sampled reward received at time $t$ . To simplify notation, we define the return $G_t$ as the discounted sum of rewards:

G_t = \sum_{k=0}^\infty \gamma^k r_{t+k}.

Then, the optimal policy is

\pi^* = \arg \max_{\pi} \mathbb{E}_{\pi} \left[ G_0 \mid s_0 \right]

Value Functions

The value function is the expected return given that we start at state $s$ and follow policy $\pi$ :

V^\pi(s) = \mathbb{E}_{\pi} \left[ G_t \mid s_t = s \right].

The optimal value function is the best expected return we can get from any policy:

V^*(s) = \max_{\pi} V^\pi(s).

Similarly, the action-value function is the expected return given that we start at state $s$ , take action $a$ , and follow policy $\pi$ :

Q^\pi(s, a) = \mathbb{E}_{\pi} \left[ G_t \mid s_t = s, a_t = a \right].

The optimal action-value function is the best expected return we can get from any policy, starting at state $s$ and taking action $a$ :

Q^*(s, a) = \max_{\pi} Q^\pi(s, a).

Model-based vs. Model-free RL

A brief aside to frame the following discussion: RL methods are often split into model-based and model-free approaches. Model-based RL learns (or assumes access to) the environment dynamics (transition probabilities $\mathcal{T}$ and reward function $R$ ), while model-free RL learns value functions or policies directly from experience (no explicit environment model).

Value Iteration

Let’s assume we know the environment dynamics. How can we recover the optimal policy?

Bellman Equations

The Bellman optimality equations say that the optimal value function $V^*(s)$ must satisfy:

V^*(s) = \max_a \left[R(s, a) + \gamma \mathbb{E}_{s' \sim \mathcal{T}(\cdot \mid s, a)} V^*(s')\right].

For the optimal action-value function $Q^*(s, a)$ ,

Q^*(s, a) = R(s, a) + \gamma \mathbb{E}_{s' \sim \mathcal{T}(\cdot \mid s, a)} \left[\max_{a'} Q^*(s', a')\right].

The Bellman optimality operator $\mathcal{B}$ maps any value function $V$ to

(\mathcal{B}V)(s) = \max_a [R(s, a) + \gamma \mathbb{E}_{s' \sim \mathcal{T}(\cdot \mid s, a)} V(s')].

Notice that the Bellman optimality equation is simply requiring that the optimal value function is a fixed point of the Bellman operator:

V^* = \mathcal{B}V^*.

To measure how good a candidate value function $V$ is, we use the Bellman error (also called the Bellman residual), i.e., how far $V$ is from being a fixed point of a Bellman operator. The Bellman error is then

\delta_B(s) = (\mathcal{B}V)(s) - V(s).

Value Iteration Algorithm

A natural idea arises out of the Bellman operator. We can initialize the value function for all states to some value, then repeatedly apply the Bellman operator until we get the optimal value function:

Initialize V(s) = 0 for all s in S
repeat
    Δ <- 0
    for each state s in S:
        v <- V(s)
        V(s) <- max_a [R(s, a) + γ * Σ_{s' in S} T(s' | s, a) V(s')]
        Δ <- max(Δ, |v - V(s)|)
until Δ < ε
return V

To better understand value iteration, one should walk through a grid world example. The textbooks in Sources walk through grid world examples.

Policy Extraction

Given the optimal value function $V^*$ , extracting the optimal policy is simple:

\pi^*(s) = \arg \max_a \left[R(s, a) + \gamma \mathbb{E}_{s' \sim \mathcal{T}(\cdot \mid s, a)} V^*(s')\right]

Framing Value Iteration

Value iteration is relatively simple and has convergence guarantees, but there are a few issues:

We often don’t know the environment dynamics.
State spaces can be large and thus make summing over all states computationally infeasible.

Learning Without a Model

Monte Carlo Control

Monte Carlo methods use random sampling to estimate certain numbers. In our case, we can use rollouts to estimate values to avoid explicitly computing Bellman updates.

Value iteration becomes a lot more simple with Monte Carlo updates. For each iteration, we sample a trajectory from the environment starting at some state $s$ and following policy $\pi$ . Then, we update the value function for each state in the trajectory:

V(s_t) \leftarrow V(s_t) + \eta \left[G_t - V(s_t)\right],

where $\eta$ is the learning rate and $G_t$ is the sampled return from time $t$ .

For Monte Carlo control, we apply the same update idea to Q-functions:

Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \eta \left[G_t - Q(s_t, a_t)\right],

Then, we improve the policy greedily:

\pi_{k+1}(s) = \arg\max_a Q_k(s, a).

We repeat these two steps until the policy converges.

Since Monte Carlo control uses rollouts, we must follow a policy during rollouts (unlike in value iteration). If our policy doesn’t improve, then our Q-function estimates also won’t improve. This is why Monte Carlo control requires us to alternate between policy evaluation (updating the Q-function estimates) and policy improvement (improving the policy). Monte Carlo control is an instantiation of policy iteration.

Temporal Difference Learning

A major issue with Monte Carlo control is that we must obtain an entire trajectory before performing updates. For policy evaluation under $\pi$ , the ideal Bellman-style update is

V(s_t) \leftarrow V(s_t) + \eta \left[(\mathcal{B}^\pi V)(s_t) - V(s_t)\right]

where

(\mathcal{B}^\pi V)(s_t) = \mathbb{E}_{\pi}[r_{t+1} + \gamma V(s_{t+1}) \mid s_t].

Instead, we can approximate the Bellman update using one sample’s transition to the next state $s_{t+1}$ :

(\mathcal{B}^\pi V)(s_t) \approx r_{t+1} + \gamma V(s_{t+1}).

Plugging this approximation into the Bellman update gives the TD(0) update:

V(s_t) \leftarrow V(s_t) + \eta \left[r_{t+1} + \gamma V(s_{t+1}) - V(s_t)\right].

Notice that TD learning relies on bootstrapping, the idea that we can use a guess to make a better guess. The term $\delta_t := r_{t+1} + \gamma V(s_{t+1}) - V(s_t)$ is called the TD error.

TD(0) is an unbiased estimator For some quantity $X$ , an estimator $\hat{X}$ is unbiased if $\mathbb{E}[\hat{X}] = X$ . of the Bellman error (conditioned on the current state):

\begin{aligned} \mathbb{E}[\delta_t \mid s_t] &= \mathbb{E}[r_{t+1} + \gamma V(s_{t+1}) - V(s_t) \mid s_t] \\ &= \mathbb{E}[r_{t+1} + \gamma V(s_{t+1}) \mid s_t] - V(s_t) \\ &= (\mathcal{B}^\pi V)(s_t) - V(s_t) \\ &= \delta_B^\pi(s_t). \end{aligned}

However, TD(0) suffers from high variance. To reduce variance, we can change the number of steps $n$ we look ahead to get n-step TD updates:

V(s_t) \leftarrow V(s_t) + \eta \left[\left(\sum_{k=0}^{n-1} \gamma^k r_{t+k+1} + \gamma^n V(s_{t+n})\right) - V(s_t)\right]