Goal-Space Planning with Subgoal Models, reproduced on FourRooms

A tabular, from-scratch reproduction of Lo, Roice, Panahi, Jordan, White, Mihucz, Aminmansour & White (JMLR 2024) with exact subgoal models · gsp v1.0.0 · generated 2026-10-04 17:05 UTC

18 → 3
value-iteration sweeps
104 states vs 5 subgoals; projected values equal v* to 1e-15
1 → 188
(s,a) pairs updated by one episode
Sarsa(0) vs GSP + Sarsa(λ), identical experience
35 → 18
episodes to near-optimal
Sarsa(λ) vs GSP + Sarsa(λ), 30 paired runs
18 → 7
episodes to re-route around lava
Sarsa vs GSP (Both), consistent initialisation

Explorer 1 · Planning over states vs. over subgoals

Sec. 5.3, Eq. 2-3, Fig. 18

Both run the same Bellman optimality backup from v = 0. Coloured cells already hold their final value v*; pale cells are still wrong. Red boxes are the hallway subgoals.

Value iteration, 104 states

Goal-space value iteration, 5 subgoals, projected (Eq. 3)

Explorer 2 · Learning, episode by episode

Sec. 4.1, Fig. 3-4

Same seed, same Sarsa(λ). Colour = the learner's value estimate (for GSP, maxa q̃ + Φ, since q̃ = q − Φ). Arrows: greedy action, white if optimal, red if not. The white line is that episode's path. Hover a cell for numbers.

Sarsa(λ)

GSP + Sarsa(λ)

Animations

gsp animate
planning.gif
learning.gif

Planning cost: states vs. subgoals

Sec. 5.3, Prop. 3, Fig. 18
What this experiment does
Planning cost: value iteration over 104 states vs. over the subgoal graph (Sec. 5.3, 7.3, Fig. 18).

Checks three claims with exact models:
  * goal-space value iteration converges in far fewer sweeps (options jump many steps per backup);
  * the projected values v_g* equal v* when subgoals lie on every optimal route (config "both");
  * the subgoal models are exactly the closed forms r = -(1-gamma^d)/(1-gamma), Gamma = gamma^d.
planning_cost.png
goal_space_graphs.png
goal_space_graph.png
value-iteration sweeps, 104 states18
value-iteration sweeps, 5 subgoals3
max |v_g* - v*| (both)1.8e-15
option models vs closed form6.2e-15

gsp run planning · 0.02 s · 2026-10-03T13:52:17+00:00

Value propagation after one episode

Sec. 4.1, Hypothesis 1, Fig. 3
What this experiment does
Fig. 3 / Hypothesis 1: GSP changes the value of more states with the same experience.

One episode under the uniform random policy (reward 0, +1 on entering the goal) is collected once
and replayed, in order, through four learners starting from q = 0. Same data for all four.
fig3_value_propagation.png
(s,a) updated after one episode: Sarsa(0)1
(s,a) updated after one episode: Sarsa(λ)103
(s,a) updated after one episode: GSP + Sarsa(0)152
(s,a) updated after one episode: GSP + Sarsa(λ)188

gsp run propagation · 0.01 s · 2026-10-03T13:52:18+00:00

Learning speed with and without GSP

Sec. 4.1, Hypothesis 2, Fig. 4
What this experiment does
Fig. 4 / Hypothesis 2: GSP lets a TD learner learn faster (reward -1 per step).

Paired comparison: run k of every learner uses the same random stream.
fig4_learning_curves.png
Sarsa(λ=0.9): episodes to near-optimal35.2 ± 1.5
Sarsa(λ=0.9): first-episode steps597 ± 33
GSP + Sarsa(λ=0.9): episodes to near-optimal17.5 ± 0.5
GSP + Sarsa(λ=0.9): first-episode steps275 ± 18
Sarsa(0): episodes to near-optimal201.0 ± 0.0
Sarsa(0): first-episode steps489 ± 38
GSP + Sarsa(0): episodes to near-optimal10.0 ± 0.4
GSP + Sarsa(0): first-episode steps120 ± 6

gsp run curves · 9.99 s · 2026-10-03T13:52:28+00:00

Subgoal placement and the region of attraction

Sec. 6.3 (remark), App. G.1, Figs. 12, 14, 23
What this experiment does
Subgoal placement and the region of attraction (Sec. 6.3 remark, App. G.1, Figs. 12, 14 top row, 23).

