class BCSACAgent(SACAgent):
"""BC with SAC agent interacting with environment.
Attrtibutes:
her (HER): hinsight experience replay
transitions_epi (list): transitions per episode (for HER)
desired_state (np.ndarray): desired state of current episode
memory (ReplayBuffer): replay memory
demo_memory (ReplayBuffer): replay memory for demo
lambda2 (float): proportion of BC loss
"""
# pylint: disable=attribute-defined-outside-init
def _initialize(self):
"""Initialize non-common things."""
# load demo replay memory
with open(self.hyper_params.demo_path, "rb") as f:
demo = list(pickle.load(f))
# HER
if self.hyper_params.use_her:
self.her = build_her(self.hyper_params.her)
print(f"[INFO] Build {str(self.her)}.")
if self.hyper_params.desired_states_from_demo:
self.her.fetch_desired_states_from_demo(demo)
self.transitions_epi: list = list()
self.desired_state = np.zeros((1,))
demo = self.her.generate_demo_transitions(demo)
if not self.her.is_goal_in_state:
self.state_dim = (self.state_dim[0] * 2,)
else:
self.her = None
if not self.is_test:
# Replay buffers
demo_batch_size = self.hyper_params.demo_batch_size
self.demo_memory = ReplayBuffer(len(demo), demo_batch_size)
self.demo_memory.extend(demo)
self.memory = ReplayBuffer(self.hyper_params.buffer_size, demo_batch_size)
# set hyper parameters
self.hyper_params["lambda2"] = 1.0 / demo_batch_size
build_args = dict(
hyper_params=self.hyper_params,
log_cfg=self.log_cfg,
env_name=self.env_info.name,
state_size=self.env_info.observation_space.shape,
output_size=self.env_info.action_space.shape[0],
is_test=self.is_test,
load_from=self.load_from,
)
self.learner = build_learner(self.learner_cfg, build_args)
def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
"""Preprocess state so that actor selects an action."""
if self.hyper_params.use_her:
self.desired_state = self.her.get_desired_state()
state = np.concatenate((state, self.desired_state), axis=-1)
state = numpy2floattensor(state, self.learner.device)
return state
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
"""Add 1 step and n step transitions to memory."""
if self.hyper_params.use_her:
self.transitions_epi.append(transition)
done = transition[-1] or self.episode_step == self.max_episode_steps
if done:
# insert generated transitions if the episode is done
transitions = self.her.generate_transitions(
self.transitions_epi,
self.desired_state,
self.hyper_params.success_score,
)
self.memory.extend(transitions)
self.transitions_epi.clear()
else:
self.memory.add(transition)
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 n_qf_mask: %d (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
loss[5], # n_qf_mask
avg_time_cost,
)
)
if self.is_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,
}
)
def train(self):
"""Train the agent."""
# logger
if self.is_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.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.is_render and self.i_episode >= self.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):
experience = self.memory.sample()
demos = self.demo_memory.sample()
experience, demo = (
numpy2floattensor(experience, self.learner.device),
numpy2floattensor(demos, self.learner.device),
)
loss = self.learner.update_model(experience, demo)
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.save_period == 0:
self.learner.save_params(self.i_episode)
self.interim_test()
# termination
self.env.close()
self.learner.save_params(self.i_episode)
self.interim_test()
class BCSACAgent2(SACAgent2):
"""BC with SAC agent interacting with environment.
Attrtibutes:
her (HER): hinsight experience replay
transitions_epi (list): transitions per episode (for HER)
desired_state (np.ndarray): desired state of current episode
memory (ReplayBuffer): replay memory
demo_memory (ReplayBuffer): replay memory for demo
lambda2 (float): proportion of BC loss
"""
# pylint: disable=attribute-defined-outside-init
def _initialize(self):
"""Initialize non-common things."""
# load demo replay memory
with open(self.args.demo_path, "rb") as f:
demo = list(pickle.load(f))
# HER
if self.hyper_params.use_her:
self.her = build_her(self.hyper_params.her)
print(f"[INFO] Build {str(self.her)}.")
if self.hyper_params.desired_states_from_demo:
self.her.fetch_desired_states_from_demo(demo)
self.transitions_epi: list = list()
self.desired_state = np.zeros((1, ))
demo = self.her.generate_demo_transitions(demo)
if not self.her.is_goal_in_state:
self.state_dim = (self.state_dim[0] * 2, )
else:
self.her = None
if not self.args.test:
# Replay buffers
demo_batch_size = self.hyper_params.demo_batch_size
self.demo_memory = ReplayBuffer(len(demo), demo_batch_size)
self.demo_memory.extend(demo)
self.memory = ReplayBuffer(self.hyper_params.sac_buffer_size,
demo_batch_size)
# set hyper parameters
self.hyper_params["lambda2"] = 1.0 / demo_batch_size
self.args.cfg_path = self.args.offer_cfg_path
self.args.load_from = self.args.load_offer_from
self.hyper_params.buffer_size = self.hyper_params.sac_buffer_size
self.hyper_params.batch_size = self.hyper_params.sac_batch_size
self.learner_cfg.type = "BCSACLearner"
self.learner_cfg.hyper_params = self.hyper_params
self.learner = build_learner(self.learner_cfg)
del self.hyper_params.buffer_size
del self.hyper_params.batch_size
# init stack
self.stack_size = self.args.stack_size
self.stack_buffer = deque(maxlen=self.args.stack_size)
self.stack_buffer_2 = deque(maxlen=self.args.stack_size)
self.scores = list()
self.utilities = list()
self.rounds = list()
self.opp_utilities = list()
def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
"""Preprocess state so that actor selects an action."""
if self.hyper_params.use_her:
self.desired_state = self.her.get_desired_state()
state = np.concatenate((state, self.desired_state), axis=-1)
state = torch.FloatTensor(state).to(device)
return state
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
"""Add 1 step and n step transitions to memory."""
if self.hyper_params.use_her:
self.transitions_epi.append(transition)
done = transition[
-1] or self.episode_step == self.args.max_episode_steps
if done:
# insert generated transitions if the episode is done
transitions = self.her.generate_transitions(
self.transitions_epi,
self.desired_state,
self.hyper_params.success_score,
)
self.memory.extend(transitions)
self.transitions_epi.clear()
else:
self.memory.add(transition)
def _update_model(self):
# training
if len(self.memory) >= self.hyper_params.sac_batch_size:
for _ in range(self.hyper_params.multiple_update):
experience = self.memory.sample()
demos = self.demo_memory.sample()
experience, demo = (
numpy2floattensor(experience),
numpy2floattensor(demos),
)
loss = self.learner.update_model(experience, demo)
self.loss_episode.append(loss) # for logging
class BCSACAgent(SACAgent):
"""BC with SAC agent interacting with environment.
Attrtibutes:
her (HER): hinsight experience replay
transitions_epi (list): transitions per episode (for HER)
desired_state (np.ndarray): desired state of current episode
memory (ReplayBuffer): replay memory
demo_memory (ReplayBuffer): replay memory for demo
lambda2 (float): proportion of BC loss
"""
# pylint: disable=attribute-defined-outside-init
def _initialize(self):
"""Initialize non-common things."""
# load demo replay memory
with open(self.args.demo_path, "rb") as f:
demo = list(pickle.load(f))
# HER
if self.hyper_params.use_her:
self.her = build_her(self.hyper_params.her)
print(f"[INFO] Build {str(self.her)}.")
if self.hyper_params.desired_states_from_demo:
self.her.fetch_desired_states_from_demo(demo)
self.transitions_epi: list = list()
self.desired_state = np.zeros((1, ))
demo = self.her.generate_demo_transitions(demo)
if not self.her.is_goal_in_state:
self.state_dim = (self.state_dim[0] * 2, )
else:
self.her = None
if not self.args.test:
# Replay buffers
demo_batch_size = self.hyper_params.demo_batch_size
self.demo_memory = ReplayBuffer(len(demo), demo_batch_size)
self.demo_memory.extend(demo)
self.memory = ReplayBuffer(self.hyper_params.buffer_size,
demo_batch_size)
# set hyper parameters
self.lambda2 = 1.0 / demo_batch_size
def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
"""Preprocess state so that actor selects an action."""
if self.hyper_params.use_her:
self.desired_state = self.her.get_desired_state()
state = np.concatenate((state, self.desired_state), axis=-1)
state = torch.FloatTensor(state).to(device)
return state
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
"""Add 1 step and n step transitions to memory."""
if self.hyper_params.use_her:
self.transitions_epi.append(transition)
done = transition[
-1] or self.episode_step == self.args.max_episode_steps
if done:
# insert generated transitions if the episode is done
transitions = self.her.generate_transitions(
self.transitions_epi,
self.desired_state,
self.hyper_params.success_score,
)
self.memory.extend(transitions)
self.transitions_epi.clear()
else:
self.memory.add(transition)
def update_model(self) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
self.update_step += 1
experiences, demos = self.memory.sample(), self.demo_memory.sample()
states, actions, rewards, next_states, dones = experiences
demo_states, demo_actions, _, _, _ = demos
new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states)
pred_actions, _, _, _, _ = self.actor(demo_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()
# update actor
actor_loss = torch.zeros(1)
n_qf_mask = 0
if self.update_step % self.hyper_params.policy_update_freq == 0:
# bc loss
qf_mask = torch.gt(
self.qf_1(torch.cat((demo_states, demo_actions), dim=-1)),
self.qf_1(torch.cat((demo_states, pred_actions), dim=-1)),
).to(device)
qf_mask = qf_mask.float()
n_qf_mask = int(qf_mask.sum().item())
if n_qf_mask == 0:
bc_loss = torch.zeros(1, device=device)
else:
bc_loss = (torch.mul(pred_actions, qf_mask) - torch.mul(
demo_actions, qf_mask)).pow(2).sum() / n_qf_mask
# actor loss
advantage = q_pred - v_pred.detach()
actor_loss = (alpha * log_prob - advantage).mean()
actor_loss = self.hyper_params.lambda1 * actor_loss + self.lambda2 * bc_loss
# 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)
return (
actor_loss.item(),
qf_1_loss.item(),
qf_2_loss.item(),
vf_loss.item(),
alpha_loss.item(),
n_qf_mask,
)
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 n_qf_mask: %d (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
loss[5], # n_qf_mask
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,
})
class BCDDPGAgent(DDPGAgent):
"""BC with DDPG agent interacting with environment.
Attributes:
her (HER): hinsight experience replay
transitions_epi (list): transitions per episode (for HER)
desired_state (np.ndarray): desired state of current episode
memory (ReplayBuffer): replay memory
demo_memory (ReplayBuffer): replay memory for demo
lambda2 (float): proportion of BC loss
"""
# pylint: disable=attribute-defined-outside-init
def _initialize(self):
"""Initialize non-common things."""
# load demo replay memory
with open(self.args.demo_path, "rb") as f:
demo = list(pickle.load(f))
# HER
if self.hyper_params.use_her:
self.her = build_her(self.hyper_params.her)
print(f"[INFO] Build {str(self.her)}.")
if self.hyper_params.desired_states_from_demo:
self.her.fetch_desired_states_from_demo(demo)
self.transitions_epi: list = list()
self.desired_state = np.zeros((1, ))
demo = self.her.generate_demo_transitions(demo)
if not self.her.is_goal_in_state:
self.state_dim = (self.state_dim[0] * 2, )
else:
self.her = None
if not self.args.test:
# Replay buffers
demo_batch_size = self.hyper_params.demo_batch_size
self.demo_memory = ReplayBuffer(len(demo), demo_batch_size)
self.demo_memory.extend(demo)
self.memory = ReplayBuffer(self.hyper_params.buffer_size,
self.hyper_params.batch_size)
# set hyper parameters
self.lambda2 = 1.0 / demo_batch_size
def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
"""Preprocess state so that actor selects an action."""
if self.hyper_params.use_her:
self.desired_state = self.her.get_desired_state()
state = np.concatenate((state, self.desired_state), axis=-1)
state = torch.FloatTensor(state).to(device)
return state
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
"""Add 1 step and n step transitions to memory."""
if self.hyper_params.use_her:
self.transitions_epi.append(transition)
done = transition[
-1] or self.episode_step == self.args.max_episode_steps
if done:
# insert generated transitions if the episode is done
transitions = self.her.generate_transitions(
self.transitions_epi,
self.desired_state,
self.hyper_params.success_score,
)
self.memory.extend(transitions)
self.transitions_epi.clear()
else:
self.memory.add(transition)
def update_model(self) -> Tuple[torch.Tensor, ...]:
"""Train the model after each episode."""
experiences = self.memory.sample()
demos = self.demo_memory.sample()
exp_states, exp_actions, exp_rewards, exp_next_states, exp_dones = experiences
demo_states, demo_actions, demo_rewards, demo_next_states, demo_dones = demos
states = torch.cat((exp_states, demo_states), dim=0)
actions = torch.cat((exp_actions, demo_actions), dim=0)
rewards = torch.cat((exp_rewards, demo_rewards), dim=0)
next_states = torch.cat((exp_next_states, demo_next_states), dim=0)
dones = torch.cat((exp_dones, demo_dones), dim=0)
# 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)
# critic loss
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)
# train critic
self.critic_optim.zero_grad()
critic_loss.backward()
nn.utils.clip_grad_norm_(self.critic.parameters(), gradient_clip_cr)
self.critic_optim.step()
# policy loss
actions = self.actor(states)
policy_loss = -self.critic(torch.cat((states, actions), dim=-1)).mean()
# bc loss
pred_actions = self.actor(demo_states)
qf_mask = torch.gt(
self.critic(torch.cat((demo_states, demo_actions), dim=-1)),
self.critic(torch.cat((demo_states, pred_actions), dim=-1)),
).to(device)
qf_mask = qf_mask.float()
n_qf_mask = int(qf_mask.sum().item())
if n_qf_mask == 0:
bc_loss = torch.zeros(1, device=device)
else:
bc_loss = (torch.mul(pred_actions, qf_mask) - torch.mul(
demo_actions, qf_mask)).pow(2).sum() / n_qf_mask
# train actor: pg loss + BC loss
actor_loss = self.hyper_params.lambda1 * policy_loss + self.lambda2 * bc_loss
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(), n_qf_mask
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, n_qf_mask: %d "
"(spent %.6f sec/step)\n" % (
i,
self.episode_step,
self.total_step,
score,
total_loss,
loss[0],
loss[1],
loss[2],
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,
})