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.

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:

1. Rollout phase. We generate random levels, rollout the policy (without updating weights) and collect transitions into memory.
2. Training phase. We then train on batches of transitions and update the agent’s weights.

After that,