DEV Community: Madhumitha Kolkar

I Rewrote My RL Agent in DeepMind's Neural Network Library - Day 9 (Haiku)

Madhumitha Kolkar — Tue, 04 Aug 2026 15:00:55 +0000

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 )

I Implemented the Algorithm Behind ChatGPT From Scratch - Day 8 (PPO).

Madhumitha Kolkar — Fri, 31 Jul 2026 16:12:33 +0000

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

When people ask "how was ChatGPT trained?", the answer usually involves RLHF - Reinforcement Learning from Human Feedback. And the RL part of RLHF is PPO.

Day 8 is PPO. And it turns out to be simpler than I expected.

The problem that PPO solves

Actor-Critic (Day 7) worked, but it had an instability problem. A single episode with an unusually high or low return would push the network weights too far in one direction. The policy would change dramatically. Sometimes it would unlearn things it already knew. Training would collapse.

The fix is elegant: put a speed limit on learning.

The one idea that is PPO

After each update, PPO compares the new policy to the old one:

ratio = new_probability / old_probability

If ratio = 1.0: nothing changed for this action.
If ratio = 1.5: the new policy is 50% more likely to take this action.
If ratio = 0.5: the new policy is 50% less likely.

Then PPO clips this ratio. It says: I will not let you change by more than 20% in one update.

clipped_ratio = clip(ratio, 0.8, 1.2)
loss = -min(ratio * advantage, clipped_ratio * advantage)

That is literally it. One clip. One min. That is the entirety of what makes PPO different from Actor-Critic.

The gradient from a clipped update becomes zero once the ratio hits the boundary. The policy stops updating further for that action in that step. Come back next batch and nudge it again if you want more.

Small stable steps, every update.

The three things PPO adds over Actor-Critic

1. Clipping (the main idea)
Keeps the policy from changing too much in one shot.

2. Multiple epochs per batch
Actor-Critic uses each episode once. PPO collects a batch of episodes and then runs 4 gradient updates on that same data. More learning per episode. More sample efficient.

3. Entropy bonus
Entropy measures how spread out the action probabilities are. High entropy means the agent is still considering many options. Low entropy means it has collapsed to always picking one action.

PPO adds a small reward for entropy:

total_loss = policy_loss + critic_loss - 0.01 * entropy

This keeps the agent exploring longer before committing. Without it, policies can collapse to one action too early and get stuck.

What the training loop looks like

for iteration in range(200):
    # collect 10 episodes with current policy
    batch = collect_batch(actor, critic, episodes=10)

    # PPO epochs: squeeze 4 updates out of this batch
    for epoch in range(4):
        actor_loss = clipped_ppo_loss(batch)
        critic_loss = value_loss(batch)
        update(actor, critic)

Two loops. Outer loop collects data. Inner loop extracts learning from that data. Actor-Critic had one loop. This is the reason PPO trains faster per episode collected.

What I saw when I ran this

The value estimates from the critic after training:

states near goal:  +0.6 to +0.8
safe middle path:  +0.1 to +0.4
states near holes: -0.3 to -0.6

The critic mapped out the whole gridworld just by watching the actor fail and succeed. Nobody told it where the holes were. It figured it out.

The policy arrows lined up cleanly. Same result as Q-learning (Day 4), DQN (Day 5), REINFORCE (Day 6), Actor-Critic (Day 7). Different algorithm every time. Same learned behavior.

That pattern keeps hitting me. The Bellman intuition runs through all of it.

Why PPO became the default

Three reasons: it is stable (the clip prevents collapse), it is simple (one extra line over Actor-Critic), and it is general (discrete actions, continuous actions, LLM fine-tuning - all the same algorithm).

When OpenAI trained ChatGPT, they had humans rank responses. Those rankings became a reward signal. PPO optimized the language model against that signal. The RL loop you just implemented is the same loop, just with a very different environment and a very different reward function.

