Refactoring a Q-learning lab turned into a retro_gamer API redesign

This started as a routine request to clean up a messy lab file, and turned into a real design discussion about retro_gamer’s observation-handling API. Recorded here close to verbatim where the back-and-forth mattered, condensed where it was just file-reading.

Two warm-up questions

Before the actual task, a couple of quick conceptual questions about the lab content itself, which turned out to matter later:

❯ Why is Q-learning called Q-learning?

⏺ The “Q” stands for “quality” — Q(s, a) represents the quality (expected future reward) of taking action a in state s. From Chris Watkins’ 1989 PhD thesis, “Learning from Delayed Rewards.”

❯ Why is it called the Bellman Equation when it’s not an equation

⏺ It is an equation — V(s) = max_a[R(s,a) + γ Σ P(s’|s,a) V(s’)] has an equals sign. What’s unusual is that V appears on both sides: it’s a recursive/fixed-point equation, not one you solve algebraically for a single unknown.

❯ Is the Q-value update rule in the lab different from the Bellman equation?

⏺ I went and read the lab’s q_learning.py. The update rule Q(s,a) <- Q(s,a) + alpha*(r + gamma*max Q(s',a') - Q(s,a)) is the standard Bellman-style update, just rearranged into incremental learning-rate form. At alpha=1 it collapses exactly to the Bellman equation.

❯ But a value assignment is not an equation

⏺ Fair correction. The Bellman equation is a proposition (true only when Q equals Q*); the update rule is an assignment (an instruction to overwrite a variable), not a claim that anything is true or false. The lab’s docstring calling the assignment “the Bellman equation” is shorthand that conflates the two.

This is the moment I actually opened the lab’s q_learning.py file, which is what led to the actual task.

The actual request

❯ I want to refactor q_learning.py, the module where students will implement q-learning. Right now, it’s too messy and will be difficult for beginners to use.

  1. Separate the q_learning algorithm from a module which imports it and uses it to train BabySnake.
  2. Is it possible to use retro_gamer.GameEnvironment as the environment, rather than writing a new environment? If so, let’s use it. (Discuss your answer before implementation)
  3. Write tests for q_learning. Update the lab so students should get the tests passing before they train.
  4. There should be a cleaner way to invoke the training.
  5. Clean up the code style.

q_learning.py mixed four things in one file: a bespoke BabySnakeEnv class, the two functions students actually implement (choose_action, update_q), a training loop, and a terminal-watching routine — with no tests, and an awkward python -c "from q_learning import watch; watch()" as the only way to run it.

Why GameEnvironment didn’t fit, at first

⏺ I read retro_gamer/env.py. GameEnvironment.reset()/step() always return a flattened, one-hot-encoded numpy.ndarray of the whole board — wrong for tabular Q-learning for three reasons: not hashable as a dict key without conversion; the wrong size (the lab’s pedagogy depends on students counting a small discrete state space — questions.md literally asks “how many rows would the full Q-table contain?”); and it drags in GameMetadata/pyproject config machinery irrelevant to the two functions being taught. My first recommendation was: keep a small custom wrapper, pulled out of q_learning.py into its own module.

❯ I would like to use retro_gamer.GameEnvironment, even if it means changing the retro_gamer package. Let’s discuss this — it seems like it would be a straightforward extension to define an optional state_function … When used within the retro_gamer framework for DQN learning, observe_state should return a numpy array. But if using it for some other reason (as here), the return value could be arbitrary.

This reframed the problem entirely: instead of working around GameEnvironment’s limitations for this one lab, generalize it so the lab’s tabular use case and the existing deep-RL use case share one mechanism.

Design discussion

What followed was several rounds of narrowing down the actual contract, each one catching something the previous round had gotten wrong:

Scope. I initially proposed adding state_function as a constructor keyword, additive alongside the existing flags, to minimize package risk.

❯ Let’s discuss this — how would a user, working on training an agent for the game, iterate state functions? I’m attracted to specifying the state function in the game’s metadata … We currently have a special case for “egocentric” which could be replaced with a function provided in a library.

This is the point where it became clear the existing egocentric/ egocentric_player/egocentric_radius flags were a symptom: a built-in special case for something that’s better expressed as ordinary code calling a library function (egocentric_board(), which already existed as a free function in observation.py).

