Tree-Backup Algorithm

Why Does This Exist?

In off-policy reinforcement learning, an agent collects experience by behaving according to an exploratory behavior policy $b(a \mid s)$, while striving to learn the optimal action-value function $Q^*$ or evaluate a distinct target policy $\pi(a \mid s)$.

When extending off-policy methods to multi-step horizons, classical algorithms encounter severe mathematical roadblocks:

1. Importance-Sampled $n$-Step SARSA explodes in variance: Standard multi-step off-policy SARSA corrects for the mismatch between behavior and target distributions by multiplying importance sampling ratios along the trajectory: $$\rho_{t:t+n-1} = \prod_{k=t}^{t+n-1} \frac{\pi(A_k \mid S_k)}{b(A_k \mid S_k)}$$ When the target policy $\pi$ differs significantly from $b$, these chained product ratios either collapse to zero (discarding informative multi-step experience) or explode to massive values (inducing devastating sample variance that causes value estimates to diverge).
2. Watkins's $Q(\lambda)$ prematurely terminates traces: Watkins's $Q(\lambda)$ resolves the off-policy dilemma by zeroing out eligibility traces whenever the agent executes an exploratory (non-greedy) action. In environments requiring continuous exploration, traces are severed after only one or two steps, effectively degrading the multi-step algorithm back into slow, single-step Q-learning.

The Tree-Backup algorithm (introduced by Precup, Sutton, and Singh) solves this fundamental dilemma by achieving multi-step off-policy learning with zero importance sampling ratios.

Instead of adjusting an entire trajectory with multiplicative weights after the fact, Tree-Backup views credit assignment as an explicit decision tree. At each visited state along the trajectory, untaken candidate actions are evaluated analytically using their expected value under the target policy $\sum_{a \neq A} \pi(a \mid s) Q(s, a)$. Meanwhile, the single action that was actually taken continues down the spine to sample future transitions.

By taking exact expectations over unchosen actions, Tree-Backup eliminates importance sampling ratios entirely, providing stable, bounded multi-step credit assignment across arbitrary behavior policies.

Think of It Like This

How It Actually Works

The Tree-Backup Return Recursion and Branching Topology

In the Tree-Backup algorithm, updates are computed for state-action pairs $(S_t, A_t)$. The backup diagram forms a tree where each visited state branches into all available actions.

Consider an agent observing a trajectory:

generated by behavior policy $b(a \mid s)$. The agent evaluates this sequence under target policy $\pi(a \mid s)$.

1Level 0:  (S_t, A_t) ──► Root action node being updated2              │3              ▼ R_{t+1}4Level 1:     S_{t+1}5            ┌──┴─────────────────────────┐6            │                            │7   Untaken a ≠ A_{t+1}          Chosen A_{t+1} (The Spine)8   Bootstrapped via Q           Sampled deeper into environment9   Weight: π(a | S_{t+1})       Weight: π(A_{t+1} | S_{t+1})10            │                            │11            ▼                            ▼ R_{t+2}12   [Terminates branch]                  S_{t+2}13                                ┌────────┴────────┐14                                ▼                 ▼15                             Leaf a_0          Leaf a_116                             Expected value under target policy π

1-Step Base Horizon (Expected SARSA)

For a horizon of $n = 1$, the tree backup return from $(S_t, A_t)$ is identical to the target in Expected SARSA:

where:

• $R_{t+1}$ is the immediate observed reward.
• $\gamma \in [0, 1]$ is the discount factor.
• $\pi(a \mid S_{t+1})$ is the probability of selecting action $a$ in state $S_{t+1}$ under the target policy.
• $Q_t(S_{t+1}, a)$ is the current estimated action-value.

Multi-Step Tree-Backup Return ($n \ge 2$)

For horizons $n \ge 2$, the tree backup return $G_{t:t+n}$ is defined recursively:

for $t < T - 1$, where:

• $\sum_{a \neq A_{t+1}} \pi(a \mid S_{t+1}) Q_{t+n-1}(S_{t+1}, a)$ is the expected value of all candidate actions that were not selected at step $t+1$.
• $\pi(A_{t+1} \mid S_{t+1}) G_{t+1:t+n}$ is the contribution of the single action that was selected, continuing recursively down the sampled trajectory.
• The leaf step at horizon boundary $t+n-1$ terminates using the complete expected value: $$G_{t+n-1:t+n} \doteq R_{t+n} + \gamma \sum_{a \in \mathcal{A}} \pi(a \mid S_{t+n}) Q_{t+n-1}(S_{t+n}, a)$$ (if $S_{t+n}$ is a terminal absorbing state, the second summation is identically zero).

