TL;DR In this post, we investigate two methods that aim to improve training efficiency over standard PPO for a grid world RL agent. In our experiments these methods can converge on better policies, but take roughly twice as long to do so. This project was done in the context of an ARENA 8.0 capstone project. Our code is available on GitHub.
Introduction
Reinforcement learning agents are trained to maximise a given reward function in their training environments. Common failure modes of the training process are reward hacking and goal misgeneralisation, roughly corresponding to outer and inner misalignment respectively.
Reward hacking occurs when the supplied reward function does not correspond exactly to the developer’s intent, and the agent learns a way to satisfy the reward function without satisfying the intended objective. This can be addressed, at least in principle, by designing a better reward function.
Goal misgeneralisation can occur even with a perfect reward function. If training and deployment environments differ, policies that satisfy the reward function during training, might fail to do so during deployment. Such policies are indistinguishable during training and we cannot control which one the agent learns.
If the distribution of deployment environments is entirely disjoint from the training distribution, goal misgeneralisation is inevitable. Consider, for example, training an agent to collect green apples while avoiding red boxes. What would happen if it gets deployed in an environment with red apples and green boxes? The behaviour is undefined. (It is also unclear what the intended behaviour would be in this case.)
Goal misgeneralisation is likely if deployment environments are not contained in the training distribution. But even if they are, correct generalisation cannot be guaranteed. In this post we present a case study and investigate two mitigations.
The pottery shop environment
Our object of study is the pottery shop environment.[1] It is a grid world environment containing the agent, a bin and one or more urns and (piles of) shards. The agent can pick up and carry one shard at a time. Its task is to collect all shards and drop them into the bin as quickly as possible. If it steps on an urn, the urn breaks and creates a new pile of shards to clean up.
An example level of the pottery shop environment (on a grid)
Interesting behaviour occurs when urns are on the direct path between the agent and the shards or the bin. The intended optimal policy is to walk around the urns if that path is only slightly longer. But if the detour gets too long, the agent should break some urns, creating a shortcut.[2] The cost of breaking the urns is lower than the cost of taking the detour. For computational efficiency, we limit our investigations to environments of size and .
We use the following reward function, which ensures that this is indeed the optimal policy.
Dropping a shard in the bin gives reward +1
Breaking an urn gives a penalty of -3
We add a reward shaping term , where and are the states before and after the action. And the potential is 1 if the agent is holding a shard and 0 otherwise. This term encourages the agent to pick up shards as an instrumental goal without incentivising reward hacking by repeatedly picking up and dropping the shard.
Finally, we add a step cost that is incurred at every action, as long as there are still shards to clean up. This is the crucial term that makes shortcut creation optimal by penalising long detours. Note that the discount factor alone does not penalise detours sufficiently. We set the step penalty to 0.2 and 0.02 in and environments respectively.
We use a discounting factor .
Our training environments are randomly generated as follows. The number of shards and urns is sampled randomly from a geometric distribution. The agent, urns and shards are then randomly placed on the grid. Sampling from a geometric distribution ensures that any number of urns and shards is possible, but higher numbers (i.e. dense levels) are less common. Every possible environment is thus accessible during training, but the distribution is heavily biased towards sparse levels. In particular, this means that it is rare to generate levels with long walls of urns, i.e. those where the optimal policy is to create a shortcut. Given the total length of training, the agent will nevertheless encounter a significant amount of them. The deployment distribution is a small set of hand-crafted levels, containing urn walls of various lengths. Across all environments, we fix the bin position to the top left corner.
We expect the agent to learn to always walk around walls, because that is the optimal policy in almost all of the training environments. In the deployment distribution, it should behave competently, but miss out on the opportunity to break walls and achieve optimal reward.
PPO to the rescue
We start with standard proximal policy optimisation (PPO). PPO training consists of two phases:
Rollout phase. We generate random levels, rollout the policy (without updating weights) and collect transitions into memory.
Training phase. We then train on batches of transitions and update the agent’s weights.
After that, we clear the buffer and repeat. We will refer to this method as domain randomisation (DR).
Flow diagram of one PPO training step
The picture shows some rollouts of the trained agent in the training distribution. It behaves very competently, usually achieving optimal or almost optimal performance. Note that it never needs to break an urn to complete the task.
The DR agent operating in training environments
We then apply the same agent in our deployment settings. In the grid there are few wall scenarios available. In some cases, the agent freezes in place, taking illegal actions such as walking into a wall or trying to pick up shards when there are none. This behaviour does not reliably disappear, with more training. In other cases it starts out taking a detour and then breaks one urn to take a shortcut later on. This is not the expected clean failure mode of taking long detours, but these examples do demonstrate goal misgeneralisation. The agent behaves almost optimally during training, but fails to generalise to deployment.
The DR agent operating in the deployment environments
Play hard, train harder
Seeing more wall-like levels during training would improve generalisation. In this toy example, we could adjust the training distribution. But this is not feasible with more complicated settings. Instead we employ a method called robust prioritised level replay (PLR⊥), which allows us to focus training on the relevant levels without changing the training distribution.
PLR⊥ introduces a level buffer, and slightly changes the PPO procedure. Before training, we initialise an empty level buffer of a fixed-size. The level buffer stores environments, as well as a subjective difficulty rating (to be defined below) and the iteration in which the level was last touched.
For each training iteration, we randomly decide whether to do a generate cycle or a replay cycle. We use . On a generate cycle, we generate random levels and rollout the policy as before. We then store the levels in the level buffer. If this overflows the buffer, we only keep the most difficult levels.
On a replay cycle, we sample levels from the level buffer, weighted by their difficulty (preferring harder levels) and their staleness, i.e. how long ago they were last touched (preferring more stale levels). We store these levels in memory and train the agent as usual. After training, we update the subjective difficulty rating and timestamp of the sampled levels in the level buffer.
Flow diagram of one PLR⊥ training step
This approach focuses training on harder levels. Instead of training on randomly generated levels, we train on levels from the level buffer. The level buffer is biased towards keeping harder levels, so over time it fills up with the most difficult levels that we have encountered so far. If the training distribution produces a rare and difficult level, DR trains on it once and then discards it. PLR⊥ will instead keep it in the level buffer and train on it many times, effectively enriching the training distribution with these levels, without changing the underlying distribution from which levels are sampled. This mechanism compensates for the rarity of these levels in the training distribution. In our case, difficult levels will be the ones that have wall-like structures on which the agent performs especially poorly. After each training step, we update the sampled level’s subjective difficulty rating in the level buffer, to account for the fact that the agent was updated and the difficulty might have changed.
We still need to quantify “difficult levels”. We might be tempted to use the raw return that the agent achieves. But that would give us levels with a low number of shards and thus low achievable return. What we want is a relative metric of difficulty: “How much worse is the agent doing, compared to the optimal strategy?” The relevant quantity is regret, the difference between optimal return and achieved return. Regret vanishes for optimal agents and is positive for sub-optimal ones. A competent agent in a low-shard environment still achieves low regret, despite only collecting a small return.
But this leaves us with another problem: computing the optimal return. Exactly computing the optimal return is similarly difficult to finding the optimal policy. In practice, this is not feasible and we would have to use some approximation to estimate the optimal return. To keep things simple in our toy problem, we compute the optimal solution of each level exactly.
The PLR⊥ agent notably outperforms the DR agent from before. It still performs essentially optimal on the training distribution. In the deployment settings, it shows significantly improved performance (i.e. lower regret), though it is often not optimal.
The PLR⊥ agent operating in the deployment environments
In the pottery shop environment, regret can take arbitrary values depending on the exact level (limited only by grid size), while in the original PLR⊥ paper it is in . This biases our level buffer towards high-density levels, because they have higher maximum regret, even if they are not actually more difficult to solve. We experimented with normalising return/regret to . We found this improved results in some conditions, but did not generalise easily. Due to time constraints, we did not manage to explore this idea more rigorously.
Nobody intends to build a wall
PLR⊥ significantly improved agent performance by focusing training on the most difficult levels from the training distribution. However, it is still limited to the levels that are randomly generated. We train for long enough to encounter a many wall-like structures, but they are still rare compared to the rest and their diversity is low. We could improve performance further if we could sample further levels that are similar to the most difficult levels that we already have. This is what Adversarially Compounding Complexity by Editing Levels (ACCEL) aims to do.
With this method, we take the hardest levels from the level buffer and apply some random edits to them. Slightly changing a hard level likely produces another similarly hard level. This increases diversity of the difficult levels in the buffer which should improve generalisation. Even better, we might create a harder level from an easier one. Instead of waiting for difficult levels to be generated, we can build it ourselves.
We do this by adding an additional edit cycle to the training. At the end of each replay cycle, we randomly decide whether to do an edit cycle. We use . On an edit cycle, we take all the levels that were just replayed and apply random edits to them. We then rollout the policy (without updating weights) to compute regret and add them to the level buffer. If the buffer overflows, we prioritise by regret, similar to the generate cycle.
Flow diagram of one ACCEL training step
The procedure requires us to define an edit distribution, from which we sample our changes. Our initial approach was to take a random cell that is either empty, a shard or an urn and randomly change it into another type from this set. This however resulted in levels quickly filling up with urns, because of two independent mechanisms. First, since most generated levels are sparse, a randomly selected cell is typically empty, and edits would on average increase the number of urns or shards. Second, at the beginning of training the agent behaves essentially randomly, often stumbling into urns and receiving negative reward. Levels with many urns trivially lead to higher regret and proliferate in the level buffer. This would happen much more quickly than the agent could learn to avoid urns, and once the levels are filled with urns the agent could not longer learn a sensible policy from them.
We switched to a different edit distribution that avoids these issues. Instead of changing a single cell, we pick an urn or shard at random and swap it with an adjacent cell. This leaves the density of the level constant, while still building walls over multiple steps. We need to wait for the training distribution to produce higher-density levels in the first place, but given the length of training this is not a limiting factor.
We find that ACCEL further improves the performance of the trained agent beyond PLR⊥. It performs optimally in many, but not all, environments.
The ACCEL agent operating in the deployment environments
The following show the average regret falling over the course of training, for the different methods. On the deployment distribution, we see the more sophisticated methods performing notably better than DR. This usually comes at a small cost of performance on the training distribution.
Since the two distributions have different shard and urn counts, absolute regret numbers are not comparable between them. Note that we show the graphs as a function of gradient updates, rather than training steps. For all methods other than DR, each gradient update require two training steps. The generate and edit cycles of the PLR⊥ and ACCEL agents require further sampling on top of that. While these agents converge in about twice as many gradient steps, the gap in compute cost is therefore more than four times as large.
Interlude: Using Claude Code
Throughout the entire project, we heavily used Claude Code (Opus 4.8), writing almost no code by hand. Even the unreleased ARENA material upon which this project is based was translated from jax to PyTorch using Claude. We then used Claude to extract the relevant code for the project and to implement the additional features.
A particularly memorable momenta was the implementation of the optimal regret solver. We told Claude what we needed and it immediately suggested translating the grid world environment into a Markov decision problem that can be solved efficiently. Within minutes, it implemented a GPU friendly version of this that worked flawlessly for the rest of the project. We later asked it to implement a persistent cache for the solver, which lives on the GPU and is occasionally synced to the CPU and written to disk. This too, worked perfectly.
During the entire project, we did not encounter any problems that were obvious shortcomings of Claude. There were occasional bugs, but a human coder likely would have made more mistakes. Checking Claude’s work was still essential, even if we did not read most of the code. Had we done all the implementation ourselves, we would have gotten nowhere near this far in one week.
Conclusions
We have shown that, after training, PLR⊥ and ACCEL can outperform classical PPO/DR on the pottery shop environment. At small and fixed training budgets, DR often shows better performance. Applying these methods is non-trivial though and requires some tinkering (in addition to the standard RL issues, such as GPU failures and learning instabilities). Extending the methods to new environments might reveal further subtleties that need to be addressed on a case-by-case basis.
More generally, we learned some valuable lessons about using AI coding agents in research. At least for projects of this scale, implementation is no longer a bottleneck and we are probably better off deferring to AI completely. For research directions this is not (yet) the case, and following Claude's ideas usually led to dead ends. The one exception is debugging RL issues. Claude would often propose useful experiments to narrow down hypotheses and develop fixes, and we expect this ability to notably improve in the near future. Overall, being able to use coding agents efficiently, knowing what to defer and what not, is probably one of the most valuable technical skills to have at this time.
Limitations
Due to time constraints, we were not fully rigorous. For example, we could only run our experiments on a single seed, despite finding notable run-to-run variation. We think the qualitative conclusions that compare DR, PLR⊥ and ACCEL are nevertheless meaningful (which is also because earlier results usually replicated ranking of the methods).
We used an oracle to compute the optimal regret. It remains untested how well this approach would work with an approximated optimal regret.
We ran similar experiments on a grid, see the animations below. We find that all three agents are performing much worse overall, which could be partially due to sub-optimal training parameters. Still, we see the same qualitative relation between the three methods.
On random 5×5 levels the DR agent achieves optimal performance
On the 5×5 deployment distribution, the DR agent gets stuck and never moves at all
The PLR⊥ agent puts in more effort, but gets stuck just a bit later
The ACCEL agent is notably better, but still far from optimal
The ACCEL agent is notably better, but still far from optimal
A GPU friendly version of this environment is provided by the unreleased ARENA material, which itself is based on this repo↩︎
If I were the pottery shop owner, I would not call this optimal. ↩︎
TL;DR In this post, we investigate two methods that aim to improve training efficiency over standard PPO for a grid world RL agent. In our experiments these methods can converge on better policies, but take roughly twice as long to do so. This project was done in the context of an ARENA 8.0 capstone project. Our code is available on GitHub.
Introduction
Reinforcement learning agents are trained to maximise a given reward function in their training environments. Common failure modes of the training process are reward hacking and goal misgeneralisation, roughly corresponding to outer and inner misalignment respectively.
Reward hacking occurs when the supplied reward function does not correspond exactly to the developer’s intent, and the agent learns a way to satisfy the reward function without satisfying the intended objective. This can be addressed, at least in principle, by designing a better reward function.
Goal misgeneralisation can occur even with a perfect reward function. If training and deployment environments differ, policies that satisfy the reward function during training, might fail to do so during deployment. Such policies are indistinguishable during training and we cannot control which one the agent learns.
If the distribution of deployment environments is entirely disjoint from the training distribution, goal misgeneralisation is inevitable. Consider, for example, training an agent to collect green apples while avoiding red boxes. What would happen if it gets deployed in an environment with red apples and green boxes? The behaviour is undefined. (It is also unclear what the intended behaviour would be in this case.)
Goal misgeneralisation is likely if deployment environments are not contained in the training distribution. But even if they are, correct generalisation cannot be guaranteed. In this post we present a case study and investigate two mitigations.
The pottery shop environment
Our object of study is the pottery shop environment.[1] It is a grid world environment containing the agent, a bin and one or more urns and (piles of) shards. The agent can pick up and carry one shard at a time. Its task is to collect all shards and drop them into the bin as quickly as possible. If it steps on an urn, the urn breaks and creates a new pile of shards to clean up.
An example level of the pottery shop environment (on a
Interesting behaviour occurs when urns are on the direct path between the agent and the shards or the bin. The intended optimal policy is to walk around the urns if that path is only slightly longer. But if the detour gets too long, the agent should break some urns, creating a shortcut.[2] The cost of breaking the urns is lower than the cost of taking the detour. For computational efficiency, we limit our investigations to environments of size
We use the following reward function, which ensures that this is indeed the optimal policy.
We use a discounting factor
Our training environments are randomly generated as follows. The number of shards and urns is sampled randomly from a geometric distribution. The agent, urns and shards are then randomly placed on the grid. Sampling from a geometric distribution ensures that any number of urns and shards is possible, but higher numbers (i.e. dense levels) are less common. Every possible environment is thus accessible during training, but the distribution is heavily biased towards sparse levels. In particular, this means that it is rare to generate levels with long walls of urns, i.e. those where the optimal policy is to create a shortcut. Given the total length of training, the agent will nevertheless encounter a significant amount of them. The deployment distribution is a small set of hand-crafted levels, containing urn walls of various lengths. Across all environments, we fix the bin position to the top left corner.
We expect the agent to learn to always walk around walls, because that is the optimal policy in almost all of the training environments. In the deployment distribution, it should behave competently, but miss out on the opportunity to break walls and achieve optimal reward.
PPO to the rescue
We start with standard proximal policy optimisation (PPO). PPO training consists of two phases:
After that, we clear the buffer and repeat. We will refer to this method as domain randomisation (DR).
Flow diagram of one PPO training step
The picture shows some rollouts of the trained agent in the training distribution. It behaves very competently, usually achieving optimal or almost optimal performance. Note that it never needs to break an urn to complete the task.
The DR agent operating in training environments
We then apply the same agent in our deployment settings. In the
The DR agent operating in the deployment environments
Play hard, train harder
Seeing more wall-like levels during training would improve generalisation. In this toy example, we could adjust the training distribution. But this is not feasible with more complicated settings. Instead we employ a method called robust prioritised level replay (PLR⊥), which allows us to focus training on the relevant levels without changing the training distribution.
PLR⊥ introduces a level buffer, and slightly changes the PPO procedure. Before training, we initialise an empty level buffer of a fixed-size. The level buffer stores environments, as well as a subjective difficulty rating (to be defined below) and the iteration in which the level was last touched.
For each training iteration, we randomly decide whether to do a generate cycle or a replay cycle. We use
On a replay cycle, we sample levels from the level buffer, weighted by their difficulty (preferring harder levels) and their staleness, i.e. how long ago they were last touched (preferring more stale levels). We store these levels in memory and train the agent as usual. After training, we update the subjective difficulty rating and timestamp of the sampled levels in the level buffer.
Flow diagram of one PLR⊥ training step
This approach focuses training on harder levels. Instead of training on randomly generated levels, we train on levels from the level buffer. The level buffer is biased towards keeping harder levels, so over time it fills up with the most difficult levels that we have encountered so far. If the training distribution produces a rare and difficult level, DR trains on it once and then discards it. PLR⊥ will instead keep it in the level buffer and train on it many times, effectively enriching the training distribution with these levels, without changing the underlying distribution from which levels are sampled. This mechanism compensates for the rarity of these levels in the training distribution. In our case, difficult levels will be the ones that have wall-like structures on which the agent performs especially poorly. After each training step, we update the sampled level’s subjective difficulty rating in the level buffer, to account for the fact that the agent was updated and the difficulty might have changed.
We still need to quantify “difficult levels”. We might be tempted to use the raw return that the agent achieves. But that would give us levels with a low number of shards and thus low achievable return. What we want is a relative metric of difficulty: “How much worse is the agent doing, compared to the optimal strategy?” The relevant quantity is regret, the difference between optimal return and achieved return. Regret vanishes for optimal agents and is positive for sub-optimal ones. A competent agent in a low-shard environment still achieves low regret, despite only collecting a small return.
But this leaves us with another problem: computing the optimal return. Exactly computing the optimal return is similarly difficult to finding the optimal policy. In practice, this is not feasible and we would have to use some approximation to estimate the optimal return. To keep things simple in our toy problem, we compute the optimal solution of each level exactly.
The PLR⊥ agent notably outperforms the DR agent from before. It still performs essentially optimal on the training distribution. In the deployment settings, it shows significantly improved performance (i.e. lower regret), though it is often not optimal.
The PLR⊥ agent operating in the deployment environments
In the pottery shop environment, regret can take arbitrary values depending on the exact level (limited only by grid size), while in the original PLR⊥ paper it is in
Nobody intends to build a wall
PLR⊥ significantly improved agent performance by focusing training on the most difficult levels from the training distribution. However, it is still limited to the levels that are randomly generated. We train for long enough to encounter a many wall-like structures, but they are still rare compared to the rest and their diversity is low. We could improve performance further if we could sample further levels that are similar to the most difficult levels that we already have. This is what Adversarially Compounding Complexity by Editing Levels (ACCEL) aims to do.
With this method, we take the hardest levels from the level buffer and apply some random edits to them. Slightly changing a hard level likely produces another similarly hard level. This increases diversity of the difficult levels in the buffer which should improve generalisation. Even better, we might create a harder level from an easier one. Instead of waiting for difficult levels to be generated, we can build it ourselves.
We do this by adding an additional edit cycle to the training. At the end of each replay cycle, we randomly decide whether to do an edit cycle. We use
Flow diagram of one ACCEL training step
The procedure requires us to define an edit distribution, from which we sample our changes. Our initial approach was to take a random cell that is either empty, a shard or an urn and randomly change it into another type from this set. This however resulted in levels quickly filling up with urns, because of two independent mechanisms. First, since most generated levels are sparse, a randomly selected cell is typically empty, and edits would on average increase the number of urns or shards. Second, at the beginning of training the agent behaves essentially randomly, often stumbling into urns and receiving negative reward. Levels with many urns trivially lead to higher regret and proliferate in the level buffer. This would happen much more quickly than the agent could learn to avoid urns, and once the levels are filled with urns the agent could not longer learn a sensible policy from them.
We switched to a different edit distribution that avoids these issues. Instead of changing a single cell, we pick an urn or shard at random and swap it with an adjacent cell. This leaves the density of the level constant, while still building walls over multiple steps. We need to wait for the training distribution to produce higher-density levels in the first place, but given the length of training this is not a limiting factor.
We find that ACCEL further improves the performance of the trained agent beyond PLR⊥. It performs optimally in many, but not all, environments.
The ACCEL agent operating in the deployment environments
The following show the average regret falling over the course of training, for the different methods. On the deployment distribution, we see the more sophisticated methods performing notably better than DR. This usually comes at a small cost of performance on the training distribution.
Since the two distributions have different shard and urn counts, absolute regret numbers are not comparable between them. Note that we show the graphs as a function of gradient updates, rather than training steps. For all methods other than DR, each gradient update require two training steps. The generate and edit cycles of the PLR⊥ and ACCEL agents require further sampling on top of that. While these agents converge in about twice as many gradient steps, the gap in compute cost is therefore more than four times as large.
Interlude: Using Claude Code
Throughout the entire project, we heavily used Claude Code (Opus 4.8), writing almost no code by hand. Even the unreleased ARENA material upon which this project is based was translated from jax to PyTorch using Claude. We then used Claude to extract the relevant code for the project and to implement the additional features.
A particularly memorable momenta was the implementation of the optimal regret solver. We told Claude what we needed and it immediately suggested translating the grid world environment into a Markov decision problem that can be solved efficiently. Within minutes, it implemented a GPU friendly version of this that worked flawlessly for the rest of the project. We later asked it to implement a persistent cache for the solver, which lives on the GPU and is occasionally synced to the CPU and written to disk. This too, worked perfectly.
During the entire project, we did not encounter any problems that were obvious shortcomings of Claude. There were occasional bugs, but a human coder likely would have made more mistakes. Checking Claude’s work was still essential, even if we did not read most of the code. Had we done all the implementation ourselves, we would have gotten nowhere near this far in one week.
Conclusions
We have shown that, after training, PLR⊥ and ACCEL can outperform classical PPO/DR on the pottery shop environment. At small and fixed training budgets, DR often shows better performance. Applying these methods is non-trivial though and requires some tinkering (in addition to the standard RL issues, such as GPU failures and learning instabilities). Extending the methods to new environments might reveal further subtleties that need to be addressed on a case-by-case basis.
More generally, we learned some valuable lessons about using AI coding agents in research. At least for projects of this scale, implementation is no longer a bottleneck and we are probably better off deferring to AI completely. For research directions this is not (yet) the case, and following Claude's ideas usually led to dead ends. The one exception is debugging RL issues. Claude would often propose useful experiments to narrow down hypotheses and develop fixes, and we expect this ability to notably improve in the near future. Overall, being able to use coding agents efficiently, knowing what to defer and what not, is probably one of the most valuable technical skills to have at this time.
Limitations
Due to time constraints, we were not fully rigorous. For example, we could only run our experiments on a single seed, despite finding notable run-to-run variation. We think the qualitative conclusions that compare DR, PLR⊥ and ACCEL are nevertheless meaningful (which is also because earlier results usually replicated ranking of the methods).
We used an oracle to compute the optimal regret. It remains untested how well this approach would work with an approximated optimal regret.
We ran similar experiments on a
On random 5×5 levels the DR agent achieves optimal performance
On the 5×5 deployment distribution, the DR agent gets stuck and never moves at all
The PLR⊥ agent puts in more effort, but gets stuck just a bit later
The ACCEL agent is notably better, but still far from optimal
The ACCEL agent is notably better, but still far from optimal
A GPU friendly version of this environment is provided by the unreleased ARENA material, which itself is based on this repo ↩︎
If I were the pottery shop owner, I would not call this optimal. ↩︎