Model-based RL: R-Max and Dyna

1Models and planning

Models come in two kinds, and either can stand in for the environment and produce simulated experience. The dynamics p(s′, r | s, a) that dynamic programming assumes are a distribution model:

A distribution model can always produce samples, so it is the stronger of the two. But a sample model is often much easier to build: a night of ten trips to the charger has 2¹⁰ = 1,024 sequences of free and taken, and a distribution model must give the probability of each, while a sample model just plays the ten trips.

The planning in these notes goes through values, as learning does:

Both look ahead at what follows a state or an action, and back the result up into a value. So a learning method can usually plan too: it only needs simulated experience in place of real.

Value iteration

With a distribution model, the agent can solve the Bellman optimality equation of the Q-learning note, which an agent without a model cannot. Written for the values of states, the equation becomes an update:

Value iteration applies it to every state, sweep after sweep, until no value changes. Each sweep gives every state the most its best action can be expected to bring, judged by the current values of what follows, so news of a reward travels back one step per sweep. The values converge to v_*(s) = max_a q_*(s, a), the optimal values of the model, and acting greedily with respect to them is the best behavior in the model. Value iteration is a dynamic programming method.

On the robot's model, leaving the untried door aside, each room has one move, so each maximum has a single term. Here is value iteration on it, starting from 0, each sweep using the values of the sweep before:

The study is worth 0.6 · 1 + 0.4 · 0 = 0.6 from the first sweep, and the news reaches the hallway one sweep later, −1 + 0.6 = −0.4; the third sweep changes nothing. These are the values batch TD(0) converged to in the TD note, the certainty-equivalence estimate, computed here straight from the model.

Q-planning

With a sample model, the agent can plan by learning from imagined steps. Random-sample one-step tabular Q-planning repeats three steps:

1. Pick a state S and an action A at random.
2. Ask the sample model for a reward R and a next state S′.
3. Apply the Q-learning update to them: Q(S, A) ← Q(S, A) + α[R + γ max_a Q(S′, a) − Q(S, A)].

It converges to the optimal policy for the model under the same conditions as Q-learning does for the environment: every pair keeps being picked, and the step size shrinks appropriately. Asked five times about charging in the study, the sample model might answer +1, 0, +1, +1 and 0. With α = 1/n at the n-th answer, Q-planning's estimate is the running average of the answers, and it ends at 0.6, the value the distribution model gives at once:

The figure runs Q-planning on the cookie grid of the Monte Carlo note, with a model that knows every move, as if the agent had been told the rules of the grid. A move there always does the same thing, so there is nothing to average: the figure takes α = 1, and each step sets the value it updates to its target.

• All 32 values start at 0 and show, since the model holds every move.
• The agent stays in A1, and a faint copy of it plays each imagined step: it appears in the state picked, with the move in yellow, and goes where the model says the move leads.
• The update is outlined in green, toward its target in red, as in the Q-learning figure.
• After the first step, each press is a whole planning step.

Each Q-planning step is small, so planning can be interrupted or redirected at any moment with little computation wasted. That matters when planning is mixed with acting, as in Dyna, in section 3.

2R-Max: optimism in the face of the unknown

R-Max learns its model by counting. It keeps N(s, a), the number of times it has taken action a in state s, and N(s, a, s′, r), the number of those times that led to s′ with reward r, and it estimates the dynamics by the fraction of times each outcome happened:

The robot's model was made this way: the study's charger was free in 3 visits of 5, so P(end, +1 | study, charge) = 3/5. It is the maximum-likelihood model of the TD note. R-Max plans on it with value iteration, at every step, and acts greedily on the values it gets. It uses its experience only to improve the model; its values come from planning.

Counting leaves a question: what to do with a pair tried only a few times, or never? The robot has N(hallway, garage) = 0, so no fraction exists, and a planner that left the door out would never go and look. R-Max answers with optimism in the face of the unknown:

In the home, R_max = 1, a free charger, so value iteration values the garage door at 1, against −1 + 0.6 = −0.4 through the study, and the robot goes to look. In the figure, the door is at the end of a corridor whose moves are known, each costing 1; drag it farther away.

Unknown pairs too far away, or behind too many costly moves, are left alone. Optimism makes exploring a matter of planning: the values send the agent to the unknown by the shortest way, and only when it might pay.

Say the robot, looking, finds a second charger in the garage, free every time. Once both are known, the garage is worth −1 + 1 = 0 from the hallway, more than the study, and the robot charges there from then on, while a planner that left the door out would still go to the study. Eventually every pair worth knowing is known, and R-Max follows the best policy of its model.

The number k sets how much each estimate rests on. Where an action always has the same result, as on the grid below, one visit tells everything, and k = 1 is enough. Where results vary, a few visits can mislead:

1. A charger free 3 times in 5 is taken on a single visit 2 times in 5.
2. With k = 1, the model then says it is never free.
3. The pair is known, so nothing sends the robot back to check, and it may settle for a worse charger.

The more visits k asks for, the less likely they all find it taken:

R-Max, introduced by Brafman and Tennenholtz in 2002, converges to a near-optimal policy, with a gap that depends on k, in a number of steps that grows only polynomially with the size of the problem.

The algorithm

For a number of visits k and a reward R_max at least as large as any reward the world can give:

1. Set every count to 0: N(s, a) ← 0 and N(s, a, s′, r) ← 0.
2. At the start of each episode, observe the first state S.
3. Run value iteration on the model, in which a known pair, with N(s, a) ≥ k, leads to s′ with reward r with probability N(s, a, s′, r)/N(s, a), and an unknown pair gives R_max and ends the episode.
4. Choose A greedily from V: the action with the largest Q(S, a) = Σ_{s′, r} P(s′, r | S, a)[r + γV(s′)], which is R_max for an unknown action.
5. Take A, and observe R and S′.
6. Count it: N(S, A) ← N(S, A) + 1 and N(S, A, S′, R) ← N(S, A, S′, R) + 1.
7. Set S ← S′, and go back to step 3 until S is terminal.