When I built the forge project (a GRPO trainer), I did not fully understand why it worked the way it did. GRPO is PPO with one change: instead of a single critic estimating advantage, it runs a group of episodes and compares their returns to each other. Same clipping. Same stability. The group comparison replaces the critic.

Five days ago I could not have explained that. Now I can.

Day 9: Haiku. DeepMind's neural network library. We rewrite all of this in a fraction of the code, and it starts looking like actual research-grade JAX.

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)

I Added a Second Brain to My RL Agent - Day 7 (Actor-Critic)

Madhumitha Kolkar — Wed, 29 Jul 2026 04:05:19 +0000

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

Day 6 ended with REINFORCE working. The agent learned to navigate the gridworld using probabilities instead of Q-values. Clean idea, clean code.

But it was noisy. Training was slow. Some episodes the agent got lucky, some unlucky, and the network kept getting confused by the difference. The signal was there but it was buried in too much variance.

Day 7 is the fix.

The basketball analogy that made it click

Imagine a basketball player scores 15 points in a game. Is that good or bad?

Depends entirely on what's normal for them. If they usually score 10, it's a great game. If they usually score 25, it's a bad one. The number alone tells you nothing. You need context.

REINFORCE has no context. It sees "return = 0.7" and tries to update the network. But is 0.7 good or bad for that particular state? It has no idea. It just treats every positive return as good and every negative one as bad.

This is why training is noisy. Lucky episodes push the network in the wrong direction. Unlucky ones push it away from good decisions.

Actor-Critic adds the context.

Two networks, two jobs

The Actor - same as REINFORCE. Takes a state, outputs action probabilities. Decides what to do.

The Critic - new. Takes a state, outputs a single number: "how good is it to be in this state?"

actor_output:   [0.1, 0.3, 0.2, 0.4]   # probabilities over 4 actions
critic_output:  0.72                     # estimated value of this state

The critic never picks actions. It just watches and judges. Its only job is to give the actor a better signal.

The advantage: the number that matters

Instead of using the raw return to update the actor, we compute:

advantage = actual_return - critic_estimate

Advantage positive: this played out better than expected. Do this more.
Advantage negative: this played out worse than expected. Do this less.
Advantage zero: exactly average. No change needed.

This one subtraction removes a huge amount of noise. The lucky episode problem disappears because "better than expected" is now a meaningful signal, not just a big number.

What the update looks like

Both networks update every episode. First the critic (so it gets better at estimating), then the actor (using the fresh estimates):

# critic: learn to predict returns accurately
critic_loss = (actual_return - critic(state)) ** 2

# actor: same as REINFORCE but now uses advantage instead of raw return
actor_loss = -advantage * log(probability_of_action_taken)

The critic trains like a simple regression. The actor trains like REINFORCE but with a cleaner signal. They improve together over time.

What I saw in the training output

The code prints the critic's value estimates as a grid at the end. After training it looked roughly like:

states near goal:    high value  (~0.7 to 0.9)
states near holes:   low value   (negative)
safe middle states:  medium value

The critic learned the layout of the gridworld just by watching the actor play. Nobody told it where the holes were. It figured it out through experience.

That was genuinely cool to watch.

Why this is the architecture that everything is built on

I keep mentioning PPO and GRPO. Here is the direct line:

PPO (used to train ChatGPT): actor-critic. Adds a clip to stop updates from overshooting.

A3C (DeepMind, 2016): actor-critic. Runs many agents in parallel to collect more diverse experience.

GRPO (DeepSeek-R1): actor-critic variant. Replaces the single critic with group comparisons across multiple episodes. Instead of "better than my estimate," it asks "better than the other episodes in this batch?"

The forge project I built earlier in this series implements GRPO. When I built it I did not fully understand why the reward normalization worked the way it did. Now I do. It is computing advantage within a group.

Understanding actor-critic was the missing piece.

The family tree in one place:

