From 5962ff5277d1ae5f6df4f1e30fd94bbb7a63e399 Mon Sep 17 00:00:00 2001 From: Whi Kwon <30768596+whikwon@users.noreply.github.com> Date: Sun, 24 Feb 2019 13:44:20 +0900 Subject: [PATCH] Add TD3 (#10) * Add td3 * Fix flake8 * Fix action clamping * Increase episode max step, add detach to actor_loss * Fix actor update freq bug * 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 * Add random initial action in ddpg (#13) * Add random initial actions in ddpg * Add reacher-v2 example of ddpg * Add soft actor critic (#12) * Add soft actor critic * Delete unnecessary examples * Add td3 * Fix flake8 * Fix action clamping * Fix flake8 * Increase episode max step, add detach to actor_loss * Fix actor update freq bug * Fix code to reflect PR * Resolve conflict --- scripts/algorithms/common/noise.py | 8 +- scripts/algorithms/td3/agent.py | 266 ++++++++++++++++++ .../examples/lunarlander_continuous_v2/td3.py | 131 +++++++++ 3 files changed, 402 insertions(+), 3 deletions(-) create mode 100644 scripts/algorithms/td3/agent.py create mode 100644 scripts/examples/lunarlander_continuous_v2/td3.py diff --git a/scripts/algorithms/common/noise.py b/scripts/algorithms/common/noise.py index b647052..8f93ef9 100644 --- a/scripts/algorithms/common/noise.py +++ b/scripts/algorithms/common/noise.py @@ -15,21 +15,23 @@ class GaussianNoise: def __init__( self, + action_dim: int, min_sigma: float = 1.0, max_sigma: float = 1.0, decay_period: int = 1000000, ): """Initialization.""" - self.max_sigma = max_sigma + self.action_dim = action_dim self.min_sigma = min_sigma + self.max_sigma = max_sigma self.decay_period = decay_period - def sample(self, action_size: int, t: int = 0) -> float: + def sample(self, t: int = 0) -> float: """Get an action with gaussian noise.""" sigma = self.max_sigma - (self.max_sigma - self.min_sigma) * min( 1.0, t / self.decay_period ) - return np.random.normal(0, sigma, size=action_size) + return np.random.normal(0, sigma, size=self.action_dim) class OUNoise: diff --git a/scripts/algorithms/td3/agent.py b/scripts/algorithms/td3/agent.py new file mode 100644 index 0000000..21e9c03 --- /dev/null +++ b/scripts/algorithms/td3/agent.py @@ -0,0 +1,266 @@ +# -*- coding: utf-8 -*- +"""TD3 agent for episodic tasks in OpenAI Gym. + +- Author: whikwon +- Contact: whikwon@gmail.com +- Paper: https://arxiv.org/pdf/1802.09477.pdf +""" + +import argparse +import os +from typing import Tuple + +import gym +import numpy as np +import torch +import torch.nn.functional as F +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): + """ActorCritic interacting with environment. + + Attributes: + memory (ReplayBuffer): replay memory + exploration_noise (GaussianNoise): random noise for exploration + target_policy_noise (GaussianNoise): random noise for regularization + hyper_params (dict): hyper-parameters + actor (nn.Module): actor model to select actions + actor_target (nn.Module): target actor model to select actions + critic1 (nn.Module): critic1 model to predict state values + critic2 (nn.Module): critic2 model to predict state values + critic1_target (nn.Module): target critic1 model to predict state values + critic2_target (nn.Module): target critic2 model to predict state values + actor_optim (Optimizer): optimizer for training actor + critic1_optim (Optimizer): optimizer for training critic1 + critic2_optim (Optimizer): optimizer for training critic2 + 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, + noises: tuple, + ): + """Initialization. + + Args: + env (gym.Env): openAI Gym environment with discrete action space + args (argparse.Namespace): arguments including hyperparameters and training settings + hyper_params (dict): hyper-parameters + models (tuple): models including actor and critics + optims (tuple): optimizers for actor and critics + noises (tuple): noises for exploration and regularization + + """ + AbstractAgent.__init__(self, env, args) + self.actor, self.actor_target, self.critic1, self.critic1_target, \ + self.critic2, self.critic2_target = models + self.actor_optim, self.critic_optim = optims + self.hyper_params = hyper_params + self.exploration_noise, self.target_policy_noise = noises + self.curr_state = np.zeros((1,)) + self.total_steps = 0 + self.episode_steps = 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) + + # 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.""" + # initial training step, try random action for exploration + random_action_count = self.hyper_params["INITIAL_RANDOM_ACTIONS"] + + self.curr_state = state + + if self.total_steps < random_action_count and not self.args.test: + return self.env.action_space.sample() + + state = torch.FloatTensor(state).to(device) + selected_action = self.actor(state) + + if not self.args.test: + noise = torch.FloatTensor(self.exploration_noise.sample()).to(device) + selected_action = (selected_action + noise).clamp(-1.0, 1.0) + + return selected_action.detach().cpu().numpy() + + def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]: + """Take an action and return the response of the env.""" + 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 + ], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Train the model after each episode.""" + states, actions, rewards, next_states, dones = experiences + + # 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 = F.mse_loss(values1, target_values) + values2 = self.critic2(torch.cat((states, actions), dim=-1)) + critic2_loss = F.mse_loss(values2, target_values) + + critic_loss = critic1_loss + critic2_loss + 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 = -self.critic1(torch.cat((states, actions), dim=-1)).mean() + 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) + else: + actor_loss = torch.zeros(1) + + return actor_loss.data, critic1_loss.data, critic2_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.critic1.load_state_dict(params["critic1_state_dict"]) + self.critic2.load_state_dict(params["critic2_state_dict"]) + self.critic1_target.load_state_dict(params["critic1_target_state_dict"]) + self.critic2_target.load_state_dict(params["critic2_target_state_dict"]) + self.actor_optim.load_state_dict(params["actor_optim_state_dict"]) + self.critic_optim.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(), + "critic1_state_dict": self.critic1.state_dict(), + "critic2_state_dict": self.critic2.state_dict(), + "critic1_target_state_dict": self.critic1_target.state_dict(), + "critic2_target_state_dict": self.critic2_target.state_dict(), + "actor_optim_state_dict": self.actor_optim.state_dict(), + "critic_optim_state_dict": self.critic_optim.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] 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]) + ) + + if self.args.log: + wandb.log( + { + "total_steps": self.total_steps, + "score": score, + "total loss": total_loss, + "actor loss": loss[0] * self.hyper_params["POLICY_UPDATE_FREQ"], + "critic1 loss": loss[1], + "critic2 loss": loss[2], + } + ) + + def train(self): + """Train the agent.""" + # logger + if self.args.log: + wandb.init() + wandb.config.update(self.hyper_params) + wandb.watch([self.actor, self.critic1, self.critic2], log="parameters") + + for i_episode in range(1, self.args.episode_num + 1): + state = self.env.reset() + done = False + score = 0 + loss_episode = list() + self.episode_steps = 0 + + 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) + self.total_steps += 1 + self.episode_steps += 1 + + if len(self.memory) >= self.hyper_params["BATCH_SIZE"]: + experiences = self.memory.sample() + loss = self.update_model(experiences) + loss_episode.append(loss) # for logging + + state = next_state + score += reward + + # 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/td3.py b/scripts/examples/lunarlander_continuous_v2/td3.py new file mode 100644 index 0000000..ec2f969 --- /dev/null +++ b/scripts/examples/lunarlander_continuous_v2/td3.py @@ -0,0 +1,131 @@ +# -*- coding: utf-8 -*- +"""Run module for TD3 on LunarLanderContinuous-v2. + +- Author: whikwon +- Contact: whikwon@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.td3.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, + "BUFFER_SIZE": int(1e6), + "BATCH_SIZE": 100, + "LR_ACTOR": 1e-3, + "LR_CRITIC": 1e-3, + "WEIGHT_DECAY": 0.0, + "EXPLORATION_NOISE": 0.1, + "TARGET_POLICY_NOISE": 0.2, + "TARGET_POLICY_NOISE_CLIP": 0.5, + "POLICY_UPDATE_FREQ": 2, + "INITIAL_RANDOM_ACTIONS": 1e4, +} + + +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["WEIGHT_DECAY"], + ) + + critic_optim = optim.Adam( + critic_parameters, + lr=hyper_params["LR_CRITIC"], + weight_decay=hyper_params["WEIGHT_DECAY"], + ) + + # 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()