I pointed out a second, independent argument for binding this to the game’s metadata rather than the training environment: model_agent.py’s TrainedPolicy.get_action() (used at inference time) was independently re-implementing the same board-encoding logic that GameEnvironment implements for training — two code paths computing the same thing, with no guarantee they’d stay in sync. Binding the function to metadata loaded identically by both fixes that for free.

❯ Yes, also migrate existing games. retro_gamer is not in use yet, so I’m not worried about regressions. There is no reason to keep messy unnecessary code paths around.

Naming. I had been calling it state_function.

❯ Shall we rename state_function to observation_function? We use the language of observations elsewhere, and state is overloaded in retro to mean a dict in addition to the board positions.

Agreed immediately — game.state already means something else in retro, and observe_state/encode_observation already establish “observation” as the package’s vocabulary.

Whether spatial/character_set/board_size survive. My first attempt at a “full migration” tried to let CNN-vs-MLP be inferred from the shape of whatever the function returned (3-D → CNN, 1-D → MLP).

❯ I think we should keep the tool.retro-gamer.spatial key; this specifies MLP/CNN, more fundamental than the state representation.

This was the correct call, and I’d missed why: GameEnvironment._observe() never actually returns a (C, H, W) tensor even in the spatial case — it returns one flat vector, and _SpatialNet.forward() slices that flat vector back into a board tensor itself, using character_set and board_size to know where to cut. The literal shape of an observation is never a usable signal for “should this be a CNN,” because it’s always flat. spatial has to stay an explicit, independent declaration, and — once that’s true — character_set/board_size have to stay too, because they’re what lets _SpatialNet do that slicing regardless of which code path produced the vector.

Two more small turns of the crank settled the rest: whether observation_function and observe_state should be allowed together (no — they’re two conflicting ways of saying the same thing, so GameEnvironment.__init__ now raises if both are set), and whether extras_size needed to stay a declared, hand-set field:

❯ Perhaps we can get rid of extras_size. We have board_size and character_set, so we can figure out how to reshape the 1-D vector in the spatial case without extras_size.

Right — since the board portion’s length (C × W × H) is always knowable from the two fields that are staying anyway, extras_size is just len(one sampled observation) − board_length, measured once. This also turned out to be exactly the right place to enforce the “must be a flat 1-D array” rule — but only inside DQNTrainer, not inside GameEnvironment itself, since GameEnvironment now also has to work for this lab’s tabular case, where the “observation” is a plain tuple used as a dict key, not an array at all.

What got built

retro_gamer (published as 0.2.0):

  • GameMetadata.observation_function: a "module:attr" dotted-path string, resolved via resolve_observation_function() with the same heading/explanation/fix error-message style the rest of the file uses.
  • GameEnvironment/TrainedPolicy both resolve it from the same metadata and call it identically — fixing the training/inference duplication.
  • egocentric/egocentric_player/egocentric_radius removed outright. Cropping is now: write a function that calls egocentric_board() + encode_board() yourself, and declare board_size to match your crop (no more auto-deriving it from a radius).
  • observe_state_sizes and _discover_observe_state_sizes removed; extras_size is now measured from one sampled observation in DQNTrainer.__init__, with friendly errors for non-array, non-1-D, or too-short-for-the-declared-board observations, plus a generic length-changed-mid-run check.
  • A sys.path fix in cli.py — running the installed retro-gamer console script doesn’t put the caller’s cwd on sys.path the way python script.py does, which would have silently broken lab-relative observation_function modules.
  • A new test suite (tests/test_observation_function.py, 16 tests; none existed for this package before), exercising resolution, delegation, mutual exclusivity, and the extras_size discovery/validation logic against a real (zero-agent) retro.game.Game, not mocks.
  • Docs (reference.rst, walkthrough.rst, background.rst, introduction.rst, api.rst) updated and deployed via make deploy.
  • Version bumped 0.1.1 → 0.2.0, committed, pushed, built, and published to PyPI; the lab’s pin bumped to match.