Actor-Critic
    |
    ├── A3C  (DeepMind, 2016)   - parallel actors
    ├── PPO  (OpenAI, 2017)     - clipped updates, trains ChatGPT
    ├── SAC  (Berkeley, 2018)   - entropy bonus, used in robotics
    └── GRPO (DeepSeek, 2024)   - group-relative advantage, trains R1

Every single one: an actor that picks actions + something that estimates quality. The differences are in how they compute advantage and stabilize training.

Day 8: we move into Haiku, DeepMind's neural network library. Same algorithms, much cleaner code. That is where the codebase starts looking like actual research code.

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)

I Taught an Agent to Act Directly - No Q-Values Needed (Day 6: REINFORCE)

Madhumitha Kolkar — Sat, 25 Jul 2026 05:05:19 +0000

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

Days 4 and 5 were value-based methods. Q-learning, DQN - the agent learns how good each action is, then picks the best one. The policy is implicit. It falls out of the Q-values as a side effect.

Day 6 changes the approach entirely. What if you just... learn the policy directly?

That is REINFORCE. And it is the foundation of every modern RL algorithm I care about - PPO, A3C, and GRPO (the algorithm behind DeepSeek-R1 that I built a small implementation of earlier in this journey).

The problem with DQN that I did not see until now :

DQN always picks the action with the highest Q-value. It is always confident. The only randomness comes from epsilon-greedy, which is literally just "sometimes ignore what you learned and act randomly."

That is a hack. A useful hack, but still a hack.

Also, DQN breaks completely for continuous actions. If your action is a steering angle between -180 and 180 degrees, you cannot have one Q-value per action - there are infinite actions. DQN has no answer for this.

REINFORCE does.

The one change that makes REINFORCE :

In DQN, the network outputs Q-values :

q_values = network(state)  # [2.1, 0.8, 1.4, 3.2]
action = argmax(q_values)  # always pick the best

In REINFORCE, the network outputs probabilities :

probs = network(state)     # [0.1, 0.3, 0.2, 0.4]
action = sample(probs)     # sample from distribution

That is the whole difference. The network now represents the policy directly. It says: given this state, here is how likely I think each action is. You sample from that - sometimes the 30% action wins, and that is fine.

How it learns :

REINFORCE runs a full episode first. Then it looks back at what happened :

If an action led to a high total reward: increase its probability.
If an action led to a low or negative reward: decrease it.
Scale the update by how large the reward was.

In math :

loss = -G_t * log π(a_t | s_t)

G_t is the total discounted return from that timestep. log π is the log probability of the action you took. The negative sign is because JAX minimizes and we want to maximize return.

In plain English: reward what worked, penalize what did not. That is it.

What I built today :

Same 4x4 gridworld. But instead of Q-values, the network outputs action probabilities via softmax. The training loop looks like this:

for episode in range(3000):
    states, actions, rewards = collect_episode(params, key)
    returns = compute_returns(rewards, gamma=0.99)

    loss, grads = grad_fn(params, states, actions, returns)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

Three key differences from DQN :

No replay buffer - we use the episode as it happened
No target network - no Bellman bootstrapping needed
Compute returns from the end of the episode backwards

After 3000 episodes, the policy arrows matched Days 4 and 5. Same gridworld. Same result. Completely different mechanism.

One thing that clicked for me today :

REINFORCE is Monte Carlo RL. It waits for the full episode, computes the actual return, then learns. DQN is temporal difference learning - it learns after every single step using an estimate of future value.

Monte Carlo is more accurate in principle (you see what actually happened, not an estimate). But it is high variance - one bad episode can throw off your update badly. This is why REINFORCE can be slow and noisy compared to DQN.

Every modern algorithm after REINFORCE is essentially trying to fix this variance problem. Actor-Critic methods combine both: a policy network that acts (REINFORCE) and a value network that estimates returns (like DQN) to reduce variance.

That is Day 7. The actor-critic architecture.

