Feat: Use memory in TNPG TRPO

Author: choru-k

PG/5-TNPG/model.py

@@ -76,9 +76,9 @@ def conjugate_gradient(net, states, loss_grad, n_step=10, residual_tol=1e-10):
             break
     return x
 
-class QNet(nn.Module):
+class TNPG(nn.Module):
     def __init__(self, num_inputs, num_outputs):
-        super(QNet, self).__init__()
+        super(TNPG, self).__init__()
         self.t = 0
         self.num_inputs = num_inputs
         self.num_outputs = num_outputs
@@ -97,7 +97,7 @@ def forward(self, input):
         return policy
 
     @classmethod
-    def train_model(cls, net, transitions, k):
+    def train_model(cls, net, transitions):
         states, actions, rewards, masks = transitions.state, transitions.action, transitions.reward, transitions.mask
 
         states = torch.stack(states)

PG/5-TNPG/train.py

@@ -7,7 +7,7 @@
 import torch
 import torch.optim as optim
 import torch.nn.functional as F
-from model import QNet
+from model import TNPG
 from tensorboardX import SummaryWriter
 
 from memory import Memory
@@ -24,18 +24,17 @@ def main():
     print('state size:', num_inputs)
     print('action size:', num_actions)
 
-    net = QNet(num_inputs, num_actions)
+    net = TNPG(num_inputs, num_actions)
     writer = SummaryWriter('logs')
 
     net.to(device)
     net.train()
     running_score = 0
     steps = 0
     loss = 0
-    k=0
     for e in range(30000):
         done = False
-        memory = Memory() 
+        memory = Memory()
 
         score = 0
         state = env.reset()
@@ -61,9 +60,7 @@ def main():
             score += reward
             state = next_state
 
-        sum_reward = 0
-        loss = QNet.train_model(net, memory.sample(), k)
-        k+=1
+        loss = TNPG.train_model(net, memory.sample())
 
         score = score if score == 500.0 else score + 1
         running_score = 0.99 * running_score + 0.01 * score

PG/6-TRPO/memory.py

@@ -0,0 +1,18 @@
+import random
+from collections import namedtuple, deque
+
+Transition = namedtuple('Transition', ('state', 'next_state', 'action', 'reward', 'mask'))
+
+class Memory(object):
+    def __init__(self):
+        self.memory = deque()
+
+    def push(self, state, next_state, action, reward, mask):
+        self.memory.append(Transition(state, next_state, action, reward, mask))
+
+    def sample(self):
+        memory = self.memory
+        return Transition(*zip(*memory)) 
+
+    def __len__(self):
+        return len(self.memory)

PG/6-TRPO/model.py

@@ -75,9 +75,9 @@ def conjugate_gradient(net, states, loss_grad, n_step=10, residual_tol=1e-10):
             break
     return x
 
-class QNet(nn.Module):
+class TRPO(nn.Module):
     def __init__(self, num_inputs, num_outputs):
-        super(QNet, self).__init__()
+        super(TRPO, self).__init__()
         self.t = 0
         self.num_inputs = num_inputs
         self.num_outputs = num_outputs
@@ -90,19 +90,27 @@ def __init__(self, num_inputs, num_outputs):
                 nn.init.xavier_uniform(m.weight)
 
     def forward(self, input):
-        x = torch.tanh(self.fc_1(input))
+        x = torch.relu(self.fc_1(input))
         policy = F.softmax(self.fc_2(x))
 
         return policy
 
     @classmethod
-    def train_model(cls, net, transitions, k):
-        states, actions, rewards, masks = transitions
+    def train_model(cls, net, transitions):
+        states, actions, rewards, masks = transitions.state, transitions.action, transitions.reward, transitions.mask
+
         states = torch.stack(states)
         actions = torch.stack(actions)
         rewards = torch.Tensor(rewards)
         masks = torch.Tensor(masks)
 
+        returns = torch.zeros_like(rewards)
+
+        running_return = 0
+        for t in reversed(range(len(rewards))):
+            running_return = rewards[t] + gamma * running_return * masks[t]
+            returns[t] = running_return
+
         policy = net(states)
         policy = policy.view(-1, net.num_outputs)
         policy_action = (policy * actions.detach()).sum(dim=1)
@@ -111,7 +119,7 @@ def train_model(cls, net, transitions, k):
         old_policy = old_policy.view(-1, net.num_outputs)
         old_policy_action = (old_policy * actions.detach()).sum(dim=1)
 
-        surrogate_loss = ((policy_action / old_policy_action) * rewards).mean()
+        surrogate_loss = ((policy_action / old_policy_action) * returns).mean()
 
         surrogate_loss_grad = torch.autograd.grad(surrogate_loss, net.parameters())
         surrogate_loss_grad = flat_grad(surrogate_loss_grad)
@@ -130,7 +138,7 @@ def train_model(cls, net, transitions, k):
             policy = net(states)
             policy = policy.view(-1, net.num_outputs)
             policy_action = (policy * actions.detach()).sum(dim=1)
-            surrogate_loss = ((policy_action / old_policy_action) * rewards).mean()
+            surrogate_loss = ((policy_action / old_policy_action) * returns).mean()
 
             kl = kl_divergence(policy, old_policy)
             kl = kl.mean()
@@ -144,6 +152,6 @@ def train_model(cls, net, transitions, k):
     def get_action(self, input):
         policy = self.forward(input)
         policy = policy[0].data.numpy()
-
+            
         action = np.random.choice(self.num_outputs, 1, p=policy)[0]
         return action

PG/6-TRPO/train.py

@@ -7,9 +7,10 @@
 import torch
 import torch.optim as optim
 import torch.nn.functional as F
-from model import QNet
+from model import TRPO 
 from tensorboardX import SummaryWriter
 
+from memory import Memory
 from config import env_name, goal_score, log_interval, device, gamma
 
 
@@ -23,19 +24,17 @@ def main():
     print('state size:', num_inputs)
     print('action size:', num_actions)
 
-    net = QNet(num_inputs, num_actions)
-
+    net = TRPO(num_inputs, num_actions)
     writer = SummaryWriter('logs')
 
     net.to(device)
     net.train()
     running_score = 0
     steps = 0
     loss = 0
-    k=0
-    for e in range(3000):
+    for e in range(30000):
         done = False
-        memory = []
+        memory = Memory()
 
         score = 0
         state = env.reset()
@@ -56,23 +55,12 @@ def main():
 
             action_one_hot = torch.zeros(2)
             action_one_hot[action] = 1
-            memory.append([state, next_state, action_one_hot, reward, mask])
+            memory.push(state, next_state, action_one_hot, reward, mask)
 
             score += reward
             state = next_state
 
-        sum_reward = 0
-        memory.reverse()
-        states, actions, rewards, masks = [], [], [], []
-        for t, transition in enumerate(memory):
-            state, next_state, action, reward, mask = transition
-            sum_reward = (reward + gamma * sum_reward)
-            states.append(state)
-            actions.append(action)
-            rewards.append(sum_reward)
-            masks.append(mask)
-        loss = QNet.train_model(net, (states, actions, rewards, masks), k)
-        k+=1
+        loss = TRPO.train_model(net, memory.sample())
 
         score = score if score == 500.0 else score + 1
         running_score = 0.99 * running_score + 0.01 * score