From ecb42d30d24d8753ae2ddcbb4a2321a2074e3970 Mon Sep 17 00:00:00 2001 From: Kyunghwan Kim Date: Sat, 16 Feb 2019 17:30:12 +0900 Subject: [PATCH] Add per (#8) * Add per and modify etc * Replace pre-commit-config.yaml and add pre-commit hook in .git * Modify .gitignore * Modify .gitignore * Modify buffer and code * Modify replay buffer and per * Modify .gitignore --- .gitignore | 9 + ...mit-config.yaml => .pre-commit-config.yaml | 0 scripts/algorithms/common/abstract/agent.py | 23 +- .../common/buffer/priortized_replay_buffer.py | 151 +++++++++++ .../algorithms/common/buffer/replay_buffer.py | 68 +++-- .../algorithms/common/buffer/segment_tree.py | 141 ++++++++++ scripts/algorithms/common/helper_functions.py | 12 + scripts/algorithms/common/noise.py | 4 +- scripts/algorithms/ddpg/agent.py | 11 +- scripts/algorithms/per/ddpg_agent.py | 251 ++++++++++++++++++ .../lunarlander_continuous_v2/ddpg.py | 16 +- .../lunarlander_continuous_v2/per-ddpg.py | 111 ++++++++ scripts/run_lunarlander_continuous.py | 8 +- 13 files changed, 750 insertions(+), 55 deletions(-) rename scripts/.pre-commit-config.yaml => .pre-commit-config.yaml (100%) create mode 100644 scripts/algorithms/common/buffer/priortized_replay_buffer.py create mode 100644 scripts/algorithms/common/buffer/segment_tree.py create mode 100644 scripts/algorithms/per/ddpg_agent.py create mode 100644 scripts/examples/lunarlander_continuous_v2/per-ddpg.py diff --git a/.gitignore b/.gitignore index 28bc0c0..df6e790 100644 --- a/.gitignore +++ b/.gitignore @@ -246,3 +246,12 @@ qtcreator-* # Catkin custom files CATKIN_IGNORE + +# Wandb +wandb + +# save file +save + +# pycharm +.idea \ No newline at end of file diff --git a/scripts/.pre-commit-config.yaml b/.pre-commit-config.yaml similarity index 100% rename from scripts/.pre-commit-config.yaml rename to .pre-commit-config.yaml diff --git a/scripts/algorithms/common/abstract/agent.py b/scripts/algorithms/common/abstract/agent.py index c52b5ba..48e4366 100644 --- a/scripts/algorithms/common/abstract/agent.py +++ b/scripts/algorithms/common/abstract/agent.py @@ -20,10 +20,9 @@ class AbstractAgent(ABC): """Abstract Agent used for all agents. Attributes: - env (gym.Env): openAI Gym environment with discrete action space + env (gym.Env): openAI Gym environment args (argparse.Namespace): arguments including hyperparameters and training settings - state_dim (int): dimension of state space - action_dim (int): dimension of action space + env_name (str) : gym env name for logging sha (str): sha code of current git commit """ @@ -32,7 +31,7 @@ class AbstractAgent(ABC): """Initialization. Args: - env (gym.Env): openAI Gym environment with discrete action space + env (gym.Env): openAI Gym environment args (argparse.Namespace): arguments including hyperparameters and training settings """ @@ -44,6 +43,7 @@ class AbstractAgent(ABC): self.args.max_episode_steps = env._max_episode_steps # for logging + self.env_name = str(self.env.env).split("<")[2].replace(">>", "") self.sha = ( subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])[:-1] .decode("ascii") @@ -67,13 +67,13 @@ class AbstractAgent(ABC): pass @abstractmethod - def save_params(self, name: str, params: dict, n_episode: int): + def save_params(self, params: dict, n_episode: int): if not os.path.exists("./save"): os.mkdir("./save") - path = os.path.join( - "./save/" + name + "_" + self.sha + "_ep_" + str(n_episode) + ".pt" - ) + save_name = self.env_name + "_" + self.args.algo + "_" + self.sha + + path = os.path.join("./save/" + save_name + "_ep_" + str(n_episode) + ".pt") torch.save(params, path) print("[INFO] Saved the model and optimizer to", path) @@ -92,6 +92,7 @@ class AbstractAgent(ABC): state = self.env.reset() done = False score = 0 + step = 0 while not done: if self.args.render and i_episode >= self.args.render_after: @@ -102,8 +103,12 @@ class AbstractAgent(ABC): state = next_state score += reward + step += 1 - print("[INFO] episode %d\ttotal score: %d" % (i_episode, score)) + print( + "[INFO] episode %d\tstep: %d\ttotal score: %d" + % (i_episode, step, score) + ) # termination self.env.close() diff --git a/scripts/algorithms/common/buffer/priortized_replay_buffer.py b/scripts/algorithms/common/buffer/priortized_replay_buffer.py new file mode 100644 index 0000000..a3cc83e --- /dev/null +++ b/scripts/algorithms/common/buffer/priortized_replay_buffer.py @@ -0,0 +1,151 @@ +# -*- coding: utf-8 -*- +"""Prioritized Replay buffer for baselines. + +- Author: Kh Kim +- Contact: kh.kim@medipixel.io +- Paper: https://arxiv.org/pdf/1511.05952.pdf + https://arxiv.org/pdf/1707.08817.pdf +""" + +import random +from typing import Tuple + +import numpy as np +import torch + +from algorithms.common.buffer.replay_buffer import ReplayBuffer +from algorithms.common.buffer.segment_tree import MinSegmentTree, SumSegmentTree + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class PrioritizedReplayBuffer(ReplayBuffer): + """Create Prioritized Replay buffer. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py + + Attributes: + buffer_size (int): size of replay buffer for experience + alpha (float): alpha parameter for prioritized replay buffer + tree_idx (int): next index of tree + sum_tree (SumSegmentTree): sum tree for prior + min_tree (MinSegmentTree): min tree for min prior to get max weight + init_priority (float): lower bound of priority + + """ + + def __init__( + self, buffer_size: int, batch_size: int, demo: list = None, alpha: float = 0.6 + ): + """Initialization. + + Args: + buffer_size (int): size of replay buffer for experience + batch_size (int): size of a batched sampled from replay buffer for training + demo (list): demonstration + alpha (float): alpha parameter for prioritized replay buffer + + """ + super(PrioritizedReplayBuffer, self).__init__(buffer_size, batch_size, demo) + assert alpha >= 0 + self.buffer_size = buffer_size + self.alpha = alpha + self.tree_idx = 0 + + # capacity must be positive and a power of 2. + tree_capacity = 1 + while tree_capacity < self.buffer_size: + tree_capacity *= 2 + + self.sum_tree = SumSegmentTree(tree_capacity) + self.min_tree = MinSegmentTree(tree_capacity) + self.init_priority = 1.0 + + # for init priority of demo + if demo: + for _ in range(len(demo)): + self.sum_tree[self.tree_idx] = self.init_priority ** self.alpha + self.min_tree[self.tree_idx] = self.init_priority ** self.alpha + self.tree_idx += 1 + + def add( + self, + state: np.ndarray, + action: np.ndarray, + reward: np.float64, + next_state: np.ndarray, + done: bool, + ): + """Add experience and priority.""" + idx = self.tree_idx + self.tree_idx = (self.tree_idx + 1) % self.buffer_size + super().add(state, action, reward, next_state, done) + + self.sum_tree[idx] = self.init_priority ** self.alpha + self.min_tree[idx] = self.init_priority ** self.alpha + + def extend(self, transitions: list): + """Add experiences to memory.""" + raise NotImplementedError + + def _sample_proportional(self, batch_size: int) -> list: + """Sample indices based on proportional.""" + indices = [] + p_total = self.sum_tree.sum(0, len(self.buffer) - 1) + segment = p_total / batch_size + for i in range(batch_size): + a = segment * i + b = segment * (i + 1) + upperbound = random.uniform(a, b) + idx = self.sum_tree.retrieve(upperbound) + indices.append(idx) + return indices + + def sample(self, beta: float = 0.4) -> Tuple[torch.Tensor, ...]: + """Sample a batch of experiences.""" + assert beta > 0 + + indices = self._sample_proportional(self.batch_size) + states, actions, rewards, next_states, dones, weights = [], [], [], [], [], [] + + # get max weight + p_min = self.min_tree.min() / self.sum_tree.sum() + max_weight = (p_min * len(self.buffer)) ** (-beta) + + for i in indices: + s, a, r, n_s, d = self.buffer[i] + states.append(np.array(s, copy=False)) + actions.append(np.array(a, copy=False)) + rewards.append(np.array(r, copy=False)) + next_states.append(np.array(n_s, copy=False)) + dones.append(np.array(float(d), copy=False)) + + # calculate weights + p_sample = self.sum_tree[i] / self.sum_tree.sum() + weight = (p_sample * len(self.buffer)) ** (-beta) + weights.append(weight / max_weight) + + states = torch.FloatTensor(np.array(states)).to(device) + actions = torch.FloatTensor(np.array(actions)).to(device) + rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device) + next_states = torch.FloatTensor(np.array(next_states)).to(device) + dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device) + weights = torch.FloatTensor(np.array(weights).reshape(-1, 1)).to(device) + + experiences = (states, actions, rewards, next_states, dones, weights, indices) + + return experiences + + def update_priorities(self, indices: list, priorities: np.ndarray): + """Update priorities of sampled transitions.""" + assert len(indices) == len(priorities) + + for idx, priority in zip(indices, priorities): + assert priority > 0 + assert 0 <= idx < len(self.buffer) + + self.sum_tree[idx] = priority ** self.alpha + self.min_tree[idx] = priority ** self.alpha + + self.init_priority = max(self.init_priority, priority) diff --git a/scripts/algorithms/common/buffer/replay_buffer.py b/scripts/algorithms/common/buffer/replay_buffer.py index a37ef45..7d4b45a 100644 --- a/scripts/algorithms/common/buffer/replay_buffer.py +++ b/scripts/algorithms/common/buffer/replay_buffer.py @@ -1,8 +1,7 @@ # -*- coding: utf-8 -*- """Replay buffer for baselines.""" -import random -from collections import deque +from typing import Tuple import numpy as np import torch @@ -18,55 +17,70 @@ class ReplayBuffer: ddpg-pendulum/ddpg_agent.py Attributes: - buffer (deque): deque of replay buffer + buffer (list): list of replay buffer batch_size (int): size of a batched sampled from replay buffer for training """ - def __init__(self, buffer_size, batch_size, seed, demo=None): + def __init__(self, buffer_size: int, batch_size: int, demo: list = None): """Initialize a ReplayBuffer object. Args: buffer_size (int): size of replay buffer for experience batch_size (int): size of a batched sampled from replay buffer for training - seed (int): random seed - demo (deque) : demonstration deque + demo (list) : demonstration list """ - self.buffer = deque(maxlen=buffer_size) if not demo else demo - + self.buffer = list() if not demo else demo + self.buffer_size = buffer_size self.batch_size = batch_size - random.seed(seed) + self.idx = 0 - def add(self, state, action, reward, next_state, done): + def add( + self, + state: np.ndarray, + action: np.ndarray, + reward: np.float64, + next_state: np.ndarray, + done: bool, + ): """Add a new experience to memory.""" - self.buffer.append((state, action, reward, next_state, done)) + data = (state, action, reward, next_state, done) - def extend(self, transitions): + if len(self.buffer) == self.buffer_size: + self.buffer[self.idx] = data + self.idx = (self.idx + 1) % self.buffer_size + else: + self.buffer.append(data) + + def extend(self, transitions: list): """Add experiences to memory.""" self.buffer.extend(transitions) - def sample(self): + def sample( + self + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Randomly sample a batch of experiences from memory.""" - experiences = random.sample(self.buffer, k=self.batch_size) + idxs = np.random.randint(0, len(self.buffer), size=self.batch_size) states, actions, rewards, next_states, dones = [], [], [], [], [] - for e in experiences: - states.append(np.expand_dims(e[0], axis=0)) - actions.append(e[1]) - rewards.append(e[2]) - next_states.append(np.expand_dims(e[3], axis=0)) - dones.append(e[4]) + for i in idxs: + s, a, r, n_s, d = self.buffer[i] + states.append(np.array(s, copy=False)) + actions.append(np.array(a, copy=False)) + rewards.append(np.array(r, copy=False)) + next_states.append(np.array(n_s, copy=False)) + dones.append(np.array(float(d), copy=False)) - states = torch.from_numpy(np.vstack(states)).float().to(device) - actions = torch.from_numpy(np.vstack(actions)).float().to(device) - rewards = torch.from_numpy(np.vstack(rewards)).float().to(device) - next_states = torch.from_numpy(np.vstack(next_states)).float().to(device) - dones = torch.from_numpy(np.vstack(dones).astype(np.uint8)).float().to(device) + states = torch.FloatTensor(np.array(states)).to(device) + actions = torch.FloatTensor(np.array(actions)).to(device) + rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device) + next_states = torch.FloatTensor(np.array(next_states)).to(device) + dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device) - return (states, actions, rewards, next_states, dones) + return states, actions, rewards, next_states, dones - def __len__(self): + def __len__(self) -> int: """Return the current size of internal memory.""" return len(self.buffer) diff --git a/scripts/algorithms/common/buffer/segment_tree.py b/scripts/algorithms/common/buffer/segment_tree.py new file mode 100644 index 0000000..0e2d44c --- /dev/null +++ b/scripts/algorithms/common/buffer/segment_tree.py @@ -0,0 +1,141 @@ +# -*- coding: utf-8 -*- +"""Segment tree for Proirtized Replay Buffer.""" + +import operator +from typing import Callable + + +class SegmentTree: + """ Create SegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + Attributes: + capacity (int) + tree (list) + operation (function) + + """ + + def __init__(self, capacity: int, operation: Callable, init_value: float): + """Initialization. + + Args: + capacity (int) + operation (function) + init_value (float) + + """ + assert ( + capacity > 0 and capacity & (capacity - 1) == 0 + ), "capacity must be positive and a power of 2." + self.capacity = capacity + self.tree = [init_value for _ in range(2 * capacity)] + self.operation = operation + + def _operate_helper( + self, start: int, end: int, node: int, node_start: int, node_end: int + ) -> float: + """Returns result of operation in segment.""" + if start == node_start and end == node_end: + return self.tree[node] + mid = (node_start + node_end) // 2 + if end <= mid: + return self._operate_helper(start, end, 2 * node, node_start, mid) + else: + if mid + 1 <= start: + return self._operate_helper(start, end, 2 * node + 1, mid + 1, node_end) + else: + return self.operation( + self._operate_helper(start, mid, 2 * node, node_start, mid), + self._operate_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end), + ) + + def operate(self, start: int = 0, end: int = 0) -> float: + """Returns result of applying `self.operation`.""" + if end <= 0: + end += self.capacity + end -= 1 + + return self._operate_helper(start, end, 1, 0, self.capacity - 1) + + def __setitem__(self, idx: int, val: float): + """Set value in tree.""" + idx += self.capacity + self.tree[idx] = val + + idx //= 2 + while idx >= 1: + self.tree[idx] = self.operation(self.tree[2 * idx], self.tree[2 * idx + 1]) + idx //= 2 + + def __getitem__(self, idx: int) -> float: + """Get real value in leaf node of tree.""" + assert 0 <= idx < self.capacity + + return self.tree[self.capacity + idx] + + +class SumSegmentTree(SegmentTree): + """ Create SumSegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + """ + + def __init__(self, capacity: int): + """Initialization. + + Args: + capacity (int) + + """ + super(SumSegmentTree, self).__init__( + capacity=capacity, operation=operator.add, init_value=0.0 + ) + + def sum(self, start: int = 0, end: int = 0) -> float: + """Returns arr[start] + ... + arr[end].""" + return super(SumSegmentTree, self).operate(start, end) + + def retrieve(self, upperbound: float) -> int: + """Find the highest index `i` about upper bound in the tree""" + assert 0 <= upperbound <= self.sum() + 1e-5 + + idx = 1 + + while idx < self.capacity: # while non-leaf + left = 2 * idx + right = left + 1 + if self.tree[left] > upperbound: + idx = 2 * idx + else: + upperbound -= self.tree[left] + idx = right + return idx - self.capacity + + +class MinSegmentTree(SegmentTree): + """ Create SegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + """ + + def __init__(self, capacity: int): + """Initialization. + + Args: + capacity (int) + + """ + super(MinSegmentTree, self).__init__( + capacity=capacity, operation=min, init_value=float("inf") + ) + + def min(self, start: int = 0, end: int = 0) -> float: + """Returns min(arr[start], ..., arr[end]).""" + return super(MinSegmentTree, self).operate(start, end) diff --git a/scripts/algorithms/common/helper_functions.py b/scripts/algorithms/common/helper_functions.py index 7a1e22f..9403580 100644 --- a/scripts/algorithms/common/helper_functions.py +++ b/scripts/algorithms/common/helper_functions.py @@ -5,6 +5,10 @@ - Contact: curt.park@medipixel.io """ +import random + +import gym +import numpy as np import torch import torch.nn as nn @@ -18,3 +22,11 @@ def soft_update(local: nn.Module, target: nn.Module, tau: float): """Soft-update: target = tau*local + (1-tau)*target.""" for t_param, l_param in zip(target.parameters(), local.parameters()): t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data) + + +def set_random_seed(seed: int, env: gym.Env): + """Set random seed""" + env.seed(seed) + torch.manual_seed(seed) + np.random.seed(seed) + random.seed(seed) diff --git a/scripts/algorithms/common/noise.py b/scripts/algorithms/common/noise.py index d99202e..fab9a69 100644 --- a/scripts/algorithms/common/noise.py +++ b/scripts/algorithms/common/noise.py @@ -15,7 +15,7 @@ class OUNoise: ddpg-pendulum/ddpg_agent.py """ - def __init__(self, size, seed, mu=0.0, theta=0.15, sigma=0.2): + def __init__(self, size, mu=0.0, theta=0.15, sigma=0.2): """Initialize parameters and noise process.""" self.state = np.float64(0.0) self.mu = mu * np.ones(size) @@ -23,8 +23,6 @@ class OUNoise: self.sigma = sigma self.reset() - random.seed(seed) - def reset(self): """Reset the internal state (= noise) to mean (mu).""" self.state = copy.copy(self.mu) diff --git a/scripts/algorithms/ddpg/agent.py b/scripts/algorithms/ddpg/agent.py index 3fa65c2..8806364 100644 --- a/scripts/algorithms/ddpg/agent.py +++ b/scripts/algorithms/ddpg/agent.py @@ -53,7 +53,7 @@ class Agent(AbstractAgent): """Initialization. Args: - env (gym.Env): openAI Gym environment with discrete action space + env (gym.Env): openAI Gym environment args (argparse.Namespace): arguments including hyperparameters and training settings hyper_params (dict): hyper-parameters models (tuple): models including actor and critic @@ -75,7 +75,7 @@ class Agent(AbstractAgent): # replay memory self.memory = ReplayBuffer( - hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"], self.args.seed + hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"] ) def select_action(self, state: np.ndarray) -> torch.Tensor: @@ -84,9 +84,10 @@ class Agent(AbstractAgent): state = torch.FloatTensor(state).to(device) selected_action = self.actor(state) - selected_action += torch.FloatTensor(self.noise.sample()).to(device) - selected_action = torch.clamp(selected_action, -1.0, 1.0) + if not self.args.test: + selected_action += torch.FloatTensor(self.noise.sample()).to(device) + selected_action = torch.clamp(selected_action, -1.0, 1.0) return selected_action @@ -163,7 +164,7 @@ class Agent(AbstractAgent): "critic_optim_state_dict": self.critic_optimizer.state_dict(), } - AbstractAgent.save_params(self, self.args.algo, params, n_episode) + AbstractAgent.save_params(self, params, n_episode) def write_log(self, i: int, loss: np.ndarray, score: int): """Write log about loss and score""" diff --git a/scripts/algorithms/per/ddpg_agent.py b/scripts/algorithms/per/ddpg_agent.py new file mode 100644 index 0000000..1ffeb0d --- /dev/null +++ b/scripts/algorithms/per/ddpg_agent.py @@ -0,0 +1,251 @@ +# -*- coding: utf-8 -*- +"""DDPG agent with PER for episodic tasks in OpenAI Gym. + +- Author: Kh Kim +- Contact: kh.kim@medipixel.io +- Paper: https://arxiv.org/pdf/1509.02971.pdf + https://arxiv.org/pdf/1511.05952.pdf +""" + +import argparse +import os +from typing import List, Tuple + +import gym +import numpy as np +import torch +import wandb + +import algorithms.common.helper_functions as common_utils +from algorithms.common.abstract.agent import AbstractAgent +from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBuffer +from algorithms.common.noise import OUNoise + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class Agent(AbstractAgent): + """ActorCritic interacting with environment. + + Attributes: + memory (PrioritizedReplayBuffer): replay memory + noise (OUNoise): random noise for exploration + actor (nn.Module): actor model to select actions + actor_target (nn.Module): target actor model to select actions + critic (nn.Module): critic model to predict state values + critic_target (nn.Module): target critic model to predict state values + actor_optimizer (Optimizer): optimizer for training actor + critic_optimizer (Optimizer): optimizer for training critic + hyper_params (dict): hyper-parameters + beta (float): beta parameter for prioritized replay buffer + curr_state (np.ndarray): temporary storage of the current state + + """ + + def __init__( + self, + env: gym.Env, + args: argparse.Namespace, + hyper_params: dict, + models: tuple, + optims: tuple, + noise: OUNoise, + ): + """Initialization. + + Args: + env (gym.Env): openAI Gym environment + args (argparse.Namespace): arguments including hyperparameters and training settings + hyper_params (dict): hyper-parameters + models (tuple): models including actor and critic + optims (tuple): optimizers for actor and critic + noise (OUNoise): random noise for exploration + + """ + AbstractAgent.__init__(self, env, args) + + self.actor, self.actor_target, self.critic, self.critic_target = models + self.actor_optimizer, self.critic_optimizer = optims + self.hyper_params = hyper_params + self.curr_state = np.zeros((1,)) + self.noise = noise + + # load the optimizer and model parameters + if args.load_from is not None and os.path.exists(args.load_from): + self.load_params(args.load_from) + + # replay memory + self.beta = self.hyper_params["PER_BETA"] + self.memory = PrioritizedReplayBuffer( + self.hyper_params["BUFFER_SIZE"], + self.hyper_params["BATCH_SIZE"], + alpha=self.hyper_params["PER_ALPHA"], + ) + + def select_action(self, state: np.ndarray) -> torch.Tensor: + """Select an action from the input space.""" + self.curr_state = state + + state = torch.FloatTensor(state).to(device) + + if not self.args.test: + selected_action = self.actor(state) + selected_action += torch.FloatTensor(self.noise.sample()).to(device) + + selected_action = torch.clamp(selected_action, -1.0, 1.0) + + return selected_action + + def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]: + """Take an action and return the response of the env.""" + action = action.detach().cpu().numpy() + next_state, reward, done, _ = self.env.step(action) + + self.memory.add(self.curr_state, action, reward, next_state, done) + + return next_state, reward, done + + def update_model( + self, + experiences: Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + List[int], + ], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Train the model after each episode.""" + + states, actions, rewards, next_states, dones, weights, indexes = experiences + + # G_t = r + gamma * v(s_{t+1}) if state != Terminal + # = r otherwise + masks = 1 - dones + next_actions = self.actor_target(next_states) + next_values = self.critic_target(torch.cat((next_states, next_actions), dim=-1)) + curr_returns = rewards + self.hyper_params["GAMMA"] * next_values * masks + curr_returns = curr_returns.to(device).detach() + + # train critic + values = self.critic(torch.cat((states, actions), dim=-1)) + critic_loss = torch.mean((values - curr_returns).pow(2) * weights) + self.critic_optimizer.zero_grad() + critic_loss.backward() + self.critic_optimizer.step() + + # train actor + actions = self.actor(states) + actor_loss_element_wise = -self.critic(torch.cat((states, actions), dim=-1)) + actor_loss = torch.mean(actor_loss_element_wise * weights) + self.actor_optimizer.zero_grad() + actor_loss.backward() + self.actor_optimizer.step() + + # update target networks + tau = self.hyper_params["TAU"] + common_utils.soft_update(self.actor, self.actor_target, tau) + common_utils.soft_update(self.critic, self.critic_target, tau) + + # update priorities in PER + new_priorities = (values - curr_returns).pow(2) + new_priorities = ( + new_priorities.data.cpu().numpy() + self.hyper_params["PER_EPS"] + ) + self.memory.update_priorities(indexes, new_priorities) + + return actor_loss.data, critic_loss.data + + def load_params(self, path: str): + """Load model and optimizer parameters.""" + if not os.path.exists(path): + print("[ERROR] the input path does not exist. ->", path) + return + + params = torch.load(path) + self.actor.load_state_dict(params["actor_state_dict"]) + self.actor_target.load_state_dict(params["actor_target_state_dict"]) + self.critic.load_state_dict(params["critic_state_dict"]) + self.critic_target.load_state_dict(params["critic_target_state_dict"]) + self.actor_optimizer.load_state_dict(params["actor_optim_state_dict"]) + self.critic_optimizer.load_state_dict(params["critic_optim_state_dict"]) + print("[INFO] loaded the model and optimizer from", path) + + def save_params(self, n_episode: int): + """Save model and optimizer parameters.""" + params = { + "actor_state_dict": self.actor.state_dict(), + "actor_target_state_dict": self.actor_target.state_dict(), + "critic_state_dict": self.critic.state_dict(), + "critic_target_state_dict": self.critic_target.state_dict(), + "actor_optim_state_dict": self.actor_optimizer.state_dict(), + "critic_optim_state_dict": self.critic_optimizer.state_dict(), + } + + AbstractAgent.save_params(self, params, n_episode) + + def write_log(self, i: int, loss: np.ndarray, score: int): + """Write log about loss and score""" + total_loss = loss.sum() + + print( + "[INFO] episode %d total score: %d, total loss: %f\n" + "actor_loss: %.3f critic_loss: %.3f\n" + % (i, score, total_loss, loss[0], loss[1]) # actor loss # critic loss + ) + + if self.args.log: + wandb.log( + { + "score": score, + "total loss": total_loss, + "actor loss": loss[0], + "critic loss": loss[1], + } + ) + + def train(self): + """Train the agent.""" + # logger + if self.args.log: + wandb.init() + wandb.config.update(self.hyper_params) + wandb.watch([self.actor, self.critic], log="parameters") + + for i_episode in range(1, self.args.episode_num + 1): + state = self.env.reset() + done = False + score = 0 + loss_episode = list() + + while not done: + if self.args.render and i_episode >= self.args.render_after: + self.env.render() + + action = self.select_action(state) + next_state, reward, done = self.step(action) + + if len(self.memory) >= self.hyper_params["BATCH_SIZE"]: + experiences = self.memory.sample(self.beta) + loss = self.update_model(experiences) + loss_episode.append(loss) # for logging + + state = next_state + score += reward + + # increase beta + fraction = min(float(i_episode) / self.args.max_episode_steps, 1.0) + self.beta = self.beta + fraction * (1.0 - self.beta) + + # logging + if loss_episode: + avg_loss = np.vstack(loss_episode).mean(axis=0) + self.write_log(i_episode, avg_loss, score) + + if i_episode % self.args.save_period == 0: + self.save_params(i_episode) + + # termination + self.env.close() diff --git a/scripts/examples/lunarlander_continuous_v2/ddpg.py b/scripts/examples/lunarlander_continuous_v2/ddpg.py index 988597e..dd917f6 100644 --- a/scripts/examples/lunarlander_continuous_v2/ddpg.py +++ b/scripts/examples/lunarlander_continuous_v2/ddpg.py @@ -41,31 +41,36 @@ def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int) action_dim (int): dimension of actions """ - hidden_sizes = [256, 256] + hidden_sizes_actor = [256, 256] + hidden_sizes_critic = [256, 256] # create actor actor = MLP( input_size=state_dim, output_size=action_dim, - hidden_sizes=hidden_sizes, + hidden_sizes=hidden_sizes_actor, output_activation=torch.tanh, ).to(device) actor_target = MLP( input_size=state_dim, output_size=action_dim, - hidden_sizes=hidden_sizes, + hidden_sizes=hidden_sizes_actor, output_activation=torch.tanh, ).to(device) actor_target.load_state_dict(actor.state_dict()) # create critic critic = MLP( - input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic_target = MLP( - input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic_target.load_state_dict(critic.state_dict()) @@ -85,7 +90,6 @@ def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int) # noise noise = OUNoise( action_dim, - args.seed, theta=hyper_params["OU_NOISE_THETA"], sigma=hyper_params["OU_NOISE_SIGMA"], ) diff --git a/scripts/examples/lunarlander_continuous_v2/per-ddpg.py b/scripts/examples/lunarlander_continuous_v2/per-ddpg.py new file mode 100644 index 0000000..08c9818 --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/per-ddpg.py @@ -0,0 +1,111 @@ +# -*- coding: utf-8 -*- +"""Run module for DDPG with PER on LunarLanderContinuous-v2. + +- Author: Curt Park +- Contact: curt.park@medipixel.io +""" + +import argparse + +import gym +import torch +import torch.optim as optim + +from algorithms.common.networks.mlp import MLP +from algorithms.common.noise import OUNoise +from algorithms.per.ddpg_agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +hyper_params = { + "GAMMA": 0.99, + "TAU": 1e-3, + "BUFFER_SIZE": int(1e5), + "BATCH_SIZE": 128, + "LR_ACTOR": 1e-4, + "LR_CRITIC": 1e-3, + "OU_NOISE_THETA": 0.0, + "OU_NOISE_SIGMA": 0.0, + "PER_ALPHA": 0.5, + "PER_BETA": 0.4, + "PER_EPS": 1e-6, + "WEIGHT_DECAY": 1e-6, +} + + +def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int): + """Run training or test. + + Args: + env (gym.Env): openAI Gym environment with continuous action space + args (argparse.Namespace): arguments including training settings + state_dim (int): dimension of states + action_dim (int): dimension of actions + + """ + hidden_sizes_actor = [256, 256] + hidden_sizes_critic = [256, 256] + + # create actor + actor = MLP( + input_size=state_dim, + output_size=action_dim, + hidden_sizes=hidden_sizes_actor, + output_activation=torch.tanh, + ).to(device) + + actor_target = MLP( + input_size=state_dim, + output_size=action_dim, + hidden_sizes=hidden_sizes_actor, + output_activation=torch.tanh, + ).to(device) + actor_target.load_state_dict(actor.state_dict()) + + # create critic + critic = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + + critic_target = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + critic_target.load_state_dict(critic.state_dict()) + + # create optimizer + actor_optim = optim.Adam( + actor.parameters(), + lr=hyper_params["LR_ACTOR"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + + critic_optim = optim.Adam( + critic.parameters(), + lr=hyper_params["LR_CRITIC"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + + # noise + noise = OUNoise( + action_dim, + theta=hyper_params["OU_NOISE_THETA"], + sigma=hyper_params["OU_NOISE_SIGMA"], + ) + + # make tuples to create an agent + models = (actor, actor_target, critic, critic_target) + optims = (actor_optim, critic_optim) + + # create an agent + agent = Agent(env, args, hyper_params, models, optims, noise) + + # run + if args.test: + agent.test() + else: + agent.train() diff --git a/scripts/run_lunarlander_continuous.py b/scripts/run_lunarlander_continuous.py index 5db4a01..a861eef 100644 --- a/scripts/run_lunarlander_continuous.py +++ b/scripts/run_lunarlander_continuous.py @@ -9,8 +9,8 @@ import argparse import importlib import gym -import numpy as np -import torch + +import algorithms.common.helper_functions as common_utils # configurations parser = argparse.ArgumentParser(description="Pytorch RL baselines") @@ -53,9 +53,7 @@ def main(): action_dim = env.action_space.shape[0] # set a random seed - env.seed(args.seed) - torch.manual_seed(args.seed) - np.random.seed(args.seed) + common_utils.set_random_seed(args.seed, env) # run module_path = "examples.lunarlander_continuous_v2." + args.algo