Also : I keep saying GRPO builds on this. Here is the direct line. GRPO runs multiple episodes in a group, computes relative returns within the group (instead of raw G_t), and uses those relative returns as the update signal. The policy gradient update structure is identical to REINFORCE. Just a smarter way of computing what "good" means.

When I built the forge project a few days ago, I did not fully appreciate why the reward computation worked the way it did. Now I do.

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)

I Replaced a Q-Table With a Neural Network and Everything Changed - Day 5 (DQN).

Madhumitha Kolkar — Fri, 24 Jul 2026 03:06:41 +0000

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

Day 4 ended with a working Q-learning agent. A table of numbers that an agent used to navigate a gridworld. Clean, simple, satisfying.

Day 5 started with one question: what happens when the state space is too large for a table?

Imagine instead of a 16-cell grid, your state is the pixels of a video game screen. There are more possible Atari frames than atoms in the universe. You cannot have a table row for each one. The Q-table breaks immediately.

DQN is the answer. And it is simpler than it sounds.

The one change that makes DQN

In Q-learning:

q_value = q_table[state, action]

In DQN:

q_values = neural_network(state)  # returns one Q-value per action

That is the whole idea. Replace the table lookup with a neural network forward pass. The Bellman equation stays the same. Epsilon-greedy stays the same. The training loop structure stays the same.

The network takes the current state as input and outputs a Q-value for every possible action. You still pick the highest one. The only difference is where the numbers come from.

Why neural networks generalise where tables cannot :

A Q-table only knows about states it has explicitly visited. State 7 ? Only useful if the agent has been to state 7 before.

A neural network learns patterns. If states 3, 7, and 11 all have similar features (say, they are all near a hole), the network learns that similarity and can make reasonable predictions for states it has never visited. This is what allows DQN to scale.

The two problems DeepMind had to solve :

Before DeepMind's 2013 paper, people had been trying to combine neural networks with RL for years. It kept failing. The network would improve for a while and then suddenly collapse and forget everything.

DeepMind solved this with two tricks that are now standard in every RL codebase.

Problem 1: Correlated samples :

In supervised learning you shuffle your dataset and sample random batches. This is important because the network needs diverse examples each step, not 32 consecutive frames from the same game session.

In RL, consecutive transitions are extremely correlated. The agent is in state 4, goes to state 5, then state 6. Training on these in sequence is like training a classifier on 32 pictures of the same cat in a row.

Fix: Experience Replay. Store every transition in a buffer. Sample randomly from it when training. Correlation broken.

replay_buffer.add(state, action, reward, next_state, done)
batch = replay_buffer.sample(batch_size=32)  # random, not sequential

Problem 2: Moving targets

The Bellman target is:

target = reward + gamma * max(Q(next_state))

But Q(next_state) comes from the same network you are updating. So every time you nudge the network, the target moves too. You are chasing a bullseye that runs away from you every time you take a step.

Fix: Target Network. Keep two copies of the network. Update only the online network each step. Every 100 steps, copy the online network weights into the target network. The target network provides stable Bellman targets while the online network trains against them.

# use target network for stable bellman target
target = reward + gamma * max(target_network(next_state))

# update only the online network
loss = (target - online_network(state)[action]) ** 2

What I built today :

Same 4x4 gridworld from Day 4. But instead of a Q-table, a three-layer neural network in raw JAX. Input: one-hot encoding of the state. Output: Q-values for all four actions.

The training loop :

for each episode:
    while not done:
        action = epsilon_greedy(online_network, state)
        next_state, reward, done = env_step(state, action)
        replay_buffer.add(state, action, reward, next_state, done)

        if len(buffer) >= 200:
            batch = replay_buffer.sample(32)
            loss, grads = grad_fn(online_params, target_params, batch)
            online_params = update(online_params, grads)

        if step % 100 == 0:
            target_params = online_params  # refresh anchor

