Documenting DS & RL Research Progress

Test a meta-reinforcement learning model (A2C) on an associative learning task.

2021-01-26T02:31:52+00:00

Test a meta-reinforcement learning model (A2C) on an associative learning task

In this post, I want to explore the power of a simple meta-reinforcement learning (meta-RL) model on simple associative learning tasks. This is the second post on my journey to understand meta-RL models and hopefully tweak them to behave like humans; and you can check out my overall project proposal here. This post assumes that you have basic knowledge on reinforcement learning (RL), and there are plenty of good RL tutorials out there.

What is meta-reinforcement learning?

The term “meta” confused me for a long time, so let’s understand “meta-learning” in a clearer term: learning to learn. Imagine the traditional way we train a RL model to play the Pong game. For example, we can decide to pick the DQN algorithm, specify the number of layers and units in the Q-network, and finally choose our favorite loss function and optimizer. We may get lucky this time and the model plays Pong well, but this combination of choices may not work for the GO game, and we need to start from scratch and test numerous other algorithms. What if we can train a model to learn the best learning algorithm for the task? Wouldn’t it be wonderful if we have such a model that can figure out the best way to learn in the current task setting? This is the essence of meta-RL models: find the best learning aglorithm for the task with reinforcement learning.

Why should we test the model on associative learning tasks?

Associative learning tasks are pretty common in the cognitive science field but not so much in computer science, but in our context they are simple multi-armed bandit tasks. You might think that multi-armed bandit task is one of the most basic problems in RL and countless algorithms are already built to solve it near-optimally, why bother testing it? Well, here are several reasons:

Most of RL algorithms have built-in memory one way or another, but meta-RL algorithms don’t. Traditional models keep information of rewarding state-action pairs in Q-tables or value networks, so the way they utilize these memory resources are designed by the human researchers. On the other hand, the learning algorithms learned by meta-RL models are complete black boxes; researchers only specifies the number of memory units that the learning algorithm can use, and how these memory units are used is totally up to the algorithm. **By testing the meta-RL algorithm on associative learning tasks, we hope to understand how they utilize the given memory resources.
While researchers can easily expand the memory resources in RL algorithms, you can’t do the same things to humans. Humans have very limited memory resources, and their performance fell short even in simple associative learning tasks as the number of state-action pairs increases. However, humans seem to recognize this limit and they can act in a near-optimal fashion under this constraint. It is therefore intriguing to see whether the learning algorithm learned by the meta-RL model can show similar behavior under limited memory resources.

Let’s quickly go through the procedure of the associative learning task. As shown in Fig. 1, the model was trained on multiple associative learning blocks which have either 3 or 6 objects. At the start of each block, each object will be randomly associated with one of the three buttons. For each trial within each block, the model observes one stimulus (one-hot encoded in the experiment), presses one of the three action buttons, and finally gets feedback on whether the object-action pair is correct or not. I tried to replicate this procedure described in this paper as they recruited human participants to do this exact same task, and this makes it possible to compare the behavior between humans and models.

The associative learning task

Fig. 1. Experiment procedure of the associative learning task. Figure adapted from [1]

Now it’s time to train some models!

First, let’s setup the environment

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
%matplotlib inline

plt.style.use('seaborn-dark-palette')
mpl.rcParams['axes.titlesize'] = 20

Second, let’s train the model

I trained an iconic meta-RL model, an A2C algorithm with 50 LSTM units, on 40 million trials. I compared the behavior between human participants and the trained meta-RL model.

To examine how meta-RL models behave under limited memory resources, I trained another model with 20 LSTM units on 20M trials. Reducing the number of LSTM units is definitely a crude way to limit memory resources, but it does the work as a first step. I will explore various capacity-limited model variants (e.g. limit the mutual information between observation and action) in the future.

from envs.Collins2018 import Collins2018Task
from utils.utils import do_experiment

# specify the agent and the environment
num_lstm_units = 20
num_train_steps = 20_000_000
policy_kwargs = {'n_lstm': num_lstm_units, 'action_noise': 0}
env = Collins2018Task(num_objects=(3, 6), num_actions=3, num_repeats=13)

do_experiment(env, num_train_steps, policy_kwargs)

Finally, let’s use the learned model to rollouts some new blocks.

To test whether the learned model can generalize to new tasks, we rollout the model on 600 new blocks where 1 to 6 objects are available. Note that this isn’t really testing whether the model can transfer knowledge from training tasks to new tasks: The model can probably solve 1 to 5-object blocks just like 6-object blocks where the remaining object are absent. It is still good to test whether the model can reach this level of generalization though.

