SERIES: Learning RL and JAX in Public - from zero to DeepMind :)
Days 4 through 8 were raw JAX. Every weight matrix initialized by hand, every forward pass written out explicitly. It was intentional - I wanted to understand exactly what was happening under the hood.
Day 9 is the payoff. Same Actor-Critic algorithm, rewritten in Haiku. DeepMind's neural network library.
What Haiku is
Haiku is what DeepMind researchers actually use. If you look at the code releases for AlphaFold, Acme (DeepMind's RL framework), or AlphaStar, the network definitions are in Haiku.
It sits on top of JAX. You get all of JAX's superpowers - grad, jit, vmap - plus a cleaner way to define networks.
The one thing that makes Haiku different from PyTorch
In PyTorch, the model stores its own weights:
model = MyNetwork()
output = model(input) # weights hidden inside
In Haiku, weights are always separate from the network:
params = network.init(key, input) # weights created here
output = network.apply(params, input) # passed in explicitly
This feels like more work at first. But it means your network is a pure function with no hidden state. JAX can jit it, vmap it, grad it - exactly like any other function. This is why DeepMind code is so composable.
The before and after
Before (raw JAX - what we wrote on Days 4-8):
def init_actor(key):
k1, k2 = jax.random.split(key)
return {
"w1": jax.random.normal(k1, (16, 64)) * 0.1,
"b1": jnp.zeros(64),
"w2": jax.random.normal(k2, (64, 4)) * 0.1,
"b2": jnp.zeros(4),
}
def actor_forward(params, x):
h = jnp.tanh(x @ params["w1"] + params["b1"])
return jax.nn.softmax(h @ params["w2"] + params["b2"])
After (Haiku):
def actor_fn(x):
return hk.Sequential([
hk.Linear(64), jax.nn.tanh,
hk.Linear(4), jax.nn.softmax,
])(x)
actor = hk.without_apply_rng(hk.transform(actor_fn))
Same network. Same output. Half the code. No manual weight shapes. No room to accidentally write (64, 16) instead of (16, 64).
What hk.transform actually does
hk.transform wraps your function into two methods:
-
actor.init(key, sample_input)- runs the network once, captures all the weights it creates, returns them as a dictionary -
actor.apply(params, input)- runs the forward pass using those params
After init, you have a plain dictionary of arrays. You can save it, load it, pass it to jax.grad, inspect individual layers, merge it with another network's params. All things that are painful with PyTorch.
The result
Same gridworld. Same Actor-Critic algorithm. Same policy arrows pointing toward the goal. But when I printed the parameter structure:
linear/w: shape (16, 64)
linear/b: shape (64,)
linear_1/w: shape (64, 4)
linear_1/b: shape (4,)
Haiku named and organized everything automatically. In the raw JAX version I had to track all of this myself.
Why I spent days on raw JAX before this
Because if you jump straight to Haiku, actor.apply(params, x) looks like magic. You do not know what params is or why it needs to be passed in.
Having built the thing manually for five days, I now know exactly what Haiku is doing behind the scenes. It is not magic. It is just a clean wrapper around the same dictionary-of-arrays pattern we wrote by hand.
That is the version of understanding that research work requires. Not "I know how to use the library." But "I know what the library is doing."
From Day 10, we start applying this to real problems. The gridworld served its purpose. Time to build something that actually goes on a resume.
All code from this series, organised by day, is on my GitHub: https://github.com/MadhumithaKolkar/jax-rl-lab
Happy learning everyone !
~ Madhumitha Kolkar ( index_0 )
Top comments (1)
Small thing hiding under "same network, same output": hk.Linear defaults w_init to a truncated normal with stddev 1/sqrt(fan_in), so Haiku starts at 0.25 on the 16-input layer and 0.125 on the head, not the flat 0.1 you hand-wrote. Same architecture, different starting policy, which in actor-critic means the early exploration isn't the same run even at a fixed seed. Did you pass an explicit w_init to match the raw version, or just compare the final arrows?