After 1500 episodes, the policy arrows matched what Q-learning produced on Day 4. Same environment, same result, different mechanism. The DQN learned it too.

The thing that actually hit me today :

When I printed the Section 7 comparison at the end of the code:

Q-learning:  Q-values stored in a table (64 numbers)
DQN:         Q-values stored in a neural network
Same Bellman equation. Same epsilon-greedy. Same goal.
The only difference is where the Q-values come from.

Six years in ML and I had never sat down and traced the line from Q-learning to DQN to AlphaGo. It is one continuous idea, each step replacing one component with something that scales better. Tables become networks. Networks become deeper. State spaces become images. Rewards become sparse. But the core Bellman intuition runs through all of it.

That is the thing about learning from first principles. You stop using tools and start understanding systems.

Also worth noting: this is the paper that put DeepMind on the map. "Playing Atari with Deep Reinforcement Learning", 2013. If you have not read it, it is surprisingly readable. The ideas are exactly what we implemented today, just applied to pixel inputs and 18 actions instead of 16 states and 4 actions.

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)

I Built My First Reinforcement Learning Agent From Scratch - Day 4.

Madhumitha Kolkar — Thu, 23 Jul 2026 03:18:05 +0000

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

Today was different. Today the agent actually learned something.

Days 1 through 3 were JAX fundamentals - arrays, gradients, jit, vmap. All of it was preparing for this. Day 4 is where reinforcement learning actually begins.

I built a Q-learning agent from scratch and watched it navigate a gridworld. By the end of 2000 episodes, it was finding the goal reliably. No hand-coded rules. No labels. Just an agent figuring out the world through trial and error.

Here is how it works.

The setup :

A 4x4 grid. The agent starts at the top-left. The goal is at the bottom-right. There are holes scattered around. Fall in a hole and the episode ends with a negative reward. Reach the goal and you get a positive reward.

 S   .   .   .
 .   H   .   H
 .   .   .   H
 H   .   .   G

S = start, G = goal, H = hole.

The agent can move up, down, left, or right. That is it. 16 states, 4 actions.

What is a Q-value ?

Before writing any code, I had to understand this properly.

A Q-value answers one question: if I am in this state and I take this action, how much total reward will I collect from here onwards if I play well?

It is a table. Rows are states. Columns are actions. Every cell is a number.

         UP    DOWN   LEFT   RIGHT
state 0:  1.2   3.4    0.8    2.1
state 1:  0.5   0.9    2.3    1.1
...

At state 0, the agent looks at the row, picks the highest number, and takes that action. That is the entire policy. The whole job of Q-learning is to fill this table with the right numbers.

The Bellman equation :

At the start, every Q-value is zero. The agent knows nothing.

After taking an action and getting a reward, it makes a small correction to the relevant cell:

Q(s, a) = Q(s, a) + alpha * (target - Q(s, a))

Where target = reward + gamma * best Q-value in the next state.

This is the Bellman equation. It says: the value of being here and taking this action should equal the immediate reward plus the best value available from where you end up.

The agent does this correction thousands of times. Slowly the table fills with values that reflect how good each action actually is.

Exploration vs exploitation :

Here is a problem. If the agent always picks the action with the highest Q-value, it never tries anything new. It might settle for a mediocre path because it never explored better ones.

The solution is epsilon-greedy:

With probability epsilon: pick a random action (explore)
With the rest: pick the best known action (exploit)

Epsilon starts at 1.0 (fully random) and decays over time toward 0.01 (mostly greedy). At the start the agent wanders. By the end it has earned the right to trust its own table.

The training loop in code :

for episode in range(2000):
    state = START_STATE

    while True:
        action = choose_action(state, q_table, epsilon)
        next_state, reward, done = step(state, action)

        # Bellman update
        target = reward + gamma * np.max(q_table[next_state])
        q_table[state, action] += alpha * (target - q_table[state, action])

        state = next_state
        if done:
            break

    epsilon = max(0.01, epsilon * 0.995)