You might encounter tons of warnings if you are running this notebook. This is because I am using the A2C implementation from stable-baselines which uses tensorflow v1.

from utils.utils import Simulation
from envs.GuessBoundary import GuessBoundaryTask
from envs.Collins2018 import Collins2018Task

env = Collins2018Task(num_objects=(1, 2, 3, 4, 5, 6))

sims = [
    Simulation('./outputs/Collins2018/A2C-LSTM50', env, load_params_file_idx=20_000_000),
    Simulation('./outputs/Collins2018/A2C-LSTM50', env, load_params_file_idx=10_000_000),
    Simulation('./outputs/Collins2018/A2C-LSTM20', env, load_params_file_idx=20_000_000)
]

for sim in sims:
    sim.evaluate(num_test_episodes=600, num_trials=100, verbose=True)

Let’s check the model’s training progress

from utils.visualization import plot_training_progress

plot_training_progress(sims[0].model_path)

The model with 50 LSTM units probably reached performance plateau after 15 million trials of training.

Let’s compare three models:

a 50-unit A2C model trained with 20M trials (possibly over-trained)
a 50-unit A2C model trained with 10M trials (possibly under-trained)
a 20-unit A2C model trained with 20M trials (possibly over-trained)

Let’s compare how human participants score against the meta-RL model

from utils.visualization import plot_associative_learning_progress, calc_exploration_and_retention
from utils.utils import behav2sim

df_behav = behav2sim()
dfs = [
    df_behav,
    sims[0].rollouts,
    sims[1].rollouts,
    sims[2].rollouts
]
mdl_names = [
    'human behavior',
    'LSTM50 / train20M',
    'LSTM50 / train10M',
    'LSTM20 / train20M'
]

fig, axs = plt.subplots(1, len(dfs), figsize=(16, 4))

for mdl in range(len(dfs)):
    # does the model learn after multiple stimulus presentations?
    plot_associative_learning_progress(axs[mdl], dfs[mdl])
    axs[mdl].set_title(mdl_names[mdl])
    # does the model learn to attempt different actions after failed trials (before the single success)? 
    #calc_exploration_and_retention([axs[1, mdl], axs[2, mdl]], dfs[mdl])
    

# insight: network should be quite larger than the optimal size

You can see how the human participants score against the three models.

The main highlight here is that human participants did worse in the 6-object blocks than the 3-object blocks (leftmost figure). This is expected since humans typically can’t remember this many object-action pairs, so they need a bit of extra practice to reach the same level of performance. The over-trained model (second figure) doesn’t show this obvious trend, but the under-trained and memory-limited model shows similar trend.

Now let’s do a more detailed analysis on how the models utilize memory.

To examine that, we can check the following two aspects:

If the make a wrong action guess, it should try another action during the next encounter. Since there is only 3 action buttons, for each object the maximum number of failed attempts before the correct action is hit is at most 2. More than 3 failed attempts shows imperfect memory on the incorrect action attempts.
After the correct action is selected, the model should stick to this action for the rest of the block. Failure to do so shows imperfect memory on the correct action attempts.

fig, axs = plt.subplots(2, len(dfs), figsize=(16, 8))
axs[0, 0].set_title('human behavior')
axs[0, 1].set_title('LSTM50 / train20M')
axs[0, 2].set_title('LSTM50 / train10M')
axs[0, 3].set_title('LSTM20 / train20M')

for mdl in range(len(dfs)):
    # does the model learn to attempt different actions after failed trials (before the single success)? 
    calc_exploration_and_retention([axs[0, mdl], axs[1, mdl]], dfs[mdl])
 

The top row shows the distribution of number of failed attempts before the correct action is found. We can easily observe that humans occasionally forget the incorrect actions they have picked in the previous trials. The models, especially the over-trained one, show almost perfect memory on previous incorrect attempts.

The bottom row shows the distribution of average reward after the correct action is found. All the points should be at 1.0 if the model can remember the correct action perfectly. We can see that the over-trained model is nearly perfect, and humans again show some level of imperfect memory.

Note the small bump at 0 in the bottom right graph. What does this mean? This means that the trained model has a noticeable chance to completely forget the answer for an object. This shows that the model fails to consistently utilize its memory resources. More training can certainly solve the problem, but this might mean the models are using memory in a fundamentally different way as humans do.

