diff --git a/scripts/algorithms/common/buffer/replay_buffer.py b/scripts/algorithms/common/buffer/replay_buffer.py index 8cd16fd..273576e 100644 --- a/scripts/algorithms/common/buffer/replay_buffer.py +++ b/scripts/algorithms/common/buffer/replay_buffer.py @@ -28,7 +28,6 @@ class ReplayBuffer: 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 list """ self.buffer: list = list() diff --git a/scripts/algorithms/common/helper_functions.py b/scripts/algorithms/common/helper_functions.py index 9403580..51a2c41 100644 --- a/scripts/algorithms/common/helper_functions.py +++ b/scripts/algorithms/common/helper_functions.py @@ -12,6 +12,8 @@ import numpy as np import torch import torch.nn as nn +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + def identity(x: torch.Tensor) -> torch.Tensor: """Return input without any change.""" diff --git a/scripts/algorithms/common/networks/mlp.py b/scripts/algorithms/common/networks/mlp.py index 449daad..86ae87c 100644 --- a/scripts/algorithms/common/networks/mlp.py +++ b/scripts/algorithms/common/networks/mlp.py @@ -93,12 +93,21 @@ class MLP(nn.Module): return output +class FlattenMLP(MLP): + """Baseline of Multilayer perceptron for Flatten input.""" + + def forward(self, *args: torch.Tensor) -> torch.Tensor: + """Forward method implementation.""" + states, actions = args + flat_inputs = torch.cat((states, actions), dim=-1) + return super(FlattenMLP, self).forward(flat_inputs) + + class GaussianDist(MLP): """Multilayer perceptron with Gaussian distribution output. Attributes: mu_activation (function): bounding function for mean - log_std_clamping (bool): whether or not to clamp log std log_std_min (float): lower bound of log std log_std_max (float): upper bound of log std mu_layer (nn.Linear): output layer for mean @@ -116,9 +125,7 @@ class GaussianDist(MLP): log_std_max: float = 2, init_w: float = 3e-3, ): - """Initialization. - - """ + """Initialization.""" super(GaussianDist, self).__init__( input_size=input_size, output_size=output_size, @@ -150,8 +157,9 @@ class GaussianDist(MLP): mu = self.mu_activation(self.mu_layer(hidden)) # get std - log_std = torch.clamp( - self.log_std_layer(hidden), self.log_std_min, self.log_std_max + log_std = torch.tanh(self.log_std_layer(hidden)) + log_std = self.log_std_min + 0.5 * (self.log_std_max - self.log_std_min) * ( + log_std + 1 ) std = torch.exp(log_std) @@ -168,16 +176,6 @@ class GaussianDist(MLP): return action, dist -class GaussianDistParams(GaussianDist): - """Multilayer perceptron with Gaussian distribution params output.""" - - def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]: - """Forward method implementation.""" - mu, log_std, std = super(GaussianDistParams, self).get_dist_params(x) - - return mu, log_std, std - - class TanhGaussianDistParams(GaussianDist): """Multilayer perceptron with Gaussian distribution output.""" diff --git a/scripts/algorithms/sac/agent.py b/scripts/algorithms/sac/agent.py new file mode 100644 index 0000000..5a2f800 --- /dev/null +++ b/scripts/algorithms/sac/agent.py @@ -0,0 +1,345 @@ +# -*- coding: utf-8 -*- +"""SAC agent 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 +""" + +import argparse +import os +from typing import Tuple + +import gym +import numpy as np +import torch +import torch.nn.functional as F +import torch.optim as optim +import wandb + +import algorithms.common.helper_functions as common_utils +from algorithms.common.abstract.agent import AbstractAgent +from algorithms.common.buffer.replay_buffer import ReplayBuffer + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +class Agent(AbstractAgent): + """SAC agent interacting with environment. + + Attrtibutes: + memory (ReplayBuffer): replay memory + actor (nn.Module): actor model to select actions + actor_target (nn.Module): target actor model to select actions + actor_optimizer (Optimizer): optimizer for training actor + critic_1 (nn.Module): critic model to predict state values + critic_2 (nn.Module): critic model to predict state values + critic_target1 (nn.Module): target critic model to predict state values + critic_target2 (nn.Module): target critic model to predict state values + critic_optimizer1 (Optimizer): optimizer for training critic_1 + critic_optimizer2 (Optimizer): optimizer for training critic_2 + curr_state (np.ndarray): temporary storage of the current state + target_entropy (int): desired entropy used for the inequality constraint + alpha (torch.Tensor): weight for entropy + alpha_optimizer (Optimizer): optimizer for alpha + hyper_params (dict): hyper-parameters + total_step (int): total step numbers + episode_step (int): step number of the current episode + + """ + + def __init__( + self, + env: gym.Env, + args: argparse.Namespace, + hyper_params: dict, + models: tuple, + optims: tuple, + target_entropy: float, + ): + """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 + target_entropy (float): target entropy for the inequality constraint + + """ + AbstractAgent.__init__(self, env, args) + + self.actor, self.vf, self.vf_target, self.qf_1, self.qf_2 = models + self.actor_optimizer, self.vf_optimizer = optims[0:2] + self.qf_1_optimizer, self.qf_2_optimizer = optims[2:4] + self.hyper_params = hyper_params + self.curr_state = np.zeros((1,)) + self.total_step = 0 + self.episode_step = 0 + + # automatic entropy tuning + if self.hyper_params["AUTO_ENTROPY_TUNING"]: + self.target_entropy = target_entropy + self.log_alpha = torch.zeros(1, requires_grad=True, device=device) + self.alpha_optimizer = optim.Adam( + [self.log_alpha], lr=self.hyper_params["LR_ENTROPY"] + ) + + # 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) + + if not self.args.test: + # replay memory + self.memory = ReplayBuffer( + hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"] + ) + + def select_action(self, state: np.ndarray) -> np.ndarray: + """Select an action from the input space.""" + self.curr_state = state + + # if initial random action should be conducted + if ( + self.total_step < self.hyper_params["INITIAL_RANDOM_ACTION"] + and not self.args.test + ): + return self.env.action_space.sample() + + state = torch.FloatTensor(state).to(device) + if self.args.test: + _, _, _, selected_action, _ = self.actor(state) + else: + selected_action, _, _, _, _ = self.actor(state) + + return selected_action.detach().cpu().numpy() + + 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 + self.episode_step += 1 + + next_state, reward, done, _ = self.env.step(action) + + if not self.args.test: + # if the last state is not a terminal state, store done as false + 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) + + return next_state, reward, done + + def update_model( + self, + experiences: Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Train the model after each episode.""" + states, actions, rewards, next_states, dones = experiences + new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states) + + # train alpha + if self.hyper_params["AUTO_ENTROPY_TUNING"]: + alpha_loss = ( + -self.log_alpha * (log_prob + self.target_entropy).detach() + ).mean() + + 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 + 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 = F.mse_loss(q_1_pred, q_target.detach()) + qf_2_loss = F.mse_loss(q_2_pred, q_target.detach()) + + # 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 + vf_loss = F.mse_loss(v_pred, v_target.detach()) + + # 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 = (alpha * log_prob - advantage).mean() + + # 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"]) + 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 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"]) + self.qf_1.load_state_dict(params["qf_1"]) + self.qf_2.load_state_dict(params["qf_2"]) + self.vf.load_state_dict(params["vf"]) + self.vf_target.load_state_dict(params["vf_target"]) + self.actor_optimizer.load_state_dict(params["actor_optim"]) + self.qf_1_optimizer.load_state_dict(params["qf_1_optim"]) + self.qf_2_optimizer.load_state_dict(params["qf_2_optim"]) + self.vf_optimizer.load_state_dict(params["vf_optim"]) + + if self.hyper_params["AUTO_ENTROPY_TUNING"]: + self.alpha_optimizer.load_state_dict(params["alpha_optim"]) + + print("[INFO] loaded the model and optimizer from", path) + + def save_params(self, n_episode: int): + """Save model and optimizer parameters.""" + params = { + "actor": self.actor.state_dict(), + "qf_1": self.qf_1.state_dict(), + "qf_2": self.qf_2.state_dict(), + "vf": self.vf.state_dict(), + "vf_target": self.vf_target.state_dict(), + "actor_optim": self.actor_optimizer.state_dict(), + "qf_1_optim": self.qf_1_optimizer.state_dict(), + "qf_2_optim": self.qf_2_optimizer.state_dict(), + "vf_optim": self.vf_optimizer.state_dict(), + } + + if self.hyper_params["AUTO_ENTROPY_TUNING"]: + params["alpha_optim"] = self.alpha_optimizer.state_dict() + + AbstractAgent.save_params(self, params, n_episode) + + def write_log( + self, i: int, loss: np.ndarray, score: float = 0.0, delayed_update: int = 1 + ): + """Write log about loss and score""" + total_loss = loss.sum() + + print( + "[INFO] episode %d, episode_step %d, total step %d, total score: %d\n" + "total loss: %.3f actor_loss: %.3f qf_1_loss: %.3f qf_2_loss: %.3f " + "vf_loss: %.3f alpha_loss: %.3f\n" + % ( + i, + self.episode_step, + self.total_step, + score, + total_loss, + loss[0] * delayed_update, # actor loss + loss[1], # qf_1 loss + loss[2], # qf_2 loss + loss[3], # vf loss + loss[4], # alpha loss + ) + ) + + if self.args.log: + wandb.log( + { + "score": score, + "total loss": total_loss, + "actor loss": loss[0] * delayed_update, + "qf_1 loss": loss[1], + "qf_2 loss": loss[2], + "vf loss": loss[3], + "alpha loss": loss[4], + } + ) + + def train(self): + """Train the agent.""" + # logger + if self.args.log: + wandb.init() + 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): + state = self.env.reset() + done = False + score = 0 + self.episode_step = 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) + + state = next_state + score += reward + + # training + if len(self.memory) >= self.hyper_params["BATCH_SIZE"]: + experiences = self.memory.sample() + loss = self.update_model(experiences) + loss_episode.append(loss) # for logging + + # logging + 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"] + ) + + 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/sac.py b/scripts/examples/lunarlander_continuous_v2/sac.py new file mode 100644 index 0000000..d250756 --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/sac.py @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- +"""Run module for SAC 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.sac.agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +hyper_params = { + "GAMMA": 0.99, + "TAU": 5e-3, + "W_ENTROPY": 1e-3, + "W_MEAN_REG": 1e-3, + "W_STD_REG": 1e-3, + "W_PRE_ACTIVATION_REG": 0.0, + "LR_ACTOR": 3e-4, + "LR_VF": 3e-4, + "LR_QF1": 3e-4, + "LR_QF2": 3e-4, + "LR_ENTROPY": 3e-4, + "DELAYED_UPDATE": 2, + "BUFFER_SIZE": int(1e6), + "BATCH_SIZE": 512, + "AUTO_ENTROPY_TUNING": True, + "WEIGHT_DECAY": 0.0, + "INITIAL_RANDOM_ACTION": 5000, +} + + +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["WEIGHT_DECAY"], + ) + vf_optim = optim.Adam( + vf.parameters(), + lr=hyper_params["LR_VF"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + qf_1_optim = optim.Adam( + qf_1.parameters(), + lr=hyper_params["LR_QF1"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + qf_2_optim = optim.Adam( + qf_2.parameters(), + lr=hyper_params["LR_QF2"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + + # 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/reacher-v2/sac.py b/scripts/examples/reacher-v2/sac.py new file mode 100644 index 0000000..58154e1 --- /dev/null +++ b/scripts/examples/reacher-v2/sac.py @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- +"""Run module for SAC on Reacher-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.sac.agent import Agent + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +# hyper parameters +hyper_params = { + "GAMMA": 0.99, + "TAU": 5e-3, + "W_ENTROPY": 1e-3, + "W_MEAN_REG": 1e-3, + "W_STD_REG": 1e-3, + "W_PRE_ACTIVATION_REG": 0.0, + "LR_ACTOR": 3e-4, + "LR_VF": 3e-4, + "LR_QF1": 3e-4, + "LR_QF2": 3e-4, + "LR_ENTROPY": 3e-4, + "DELAYED_UPDATE": 2, + "BUFFER_SIZE": int(1e6), + "BATCH_SIZE": 512, + "AUTO_ENTROPY_TUNING": True, + "WEIGHT_DECAY": 0.0, + "INITIAL_RANDOM_ACTION": 20000, +} + + +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["WEIGHT_DECAY"], + ) + vf_optim = optim.Adam( + vf.parameters(), + lr=hyper_params["LR_VF"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + qf_1_optim = optim.Adam( + qf_1.parameters(), + lr=hyper_params["LR_QF1"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + qf_2_optim = optim.Adam( + qf_2.parameters(), + lr=hyper_params["LR_QF2"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + + # 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()