How to use GRU/LSTM is RL?

Ayman_Al_Jabri: <p>Hi all,<br>
I want to add memory cell/layer to my model to improve performance on Atari games. But I’m not sure if I’m doing it right!<br>
If I understood recurrent networks correctly, they take a sequence of observations from the environment. This defies the i.i.d assumption as the observations in the batch become highly correlated, but that is fine since the memory cells are made for that specific task.</p>
<p>That said, what I’m struggling with is how to use the hidden state output in my training. Should I:<br>
1- Save it after each step and feed it back to the network in step + 1 until episode end?<br>
2- OR ignore it and just initiate zero tensor at each step?</p>
<p>Secondly:<br>
Assuming the GRU is feeding a dense layers below. Which output should I pass to the FC layer (output or hn)?</p>

raharth: <p>Hi Ayman,</p>
<p>If you don’t feed back the hidden state of t to the step t+1, you basically use a fancy version of a linear layer with shared weights for all time steps, but you loose the recurrent property of the model (I think this approach is called time distributed layer, which is not a recurrent one). Whenever you have a time series the states of a sequence are highly correlated so I don’t see a problem with that (since that is an inherent property to all series, there is nothing specific about RL in this case). Or am I getting your question wrong?</p>
<p>Also, you always have to feed the output to the next layer not the hidden state.</p>

Ayman_Al_Jabri: <p>Hey,</p>
<p>That’s exactly the question I needed confirmation on. However, I tried initiating h0 as zeros at the beginning of each episode and make the agent pass it from one step to another until episode end. But the agent didn’t learn. Not even after couple million frames on any of the Atari games.<br>
I must be doing something wrong!</p>
<p>On the second point, the output is nothing but hidden states from all sequences stacked. i.e. last value of output == h_n<br>
If that is the case why not pass h_n to the dense layer instead of output. it’s computationally cheaper and backpropogation will go through the other hidden layers by default since h_n was calculated from them. What do you think?</p>

raharth: <aside class="quote no-group" data-username="Ayman_Al_Jabri" data-post="3" data-topic="124650">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ayman_al_jabri/48/20022_2.png" class="avatar"> Ayman_Al_Jabri:</div>
<blockquote>
<p>On the second point, the output is nothing but hidden states from all sequences stacked. i.e. last value of output == h_n</p>
</blockquote>
</aside>
<p>(Ok, first of all a disclaimer: I’m for sure no expert, but I have worked with several RL algorithms in Deep Learning)<br>
So having this out of the way <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"> I see what you mean, interesting question… I guess the difference is that once the linear layer only knows about the last step, in the other case it has access to the entire sequence. I think both ways are valid, but in the context with RL with unknown sequence length you will probably have to go with with the hidden state. But I also assume that you will not need any Linear non-recurrent layer, so you could probably just use the hidden state of the last layer directly as output.</p>
<p>Regarding the first question what algorithm are you using? DQN, DDQN, or a Policy Gradient? Which environment are you trying to solve and what is the architecture you are using?</p>
<p>If you are interested, I’m also happy to discuss that with our outside of that threat!</p>

Ayman_Al_Jabri: <p>You know what I’m going to try both approaches on a simple environment and see which one learns faster! I think sequence number can be fixed with pad_sequance.</p>
<aside class="quote no-group quote-modified" data-username="raharth" data-post="4" data-topic="124650">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/raharth/48/14974_2.png" class="avatar"> raharth:</div>
<blockquote>
<p>DQN, DDQN, Policy Gradient?</p>
</blockquote>
</aside>
<p>I’m using A3C for this. I’d created the “fancy linear version” by adding GRU cell without feeding back hidden states from last step :). It gave much better results than A3C alone but still not the level I want!!</p>
<p>here’s the model I’m upgrading:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/ayjabri/refresh/blob/master/09_A3C_GRU.py">
  <header class="source">

      <a href="https://github.com/ayjabri/refresh/blob/master/09_A3C_GRU.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/ayjabri/refresh/blob/master/09_A3C_GRU.py" target="_blank" rel="noopener nofollow ugc">ayjabri/refresh/blob/master/09_A3C_GRU.py</a></h4>


      <pre><code class="lang-py">#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Attempting to solve Atari games using A2C

