Compare Count and Uncertainty Bonuses in a Small Reinforcement Learning Task

You add a curiosity bonus. Coverage climbs. The return curve goes flat. That gap is the whole lesson: a bonus is a second reward channel, and the only question that matters is whether the states it rewards carry task information.

You already know what each bonus family rewards from the earlier comparison. Here we make them compete on the same task, with the same seed, the same episode budget, and the same learning rate. The only variable is the bonus.

What You Are Actually Adding When You Add a Bonus

An exploration bonus is not a mood. It is a term in the objective.

r_total = r_task + beta * b(s)

The agent optimizes r_total. If beta is large enough, the agent will happily chase b(s) while the task reward sits in the corner, ignored. That is not a bug in the agent. That is the agent doing exactly what you asked.

For this experiment we use two transparent proxies:

• Count bonus: decays as a state is visited, roughly beta / sqrt(N(s)). Rewards raw visitation.
• Uncertainty bonus: built from disagreement across a small ensemble of predictors. Rewards states where the agent's own estimates are unstable.

These two diverge exactly where it matters. A state can be rarely visited and still boring. A state can be visited often and still confusing. Counts cannot see that difference. Disagreement can.

One assumption up front: this is a tabular toy environment, so counts are exact and the ensemble is cheap. Both proxies get harder in high dimensions. That difficulty is part of the lesson, not a footnote.

The Toy Task and Why It Is Built to Expose the Difference

We need an environment where novelty and usefulness come apart. A small gridworld does it.

• A sparse goal reward. The agent gets nothing until it reaches the goal.
• A dead-end region that is easy to reach, high in novelty, and useless for the task.

That dead-end region is the trap. It separates a bonus that rewards surprise from a bonus that rewards information.

Hold everything else fixed across three runs: no bonus, count bonus, uncertainty bonus. Same seed, same episode budget, same learning rate, same epsilon schedule. If you change two things, you learn nothing.

Define your measurements before you run anything:

• Task return per episode. Did the agent get better at the task?
• Task-relevant discovery. The fraction of episodes in which the agent reaches the goal-adjacent informative states. Not the fraction of states visited. That distinction is the point.

Note: A single toy task is evidence about this mechanism, not a benchmark. Say so plainly when you report it.

Implementing the Count Bonus

The count bonus is a dictionary and one line of arithmetic.

visit_counts = {}

def count_bonus(state, beta=0.1):
    visit_counts[state] = visit_counts.get(state, 0) + 1
    return beta / (visit_counts[state] ** 0.5)

Plug it in before the Q-update:

bonus = count_bonus(state)
r_total = r_task + bonus
# Q-update uses r_total, not r_task

Why the square root? It keeps the bonus from collapsing to zero after a handful of visits while still decaying. Early exploration stays aggressive. Later exploration fades. A linear decay would kill the signal too fast; no decay would drown the task reward forever.

After 50 episodes, print the visit-count table and the bonus value for a few states. You should see the goal-adjacent states still carrying a bonus while the dead-end region has gone quiet. That is the count bonus working as designed.

And that is also its limitation. This bonus is a pure function of visitation. It cannot distinguish a useful new state from a useless new state. For coverage, that is a feature. For progress, it is a bug.

Implementing the Uncertainty Bonus

Now build a structurally different bonus. Train a small ensemble of predictors — three to five members — on observed next-state or reward targets. Use the spread of their predictions as the bonus.

def uncertainty_bonus(state, ensemble, beta=0.1):
    preds = [m.predict(state) for m in ensemble]
    disagreement = np.var(preds)
    return beta * disagreement

Why disagreement is a usable uncertainty proxy: where the ensemble agrees, the agent has seen enough similar experience. Where it disagrees, its model is still unstable. The bonus points at the unstable states.

Print the per-state disagreement for the same states you printed for counts. The dead-end region should show low disagreement after a few visits even though its count is also low. That is the divergence you are looking for.

Warning: An ensemble can be confidently wrong. Disagreement can spike on stochastic transitions that carry no task information. Uncertainty is a proxy, not a guarantee.

Keep the implementation deliberately small. The goal is a transparent proxy you can reason about, not a production intrinsic-motivation stack.

Running the Comparison and Reading the Curves

Plot task return and task-relevant discovery on the same axes for all three agents.

Expect the count bonus to win on coverage and sometimes lose on return. The signature result to look for: a bonus agent that visits more states but reaches the goal no more often than the no-bonus baseline. That is novelty without progress, made visible.

The mechanism behind that outcome is simple. The bonus term competes with the task reward. beta controls how much of the agent's attention the second channel steals.

Common mistakes to name before you make them:

• Tuning beta until the bonus dominates the task reward.
• Comparing runs with different seeds or episode budgets.
• Reading a single lucky run as a result.

When the Bonus Helps and When It Just Adds Noise

The bonus helps when the task reward is sparse and the informative states are reachable but rarely stumbled into by epsilon-greedy exploration. It hurts when the environment contains large regions of high novelty and low task relevance, or when the bonus scale is not decayed as the agent learns.

Decision rule: before adding a bonus, ask what information the rewarded states carry about the task. If you cannot answer that, you are buying coverage, not competence.

Common mistake: Assuming that more states visited means more learning happened. Coverage is a measurement of where the agent went, not what it understood.

This experiment stops being exact when counts no longer exist. In continuous or high-dimensional state spaces, exact counts are unavailable, and the count-versus-uncertainty distinction blurs once you move to learned representations. The deeper version of the problem — propagating uncertainty across multiple steps rather than one — is a separate topic.

One Modification Worth Trying Next

Add a decay schedule to beta so the bonus fades as episodes accumulate. Re-run all three agents and compare against the fixed-beta results.

Prediction to check: the bonus agent should keep its early coverage advantage while recovering the return it lost to a permanently loud bonus channel.

If you want a sharper test, shrink the useless region and see whether the count bonus's coverage advantage stops translating into any measurable difference.

Before you run either version, write down which curve you expect to move and why. A prediction that fails teaches you more than a plot that merely looks fine.

The habit worth keeping: a bonus is a second reward channel, and its value depends entirely on whether the states it rewards carry task information. Ask that question before you tune beta.