The figure runs R-Max on the cookie grid, with k = 1, R_max = 1, the cookie's reward, and γ = 0.9. Each triangle holds the value of a move, Q(s, a) as step 4 computes it:

• an untried move is shaded gray and holds the 1 that R-Max assumes for it;
• a known move holds its reward plus 0.9 times the value of the cell it leads to.

Each press is one line of the algorithm: the agent chooses a move, in yellow, takes it, counts it, in green, and value iteration runs on the new model, which lifts the shade.

In the figure, the model changes only when a move becomes known, so at all the other steps value iteration gives back the values it already had. Running it at every step is the simplest way to write R-Max, not the cheapest: an efficient version plans again only when the model changes.

3Dyna: learning, planning and acting together

An agent that plans while it acts can use each real step in two ways: to improve its model, so that the model matches the world more closely, and to improve its values and policy directly, with the learning methods of the earlier notes. The first is model learning, and the second direct reinforcement learning, or direct RL. Through the model, experience also improves the values indirectly, by planning.

Each way of using experience has its strengths:

• planning gets more out of a limited amount of experience, so it reaches a good policy with fewer real steps;
• direct learning is simpler, and a wrong model cannot mislead it.

Dyna does all of it at once: acting, model learning, direct RL and planning. Its planning applies the same learning method to simulated steps as to real ones, so learning and planning share almost everything except where their experience comes from. The agent is always reactive and always deliberative: it acts at once on its latest values, while planning goes on in the background. When the model learns something new, the planning that follows works it into the values.

Dyna-Q, the simplest Dyna, is one-step tabular Q-learning, its direct RL, with two parts added:

• a model, a table that keeps, for each pair the agent has tried, the reward and next state last seen after it, which is all there is to know in a world where an action always has the same result;
• planning, random-sample one-step Q-planning, n steps after every real step. Its search control, the rule for which simulated steps to take, picks states the agent has been in and actions it has taken there, so the model is never asked about a pair it knows nothing about.

The algorithm

Tabular Dyna-Q, for a step size α in (0, 1], a small ε > 0 and a number of planning steps n, with every Q(s, a) at 0 and the model empty:

1. Observe the current state S, and choose A ε-greedily with respect to Q.
2. Take A, and observe the reward R and the next state S′.
3. Update Q(S, A) ← Q(S, A) + α[R + γ max_a Q(S′, a) − Q(S, A)].
4. Remember the outcome: Model(S, A) ← R, S′.
5. Plan, n times: pick at random a state seen before and an action taken there before, get R and S′ from Model, and apply the update of step 3 to them.
6. If S′ is terminal, a new episode starts at its first state. Go back to step 1.

The figure runs Dyna-Q on the cookie grid with α = 0.5, ε = 0.1, γ = 0.9 and n = 10. Its marks are those of the Q-learning figure, and its planning steps look as in the figure of section 1. As there, the triangles with a number are the moves the model holds, 0 included, but this model starts empty: a move shows its value once the agent has tried it. The move it has just remembered is outlined in blue.

After the figure's first episode, methods differ in how much of the episode they credit:

Planning carries the cookie's value back through every remembered move, and pieces of different walks join into routes never walked in one go. With more planning steps per move, the values settle in fewer episodes, at the price of more computation per move.

R-Max and Dyna

The two figures show the difference in exploring. R-Max tried all 32 moves of the grid in its first three episodes, and then never explored again. Dyna-Q tried 15 of them in 30 episodes, and it keeps making a random move one time in ten.

4When the model is wrong

So far the models have been right: on the grid, a known move does what the model says. A model can be wrong for three reasons:

• the world is random and has been seen only a few times, like a charger visited once;
• the model generalizes from some states to others, imperfectly;
• or the world has changed since the model learned it.

Planning with a wrong model computes a policy that is good for the model, and possibly bad for the world. Some errors correct themselves; others are harder:

• When a model is optimistic, predicting more reward or better moves than the world allows, the planned policy goes to collect them, finds they are not there, and the experience it gets corrects the model. R-Max's optimism in section 2 is such an error, made on purpose: its model promises R_max for every move it has not tried, and trying the move corrects it.
• When the world becomes better than the model believes, the policy the model recommends may never go where the improvement is, and the error can last a long time, or forever.

Back in the robot's home, the world changes both ways. One night the garage door is locked, so the move that always led to the garage now leaves the robot in the hallway, at a cost of 1. And one day a back door in the hallway, which the robot tried once and found locked, is opened, with a charger right behind it; nothing in either model points there.

Dyna-Q's model holds the outcome last seen, and R-Max's counts every outcome since the start, old and new:

The figure puts both methods on the cliff of the TD note, γ = 1: R-Max with k = 1 and R_max = −1, the largest reward there, and Dyna-Q with the settings of section 3's figure. Its first press plays their first 10 episodes quickly, after which both follow the edge. A wall then goes up between A2 and B2, across the edge, before episode 11, and comes down before episode 21. The buttons above the grid switch between the two methods; their marks are those of sections 2 and 3, and the wall is drawn in ink.

5Summary

The key ideas next to their NumPy equivalents. In R-Max, n_sa[s, a] holds N(s, a), and n_sasr[s, a, s2, i] holds N(s, a, s′, r) for the i-th possible reward, rewards[i]; in Dyna-Q, model maps a pair to what followed it, seen lists the states visited, and taken[s] the actions taken in s.