Markov Decision Process  

Basic Theoretical Results in Model-Free Reinforcement Learning are proved for Markov Decision Processes.

Markovian propery:

In other words:   "The future is independent of the past given the present."

Memory in RL

Why We Need Working Memory?

. . .

. . .

Recurrent Memory:

\(h_{10}\)

. . .

. . .

\(h_9\)

\(h_{19}\)

. . .

. . .

. . .

Window Based Memory

. . .

. . .

Information

Gradients

Problems with RNN and Transformers

• Temporal dependency should fully fit into TBPTT or Attention Span to be learned

• You need to process all intermediate steps to implement backpropagation over Attention Span/TBPTT

Problem with RNN and Transformers

Idea: long version

Can we find temporal dependencies locally?

Information from red timestep could help at blue timestep.

Can we find temporal dependencies locally?

Information from red timestep could help at blue timestep.

Markov property doesn't work for these timesteps: 

How important it is to remember the red timestep?

\(f_t\) can be Q-value, \(s_{t+1}\), \(r_{t+1}\), etc.

Memory's Objective

The best memory would maximize the following sum: 

Memory's Objective

The best memory would maximize the following sum: 

Memory's Objective

Train memory to maximize MI for all future steps: 

Finding Important Moments

Goal: Find moments where memory is the most important for model's predictions!

Resume

Before:

Train memory to maximize MI for all future steps: 

Idea: short version

Imagine we have a perfect memory state \(m^*_t\) for each t!

Can we find temporal dependencies locally?

Imagine we have a perfect memory state \(m^*_t\) for each t!

Can we find temporal dependencies locally?

Memory's Objective

Train memory to maximize MI at moments that would benefit the most from using memory:​

   where \(U_t\) is a set of steps with the highest                     

Long-Term memory objective

Train memory to maximize MI at moments that would benefit the most from using memory:​

   where \(U_t\) is a set of steps with the highest                     

Intuition

T-Maze Example

hint: stores information about reward placement

• black wall: reward is on the right
• red wall: reward is on the left

Intuition

Learning Algorithm:

1. Estimate Uncertainty for each step in a trajectory
2. Select K steps with highest uncertainty
3. Traverse trajectory with memory
   • Train memory to predict future at selected high uncertainty steps

Intuition:

Two main steps:

• First we find the most surprising/uncertain events in a sequence
• Then we search sequence for information that explains these events and store it in memory

Implementation

Practical Implementation for RNN

Learning is divided in two phases:

• Memory Pretraining Phase - we train memory without improving agent's policy
• Policy Learning Phase - We train PPO agent, that receives output of the pretrained memory as an additional input.

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode
2. Select top K moments in episode with highest uncertainty

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode
2. Select top K moments in episode with highest uncertainty
3. Pass through the episode and train a memory model to store information that improves predictions at the selected moments

Experiments

Experiments

Baselines

We compare MemUP (Memory via Uncertainty Prediction) with the following baselines:

• PPO-LSTM is a recurrent version of PPO with a single LSTM-layer.
• IMPALA-ST is an IMPALA agent with Transformer-based architecture. The Stabilized Transformer architecture was presented by Parisotto et al. (2020)
• AMRL is proposed by Beck et al. (2020) for the Noisy T-Maze Task.  AMRL
  is similar to PPO-LSTM baseline, but extends LSTM with AMRL Layer.

IMPALA-ST

Experiments: Noisy T-Maze

Noisy T-Maze Environment:

• Hard Task in terms of Memory
• Simple Task in every other aspect 

Env Details:

• This version was proposed in AMRL paper
• Noisy Observations
• Determine the long-term dependency between hint at the start and reward at the end
• Learn to remember the hint, to achieve maximum reward:
  • +4 for matching the hint, -3 otherwise 

Experiments: Noisy T-Maze

Noisy T-Maze-100

Noisy T-Maze-1000

• All algorithms process sequences shorter than the temporal dependency
• Only MemUp solves the problem consistently

Experiments: Vizdoom

Details:

• Agent walks on acid (-4 health)
• +25 health for objects with a right color, otherwise -25 health
• The column hints the right color of objects 
• The column disappears after 45 steps
• +1 reward for the right color, -1 otherwise

ViZDoom-two-colors:

• 3D environment
• Complex observation space
• ~100 steps long temporal dependency
• Proposed in Breeching et al (2019)

Experiments: Vizdoom

Uncertainty Detector Demonstration

Experiments: Vizdoom

Results

• MemUP solves Vizdoom-Two-Colors

Experiments: Vizdoom

Memory Demonstration

Experiments: Ablation

Ablation Baselines

Memory Baselines:

• MemUP - our proposed algorithm
• Rnd-Pred - Same as MemUP, but events for prediction are selected randomly and uniformly among all future timesteps.
• Default - Same as MemUP, but it is trained to predict return R t at each step t, as oposed to an arbitrarily distant future events as in MemUP and Rnd-Pred.

Experiment Setup:

• Noisy T-maze environment, but we ignore policy learning phase
• Baselines are compared in terms of their ability to predict \(f_t\)
• Separate test set of 100 episodes

Experiments: Ablation

Results