7.Unifying Learning & Planning: n-step, Traces, Dyna & MCTS

7.1 The space between one step and all of them

Rusty delivers a package and receives $+25$. Which of the forty decisions that got him there deserves credit?

TD(0) credits exactly one: the state immediately before the reward. The state two steps back learns nothing this episode — it must wait until its successor's value improves, then inherit a little. Monte Carlo credits all forty at once, but every one of them absorbs the noise of the entire trajectory.

The $n$-step return interpolates by looking ahead $n$ steps before bootstrapping:

$G_{t:t+n} \doteq R_{t+1} + \gamma R_{t+2} + \cdots + \gamma^{n-1} R_{t+n} + \gamma^n V(S_{t+n}).$

At $n=1$ this is the TD target; at $n = \infty$ (or the episode end) it is the Monte Carlo return. The update is the usual one, $V(S_t) \leftarrow V(S_t) + \alpha[G_{t:t+n} - V(S_t)]$.

Intermediate $n$ is not a compromise — it is usually strictly better than either extreme, and there is a theorem saying why.

The expected $n$-step return is a better estimate of $v_\pi$ than $V$ itself, by a factor $\gamma^n$. Larger $n$ contracts the bias faster — while adding the variance of $n$ extra sampled rewards. That is the whole trade, and it has no universal optimum.

7.2 The λ-return: averaging all of them

Rather than choosing one $n$, average them all, with geometrically decaying weights:

$G_t^\lambda \doteq (1-\lambda) \sum_{n=1}^{\infty} \lambda^{n-1} G_{t:t+n}.$

The factor $(1-\lambda)$ normalizes the weights to sum to one. At $\lambda = 0$ only the $n=1$ term survives — TD(0). At $\lambda = 1$ all the weight slides onto the full return — Monte Carlo. In between, a smooth blend.

Drag the dial and watch both panels. On the left, the weight distribution over $n$-step returns; on the right, the eligibility trace decaying backwards along the corridor Rusty just drove. Those are two views of the same number.

7.3 Forward and backward views

The definition above is the forward view: to update $V(S_t)$ you look ahead over future returns. It is conceptually clean and computationally useless online — you would have to wait for the episode to finish.

The backward view achieves the same thing causally. Maintain an eligibility trace $E_t(s)$ for every state, decaying by $\gamma\lambda$ each step and incremented when the state is visited:

$E_t(s) = \gamma\lambda\, E_{t-1}(s) + \mathbb{1}[S_t = s].$

Then at every step, compute the one-step TD error $\delta_t$ and apply it to every state in proportion to its trace:

$V(s) \leftarrow V(s) + \alpha\, \delta_t\, E_t(s) \quad \text{for all } s.$

The trace is a short-term memory of who is responsible for what happens now. A cell visited recently has a large trace and receives most of the credit; a cell visited long ago has a small one.

7.4 Dyna: planning with a learned model

Now the second unification. Chapter 5 planned with a given model; Chapter 6 learned with no model. Dyna does both: learn a model from real experience, then use it to generate simulated experience and learn from that too.

