Notebook Overview

In this notebook we show the results of a Leaky Integrator Echo State Network Reinforcement Learning algorithm closely followed from the work by Zhang et al. on the cartpole environment:

  • Reservoir nodes are initialized to 0 and updated according to $\bf{h}_t^{(k)} = \bf{h}_t^{(k-1)} + max(0,\bf{W}_{in} \cdot \bf{x}_t^{(k)} + \bf{b})$, where input weights $\bf{W}_{in}$ and the bias are initialized randomly. 't' is the time step until the environment returns 'done' and k stands for the $k^{th}$ entry of the $i^{th}$ series element $(\bf{s}_{t,i}^{(k)},\bf{a}_{t,i}^{(k)},\bf{s'}_{t,i}^{(k)},\bf{r}_{t,i}^{(k)})$ in the minibatch $\{(\bf{s}_{t,i}^{(k)},\bf{a}_{t,i}^{(k)},\bf{s'}_{t,i}^{(k)},\bf{r}_{t,i}^{(k)})\}_{i=1,...,M;k=1,...,T}$. Batch size $M$ is set to 64 and length of each series $T$ is 5

  • The rest of the algorithm is analogous to the well-known DQN algorithm. The Q values of the action space are calculated by the policy network $Q(\bf{S_t},\mathcal{A},\Theta_{t-1}) = \Theta_{t-1} \cdot \bf{U}_t^{k}$ and the optimal Q values are approximated by the target network $Q^{\pi}(\bf{S'_t},\mathcal{A},\tilde{\Theta}) = \tilde{\Theta} \cdot \bf{U'}_t^{k}$, where $\Theta_{t-1}$ and $\tilde{\Theta}$ consist of the randomly initialized output weights and output bias $[W_{out};b^{out}]$ and $\bf{U}_t^{k} = [\bf{S}_{t}^{(k)};\bf{H}_{t}^{(k)};\bf{1}]$ with $\bf{S}_{t}^{(k)}$ containing the states in the batch as column vectors. $Q^{\pi}$ is later updated according to Bellman equation with the rewards in the batch $r_t^{(k)}$ : $Q^{\pi}(\bf{S}_t^{(k)},\mathcal{A}) = r_t^{(k)} + \gamma \cdot Q^{\pi}(\bf{S'_t},\mathcal{A},\tilde{\Theta})$

  • Another key aspect of the algorithm is the mean approximation in the gradient calculation. We average over the series elements and batch elements in the following to acquire one $Q^{\pi}$ and $Q$ at each iteration: \begin{align} \bf{\bar{u}}_t &= \frac{1}{MT} \sum_{k=1}^{T} \sum_{i=1}^{M} \bf{U}_{t,i}^{k}, \text{where $\bf{U}_{t,i}^{k}$ is the $i^{th}$ column vector of $\bf{U}_t^{k}$} \\ \bar{Q}^{\pi}_t &= \frac{1}{MT} \sum_{k=1}^{T} \sum_{i=1}^{M} Q^{\pi}(\bf{s}_t^{(k)},\mathcal{A}) \\ \bar{Q}_t &= \frac{1}{MT} \sum_{k=1}^{T} \sum_{i=1}^{M} Q(\bf{s}_{t,i}^{(k)},\mathcal{A},\Theta_{t-1}) \end{align}

The output weights are then updated with the learning rate \begin{align} \alpha_t &= \frac{P_{t-1}^T}{\lambda + \bf{v}_t^T \bf{\bar{u}}_t} \end{align} and the gradient \begin{align} \nabla_{t-1} &= (\bar{Q}^{\pi}_t - \bar{Q}_t) \bf{\bar{u}}_t^T \end{align} and the regularization term \begin{align} L_1 &= - \kappa \cdot sgn(\Theta_{t-1})P_{t-1}^T \end{align}

yielding

\begin{equation} \Theta_{t} = \Theta_{t-1} + \nabla_{t-1} \alpha_t + L_1 \end{equation}

, where

\begin{align} P_{t} &= \frac{1}{\lambda}(P_{t-1} - \bf{g_t}\bf{v}_t^T) \\ \bf{v}_t &= P_{t-1}\bf{\bar{u}}_t \\ \bf{g}_t &= \frac{\bf{v}_t}{\lambda + \bf{v}_t^T \bf{\bar{u}}_t} \end{align}
In [11]:
import sys,os
import pandas as pd
import numpy as np
import gym
import numpy as np
import torch
from IPython.display import clear_output
from tqdm.notebook import tqdm,trange
sys.path.append("./../../../")
from collections import namedtuple, deque
from itertools import count
import random,math
from matplotlib import pyplot as plt

def at_least_2d(arr):
    if len(arr.shape)==1:
        return arr[:,None]
    elif len(arr.shape)==2:
        return arr
    else:
        raise Exception("Unsupported array shape.")

env = gym.make("CartPole-v0")
n_actions = env.action_space.n

relu = np.vectorize(lambda x: x if x>=0 else 0*x)#,otypes=[np.float64])

random.seed(42)
np.random.seed(42)
env.seed(42);

Hyperparameters

In [12]:
BATCH_SIZE = 64
MEMORY_SIZE = int(1e5)
GAMMA = 0.99
EPS_START = 0.01
EPS_END = 0.01
EPS_DECAY = 256
TARGET_UPDATE = 1
resSize = 256
bias=1
T_transient = 5
T_opt = 5
T = T_transient+T_opt
forgetting_factor = 0.99999
kappa = 1e-5
max_episodes=100
omega = 1
P_alpha = 0.4

Replay Memory

As the agent chooses actions the states, actions, next states and rewards will be recorded into replay memory.

In [13]:
Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward', 'done'))

class ReplayMemory(object):

    def __init__(self, capacity,series_length):
        self.series_memory = deque([],maxlen=series_length)
        self.memory = deque([],maxlen=capacity)
        self.length = 0
        self.series_length = series_length

    def push(self, *args):
        """Save a transition"""
        if args[4]:
            while len(self.series_memory)<self.series_length:
                self.series_memory.append(Transition(*args))
                self.length +=1

        else:
            self.series_memory.append(Transition(*args))
            self.length +=1
        if len(self.series_memory)==self.series_length:
            self.memory.append(self.series_memory)
            self.series_memory = deque([],maxlen=self.series_length)
            

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

One Step Optimizer

At each iteration the policy network weights are optimized via gradient ascent.

In [14]:
def optimize_model():
    global P
    global policy_train_res_layer
    global policy_train_Wout
    global target_train_res_layer
    
    if len(memory) < BATCH_SIZE:
        return False
    
    transition_series = memory.sample(BATCH_SIZE)
    batch_series = pd.DataFrame(transition_series)
    q_diff = np.zeros((env.action_space.n,BATCH_SIZE))
    u = np.zeros((env.observation_space.shape[0]+resSize+bias,BATCH_SIZE))
    
    for i in range(T):

        batch = Transition(*zip(*batch_series[i].to_list()))

        state_batch = np.array(batch.state).T
        action_batch = np.array(batch.action)
        reward_batch = np.array(batch.reward)

        next_state_batch = np.array([s if s is not None else env.observation_space.shape[0]*[0] for s in batch.next_state]).T
        
        done_batch = np.array(batch.done)
        
        policy_train_res_layer = update_f(policy_train_res_layer,state_batch)
        target_train_res_layer = update_f(target_train_res_layer,next_state_batch)
        if i > T_transient:
            U = get_U(policy_train_res_layer,state_batch)
            state_action_values = np.dot(policy_train_Wout,U)
            u += U

            #target_train_res_layer = update_f(target_train_res_layer,next_state_batch)
            next_state_values = np.dot(target_train_Wout,get_U(target_train_res_layer,next_state_batch))

            next_state_values = (1-done_batch)*next_state_values

            # Compute the expected Q values
            expected_state_action_values = (next_state_values * GAMMA) + reward_batch

            q_diff += expected_state_action_values - state_action_values

    ### UPDATE SECTION ###
    #Necessary for updates
    q_diff = q_diff.sum(1,keepdims=1)/BATCH_SIZE/T
    u = u.sum(1,keepdims=1)/BATCH_SIZE/T
    v = np.dot(P,u)
    g = v/(forgetting_factor+np.dot(v.T,u))

    # Updates
    policy_train_Wout+= np.dot(q_diff,g.T) - kappa*np.dot(np.sign(policy_train_Wout),P.T)
    P = (P-np.dot(g,v.T))/forgetting_factor
    policy_train_res_layer = np.zeros((resSize,BATCH_SIZE))
    target_train_res_layer = np.zeros((resSize,1))
    #####################
    
    return True

Training

In [15]:
Wout_0 = np.random.rand(n_actions,env.observation_space.shape[0]+resSize+bias)
Win = np.random.rand(resSize,bias+env.observation_space.shape[0]) - 0.5

def update_f(x,in_):
    _in_ = at_least_2d(in_)
    _bias = at_least_2d(np.array(_in_.shape[1]*[bias]))
    return x + relu(np.dot(Win, np.vstack((_bias.T,_in_))))

def get_U(x,in_):
    _in_ = at_least_2d(in_)
    _bias = at_least_2d(np.array(_in_.shape[1]*[bias]))
    return np.hstack((_in_.T,x.T,_bias)).T

policy_res_layer = np.zeros((resSize,1))
policy_train_res_layer = np.zeros((resSize,BATCH_SIZE))
target_train_res_layer = np.zeros((resSize,1))

policy_Wout = Wout_0.copy()
policy_train_Wout = Wout_0.copy()
target_train_Wout = Wout_0.copy()

# P from Sherman-Morrison formula
P_0 = np.identity(env.observation_space.shape[0]+resSize+1)*P_alpha
P = P_0.copy()

# Parameter to track no of ALL steps taken during the WHOLE training.
steps_done = 0


episode_durations = []

highest_score = 0


def select_action(state):
    global steps_done
    global policy_res_layer
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        policy_res_layer = update_f(policy_res_layer,state)
        return np.argmax(np.dot(policy_Wout,get_U(policy_res_layer,state)))
    else:
        return random.randrange(n_actions)
    
    
memory = ReplayMemory(MEMORY_SIZE,T)

i_episode = tqdm(total=max_episodes)

while i_episode.last_print_n < max_episodes:
    assert np.all(policy_Wout == Wout_0.copy())
    assert np.all(policy_train_Wout == Wout_0.copy())
    assert np.all(policy_res_layer == np.zeros((resSize,1)))
    assert np.all(policy_train_res_layer == np.zeros((resSize,BATCH_SIZE)))
    assert np.all(target_train_res_layer == np.zeros((resSize,1)))
    assert np.all(P_0 == np.identity(env.observation_space.shape[0]+resSize+1)*P_alpha)
    assert np.all(P==P_0)
    
    # Initialize the environment and state
    state = env.reset()
    for t in count():
        assert np.all(policy_train_Wout == policy_Wout)
        # Select and perform an action
        action = select_action(state)
        next_state, reward, done, _ = env.step(action)

        # Observe new state
        if done:
            next_state = None
            if t < 199:
                reward=-10

        # Store the transition in memory
        memory.push(state, action, next_state, reward, done)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the policy network)
        optimized = optimize_model()
        if optimized:
            policy_Wout = policy_train_Wout.copy()
            if i_episode.last_print_n % TARGET_UPDATE == 0:
                target_train_Wout = policy_train_Wout.copy()
        
        if done:
            policy_res_layer = np.zeros((resSize,1))
            if optimized:
                episode_durations.append(t + 1)
                i_episode.update()
                P = P_0.copy()
                policy_train_Wout = Wout_0.copy()
                policy_Wout = Wout_0.copy()
            break

print('Complete')
plt.plot(episode_durations)
plt.show()

Complete

Validation

We let the agent play the game 100 times and collect the results.

In [19]:
ep_lengths = []
target_res_layer = np.zeros((resSize,1))
target_Wout = target_train_Wout.copy()
for i in range(100):
    state = env.reset()
    done=0
    rewards=[]
    target_res_layer = np.zeros((resSize,1))
    while not done:
        target_res_layer = update_f(target_res_layer,state)
        action = np.argmax(np.dot(target_Wout,get_U(target_res_layer,state)))
        next_state, reward, done, _ = env.step(action)
        state = next_state
        rewards.append(reward)
    ep_lengths.append(len(rewards))
np.mean(ep_lengths)
Out[19]:
194.13

The average episode length is ~194. More than 80% of the time, the agent is able to complete the game succesfully.

In [20]:
plt.hist(ep_lengths,align='right')
plt.show()