That is the whole thing. No neural network. No backpropagation. Just a table being updated one cell at a time.

What the agent learned :

After 2000 episodes, the policy arrows looked like this:

  >   >   v   v
  ^   H   v   H
  ^   >   v   H
  H   >   >   G

It learned to avoid holes. It learned the shortest path. It built this entirely from reward signals, no supervision.

The reward curve told the whole story. Early episodes: negative reward, the agent kept falling in holes. Around episode 500: reward starts climbing. By episode 2000: consistently positive, consistently reaching the goal.

Why does this matter beyond a gridworld ?

Q-learning is the foundation of DQN, which is the foundation of AlphaGo, which is the foundation of everything DeepMind built in the RL space. The gridworld feels toy-sized but the mechanics are identical. The only difference in modern deep RL is that the Q-table gets replaced by a neural network because the state space is too large to enumerate.

That is literally the next step in this series.

One thing I genuinely did not expect: how satisfying it is to watch an agent learn. There is something different about seeing a system figure something out on its own versus supervising it toward a known answer. I think this is why RL has such a hold on people who work in it.

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)

The JAX Superpower Nobody Talks About Enough - vmap - Day 3.

Madhumitha Kolkar — Wed, 22 Jul 2026 06:02:06 +0000

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

Everyone talks about JAX's automatic differentiation and JIT compilation. Those are great. But vmap is the one that made me actually understand why research teams at places like DeepMind and Google Brain write everything in JAX. It is also the one that is hardest to explain with a one-liner, so let me take a proper shot at it.

The problem :

You have a function that works on one input. A single vector, a single example, a single data point.

def score(x):
    return jnp.sum(x ** 2)

Now you have a batch of 1000 inputs and you need to run this on all of them.

The Python way:

results = [score(x) for x in batch]  # slow loop

The PyTorch way: rewrite your function to handle batches explicitly, add batch dimensions, make sure everything broadcasts correctly. Tedious. Error-prone.

The JAX way: vmap !!

What vmap does :

vmap takes a function written for a single input and returns a new function that runs on a whole batch, in parallel, without you changing anything.

batch_score = jax.vmap(score)
results = batch_score(batch)  # runs on all 1000 inputs at once

You wrote score thinking about one input. JAX handles the batching. Same results as the loop, but running as a single vectorised operation on the hardware.

I benchmarked this today. The vmap version was 32 times faster than a Python loop on CPU. On a GPU, that gap is even wider.

The in_axes argument :

This is the one thing about vmap that needs a second read. Some arguments you want to batch over, and some you want to keep fixed across the whole batch.

Example: you have a loss function with a shared weight vector w, but different inputs x and targets y for each example.

def loss(w, x, y):
    prediction = jnp.dot(w, x)
    return (prediction - y) ** 2

batch_loss = jax.vmap(loss, in_axes=(None, 0, 0))

in_axes=(None, 0, 0) means:

None: do not batch over w. It is shared across all examples.
0: batch over the 0th axis of x. One x per example.
0: batch over the 0th axis of y. One y per example.

That is it. You tell vmap which arguments vary across the batch and which stay fixed.

The real JAX pattern: jit + vmap + grad together :

This is what DeepMind code looks like. One line, all three superpowers:

fast_batched_grad = jax.jit(jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0)))

Your loss function is now differentiable, batched across all examples simultaneously, and compiled for speed. In PyTorch this would take 10-15 lines and careful attention to tensor shapes.

I ran a full batched gradient descent loop today with this pattern. w started at [0, 0, 0] and converged to [2, 3, 5] in 115 steps with loss hitting exactly 0.0. I played with the iteration count to find that number and found the convergence satisfying to watch.

Why this matters for research :

When you are running experiments on large models, you often want to compute gradients for an entire batch of examples simultaneously. Without vmap you either loop (slow) or rewrite your functions to be batch-aware (messy and error-prone at scale).

