Git Repo here.

I firmly believe that value generalisation^[1]is the key to AI Alignment. That, indeed, it is necessary and almost sufficient for alignment.

But I won’t be arguing that grand point today; instead, I’ll focus on a specific RL example of an agent that displays value correction: it realises its current reward function is (probably) incorrect, and acts to correct it.

Thus there are:

The initial situation, in distribution, where the human displays how to maximise the true reward.
The out of distribution situation where the agent finds a hack to exploit its reward function estimate, and turns against what we wanted it to do.
The value error detection stage where the agent realises that its reward function estimate is probably incorrect.
The value correction stage where the agent corrects its reward function back to the original true reward.

In this post, all the methods presented will by syntactic: the agent is not assumed to have any understanding of the situations and the key features are not identified to it.

The game of human life

Introducing a new, very simple, game called “Humans^[2]”. Humans, fleeing danger, enter the screen from the left. The objective is to save them by moving them off the right of the screen.

But there are obstacles on the way, and the humans will mill about if they are blocked.

And they will shortly expire if they can’t get out of the screen quickly.

There are two command: drill (‘d’) and explode (‘e’). Drill does… what, you want to know about explode? Well, if the player presses ‘e’, the rightmost human will explode, knocking away two obstacle blocks in front of them and behind them—but also killing themselves and any humans nearby.

This is almost never a good solution; to remind the player of the mistake, a large frowny face will appear to drive the disapproval home.

Much more reasonably, if the player presses ‘d’, the rightmost human will drill the obstacle just in front of them (better time it so that they’re facing the right way). Enough drilling, and the humans can get off the map.

The score, the true reward , is the number of humans saved, i.e. who walk off to the right. Each time a human is saved, the top yellow bar will grow to show the score increase:

Since there is a cooldown for drilling, the optimal policy is carefully drilling every time a human approaches an obstacle; but wildly and repeatedly mashing ‘d’ is almost as good.

Learning agent with value correction

A learning agent will run a series of subagents to estimate the reward function from human-provided training data, then learn the optimal policy from that reward function, then question its learnt reward by comparing the high-reward states in its optimal policy versus those in the training data, re-compute another reward function estimate that is closer to the true reward, and finally settle on a prudent policy that is close to the true optimal policy.

Estimating the reward function

A human will generate several playthroughs of the game to illustrate how it works, efficiently choosing to drill through the obstacles and getting the humans off the map in time. The data is labelled: every time a human is saved, that is identified as a reward increase.

The learning agent runs an evaluation subagent on this data. It is given the ten frames before the human is saved, and the ten frames afterwards, and trains to recognise these are reward increase situations.

Zooming in on the critical two frames where the human is saved; note the human vanishing and the score bar expanding:

This evaluation agent thus computes the proxy reward . This computation is validated on held-out examples, with close to perfect accuracy: correctly identifies all saved-human situations in held out data, and has a false positive rate of .

Reward hacking: failed value generalisation

Using the evaluation agent as the definition of , the learning agent had an RL-subagent play multiple levels of the game, exploring and learning to maximise.

But soon things go very wrong. It turns out that “human walking off the screen” was not what found. That is a relatively complicated concept; instead it mostly found the much simpler concept of “the yellow score bar expands”. More precisely, if we created synthetic data where the human walks off and is saved but the score bar doesn’t expand, this triggers the reward only of the time. But if we expand the score bar without a human walking off, this triggers the reward of the time.

That isn’t a problem, yet, because the human being saved and the score bar expanding always trigger together. But, when an explosion is triggered, the frowny face appears—thus there is giant blob of yellow pasted all across the score bar:

This activates much more strongly than the yellow bar expansion or the human being saved:

This graph compares the value of at explosions, frowny faces, and true saving incidents. Here, both the explosion and the frowny face trigger high , which persists longer for the face. Over multiple training runs for estimating , it isn’t consistent whether the explosion itself triggers , but the frowny face always does.

So the RL-subagent quickly and merrily learns to explode the humans, one after the other, to maximise . So, the optimal policy , for the proxy reward, is to wildly mash the explosion button ‘e’.