@author: Ayman Jabri
"""
import os
import torch
import ptan
from lib import data, utils, mp_utils, model
from tensorboardX import SummaryWriter
import torch.multiprocessing as mp


ALGORITHM = 'A3C_GRU'
FORKS = 4
MINI_BATCH = 16


</code></pre>



  This file has been truncated. <a href="https://github.com/ayjabri/refresh/blob/master/09_A3C_GRU.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


raharth: <aside class="quote no-group" data-username="Ayman_Al_Jabri" data-post="5" data-topic="124650">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ayman_al_jabri/48/20022_2.png" class="avatar"> Ayman_Al_Jabri:</div>
<blockquote>
<p>You know what I’m going to try both approaches on a simple environment and see which one learns faster! I think sequence number can be fixed with pad_sequance</p>
</blockquote>
</aside>
<p>Sounds good, let me know what you found! It is for sure possible, but I think it is computationally less efficient, since you then have to process the entire length even though it is empty?</p>
<aside class="quote no-group" data-username="Ayman_Al_Jabri" data-post="5" data-topic="124650">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ayman_al_jabri/48/20022_2.png" class="avatar"> Ayman_Al_Jabri:</div>
<blockquote>
<p>I’d created the “fancy linear version” by adding GRU cell without feeding back hidden states from last step <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>I’m sorry, I didn’t mean it as an insult! <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
<p>Looks interesting I will have a look at it! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Ayman_Al_Jabri: <aside class="quote no-group" data-username="raharth" data-post="6" data-topic="124650">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/raharth/48/14974_2.png" class="avatar"> raharth:</div>
<blockquote>
<p>I’m sorry, I didn’t mean it as an insult! <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>Lol… non taken man! It must be my dark sense of humor <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<p>As for the recurrent model; I’ve experimented yesterday with OpenAI game ‘CartPole-v0’, and a basic version of Policy Gradient to keep things simple until I understand the dynamics of recurrent networks in RL and they’re trained!</p>
<p>I’m going to  <strong>benchmark</strong> it against this linear model, which solves the environment in 7 seconds.</p>
<pre><code class="lang-auto">class PGNet(nn.Module):
    r"""Plain vanilla policy gradient network."""

    def __init__(self, obs_size, act_size, hid_size=128):
        super().__init__()
        self.fc1 = nn.Linear(obs_size, hid_size)
        self.output = nn.Linear(hid_size, act_size, bias=True)

    def forward(self, x):
        """Feed forward."""
        y = F.relu(self.fc1(x))
        return self.output(y)
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/d/5d6a046d3f7d00735f27239b12025fd33d01c5fd.png" alt="Screenshot from 2021-06-24 07-43-19" data-base62-sha1="dknxPS3eZk80p4z5SqnegW5z8Rf" width="655" height="226"></p>
<p>Benchmark code on :<a href="https://github.com/ayjabri/DeepRL/tree/main/05_PolicyGradient" class="inline-onebox" rel="noopener nofollow ugc">DeepRL/05_PolicyGradient at main · ayjabri/DeepRL · GitHub</a></p>
<p>I’ll post the logic and code later today…tbc</p>

Ayman_Al_Jabri: <p>This is the Recurrent Policy Gradient model I’m gonna use. It has two GRU layers and one dense for policy output (just like the one above).</p>
<p>The network takes a sequence of 5 observations (.i.e. input size (batch, 5, features)) and pads it to the desired sequence number. It takes into account episode end</p>
<p><strong>For example:</strong><br>
At the beginning of the episode the input is only one observation of shape (batch, 1, features). The functions adds zeros <strong>before</strong> the observation. But at episode end it adds them <strong>after</strong> the last observation.</p>
<pre><code class="lang-auto">class RNNPG(nn.Module):
    r"""Recurrent plain vanila policy gradient model."""

    def __init__(self, obs_size, act_size, sequence=5, hid_size=128, num_layers=2):
        super().__init__()
        self.sequence = sequence

        self.gru = nn.GRU(obs_size, hid_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hid_size*sequence, act_size)

    @staticmethod
    def pad_sequence(states, episode_end=False, sequence=5, padding_value=0.0):
        r"""Pad a sequence of observations upto sequence size by adding `padding_value` above or below the input.

        It pads above states by default unless episode end is `True` then it pads after (see example).

        Args:
            states: tensor of size (batch, sequence_0, features)

            episode_end: bool   

            sequence = 5: int. desired sequence size    

            padding_value = 0.0: float. value to pad tensor with     

        Example:
            &gt;&gt;&gt; a = a = torch.ones(1,2,4)
            &gt;&gt;&gt; print(a) 
            tensor([[[1., 1., 1., 1.],
                 [1., 1., 1., 1.]]])
            &gt;&gt;&gt; a = pad_sequence(a,episode_end=False, sequence=5)
            &gt;&gt;&gt; print(a)
            tensor([[[0., 0., 0., 0.],
                     [0., 0., 0., 0.],
                     [0., 0., 0., 0.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.]]])
        """
        # assuming batch is always first
        out_dims = (len(states), sequence, states.size(2))
        out_tensor = states[0].data.new(*out_dims).fill_(padding_value)
        length = states.size(1)
        if episode_end:
            out_tensor[:, :length, :] = states
        else:
            out_tensor[:, -length:, :] = states
        return out_tensor

    def forward(self, x, hx=None, episode_end=False):
        r"""Return output and pad the input if observation is less than sequence."""
        if x.size(1) &lt; self.sequence:
            x = self.pad_sequence(x, episode_end, self.sequence)
        out, hx = self.gru(x, hx)
        out = F.relu(out.flatten(1))
        return self.fc(out), hx
</code></pre>

Ayman_Al_Jabri: <p>Unfortunately the GRU (or LSTM) layer model is not learning! I’ve tried many combinations and feeding techniques, but it didn’t learn at all <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"><br>
However, when I replace <strong>GRU</strong> layer with a <strong>GRUCell</strong>, the model works fine and solves the environment (in 6 seconds)</p>
<p>** <strong>My question is:</strong> what’s the difference between <strong>GRUCell</strong> and <strong>GRU</strong> Layer? And why this is happening?</p>
<p>=====================================================<br>
<em>Here’s the logic I’m using:</em></p>
<p><strong>Reinforce algorithm:</strong><br>
1 - initiate the model using random weights<br>
2 - play N full episodes gathering (state, action, reward)<br>
3- for each episode: calculate the discounted rewards for subsequent steps:</p>
<blockquote>
<p>𝑄(𝑘,𝑡) = ∑ 𝛾^𝑖 * 𝑟𝑖</p>
</blockquote>
<p>4-  Calculate the loss function and add entropy to improve exploration:</p>
<blockquote>
<p>ℒ = − ∑ 𝑄(𝑘,𝑡) log (𝜋(𝑠(𝑘,𝑡) , 𝑎(𝑘,𝑡)))</p>
</blockquote>
<p>5- perform SGD update of weights, minimizing loss<br>
6- repeat until solved</p>
<p><strong>When GRU layer:</strong><br>
a. hn = None at the beginning of each episode t=0<br>
b. hn from t-1 is then fed to the network at step t</p>
<p>=========================================================</p>
<p>Here’s the code. To run it as a GRUCell mode just add <code>--cell</code> argument</p>
<pre><code class="lang-auto">#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jun 26 14:39:33 2021

@author: Ayman Al Jabri
"""