With vmap you write clean single-example functions and compose them up to batch scale. Your code stays readable. Your functions are testable on one example at a time. And vmap handles the rest.

This is why JAX code in research papers is often remarkably compact compared to equivalent PyTorch. The functions stay simple. The composability does the scaling.

The summary I keep in my notes :

What you want	How you get it
GPU acceleration	automatic
Gradients	jax.grad(fn)
Speed	jax.jit(fn)
Batch without a loop	jax.vmap(fn)
All three	jax.jit(jax.vmap(jax.grad(fn)))

Next up: Day 4, which is where things get genuinely exciting. Q-learning from scratch. First real reinforcement learning algorithm. I built a gridworld environment in pure JAX and trained an agent to navigate it. See you there.

If you have been following along: which of the four superpowers surprised you most? For me it was vmap. I expected grad to be the revelation but vmap is the one that changed how I think about writing ML code.

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)

I Built a Neural Network Training Loop in 5 Lines Using JAX - Day 2

Madhumitha Kolkar — Tue, 21 Jul 2026 04:24:25 +0000

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.

You turn the dial a tiny bit. The volume goes up.

The derivative answers one question: if I nudge this dial by a tiny amount, how much does the output change?

That is it. In math:

f(x) = x * x
f(3) = 9
f(3.001) = 9.006001

nudge of 0.001 caused a change of 0.006
rate of change = 0.006 / 0.001 = 6

The derivative of x squared at x equals 3 is 6. Meaning: at this point, if x increases by a tiny amount, the output increases 6 times as fast.

What is a gradient?

A derivative is for a function with one input. A gradient is the same idea for a function with many inputs. It is just a list of derivatives, one per input.

function with 1 input  ->  derivative  (a single number)
function with N inputs ->  gradient    (a list of N numbers)

Why does this matter for ML?

Every neural network has weights. Training means finding the weights that make the model least wrong. "How wrong" is measured by a number called the loss.

The gradient of the loss tells you: for each weight, does increasing it make the loss bigger or smaller, and by how much?

Nudge every weight in the direction that reduces the loss. Repeat thousands of times. That is gradient descent. That is literally how every neural network on the planet trains.

Now the code

First I verified grad by hand:

import jax

def square(x):
    return x * x

grad_of_square = jax.grad(square)
print(grad_of_square(3.0))  # 6.0

Derivative of x squared is 2x. At x equals 3, that is 6. JAX says 6.0. Verified.

Then I wrote a real loss function:

def loss(w):
    x = 2.0
    target = 10.0
    prediction = w * x
    return (prediction - target) ** 2

grad_of_loss = jax.grad(loss)

This loss measures how wrong a prediction is. The ideal weight is 5.0 because 5 times 2 equals 10.

Then gradient descent from scratch:

learning_rate = 0.1
w = 1.0

for step in range(5):
    grad = grad_of_loss(w)
    w = w - learning_rate * grad
    print(f"step {step + 1}: w = {w:.4f}, loss = {loss(w):.4f}")

Output:

step 1: w = 2.6000, loss = 21.1600
step 2: w = 3.6640, loss = 8.6491
step 3: w = 4.3302, loss = 3.5369
step 4: w = 4.7481, loss = 1.4461
step 5: w = 5.0105, loss = 0.0044 (almost zero, w approaching 5.0)

That loop is the core of every neural network training run. More weights, more complex functions, but the same mechanic.

The jit speedup

I benchmarked a large computation with and without jit over 100 runs:

def slow_computation(x):
    return jnp.sum(jnp.sin(x) ** 2 + jnp.cos(x) ** 2)

fast_computation = jax.jit(slow_computation)

Result: 32x speedup on my MacBook CPU. On a GPU this goes to 100x or more.

The first call is always slower because that is when JAX compiles. Every call after uses the cached compiled version. One catch: if you change the shape of your input, JAX recompiles. So keep your array shapes consistent.

