N-Step Q-Learning and One-Step Sarsa

README.md

@@ -2,6 +2,8 @@
 > Highly modularized implementation of popular deep RL algorithms powered by PyTorch
 * Deep Q-Learning
 * Asynchronous One-Step Q-Learning
+* Asynchronous One-Step Sarsa 
+* Asynchronous N-Step Q-Learning
 
 >Benchmarked by classical control tasks (CartPole, LunarLander). Atari games will make it difficult to replicate in a regular laptop without a good GPU. However it's fairly easy to adapt the components to fit Atari games.
 

async_agent.py

@@ -10,9 +10,10 @@ import numpy as np
 import torch.multiprocessing as mp
 from task import *
 from network import *
+from bootstrap import *
 
 class AsyncAgent:
-    def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, discount, step_limit,
+    def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, bootstrap_fn, discount, step_limit,
                  target_network_update_freq, n_workers, batch_size, test_interval, test_repeats):
         self.network_fn = network_fn
         self.learning_network = network_fn()
@@ -20,6 +21,7 @@ class AsyncAgent:
         self.target_network = network_fn()
         self.target_network.share_memory()
         self.target_network.load_state_dict(self.learning_network.state_dict())
+        self.bootstrap_fn = bootstrap_fn
 
         self.optimizer_fn = optimizer_fn
         self.task_fn = task_fn
@@ -81,22 +83,23 @@ class AsyncAgent:
                 terminal = False
                 state = task.reset()
                 state = state.reshape([1, -1])
+                value = worker_network.predict(state)
+                action = policy.sample(value.flatten())
             while not terminal and len(batch_states) < self.batch_size:
                 episode_steps += 1
                 with self.steps_lock:
                     self.total_steps.value += 1
                 batch_states.append(state)
-                value = worker_network.predict(state)
-                action = policy.sample(value.flatten())
                 batch_actions.append(action)
                 state, reward, terminal, _ = task.step(action)
+                batch_rewards.append(reward)
                 episode_return += reward
                 state = state.reshape([1, -1])
-                if not terminal:
+                value = worker_network.predict(state)
-                    with self.network_lock:
+                action = policy.sample(value.flatten())
-                        q_next = np.max(self.target_network.predict(state))
+
-                    reward += self.discount * q_next
+            batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards,
-                batch_rewards.append(reward)
+                                              state, action, terminal, self)
 
             if episode_steps > self.step_limit:
                 terminal = True

bootstrap.py

@@ -0,0 +1,45 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+import numpy as np
+
+def NStepQLearning(batch_states, batch_actions, batch_rewards,
+                   tailing_state, tailing_action, terminal, agent):
+    if terminal:
+        reward = 0
+    else:
+        with agent.network_lock:
+            reward = np.max(agent.target_network.predict(tailing_state))
+    rewards = []
+    for r in reversed(batch_rewards):
+        reward = r + agent.discount * reward
+        rewards.append(reward)
+    return rewards
+
+def OneStepQLearning(batch_states, batch_actions, batch_rewards,
+                     tailing_state, tailing_action, terminal, agent):
+    batch_states.append(tailing_state)
+    with agent.network_lock:
+        q_next = agent.target_network.predict(np.vstack(batch_states[1:]))
+    q_next = np.max(q_next, axis=1)
+    if terminal:
+        q_next[-1] = 0
+    batch_states.pop(-1)
+    batch_rewards = np.asarray(batch_rewards) + agent.discount * q_next
+    return batch_rewards
+
+def OneStepSarsa(batch_states, batch_actions, batch_rewards,
+                 tailing_state, tailing_action, terminal, agent):
+    batch_states.append(tailing_state)
+    batch_actions.append(tailing_action)
+    with agent.network_lock:
+        q_next = agent.target_network.predict(np.vstack(batch_states[1:]))
+    q_next = q_next[np.arange(len(batch_actions[1:])), batch_actions[1:]]
+    if terminal:
+        q_next[-1] = 0
+    batch_states.pop(-1)
+    batch_actions.pop(-1)
+    batch_rewards = np.asarray(batch_rewards) + agent.discount * q_next
+    return batch_rewards

main.py

@@ -7,6 +7,9 @@ def async_cart_pole():
     config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001)
     config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2])
     config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1)
+    # config['bootstrap_fn'] = OneStepQLearning
+    # config['bootstrap_fn'] = NStepQLearning
+    config['bootstrap_fn'] = OneStepSarsa
     config['discount'] = 0.99
     config['target_network_update_freq'] = 200
     config['step_limit'] = 300
@@ -17,6 +20,23 @@ def async_cart_pole():
     agent = AsyncAgent(**config)
     agent.run()
 
+def async_lunar_lander():
+    config = dict()
+    config['task_fn'] = lambda: LunarLander()
+    config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
+    config['network_fn'] = lambda: FullyConnectedNet([8, 50, 200, 4])
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=2000, min_epsilon=0.05)
+    config['bootstrap_fn'] = OneStepQLearning
+    config['discount'] = 0.99
+    config['target_network_update_freq'] = 200
+    config['step_limit'] = 5000
+    config['n_workers'] = 8
+    config['batch_size'] = 10
+    config['test_interval'] = 1000
+    config['test_repeats'] = 5
+    agent = AsyncAgent(**config)
+    agent.run()
+
 # Mountain Car is fairly unstable
 def dqn_mountain_car():
     config = dict()
@@ -31,22 +51,6 @@ def dqn_mountain_car():
     agent = DQNAgent(**config)
     agent.run()
 
-def async_lunar_lander():
-    config = dict()
-    config['task_fn'] = lambda: LunarLander()
-    config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
-    config['network_fn'] = lambda: FullyConnectedNet([8, 50, 200, 4])
-    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=2000, min_epsilon=0.05)
-    config['discount'] = 0.99
-    config['target_network_update_freq'] = 200
-    config['step_limit'] = 5000
-    config['n_workers'] = 8
-    config['batch_size'] = 10
-    config['test_interval'] = 1000
-    config['test_repeats'] = 5
-    agent = AsyncAgent(**config)
-    agent.run()
-
 def dqn_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()