Planning to Learn
Ian Osband
cs.LG, math.OC, stat.ML
2026-10-03
Exact policy gradient hits 4.2% expected accuracy on ImageNet vs 62% for cross-entropy. Horizon loss, tied to budget left, adds 1.0-2.4 top-1 points, more under label noise.
Policy-gradient post-training samples an action and raises its probability when the reward comes back high. When that update struggles, the usual blame is exploration, credit assignment, or noise from sampling the action. Classification has none of the three. The label is known, the reward is an immediate right or wrong, and every class can be summed over, so the gradient of expected accuracy is exact.
It still loses to cross-entropy. On ImageNet with a ResNet-50 and one shared flat learning rate, exact policy gradient reaches 4.2% expected accuracy. Cross-entropy reaches 62%. The objective in both cases is expected accuracy.
A gradient only prices the next infinitesimal step. Take two examples with correct-class probabilities 0.9 and 0.01, and charge the same cost to raise either log-odds. That gap is the correct logit minus the log-sum-exp of the others. Probability rises with slope p(1-p), about 0.09 against about 0.01, so the exact gradient gives the already-confident example about nine times the weight. With almost no training left, that is the right spend. With a long run left, the budget belongs on the example near 0.01: it barely moves at first and has much more room. The value of a step depends on how much learning remains, and the exact gradient does not contain that number.
Drop the network and each example is one gap z, with correct-class probability p = σ(z). Exact policy gradient weights it by p(1-p). Cross-entropy weights it by 1-p. The direction matches. The split of the budget does not.
Cross-entropy is an integral along a made-up path. Send every later unit of learning to this one example, let z rise at unit speed, and add up the error 1-σ(z+s) forever. That total is the cross-entropy. The loss prices training as if it never stopped, so the largest weight falls on the examples furthest from success, including ones this run cannot repair.
Truncate the integral at the learning still left, H. The weight becomes σ(z+H)·(1-p). The new factor is the chance the example would be solved if it received the whole remaining budget. H = 0 recovers exact policy gradient. Large H recovers cross-entropy. A solved example has little error left. An example past the horizon has almost no chance of success inside it. Both weights fall.
A network does not announce how much logit progress is left. Estimate it by the learning rates still to come, times a calibration constant κ. Under a flat learning rate the sum falls linearly to zero. κ is set once per dataset and schedule: 1 on MNIST and 3 on ImageNet at a flat rate, 0.3 under cosine decay on both. The ImageNet value is chosen on ResNet-50 and reused for every other architecture and every noise level.
Two small allocation problems make the endpoint failures exact. At gaps (a, -2a) with budget H = 2a, the optimum spends everything on the hard example. Exact policy gradient has no H in its weights and cannot tell a positive gap from a negative one, so it cannot make that reversal. It keeps only a fraction 4σ'(a) of the available gain, and that fraction dies exponentially in a. Seat one finishable example beside M lost causes outside the horizon, and cross-entropy gives the finishable one at most 1/(M+1) of the budget. Horizon-loss regret stays within 2ae^{-a} and HMe^{H-a}. At a fixed H the expression is old: H = 0 is categorical absolute error, and generalized cross-entropy mixes horizons. The new piece is the schedule. H is the budget still left, and it shrinks as training proceeds.
Only the loss changes. MNIST is a two-hidden-layer MLP, 30 seeds. ImageNet is 30 epochs of ResNet-50, ResNet-101, and ViT-S/16, with 3 to 10 seeds. The default is a flat learning rate, so the schedule itself does not encode time remaining.
| Setup | Metric | Result |
| ResNet-50, learning rate 10^{-3} | Expected accuracy | Exact gradient 4.2%, cross-entropy 62%, horizon loss 67.5% |
| ImageNet, three models, flat rate | Top-1 | Error 1.0 to 2.4 points below cross-entropy |
| Same | Expected accuracy | 4.7 to 5.4 points above cross-entropy |
| MNIST | Top-1 | Error 0.16±0.04 below cross-entropy |
| MNIST | Versus exact gradient | Expected error down 24%, top-1 error down 29% |
| ImageNet, cosine decay | Top-1 | Gain about 0.6 points, still ahead on every model |
| Flipped labels, ResNet-50 | Top-1 gain | 3.3 points at 10% noise, 7.3 at 50% |
| Flipped labels, MNIST | Top-1 gain | From 0.1 points to 2.7 |
| Horizon cross-entropy, κ=10 | Test NLL | MNIST 0.103 to 0.089, ResNet-50 1.467 to 1.364 |
Part of the expected-accuracy gain is sharpening. Probabilities get peakier, and the predicted class can stay put, so top-1 moves less.
On ResNet-50 a flat rate of 10^{-3} over about 150k steps starts at H near 450, well above the gaps still moving. Through most of training the horizon loss is nearly cross-entropy. The curves split in the last few percent of steps. MNIST starts near H = 10, already on the scale of its gaps, and splits earlier.
At 30% label noise on MNIST, the probability cross-entropy puts on the flipped label dips and then climbs as the net memorizes the corruption. The horizon loss follows for the first half, then pushes that probability to about a third of cross-entropy's final value. With clean labels, exact policy gradient has worse MNIST top-1 than cross-entropy. From 20% noise upward it is better. The horizon loss beats both at every noise rate tested. The ImageNet noise sweep has no exact-gradient arm.
On MNIST, constant and growing horizons, matched to the shrinking one on average H, manage at most a marginal gain and lose at small κ. Only the shrinking horizon is consistently better. On ResNet-50, freezing H for the whole run ties cross-entropy once H is at least 10, and H of 3 or less is clearly worse from the first step.
The same integral applied to cross-entropy targets finite-budget negative log-likelihood. Test NLL falls, while ImageNet expected error rises from 0.378 to 0.393. The accuracy-oriented horizon loss moves NLL the other way, to 0.122 on MNIST and 1.77 on ResNet-50. Cosine decay shrinks the NLL gap to 1.270 against 1.289. At κ = 1, with no retuning, expected accuracy still beats cross-entropy on every ImageNet model and schedule reported. Top-1 fails to do so only for ViT-S/16 at a flat learning rate.
On a prompt with binary reward, the policy gradient carries the same p(1-p) factor. DAPO and Dr. GRPO change how that gradient is spread across prompts. They do not use how many steps remain. Strip away exploration, credit assignment, and sampling noise, and the myopic weight still fails. The allocation bug is not special to reinforcement learning.
In a classifier the change is one line. The useful regime is a flat learning rate and dirty labels: 7.3 top-1 points on ResNet-50 when half the training labels are flipped. Clean ImageNet with cosine decay, the schedule people actually run, leaves about 0.6 points.
A hand-built recipe, cross-entropy first and exact policy gradient late, roughly ties the horizon loss on MNIST and slightly beats it on ResNet-50 under cosine decay. Most of the measured gain is that late move toward the exact gradient.
κ is tuned on test accuracy. There is no validation split. The κ = 1 rerun still improves expected accuracy, so the result is not only a test-set fit of κ, but flat-rate ViT-S/16 top-1 remains an exception.
The theory assumes orthonormal example gradients and moves each gap at unit speed on its own. The nets share parameters and train with Adam or AdamW. Nothing in the experiments separates a better split across examples from a change in effective step size late in the run. A late switch to exact policy gradient recovers most of the gain, so the exponential trap bounds are not isolated on ImageNet.
The loss is separable. Every example is priced as if the remaining budget were its alone. Identical examples break every separable rule. A planner that couples examples could do better, and this paper does not build one. κ is not derived from optimizer or representation dynamics either.
Classification computes the correct-class probability in closed form. Sampled-action reinforcement learning has to estimate both that probability and the value of further learning from partial feedback. How to plan to the remaining budget once exploration, credit assignment, and action noise are back is left open.