Tabular Action-Value Update

Once $n$ steps of experience have been observed, the Q-table entry for the root state-action pair $(S_t, A_t)$ is updated:

where $\alpha \in (0, 1]$ is the step-size learning rate parameter.

Why Importance Sampling is Completely Eliminated

In standard off-policy SARSA, the agent must correct for sampling action $A_{t+1}$ from behavior policy $b$ rather than target policy $\pi$, requiring multiplication by the importance sampling ratio $\rho_{t+1} = \frac{\pi(A_{t+1} \mid S_{t+1})}{b(A_{t+1} \mid S_{t+1})}$.

In Tree-Backup, no importance sampling ratio is needed because the expectation over candidate actions is computed explicitly:

1. Untaken actions are multiplied directly by their exact target policy weights $\pi(a \mid S_{t+1})$.
2. The taken action's continuation $G_{t+1:t+n}$ is scaled directly by its target policy weight $\pi(A_{t+1} \mid S_{t+1})$.

Because every branch is scaled by $\pi$ and all branches sum to $\sum_{a} \pi(a \mid S_{t+1}) = 1$, the distribution is mathematically exact.

The Unified Spectrum: $Q(\sigma)$

The Tree-Backup algorithm serves as an anchor in the broader unification of multi-step reinforcement learning. Asis et al. (2017) demonstrated that multi-step SARSA (which fully samples the next action, $\sigma = 1$) and Tree-Backup (which computes expectations over candidate actions, $\sigma = 0$) can be unified into a continuous spectrum known as $Q(\sigma)$, allowing practitioners to dynamically interpolate between sampling and expectation at every step.

Worked numerical calculation

We compute a concrete 2-step Tree-Backup return ($n = 2$) for state-action pair $(S_0, A_0)$.

Trajectory observed:

Parameters:

• Discount factor: $\gamma = 0.9$
• Learning rate: $\alpha = 0.2$
• Initial root value: $Q(S_0, A_0) = 4.0000$

Available actions at all states: $\mathcal{A} = \{a_0, a_1\}$.

State $S_1$ Dynamics:

• Taken action: $A_1 = a_0$
• Untaken action: $a_1$
• Target policy: $\pi(a_0 \mid S_1) = 0.60$, $\pi(a_1 \mid S_1) = 0.40$
• Estimated Q-values: $Q(S_1, a_0) = 5.0000$, $Q(S_1, a_1) = 3.0000$

State $S_2$ Leaf Dynamics:

• Target policy: $\pi(a_0 \mid S_2) = 0.70$, $\pi(a_1 \mid S_2) = 0.30$
• Estimated Q-values: $Q(S_2, a_0) = 8.0000$, $Q(S_2, a_1) = 6.0000$

Step 1: Compute Expected Value at Leaf State $S_2$

At leaf state $S_2$, we take the expectation across all candidate actions under target policy $\pi$:

Step 2: Compute 1-Step Return $G_{1:2}$ from $(S_1, A_1)$

Discounting the leaf value and adding observed reward $R_2 = 4.0$:

Step 3: Compute Expected Value of Untaken Branches at $S_1$

The action taken at $S_1$ was $A_1 = a_0$. The untaken action is $a_1$:

Step 4: Weight the Chosen Spine Action at $S_1$

The recursive multi-step return $G_{1:2}$ along the chosen branch is weighted by its target probability:

Combining untaken and chosen branches at $S_1$:

Step 5: Compute Full 2-Step Tree-Backup Return $G_{0:2}$

Adding immediate reward $R_1 = 2.0$ and discounting the combined branches:

Step 6: In-Place Q-Value Update for $(S_0, A_0)$

With learning rate $\alpha = 0.2$ and prior estimate $Q(S_0, A_0) = 4.0000$:

Every step of the computation operates with bounded values without any probability ratios.

Code

