"""My SAC continuous demo"""
import argparse
import copy
import os
import random
import gym
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.distributions import Normal
def parse_args() -> argparse.Namespace:
"""Parse arguments."""
parser = argparse.ArgumentParser(description="Training")
parser.add_argument("--seed", type=int, help="Fix random seed", default=0)
parser.add_argument(
"--log_path", type=str, help="Model path", default="./training_log/"
)
parser.add_argument(
"--device", type=str, help="Run on which device", default="cuda"
)
parser.add_argument(
"--max_buffer_size", type=int, help="Max buffer size", default=1e7
)
parser.add_argument(
"--min_buffer_size", type=int, help="Min buffer size", default=5e4
)
parser.add_argument("--hidden_width", type=int, help="Hidden width", default=256)
parser.add_argument("--gamma", type=float, help="gamma", default=0.99)
parser.add_argument("--tau", type=float, help="tau", default=0.005)
parser.add_argument(
"--learning_rate", type=float, help="Learning rate", default=1e-3
)
parser.add_argument(
"--max_train_steps", type=int, help="Max training steps", default=1e7
)
parser.add_argument("--batch_size", type=int, help="Batch size", default=256)
parser.add_argument(
"--evaluate_freqency", type=int, help="Evaluate freqency", default=1e6
)
return parser.parse_args()
def set_seed(seed: int) -> None:
"""Set seed for reproducibility."""
os.environ["PYTHONHASHSEED"] = str(seed)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True)
class ReplayBuffer:
"""Replay buffer for storing transitions."""
def __init__(self, state_dim: int, action_dim: int) -> None:
self.max_size = int(args.max_buffer_size)
self.count = 0
self.size = 0
self.state = np.zeros((self.max_size, state_dim))
self.action = np.zeros((self.max_size, action_dim))
self.reward = np.zeros((self.max_size, 1))
self.next_state = np.zeros((self.max_size, state_dim))
self.done = np.zeros((self.max_size, 1))
def store(
self,
state: np.ndarray,
action: np.ndarray,
reward: np.ndarray,
next_state: np.ndarray,
done: np.ndarray,
) -> None:
"""Store a transition in the replay buffer."""
self.state[self.count] = state
self.action[self.count] = action
self.reward[self.count] = reward
self.next_state[self.count] = next_state
self.done[self.count] = done
self.count = (self.count + 1) % self.max_size
self.size = min(self.size + 1, self.max_size)
def sample(self, batch_size: int) -> tuple:
"""Sample a batch of transitions."""
index = np.random.choice(self.size, size=batch_size)
batch_state = torch.tensor(self.state[index], dtype=torch.float).to(args.device)
batch_action = torch.tensor(self.action[index], dtype=torch.float).to(
args.device
)
batch_reward = torch.tensor(self.reward[index], dtype=torch.float).to(
args.device
)
batch_next_state = torch.tensor(self.next_state[index], dtype=torch.float).to(
args.device
)
batch_done = torch.tensor(self.done[index], dtype=torch.float).to(args.device)
return batch_state, batch_action, batch_reward, batch_next_state, batch_done
class Actor(nn.Module):
"""Actor network."""
def __init__(
self, state_dim: int, action_dim: int, hidden_width: int, max_action: float
) -> None:
super().__init__()
self.max_action = max_action
self.in_layer = nn.Sequential(
nn.Linear(state_dim, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
)
self.res_layer = nn.Sequential(
nn.Linear(hidden_width, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
nn.Linear(hidden_width, hidden_width),
)
self.out_layer = nn.Sequential(
nn.Linear(hidden_width, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
)
self.mean_layer = nn.Sequential(nn.ReLU(), nn.Linear(hidden_width, action_dim))
self.log_std_layer = nn.Sequential(
nn.ReLU(inplace=True), nn.Linear(hidden_width, action_dim)
)
def forward(self, x: torch.Tensor, deterministic: bool = False) -> tuple:
"""Forward pass."""
x = self.in_layer(x)
x = self.out_layer(x + self.res_layer(x))
mean = self.mean_layer(x)
log_std = self.log_std_layer(x)
log_std = torch.clamp(log_std, -20, 2)
std = torch.exp(log_std)
dist = Normal(mean, std)
if deterministic:
action = mean
else:
action = dist.rsample()
log_pi = dist.log_prob(action).sum(dim=1, keepdim=True)
log_pi -= (2 * (np.log(2) - action - F.softplus(-2 * action))).sum(
dim=1, keepdim=True
)
action = self.max_action * torch.tanh(action)
return action, log_pi
class Critic(nn.Module):
"""Critic network."""
def __init__(self, state_dim: int, action_dim: int, hidden_width: int) -> None:
super().__init__()
self.in_layer1 = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
)
self.res_layer1 = nn.Sequential(
nn.Linear(hidden_width, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
nn.Linear(hidden_width, hidden_width),
)
self.out_layer1 = nn.Sequential(
nn.ReLU(inplace=True), nn.Linear(hidden_width, 1)
)
self.in_layer2 = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
)
self.res_layer2 = nn.Sequential(
nn.Linear(hidden_width, hidden_width),
nn.ReLU(inplace=True),
nn.LayerNorm(hidden_width),
nn.Linear(hidden_width, hidden_width),
)
self.out_layer2 = nn.Sequential(
nn.ReLU(inplace=True), nn.Linear(hidden_width, 1)
)
def forward(self, state: torch.Tensor, action: torch.Tensor) -> tuple:
"""Forward pass."""
state_action = torch.cat([state, action], 1)
q1 = self.in_layer1(state_action)
q1 = self.out_layer1(q1 + self.res_layer1(q1))
q2 = self.in_layer2(state_action)
q2 = self.out_layer2(q2 + self.res_layer2(q2))
return q1, q2
class SACContinuous:
"""Soft Actor-Critic for continuous action space."""
def __init__(self, state_dim: int, action_dim: int, max_action: float) -> None:
self.gamma = args.gamma
self.tau = args.tau
self.batch_size = args.batch_size
self.learning_rate = args.learning_rate
self.hidden_width = args.hidden_width
self.max_action = max_action
self.target_entropy = -np.log(2 * action_dim)
self.log_alpha = torch.tensor(1.0).to(args.device)
self.log_alpha.requires_grad = True
self.alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=self.learning_rate)
self.actor = Actor(state_dim, action_dim, self.hidden_width, max_action).to(
args.device
)
self.actor_optimizer = torch.optim.Adam(
self.actor.parameters(), lr=self.learning_rate
)
self.critic = Critic(state_dim, action_dim, self.hidden_width).to(args.device)
self.critic_target = copy.deepcopy(self.critic).to(args.device)
self.critic_optimizer = torch.optim.Adam(
self.critic.parameters(), lr=self.learning_rate
)
def choose_action(
self, state: np.ndarray, deterministic: bool = False
) -> np.ndarray:
"""Choose action."""
state = torch.unsqueeze(torch.tensor(state, dtype=torch.float), 0).to(
args.device
)
action, _ = self.actor(state, deterministic)
return action.data.cpu().numpy().flatten()
def learn(self, relay_buffer: ReplayBuffer) -> None:
"""Learn."""
batch_state, batch_action, batch_reward, batch_next_state, batch_done = (
relay_buffer.sample(self.batch_size)
)
batch_next_action, log_pi_ = self.actor(batch_next_state)
target_q1, target_q2 = self.critic_target(batch_next_state, batch_next_action)
target_q = batch_reward + self.gamma * (1 - batch_done) * (
torch.min(target_q1, target_q2) - self.log_alpha.exp() * log_pi_
)
current_q1, current_q2 = self.critic(batch_state, batch_action)
critic_loss = F.mse_loss(current_q1, target_q.detach()) + F.mse_loss(
current_q2, target_q.detach()
)
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
for params in self.critic.parameters():
params.requires_grad = False
action, log_pi = self.actor(batch_state)
q1, q2 = self.critic(batch_state, action)
q = torch.min(q1, q2)
actor_loss = (self.log_alpha.exp() * log_pi - q).mean()
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
for params in self.critic.parameters():
params.requires_grad = True
alpha_loss = -(
self.log_alpha.exp() * (log_pi + self.target_entropy).detach()
).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
for param, target_param in zip(
self.critic.parameters(), self.critic_target.parameters()
):
target_param.data.copy_(
self.tau * param.data + (1 - self.tau) * target_param.data
)
def evaluate_policy(env, agent: SACContinuous) -> float:
"""Evaluate the policy."""
state = env.reset()[0]
done = False
episode_reward = 0
action_num = 0
agent.actor.eval()
while not done:
action = agent.choose_action(state, deterministic=True)
next_statue, reward, done, _, _ = env.step(action)
episode_reward += reward
state = next_statue
action_num += 1
if action_num >= 1e6:
print("action_num too large.")
break
if episode_reward <= -1e6:
print("episode_reward too small.")
break
return episode_reward
def training() -> None:
"""My demo training function."""
env_name = "Pendulum-v1"
env = gym.make(env_name)
env_evaluate = gym.make(env_name)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
max_action = float(env.action_space.high[0])
agent = SACContinuous(state_dim, action_dim, max_action)
replay_buffer = ReplayBuffer(state_dim, action_dim)
evaluate_num = 0
total_steps = 0
while total_steps < args.max_train_steps:
state = env.reset()[0]
episode_steps = 0
done = False
while not done:
episode_steps += 1
action = agent.choose_action(state)
next_state, reward, done, _, _ = env.step(action)
replay_buffer.store(state, action, reward, next_state, done)
state = next_state
if total_steps >= args.min_buffer_size:
agent.learn(replay_buffer)
if (total_steps + 1) % args.evaluate_freqency == 0:
evaluate_num += 1
evaluate_reward = evaluate_policy(env_evaluate, agent)
print(
f"evaluate_num: {evaluate_num} \t evaluate_reward: {evaluate_reward}"
)
total_steps += 1
if total_steps >= args.max_train_steps:
break
env.close()
torch.save(agent.actor.state_dict(), f"{args.log_path}/trained_model.pth")
def testing() -> None:
"""My demo testing function."""
env_name = "Pendulum-v1"
env = gym.make(env_name)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
max_action = float(env.action_space.high[0])
agent = SACContinuous(state_dim, action_dim, max_action)
agent.actor.load_state_dict(torch.load(f"{args.log_path}/trained_model.pth"))
agent.actor.eval()
state = env.reset()[0]
total_rewards = 0
with torch.no_grad():
for _ in range(10000):
env.render()
action = agent.choose_action(state, deterministic=True)
new_state, reward, _, _, _ = env.step(action)
total_rewards += reward
state = new_state
env.close()
print(f"SAC actor scores: {total_rewards}")
if __name__ == "__main__":
args = parse_args()
set_seed(args.seed)
training()
testing()
