DEV Community

Cover image for I Added a Second Brain to My RL Agent - Day 7 (Actor-Critic)
Madhumitha Kolkar
Madhumitha Kolkar

Posted on

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

#python #machinelearning #deeplearning #ai

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode
  • 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)
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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)

Top comments (0)