Reinforcement Learning — States, Returns and the Bellman Equation

Machine Learning Notes · reinforcement learning 1 of 2 · Foundations ← Collaborative filtering Deep Q-learning →

Everything until now started with a dataset. Supervised learning had examples and answers; unsupervised learning had examples; recommenders had a matrix of ratings. Reinforcement learning has none of them. There is an agent, a world it can act in, and a number that occasionally arrives saying how things are going. Nobody ever says which action was correct.

That last part is the whole difficulty. A rover that drives left for three steps and then finds a +100 crater was not told "left" three times — it was told nothing, twice, and then told 100. Working out which of those decisions deserves the credit is what this subject is about.

In one line. The value of an action is the return it eventually leads to — the sum of future rewards, each discounted by $\gamma$ per step into the future. Write that value as $Q(s,a)$, and it satisfies a recursion, $Q(s,a) = R(s) + \gamma \max_{a'} Q(s',a')$: one step of reward now, plus the discounted value of behaving optimally from wherever you land. The best policy is then just $\pi(s) = \arg\max_a Q(s,a)$.

Contents
  1. What makes it a different problem
  2. The Mars rover
  3. The return
  4. Policies
  5. The Markov decision process
  6. The state-action value function
  7. The Bellman equation
  8. When actions do not always work
  9. Where this goes next

What makes it a different problem

Three things are new, and each breaks an assumption used everywhere else in these notes:


The Mars rover

States, actions, rewards

The example is a rover on a six-cell strip of ground:

State123456
Descriptionleft crater
terminal
flatflatstartflatright crater
terminal
Reward $R(s)$+1000000+40
Table 1: The Mars rover environment. Both ends are terminal — reaching either ends the episode — and every cell in between pays nothing. The rover starts in the middle, with the large reward three steps to its left and the small one two steps to its right.

At every state the rover picks one of two actions: left ($\leftarrow$) or right ($\rightarrow$). Each action moves it one cell.

One step, written down

Every interaction is recorded as the same four-part tuple, and this is the only data reinforcement learning ever gets:

$$ \left(s,\; a,\; R(s),\; s'\right) \qquad\text{for example}\qquad \left(4,\; \leftarrow,\; 0,\; 3\right). \tag{1} $$
Equation 1: One step of experience. The current state, the action taken, the reward collected, and where the agent ended up. A whole episode is a chain of these, and the next note stores millions of them in a replay buffer.

That example says: from state 4 the rover went left, received nothing, and is now in state 3.


The return

Definition

The return $G$ is the total future reward, with each step discounted by a factor $\gamma$ per unit of delay:

$$ G_t = R_t + \gamma R_{t+1} + \gamma^2 R_{t+2} + \gamma^3 R_{t+3} + \cdots = \sum_{k=0}^{\infty} \gamma^k R_{t+k} , \qquad 0 \le \gamma \le 1 . \tag{2} $$
Equation 2: The discounted return. $\gamma$ between 0 and 1 makes a reward worth less the longer you have to wait for it; with $\gamma = 1$ nothing is discounted and with $\gamma = 0$ only the immediate reward counts.

Discounting is not only impatience. It also keeps the sum finite when an episode might never end, and it encodes a real preference: a reward now is worth more than the same reward after a hundred risky steps.

What the discount factor does

Start at state 4 and drive left. The +100 arrives three steps later, so the return is $\gamma^3 \times 100$:

$$ \gamma = 0.9: \quad G = 0.9^3 \times 100 = 72.9 \qquad\qquad \gamma = 0.5: \quad G = 0.5^3 \times 100 = 12.5 . $$

Same path, same reward, and a factor of nearly six between them.

Figure 1: How much a reward of 100, three steps away, is worth today — as a function of $\gamma$. At $\gamma = 0.9$ the agent still sees 72.9 of it and will happily walk; at $\gamma = 0.5$ only 12.5 survives, and the curve collapses toward zero below about 0.4. The two marked points are the cases worked above.

Returns from every state

With $\gamma = 0.5$ fixed, the return depends on where you start and which way you commit to going:

State123456
Reward $R(s)$100000040
Return, always $\leftarrow$100502512.56.2540
Return, always $\rightarrow$1002.55102040
Return, mixed ($\leftarrow$ then $\rightarrow$)100502512.52040
Table 2: Returns under three fixed behaviours, with $\gamma = 0.5$. Terminal states pay their own reward. Going left beats going right from states 2, 3 and 4; from state 5 the near +40 wins, because two steps of discounting cost less than four. The third row is a mixed policy — left from 2, 3, 4 and right from 5 — and it takes the better of the two everywhere.
Figure 2: The same three behaviours as bars. The two terminal states are fixed points — every behaviour collects the same thing there — so read the middle four. Going left wins at 2, 3 and 4 by a widening margin as you approach the +100, and loses at 5, where the +40 is one step away instead of four. The green bars are the pointwise maximum of the other two, which is what an optimal policy takes.

The mixed row is the interesting one: the best behaviour is not the same everywhere. From state 5 it is right; from state 4, one cell away, it is left. Which is exactly why the thing being learned has to be a function of the state.


Policies

A policy $\pi$ is that function — it maps each state to the action to take there:

$$ \pi(s) = a, \qquad\text{written}\qquad s \;\xrightarrow{\;\pi\;}\; a . \tag{3} $$
Equation 3: A policy. Nothing more than a lookup: given the situation, what do I do? The goal of reinforcement learning is stated entirely in terms of this object — find the $\pi$ whose expected return is largest.

For the rover, the policy that the return table above recommends is $\pi(2) = \pi(3) = \pi(4) = {\leftarrow}$ and $\pi(5) = {\rightarrow}$. Drawn on the strip it is just an arrow in each cell — a complete description of the agent's behaviour, with no memory of how it got there.


The Markov decision process

The loop

Everything above is an instance of one formal object, the Markov decision process. Each step: the agent is in state $s_t$, its policy picks $a_t = \pi(s_t)$, and the environment answers with a reward $R_t$ and a next state $s_{t+1}$.

Figure 3: The agent-environment loop. The only channel from agent to world is the action; the only channel back is the reward and the next state. Markov means that channel is enough — the next state depends on the current state and action, not on anything earlier.

The word Markov is the assumption that makes this tractable: the future depends only on the present state and action, not on the path taken to reach it. Choosing a state representation that actually has this property is a design decision, and the next note shows what it looks like when the state is a vector of positions and velocities.

Three domains

DomainStatesActions$\gamma$Rewards
Mars rover6 grid positionsmove $\leftarrow$ or $\rightarrow$0.5terminal 100 and 40, else 0
Helicopterposition and attitudecontrol stick0.99$+1$ flying well, $-1000$ crash
Chesspieces on the boardany legal move0.995$+1$ win, $0$ draw, $-1$ loss
Table 3: The same formalism at three scales. Only the discount factor is worth staring at: a rover two steps from a crater can afford 0.5, while a helicopter making a hundred control decisions per second needs 0.99 or its whole flight would be worth nothing, and chess at 0.995 must value a win eighty moves away.

That $\gamma$ column is worth one picture, because the difference between $0.5$ and $0.995$ is not a matter of degree — it is the difference between an agent that can see two steps and one that can see two hundred:

Figure 4: What a reward $k$ steps away is still worth, for the three discount factors in the table above. The x-axis is logarithmic, so each curve is the same shape slid sideways. The dashed line at 0.5 marks the half-life — the delay at which a reward has lost half its value — and it lands at 1 step for $\gamma = 0.5$, 6.6 steps for $\gamma = 0.9$, 69 steps for $\gamma = 0.99$ and 138 steps for $\gamma = 0.995$. The useful rule of thumb is that an agent effectively plans about $1/(1-\gamma)$ steps ahead: 2, 10, 100, 200.

The state-action value function

What $Q$ means

Here is the central definition, and it is worth reading twice because of the odd conditional in it:

$$ Q^\ast(s,a) = \mathbb{E}\left[\sum_{t=0}^{\infty} \gamma^t R_{t+1} \;\middle|\; s_0 = s,\; a_0 = a,\; \text{optimal thereafter}\right], \qquad V^\ast(s) = \max_a Q^\ast(s,a). \tag{4} $$
Equation 4: The optimal state-action value. The return if you start in state $s$, take action $a$ once — even if that action is a mistake — and behave optimally from then on. The state value $V^\ast$ is what you get by taking the best action in the first place.

The "take $a$ once, then be optimal" construction is what makes $Q$ useful. It lets you compare actions from the same state on equal terms, and the comparison immediately gives the policy.

Working it out on the chain

Because all intermediate rewards are zero, $V^\ast(s)$ is just the best discounted terminal reward reachable from $s$:

$$ V^\ast(1) = 100, \quad V^\ast(2) = 0.5 \times 100 = 50, \quad V^\ast(3) = 0.5^2 \times 100 = 25, $$ $$ V^\ast(4) = 0.5^3 \times 100 = 12.5, \quad V^\ast(5) = 0.5 \times 40 = 20, \quad V^\ast(6) = 40 . $$

And each $Q^\ast$ is one step of that: move, then collect the value of where you land.

$$ Q^\ast(2, \rightarrow) = 0 + \gamma\, V^\ast(3) = 0.5 \times 25 = 12.5, \qquad Q^\ast(2, \leftarrow) = 0 + \gamma\, V^\ast(1) = 0.5 \times 100 = 50 . $$

Reading the policy off $Q$

State$R(s)$$Q^\ast(s, \leftarrow)$$Q^\ast(s, \rightarrow)$$V^\ast(s)$greedy action
1100100100100terminal
205012.550$\leftarrow$
30256.2525$\leftarrow$
4012.51012.5$\leftarrow$
506.252020$\rightarrow$
640404040terminal
Table 4: Every value in the chain, with $\gamma = 0.5$. The greedy column is simply whichever of the two $Q$ values is larger, and it reproduces the mixed policy found by inspection earlier — left in states 2 to 4, right in state 5. Note state 4: 12.5 against 10 is the closest call in the table.
Figure 5: The two action values side by side in each non-terminal state; the taller bar is the greedy action, and $V^\ast$ is its height. Three states prefer left and one prefers right, and the picture makes the reason plain — the left bars decay by a factor of 2 per cell of distance from the +100, the right bars grow by the same factor as the +40 approaches, and they cross between state 4 and state 5. State 4 is where the two are closest: 12.5 against 10.

This is the payoff of the whole construction. Once $Q^\ast$ is known, the optimal policy needs no search, no planning and no model of the world — pick the larger number:

$$ \pi^\ast(s) = \arg\max_a Q^\ast(s,a). $$

Everything in the next note is an attempt to get $Q$ for problems where this table cannot be written down.


The Bellman equation

The equation

The pattern used in §6.2 — one step of reward, then the value of the next state — is general, and it is the equation the entire subject is built on:

$$ Q^\ast(s,a) = R(s) + \gamma \max_{a'} Q^\ast(s', a'), \tag{5} $$
Equation 5: The Bellman optimality equation. The reward for being where you are, plus the discounted value of behaving optimally from wherever the action takes you. It is a definition of $Q^\ast$ in terms of itself, which is exactly why it can be turned into a training target.

where $s'$ is the state reached from $s$ by taking $a$. For terminal states there is no next state and $Q^\ast(s,a) = R(s)$.

Why one step is enough

It looks as though this can only see one step ahead, and the whole point was that rewards arrive late. The resolution is that $\max_{a'} Q^\ast(s',a')$ is itself defined by the same equation, so unrolling it gives

$$ Q^\ast(s,a) = R(s) + \gamma R(s') + \gamma^2 R(s'') + \cdots $$

— the full discounted return, reconstructed one link at a time. The recursion has folded an infinite sum into a relation between neighbours, and that is what makes it computable.

Check it on a row. Take $Q^\ast(4, \leftarrow)$. Bellman says $R(4) + \gamma \max_{a'} Q^\ast(3, a') = 0 + 0.5 \times 25 = 12.5$, which is what the table says. The 25 already contains the discounted 100 two steps further on.


When actions do not always work

Expected return

Real actuators slip. Suppose each command succeeds with probability $1 - p$ and sends the rover the wrong way with probability $p = 0.1$. Now the return is a random variable, and what gets maximised is its expectation:

$$ \begin{align} Q(s, \leftarrow) &= R(s) + \gamma\Big[(1-p)\,V(s-1) + p\,V(s+1)\Big], \tag{6.1} \\ Q(s, \rightarrow) &= R(s) + \gamma\Big[(1-p)\,V(s+1) + p\,V(s-1)\Big], \tag{6.2} \\ V(s) &= \max\big\{Q(s,\leftarrow),\, Q(s,\rightarrow)\big\} . \tag{6.3} \end{align} $$
Equation 6: The Bellman equation with unreliable actions. Each action produces a distribution over next states, so the value of the next state is replaced by its expectation — 90 percent of the intended cell, 10 percent of the other one.

Solving the chain with missteps

With the optimal actions being left at 2, 3, 4 and right at 5, the four unknown values satisfy a small linear system:

$$ \begin{aligned} V_2 &= 0.5\left(0.9 \times 100 + 0.1 V_3\right) = 45 + 0.05 V_3, & V_3 &= 0.5\left(0.9 V_2 + 0.1 V_4\right) = 0.45 V_2 + 0.05 V_4, \[4pt] V_4 &= 0.5\left(0.9 V_3 + 0.1 V_5\right) = 0.45 V_3 + 0.05 V_5, & V_5 &= 0.5\left(0.9 \times 40 + 0.1 V_4\right) = 18 + 0.05 V_4 . \end{aligned} $$

Solving gives the values, and from them every $Q$:

State$Q(s,\leftarrow)$$Q(s,\rightarrow)$$V(s)$greedymargin
1100100100terminal
246.0614.5646.06$\leftarrow$31.50
321.257.0221.25$\leftarrow$14.23
410.499.4010.49$\leftarrow$1.09
56.7218.5218.52$\rightarrow$11.80
6404040terminal
Table 5: The same chain with a 10 percent chance of slipping the wrong way, $\gamma = 0.5$. Every value has fallen — 50 becomes 46.06, 25 becomes 21.25 — because some of the probability now leaks toward the wrong crater. The policy is unchanged, but the margin column shows how much less confident it is.

What the numbers say

State 4 is where the interesting thing happens. In the reliable world the choice was $12.5$ against $10$ — a margin of $2.5$, or a fifth of the winning value. With a ten percent misstep rate it is $10.49$ against $9.40$, a margin of $1.09$: the decision has become nearly a coin flip. Unreliability hurts the long journey to the +100 more than the short one to the +40, because the long path has three chances to slip instead of two.

Sweeping the misstep rate rather than fixing it at ten percent shows where that is heading:

Figure 6: The two action values at state 4 as the misstep rate runs from perfectly reliable to a coin flip, each computed by value iteration so that the behaviour everywhere else is optimal for that $p$ too. Both fall, but the left value falls faster, and at $p = 0.2386$ they cross: past that point the optimal action at state 4 is right, and the rover abandons the +100 for the +40. At $p = 0.5$ the two meet again at 4.785, because an action that goes the wrong way half the time carries no information at all.

That is a general lesson, not an artefact of this toy. Adding noise to the dynamics systematically favours the near, safe reward over the distant, larger one — and it does so without changing a single reward value. What Figure 6 adds is that the effect has a threshold: the preference does not merely weaken, it eventually reverses. A rover reliable enough is worth sending across the map for the big prize; the same rover with a one-in-four chance of stepping wrong should take the small one nearby.


Where this goes next

Everything above assumed the state was one of six numbered cells, so $Q$ could be written as a table with twelve entries. That assumption dies immediately in any real problem:

All three are the subject of the next note, which builds a deep Q-network for the lunar lander.

Backwards from here: linear regression has the gradient descent that fits the network there, and model development has the diagnostics that still apply once the network exists.

Machine Learning Notes · reinforcement learning 1 of 2 · Foundations ← Collaborative filtering Deep Q-learning →