InfODist: online distillation with informative rewards

A workshop paper, not a main-conference one: it appeared at the Deep Reinforcement Learning Workshop at NeurIPS 2022, where NeurIPS’s own listing records it as a workshop poster. Led by Rahul Siripurapu at IARAI, with Vihang Patil, Kajetan Schweighofer, Thomas Schmied, Luis Ferro, Markus Holzleitner, Hamid Eghbal-Zadeh, Michael Kopp and Sepp Hochreiter. I am the fourth of ten authors.

The problem

Curriculum learning is the obvious way to teach an agent something hard: train on an easy version of the task, then move to a harder one. It underpins much of what deep reinforcement learning has achieved through self-play and domain randomisation.

It also generalises worse than it should. An agent trained to competence on the easy task often transfers badly to the harder one, and the accepted explanation was non-stationarity. Deep RL is non-stationary by construction — the policy changes, so the value targets change, and bootstrapping compounds both — and this shows up measurably as capacity loss: the rank of the features in the penultimate layer of the value network collapses. A low-rank feature layer means the network can no longer tell as many states apart as its parameter count should allow. It aliases them.

This paper’s argument is that the moving targets are only half the story. The other half is the inputs. Exploration, by design, drags the agent through a large and varied set of task-irrelevant states, and fitting all of them consumes capacity that later tasks will need.

The intuition is worth stating plainly. Exploration teaches you what to avoid, but modelling each failure mode in detail is usually wasted effort. It is enough to know an approach is wrong without learning precisely why.

The idea

Give the agent a reward that is informative: positive when it moves closer to the goal, negative otherwise.

The distinction from a merely dense reward matters, and the paper isolates it: a control condition handing out Gaussian noise at every step is worse than the original sparse reward. Density is not the active ingredient.

The concrete instrument is a proximity function borrowed from prior work in imitation learning, learned from a handful of expert demonstrations. It estimates an exponentially discounted temporal distance to the goal — near 1 at the goal, near 0 far away — so the difference in proximity between consecutive states measures progress. An ensemble of these functions supplies a variance term that is subtracted as an uncertainty penalty, discouraging the agent from trusting the estimate where the ensemble disagrees.

One adaptation matters. In imitation learning you have only the expert’s trajectories; here the goal termination condition is available, so the agent’s own successful runs fold into the regression target. That is what makes the method practical on procedurally generated environments, where a few demonstrations cannot cover the state space.

Around that sits online distillation. Train on the easy task, learn the informative reward from what worked, then retrain an agent from scratch under that reward before transferring. The retrained agent reaches the same competence having visited far fewer irrelevant states.

There is a deliberate concession. RUDDER, from the same group, redistributes reward under a return-equivalence constraint guaranteeing the underlying task is unchanged. InfODist drops that guarantee, arguing that for most real tasks nobody can write a reward capturing the desired behaviour anyway, so freedom to reshape it is worth more. The paper says so rather than glossing it.

Why it matters, and what it does not show

The testbed is MiniGrid. The easy task is two procedurally generated rooms; the six harder tasks add keys, boxes, balls and lava, and require actions — pick up, drop, avoid — the agent has never taken. Everything is regenerated each episode.

InfODist transfers best. Below it comes a method that reduces irrelevant exploration indirectly by starting the agent near the goal, then a method that attacks non-stationarity through offline distillation, then plain curriculum learning. The ordering supports the paper’s thesis: exploration is the lever.

The genuinely interesting result is the one at the bottom of that list. Regularising feature rank directly — the technique built to prevent capacity collapse — transfers worse than doing nothing. It succeeds at its stated job, keeping rank from collapsing, and still hurts. The proposed explanation is that a network can inflate its rank cheaply by memorising irrelevant state information. Feature rank is a good capacity measure only for as long as nobody optimises it.

The scale is modest and the paper does not pretend otherwise: three seeds per baseline, with the best seed transferred, in discrete grid worlds on goal-reaching tasks. The authors also note that the proximity function does not reliably produce students that outperform their teachers, and that a better way to learn informative rewards is left open. The curriculum itself is still written by hand.

Where it applies

The setting to reach for is one where you can demonstrate success a few times but cannot specify it, and where today’s task rehearses a harder one — robotic manipulation under a fixed capacity budget is the natural case.

The transferable claim is that what a bounded network spends its capacity on matters more than how much capacity it has, and that a per-step signal distinguishing progress from non-progress is a cheap way to steer that spending. The warning about feature rank travels further still, to anyone tracking a diagnostic and tempted to optimise it.