Are there more artifacts unique to the models?

You bet. Since the objects are coded 0-2 and 0-5 in the 3- and 6-object tasks respectively, the model may be trained more on the first three objects and less the last three. On the other hand, humans subjects might show similar trends since their working memory is already pre-occupied by earlier objects, but the effect shouldn’t but huge.

Here is the comparison between the human behavior and the model:

# is there training artifact that the stimulus 0-2 are better trained?
# Ans: yes
fig, axs = plt.subplots(1, len(dfs), figsize=(16, 6))

for mdl in range(len(dfs)):
    obj_rewards = dfs[mdl].groupby('observations').rewards.mean().tolist()
    axs[mdl].bar(np.arange(len(obj_rewards)), obj_rewards)
    axs[mdl].set_title(mdl_names[mdl])
    axs[mdl].set_xlabel('Object ID')
    axs[mdl].set_ylabel('Mean Reward')
    axs[mdl].set_ylim([0.5, 1])

Hmm, that is messy. As expected, humans are pretty consistent across different objects; as long as a red circle looks different from a green one, we should have no trouble associating them with different actions. For the under-trained and the memory-limited models we can see a drop for the last three objects, probably because they received 50% less training. The third object in the over-trained model oddly shows lower performance, and I haven’t figured out why yet.

Nevertheless, one thing is for sure: while humans can easily identify that the same rule can be applied to all objects regardless of how many they are, models probably learn the rule of each object independently. This partially explains why it takes 20 million trials to train a meta-RL model to learn such a simple associative learning task. This is one of the advantages of graph neural networks (GNNs) over plain LSTM units: if the memory resources given to a meta-RL model is a graph instead of a fixed-sized network, the model may opt to create a memory node for each object it encounters and learn a single operation that applies to all nodes. The latter usage of memory is much closer to humans and it needs way fewer parameters. In the next post, I will explore how such networks can be implemented and how they behave on this associative learning task.

Conclusion

In this post, I examined how a meta-RL model utilizes its memory resources and compared it to human behavior. I trained multiple models on a simple associative learning task and found that while models often memorize important information better than humans do, they require a huge amount of training data and they tend to learn the rule for each object independently. In the future posts, I will explore how to enforce stronger inductive biases into a meta-RL model and force it to learn in a human-like manner.

This is only a work in progress and there are still quite a lot of TODOs to clean up:

While behaviors can tell us a lot about how models utilize memory, it is even better if we can directly infer what is stored inside the LSTM hidden units. One way to do this is to correlate task-relevant information (e.g. which objects are already associated with action buttons) through multiple blocks with the hidden state values. This yields mixed results in my previous attempts, but I think this might be worth a shot.
Limiting the memory resources of a model by reducing the number of LSTM units is definitely a crude approach, and there are more elegant information-theoretic approaches that I have yet explored. A meta-RL model can be interpreted as a capacity-limited channel that maps observations to actions, and this capacity can be tuned by limited the mutual information between the model input and output. It will be interesting to see whether enforcing such mutual information constraints can make the meta-RL model behave more like humans.

References:

1.Collins AGE. The Tortoise and the Hare: Interactions between Reinforcement Learning and Working Memory. J Cogn Neurosci. 2018 Oct;30(10):1422-1432. doi: 10.1162/jocn_a_01238. Epub 2018 Jan 18. PMID: 29346018.

Learning to form a structured hypothesis space and infer optimally.

2021-01-25T02:31:52+00:00

How do humans achieve planning and reasoning?

In my leisure time, I play a lot of brain teasers. A quick online search will reveal countless of these games, and you can always find completely new games or those that share some rules that you have encountered previously. However different they are, one thing these games have in common is that with a quick read of the rules or a simple tutorial, most humans can have a good grasp of them and beat simple levels with near-perfect strategies. How do we do that, and can we build artificial intelligence (AI) agents to learn like that?

To get a rough understanding of what properties we want in our intelligent agent, let’s walk through a really simple game and think about the thought processes of a typical human player. In a basic “Guess A Number” game, the agent is given a range of numbers (e.g. from 1 to 100); in each trial, the agent chooses a number, gets told whether he should go higher or lower, and this process repeats until he selects the correct number. Here are some amazing properties that lead to the insights a human player would have:

