diff --git a/README.md b/README.md index 977c6db..57a772a 100644 --- a/README.md +++ b/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. diff --git a/async_agent.py b/async_agent.py index d607f35..4e7ed94 100644 --- a/async_agent.py +++ b/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: - with self.network_lock: - q_next = np.max(self.target_network.predict(state)) - reward += self.discount * q_next - batch_rewards.append(reward) + value = worker_network.predict(state) + action = policy.sample(value.flatten()) + + batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards, + state, action, terminal, self) if episode_steps > self.step_limit: terminal = True diff --git a/bootstrap.py b/bootstrap.py new file mode 100644 index 0000000..788a91d --- /dev/null +++ b/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 diff --git a/main.py b/main.py index 0bea53d..be0bcdf 100644 --- a/main.py +++ b/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()