The architecture is small enough to state completely. After each real transition $(S, A, R, S')$:

1. Direct RL — apply a Q-learning update from the real transition.
2. Model learning — store $\text{Model}(S,A) \leftarrow (R, S')$.
3. Planning — repeat $n$ times: sample a previously-seen $(S,A)$, look up the model's prediction, and apply the same Q-learning update to that imagined transition.

The planning updates are indistinguishable from real ones as far as the learner is concerned. That is the insight in the epigraph: planning and learning differ only in where the experience comes from.

The payoff is sample efficiency, which for a robot is the currency that matters. With $n = 50$ planning steps per real step, Dyna-Q typically reaches good policies in an order of magnitude fewer real interactions than Q-learning — because each expensive real transition is mined fifty times.

Prioritized sweeping improves on uniform sampling of imagined transitions. Instead of picking $(S,A)$ at random, maintain a priority queue ordered by the magnitude of the change a backup would produce, and always work on the largest. When a state's value changes, all its predecessors are enqueued with priority proportional to the change. Computation flows to where it matters, and the speedup over uniform Dyna is often an order of magnitude.

use std::collections::HashMap;
use rl_core::Transition;
use crate::td::QLearning;
 
pub struct DynaQPlus {
    learner: QLearning,
    model: HashMap<(usize, usize), (f64, usize)>,
    /// Steps since each (s, a) was tried for real — drives the bonus.
    last_tried: HashMap<(usize, usize), u64>,
    planning_steps: usize,
    kappa: f64,
    step: u64,
}
 
impl DynaQPlus {
    pub fn observe<R: rand::Rng>(&mut self, s: usize, a: usize, t: &Transition<usize>, rng: &mut R) {
        self.step += 1;
 
        // 1. Direct RL from real experience.
        self.learner.update(s, a, t);
 
        // 2. Model learning.
        self.model.insert((s, a), (t.reward, t.next_state));
        self.last_tried.insert((s, a), self.step);
 
        // 3. Planning — the same update, on remembered transitions.
        let keys: Vec<_> = self.model.keys().copied().collect();
        for _ in 0..self.planning_steps {
            let &(ps, pa) = &keys[rng.gen_range(0..keys.len())];
            let (reward, next) = self.model[&(ps, pa)];
 
            // Dyna-Q+ bonus: reward staleness, so the agent re-checks old beliefs.
            let tau = (self.step - self.last_tried[&(ps, pa)]) as f64;
            let imagined = Transition {
                next_state: next,
                reward: reward + self.kappa * tau.sqrt(),
                done: false,
            };
            self.learner.update(ps, pa, &imagined);
        }
    }
}

7.5 Decision-time planning and MCTS

Everything so far uses planning to improve a stored policy — Sutton and Barto call this background planning. There is another use: plan at the moment of decision, for the state you are actually in, then throw the computation away.

Monte Carlo tree search is the canonical method. From the current state, repeat four steps:

1. Selection — descend the tree using a rule that balances exploitation and exploration. The standard is UCT, which is Chapter 3's UCB applied at each node: $\arg\max_a \left[ Q(s,a) + c\sqrt{\ln N(s)/N(s,a)} \right]$.
2. Expansion — add a child node for an untried action.
3. Simulation — roll out to a terminal state with a cheap default policy.
4. Backpropagation — propagate the outcome up the path, updating $N$ and $Q$.

Given enough iterations, UCT converges to the optimal action at the root. The tree grows asymmetrically, concentrating depth along promising lines — which is exactly how you would want a finite computation budget spent.

For robots, decision-time planning is most familiar as model predictive control: at each control step, optimize a short action sequence against a model, execute the first action, discard the rest, and re-plan. Chapter 12 derives it properly and Chapter 13 gives it its classical-control context.

7.6 Chapter bridge

Part I is complete, and it is worth naming what has been assembled.

We have a formalism (Chapter 4), an exact planning method for when the model is known (Chapter 5), sampling methods for when it is not (Chapter 6), a continuous dial between one-step and full-return credit assignment (this chapter), and an architecture that learns a model and plans with it while continuing to learn (also this chapter). All of it is unified by generalized policy iteration and underwritten by the contraction mathematics of Chapter 2.

All of it also assumes tables. Every method here stores one number per state or per state–action pair, and Chapter 5's arithmetic already showed that a modest robot arm has more states than a sweep can touch before the sun burns out.

Part II abandons tables. We will represent value functions and policies as parameterized functions — linear in features, then deep networks — and discover that the guarantees do not survive the transition intact. Chapter 8 begins with what breaks, because knowing precisely how approximation can diverge is what makes the engineering of Chapters 9 through 12 comprehensible rather than superstitious.

References