Can form the right hypothesis space: The player should figure out that he doesn’t have to check each number to find the answer. Each guess he makes may eliminate multiple candidate numbers, so he should prefer this strategy to the first and exhaustive search one. Furthermore, the range of possible numbers is always between the largest “higher” guess and the smallest “lower” guess, so only two numbers are necessary to keep track of the current “hypothesis”.
Can store only the relevant information in memory: The player should be able to keep track of the current hypothesis while avoiding irrelevant previous trials from interference.
Can form a near-optimal strategy by eliminating as many hypotheses as possible: The player should figure out that the number in the middle of the range is always the best guess, as the feedback for this number can narrow the range down by half.

While I use brain teasers as motivation, most of the human reasoning and planning behaviors also fit into this framework, and these abilities are the cornerstones of our intelligent behavior. More importantly, humans can reason with these abilities and develop strategies with very little experience with the game (i.e. learn in a few-shot manner), and they can easily deploy the same strategy to other games with a different range of numbers (i.e. flexible task transfer). It is therefore desirable to build an intelligent agent that behaves in similar ways and possesses all of these properties.

How do current DNN models achieve planning and reasoning?

DNN can be trained to play GO and a variety of games

Does single-task DRL models satisfy the 5 properties?

No, the hypothesis space is usually hard-coded in the model. In a non-hierarchical RL model, this H space grows exponentially (i.e. num_actions ** num_time_steps)
Yes, most RL networks store information either in the form of memory units or value table/networks.
No, DRL models usually don’t consider “hypotheses” explicitly, so they don’t perform inference
No, typical DRL models require huge amount of data
No, typical DRL models generalize poorly to other tasks

Meta-RL model has shown promising results on rapid and flexible task transfer

Does meta-RL model fill in some shortcomings compared to single-task DRL models?

Yes, meta-RL models can learn better forms of hypothesis space after being trained on numerous training tasks. However, the hypothesis space can be implicit in the model weights, so it is challenging to compare this space with the space humans use to reason in.
Yes, meta-RL models can store task-specific information in memory units.
Yes, meta-RL models have shown near-optimal inference results on some tasks, but this needs to be further investigated.
Yes, after being trained on numerous training tasks, meta-RL models can solve a novel testing tasks fairly optimally.
Yes, if the training and testing tasks are sampled from the same distribution, flexible task transfer is possible.

Problems solved. Now what is the gap?

Yes, it is easy to claim that we are able to use meta-RL models to satisfy all 5 properties in human reasoning. However, there are still quite a lot of fundamental differences between the two.

lack of structure in the hypothesis space Humans can often identifty a good hypothesis space to tackle the problem; usually, this space can be described explicitly in the form of programs or graphs. However, the hypothesis space formed by meta-RL models is implicitly described in model parameters. While it is possible that the meta-RL model can form some fascinating hypothesis spaces beyond human imagination, we as scientists wants to constrain the space such that it is explainable.
human memory capacity is limited One of the most obvious things we understand about human reasoning behavior is that humans reason with very limited memory. This means that humans starts to forget crucial information or fail to consider all hypotheses when the task grows in size. On the other hand, meta-RL models are usually trained with ample memory resource. Moreover, previous results showed that while the models can store task-relevant information in memory, they often can’t get rid of the task-irrelevant bits.
models don’t actually eliminate hypotheses like humans

Graph Neural Networks to the Rescue

We propose to replace plain LSTM units with a graph neural network. How does this fill in the gap?

GNN provides structure It is easy to visualize the model on graphs. Furthermore, we hypothesize that the model can adjust the number of nodes based on the size of the task input, and this flexibility may lead to better generalization to novel tasks.
We can simulate the “low-memory” scenario in human reasoning behavior As mentioned previously, human memory resources are limited. With a graph-structured hypothesis space, we can constrain the model in various ways to examine the possible causes of human error. For example, humans might only form 5 nodes at a time, and any additional information must be discarded or compressed into existing nodes. As another example, when the hypothesis space is a huge tree structure, humans may only consider several hypothesis that are near each other.
We can hopefully recreate classic divide-and-conquer strategy from GNN

Model Structure

Tasks

Associative Learning Task (Collins 2018): This task tests the model’s ability to utilize memory.
- H1: Once the model gets a reward, it should remember the state-action pair and always act accordingly.
- H2: If the model didn’t get rewarded, it should “avoid” this state-action pair and try out other actions.
- H3: When the number of nodes increases, the model should start to forget information, and it should be consistent with the Collins 2018 and Gershman 2020 model.