A few years ago, I worked through Sutton & Barto as part of the University of Alberta's Coursera course on reinforcement learning. I found the math to be incredibly complex, with various branches into algorithms that seemed to be ancillary to the main progression toward deep RL. Some of the most important concepts were left as "exercises for the student."

I'm currently working on a video series that teaches the basics of using deep RL (namely PPO) to train an ESP32-based balance bot. Additionally, I wanted to solidify my understanding of the underlying math. As a result, I am actively working on a set of blog posts that build up to PPO, starting from the very basics.

In other words, my goal is to create a quicker and more approachable text to bring newcomers up to speed on the math behind modern deep RL algorithms. I've listed the posts below, and I would appreciate any feedback you might have! I'll update this list as I continue to add articles.

1. What is Reinforcement Learning?
2. Rewards, Returns, and the Discount Factor
3. Policies, Markov Decision Processes (MDPs), and Trajectories
4. Expected Return, Value Functions, and Bellman Equations
5. The Bellman Optimality Equations
6. Dynamic Programming
7. Monte Carlo Methods
8. Temporal-Difference (TD) Learning
9. TD(λ) and Eligibility Traces
10. Q-Learning
11. Deep Q-Networks (DQN)

I'm training an AlphaZero-style agent (Gumbel MuZero via DeepMind's mctx, in JAX, single RTX 5070) and found that the MCTS search cost grows superlinearly with the simulation budget: at 16 / 32 / 64 simulations per move the training throughput was 143 / 47 / 9 episodes per second. Doubling the budget should roughly double the cost, not triple it — so something in the search was superlinear.

Isolating the tree machinery (network replaced by constants, then the environment removed too) showed it is O(N²⁾ all by itself: the pure per-simulation tree overhead measured 0.47 / 0.52 / 1.05 / 1.95 ms at 8 / 16 / 32 / 64 sims — the per-simulation cost doubles every time the number of simulations doubles.

The compiled HLO showed why. mctx's backward pass carries the whole search tree through a lax.while_loop and, within one iteration, both gathers from and scatters into children_values / children_visits. XLA:GPU cannot prove the scatter can alias, so the compiled loop body contains a full copy of both [batch, nodes, actions] buffers on every step of the leaf-to-root walk. Buffer size grows with the simulation count and the walk depth grows with it too, so the whole thing is quadratic.

The fix carries only an O(num_nodes) path through the loop and applies one scatter per array after the loop. Search results are bitwise-identical (same rng gives the same actions and action_weights), and the full upstream test suite passes, including the golden-tree comparisons. End-to-end training speedup: x1.10 at 16 sims, x1.58 at 32, x3.27 at 64; XLA compile time at 64 sims dropped from 63 s to 19 s.

Two gotchas that cost me time: 1. Top-K action sampling — the "textbook" fix for large action spaces — gained nothing once the copies were gone; the copies were the A-scaled term, not the per-action math. 2. Consumer-GPU clock ramp-up (495 -> 2932 MHz) makes cross-process benchmarks lie by up to x1.7 — all A/B numbers above are interleaved in a single process.

Full write-up with HLO dumps, benchmark methodology and the things that didn't work, plus the PR, are linked in a comment below.