Reinforcement learning in a gridworld

In reinforcement learning (RL) an agent acts in an environment. At each time step it is in a state $s$, takes an action $a$, receives a reward $r$ and moves to a next state $s'$. A policy $\pi$ is a rule that chooses the action in each state. The goal is a policy that collects as much reward as possible over time.

This page has two simulations of the same small world:

• Planning. The agent knows the rules of the world (the model). Value iteration and policy iteration compute the best policy exactly, one step at a time.
• Learning. The agent does not know the rules. It walks through the world, tries actions and learns from the rewards it receives, with Q-learning or SARSA.

Play with both figures first, then do the experiments in section 3. Section 4 explains the math.

1. Planning with a known model #

The value $V(s)$ of a cell is the expected discounted sum of all future rewards when the agent starts in that cell. The action value $Q(s,a)$ is the same, but the first action is fixed to $a$. In the figure the colour of a cell is its value: blue for positive values, magenta for negative ones, stronger for larger size. The orange arrows are the greedy policy: in each cell, the action with the largest $Q(s,a)$. A small dot means that all four actions are equally good.

Interactive figure: loading…

2. Learning from experience #

Now the agent does not know the slip probability or where the rewards are. It only sees what happens when it acts. It keeps a table of action values $Q(s,a)$, one number for each cell and action. The figure draws them as four triangles in each cell. After every step the agent moves one number a little toward a better estimate, the TD target (TD means temporal difference). To find out what works, it sometimes acts at random: with probability $\varepsilon$ it picks a random action, otherwise the greedy one. This rule is called ε-greedy.

Interactive figure: loading…

3. Try this #

1. The policy settles before the values. In the planning figure (world Two exits), press Reset, then press Step many times. Question: what stops changing first, the arrows or the numbers? What to see: the arrows do not change after sweep 11, but value iteration needs 52 sweeps until the largest change is below $10^{-6}$. A greedy policy is often optimal long before the values are exact. Policy iteration (switch the algorithm) needs only 3 evaluations for the same world.
2. The discount changes the destination. Set the slip probability to 0. With $\gamma = 0.9$ the arrows from S go along the bottom row to +10. Now move $\gamma$ down slowly. Question: at which $\gamma$ does the agent prefer the +1 exit? Hint: +1 is 3 moves from S and +10 is 5 moves from S. What to see: the arrows switch between $\gamma = 0.32$ and $\gamma = 0.31$. The reward of the third move is discounted by $\gamma^2$ and the reward of the fifth by $\gamma^4$, and $\gamma^2 \cdot 1 = \gamma^4 \cdot 10$ gives $\gamma = \sqrt{0.1} \approx 0.316$.
3. Slip makes the agent careful. Go back to Two exits (slip 0.2, $\gamma = 0.9$). Question: which way does the agent go from S? What to see: it takes the long way over the top (9 moves) to +10. On the bottom row every move has a 10% chance to slip down into a trap, so the short way is risky. Set the slip to 0 and the short way becomes best. Then open the world Bridge: with slip 0 the agent crosses the bridge to +10, but with slip 0.02 it already prefers the +1 exit next to the start.
4. The step reward. In Two exits (slip 0.2, $\gamma = 0.9$) set the step reward $c$ to −1. Question: what happens to the long safe way? What to see: every move now costs 1, so the long way is too expensive and the agent takes the short way along the traps (5 moves). Now set $c = +0.5$ and $\gamma = 0.99$. The agent avoids both exits: from S it moves up, away from the traps, and in the top row all actions are equally good (dots). Staying in the world forever is worth up to $0.5 / (1 - 0.99) = 50$, more than any exit.
5. Exploration against exploitation. In the learning figure choose Two exits, set the slip to 0 and $\varepsilon = 0$. Press +100 episodes three times. Question: does the agent find the +10 exit? What to see: in most runs it finds +1 first and then goes there every time. With $\varepsilon = 0$ it never tries anything new, so the triangles near +10 stay pale. Press Reset, set $\varepsilon = 0.5$ and press +100 episodes ten times: now, in most runs, the arrows from S point toward +10. The agent must explore (try actions that look worse) to find better rewards, and exploit (use what it knows) to collect them.
6. Q-learning or SARSA on the cliff. In the world Cliff ($\varepsilon = 0.1$) press +100 episodes five times with Q-learning and look at the arrows and the chart. Then choose SARSA and do the same. Question: which algorithm learns the path right along the cliff, and which one collects more reward while it learns? What to see: Q-learning learns the shortest path, next to the cliff, but its random moves sometimes throw it into the cliff, so many of its returns are far below the dashed line and its mean return is low. SARSA learns a path one or two rows away from the cliff and has a higher mean return. Finally choose Q-learning again, train it, set $\varepsilon = 0$ and press +1 episode: without random moves its return equals $V^*(s_0) \approx -2.64$.

4. The math behind the figures #

Returns and values #

The return is the discounted sum of the rewards from time $t$ on:

Here $r_t$ is the reward received after the action $a_t$ (Sutton and Barto write $R_{t+1}$ for the same reward). The values of a policy $\pi$ are expected returns:

