From ceaf83bca34148a3a5d130f028f3143652085605 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 31 Jan 2018 16:19:33 -0700 Subject: [PATCH] Support A2C --- agent/A2C_agent.py | 118 ++++++++++++++++++++++++++ agent/__init__.py | 3 +- component/task.py | 178 +++++++++++++--------------------------- main.py | 21 ++++- network/base_network.py | 2 + utils/config.py | 1 + utils/misc.py | 1 + 7 files changed, 201 insertions(+), 123 deletions(-) create mode 100644 agent/A2C_agent.py diff --git a/agent/A2C_agent.py b/agent/A2C_agent.py new file mode 100644 index 0000000..728accf --- /dev/null +++ b/agent/A2C_agent.py @@ -0,0 +1,118 @@ +####################################################################### +# 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 +import torch.multiprocessing as mp +from network import * +from utils import * +from component import * +import pickle +import os +import time +import gym.monitoring + +class A2CAgent: + def __init__(self, config): + self.config = config + self.task = config.task_fn() + self.evaluator = self.task.task_fn() + self.network = config.network_fn() + self.optimizer = config.optimizer_fn(self.network.parameters()) + self.policy = config.policy_fn() + self.total_steps = 0 + self.states = self.task.reset() + + self.episode_counts = np.zeros(config.num_workers) + self.episode_rewards = np.zeros(config.num_workers) + self.total_rewards = np.zeros(config.num_workers) + self.prev_episode_counts = 0.0 + self.prev_total_rewards = 0.0 + + def close(self): + self.task.close() + + def save(self, file_name): + with open(file_name, 'wb') as f: + torch.save(self.network.state_dict(), f) + + def evaluate(self): + state = self.evaluator.reset() + total_rewards = 0 + steps = 0 + while True: + prob, _, _ = self.network.predict(np.stack([state])) + action = self.policy.sample(prob.data.numpy().flatten(), True) + state, reward, done, _ = self.evaluator.step(action) + total_rewards += reward + steps += 1 + if done: + break + return total_rewards, steps + + def episode(self, deterministic=False): + if deterministic: + return self.evaluate() + + config = self.config + rollout = [] + states = self.states + for i in range(config.rollout_length): + prob, log_prob, value = self.network.predict(states) + actions = [self.policy.sample(p, deterministic) for p in prob.data.numpy()] + actions = config.action_shift_fn(actions) + next_states, rewards, terminals, _ = self.task.step(actions) + self.episode_rewards += rewards + rewards = config.reward_shift_fn(rewards) + for i, terminal in enumerate(terminals): + if terminals[i]: + next_states[i] = self.task.reset(i) + self.episode_counts[i] += 1 + self.total_rewards[i] += self.episode_rewards[i] + self.episode_rewards[i] = 0 + + rollout.append([prob, log_prob, value, actions, rewards, 1 - terminals]) + states = next_states + + self.states = states + _, _, pending_value = self.network.predict(states) + rollout.append([None, None, pending_value, None, None, None]) + + processed_rollout = [None] * (len(rollout) - 1) + advantages = self.network.FloatTensor(np.zeros((config.num_workers, 1))) + returns = pending_value.data + for i in reversed(range(len(rollout) - 1)): + prob, log_prob, value, actions, rewards, terminals = rollout[i] + terminals = self.network.FloatTensor(terminals).unsqueeze(1) + rewards = self.network.FloatTensor(rewards).unsqueeze(1) + actions = self.network.LongTensor(actions).unsqueeze(1) + next_value = rollout[i + 1][2] + returns = rewards + terminals * config.discount * returns + td_error = rewards + config.discount * terminals * next_value.data - value.data + advantages = advantages * config.gae_tau * config.discount * terminals + td_error + processed_rollout[i] = [prob, log_prob, value, actions, returns, advantages] + + prob, log_prob, value, actions, returns, advantages = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout)) + policy_loss = -log_prob.gather(1, Variable(actions)) * Variable(advantages) + policy_loss += config.entropy_weight * torch.sum(prob * log_prob, dim=1, keepdim=True) + value_loss = 0.5 * (Variable(returns) - value).pow(2) + + self.optimizer.zero_grad() + (policy_loss + value_loss).mean().backward() + nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip) + self.optimizer.step() + + steps = config.rollout_length * config.num_workers + self.total_steps += steps + new_episode_counts = np.sum(self.episode_counts) + new_total_rewards = np.sum(self.total_rewards) + avg_reward = (new_total_rewards - self.prev_total_rewards) / \ + (new_episode_counts - self.prev_episode_counts + 1e-5) + self.prev_total_rewards = new_total_rewards + self.prev_episode_counts = new_episode_counts + return avg_reward, steps + + + diff --git a/agent/__init__.py b/agent/__init__.py index ab0c05d..b4a371e 100644 --- a/agent/__init__.py +++ b/agent/__init__.py @@ -1,3 +1,4 @@ from .async_agent import * from .DQN_agent import * -from .DDPG_agent import * \ No newline at end of file +from .DDPG_agent import * +from .A2C_agent import * \ No newline at end of file diff --git a/component/task.py b/component/task.py index 8079d34..d08a4ea 100644 --- a/component/task.py +++ b/component/task.py @@ -7,15 +7,20 @@ import gym import sys import numpy as np from .atari_wrapper import * +import torch.multiprocessing as mp +import sys class BasicTask: - def __init__(self): + def __init__(self, max_steps=sys.maxsize): self.normalized_state = True + self.steps = 0 + self.max_steps = max_steps def normalize_state(self, state): return state def reset(self): + self.steps = 0 state = self.env.reset() if self.normalized_state: return self.normalize_state(state) @@ -23,6 +28,8 @@ class BasicTask: def step(self, action): next_state, reward, done, info = self.env.step(action) + self.steps += 1 + done = (done or self.steps >= self.max_steps) if self.normalized_state: next_state = self.normalize_state(next_state) return next_state, np.sign(reward), done, info @@ -34,8 +41,8 @@ class MountainCar(BasicTask): name = 'MountainCar-v0' success_threshold = -110 - def __init__(self): - BasicTask.__init__(self) + def __init__(self, max_steps=200): + BasicTask.__init__(self, max_steps) self.env = gym.make(self.name) self.env._max_episode_steps = sys.maxsize @@ -43,8 +50,8 @@ class CartPole(BasicTask): name = 'CartPole-v0' success_threshold = 195 - def __init__(self): - BasicTask.__init__(self) + def __init__(self, max_steps=200): + BasicTask.__init__(self, max_steps) self.env = gym.make(self.name) self.env._max_episode_steps = sys.maxsize @@ -182,121 +189,50 @@ class Roboschool(BasicTask): next_state, reward, done, info = self.env.step(action) return next_state, reward, done, info -class Fruit(BasicTask): - def __init__(self, hybrid_reward=False, pseudo_reward=False, atomic_state=True): - self.hybrid_reward = hybrid_reward - self.atomic_state = atomic_state - self.pseudo_reward = pseudo_reward - self.name = "Fruit" - self.success_threshold = 5 - self.width = 10 - self.height = 10 - self.possible_fruits = 10 - self.actual_fruits = 5 - xs = np.random.randint(0, self.width, size=self.possible_fruits) - ys = np.random.randint(0, self.height, size=self.possible_fruits) - self.possible_locations = list(zip(xs, ys)) - self.x = 0 - self.y = 0 - self.indices = np.arange(self.possible_fruits) - self.taken = [] - self.remaining_fruits = 0 - - def get_nearest(self): - def distance(i): - x, y = self.possible_locations[i] - return np.abs(self.x - x) + np.abs(self.y - y) - pool = [] - for i in range(self.possible_fruits): - if not self.taken[i]: - pool.append([i, distance(i)]) - pool = sorted(pool, key=lambda x:x[1]) - return pool[0][0] - - def encode_pos(self, x, y): - return '{:04b}'.format(x) + '{:04b}'.format(y) - - def encode_atomic_state(self): - offset = 8 * self.possible_fruits - state = np.copy(self.base_state) - str = self.encode_pos(self.x, self.y) - for i in range(len(str)): - state[offset + i] = int(str[i]) - offset += 8 - for i in range(len(self.taken)): - state[offset + i] = self.taken[i] - return state - - def encode_decomposed_state(self): - state_size = (4 + 4) * 2 + 1 - base_state = np.zeros(state_size) - str = self.encode_pos(self.x, self.y) - for i in range(len(str)): - base_state[i] = int(str[i]) - states = [] - for i in range(self.possible_fruits): - states.append(np.copy(base_state)) - str = self.encode_pos(*self.possible_locations[i]) - for j in range(len(str)): - states[-1][8 + j] = int(str[j]) - states[-1][-1] = self.taken[i] - return np.asarray(states) - - def encode_state(self): - if self.atomic_state: - return self.encode_atomic_state() - return self.encode_decomposed_state() - - def reset(self): - self.x = np.random.randint(0, self.width) - self.y = np.random.randint(0, self.height) - np.random.shuffle(self.indices) - self.taken = np.ones(self.possible_fruits, dtype=np.bool) - self.taken[self.indices[: self.actual_fruits]] = False - self.remaining_fruits = self.actual_fruits - state_size = (4 + 4) * (self.possible_fruits + 1) + self.possible_fruits - self.base_state = np.zeros(state_size) - offset = 0 - for x, y in self.possible_locations: - str = self.encode_pos(x, y) - for i in range(len(str)): - self.base_state[offset + i] = int(str[i]) - offset += 8 - return self.encode_state() - - def step(self, action): - # action = action[0] - if action == 0: - self.x -= 1 - elif action == 1: - self.x += 1 - elif action == 2: - self.y -= 1 - elif action == 3: - self.y += 1 +def sub_task(parent_pipe, pipe, task_fn): + parent_pipe.close() + task = task_fn() + while True: + op, data = pipe.recv() + if op == 'step': + pipe.send(task.step(data)) + elif op == 'reset': + pipe.send(task.reset()) + elif op == 'exit': + pipe.close() + return else: - assert False - self.x = min(max(self.x, 0), self.width - 1) - self.y = min(max(self.y, 0), self.height - 1) - try: - pos = self.possible_locations.index((self.x, self.y)) - except ValueError: - pos = -1 - if self.hybrid_reward: - reward = np.zeros(self.possible_fruits) - if pos >= 0 and not self.taken[pos]: - reward[pos] = 10 - self.taken[pos] = True - self.remaining_fruits -= 1 - if self.pseudo_reward: - pseudo_reward = np.zeros(self.possible_fruits) - if pos >= 0: - pseudo_reward[pos] = 1 - reward = (reward, pseudo_reward) + assert False, 'Unknown Operation' + +class ParallelizedTask: + def __init__(self, task_fn, num_workers): + self.task_fn = task_fn + self.task = task_fn() + self.name = self.task.name + self.pipes, worker_pipes = zip(*[mp.Pipe() for _ in range(num_workers)]) + args = [(p, wp, task_fn) for p, wp in zip(self.pipes, worker_pipes)] + self.workers = [mp.Process(target=sub_task, args=arg) for arg in args] + for p in self.workers: p.start() + for p in worker_pipes: p.close() + + def step(self, actions): + for pipe, action in zip(self.pipes, actions): + pipe.send(('step', action)) + results = [p.recv() for p in self.pipes] + results = map(lambda x: np.stack(x), zip(*results)) + return results + + def reset(self, i=None): + if i is None: + for pipe in self.pipes: + pipe.send(('reset', None)) + results = [p.recv() for p in self.pipes] else: - reward = 0.0 - if pos >= 0 and not self.taken[pos]: - reward = 1.0 - self.taken[pos] = True - self.remaining_fruits -= 1 - return self.encode_state(), reward, not self.remaining_fruits, self.taken \ No newline at end of file + self.pipes[i].send(('reset', None)) + results = self.pipes[i].recv() + return np.stack(results) + + def close(self): + for pipe in self.pipes: + pipe.send(('exit', None)) + for p in self.workers: p.join() \ No newline at end of file diff --git a/main.py b/main.py index 9619f87..559f8af 100644 --- a/main.py +++ b/main.py @@ -69,6 +69,24 @@ def a3c_cart_pole(): agent = AsyncAgent(config) agent.run() +def a2c_cart_pole(): + config = Config() + task_fn = lambda: CartPole(max_steps=200) + config.num_workers = 3 + config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers) + config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001) + config.network_fn = lambda: ActorCriticFCNet(4, 2) + config.policy_fn = SamplePolicy + config.discount = 0.99 + config.test_interval = 20 + config.test_repetitions = 10 + config.logger = Logger('./log', logger) + config.gae_tau = 1.0 + config.entropy_weight = 0.01 + config.rollout_length = 50 + config.success_threshold = 195 + run_episodes(A2CAgent(config)) + def dqn_pixel_atari(name): config = Config() config.history_length = 4 @@ -317,9 +335,10 @@ if __name__ == '__main__': # logger.setLevel(logging.DEBUG) logger.setLevel(logging.INFO) - dqn_cart_pole() + # dqn_cart_pole() # async_cart_pole() # a3c_cart_pole() + a2c_cart_pole() # a3c_continuous() # p3o_continuous() # d3pg_continuous() diff --git a/network/base_network.py b/network/base_network.py index 3e270e5..c884d88 100644 --- a/network/base_network.py +++ b/network/base_network.py @@ -18,8 +18,10 @@ class BasicNet: if self.gpu: self.cuda() self.FloatTensor = torch.cuda.FloatTensor + self.LongTensor = torch.cuda.LongTensor else: self.FloatTensor = torch.FloatTensor + self.LongTensor = torch.LongTensor def to_torch_variable(self, x, dtype='float32'): if isinstance(x, Variable): diff --git a/utils/config.py b/utils/config.py index 9df817d..781cb75 100644 --- a/utils/config.py +++ b/utils/config.py @@ -51,3 +51,4 @@ class Config: self.max_steps = 0 self.success_threshold = float('inf') self.render_episode_freq = 0 + self.rollout_length = None diff --git a/utils/misc.py b/utils/misc.py index 8510196..feb3d8d 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -60,6 +60,7 @@ def run_episodes(agent): if avg_reward > config.success_threshold: break + agent.close() return steps, rewards, avg_test_rewards def sync_grad(target_network, src_network):