Self-Attentional Credit Assignment in Reinforcement Learning
By: Nick Baliesnyi
Tags: Attention Mechanism, Artificial Intelligence, Reinforcement Learning, University Project

by Mykyta Baliesnyi, supervisor: Oriol Corcoll

GitHub repository

This is a study project for the Computational Neuroscience course at the University of Tartu.

Motivation

Humans and other intelligent animals learn about causality in the world pretty quickly by automatically assigning credit to the events and actions that we think caused something to happen.

For example, as a kid, when we first get burnt by fire, we quickly assign the feeling of pain to our action of putting the hand above the fire. When the delay between the past action and the current result gets larger, as in understanding that eating an unripe apple is what caused a gurgling and hurting stomach, assigning credit becomes harder — but we can still usually do it successfully.

But how exactly does the brain identify events that are relevant among all the millions things that could have happened in between?

Effective credit assignment in artificial intelligence has been a long-standing problem. If one of the current machine learning approaches was that person who ate an unripe apple, it could easily take several dozens of trials until it understands the connection between the events.

In this project, I would like to better understand how credit assignment might be implemented by modeling this process in the context of reinforcement learning (RL).

Methodology

Reinforcement Learning

Reinforcement learning is a type of machine learning, like supervised and unsupervised, but with a focus on learning through continuous interaction with an environment. Instead of learning to make the most accurate predictions on a static dataset, reinforcement learning is about learning, through trial and error, how to behave in a world so as to get the most reward.

Reinforcement learning is about continuous interaction

This paradigm is inspired by how humans and other intelligent animals learn, and so, many think of it as one of the most promising ways to understanding decision making and learning, and eventually developing general artificial intelligence (AGI).

Self-Attention

The self-attention mechanism is an alternative to recurrent neural networks, which allows learning from sequential data by adaptively “paying attention” to different moments of the past at each step. It often works much better for capturing long-term dependencies than recurrent neural network architectures such as LSTMs and GRUs. This is especially useful in the context of reinforcement learning, where interaction with an environment can last anywhere from hundreds to thousands of steps, and some event that happened a very long time ago can determine the agent’s reward in the present moment.

In this paper, the authors used a self-attentional neural network to predict the agent’s rewards just from seeing what the agent observes and does in the environment. They showed that using this attentional model to create a denser reward signal improved the sample efficiency of an arbitrary RL agent even in relatively novels environments. For this project, I’m going to re-implement the reward prediction model they used and apply it in a different environment to see if how generalizable this approach is. The key question I’d like to explore is:

Can a self-attentional neural network identify important past events?

Environment

For an environment, I chose the Mayan Adventure environment for MLAgents in Unity. This is a small 3-D virtual world with a set of traps and a cute cubic agent in a cool Indiana Jones hat.

The goal in the environment is to stay alive and obtain the golden statue of Maya. Along the way, there are buttons that can turn our Indiana Jones agent into one of two states: Rock or Wood. Turned into Wood, the agent is light enough to go over bridges but will burn the instant it touches fire. Turned into Rock, the agent can go through the fire untouched but if it steps on a bridge, it will collapse and the agent will fall down.

To make training easier and incentivize the agent, there are some randomly located diamonds that act as intermediate rewards. Stepping on the diamonds and getting the golden statue all give +1 reward, while falling off the area or getting burned by fire give a -1 reward, instantly finishing the episode. A tiny negative reward is given to the agent at each step to incentivize it to finish the episode faster.

For observations, the environment uses a vector of what’s called ray casts. Think of those as invisible lasers shooting in several directions and telling the agent what objects are around it and how far away. The observations were stacked together so that a single observation contains information from the 6 most recent timesteps. In addition, the agent has access to the state it is turned into, either rock or wood, and to its speed and direction of movement.

Approach

To answer the key question, I first trained a reinforcement learning agent to get moderately good at solving the environment. The agent is trained with the Proximal Policy Optimization (PPO) algorithm. PPO is a very common learning algorithm, being easy to tune and giving good performance with relatively little hyperparameter tinkering. I didn’t have to implement the learning algorithm, as it was provided by the MLAgents Python package.

Learning the basics — how to get the statue without falling off

