The technical pivot is simple: instead of the network outputting Q-values and using an argmax
to pick the "best" action, the network outputs a probability distribution.
q_values = network(state) # [2.1, 0.8, 1.4, 3.2]
action = argmax(q_values) # deterministic choice
probs = network(state) # [0.1, 0.3, 0.2, 0.4]
action = sample(probs) # stochastic sampling
By sampling from a distribution, the agent maintains natural exploration. The learning mechanism is a straightforward reward/penalty system: if an action leads to a high total reward, the probability of taking that action again increases.
The loss function handles this via the log probability of the action taken, scaled by the discounted return:
loss = -G_t * log π(a_t | s_t)
I put this to the test in a 4x4 gridworld. Unlike a DQN setup, this AI workflow requires no replay buffer and no target network because it doesn't rely on Bellman bootstrapping. The training loop is a clean cycle of collecting a full episode and updating parameters based on actual returns:
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)
Comparing the two from a performance standpoint:
Accuracy: REINFORCE is Monte Carlo based, meaning it uses actual observed returns. DQN uses estimates (Temporal Difference), which can be biased.Stability: REINFORCE suffers from high variance. A single outlier episode can swing the gradients wildly, making it noisier and often slower to converge than DQN.Versatility: REINFORCE handles continuous action spaces effortlessly, whereas DQN is locked into discrete sets.
This trade-off is exactly why Actor-Critic architectures exist—they essentially merge the two, using a value network to stabilize the policy network's variance. This is the fundamental logic that eventually leads to advanced LLM agent optimizations like PPO or GRPO.
Next LLM Benchmarking: Stop Celebrating 0.000 Scores →