1from typing import Dict, List, Tuple2
3
4def compute_tree_backup_return(5    rewards: List[float],6    states: List[str],7    actions: List[str],8    action_space: List[str],9    policy: Dict[str, Dict[str, float]],10    q_table: Dict[Tuple[str, str], float],11    gamma: float,12    start_t: int,13    n: int,14) -> float:15    """Compute the n-step Tree-Backup return G_{t:t+n} for starting pair (S_t, A_t).16
17    Args:18        rewards: Sequence of observed rewards [R_{t+1}, R_{t+2}, ...].19        states: Sequence of visited states [S_t, S_{t+1}, ..., S_{t+n}].20        actions: Sequence of taken actions [A_t, A_{t+1}, ..., A_{t+n-1}].21        action_space: Universal discrete action space.22        policy: Mapping state -> {action -> probability under target policy pi}.23        q_table: Action-value table mapping (state, action) -> estimated Q-value.24        gamma: Discount factor in [0, 1].25        start_t: Origin time index t.26        n: Number of backup steps (horizon).27
28    Returns:29        Calculated scalar multi-step tree-backup return G_{t:t+n}.30    """31    horizon = min(len(rewards), start_t + n)32
33    # 1. Base case: leaf state S_{horizon} expected value over all candidate actions34    leaf_state = states[horizon]35    expected_leaf_val = sum(36        policy[leaf_state][a] * q_table.get((leaf_state, a), 0.0)37        for a in action_space38    )39
40    # Return at the final leaf level41    current_g = rewards[horizon - 1] + gamma * expected_leaf_val42
43    # 2. Fold backward from horizon - 2 down to start_t44    for k in range(horizon - 2, start_t - 1, -1):45        s_kplus1 = states[k + 1]46        a_kplus1 = actions[k + 1]47
48        # Sum expected values of all untaken actions at S_{k+1}49        untaken_expected_sum = sum(50            policy[s_kplus1][a] * q_table.get((s_kplus1, a), 0.0)51            for a in action_space52            if a != a_kplus153        )54
55        # Chosen action branch weighted by target policy probability pi(A_{k+1} | S_{k+1})56        taken_prob = policy[s_kplus1][a_kplus1]57        taken_branch_val = taken_prob * current_g58
59        # Combine and discount back to step k60        current_g = rewards[k] + gamma * (untaken_expected_sum + taken_branch_val)61
62    return current_g63
64
65# Verification of 2-step Tree Backup matching the numerical walkthrough66rewards_seq = [2.0, 4.0]67states_seq = ["S0", "S1", "S2"]68actions_seq = ["a0", "a0"]69actions_all = ["a0", "a1"]70
71target_pi = {72    "S0": {"a0": 0.5, "a1": 0.5},73    "S1": {"a0": 0.6, "a1": 0.4},74    "S2": {"a0": 0.7, "a1": 0.3},75}76
77q_values = {78    ("S0", "a0"): 4.0,79    ("S0", "a1"): 2.0,80    ("S1", "a0"): 5.0,81    ("S1", "a1"): 3.0,82    ("S2", "a0"): 8.0,83    ("S2", "a1"): 6.0,84}85
86# Compute 2-step Tree-Backup return G_{0:2}87ret_2step = compute_tree_backup_return(88    rewards=rewards_seq,89    states=states_seq,90    actions=actions_seq,91    action_space=actions_all,92    policy=target_pi,93    q_table=q_values,94    gamma=0.9,95    start_t=0,96    n=2,97)98
99assert abs(ret_2step - 8.8364) < 1e-4100
101# Apply in-place tabular Q-value update102alpha = 0.2103old_q = q_values[("S0", "a0")]104new_q = old_q + alpha * (ret_2step - old_q)105assert abs(new_q - 4.9673) < 1e-4106
107print(f"2-Step Tree-Backup Return G_0:2: {ret_2step:.4f}")108print(f"Updated Q(S0, a0): {new_q:.4f}")109
110# -> 2-Step Tree-Backup Return G_0:2: 8.8364111# -> Updated Q(S0, a0): 4.9673

Watch Out For

The Quick Version

• The Tree-Backup algorithm enables multi-step off-policy reinforcement learning without using importance sampling ratios.
• At each step along an exploratory trajectory, untaken candidate actions are backed up analytically by their target policy expectation $\sum_{a \neq A} \pi(a \mid s) Q(s, a)$.
• The single chosen action continues down the spine, weighted by its target policy probability $\pi(A \mid s)$ times the recursive return.
• Tree-Backup eliminates importance sampling variance explosion and forms the expectation anchor ($\sigma = 0$) of the unified $Q(\sigma)$ multi-step continuum.

What to Read Next

• Watkins's Q(lambda)
• n-Step TD Prediction
• Expected SARSA