DEV Community

Cover image for I Taught an Agent to Act Directly - No Q-Values Needed (Day 6: REINFORCE)
Madhumitha Kolkar
Madhumitha Kolkar

Posted on

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

#python #machinelearning #deeplearning #opensource

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

In REINFORCE, the network outputs probabilities :

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

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

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

Three key differences from DQN :

  1. No replay buffer - we use the episode as it happened
  2. No target network - no Bellman bootstrapping needed
  3. 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)

Top comments (0)