The four configurations differ only in the hallway next to the goal room. The abstract MDP then knows
only some routes, v_g* encodes "the best route through known subgoals", and an epsilon-greedy learner
is pulled towards that route even when it is not the shortest.
subgoal_potentials.png
subgoal_placement.png
Sarsa(λ), no GSP: mean steps, last 50 episodes19.6 ± 0.2
GSP: Both: mean steps, last 50 episodes19.7 ± 0.1
GSP: On Initial: mean steps, last 50 episodes19.9 ± 0.2
GSP: On Alternate: mean steps, last 50 episodes21.8 ± 0.2
GSP: No Near: mean steps, last 50 episodes26.5 ± 0.9

gsp run subgoals · 10.95 s · 2026-10-03T13:52:40+00:00

Adapting to a lava pool

Sec. 6.3, Hypothesis 4, Figs. 13-14
What this experiment does
Adapting to change: a lava pool appears on the optimal route (Sec. 6.3, Hypothesis 4, Figs. 13-14).

Setup, following the paper:
  * The learner starts from the optimal action values of the lava-free maze and is then placed in the
    maze with lava on the cells just below the east hallway (each lava cell: reward -20).
  * Only the reward changed, so option policies and Gamma stay fixed. The reward model is represented
    with successor features, r_gamma(s, g) = psi_g(s) . w, and w (one weight per cell entered) is
    re-estimated online by SGD: w[s'] += eta * (R - w[s']). Whenever w moves, Eq. 1-3 are recomputed.
  * Base learner: tabular Sarsa (lambda = 0), alpha = 0.1, epsilon-greedy 0.1.

Initialisation, run both ways:
  * ``paper``:      every learner starts from q*_old, as the paper states.
  * ``consistent``: GSP learners start from q*_old - v_g*_old. A shaped learner's optimum is
                    q* - Phi(s), not q*, so this is the shaped learner's own optimum for the old maze.
With ``paper`` initialisation the shaping term is a non-zero TD error on every step from the start,
so GSP immediately pushes the learner towards the route the abstract MDP prefers. With ``consistent``
initialisation that push is already "paid for" and only *changes* in v_g* produce TD errors. This may explain the paper's App. G.1 puzzle (On Alternate switches fast although v_g* barely moves).
lava_adaptation.png
lava_value_change.png
[paper init] Sarsa: mean return-142.3 ± 0.3
[paper init] Sarsa: episode when most runs take the S route18
[paper init] GSP: No Near: mean return-40.9 ± 0.1
[paper init] GSP: No Near: episode when most runs take the S routenever
[paper init] GSP: On Initial: mean return-41.1 ± 0.1
[paper init] GSP: On Initial: episode when most runs take the S routenever
[paper init] GSP: On Alternate: mean return-29.8 ± 0.2
[paper init] GSP: On Alternate: episode when most runs take the S route27
[paper init] GSP: Both: mean return-30.1 ± 0.2
[paper init] GSP: Both: episode when most runs take the S route31
[consistent init] Sarsa: mean return-142.3 ± 0.3
[consistent init] Sarsa: episode when most runs take the S route18
[consistent init] GSP: No Near: mean return-139.1 ± 0.2
[consistent init] GSP: No Near: episode when most runs take the S route12
[consistent init] GSP: On Initial: mean return-42.4 ± 0.2
[consistent init] GSP: On Initial: episode when most runs take the S routenever
[consistent init] GSP: On Alternate: mean return-138.9 ± 0.3
[consistent init] GSP: On Alternate: episode when most runs take the S route12
[consistent init] GSP: Both: mean return-31.3 ± 0.2
[consistent init] GSP: Both: episode when most runs take the S route7

gsp run lava · 44.98 s · 2026-10-03T13:53:28+00:00

Which potential? v_g* versus heuristics

Sec. 6.4, Hypothesis 5, Fig. 15; App. F, Prop. 5
What this experiment does
Which potential? v_g* against heuristic and random potentials (Sec. 6.4, Hypothesis 5, Fig. 15; Prop. 5).

Every potential keeps the optimal policy (Ng et al., 1999); they differ in how fast learning gets there.
Prop. 5 says why a good potential helps: with Phi = v*, ONE Bellman backup from q = 0 already yields
the optimal greedy policy.
potentials.png
No potential: total steps over all episodes9621 ± 81
No potential: optimal after one backup0%
Random U[−100, 0]: total steps over all episodes17181 ± 1858
Random U[−100, 0]: optimal after one backup49%
−Manhattan distance to goal: total steps over all episodes6121 ± 68
−Manhattan distance to goal: optimal after one backup78%
v_g* (GSP): total steps over all episodes6101 ± 63
v_g* (GSP): optimal after one backup100%