import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import argparse
import numpy as np
from time import time
from collections import namedtuple
from itertools import count
from datetime import datetime, timedelta


Experience = namedtuple('Experience',['state','action','reward'])


class GRULayer(nn.Module):
    r"""GRU plain vanila policy gradient model."""

    def __init__(self, obs_size, act_size, hid_size=128, num_layers=2):
        super().__init__()
        self.hid_size = hid_size
        self.num_layers = num_layers
        
        self.input = nn.Linear(obs_size, 32)
        self.gru = nn.GRU(32, hid_size, num_layers, batch_first=True)
        self.output = nn.Linear(hid_size, act_size)

    def forward(self, x, hx=None):
        r"""Return output and pad the input if observation is less than sequence."""
        batch_size = x.size(0)
        y = self.input(x)
        y = y.view(batch_size, 1, -1)
        if hx is None:
            hx = torch.zeros((self.num_layers, x.size(0), self.hid_size))
            y, hx = self.gru(y, hx)
        else:
            y, hx = self.gru(y, hx)
        y = self.output(F.relu(y.flatten(1)))
        return y, hx
    
    
class GRUCell(nn.Module):
    r"""Plain vanilla policy gradient network."""

    def __init__(self, obs_size, act_size, hid_size=128):
        super().__init__()
        self.fc1 = nn.GRUCell(obs_size, hid_size)
        self.output = nn.Linear(hid_size, act_size, bias=True)

    def forward(self, x):
        """Feed forward."""
        y = F.relu(self.fc1(x))
        return self.output(y)
    

def discount_rewards(rewards, gamma):
    r"""
    Summary: calculate the discounted future rewards.

    Takes in list of rewards and discount rate
    Returns the accumlated future values of these rewards

    Example:
        &gt;&gt;&gt; r = [1,1,1,1,1,1]
        &gt;&gt;&gt; gamma = 0.9
        &gt;&gt;&gt; [4.68559, 4.0951, 3.439, 2.71, 1.9, 1.0]
    """
    sum_r = 0.0
    res = []
    for r in reversed(rewards):
        sum_r *= gamma
        sum_r += r
        res.append(sum_r)
    return list(reversed(res))