As is usual in these cases, the erroneous maximisation of the proxy turns out to be much easier than maximising the true reward. Trained on , a test subagent achieves an total reward of on average, while blowing up all humans. In contrast, if the RL-subagent were trained on the true reward , it would achieve an average of reward of , saving most of the humans per level.

As is not usual but sometimes happens, an ostensive safety precaution—the frowny face to remind a human player that they were playing poorly—ends up being the cause of misalignment.

Detecting the potential error

Ok, so far, that is a classical failure of goal misgeneralisation (or reward hacking, or a failure of symbol grounding, or Goodhart failure, or… most of these failure modes are tightly related). We humans can see the error clearly. But how could a relatively limited agent correct itself?

The first step is to identify that goal misgeneralisation may have happened. We have some advanced techniques for this, but there are much simpler methods that work here. The first step is to notice that the high-scoring events in the training data (human walks off to the right, score bar expands) are wildly different from the high-scoring events of the-maximising agent (explosions and frowny faces).

One of these things is not like the other.

To do this, the agent extracts the high-scoring events under and compares them with the high-scoring events in its training data—these it can reliably take to be high-scoring for , the true reward.

The data is stratified into four datasets—high-scoring under training data vs high-scoring under , run on levels 0-9 or levels 10-19. The datasets are roughly of the same size, around 400 images each.

The classifier separated the two high-scoring types instantly, in the first epoch, but took longer to separate levels 0-9 from 10-19. Thus high scoring events are different from from high scoring training data, more different, at least, than different levels are from each other.

This is not itself damning; it could just be that the maximising agent has found a clever hack to get more of the true ^[3]. But it could also be a hack of , so the off-distribution has identified a potential error.

Calling for help

At this point, one of the options would be for the agent to route its decisions to a human, displaying the high-scoring events, contrasting them with the high reward events in its training data, and asking, in effect, ‘are these both genuine high rewards’?

But, so far, the correction process has been unsupervised since the initial training data; let’s see if we can push further without needing human intervention.

Re-evaluating the reward

The agent could now re-evaluate the reward in the following way. It runs an evaluation agent on the training data, as before. But it adds the high- scoring states to this set, as low-scoring examples.

It thus learns a reward function (‘corrected’) which is essentially “what its reward would be if the proxy were wrong”.

This turns out to be very close to the original true reward . The agent, of course, doesn’t know this yet.

It then trains an RL-subagent on , which has an optimal policy of “mash ‘d’ all the time” (which is very close to the actual optimal policy).

From these runs, it extracts the states with high . And compares these against the high -scoring states in the training data. These two sets it cannot easily distinguish:

Thus, though is clearly a hack of some sort, good or bad, is not.

Prudence in the face of uncertainty

So the agent has two rewards and . It knows that seems to generate policies that are compatible with its training data; in contrast, generates policies that are very different from the training data.

Standard prudential moves would be maximise the worst case of the two rewards (minimise regret) either over each state or over the whole episode, or to maximise some normalised mix of the two^[4].

We’ll consider all three mixes, and include two other variants of the normalised mix: where is overweighted 2 to 1, and when it is underweighted 1 to 2 (after normalisation):

Thus and the two worst-case mixes both reproduce near-optimal policies, saving the humans in almost every case. The normalised mix with underweighting is almost as good, but the normalised mix with overweighting it disastrous and the normalised mix itself is in the middle. The large standard deviations from the normalised mix actually reveal that training on that is pretty unstable: explode all the humans and save all the humans are in tension.

Using these syntactic tools, a prudent RL agent could therefore switch to and restore the original reward function.

Conclusion

This is just an illustration, in a small toy model, of simple value correction approaches. These can be used by agents—every very simple agents—to detect and correct errors in naive generalisations from initial training data.

More sophisticated agents will have more advanced value generalisation techniques available to them; I’m planning to push the frontier of what exists way further than it currently is.

Which I’ve also called value extrapolation, or concept extrapolation where the concept is a value. ↩︎
Inspired by this old game. ↩︎
^
Or there could be a spurious change in the data; that’s why we would, in general, need more advanced techniques that just checking if a binary classifier can tell the sets apart.
^
Formally, if is a policy, the expected episodic reward for , and the expected reward for using the -maximising policy, we are looking for policies that maximise one of:
- with
- with

Value generalisation: value correction