gsp run potentials · 9.6 s · 2026-10-03T13:53:38+00:00

Shaping versus bootstrapping from v_g* (approximate LAVI)

Sec. 6.5, Fig. 16; App. F; Sec. 6.2
What this experiment does
Shaping versus bootstrapping from v_g* (Sec. 6.5 "approximate LAVI", Fig. 16; App. F; robustness, Sec. 6.2).

Bootstrapping straight from v_g* (beta = 1) inherits every error in v_g*; potential-based shaping cannot
change the optimal policy whatever v_g* is. The base learner here is one-step Sarsa (lambda = 0): a target
that bootstraps from a fixed function is a one-step target, and combining it with eligibility traces is a
mismatch (the trace sum no longer telescopes to a lambda-return). Panel (a) adds that combination, as the
paper's Sarsa(lambda) setting would have it, to show the effect. Three views:
  (a) exact, correct v_g* (config "both"):         bootstrapping is fine, even fast;
  (b) exact but biased v_g* (config "on_alternate"): bootstrapping locks onto the 19-step route;
  (c) noisy option models (multiplicative noise on r_gamma and Gamma, as learned models would have):
      shaping degrades gracefully, bootstrapping does not.
approx_lavi.png
Both / Sarsa(0): steps, last 30 episodes75.1
Both / GSP: potential shaping (Eq. 4): steps, last 30 episodes19.5
Both / bootstrap β = 0.5 (App. F): steps, last 30 episodes27.3
Both / bootstrap β = 1 (approx. LAVI): steps, last 30 episodes20.9
Both / approx. LAVI with traces, λ = 0.9: steps, last 30 episodes33.2
On Alternate / Sarsa(0): steps, last 30 episodes75.1
On Alternate / GSP: potential shaping (Eq. 4): steps, last 30 episodes21.8
On Alternate / bootstrap β = 0.5 (App. F): steps, last 30 episodes26.3
On Alternate / bootstrap β = 1 (approx. LAVI): steps, last 30 episodes23.2
noise 0% / GSP: potential shaping (Eq. 4): steps, last 3019.3 ± 0.1
noise 10% / GSP: potential shaping (Eq. 4): steps, last 3021.7 ± 0.5
noise 20% / GSP: potential shaping (Eq. 4): steps, last 3026.6 ± 1.3
noise 40% / GSP: potential shaping (Eq. 4): steps, last 3038.0 ± 2.4
noise 0% / bootstrap β = 1 (approx. LAVI): steps, last 3022.4 ± 0.4
noise 10% / bootstrap β = 1 (approx. LAVI): steps, last 30137.6 ± 42.3
noise 20% / bootstrap β = 1 (approx. LAVI): steps, last 30670.1 ± 70.3
noise 40% / bootstrap β = 1 (approx. LAVI): steps, last 30962.5 ± 11.2

gsp run lavi · 51.77 s · 2026-10-03T13:54:30+00:00

What planning computation buys: Dyna-Q versus GSP

Sec. 7.2-7.3, Fig. 17 (hypothetical in the paper)
What this experiment does
What planning computation buys: Dyna-Q with n planning updates per step versus GSP (Sec. 7.2-7.3, Fig. 17).

The paper's Fig. 17 is explicitly hypothetical. This puts measured points on it for the tabular case:
Dyna-Q's learning speed as a function of planning updates per real step, against GSP (on the same
Q-learning base learner), whose planning is a one-off goal-space value iteration.
dyna_vs_gsp.png
Dyna-Q n=0: environment steps for all episodes13314
Dyna-Q n=1: environment steps for all episodes9151
Dyna-Q n=3: environment steps for all episodes6211
Dyna-Q n=10: environment steps for all episodes4070
Dyna-Q n=30: environment steps for all episodes2732
GSP + Q-learning: environment steps for all episodes1231
GSP one-off planning backups30

gsp run dyna · 15.24 s · 2026-10-03T13:54:46+00:00

Reproduce

uv sync --extra dev
uv run gsp check          # numerical sanity checks, ~2 s
uv run gsp reproduce      # every experiment, figure, animation and this report
uv run gsp live           # live terminal dashboard
uv run pytest