N-Step Return - Reinforcement Learning

N-Step Return

What you'll build

The n-step bootstrapped return G_t^(n): n real rewards, then the critic's estimate of the state reached.

Background

TD(0) bootstraps after one step (low variance, biased by the critic); Monte Carlo waits for the end (unbiased, high variance). n-step returns sit in between and are used by A3C/A2C, Rainbow DQN and as the building block that GAE averages over.

Formula

G_t^(n) = r_t + γ·r_{t+1} + ... + γ^(n-1)·r_{t+n-1} + γ^n · V(s_{t+n})

n	Method	Bias	Variance
1	TD(0)	highest	lowest
n	n-step	medium	medium
∞	Monte Carlo	none	highest

Inputs

• rewards: (T,) rewards r_0 ... r_{T-1}
• values: (T+1,) values V(s_0) ... V(s_T) (V(s_T) is 0 if the episode terminated)
• t: start step, 0 ≤ t < T
• n: number of steps, n ≥ 1
• gamma: discount factor

Expected return format

A Python float: G_t^(n), rounded to 2 decimals.

Edge cases and conventions

• Truncation: if t + n > T, the episode ends first. Sum only the rewards up to r_{T-1}, and add no bootstrap term.
• If t + n = T exactly, bootstrap with values[T] (0 for a terminated episode).
• In code: add γ^n · values[t + n] only when t + n < len(values).

Constraints

NumPy only.

Follow-up

• For which n does G_t^(n) become the Monte Carlo return?
• GAE is an exponentially weighted average of all n-step advantages. Write the weights.

Input: rewards = [1, 2, 3, 4], values = [0, 0, 0, 0, 10], t = 0, n = 2, gamma = 0.9
Output: 2.8

Input: rewards = [1, 1, 1, 1], values = [0, 0, 0, 5, 0], t = 0, n = 3, gamma = 0.9
Output: 6.36

Input: rewards = [1, 2, 3], values = [0, 0, 0, 0], t = 1, n = 5, gamma = 0.5
Output: 3.5

import numpy as np

def n_step_return(rewards: np.ndarray, values: np.ndarray,
                  t: int, n: int, gamma: float) -> float:
    """
    Compute the n-step return from timestep t.

    Args:
        rewards: (T,) rewards [r_0, ..., r_{T-1}]
        values: (T+1,) values [V(s_0), ..., V(s_T)]
        t: starting timestep
        n: number of steps
        gamma: discount factor

    Returns:
        G_t^(n) as a float, rounded to 2 decimals (no bootstrap if the
        episode ends before t + n)
    """
    rewards = np.array(rewards, dtype=float)
    values = np.array(values, dtype=float)
    T = len(rewards)

    # Your code here
    pass