One question I had: do I need to memorise all this syntax?

No. The important thing is knowing what to use and when. Nobody in a research role writes jax.grad from memory every time - they know it exists, they know what it does, they look up the exact syntax when needed.

What you do need to carry in your head: grad differentiates, jit compiles, and they compose together. That mental model handles 90% of situations.

Day 3 tomorrow: vmap, the fourth JAX superpower. This one genuinely surprised me.

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)

I Started Learning JAX as a Senior ML Engineer - Here's My First Impression.

Madhumitha Kolkar — Mon, 20 Jul 2026 05:41:20 +0000

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

I have been a machine learning engineer for six years. PyTorch is basically muscle memory at this point. So when I decided to seriously learn JAX, I expected it to feel like switching keyboards - same idea, different layout.

It is not that. JAX is a genuinely different way of thinking about numerical computing. And on Day 1, I want to share the mental model that made it click for me.

If you already know NumPy, you already know 80% of JAX's syntax. That part is intentional. The JAX team designed the API to mirror NumPy almost exactly so the learning curve stays shallow.

# NumPy
import numpy as np
x = np.array([1.0, 2.0, 3.0])
print(np.sum(x))  # 6.0

# JAX - identical
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
print(jnp.sum(x))  # 6.0

Same thing. Where JAX diverges is in three superpowers that NumPy simply does not have.

Superpower 1: It runs on GPU and TPU automatically

No .to('cuda'). No .cuda(). JAX looks for available accelerators and uses them. You write the same code whether you are on a laptop CPU or a cloud TPU. This alone is huge if you work across different environments.

Superpower 2: jax.grad - automatic differentiation as a function

In PyTorch, you call .backward() on a loss tensor. In JAX, grad is a standalone function transformer. You hand it any Python function, it hands you back a new function that computes the gradient of the original.

import jax

def square(x):
    return x ** 2

derivative = jax.grad(square)
print(derivative(3.0))  # 6.0

The derivative of x squared is 2x. At x equals 3, that is 6. JAX confirms it instantly. The fact that this works on any arbitrary Python function is what makes JAX powerful for research.

Superpower 3: jax.jit - compile your code for speed

Python is slow by default. jit (Just-In-Time compilation) takes your function, compiles it to optimised machine code on the first call, and every subsequent call skips Python entirely.

fast_fn = jax.jit(my_function)

First call is slower because of compilation. Every call after that is significantly faster. On Day 2 I benchmarked this and got a 32x speedup. More on that in the next post.

The one rule that surprised me: no in-place mutation

PyTorch lets you modify arrays in place. JAX does not. Every operation creates a new array. This feels wrong for about two days. Then you realise it is what makes grad and jit mathematically clean - they need functions to be pure (same input always gives same output). Mutation breaks that.

# This errors in JAX
x[0] = 99

# Do this instead
x = x.at[0].set(99)  # creates a new array

Why top research labs use JAX

The more I dig into JAX, the more I understand why the most ambitious ML research happening today is built on it. Labs like Google DeepMind run their entire research stack on JAX and Flax, and it is not arbitrary. When you are training models across hundreds of TPUs, you need code that is compiled, mathematically clean, and composable at scale. The combination of jit and grad doing exactly that is genuinely elegant. I find myself wanting to understand not just how to use these tools but why they were designed this way, and that curiosity is pushing me to go deeper into the fundamentals than I ever did with PyTorch. There is a certain level of ML engineering where understanding the infrastructure is just as important as understanding the models, and JAX sits right at that intersection.

Day 1 was arrays and basic operations. Simple stuff. But getting the mental model right before writing complex code matters more than rushing into the deep end.

Day 2 is grad, jit, and a gradient descent loop from scratch. See you there.

If you are on a similar journey from PyTorch to JAX, I would love to hear what tripped you up first. Drop it in the comments.

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)