class TD3Agent(Agent):
"""ActorCritic interacting with environment.
Attributes:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
memory (ReplayBuffer): replay memory
exploration_noise (GaussianNoise): random noise for exploration
target_policy_noise (GaussianNoise): random noise for target values
actor (nn.Module): actor model to select actions
critic1 (nn.Module): critic model to predict state values
critic2 (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
actor_target (nn.Module): target actor model to select actions
critic_optim (Optimizer): optimizer for training critic
actor_optim (Optimizer): optimizer for training actor
curr_state (np.ndarray): temporary storage of the current state
total_steps (int): total step numbers
episode_steps (int): step number of the current episode
i_episode (int): current episode number
noise_cfg (ConfigDict): config of noise
"""
def __init__(
self,
env: gym.Env,
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
noise_cfg: ConfigDict,
):
"""Initialize.
Args:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
"""
Agent.__init__(self, env, args, log_cfg)
self.curr_state = np.zeros((1, ))
self.total_step = 0
self.episode_step = 0
self.update_step = 0
self.i_episode = 0
self.hyper_params = hyper_params
self.noise_cfg = noise_cfg
self.backbone_cfg = backbone
self.head_cfg = head
self.optim_cfg = optim_cfg
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.shape[0]
# noise instance to make randomness of action
self.exploration_noise = GaussianNoise(self.action_dim,
noise_cfg.exploration_noise,
noise_cfg.exploration_noise)
self.target_policy_noise = GaussianNoise(
self.action_dim,
noise_cfg.target_policy_noise,
noise_cfg.target_policy_noise,
)
if not self.args.test:
# replay memory
self.memory = ReplayBuffer(self.hyper_params.buffer_size,
self.hyper_params.batch_size)
self._init_network()
def _init_network(self):
"""Initialize networks and optimizers."""
self.head_cfg.actor.configs.state_size = self.state_dim
self.head_cfg.critic.configs.state_size = (self.state_dim[0] +
self.action_dim, )
self.head_cfg.actor.configs.output_size = self.action_dim
# create actor
self.actor = BaseNetwork(self.backbone_cfg.actor,
self.head_cfg.actor).to(device)
self.actor_target = BaseNetwork(self.backbone_cfg.actor,
self.head_cfg.actor).to(device)
self.actor_target.load_state_dict(self.actor.state_dict())
# create q_critic
self.critic1 = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
self.critic2 = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
self.critic_target1 = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
self.critic_target2 = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
self.critic_target1.load_state_dict(self.critic1.state_dict())
self.critic_target2.load_state_dict(self.critic2.state_dict())
# concat critic parameters to use one optim
critic_parameters = list(self.critic1.parameters()) + list(
self.critic2.parameters())
# create optimizers
self.actor_optim = optim.Adam(
self.actor.parameters(),
lr=self.optim_cfg.lr_actor,
weight_decay=self.optim_cfg.weight_decay,
)
self.critic_optim = optim.Adam(
critic_parameters,
lr=self.optim_cfg.lr_critic,
weight_decay=self.optim_cfg.weight_decay,
)
# load the optimizer and model parameters
if self.args.load_from is not None:
self.load_params(self.args.load_from)
def select_action(self, state: np.ndarray) -> np.ndarray:
"""Select an action from the input space."""
# initial training step, try random action for exploration
self.curr_state = state
if (self.total_step < self.hyper_params.initial_random_action
and not self.args.test):
return np.array(self.env.action_space.sample())
state = torch.FloatTensor(state).to(device)
selected_action = self.actor(state).detach().cpu().numpy()
if not self.args.test:
noise = self.exploration_noise.sample()
selected_action = np.clip(selected_action + noise, -1.0, 1.0)
return selected_action
def step(self,
action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool, dict]:
"""Take an action and return the response of the env."""
next_state, reward, done, info = self.env.step(action)
if not self.args.test:
# if last state is not terminal state in episode, done is 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, info
def update_model(self) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
self.update_step += 1
experiences = self.memory.sample()
states, actions, rewards, next_states, dones = experiences
masks = 1 - dones
# get actions with noise
noise = torch.FloatTensor(self.target_policy_noise.sample()).to(device)
clipped_noise = torch.clamp(
noise,
-self.noise_cfg.target_policy_noise_clip,
self.noise_cfg.target_policy_noise_clip,
)
next_actions = (self.actor_target(next_states) + clipped_noise).clamp(
-1.0, 1.0)
# min (Q_1', Q_2')
next_states_actions = torch.cat((next_states, next_actions), dim=-1)
next_values1 = self.critic_target1(next_states_actions)
next_values2 = self.critic_target2(next_states_actions)
next_values = torch.min(next_values1, next_values2)
# G_t = r + gamma * v(s_{t+1}) if state != Terminal
# = r otherwise
curr_returns = rewards + self.hyper_params.gamma * next_values * masks
curr_returns = curr_returns.detach()
# critic loss
state_actions = torch.cat((states, actions), dim=-1)
values1 = self.critic1(state_actions)
values2 = self.critic2(state_actions)
critic1_loss = F.mse_loss(values1, curr_returns)
critic2_loss = F.mse_loss(values2, curr_returns)
# train critic
critic_loss = critic1_loss + critic2_loss
self.critic_optim.zero_grad()
critic_loss.backward()
self.critic_optim.step()
if self.update_step % self.hyper_params.policy_update_freq == 0:
# policy loss
actions = self.actor(states)
state_actions = torch.cat((states, actions), dim=-1)
actor_loss = -self.critic1(state_actions).mean()
# train actor
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.critic1, self.critic_target1, tau)
common_utils.soft_update(self.critic2, self.critic_target2, tau)
common_utils.soft_update(self.actor, self.actor_target, tau)
else:
actor_loss = torch.zeros(1)
return actor_loss.item(), critic1_loss.item(), critic2_loss.item()
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
params = torch.load(path)
self.critic1.load_state_dict(params["critic1"])
self.critic2.load_state_dict(params["critic2"])
self.critic_target1.load_state_dict(params["critic_target1"])
self.critic_target2.load_state_dict(params["critic_target2"])
self.critic_optim.load_state_dict(params["critic_optim"])
self.actor.load_state_dict(params["actor"])
self.actor_target.load_state_dict(params["actor_target"])
self.actor_optim.load_state_dict(params["actor_optim"])
print("[INFO] loaded the model and optimizer from", path)
def save_params(self, n_episode: int): # type: ignore
"""Save model and optimizer parameters."""
params = {
"actor": self.actor.state_dict(),
"actor_target": self.actor_target.state_dict(),
"actor_optim": self.actor_optim.state_dict(),
"critic1": self.critic1.state_dict(),
"critic2": self.critic2.state_dict(),
"critic_target1": self.critic_target1.state_dict(),
"critic_target2": self.critic_target2.state_dict(),
"critic_optim": self.critic_optim.state_dict(),
}
Agent.save_params(self, params, n_episode)
def write_log(self, log_value: tuple):
"""Write log about loss and score"""
i, loss, score, policy_update_freq, avg_time_cost = log_value
total_loss = loss.sum()
print(
"[INFO] episode %d, episode_step: %d, total_step: %d, total score: %d\n"
"total loss: %f actor_loss: %.3f critic1_loss: %.3f critic2_loss: %.3f "
"(spent %.6f sec/step)\n" % (
i,
self.episode_step,
self.total_step,
score,
total_loss,
loss[0] * policy_update_freq, # actor loss
loss[1], # critic1 loss
loss[2], # critic2 loss
avg_time_cost,
))
if self.args.log:
wandb.log({
"score": score,
"total loss": total_loss,
"actor loss": loss[0] * policy_update_freq,
"critic1 loss": loss[1],
"critic2 loss": loss[2],
"time per each step": avg_time_cost,
})
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.actor, self.critic1, self.critic2], log="parameters")
for self.i_episode in range(1, self.args.episode_num + 1):
state = self.env.reset()
done = False
score = 0
loss_episode = list()
self.episode_step = 0
t_begin = time.time()
while not done:
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.total_step += 1
self.episode_step += 1
state = next_state
score += reward
if len(self.memory) >= self.hyper_params.batch_size:
loss = self.update_model()
loss_episode.append(loss) # for logging
t_end = time.time()
avg_time_cost = (t_end - t_begin) / self.episode_step
# logging
if loss_episode:
avg_loss = np.vstack(loss_episode).mean(axis=0)
log_value = (
self.i_episode,
avg_loss,
score,
self.hyper_params.policy_update_freq,
avg_time_cost,
)
self.write_log(log_value)
if self.i_episode % self.args.save_period == 0:
self.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.save_params(self.i_episode)
self.interim_test()
class PPOAgent(Agent):
"""PPO Agent.
Attributes:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
actor (nn.Module): policy gradient model to select actions
critic (nn.Module): policy gradient model to predict values
actor_optim (Optimizer): optimizer for training actor
critic_optim (Optimizer): optimizer for training critic
episode_steps (np.ndarray): step numbers of the current episode
states (list): memory for experienced states
actions (list): memory for experienced actions
rewards (list): memory for experienced rewards
values (list): memory for experienced values
masks (list): memory for masks
log_probs (list): memory for log_probs
i_episode (int): current episode number
epsilon (float): value for clipping loss
"""
def __init__(
self,
env: gym.Env, # for testing
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
):
"""Initialize.
Args:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
"""
env_gen = env_generator(env.spec.id, args)
env_multi = make_envs(env_gen, n_envs=hyper_params.n_workers)
Agent.__init__(self, env, args, log_cfg)
self.episode_steps = np.zeros(hyper_params.n_workers, dtype=np.int)
self.states: list = []
self.actions: list = []
self.rewards: list = []
self.values: list = []
self.masks: list = []
self.log_probs: list = []
self.i_episode = 0
self.next_state = np.zeros((1, ))
self.hyper_params = hyper_params
self.backbone_cfg = backbone
self.head_cfg = head
self.optim_cfg = optim_cfg
if not self.args.test:
self.env = env_multi
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.shape[0]
self.epsilon = hyper_params.max_epsilon
self._init_network()
def _init_network(self):
"""Initialize networks and optimizers."""
self.head_cfg.actor.configs.state_size = (
self.head_cfg.critic.configs.state_size) = self.state_dim
self.head_cfg.actor.configs.output_size = self.action_dim
# create actor
self.actor = BaseNetwork(self.backbone_cfg.actor,
self.head_cfg.actor).to(device)
self.critic = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
# create optimizer
self.actor_optim = optim.Adam(
self.actor.parameters(),
lr=self.optim_cfg.lr_actor,
weight_decay=self.optim_cfg.weight_decay,
)
self.critic_optim = optim.Adam(
self.critic.parameters(),
lr=self.optim_cfg.lr_critic,
weight_decay=self.optim_cfg.weight_decay,
)
# load model parameters
if self.args.load_from is not None:
self.load_params(self.args.load_from)
def select_action(self, state: np.ndarray) -> torch.Tensor:
"""Select an action from the input space."""
state = torch.FloatTensor(state).to(device)
selected_action, dist = self.actor(state)
if self.args.test and not self.is_discrete:
selected_action = dist.mean
if not self.args.test:
value = self.critic(state)
self.states.append(state)
self.actions.append(selected_action)
self.values.append(value)
self.log_probs.append(dist.log_prob(selected_action))
return selected_action
def step(
self,
action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool, dict]:
next_state, reward, done, info = self.env.step(
action.detach().cpu().numpy())
if not self.args.test:
# if the last state is not a terminal state, store done as false
done_bool = done.copy()
done_bool[np.where(
self.episode_steps == self.args.max_episode_steps)] = False
self.rewards.append(
torch.FloatTensor(reward).unsqueeze(1).to(device))
self.masks.append(
torch.FloatTensor(1 - done_bool).unsqueeze(1).to(device))
return next_state, reward, done, info
def update_model(self) -> Tuple[torch.Tensor, ...]:
"""Train the model after every N episodes."""
next_state = torch.FloatTensor(self.next_state).to(device)
next_value = self.critic(next_state)
returns = ppo_utils.compute_gae(
next_value,
self.rewards,
self.masks,
self.values,
self.hyper_params.gamma,
self.hyper_params.tau,
)
states = torch.cat(self.states)
actions = torch.cat(self.actions)
returns = torch.cat(returns).detach()
values = torch.cat(self.values).detach()
log_probs = torch.cat(self.log_probs).detach()
advantages = returns - values
if self.is_discrete:
actions = actions.unsqueeze(1)
log_probs = log_probs.unsqueeze(1)
if self.hyper_params.standardize_advantage:
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-7)
actor_losses, critic_losses, total_losses = [], [], []
for state, action, old_value, old_log_prob, return_, adv in ppo_utils.ppo_iter(
self.hyper_params.epoch,
self.hyper_params.batch_size,
states,
actions,
values,
log_probs,
returns,
advantages,
):
# calculate ratios
_, dist = self.actor(state)
log_prob = dist.log_prob(action)
ratio = (log_prob - old_log_prob).exp()
# actor_loss
surr_loss = ratio * adv
clipped_surr_loss = (
torch.clamp(ratio, 1.0 - self.epsilon, 1.0 + self.epsilon) *
adv)
actor_loss = -torch.min(surr_loss, clipped_surr_loss).mean()
# critic_loss
value = self.critic(state)
if self.hyper_params.use_clipped_value_loss:
value_pred_clipped = old_value + torch.clamp(
(value - old_value), -self.epsilon, self.epsilon)
value_loss_clipped = (return_ - value_pred_clipped).pow(2)
value_loss = (return_ - value).pow(2)
critic_loss = 0.5 * torch.max(value_loss,
value_loss_clipped).mean()
else:
critic_loss = 0.5 * (return_ - value).pow(2).mean()
# entropy
entropy = dist.entropy().mean()
# total_loss
w_value = self.hyper_params.w_value
w_entropy = self.hyper_params.w_entropy
total_loss = actor_loss + w_value * critic_loss - w_entropy * entropy
# train critic
gradient_clip_ac = self.hyper_params.gradient_clip_ac
gradient_clip_cr = self.hyper_params.gradient_clip_cr
self.critic_optim.zero_grad()
total_loss.backward(retain_graph=True)
nn.utils.clip_grad_norm_(self.critic.parameters(),
gradient_clip_ac)
self.critic_optim.step()
# train actor
self.actor_optim.zero_grad()
total_loss.backward()
nn.utils.clip_grad_norm_(self.critic.parameters(),
gradient_clip_cr)
self.actor_optim.step()
actor_losses.append(actor_loss.item())
critic_losses.append(critic_loss.item())
total_losses.append(total_loss.item())
self.states, self.actions, self.rewards = [], [], []
self.values, self.masks, self.log_probs = [], [], []
actor_loss = sum(actor_losses) / len(actor_losses)
critic_loss = sum(critic_losses) / len(critic_losses)
total_loss = sum(total_losses) / len(total_losses)
return actor_loss, critic_loss, total_loss
def decay_epsilon(self, t: int = 0):
"""Decay epsilon until reaching the minimum value."""
max_epsilon = self.hyper_params.max_epsilon
min_epsilon = self.hyper_params.min_epsilon
epsilon_decay_period = self.hyper_params.epsilon_decay_period
self.epsilon = self.epsilon - (max_epsilon - min_epsilon) * min(
1.0, t / (epsilon_decay_period + 1e-7))
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
params = torch.load(path)
self.actor.load_state_dict(params["actor_state_dict"])
self.critic.load_state_dict(params["critic_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): # type: ignore
"""Save model and optimizer parameters."""
params = {
"actor_state_dict": self.actor.state_dict(),
"critic_state_dict": self.critic.state_dict(),
"actor_optim_state_dict": self.actor_optim.state_dict(),
"critic_optim_state_dict": self.critic_optim.state_dict(),
}
Agent.save_params(self, params, n_episode)
def write_log(
self,
log_value: tuple,
):
i_episode, n_step, score, actor_loss, critic_loss, total_loss = log_value
print("[INFO] episode %d\tepisode steps: %d\ttotal score: %d\n"
"total loss: %f\tActor loss: %f\tCritic loss: %f\n" %
(i_episode, n_step, score, total_loss, actor_loss, critic_loss))
if self.args.log:
wandb.log({
"total loss": total_loss,
"actor loss": actor_loss,
"critic loss": critic_loss,
"score": score,
})
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.actor, self.critic], log="parameters")
score = 0
i_episode_prev = 0
loss = [0.0, 0.0, 0.0]
state = self.env.reset()
while self.i_episode <= self.args.episode_num:
for _ in range(self.hyper_params.rollout_len):
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.episode_steps += 1
state = next_state
score += reward[0]
i_episode_prev = self.i_episode
self.i_episode += done.sum()
if (self.i_episode // self.args.save_period) != (
i_episode_prev // self.args.save_period):
self.save_params(self.i_episode)
if done[0]:
n_step = self.episode_steps[0]
log_value = (
self.i_episode,
n_step,
score,
loss[0],
loss[1],
loss[2],
)
self.write_log(log_value)
score = 0
self.episode_steps[np.where(done)] = 0
self.next_state = next_state
loss = self.update_model()
self.decay_epsilon(self.i_episode)
# termination
self.env.close()
self.save_params(self.i_episode)
class DDPGAgent(Agent):
"""ActorCritic interacting with environment.
Attributes:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
memory (ReplayBuffer): replay memory
noise (OUNoise): random noise for exploration
actor (nn.Module): actor model to select actions
actor_target (nn.Module): target actor model to select actions
critic (nn.Module): critic model to predict state values
critic_target (nn.Module): target critic model to predict state values
actor_optim (Optimizer): optimizer for training actor
critic_optim (Optimizer): optimizer for training critic
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
"""
def __init__(
self,
env: gym.Env,
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
noise_cfg: ConfigDict,
):
"""Initialize."""
Agent.__init__(self, env, args, log_cfg)
self.curr_state = np.zeros((1,))
self.total_step = 0
self.episode_step = 0
self.i_episode = 0
self.hyper_params = hyper_params
self.backbone_cfg = backbone
self.head_cfg = head
self.optim_cfg = optim_cfg
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.shape[0]
# set noise
self.noise = OUNoise(
self.action_dim,
theta=noise_cfg.ou_noise_theta,
sigma=noise_cfg.ou_noise_sigma,
)
self._initialize()
self._init_network()
# pylint: disable=attribute-defined-outside-init
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
)
# pylint: disable=attribute-defined-outside-init
def _init_network(self):
"""Initialize networks and optimizers."""
self.head_cfg.actor.configs.state_size = self.state_dim
# ddpg critic gets state & action as input,
# and make the type to tuple to conform the gym action_space type.
self.head_cfg.critic.configs.state_size = (self.state_dim[0] + self.action_dim,)
self.head_cfg.actor.configs.output_size = self.action_dim
# create actor
self.actor = BaseNetwork(self.backbone_cfg.actor, self.head_cfg.actor).to(
device
)
self.actor_target = BaseNetwork(
self.backbone_cfg.actor, self.head_cfg.actor
).to(device)
self.actor_target.load_state_dict(self.actor.state_dict())
# create critic
self.critic = BaseNetwork(self.backbone_cfg.critic, self.head_cfg.critic).to(
device
)
self.critic_target = BaseNetwork(
self.backbone_cfg.critic, self.head_cfg.critic
).to(device)
self.critic_target.load_state_dict(self.critic.state_dict())
# create optimizer
self.actor_optim = optim.Adam(
self.actor.parameters(),
lr=self.optim_cfg.lr_actor,
weight_decay=self.optim_cfg.weight_decay,
)
self.critic_optim = optim.Adam(
self.critic.parameters(),
lr=self.optim_cfg.lr_critic,
weight_decay=self.optim_cfg.weight_decay,
)
# load the optimizer and model parameters
if self.args.load_from is not None:
self.load_params(self.args.load_from)
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 (
self.total_step < self.hyper_params.initial_random_action
and not self.args.test
):
return np.array(self.env.action_space.sample())
selected_action = self.actor(state).detach().cpu().numpy()
if not self.args.test:
noise = self.noise.sample()
selected_action = np.clip(selected_action + noise, -1.0, 1.0)
return selected_action
# pylint: disable=no-self-use
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, dict]:
"""Take an action and return the response of the env."""
next_state, reward, done, info = 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
)
transition = (self.curr_state, action, reward, next_state, done_bool)
self._add_transition_to_memory(transition)
return next_state, reward, done, info
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) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
experiences = self.memory.sample()
states, actions, rewards, next_states, dones = experiences
# G_t = r + gamma * v(s_{t+1}) if state != Terminal
# = r otherwise
masks = 1 - dones
next_actions = self.actor_target(next_states)
next_values = self.critic_target(torch.cat((next_states, next_actions), dim=-1))
curr_returns = rewards + self.hyper_params.gamma * next_values * masks
curr_returns = curr_returns.to(device)
# train critic
gradient_clip_ac = self.hyper_params.gradient_clip_ac
gradient_clip_cr = self.hyper_params.gradient_clip_cr
values = self.critic(torch.cat((states, actions), dim=-1))
critic_loss = F.mse_loss(values, curr_returns)
self.critic_optim.zero_grad()
critic_loss.backward()
nn.utils.clip_grad_norm_(self.critic.parameters(), gradient_clip_cr)
self.critic_optim.step()
# train actor
actions = self.actor(states)
actor_loss = -self.critic(torch.cat((states, actions), dim=-1)).mean()
self.actor_optim.zero_grad()
actor_loss.backward()
nn.utils.clip_grad_norm_(self.actor.parameters(), gradient_clip_ac)
self.actor_optim.step()
# update target networks
common_utils.soft_update(self.actor, self.actor_target, self.hyper_params.tau)
common_utils.soft_update(self.critic, self.critic_target, self.hyper_params.tau)
return actor_loss.item(), critic_loss.item()
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
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.critic.load_state_dict(params["critic_state_dict"])
self.critic_target.load_state_dict(params["critic_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(),
"critic_state_dict": self.critic.state_dict(),
"critic_target_state_dict": self.critic_target.state_dict(),
"actor_optim_state_dict": self.actor_optim.state_dict(),
"critic_optim_state_dict": self.critic_optim.state_dict(),
}
Agent._save_params(self, params, n_episode)
def write_log(self, log_value: tuple):
"""Write log about loss and score"""
i, loss, score, avg_time_cost = log_value
total_loss = loss.sum()
print(
"[INFO] episode %d, episode step: %d, total step: %d, total score: %d\n"
"total loss: %f actor_loss: %.3f critic_loss: %.3f (spent %.6f sec/step)\n"
% (
i,
self.episode_step,
self.total_step,
score,
total_loss,
loss[0],
loss[1],
avg_time_cost,
) # actor loss # critic loss
)
if self.args.log:
wandb.log(
{
"score": score,
"total loss": total_loss,
"actor loss": loss[0],
"critic loss": loss[1],
"time per each step": avg_time_cost,
}
)
# pylint: disable=no-self-use, unnecessary-pass
def pretrain(self):
"""Pretraining steps."""
pass
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.actor, self.critic], log="parameters")
# pre-training if needed
self.pretrain()
for self.i_episode in range(1, self.args.episode_num + 1):
state = self.env.reset()
done = False
score = 0
self.episode_step = 0
losses = list()
t_begin = time.time()
while not done:
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.total_step += 1
self.episode_step += 1
if len(self.memory) >= self.hyper_params.batch_size:
for _ in range(self.hyper_params.multiple_update):
loss = self.update_model()
losses.append(loss) # for logging
state = next_state
score += reward
t_end = time.time()
avg_time_cost = (t_end - t_begin) / self.episode_step
# logging
if losses:
avg_loss = np.vstack(losses).mean(axis=0)
log_value = (self.i_episode, avg_loss, score, avg_time_cost)
self.write_log(log_value)
losses.clear()
if self.i_episode % self.args.save_period == 0:
self.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.save_params(self.i_episode)
self.interim_test()
class SACAgent(Agent):
"""SAC agent interacting with environment.
Attrtibutes:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
memory (ReplayBuffer): replay memory
actor (nn.Module): actor model to select actions
actor_target (nn.Module): target actor model to select actions
actor_optim (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_optim1 (Optimizer): optimizer for training critic_1
critic_optim2 (Optimizer): optimizer for training critic_2
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
update_step (int): step number of updates
i_episode (int): current episode number
target_entropy (int): desired entropy used for the inequality constraint
log_alpha (torch.Tensor): weight for entropy
alpha_optim (Optimizer): optimizer for alpha
"""
def __init__(
self,
env: gym.Env,
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
):
"""Initialize.
Args:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
"""
Agent.__init__(self, env, args, log_cfg)
self.curr_state = np.zeros((1, ))
self.total_step = 0
self.episode_step = 0
self.update_step = 0
self.i_episode = 0
self.hyper_params = hyper_params
self.backbone_cfg = backbone
self.head_cfg = head
self.optim_cfg = optim_cfg
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.shape[0]
# target entropy
target_entropy = -np.prod((self.action_dim, )).item() # heuristic
# automatic entropy tuning
if hyper_params.auto_entropy_tuning:
self.target_entropy = target_entropy
self.log_alpha = torch.zeros(1, requires_grad=True, device=device)
self.alpha_optim = optim.Adam([self.log_alpha],
lr=optim_cfg.lr_entropy)
self._initialize()
self._init_network()
# pylint: disable=attribute-defined-outside-init
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)
# pylint: disable=attribute-defined-outside-init
def _init_network(self):
"""Initialize networks and optimizers."""
self.head_cfg.actor.configs.state_size = (
self.head_cfg.critic_vf.configs.state_size) = self.state_dim
self.head_cfg.critic_qf.configs.state_size = (self.state_dim[0] +
self.action_dim, )
self.head_cfg.actor.configs.output_size = self.action_dim
# create actor
self.actor = BaseNetwork(self.backbone_cfg.actor,
self.head_cfg.actor).to(device)
# create v_critic
self.vf = BaseNetwork(self.backbone_cfg.critic_vf,
self.head_cfg.critic_vf).to(device)
self.vf_target = BaseNetwork(self.backbone_cfg.critic_vf,
self.head_cfg.critic_vf).to(device)
self.vf_target.load_state_dict(self.vf.state_dict())
# create q_critic
self.qf_1 = BaseNetwork(self.backbone_cfg.critic_qf,
self.head_cfg.critic_qf).to(device)
self.qf_2 = BaseNetwork(self.backbone_cfg.critic_qf,
self.head_cfg.critic_qf).to(device)
# create optimizers
self.actor_optim = optim.Adam(
self.actor.parameters(),
lr=self.optim_cfg.lr_actor,
weight_decay=self.optim_cfg.weight_decay,
)
self.vf_optim = optim.Adam(
self.vf.parameters(),
lr=self.optim_cfg.lr_vf,
weight_decay=self.optim_cfg.weight_decay,
)
self.qf_1_optim = optim.Adam(
self.qf_1.parameters(),
lr=self.optim_cfg.lr_qf1,
weight_decay=self.optim_cfg.weight_decay,
)
self.qf_2_optim = optim.Adam(
self.qf_2.parameters(),
lr=self.optim_cfg.lr_qf2,
weight_decay=self.optim_cfg.weight_decay,
)
# load the optimizer and model parameters
if self.args.load_from is not None:
self.load_params(self.args.load_from)
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 (self.total_step < self.hyper_params.initial_random_action
and not self.args.test):
return np.array(self.env.action_space.sample())
if self.args.test:
_, _, _, selected_action, _ = self.actor(state)
else:
selected_action, _, _, _, _ = self.actor(state)
return selected_action.detach().cpu().numpy()
# pylint: disable=no-self-use
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, dict]:
"""Take an action and return the response of the env."""
next_state, reward, done, info = 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)
transition = (self.curr_state, action, reward, next_state,
done_bool)
self._add_transition_to_memory(transition)
return next_state, reward, done, info
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) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
self.update_step += 1
experiences = self.memory.sample()
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_optim.zero_grad()
alpha_loss.backward()
self.alpha_optim.step()
alpha = self.log_alpha.exp()
else:
alpha_loss = torch.zeros(1)
alpha = self.hyper_params.w_entropy
# Q function loss
masks = 1 - dones
states_actions = torch.cat((states, actions), dim=-1)
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
states_actions = torch.cat((states, new_actions), dim=-1)
v_pred = self.vf(states)
q_pred = torch.min(self.qf_1(states_actions),
self.qf_2(states_actions))
v_target = q_pred - alpha * log_prob
vf_loss = F.mse_loss(v_pred, v_target.detach())
# train Q functions
self.qf_1_optim.zero_grad()
qf_1_loss.backward()
self.qf_1_optim.step()
self.qf_2_optim.zero_grad()
qf_2_loss.backward()
self.qf_2_optim.step()
# train V function
self.vf_optim.zero_grad()
vf_loss.backward()
self.vf_optim.step()
if self.update_step % self.hyper_params.policy_update_freq == 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_optim.zero_grad()
actor_loss.backward()
self.actor_optim.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.item(),
qf_1_loss.item(),
qf_2_loss.item(),
vf_loss.item(),
alpha_loss.item(),
)
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
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_optim.load_state_dict(params["actor_optim"])
self.qf_1_optim.load_state_dict(params["qf_1_optim"])
self.qf_2_optim.load_state_dict(params["qf_2_optim"])
self.vf_optim.load_state_dict(params["vf_optim"])
if self.hyper_params.auto_entropy_tuning:
self.alpha_optim.load_state_dict(params["alpha_optim"])
print("[INFO] loaded the model and optimizer from", path)
def save_params(self, n_episode: int): # type: ignore
"""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_optim.state_dict(),
"qf_1_optim": self.qf_1_optim.state_dict(),
"qf_2_optim": self.qf_2_optim.state_dict(),
"vf_optim": self.vf_optim.state_dict(),
}
if self.hyper_params.auto_entropy_tuning:
params["alpha_optim"] = self.alpha_optim.state_dict()
Agent.save_params(self, params, n_episode)
def write_log(self, log_value: tuple):
"""Write log about loss and score"""
i, loss, score, policy_update_freq, avg_time_cost = log_value
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 (spent %.6f sec/step)\n" % (
i,
self.episode_step,
self.total_step,
score,
total_loss,
loss[0] * policy_update_freq, # actor loss
loss[1], # qf_1 loss
loss[2], # qf_2 loss
loss[3], # vf loss
loss[4], # alpha loss
avg_time_cost,
))
if self.args.log:
wandb.log({
"score": score,
"total loss": total_loss,
"actor loss": loss[0] * policy_update_freq,
"qf_1 loss": loss[1],
"qf_2 loss": loss[2],
"vf loss": loss[3],
"alpha loss": loss[4],
"time per each step": avg_time_cost,
})
# pylint: disable=no-self-use, unnecessary-pass
def pretrain(self):
"""Pretraining steps."""
pass
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.actor, self.vf, self.qf_1, self.qf_2], log="parameters")
# pre-training if needed
self.pretrain()
for self.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()
t_begin = time.time()
while not done:
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.total_step += 1
self.episode_step += 1
state = next_state
score += reward
# training
if len(self.memory) >= self.hyper_params.batch_size:
for _ in range(self.hyper_params.multiple_update):
loss = self.update_model()
loss_episode.append(loss) # for logging
t_end = time.time()
avg_time_cost = (t_end - t_begin) / self.episode_step
# logging
if loss_episode:
avg_loss = np.vstack(loss_episode).mean(axis=0)
log_value = (
self.i_episode,
avg_loss,
score,
self.hyper_params.policy_update_freq,
avg_time_cost,
)
self.write_log(log_value)
if self.i_episode % self.args.save_period == 0:
self.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.save_params(self.i_episode)
self.interim_test()
class A2CAgent(Agent):
"""1-Step Advantage Actor-Critic interacting with environment.
Attributes:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
actor (nn.Module): policy model to select actions
critic (nn.Module): critic model to evaluate states
actor_optim (Optimizer): optimizer for actor
critic_optim (Optimizer): optimizer for critic
episode_step (int): step number of the current episode
i_episode (int): current episode number
transition (list): recent transition information
"""
def __init__(
self,
env: gym.Env,
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
):
"""Initialize."""
Agent.__init__(self, env, args, log_cfg)
self.transition: list = list()
self.episode_step = 0
self.i_episode = 0
self.hyper_params = hyper_params
self.backbone_cfg = backbone
self.head_cfg = head
self.optim_cfg = optim_cfg
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.shape[0]
self._init_network()
def _init_network(self):
# create models
self.head_cfg.actor.configs.state_size = (
self.head_cfg.critic.configs.state_size) = self.state_dim
self.head_cfg.actor.configs.output_size = self.action_dim
self.actor = BaseNetwork(self.backbone_cfg.actor,
self.head_cfg.actor).to(device)
self.critic = BaseNetwork(self.backbone_cfg.critic,
self.head_cfg.critic).to(device)
# create optimizer
self.actor_optim = optim.Adam(
self.actor.parameters(),
lr=self.optim_cfg.lr_actor,
weight_decay=self.optim_cfg.weight_decay,
)
self.critic_optim = optim.Adam(
self.critic.parameters(),
lr=self.optim_cfg.lr_critic,
weight_decay=self.optim_cfg.weight_decay,
)
if self.args.load_from is not None:
self.load_params(self.args.load_from)
def select_action(self, state: np.ndarray) -> torch.Tensor:
"""Select an action from the input space."""
state = torch.FloatTensor(state).to(device)
selected_action, dist = self.actor(state)
if self.args.test:
selected_action = dist.mean
else:
predicted_value = self.critic(state)
log_prob = dist.log_prob(selected_action).sum(dim=-1)
self.transition = []
self.transition.extend([log_prob, predicted_value])
return selected_action
def step(
self,
action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool, dict]:
"""Take an action and return the response of the env."""
action = action.detach().cpu().numpy()
next_state, reward, done, info = self.env.step(action)
if not self.args.test:
done_bool = done
if self.episode_step == self.args.max_episode_steps:
done_bool = False
self.transition.extend([next_state, reward, done_bool])
return next_state, reward, done, info
def update_model(self) -> Tuple[torch.Tensor, ...]:
log_prob, pred_value, next_state, reward, done = self.transition
next_state = torch.FloatTensor(next_state).to(device)
# Q_t = r + gamma * V(s_{t+1}) if state != Terminal
# = r otherwise
mask = 1 - done
next_value = self.critic(next_state).detach()
q_value = reward + self.hyper_params.gamma * next_value * mask
q_value = q_value.to(device)
# advantage = Q_t - V(s_t)
advantage = q_value - pred_value
# calculate loss at the current step
policy_loss = -advantage.detach(
) * log_prob # adv. is not backpropagated
policy_loss += self.hyper_params.w_entropy * -log_prob # entropy
value_loss = F.smooth_l1_loss(pred_value, q_value.detach())
# train
gradient_clip_ac = self.hyper_params.gradient_clip_ac
gradient_clip_cr = self.hyper_params.gradient_clip_cr
self.actor_optim.zero_grad()
policy_loss.backward()
nn.utils.clip_grad_norm_(self.actor.parameters(), gradient_clip_ac)
self.actor_optim.step()
self.critic_optim.zero_grad()
value_loss.backward()
nn.utils.clip_grad_norm_(self.critic.parameters(), gradient_clip_cr)
self.critic_optim.step()
return policy_loss.item(), value_loss.item()
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
params = torch.load(path)
self.actor.load_state_dict(params["actor_state_dict"])
self.critic.load_state_dict(params["critic_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(),
"critic_state_dict": self.critic.state_dict(),
"actor_optim_state_dict": self.actor_optim.state_dict(),
"critic_optim_state_dict": self.critic_optim.state_dict(),
}
Agent._save_params(self, params, n_episode)
def write_log(self, log_value: tuple):
i, score, policy_loss, value_loss = log_value
total_loss = policy_loss + value_loss
print(
"[INFO] episode %d\tepisode step: %d\ttotal score: %d\n"
"total loss: %.4f\tpolicy loss: %.4f\tvalue loss: %.4f\n" %
(i, self.episode_step, score, total_loss, policy_loss, value_loss))
if self.args.log:
wandb.log({
"total loss": total_loss,
"policy loss": policy_loss,
"value loss": value_loss,
"score": score,
})
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.actor, self.critic], log="parameters")
for self.i_episode in range(1, self.args.episode_num + 1):
state = self.env.reset()
done = False
score = 0
policy_loss_episode = list()
value_loss_episode = list()
self.episode_step = 0
while not done:
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.episode_step += 1
policy_loss, value_loss = self.update_model()
policy_loss_episode.append(policy_loss)
value_loss_episode.append(value_loss)
state = next_state
score += reward
# logging
policy_loss = np.array(policy_loss_episode).mean()
value_loss = np.array(value_loss_episode).mean()
log_value = (self.i_episode, score, policy_loss, value_loss)
self.write_log(log_value)
if self.i_episode % self.args.save_period == 0:
self.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.save_params(self.i_episode)
self.interim_test()
class DQNAgent(Agent):
"""DQN interacting with environment.
Attribute:
env (gym.Env): openAI Gym environment
args (argparse.Namespace): arguments including hyperparameters and training settings
hyper_params (ConfigDict): hyper-parameters
network_cfg (ConfigDict): config of network for training agent
optim_cfg (ConfigDict): config of optimizer
state_dim (int): state size of env
action_dim (int): action size of env
memory (PrioritizedReplayBuffer): replay memory
dqn (nn.Module): actor model to select actions
dqn_target (nn.Module): target actor model to select actions
dqn_optim (Optimizer): optimizer for training actor
curr_state (np.ndarray): temporary storage of the current state
total_step (int): total step number
episode_step (int): step number of the current episode
i_episode (int): current episode number
epsilon (float): parameter for epsilon greedy policy
n_step_buffer (deque): n-size buffer to calculate n-step returns
per_beta (float): beta parameter for prioritized replay buffer
use_conv (bool): whether or not to use convolution layer
use_n_step (bool): whether or not to use n-step returns
"""
def __init__(
self,
env: gym.Env,
args: argparse.Namespace,
log_cfg: ConfigDict,
hyper_params: ConfigDict,
backbone: ConfigDict,
head: ConfigDict,
optim_cfg: ConfigDict,
):
"""Initialize."""
Agent.__init__(self, env, args, log_cfg)
self.curr_state = np.zeros(1)
self.episode_step = 0
self.total_step = 0
self.i_episode = 0
self.hyper_params = hyper_params
self.optim_cfg = optim_cfg
self.backbone_cfg = backbone
self.head_cfg = head
self.state_dim = self.env.observation_space.shape
self.action_dim = self.env.action_space.n
self.per_beta = hyper_params.per_beta
self.use_conv = len(self.state_dim) > 1
self.use_n_step = hyper_params.n_step > 1
if head.configs.use_noisy_net:
self.max_epsilon = 0.0
self.min_epsilon = 0.0
self.epsilon = 0.0
else:
self.max_epsilon = hyper_params.max_epsilon
self.min_epsilon = hyper_params.min_epsilon
self.epsilon = hyper_params.max_epsilon
self._initialize()
self._init_network()
# pylint: disable=attribute-defined-outside-init
def _initialize(self):
"""Initialize non-common things."""
if not self.args.test:
# replay memory for a single step
self.memory = PrioritizedReplayBuffer(
self.hyper_params.buffer_size,
self.hyper_params.batch_size,
alpha=self.hyper_params.per_alpha,
)
# replay memory for multi-steps
if self.use_n_step:
self.memory_n = ReplayBuffer(
self.hyper_params.buffer_size,
self.hyper_params.batch_size,
n_step=self.hyper_params.n_step,
gamma=self.hyper_params.gamma,
)
# pylint: disable=attribute-defined-outside-init
def _init_network(self):
"""Initialize networks and optimizers."""
self.head_cfg.configs.state_size = self.state_dim
self.head_cfg.configs.output_size = self.action_dim
self.dqn = BaseNetwork(self.backbone_cfg, self.head_cfg).to(device)
self.dqn_target = BaseNetwork(self.backbone_cfg, self.head_cfg).to(device)
self.loss_fn = build_loss(self.hyper_params.loss_type)
self.dqn_target.load_state_dict(self.dqn.state_dict())
# create optimizer
self.dqn_optim = optim.Adam(
self.dqn.parameters(),
lr=self.optim_cfg.lr_dqn,
weight_decay=self.optim_cfg.weight_decay,
eps=self.optim_cfg.adam_eps,
)
# load the optimizer and model parameters
if self.args.load_from is not None:
self.load_params(self.args.load_from)
def select_action(self, state: np.ndarray) -> np.ndarray:
"""Select an action from the input space."""
self.curr_state = state
# epsilon greedy policy
# pylint: disable=comparison-with-callable
if not self.args.test and self.epsilon > np.random.random():
selected_action = np.array(self.env.action_space.sample())
else:
state = self._preprocess_state(state)
selected_action = self.dqn(state).argmax()
selected_action = selected_action.detach().cpu().numpy()
return selected_action
# pylint: disable=no-self-use
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, dict]:
"""Take an action and return the response of the env."""
next_state, reward, done, info = 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
)
transition = (self.curr_state, action, reward, next_state, done_bool)
self._add_transition_to_memory(transition)
return next_state, reward, done, info
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 update_model(self) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
# 1 step loss
experiences_1 = self.memory.sample(self.per_beta)
weights, indices = experiences_1[-3:-1]
gamma = self.hyper_params.gamma
dq_loss_element_wise, q_values = self.loss_fn(
self.dqn, self.dqn_target, experiences_1, gamma, self.head_cfg
)
dq_loss = torch.mean(dq_loss_element_wise * weights)
# n step loss
if self.use_n_step:
experiences_n = self.memory_n.sample(indices)
gamma = self.hyper_params.gamma ** self.hyper_params.n_step
dq_loss_n_element_wise, q_values_n = self.loss_fn(
self.dqn, self.dqn_target, experiences_n, gamma, self.head_cfg
)
# to update loss and priorities
q_values = 0.5 * (q_values + q_values_n)
dq_loss_element_wise += dq_loss_n_element_wise * self.hyper_params.w_n_step
dq_loss = torch.mean(dq_loss_element_wise * weights)
# q_value regularization
q_regular = torch.norm(q_values, 2).mean() * self.hyper_params.w_q_reg
# total loss
loss = dq_loss + q_regular
self.dqn_optim.zero_grad()
loss.backward()
clip_grad_norm_(self.dqn.parameters(), self.hyper_params.gradient_clip)
self.dqn_optim.step()
# update target networks
common_utils.soft_update(self.dqn, self.dqn_target, self.hyper_params.tau)
# update priorities in PER
loss_for_prior = dq_loss_element_wise.detach().cpu().numpy()
new_priorities = loss_for_prior + self.hyper_params.per_eps
self.memory.update_priorities(indices, new_priorities)
# increase beta
fraction = min(float(self.i_episode) / self.args.episode_num, 1.0)
self.per_beta = self.per_beta + fraction * (1.0 - self.per_beta)
if self.head_cfg.configs.use_noisy_net:
self.dqn.head.reset_noise()
self.dqn_target.head.reset_noise()
return loss.item(), q_values.mean().item()
def load_params(self, path: str):
"""Load model and optimizer parameters."""
Agent.load_params(self, path)
params = torch.load(path)
self.dqn.load_state_dict(params["dqn_state_dict"])
self.dqn_target.load_state_dict(params["dqn_target_state_dict"])
self.dqn_optim.load_state_dict(params["dqn_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 = {
"dqn_state_dict": self.dqn.state_dict(),
"dqn_target_state_dict": self.dqn_target.state_dict(),
"dqn_optim_state_dict": self.dqn_optim.state_dict(),
}
Agent._save_params(self, params, n_episode)
def write_log(self, log_value: tuple):
"""Write log about loss and score"""
i, loss, score, avg_time_cost = log_value
print(
"[INFO] episode %d, episode step: %d, total step: %d, total score: %f\n"
"epsilon: %f, loss: %f, avg q-value: %f (spent %.6f sec/step)\n"
% (
i,
self.episode_step,
self.total_step,
score,
self.epsilon,
loss[0],
loss[1],
avg_time_cost,
)
)
if self.args.log:
wandb.log(
{
"score": score,
"epsilon": self.epsilon,
"dqn loss": loss[0],
"avg q values": loss[1],
"time per each step": avg_time_cost,
}
)
# pylint: disable=no-self-use, unnecessary-pass
def pretrain(self):
"""Pretraining steps."""
pass
def train(self):
"""Train the agent."""
# logger
if self.args.log:
self.set_wandb()
# wandb.watch([self.dqn], log="parameters")
# pre-training if needed
self.pretrain()
for self.i_episode in range(1, self.args.episode_num + 1):
state = self.env.reset()
self.episode_step = 0
losses = list()
done = False
score = 0
t_begin = time.time()
while not done:
if self.args.render and self.i_episode >= self.args.render_after:
self.env.render()
action = self.select_action(state)
next_state, reward, done, _ = self.step(action)
self.total_step += 1
self.episode_step += 1
if len(self.memory) >= self.hyper_params.update_starts_from:
if self.total_step % self.hyper_params.train_freq == 0:
for _ in range(self.hyper_params.multiple_update):
loss = self.update_model()
losses.append(loss) # for logging
# decrease epsilon
self.epsilon = max(
self.epsilon
- (self.max_epsilon - self.min_epsilon)
* self.hyper_params.epsilon_decay,
self.min_epsilon,
)
state = next_state
score += reward
t_end = time.time()
avg_time_cost = (t_end - t_begin) / self.episode_step
if losses:
avg_loss = np.vstack(losses).mean(axis=0)
log_value = (self.i_episode, avg_loss, score, avg_time_cost)
self.write_log(log_value)
if self.i_episode % self.args.save_period == 0:
self.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.save_params(self.i_episode)
self.interim_test()