The training process was organized in 4 levels of curriculum learning, with progressively higher complexity. This is an effective method of regularizing and speeding up the learning of the agent that tries to master a complex task. It’s also just a natural way to learn for humans, starting from the very basics, and moving on to more high-level and complicated things as you get better.

Learning more complex stuff — how to go through the fire without dying

Next, I implemented and trained the self-attentional model to predict rewards based on a sample of episodes of the agent interacting with the environment.

The architecture of the reward prediction model

For implementing and training the model, I used PyTorch. I modified the original model from the paper by replacing the convolutional block with a simple linear layer. We don’t need convolutional layers because the observations in the environment are in 1-D vector form.

Since most of the time, the agent gets the reward of zero(I treated the tiny negative reward at each step as zero), with few events giving positive or negative rewards, the dataset for the reward prediction task is extremely imbalanced. I used a weighted Sequential CrossEntropy loss with class weights of 0.499 for positive and negative rewards and 0.02 for zero rewards, in accordance with the original paper.

Since we feed the full episode to the model at once, we need to prevent the model from “cheating” by looking into the future to know which reward it should predict now. This is done by applying a mask to the output of the self-attention that zeroes out future attention at future timesteps, such that the model can only learn from the past.

Expectations

Ideally, we would like the model to identify past events of pressing the buttons as important when predicting the agent’s reward as it steps onto a bridge or into the fire. Specifically, the really important event should be around the button that was pressed the last.

However, our observation space includes the information about the full agent state, including whether the agent is currently Rock or Wood, following what’s called the Markov Property. This makes it really unnecessary to remember or attend to those past events, as all the information the agent needs, was with it in the most recent moment. Because of this, the model might pay most of its attention to just the recent moment, and we would need to artificially limit what information the model has access to in order to force it to look into the past.

Results

Models’ performance

RL agent learning metrics

I trained the agent to a level of moderately good performance where it was successful in most but not all episodes. As you can see from the cumulative reward plot, the agent’s performance was still continually improving. I stopped the training at this moderately-good performance level to make sure the agent is not perfect at solving the environment and sometimes gets negative rewards. Otherwise, the model would only have to learn to assign positive rewards and probably would not pay attention to events that are important for such good performance. Think of it this way: when somebody is good at playing tennis, it’s often hard for them to explain what is important to make a great shot — they just do it automatically.

The training took about 2 hours on my laptop’s CPU, which is quite fast and is in part due to the small observation size, and running 32 copies of the agent in parallel. Here is a video of the agent successfully solving the task:

To train the attention model, I used 33 thousand episodes of the agent above interacting with the environment and chose the model with the lowest validation set loss after ten epochs.

Following are per-class F-1 score accuracy metrics of the resulting model on a separate held-out evaluation set:

F1 Scores for the three reward classes

We can see that the F-1 score for zero rewards is almost perfect, while negative and positive rewards have lower values.

Precision and recall, which make up the f-1 score, show why the results are this way. All classes have a very high recall, but negative and positive rewards have lower precision — which basically means the model assigns negative and positive values more often than it should.

This happens, for example, when the agent steps on a diamond which gives a positive reward. The model might get confused about whether the agent picked up the diamond just now or in the previous step, and predicts the positive rewards twice. A similar thing happens with negative rewards when the agent falls or steps into the fire.

Still, the confusion matrix shows quite great performance

The confusion matrix shows that most of the time, however, the classes are predicted correctly.

Attention to important events

As expected, most of the attention is usually concentrated in the few most recent moments, because they are enough to predict the reward well (the Markov property). However, sometimes the model does attend to earlier moments as well:

Attention weights when predicting the reward at moment B

We can see that although the attention value is quite small, there’s relatively more attention paid to moment A than to the rest of the past. Below you can see the screenshots from steps A and B.

As hoped for, there is increased attention to the past moment of pressing the Rock button (A) when walking into the fire (B).

Now, here is an example where attention is not concentrated around the most recent moment:

Attention weights when predicting reward at moment C

In this case, the agent was about to step on a bridge (C) and, as we would hope, the model paid increased attention to the moment the agent turned into Wood (B). Somewhat interestingly, the model also paid a lot of attention to an early moment in the episode, when the agent was about to step on another bridge (A).