@torch.no_grad()
def generate_eipsodes(env, gamma, cell, n=3):
    r"""Yield n number of episodes' observations:(state,action,discounted_rewards,total_rewards,frames)"""
    episode = 0
    batch_total_rewards = 0
    hc = None
    act_size = env.action_space.n
    states,actions,rewards = [],[],[]
    dis_r = []
    state = env.reset()
    for frame in count():
        state_v = torch.FloatTensor([state])
        # if np.random.random() &gt; 0.5: # placeholder to injects noise
        #     state_v.fill_(0.0)
        #     pass
        if cell:
            prob = net(state_v) ##Linear
        else:
            prob, hc = net(state_v, hc) ##GRU
        prob = torch.softmax(prob,dim=1).data.cpu().numpy()[0]
        action = np.random.choice(act_size, p=prob)
        last_state, reward, done, _ = env.step(action)
        states.append(state)
        actions.append(action)
        rewards.append(reward)
        if done:
            dis_r.extend(discount_rewards(rewards, gamma))
            batch_total_rewards += sum(rewards)
            rewards.clear()
            episode += 1
            state=env.reset()
            hc = None
            if episode ==n:
                yield (np.array(states,copy=False),
                       np.array(actions),
                       np.array(dis_r), 
                       batch_total_rewards/n,
                       frame)
                states.clear()
                actions.clear()
                rewards.clear()
                dis_r.clear()
                batch_total_rewards = 0
                episode = 0
        state = last_state

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cell', action='store_true', help='Use GRUCell instead of GRU layer')
    args = parser.parse_args()
    
    ENTROPY_BETA = 0.02
    GAMMA = 0.99
    HID_SIZE = 64
    NUM_LAYERS = 2
    SOLVE = 195
    LR = 0.01
    N_EPS = 1
    SEED = 155

    env = gym.make('CartPole-v0') 
    env.seed(SEED)
    torch.manual_seed(SEED)
    np.random.seed(SEED)
    obs_size = env.observation_space.shape[0]
    act_size = env.action_space.n
    
    if args.cell:
        net = GRUCell(obs_size, act_size, HID_SIZE) #Linear
    else:        
        net = GRULayer(obs_size, act_size,HID_SIZE,NUM_LAYERS) #GRU
        
    print(net)
    
    optimizer = torch.optim.Adam(net.parameters(), lr=LR)

    total_rewards = []
    print_time = time()
    start = datetime.now()
    frame = 0
    prev_frame = 0
    mean = None
    for episode,batch in enumerate(generate_eipsodes(env, GAMMA, args.cell, n=N_EPS)):
        states, actions, rewards, batch_total_rewards,frame = batch
        total_rewards.append(batch_total_rewards)
        mean_reward = np.mean(total_rewards[-100:])
        if time() - print_time &gt; 1:
            speed = (frame - prev_frame)/(time()-print_time)
            prev_frame = frame
            print(
                f"{frame:,}: done {episode} episodes, mean reward {mean_reward:6.3f}, speed:{speed:.0f} fps", flush=True)
            print_time = time()

        if mean_reward &gt; SOLVE:
            duration = timedelta(seconds=(datetime.now()-start).seconds)
            print(f'Solved in {duration}')
            break
        
        ### training ###
        states_v = torch.FloatTensor(states)
        batch_scale_v = torch.FloatTensor(discount_rewards(rewards, GAMMA))
        actions_v = torch.LongTensor(actions)
        optimizer.zero_grad()
    
        # policy loss
        if args.cell:
          logit = net(states_v) #Linear
        else:
            logit, hn = net(states_v) ##GRU
        log_p = F.log_softmax(logit, dim=1)

        # Gather probabilities with taken actions
        log_p_a = batch_scale_v * log_p[range(len(actions)), actions]
        policy_loss = - log_p_a.mean()

        # entropy loss
        probs_v = F.softmax(logit, dim=1)
        entropy = - (probs_v * log_p).sum(dim=1).mean()
        entropy_loss = - ENTROPY_BETA * entropy
        
        #total loss
        loss = policy_loss + entropy_loss
        loss.backward()
        optimizer.step()

</code></pre>

ranon-rat: <p>I just started to build things with reinforcement learning and pytorch and i have been having this same problem.<br>
I tried using lstmCells with a custom model that i made but it still not working nor learning.<br>
This is just weird<br>
From what i have read its just a limitation that lstm have, that their output becomes smaller and smaller as the layers increases and i have been testing with other stuff.<br>
For some reason the lstm will just repeat the same patterns again in every escenario i was thinking that it may be caused by the hidden cells.</p>