diff --git a/scripts/algorithms/common/abstract/agent.py b/scripts/algorithms/common/abstract/agent.py index 48e4366..23846e2 100644 --- a/scripts/algorithms/common/abstract/agent.py +++ b/scripts/algorithms/common/abstract/agent.py @@ -37,6 +37,7 @@ class AbstractAgent(ABC): """ self.args = args self.env = NormalizedActions(env) + if self.args.max_episode_steps > 0: env._max_episode_steps = self.args.max_episode_steps else: diff --git a/scripts/algorithms/common/buffer/priortized_replay_buffer.py b/scripts/algorithms/common/buffer/priortized_replay_buffer.py index b7c3361..c45aaed 100644 --- a/scripts/algorithms/common/buffer/priortized_replay_buffer.py +++ b/scripts/algorithms/common/buffer/priortized_replay_buffer.py @@ -31,13 +31,11 @@ class PrioritizedReplayBuffer(ReplayBuffer): 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 + _max_priority (float): max priority """ - def __init__( - self, buffer_size: int, batch_size: int, alpha: float = 0.6 - ): + def __init__(self, buffer_size: int, batch_size: int, alpha: float = 0.6): """Initialization. Args: @@ -59,7 +57,7 @@ class PrioritizedReplayBuffer(ReplayBuffer): self.sum_tree = SumSegmentTree(tree_capacity) self.min_tree = MinSegmentTree(tree_capacity) - self.init_priority = 1.0 + self._max_priority = 1.0 def add( self, @@ -74,8 +72,8 @@ class PrioritizedReplayBuffer(ReplayBuffer): 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 + self.sum_tree[idx] = self._max_priority ** self.alpha + self.min_tree[idx] = self._max_priority ** self.alpha def extend(self, transitions: list): """Add experiences to memory.""" @@ -140,4 +138,143 @@ class PrioritizedReplayBuffer(ReplayBuffer): self.sum_tree[idx] = priority ** self.alpha self.min_tree[idx] = priority ** self.alpha - self.init_priority = max(self.init_priority, priority) + self._max_priority = max(self._max_priority, priority) + + +class PrioritizedReplayBufferfD(PrioritizedReplayBuffer): + """Create Prioritized Replay buffer with demo. + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py + Attributes: + demo (list): list of demo replay buffer + buffer_size (int): size of replay buffer for experience + demo_size (int): size of replay buffer for demonstration + total_size (int): sum of demo size and number of samples of experience + epsilon_d (float) : epsilon_d parameter to update priority using demo + """ + + def __init__( + self, + buffer_size: int, + batch_size: int, + demo: list, + alpha: float = 0.6, + epsilon_d: float = 1.0, + ): + """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 + epsilon_d (float) : epsilon_d parameter to update priority using demo + """ + super(PrioritizedReplayBufferfD, self).__init__(buffer_size, batch_size, alpha) + self.demo = demo + self.demo_size = len(demo) + self.total_size = self.demo_size + len(self.buffer) + self.epsilon_d = epsilon_d + + # for init priority of demo + for _ in range(self.demo_size): + self.sum_tree[self.tree_idx] = self._max_priority ** self.alpha + self.min_tree[self.tree_idx] = self._max_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 + # buffer is full + if (self.tree_idx + 1) % (self.buffer_size + self.demo_size) == 0: + self.tree_idx = self.demo_size + else: + self.tree_idx = self.tree_idx + 1 + super().add(state, action, reward, next_state, done) + + self.sum_tree[idx] = self._max_priority ** self.alpha + self.min_tree[idx] = self._max_priority ** self.alpha + + # update current total size + self.total_size = self.demo_size + len(self.buffer) + + 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, eps_d = [], [] + + # get max weight + p_min = self.min_tree.min() / self.sum_tree.sum() + max_weight = (p_min * self.total_size) ** (-beta) + + for i in indices: + # sample from buffer + if i < self.demo_size: + s, a, r, n_s, d = self.demo[i] + eps_d.append(self.epsilon_d) + else: + s, a, r, n_s, d = self.buffer[i - self.demo_size] + eps_d.append(0.0) + + # append transition info + 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 * self.total_size) ** (-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) + eps_d = np.array(eps_d) + + if torch.cuda.is_available(): + states_ = states_.cuda(non_blocking=True) + actions_ = actions_.cuda(non_blocking=True) + rewards_ = rewards_.cuda(non_blocking=True) + next_states_ = next_states_.cuda(non_blocking=True) + dones_ = dones_.cuda(non_blocking=True) + weights_ = weights_.cuda(non_blocking=True) + + experiences = ( + states_, + actions_, + rewards_, + next_states_, + dones_, + weights_, + indices, + eps_d, + ) + + 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 < self.total_size + + self.sum_tree[idx] = priority ** self.alpha + self.min_tree[idx] = priority ** self.alpha + + self._max_priority = max(self._max_priority, priority) diff --git a/scripts/algorithms/common/buffer/replay_buffer.py b/scripts/algorithms/common/buffer/replay_buffer.py index 273576e..db89ee4 100644 --- a/scripts/algorithms/common/buffer/replay_buffer.py +++ b/scripts/algorithms/common/buffer/replay_buffer.py @@ -1,11 +1,14 @@ # -*- coding: utf-8 -*- """Replay buffer for baselines.""" -from typing import Tuple +from collections import deque +from typing import Any, Deque, List, Tuple import numpy as np import torch +from algorithms.common.helper_functions import get_n_step_info + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") @@ -84,3 +87,88 @@ class ReplayBuffer: def __len__(self) -> int: """Return the current size of internal memory.""" return len(self.buffer) + + +class NStepTransitionBuffer: + """Fixed-size buffer to store experience tuples. + + Attributes: + buffer (list): list of replay buffer + buffer_size (int): buffer size not storing demos + demo_size (int): size of a demo to permanently store in the buffer + cursor (int): position to store next transition coming in + + """ + + def __init__(self, buffer_size: int, n_step: int, gamma: float, demo: list = None): + """Initialize a ReplayBuffer object. + + Args: + buffer_size (int): size of replay buffer for experience + demo (list): demonstration transitions + + """ + assert buffer_size > 0 + + self.n_step_buffer: Deque = deque(maxlen=n_step) + self.buffer_size = buffer_size + self.buffer: list = list() + self.n_step = n_step + self.gamma = gamma + self.demo_size = 0 + self.cursor = 0 + + # if demo exists + if demo: + self.demo_size = len(demo) + self.buffer.extend(demo) + + self.buffer.extend([None] * self.buffer_size) + + def add(self, transition: Tuple[np.ndarray, ...]) -> Tuple[Any, ...]: + """Add a new transition to memory.""" + self.n_step_buffer.append(transition) + + # single step transition is not ready + if len(self.n_step_buffer) < self.n_step: + return () + + # add a multi step transition + reward, next_state, done = get_n_step_info(self.n_step_buffer, self.gamma) + curr_state, action = self.n_step_buffer[0][:2] + new_transition = (curr_state, action, reward, next_state, done) + + # insert the new transition to buffer + idx = self.demo_size + self.cursor + self.buffer[idx] = new_transition + self.cursor = (self.cursor + 1) % self.buffer_size + + # return a single step transition to insert to replay buffer + return self.n_step_buffer[0] + + def sample(self, indices: List[int]) -> Tuple[torch.Tensor, ...]: + """Randomly sample a batch of experiences from memory.""" + states, actions, rewards, next_states, dones = [], [], [], [], [] + + 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)) + + 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) + + if torch.cuda.is_available(): + states_ = states_.cuda(non_blocking=True) + actions_ = actions_.cuda(non_blocking=True) + rewards_ = rewards_.cuda(non_blocking=True) + next_states_ = next_states_.cuda(non_blocking=True) + dones_ = dones_.cuda(non_blocking=True) + + return states_, actions_, rewards_, next_states_, dones_ diff --git a/scripts/algorithms/common/helper_functions.py b/scripts/algorithms/common/helper_functions.py index 51a2c41..2347722 100644 --- a/scripts/algorithms/common/helper_functions.py +++ b/scripts/algorithms/common/helper_functions.py @@ -6,6 +6,8 @@ """ import random +from collections import deque +from typing import Deque, List, Tuple import gym import numpy as np @@ -32,3 +34,46 @@ def set_random_seed(seed: int, env: gym.Env): torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) + + +def get_n_step_info_from_demo( + demo: List, n_step: int, gamma: float +) -> Tuple[List, List]: + """Return 1 step and n step demos.""" + assert demo + assert n_step > 1 + + demos_1_step = list() + demos_n_step = list() + n_step_buffer: Deque = deque(maxlen=n_step) + + for transition in demo: + n_step_buffer.append(transition) + + if len(n_step_buffer) == n_step: + # add a single step transition + demos_1_step.append(n_step_buffer[0]) + + # add a multi step transition + curr_state, action = n_step_buffer[0][:2] + reward, next_state, done = get_n_step_info(n_step_buffer, gamma) + transition = (curr_state, action, reward, next_state, done) + demos_n_step.append(transition) + + return demos_1_step, demos_n_step + + +def get_n_step_info( + n_step_buffer: Deque, gamma: float +) -> Tuple[np.int64, np.ndarray, bool]: + """Return n step reward, next state, and done.""" + # info of the last transition + reward, next_state, done = n_step_buffer[-1][-3:] + + for transition in reversed(list(n_step_buffer)[:-1]): + r, n_s, d = transition[-3:] + + reward = r + gamma * reward * (1 - d) + next_state, done = (n_s, d) if d else (next_state, done) + + return reward, next_state, done diff --git a/scripts/algorithms/ddpg/agent.py b/scripts/algorithms/ddpg/agent.py index a550497..729aeea 100644 --- a/scripts/algorithms/ddpg/agent.py +++ b/scripts/algorithms/ddpg/agent.py @@ -40,6 +40,7 @@ class Agent(AbstractAgent): curr_state (np.ndarray): temporary storage of the current state total_step (int): total step numbers episode_step (int): step number of the current episode + i_episode (int): current episode number """ @@ -72,20 +73,26 @@ class Agent(AbstractAgent): self.noise = noise self.total_step = 0 self.episode_step = 0 + self.i_episode = 0 # 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) + self._initialize() + + def _initialize(self): + """Initialize non-common things.""" if not self.args.test: # replay memory self.memory = ReplayBuffer( - hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"] + self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"] ) def select_action(self, state: np.ndarray) -> np.ndarray: """Select an action from the input space.""" self.curr_state = state + state = self._preprocess_state(state) # if initial random action should be conducted if ( @@ -94,7 +101,6 @@ class Agent(AbstractAgent): ): return self.env.action_space.sample() - state = torch.FloatTensor(state).to(device) selected_action = self.actor(state) if not self.args.test: @@ -103,6 +109,11 @@ class Agent(AbstractAgent): return selected_action.detach().cpu().numpy() + def _preprocess_state(self, state: np.ndarray) -> torch.Tensor: + """Preprocess state so that actor selects an action.""" + state = torch.FloatTensor(state).to(device) + return state + def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool]: """Take an action and return the response of the env.""" self.total_step += 1 @@ -115,10 +126,15 @@ class Agent(AbstractAgent): done_bool = ( False if self.episode_step == self.args.max_episode_steps else done ) - self.memory.add(self.curr_state, action, reward, next_state, done_bool) + transition = (self.curr_state, action, reward, next_state, done_bool) + self._add_transition_to_memory(transition) return next_state, reward, done + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + self.memory.add(*transition) + def update_model( self, experiences: Tuple[ @@ -221,7 +237,7 @@ class Agent(AbstractAgent): 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): + for self.i_episode in range(1, self.args.episode_num + 1): state = self.env.reset() done = False score = 0 @@ -229,7 +245,7 @@ class Agent(AbstractAgent): loss_episode = list() while not done: - if self.args.render and i_episode >= self.args.render_after: + if self.args.render and self.i_episode >= self.args.render_after: self.env.render() action = self.select_action(state) @@ -246,10 +262,10 @@ class Agent(AbstractAgent): # logging if loss_episode: avg_loss = np.vstack(loss_episode).mean(axis=0) - self.write_log(i_episode, avg_loss, score) + self.write_log(self.i_episode, avg_loss, score) - if i_episode % self.args.save_period == 0: - self.save_params(i_episode) + if self.i_episode % self.args.save_period == 0: + self.save_params(self.i_episode) # termination self.env.close() diff --git a/scripts/algorithms/fd/ddpg_agent.py b/scripts/algorithms/fd/ddpg_agent.py new file mode 100644 index 0000000..487e52f --- /dev/null +++ b/scripts/algorithms/fd/ddpg_agent.py @@ -0,0 +1,176 @@ +# -*- coding: utf-8 -*- +"""DDPGfD agent using demo agent 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 + https://arxiv.org/pdf/1707.08817.pdf +""" + +import pickle +from typing import List, Tuple + +import numpy as np +import torch + +import algorithms.common.helper_functions as common_utils +from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD +from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer +from algorithms.ddpg.agent import Agent as DDPGAgent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class Agent(DDPGAgent): + """ActorCritic interacting with environment. + + Attributes: + memory (PrioritizedReplayBufferfD): replay memory + beta (float): beta parameter for prioritized replay buffer + + """ + + # pylint: disable=attribute-defined-outside-init + def _initialize(self): + """Initialize non-common things.""" + self.use_n_step = self.hyper_params["N_STEP"] > 1 + + if not self.args.test: + # load demo replay memory + with open(self.args.demo_path, "rb") as f: + demos = pickle.load(f) + + if self.use_n_step: + demos, demos_n_step = common_utils.get_n_step_info_from_demo( + demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"] + ) + + # replay memory for multi-steps + self.memory_n = NStepTransitionBuffer( + buffer_size=self.hyper_params["BUFFER_SIZE"], + n_step=self.hyper_params["N_STEP"], + gamma=self.hyper_params["GAMMA"], + demo=demos_n_step, + ) + + # replay memory for a single step + self.beta = self.hyper_params["PER_BETA"] + self.memory = PrioritizedReplayBufferfD( + self.hyper_params["BUFFER_SIZE"], + self.hyper_params["BATCH_SIZE"], + demo=list(demos), + alpha=self.hyper_params["PER_ALPHA"], + epsilon_d=self.hyper_params["PER_EPS_DEMO"], + ) + + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + # add n-step transition + if self.use_n_step: + transition = self.memory_n.add(transition) + + # add a single step transition + # if transition is not an empty tuple + if transition: + self.memory.add(*transition) + + def _get_critic_loss( + self, experiences: Tuple[torch.Tensor, ...], gamma: float + ) -> torch.Tensor: + """Return element-wise critic loss.""" + states, actions, rewards, next_states, dones = experiences[:5] + + # G_t = r + gamma * v(s_{t+1}) if state != Terminal + # = r otherwise + masks = 1 - dones + next_actions = self.actor_target(next_states) + next_states_actions = torch.cat((next_states, next_actions), dim=-1) + next_values = self.critic_target(next_states_actions) + curr_returns = rewards + gamma * next_values * masks + curr_returns = curr_returns.to(device).detach() + + # train critic + values = self.critic(torch.cat((states, actions), dim=-1)) + critic_loss_element_wise = (values - curr_returns).pow(2) + + return critic_loss_element_wise + + def update_model( + self, + experiences: Tuple[ + torch.Tensor, + 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.""" + # NOTE This is for old update_model() interface. + # experiences_1 = self.memory.sample(self.beta) + experiences_1 = experiences + states, actions = experiences_1[:2] + weights, indices, eps_d = experiences_1[-3:] + gamma = self.hyper_params["GAMMA"] + + # train critic + critic_loss_element_wise = self._get_critic_loss(experiences_1, gamma) + critic_loss = torch.mean(critic_loss_element_wise * weights) + + if self.use_n_step: + experiences_n = self.memory_n.sample(indices) + gamma = gamma ** self.hyper_params["N_STEP"] + critic_loss_n_element_wise = self._get_critic_loss(experiences_n, gamma) + # to update loss and priorities + lambda1 = self.hyper_params["LAMBDA1"] + critic_loss_element_wise += critic_loss_n_element_wise * lambda1 + critic_loss = torch.mean(critic_loss_element_wise * 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 + new_priorities = critic_loss_element_wise + new_priorities += self.hyper_params["LAMBDA3"] * actor_loss_element_wise.pow(2) + new_priorities += self.hyper_params["PER_EPS"] + new_priorities = new_priorities.data.cpu().numpy().squeeze() + new_priorities += eps_d + self.memory.update_priorities(indices, new_priorities) + + # increase beta + fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) + self.beta = self.beta + fraction * (1.0 - self.beta) + + return actor_loss.data, critic_loss.data + + def pretrain(self): + """Pretraining steps.""" + pretrain_loss = list() + print("[INFO] Pre-Train %d step." % self.hyper_params["PRETRAIN_STEP"]) + for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1): + loss = self.update_model() + pretrain_loss.append(loss) # for logging + + # logging + if i_step == 1 or i_step % 100 == 0: + avg_loss = np.vstack(pretrain_loss).mean(axis=0) + pretrain_loss.clear() + self.write_log(0, avg_loss, 0) diff --git a/scripts/algorithms/fd/sac_agent.py b/scripts/algorithms/fd/sac_agent.py new file mode 100644 index 0000000..cad4ef4 --- /dev/null +++ b/scripts/algorithms/fd/sac_agent.py @@ -0,0 +1,226 @@ +# -*- coding: utf-8 -*- +"""SAC agent from demonstration for episodic tasks in OpenAI Gym. + +- Author: Curt Park +- Contact: curt.park@medipixel.io +- Paper: https://arxiv.org/pdf/1801.01290.pdf + https://arxiv.org/pdf/1812.05905.pdf + https://arxiv.org/pdf/1511.05952.pdf + https://arxiv.org/pdf/1707.08817.pdf +""" + +import pickle +from typing import List, Tuple + +import numpy as np +import torch + +import algorithms.common.helper_functions as common_utils +from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD +from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer +from algorithms.sac.agent import Agent as SACAgent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class Agent(SACAgent): + """SAC agent interacting with environment. + + Attrtibutes: + memory (PrioritizedReplayBufferfD): replay memory + beta (float): beta parameter for prioritized replay buffer + + """ + + # pylint: disable=attribute-defined-outside-init + def _initialize(self): + """Initialize non-common things.""" + self.use_n_step = self.hyper_params["N_STEP"] > 1 + + if not self.args.test: + # load demo replay memory + with open(self.args.demo_path, "rb") as f: + demos = pickle.load(f) + + if self.use_n_step: + demos, demos_n_step = common_utils.get_n_step_info_from_demo( + demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"] + ) + + # replay memory for multi-steps + self.memory_n = NStepTransitionBuffer( + buffer_size=self.hyper_params["BUFFER_SIZE"], + n_step=self.hyper_params["N_STEP"], + gamma=self.hyper_params["GAMMA"], + demo=demos_n_step, + ) + + # replay memory + self.beta = self.hyper_params["PER_BETA"] + self.memory = PrioritizedReplayBufferfD( + self.hyper_params["BUFFER_SIZE"], + self.hyper_params["BATCH_SIZE"], + demo=demos, + alpha=self.hyper_params["PER_ALPHA"], + epsilon_d=self.hyper_params["PER_EPS_DEMO"], + ) + + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + # add n-step transition + if self.use_n_step: + transition = self.memory_n.add(transition) + + # add a single step transition + # if transition is not an empty tuple + if transition: + self.memory.add(*transition) + + # pylint: disable=too-many-statements + def update_model( + self, + experiences: Tuple[ + torch.Tensor, + 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, indices, eps_d = ( + experiences + ) + new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states) + + # train alpha + if self.hyper_params["AUTO_ENTROPY_TUNING"]: + alpha_loss = torch.mean( + (-self.log_alpha * (log_prob + self.target_entropy).detach()) * weights + ) + + self.alpha_optimizer.zero_grad() + alpha_loss.backward() + self.alpha_optimizer.step() + + alpha = self.log_alpha.exp() + else: + alpha_loss = torch.zeros(1) + alpha = self.hyper_params["W_ENTROPY"] + + # Q function loss + masks = 1 - dones + gamma = self.hyper_params["GAMMA"] + q_1_pred = self.qf_1(states, actions) + q_2_pred = self.qf_2(states, actions) + v_target = self.vf_target(next_states) + q_target = rewards + self.hyper_params["GAMMA"] * v_target * masks + qf_1_loss = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights) + qf_2_loss = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights) + + if self.use_n_step: + experiences_n = self.memory_n.sample(indices) + _, _, rewards, next_states, dones = experiences_n + gamma = gamma ** self.hyper_params["N_STEP"] + lambda1 = self.hyper_params["LAMBDA1"] + masks = 1 - dones + + v_target = self.vf_target(next_states) + q_target = rewards + gamma * v_target * masks + qf_1_loss_n = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights) + qf_2_loss_n = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights) + + # to update loss and priorities + qf_1_loss = qf_1_loss + qf_1_loss_n * lambda1 + qf_2_loss = qf_2_loss + qf_2_loss_n * lambda1 + + # V function loss + v_pred = self.vf(states) + q_pred = torch.min( + self.qf_1(states, new_actions), self.qf_2(states, new_actions) + ) + v_target = (q_pred - alpha * log_prob).detach() + vf_loss_element_wise = (v_pred - v_target).pow(2) + vf_loss = torch.mean(vf_loss_element_wise * weights) + + # train Q functions + self.qf_1_optimizer.zero_grad() + qf_1_loss.backward() + self.qf_1_optimizer.step() + + self.qf_2_optimizer.zero_grad() + qf_2_loss.backward() + self.qf_2_optimizer.step() + + # train V function + self.vf_optimizer.zero_grad() + vf_loss.backward() + self.vf_optimizer.step() + + if self.total_step % self.hyper_params["DELAYED_UPDATE"] == 0: + # actor loss + advantage = q_pred - v_pred.detach() + actor_loss_element_wise = alpha * log_prob - advantage + actor_loss = torch.mean(actor_loss_element_wise * weights) + + # regularization + mean_reg = self.hyper_params["W_MEAN_REG"] * mu.pow(2).mean() + std_reg = self.hyper_params["W_STD_REG"] * std.pow(2).mean() + pre_activation_reg = self.hyper_params["W_PRE_ACTIVATION_REG"] * ( + pre_tanh_value.pow(2).sum(dim=-1).mean() + ) + actor_reg = mean_reg + std_reg + pre_activation_reg + + # actor loss + regularization + actor_loss += actor_reg + + # train actor + self.actor_optimizer.zero_grad() + actor_loss.backward() + self.actor_optimizer.step() + + # update target networks + common_utils.soft_update(self.vf, self.vf_target, self.hyper_params["TAU"]) + + # update priorities + new_priorities = vf_loss_element_wise + new_priorities += self.hyper_params[ + "LAMBDA3" + ] * actor_loss_element_wise.pow(2) + new_priorities += self.hyper_params["PER_EPS"] + new_priorities = new_priorities.data.cpu().numpy().squeeze() + new_priorities += eps_d + self.memory.update_priorities(indices, new_priorities) + + # increase beta + fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) + self.beta = self.beta + fraction * (1.0 - self.beta) + else: + actor_loss = torch.zeros(1) + + return ( + actor_loss.data, + qf_1_loss.data, + qf_2_loss.data, + vf_loss.data, + alpha_loss.data, + ) + + def pretrain(self): + """Pretraining steps.""" + pretrain_loss = list() + print("[INFO] Pre-Train %d steps." % self.hyper_params["PRETRAIN_STEP"]) + for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1): + loss = self.update_model() + pretrain_loss.append(loss) # for logging + + # logging + if i_step == 1 or i_step % 100 == 0: + avg_loss = np.vstack(pretrain_loss).mean(axis=0) + pretrain_loss.clear() + self.write_log( + 0, avg_loss, 0, delayed_update=self.hyper_params["DELAYED_UPDATE"] + ) diff --git a/scripts/algorithms/fd/td3_agent.py b/scripts/algorithms/fd/td3_agent.py new file mode 100644 index 0000000..63069a9 --- /dev/null +++ b/scripts/algorithms/fd/td3_agent.py @@ -0,0 +1,209 @@ +# -*- coding: utf-8 -*- +"""TD3 agent from demonstration for episodic tasks in OpenAI Gym. + +- Author: Seungjae Ryan Lee +- Contact: seungjaeryanlee@gmail.com +- Paper: https://arxiv.org/pdf/1802.09477.pdf (TD3) + https://arxiv.org/pdf/1511.05952.pdf (PER) + https://arxiv.org/pdf/1707.08817.pdf (DDPGfD) +""" + +import pickle +from typing import List, Tuple + +import numpy as np +import torch + +import algorithms.common.helper_functions as common_utils +from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD +from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer +from algorithms.td3.agent import Agent as TD3Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class Agent(TD3Agent): + """TD3 agent interacting with environment. + + Attrtibutes: + memory (PrioritizedReplayBufferfD): replay memory + beta (float): beta parameter for prioritized replay buffer + + """ + + # pylint: disable=attribute-defined-outside-init + def _initialize(self): + """Initialize non-common things.""" + self.use_n_step = self.hyper_params["N_STEP"] > 1 + + if not self.args.test: + # load demo replay memory + with open(self.args.demo_path, "rb") as f: + demos = pickle.load(f) + + if self.use_n_step: + demos, demos_n_step = common_utils.get_n_step_info_from_demo( + demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"] + ) + + # replay memory for multi-steps + self.memory_n = NStepTransitionBuffer( + buffer_size=self.hyper_params["BUFFER_SIZE"], + n_step=self.hyper_params["N_STEP"], + gamma=self.hyper_params["GAMMA"], + demo=demos_n_step, + ) + + # replay memory + self.beta = self.hyper_params["PER_BETA"] + self.memory = PrioritizedReplayBufferfD( + self.hyper_params["BUFFER_SIZE"], + self.hyper_params["BATCH_SIZE"], + demo=demos, + alpha=self.hyper_params["PER_ALPHA"], + epsilon_d=self.hyper_params["PER_EPS_DEMO"], + ) + + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + # add n-step transition + if self.use_n_step: + transition = self.memory_n.add(transition) + + # add a single step transition + # if transition is not an empty tuple + if transition: + self.memory.add(*transition) + + def _get_critic_loss( + self, experiences: Tuple[torch.Tensor, ...], gamma: float + ) -> torch.Tensor: + """Return element-wise critic loss.""" + states, actions, rewards, next_states, dones = experiences[:5] + + # G_t = r + gamma * v(s_{t+1}) if state != Terminal + # = r otherwise + masks = 1 - dones + noise = torch.FloatTensor(self.target_policy_noise.sample()).to(device) + clipped_noise = torch.clamp( + noise, + -self.hyper_params["TARGET_POLICY_NOISE_CLIP"], + self.hyper_params["TARGET_POLICY_NOISE_CLIP"], + ) + next_actions = (self.actor_target(next_states) + clipped_noise).clamp(-1.0, 1.0) + + target_values1 = self.critic1_target( + torch.cat((next_states, next_actions), dim=-1) + ) + target_values2 = self.critic2_target( + torch.cat((next_states, next_actions), dim=-1) + ) + target_values = torch.min(target_values1, target_values2) + target_values = ( + rewards + (self.hyper_params["GAMMA"] * target_values * masks).detach() + ) + + # train critic + values1 = self.critic1(torch.cat((states, actions), dim=-1)) + critic1_loss_element_wise = (values1 - target_values.detach()).pow(2) + + values2 = self.critic2(torch.cat((states, actions), dim=-1)) + critic2_loss_element_wise = (values2 - target_values.detach()).pow(2) + + return critic1_loss_element_wise, critic2_loss_element_wise + + # pylint: disable=too-many-statements + def update_model( + self, + experiences: Tuple[ + torch.Tensor, + 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, indices, eps_d = ( + experiences + ) + + gamma = self.hyper_params["GAMMA"] + critic1_loss_element_wise, critic2_loss_element_wise = self._get_critic_loss( + experiences, gamma + ) + critic_loss_element_wise = critic1_loss_element_wise + critic2_loss_element_wise + critic1_loss = torch.mean(critic1_loss_element_wise * weights) + critic2_loss = torch.mean(critic2_loss_element_wise * weights) + critic_loss = critic1_loss + critic2_loss + + if self.use_n_step: + experiences_n = self.memory_n.sample(indices) + gamma = self.hyper_params["GAMMA"] ** self.hyper_params["N_STEP"] + critic1_loss_n_element_wise, critic2_loss_n_element_wise = self._get_critic_loss( + experiences_n, gamma + ) + critic_loss_n_element_wise = ( + critic1_loss_n_element_wise + critic2_loss_n_element_wise + ) + critic1_loss_n = torch.mean(critic1_loss_n_element_wise * weights) + critic2_loss_n = torch.mean(critic2_loss_n_element_wise * weights) + critic_loss_n = critic1_loss_n + critic2_loss_n + + lambda1 = self.hyper_params["LAMBDA1"] + critic_loss_element_wise += lambda1 * critic_loss_n_element_wise + critic_loss += lambda1 * critic_loss_n + + self.critic_optim.zero_grad() + critic_loss.backward() + self.critic_optim.step() + + if self.episode_steps % self.hyper_params["POLICY_UPDATE_FREQ"] == 0: + # train actor + actions = self.actor(states) + actor_loss_element_wise = -self.critic1( + torch.cat((states, actions), dim=-1) + ) + actor_loss = torch.mean(actor_loss_element_wise * weights) + self.actor_optim.zero_grad() + actor_loss.backward() + self.actor_optim.step() + + # update target networks + tau = self.hyper_params["TAU"] + common_utils.soft_update(self.actor, self.actor_target, tau) + common_utils.soft_update(self.critic1, self.critic1_target, tau) + common_utils.soft_update(self.critic2, self.critic2_target, tau) + + # update priorities + new_priorities = critic_loss_element_wise + new_priorities += self.hyper_params[ + "LAMBDA3" + ] * actor_loss_element_wise.pow(2) + new_priorities += self.hyper_params["PER_EPS"] + new_priorities = new_priorities.data.cpu().numpy().squeeze() + new_priorities += eps_d + self.memory.update_priorities(indices, new_priorities) + else: + actor_loss = torch.zeros(1) + + return actor_loss.data, critic1_loss.data, critic2_loss.data + + def pretrain(self): + """Pretraining steps.""" + pretrain_loss = list() + print("[INFO] Pre-Train %d steps." % self.hyper_params["PRETRAIN_STEP"]) + for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1): + loss = self.update_model() + pretrain_loss.append(loss) # for logging + + # logging + if i_step == 1 or i_step % 100 == 0: + avg_loss = np.vstack(pretrain_loss).mean(axis=0) + pretrain_loss.clear() + self.write_log( + 0, avg_loss, 0, delayed_update=self.hyper_params["DELAYED_UPDATE"] + ) diff --git a/scripts/algorithms/sac/agent.py b/scripts/algorithms/sac/agent.py index 5a2f800..86a0bb0 100644 --- a/scripts/algorithms/sac/agent.py +++ b/scripts/algorithms/sac/agent.py @@ -46,6 +46,7 @@ class Agent(AbstractAgent): hyper_params (dict): hyper-parameters total_step (int): total step numbers episode_step (int): step number of the current episode + i_episode (int): current episode number """ @@ -78,6 +79,7 @@ class Agent(AbstractAgent): self.curr_state = np.zeros((1,)) self.total_step = 0 self.episode_step = 0 + self.i_episode = 0 # automatic entropy tuning if self.hyper_params["AUTO_ENTROPY_TUNING"]: @@ -91,15 +93,20 @@ class Agent(AbstractAgent): if args.load_from is not None and os.path.exists(args.load_from): self.load_params(args.load_from) + self._initialize() + + def _initialize(self): + """Initialize non-common things.""" if not self.args.test: # replay memory self.memory = ReplayBuffer( - hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"] + self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"] ) def select_action(self, state: np.ndarray) -> np.ndarray: """Select an action from the input space.""" self.curr_state = state + state = self._preprocess_state(state) # if initial random action should be conducted if ( @@ -108,7 +115,6 @@ class Agent(AbstractAgent): ): return self.env.action_space.sample() - state = torch.FloatTensor(state).to(device) if self.args.test: _, _, _, selected_action, _ = self.actor(state) else: @@ -116,6 +122,11 @@ class Agent(AbstractAgent): return selected_action.detach().cpu().numpy() + def _preprocess_state(self, state: np.ndarray) -> torch.Tensor: + """Preprocess state so that actor selects an action.""" + state = torch.FloatTensor(state).to(device) + return state + def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool]: """Take an action and return the response of the env.""" self.total_step += 1 @@ -128,10 +139,15 @@ class Agent(AbstractAgent): done_bool = ( False if self.episode_step == self.args.max_episode_steps else done ) - self.memory.add(self.curr_state, action, reward, next_state, done_bool) + transition = (self.curr_state, action, reward, next_state, done_bool) + self._add_transition_to_memory(transition) return next_state, reward, done + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + self.memory.add(*transition) + def update_model( self, experiences: Tuple[ @@ -308,7 +324,7 @@ class Agent(AbstractAgent): wandb.config.update(self.hyper_params) wandb.watch([self.actor, self.vf, self.qf_1, self.qf_2], log="parameters") - for i_episode in range(1, self.args.episode_num + 1): + for self.i_episode in range(1, self.args.episode_num + 1): state = self.env.reset() done = False score = 0 @@ -316,7 +332,7 @@ class Agent(AbstractAgent): loss_episode = list() while not done: - if self.args.render and i_episode >= self.args.render_after: + if self.args.render and self.i_episode >= self.args.render_after: self.env.render() action = self.select_action(state) @@ -335,11 +351,11 @@ class Agent(AbstractAgent): if loss_episode: avg_loss = np.vstack(loss_episode).mean(axis=0) self.write_log( - i_episode, avg_loss, score, self.hyper_params["DELAYED_UPDATE"] + self.i_episode, avg_loss, score, self.hyper_params["DELAYED_UPDATE"] ) - if i_episode % self.args.save_period == 0: - self.save_params(i_episode) + if self.i_episode % self.args.save_period == 0: + self.save_params(self.i_episode) # termination self.env.close() diff --git a/scripts/algorithms/td3/agent.py b/scripts/algorithms/td3/agent.py index 21e9c03..882e1e4 100644 --- a/scripts/algorithms/td3/agent.py +++ b/scripts/algorithms/td3/agent.py @@ -65,8 +65,9 @@ class Agent(AbstractAgent): """ AbstractAgent.__init__(self, env, args) - self.actor, self.actor_target, self.critic1, self.critic1_target, \ - self.critic2, self.critic2_target = models + self.actor, self.actor_target, self.critic1, self.critic1_target, self.critic2, self.critic2_target = ( # noqa: B950 + models + ) self.actor_optim, self.critic_optim = optims self.hyper_params = hyper_params self.exploration_noise, self.target_policy_noise = noises @@ -78,10 +79,15 @@ class Agent(AbstractAgent): if args.load_from is not None and os.path.exists(args.load_from): self.load_params(args.load_from) - # replay memory - self.memory = ReplayBuffer( - hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"], - ) + self._initialize() + + def _initialize(self): + """Initialize non-common things.""" + if not self.args.test: + # replay memory + self.memory = ReplayBuffer( + self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"] + ) def select_action(self, state: np.ndarray) -> np.ndarray: """Select an action from the input space.""" @@ -104,12 +110,25 @@ class Agent(AbstractAgent): def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]: """Take an action and return the response of the env.""" + self.total_steps += 1 + self.episode_steps += 1 + next_state, reward, done, _ = self.env.step(action) - self.memory.add(self.curr_state, action, reward, next_state, done) + if not self.args.test: + # if the last state is not a terminal state, store done as false + done_bool = ( + False if self.episode_steps == self.args.max_episode_steps else done + ) + transition = (self.curr_state, action, reward, next_state, done_bool) + self._add_transition_to_memory(transition) return next_state, reward, done + def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): + """Add 1 step and n step transitions to memory.""" + self.memory.add(*transition) + def update_model( self, experiences: Tuple[ @@ -137,7 +156,9 @@ class Agent(AbstractAgent): torch.cat((next_states, next_actions), dim=-1) ) target_values = torch.min(target_values1, target_values2) - target_values = rewards + (self.hyper_params["GAMMA"] * target_values * masks).detach() + target_values = ( + rewards + (self.hyper_params["GAMMA"] * target_values * masks).detach() + ) # train critic values1 = self.critic1(torch.cat((states, actions), dim=-1)) @@ -206,8 +227,8 @@ class Agent(AbstractAgent): print( "[INFO] total_steps: %d episode: %d total score: %d, total loss: %f\n" - "actor_loss: %.3f critic_loss: %.3f\n" - % (self.total_steps, i, score, total_loss, loss[0], loss[1]) + "actor_loss: %.3f critic1_loss: %.3f critic2_loss: %.3f\n" + % (self.total_steps, i, score, total_loss, loss[0], loss[1], loss[2]) ) if self.args.log: @@ -243,8 +264,6 @@ class Agent(AbstractAgent): action = self.select_action(state) next_state, reward, done = self.step(action) - self.total_steps += 1 - self.episode_steps += 1 if len(self.memory) >= self.hyper_params["BATCH_SIZE"]: experiences = self.memory.sample() diff --git a/scripts/data/lunarlander_continuous_demo.pkl b/scripts/data/lunarlander_continuous_demo.pkl new file mode 100644 index 0000000..fe52886 Binary files /dev/null and b/scripts/data/lunarlander_continuous_demo.pkl differ diff --git a/scripts/examples/lunarlander_continuous_v2/ddpgfd.py b/scripts/examples/lunarlander_continuous_v2/ddpgfd.py new file mode 100644 index 0000000..c4f6c4b --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/ddpgfd.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- +"""Run module for DDPGfD 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.fd.ddpg_agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +hyper_params = { + "N_STEP": 1, + "GAMMA": 0.99, + "TAU": 5e-3, + "BUFFER_SIZE": int(1e5), + "BATCH_SIZE": 64, + "LR_ACTOR": 3e-4, + "LR_CRITIC": 3e-4, + "OU_NOISE_THETA": 0.0, + "OU_NOISE_SIGMA": 0.0, + "PRETRAIN_STEP": 0, + "MULTIPLE_LEARN": 2, # multiple learning updates + "LAMBDA1": 1.0, # N-step return weight + "LAMBDA2": 1e-5, # l2 regularization weight + "LAMBDA3": 1.0, # actor loss contribution of prior weight + "PER_ALPHA": 0.3, + "PER_BETA": 1.0, + "PER_EPS": 1e-6, + "PER_EPS_DEMO": 1.0, + "INITIAL_RANDOM_ACTION": int(5e3), +} + + +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["LAMBDA2"], + ) + + critic_optim = optim.Adam( + critic.parameters(), + lr=hyper_params["LR_CRITIC"], + weight_decay=hyper_params["LAMBDA2"], + ) + + # 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/examples/lunarlander_continuous_v2/sacfd.py b/scripts/examples/lunarlander_continuous_v2/sacfd.py new file mode 100644 index 0000000..82fdbe4 --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/sacfd.py @@ -0,0 +1,121 @@ +# -*- coding: utf-8 -*- +"""Run module for SACfD on LunarLanderContinuous-v2. + +- Author: Curt Park +- Contact: curt.park@medipixel.io +""" + +import argparse + +import gym +import numpy as np +import torch +import torch.optim as optim + +from algorithms.common.networks.mlp import MLP, FlattenMLP, TanhGaussianDistParams +from algorithms.fd.sac_agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +hyper_params = { + "N_STEP": 3, + "GAMMA": 0.99, + "TAU": 1e-3, + "BUFFER_SIZE": int(1e5), + "BATCH_SIZE": 64, + "AUTO_ENTROPY_TUNING": True, + "LR_ACTOR": 3e-4, + "LR_VF": 3e-4, + "LR_QF1": 3e-4, + "LR_QF2": 3e-4, + "LR_ENTROPY": 3e-4, + "W_ENTROPY": 1e-3, + "W_MEAN_REG": 1e-3, + "W_STD_REG": 1e-3, + "W_PRE_ACTIVATION_REG": 0.0, + "DELAYED_UPDATE": 2, + "PRETRAIN_STEP": 100, + "MULTIPLE_LEARN": 2, # multiple learning updates + "LAMBDA1": 1.0, # N-step return weight + "LAMBDA2": 1e-5, # l2 regularization weight + "LAMBDA3": 1.0, # actor loss contribution of prior weight + "PER_ALPHA": 0.6, + "PER_BETA": 0.4, + "PER_EPS": 1e-6, + "PER_EPS_DEMO": 1.0, + "INITIAL_RANDOM_ACTION": int(5e3), +} + + +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_vf = [256, 256] + hidden_sizes_qf = [256, 256] + + # target entropy + target_entropy = -np.prod((action_dim,)).item() # heuristic + + # create actor + actor = TanhGaussianDistParams( + input_size=state_dim, output_size=action_dim, hidden_sizes=hidden_sizes_actor + ).to(device) + + # create v_critic + vf = MLP(input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf).to( + device + ) + vf_target = MLP( + input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf + ).to(device) + vf_target.load_state_dict(vf.state_dict()) + + # create q_critic + qf_1 = FlattenMLP( + input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf + ).to(device) + qf_2 = FlattenMLP( + input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf + ).to(device) + + # create optimizers + actor_optim = optim.Adam( + actor.parameters(), + lr=hyper_params["LR_ACTOR"], + weight_decay=hyper_params["LAMBDA2"], + ) + vf_optim = optim.Adam( + vf.parameters(), lr=hyper_params["LR_VF"], weight_decay=hyper_params["LAMBDA2"] + ) + qf_1_optim = optim.Adam( + qf_1.parameters(), + lr=hyper_params["LR_QF1"], + weight_decay=hyper_params["LAMBDA2"], + ) + qf_2_optim = optim.Adam( + qf_2.parameters(), + lr=hyper_params["LR_QF2"], + weight_decay=hyper_params["LAMBDA2"], + ) + + # make tuples to create an agent + models = (actor, vf, vf_target, qf_1, qf_2) + optims = (actor_optim, vf_optim, qf_1_optim, qf_2_optim) + + # create an agent + agent = Agent(env, args, hyper_params, models, optims, target_entropy) + + # run + if args.test: + agent.test() + else: + agent.train() diff --git a/scripts/examples/lunarlander_continuous_v2/td3.py b/scripts/examples/lunarlander_continuous_v2/td3.py index ec2f969..a7f8200 100644 --- a/scripts/examples/lunarlander_continuous_v2/td3.py +++ b/scripts/examples/lunarlander_continuous_v2/td3.py @@ -65,25 +65,29 @@ def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int) # create critic1 critic1 = MLP( - input_size=state_dim + action_dim, output_size=1, - hidden_sizes=hidden_sizes_critic + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic1_target = MLP( - input_size=state_dim + action_dim, output_size=1, - hidden_sizes=hidden_sizes_critic + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic1_target.load_state_dict(critic1.state_dict()) # create critic2 critic2 = MLP( - input_size=state_dim + action_dim, output_size=1, - hidden_sizes=hidden_sizes_critic + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic2_target = MLP( - input_size=state_dim + action_dim, output_size=1, - hidden_sizes=hidden_sizes_critic + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, ).to(device) critic2_target.load_state_dict(critic2.state_dict()) diff --git a/scripts/examples/lunarlander_continuous_v2/td3fd.py b/scripts/examples/lunarlander_continuous_v2/td3fd.py new file mode 100644 index 0000000..c24cae2 --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/td3fd.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- +"""Run module for SACfD on LunarLanderContinuous-v2. + +- Author: Seungjae Ryan Lee +- Contact: seungjaeryanlee@gmail.com +""" + +import argparse + +import gym +import torch +import torch.optim as optim + +from algorithms.common.networks.mlp import MLP +from algorithms.common.noise import GaussianNoise +from algorithms.fd.td3_agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +# TODO Tune hyperparameters on LunarLander-v2 + +# hyper parameters +hyper_params = { + "N_STEP": 3, + "GAMMA": 0.99, + "TAU": 1e-3, + "BUFFER_SIZE": int(1e5), + "BATCH_SIZE": 64, + "LR_ACTOR": 3e-4, + "LR_CRITIC": 3e-4, + "EXPLORATION_NOISE": 0.1, + "TARGET_POLICY_NOISE": 0.2, + "TARGET_POLICY_NOISE_CLIP": 0.5, + "POLICY_UPDATE_FREQ": 2, + "INITIAL_RANDOM_ACTIONS": 5e3, + "PRETRAIN_STEP": 100, + "MULTIPLE_LEARN": 2, # multiple learning updates + "LAMBDA1": 1.0, # N-step return weight + "LAMBDA2": 1e-5, # l2 regularization weight + "LAMBDA3": 1.0, # actor loss contribution of prior weight + "PER_ALPHA": 0.3, + "PER_BETA": 1.0, + "PER_EPS": 1e-6, + "PER_EPS_DEMO": 1.0, +} + + +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 = [400, 300] + hidden_sizes_critic = [400, 300] + + # 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 critic1 + critic1 = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + + critic1_target = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + critic1_target.load_state_dict(critic1.state_dict()) + + # create critic2 + critic2 = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + + critic2_target = MLP( + input_size=state_dim + action_dim, + output_size=1, + hidden_sizes=hidden_sizes_critic, + ).to(device) + critic2_target.load_state_dict(critic2.state_dict()) + + # concat critic parameters to use one optim + critic_parameters = list(critic1.parameters()) + list(critic2.parameters()) + + # create optimizer + actor_optim = optim.Adam( + actor.parameters(), + lr=hyper_params["LR_ACTOR"], + weight_decay=hyper_params["LAMBDA2"], + ) + + critic_optim = optim.Adam( + critic_parameters, + lr=hyper_params["LR_CRITIC"], + weight_decay=hyper_params["LAMBDA2"], + ) + + # noise + exploration_noise = GaussianNoise( + action_dim, + min_sigma=hyper_params["EXPLORATION_NOISE"], + max_sigma=hyper_params["EXPLORATION_NOISE"], + ) + + target_policy_noise = GaussianNoise( + action_dim, + min_sigma=hyper_params["TARGET_POLICY_NOISE"], + max_sigma=hyper_params["TARGET_POLICY_NOISE"], + ) + + # make tuples to create an agent + models = (actor, actor_target, critic1, critic1_target, critic2, critic2_target) + optims = (actor_optim, critic_optim) + noises = (exploration_noise, target_policy_noise) + + # create an agent + agent = Agent(env, args, hyper_params, models, optims, noises) + + # run + if args.test: + agent.test() + else: + agent.train() diff --git a/scripts/run_lunarlander_continuous.py b/scripts/run_lunarlander_continuous.py index a861eef..55987a2 100644 --- a/scripts/run_lunarlander_continuous.py +++ b/scripts/run_lunarlander_continuous.py @@ -40,6 +40,12 @@ parser.add_argument( parser.add_argument("--save-period", type=int, default=100, help="save model period") parser.add_argument("--log", action="store_true", help="turn on logging") parser.add_argument("--test", action="store_true", help="test mode (no training)") +parser.add_argument( + "--demo-path", + type=str, + default="data/lunarlander_continuous_demo.pkl", + help="demonstration path", +) parser.set_defaults(render=True) args = parser.parse_args() diff --git a/scripts/test_lstm.py b/scripts/test_lstm.py index ecbe6d8..311889a 100644 --- a/scripts/test_lstm.py +++ b/scripts/test_lstm.py @@ -1,13 +1,12 @@ # -*- coding: utf-8 -*- """Simple test case(sine+noise->cos) for LSTM network.""" -import numpy as np import matplotlib.pyplot as plt - +import numpy as np import torch -from torch.utils.data import Dataset, DataLoader -import torch.optim as optim import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset from algorithms.common.networks.lstm import LSTM