The expectation is over the random slips and over the policy's own choices. In this world the reward of one move is $r(s,a,s') = c + R(s')$, where $R(s')$ is the reward of the cell the agent enters (0 for a free cell), and $P(s' \mid s,a)$ is $1-p$ for the chosen direction and $p/2$ for each sideways one.

The Bellman equations #

Use $G_t = r_t + \gamma G_{t+1}$ inside the expectation. The value of a state is the expected reward of one step plus the discounted value of the state where the agent lands. This is the Bellman expectation equation:

The best possible values $V^*$ and $Q^*$ satisfy the Bellman optimality equations, where the average over the policy becomes a maximum over actions:

The greedy policy with respect to $Q^*$, $\pi^*(s) = \arg\max_a Q^*(s,a)$, is optimal. The inspector of the first figure shows exactly one such backup: the sum over the (at most three) cells where the agent can land.

Value iteration #

Value iteration turns the optimality equation into an update. Start with $V_0 = 0$ and apply to every cell at once (one sweep):

The Bellman optimality operator $T$ is a contraction in the max norm: for any two value tables, $\|TV - TU\|_\infty \le \gamma \|V - U\|_\infty$. So every sweep shrinks the error at least by the factor $\gamma$: $\|V_k - V^*\|_\infty \le \gamma^k \|V_0 - V^*\|_\infty$. It also gives a stopping rule. If the largest change of a sweep is $\delta = \|V_{k+1} - V_k\|_\infty$, then

This is the readout "distance to V* at most". In the chart, the largest change of each sweep stays below the dashed line $\delta_1 \gamma^{k-1}$, where $\delta_1$ is the change of the first sweep.

Policy iteration #

Policy iteration starts from any policy and repeats two steps:

1. Policy evaluation. For a fixed deterministic policy the Bellman expectation equation is linear: $V^\pi = r_\pi + \gamma P_\pi V^\pi$, where $P_\pi$ is the matrix of transition probabilities and $r_\pi$ the vector of expected rewards when the agent follows $\pi$. So $V^\pi = (I - \gamma P_\pi)^{-1} r_\pi$. The matrix $I - \gamma P_\pi$ is invertible for $\gamma < 1$ because every row of $P_\pi$ sums to at most 1. The figure solves this system with Gaussian elimination.
2. Policy improvement. Make the policy greedy: $\pi'(s) = \arg\max_a Q^\pi(s,a)$. The policy improvement theorem says that then $V^{\pi'}(s) \ge V^\pi(s)$ in every state.

There are only finitely many deterministic policies. Each improvement that changes an action makes the policy strictly better in at least one state, so no policy comes back and the loop stops. (This needs a tie rule: keep the old action if it is still among the best. The figure does this.) When no action changes, $\pi$ is greedy with respect to its own values, so $V^\pi$ satisfies the Bellman optimality equation and $\pi$ is optimal. Policy iteration needs few iterations, but each one solves a linear system. Value iteration does cheap sweeps, but many of them.

Q-learning and SARSA #

Without a model the agent cannot compute the sum over $s'$. It uses the one next state it really sees instead. After the step $(s, a, r, s')$ it computes a TD target and moves $Q(s,a)$ toward it with the learning rate $\alpha$:

The bracket is the TD error $\delta$. If $s'$ is terminal, the target is only $r$. On average over $s'$, the Q-learning target is the right-hand side of the optimality equation for $Q^*$, with the current $Q$ in place of $Q^*$. So Q-learning is a sampled version of value iteration on action values.

• Q-learning is off-policy. Its target uses the best next action, whatever the agent does next. It learns $Q^*$, the values of the greedy policy, even while it explores. (The guarantee needs every state and action to be tried infinitely often and step sizes with $\sum_t \alpha_t = \infty$ and $\sum_t \alpha_t^2 < \infty$. With a constant $\alpha$, as in the figure, the values keep moving a little in a world with slip.)
• SARSA is on-policy. Its target uses the action $a'$ that the agent really takes next (the name comes from the five quantities $s, a, r, s', a'$). It learns the values of the ε-greedy policy it follows, random moves included. That is why it keeps away from the cliff when $\varepsilon > 0$.

Exploration and exploitation #

To exploit means to take the action that looks best now. To explore means to try another action to learn more about it. Pure exploitation ($\varepsilon = 0$) can lock the agent onto the first reward it finds (experiment 5). Too much exploration wastes reward on random moves (experiment 6). ε-greedy is the simplest compromise. Other methods, such as optimistic initial values or upper confidence bounds (UCB), explore in a more directed way.

Why the discount changes the policy #

A reward that comes $k$ steps in the future is multiplied by $\gamma^k$. A small $\gamma$ makes the agent short-sighted; a $\gamma$ near 1 makes it patient. A rough measure is the effective horizon $1/(1-\gamma)$: about 10 steps for $\gamma = 0.9$ and 100 steps for $\gamma = 0.99$. The discount also interacts with the other parts of the task:

• Distance: a far, large reward beats a near, small one only if $\gamma$ is large enough (experiment 2).
• Risk: with slip, a short path next to a trap risks a large penalty soon, while the long safe path delays the reward. The discount sets the price of that delay (experiment 3).
• Step reward: with $c < 0$ every extra step costs; with $c > 0$ every step pays, and for $\gamma$ near 1 the endless sum $c/(1-\gamma)$ can be worth more than any exit (experiment 4).

The discount is part of the problem definition. Changing $\gamma$ changes which policy is optimal, not only how fast the algorithms converge.

5. Related lessons #

• Reinforcement Learning: the full track. Its lessons on MDPs and the Bellman equations, dynamic programming, bandits and exploration, and TD control with SARSA and Q-learning go deeper into everything on this page.
• Inverse, determinant and projection: policy evaluation solves the linear system $(I - \gamma P_\pi) V = r_\pi$.
• Norms of vectors and matrices: the max norm $\|\cdot\|_\infty$ behind the contraction argument.
• SGD and batch size: the Q-learning update is a stochastic approximation step with step size $\alpha$, like a step of SGD.
• Probability and Statistics: expectations and conditional expectations, the language of values and Bellman equations.