SERIES: Learning RL and JAX in Public - from zero to DeepMind :)

On Day 1 I covered what JAX is and the three superpowers: GPU acceleration, grad, and jit. Today I actually used them. And the moment where I manually verified a gradient by hand and watched JAX confirm it is the one that made everything real for me.

But before the code, I had to understand what a gradient actually is. Because I realised I had been using them for years without a clean mental model.

What is a derivative, really?

Imagine you are adjusting a dial on a machine. The dial controls the volume. Your goal is to get the volume to exactly 50. Right now it is at 30.