← Blog

Giving an RL agent a hunch about its world

Based on Augmenting Decision with Hypothesis in Reinforcement Learning (ICML 2024, with Hady W. Lauw). Paper · Code

Background

Value-based RL (DDPG, TD3) trains a critic that scores how good each action is in each state, and an actor that picks the action the critic likes best. It is sample-efficient and cheaper than learning a full model of the world, which is why it's a default for continuous control.

The paper starts from two weaknesses of this recipe.

Weakness 1: slow exploitation

The critic's estimates are noisy, and the noise is larger on states it hasn't visited much. The Bellman update takes a max over next actions, so unfamiliar, high-uncertainty states tend to get inflated values. The result: the agent is pulled toward the unknown even when a known state is just as good. I show this formally (optimistic estimates prioritize exploration) and in a toy world.

The toy world, MultiNormEnv, is a line from 0 to 600 where reward is a sum of bumps, with a few good peaks. After two million steps, TD3's critic still can't clearly separate all the peaks, and the policy sits in the first local peak for over 1.6 million steps.

state (0 to 600) reward bumps (schematic) TD3: stays here ~1.6M steps ALH: uses all peaks early
The toy environment. A standard critic is slow to learn which peaks are best, so the policy lingers in the first one it finds.

Weakness 2: sensitivity to critic bias

If the critic is trained on unrepresentative data, say starts confined to a narrow range, it is wrong in a way that the actor can't fix, because the actor only ever maximizes the critic. In the toy world a biased TD3 critic puts its maxima at the edges of the state space, and the policy wanders over the whole space with no target.

The idea: a hypothesis

A world model could provide additional information, but learning one introduces its own prediction errors and costs. So the paper asks for something weaker. A hypothesis is a compact vector h describing the environment, learned from sets of observed (state, action, reward) tuples. The actor and critic both take h as an extra input.

A hypothesis is judged on two properties. Validity: it explains the observations better than noise does. Robustness: the more observations support it, the stronger it is.

Part A: weak-model-based learning

A separate module M maps a batch of observations to a hypothesis vector. It is trained to reconstruct the observations from h (a denoising objective), with a diversity penalty that stops every local hypothesis from collapsing into one global h0. Crucially, M is trained separately from the actor and critic (gradients are detached), so the critic's gradients do not directly shape it. It still depends on observations collected by the agent. It is "weak" because it never has to predict the next state or reward, only to summarize what has been seen.

To avoid overfitting to one global hypothesis, each training step builds hypotheses locally: sub-sample small sets from the batch and compute one h per sub-sample.

Effect in the toy world: agents conditioned on the hypothesis reach the optimal behavior from very early in training, instead of after 1.6 million steps.

Replay buffer(s, a, r, s') M (weak model)summarizes batches h TD3 learneractor(s, h)critic(s, a, h) Environment action latest (s, a, r) h' = norm(h' + M(new)) no gradient flows back from actor or critic
M is trained on its own from the replay buffer. While acting, the agent nudges its hypothesis with what it just saw, with no weight update.

Part B: adaptive rollout

Under bias, even a learned hypothesis can be out of date. So at rollout time the agent updates its hypothesis with what it has just seen: 

h' = norm( h' + M(latest obs) )   # no gradient step
reset h' to the global h0 when the episode ends

The policy changes its behavior as h' changes, without changing any network weights. With this, ALH policies reach good behavior even when the critic's value estimates are still inaccurate.

Putting it together

ALH sits on top of TD3. Variants: ALH-g collects data using the global hypothesis; ALH-a collects data with the adaptive one. Model size stays close to TD3's.

On eight MuJoCo tasks (HalfCheetah, Hopper, Walker2d, Ant, Humanoid, Reacher and two pendulums), ALH is best on 6 of 8 and runner-up on the rest. Averaged over tasks with min-max normalization, both ALH variants score about 85 to 86, against 67 for TD3, 45 for DDPG, 37 for the model-based MBPO and 23 for PPO. Ablations show each piece pulls its weight: detaching M, localizing hypotheses and adapting at rollout.

Takeaway

You don't need a world simulator to help an RL agent. A cheap, separately trained summary of what it has seen, refreshed as it acts, is enough to speed up exploitation and soften the damage of a biased critic.