The lab (modules/lab_reinforcement_learning/):

  • q_learning.py is now exactly the two functions students implement — choose_action(q_table, state, actions, epsilon) and update_q(q_table, state, action, reward, next_state, actions, alpha, gamma) — with actions as an explicit parameter instead of a module-global, which is what actually makes the file game-agnostic. No babysnake/retro/retro_gamer imports at all.
  • babysnake_env.py holds ACTIONS and get_state(game) — BabySnake’s observation_function, a plain (agent_x, agent_y, food_x, food_y) tuple, referenced from babysnake/pyproject.toml.
  • train_babysnake.py builds GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("babysnake")) directly — the bespoke environment wrapper class is gone entirely, replaced by the now-properly- generalized GameEnvironment. watch() reuses q_learning.choose_action(..., epsilon=0.0) instead of re-implementing “pick the best action” a second time. One command, python train_babysnake.py, trains and watches.
  • test_q_learning.py — 8 fast unittest tests, no game dependency at all, matching this repo’s existing test-file convention. Confirmed they fail clearly against the unfilled stub and pass against a filled-in solution.
  • questions.md updated: a checkpoint instructing students to get the tests passing before training, and the post-training instructions now point at python train_babysnake.py.
  • Found and fixed (at the user’s prompt) a real latent issue surfaced during verification: the training loop had no per-episode step cap, so a lucky random walk that kept finding food (each pickup restores more energy than a turn costs) could produce an effectively unbounded episode — confirmed by one episode hitting 5000+ steps from pure randomness in 0.03 seconds. This existed in the original q_learning.py too; not something the refactor introduced, but worth a max_steps_per_episode cap regardless.
  • runs/snake/ migrated to observation_function (a hand-written egocentric_observation() reproducing the old egocentric+board+extras behavior via the library functions). Its prior 12,000-episode checkpoint history became incompatible the moment the config declared observation_function (the checkpoint compatibility checker has no way to verify the new function is behaviorally equivalent to the old flags, so it conservatively refuses to resume) — at the user’s direction, those checkpoints were deleted and a fresh 20,000-episode run launched.

A genuinely useful constraint: the checkpoint compatibility checker

Worth noting on its own: DQNTrainer snapshots the parts of its config that determine network shape/meaning into every checkpoint, and refuses to resume if they’ve changed, with a message naming exactly which field changed and why it matters. That’s why migrating runs/snake/’s config to observation_function couldn’t just be a silent, free win — the checker did its job correctly, flagging a real representational change even though the new function was designed to be numerically equivalent to the old one. Conservative-by-construction compatibility checking like this is a good pattern in general: it would have been easy to build something that just let stale checkpoints load with a corrupted/mismatched network instead.

Results

The fresh 20,000-episode run on the migrated config (observation_function = "snake_observation:egocentric_observation") finished in 56m17s on Apple Silicon (MPS), saving 200 checkpoints. The architecture line confirms the migration preserved the intended shape:

[INIT] Board: 17×17, character set: 6 chars (one-hot per cell)
[INIT] Observed state features: 2  |  Actions (incl. no-op): 5
[INIT] spatial=False → using MLP architecture
[INIT] MLP: 1736 → 128 → 64 → 5

1736 = 17×17×6 (board) + 2 (apple_dx, apple_dy) — extras_size = 2 was discovered automatically from one sampled observation, exactly as designed, with no observe_state_sizes declared anywhere.

The new run reproduced the same qualitative arc the lab’s snake_training.md was already built to teach — reward climbs, dips, recovers with a sharp drop in episode length, then climbs to a stable high plateau — just with new numbers, since training is stochastic:

[ep_1300]   avg_reward= +31.6  avg_steps=49
[ep_2300]   avg_reward=  -3.8  avg_steps=87
[ep_4000]   avg_reward= +53.2  avg_steps=42
[ep_20000]  avg_reward=+335.6  avg_steps=94

The dip-and-recovery is real, not an artifact: avg_reward bottoms out around episode 2,300 (agent wandering, low reward, high step count) and recovers by episode 4,000 with avg_steps dropping sharply (87 → 42) at the same moment reward jumps — the same signature of “agent stops wandering and starts beelining for the apple” as the original run showed, just centered on different episode numbers. snake_training.md’s Q5 (the training-curve table) and Q6 (the three --checkpoint episodes to watch and compare) were updated to ep_1300/ep_2300/ep_4000/ep_20000 to match. Loaded each of the three Q6 checkpoints through TrainedPolicy directly (not just visually via retro-gamer play) to confirm they all produce valid actions before handing this back.

One thing this surfaced that’s worth remembering for next time: I launched the long training run with nohup ... & disown, which detached it from the harness’s own task tracking — so the “task complete” notification fired for the launcher command almost immediately, not for the actual 56-minute training process. The training itself kept running correctly; I just had no way to be notified when it finished. Fixed by leaving the real process alone and wrapping a new, properly-tracked background command around it (while kill -0 $PID; do sleep 15; done), which got a real completion notification instead of requiring manual polling.