mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-09-09 11:25:10 +08:00
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
This commit is contained in:
Whi Kwon
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user