Full Return-to-Go Sequence

Difficulty: Medium · Concept: RL Fundamentals — Markov Decision Processes (MDPs)

Discounted Return computed a single number, $G_0$, from the start of an episode. But every value function needs $G_t$ for every timestep $t$ in a trajectory, not just the first — training a value estimator means comparing its prediction at each state against the actual return-to-go from that point onward.

The naive approach re-runs the whole sum from scratch at every $t$ — $O(n)$ work per timestep, $O(n^2)$ total across the trajectory. There's a much cheaper way, using the same recursive identity noted at the end of the original problem: $G_t = r_t + \gamma G_{t+1}$. Compute backward from the end of the episode, and each $G_t$ costs $O(1)$ given $G_{t+1}$ — no re-summing.

Your task: implement returns_to_go(rewards, gamma) in $O(n)$ time, returning a list the same length as rewards, where entry $t$ is $G_t$.

Next: Design a Min Stack