class Agent():
def __init__(self, state_size, action_size, config):
self.seed = config["seed"]
torch.manual_seed(self.seed)
np.random.seed(seed=self.seed)
random.seed(self.seed)
env = gym.make(config["env_name"])
self.env = FrameStack(env, config)
self.env.seed(self.seed)
self.state_size = state_size
self.action_size = action_size
self.clip = config["clip"]
self.device = 'cuda'
self.double_dqn = config["DDQN"]
self.lr_pre = config["lr_pre"]
self.batch_size = config["batch_size"]
self.lr = config["lr"]
self.tau = config["tau"]
self.gamma = 0.99
self.fc1 = config["fc1_units"]
self.fc2 = config["fc2_units"]
self.qnetwork_local = QNetwork(state_size, action_size, self.fc1,
self.fc2, self.seed).to(self.device)
self.qnetwork_target = QNetwork(state_size, action_size, self.fc1,
self.fc2, self.seed).to(self.device)
self.optimizer = optim.Adam(self.qnetwork_local.parameters(),
lr=self.lr)
self.soft_update(self.qnetwork_local, self.qnetwork_target, 1)
self.q_shift_local = QNetwork(state_size, action_size, self.fc1,
self.fc2, self.seed).to(self.device)
self.q_shift_target = QNetwork(state_size, action_size, self.fc1,
self.fc2, self.seed).to(self.device)
self.optimizer_shift = optim.Adam(self.q_shift_local.parameters(),
lr=self.lr)
self.soft_update(self.q_shift_local, self.q_shift_target, 1)
self.R_local = QNetwork(state_size, action_size, self.fc1, self.fc2,
self.seed).to(self.device)
self.R_target = QNetwork(state_size, action_size, self.fc1, self.fc2,
self.seed).to(self.device)
self.optimizer_r = optim.Adam(self.R_local.parameters(), lr=self.lr)
self.soft_update(self.R_local, self.R_target, 1)
self.steps = 0
self.predicter = QNetwork(state_size, action_size, self.fc1, self.fc2,
self.seed).to(self.device)
self.optimizer_pre = optim.Adam(self.predicter.parameters(),
lr=self.lr_pre)
#self.encoder_freq = Encoder(config).to(self.device)
#self.encoder_optimizer_frq = torch.optim.Adam(self.encoder_freq.parameters(), self.lr)
self.encoder = Encoder(config).to(self.device)
self.encoder_optimizer = torch.optim.Adam(self.encoder.parameters(),
self.lr)
pathname = "lr_{}_batch_size_{}_fc1_{}_fc2_{}_seed_{}".format(
self.lr, self.batch_size, self.fc1, self.fc2, self.seed)
pathname += "_clip_{}".format(config["clip"])
pathname += "_tau_{}".format(config["tau"])
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H:%M:%S")
pathname += dt_string
tensorboard_name = str(config["locexp"]) + '/runs/' + pathname
self.vid_path = str(config["locexp"]) + '/vid'
self.writer = SummaryWriter(tensorboard_name)
self.average_prediction = deque(maxlen=100)
self.average_same_action = deque(maxlen=100)
self.all_actions = []
for a in range(self.action_size):
action = torch.Tensor(1) * 0 + a
self.all_actions.append(action.to(self.device))
def learn(self, memory_ex):
logging.debug(
"--------------------------New update-----------------------------------------------"
)
self.steps += 1
states, next_states, actions, dones = memory_ex.expert_policy(
self.batch_size)
states = states.type(torch.float32).div_(255)
states = self.encoder.create_vector(states)
next_states = next_states.type(torch.float32).div_(255)
next_states = self.encoder.create_vector(next_states)
self.state_action_frq(states, actions)
actions = torch.randint(0,
3, (self.batch_size, 1),
dtype=torch.int64,
device=self.device)
self.compute_shift_function(states.detach(), next_states, actions,
dones)
self.compute_r_function(states.detach(), actions)
self.compute_q_function(states.detach(), next_states, actions, dones)
self.soft_update(self.R_local, self.R_target, self.tau)
self.soft_update(self.q_shift_local, self.q_shift_target, self.tau)
self.soft_update(self.qnetwork_local, self.qnetwork_target, self.tau)
return
def compute_q_function(self, states, next_states, actions, dones):
"""Update value parameters using given batch of experience tuples. """
actions = actions.type(torch.int64)
# Get max predicted Q values (for next states) from target model
if self.double_dqn:
q_values = self.qnetwork_local(next_states).detach()
_, best_action = q_values.max(1)
best_action = best_action.unsqueeze(1)
Q_targets_next = self.qnetwork_target(next_states).detach()
Q_targets_next = Q_targets_next.gather(1, best_action)
else:
Q_targets_next = self.qnetwork_target(next_states).detach().max(
1)[0].unsqueeze(1)
# Compute Q targets for current states
# Get expected Q values from local model
# Compute loss
rewards = self.R_target(states).detach().gather(
1, actions.detach()).squeeze(0)
Q_targets = rewards + (self.gamma * Q_targets_next * (dones))
Q_expected = self.qnetwork_local(states).gather(1, actions)
loss = F.mse_loss(Q_expected, Q_targets.detach())
# Get max predicted Q values (for next states) from target model
self.writer.add_scalar('Q_loss', loss, self.steps)
# Minimize the loss
self.optimizer.zero_grad()
loss.backward()
self.encoder_optimizer.step()
# torch.nn.utils.clip_grad_norm_(self.qnetwork_local.parameters(), 1)
self.optimizer.step()
def compute_shift_function(self, states, next_states, actions, dones):
"""Update Q shift parameters using given batch of experience tuples """
actions = actions.type(torch.int64)
with torch.no_grad():
# Get max predicted Q values (for next states) from target model
if self.double_dqn:
q_shift = self.q_shift_local(next_states)
max_q, max_actions = q_shift.max(1)
Q_targets_next = self.qnetwork_target(next_states).gather(
1, max_actions.unsqueeze(1))
else:
Q_targets_next = self.qnetwork_target(
next_states).detach().max(1)[0].unsqueeze(1)
# Compute Q targets for current states
Q_targets = self.gamma * Q_targets_next
# Get expected Q values from local model
Q_expected = self.q_shift_local(states).gather(1, actions)
# Compute loss
loss = F.mse_loss(Q_expected, Q_targets.detach())
# Minimize the loss
self.optimizer_shift.zero_grad()
loss.backward()
self.writer.add_scalar('Shift_loss', loss, self.steps)
self.optimizer_shift.step()
def compute_r_function(self, states, actions, debug=False, log=False):
""" compute reward for the state action pair """
actions = actions.type(torch.int64)
# sum all other actions
size = states.shape[0]
idx = 0
all_zeros = [1 for i in range(actions.shape[0])]
zeros = False
y_shift = self.q_shift_target(states).gather(1, actions).detach()
log_a = self.get_action_prob(states, actions).detach()
y_r_part1 = log_a - y_shift
y_r_part2 = torch.empty((size, 1), dtype=torch.float32).to(self.device)
for a, s in zip(actions, states):
y_h = 0
taken_actions = 0
for b in self.all_actions:
b = b.type(torch.int64).unsqueeze(1)
n_b = self.get_action_prob(s.unsqueeze(0), b)
if torch.eq(a, b) or n_b is None:
continue
taken_actions += 1
y_s = self.q_shift_target(s.unsqueeze(0)).detach().gather(
1, b).item()
n_b = n_b.data.item() - y_s
r_hat = self.R_target(s.unsqueeze(0)).gather(1, b).item()
y_h += (r_hat - n_b)
if log:
text = "a {} r _hat {:.2f} - n_b {:.2f} | sh {:.2f} ".format(
b.item(), r_hat, n_b, y_s)
logging.debug(text)
if taken_actions == 0:
all_zeros[idx] = 0
zeros = True
y_r_part2[idx] = 0.0
else:
y_r_part2[idx] = (1. / taken_actions) * y_h
idx += 1
y_r = y_r_part1 + y_r_part2
# check if there are zeros (no update for this tuble) remove them from states and
if zeros:
mask = torch.BoolTensor(all_zeros)
states = states[mask]
actions = actions[mask]
y_r = y_r[mask]
y = self.R_local(states).gather(1, actions)
if log:
text = "Action {:.2f} r target {:.2f} = n_a {:.2f} + n_b {:.2f} y {:.2f}".format(
actions[0].item(), y_r[0].item(), y_r_part1[0].item(),
y_r_part2[0].item(), y[0].item())
logging.debug(text)
r_loss = F.mse_loss(y, y_r.detach())
# Minimize the loss
self.optimizer_r.zero_grad()
r_loss.backward()
# torch.nn.utils.clip_grad_norm_(self.R_local.parameters(), 5)
self.optimizer_r.step()
self.writer.add_scalar('Reward_loss', r_loss, self.steps)
def get_action_prob(self, states, actions):
""" compute prob for state action pair """
actions = actions.type(torch.long)
# check if action prob is zero
output = self.predicter(states)
output = F.softmax(output, dim=1)
action_prob = output.gather(1, actions)
action_prob = action_prob + torch.finfo(torch.float32).eps
# check if one action if its to small
if action_prob.shape[0] == 1:
if action_prob.cpu().detach().numpy()[0][0] < 1e-4:
return None
action_prob = torch.log(action_prob)
action_prob = torch.clamp(action_prob, min=self.clip, max=0)
return action_prob
def state_action_frq(self, states, action):
""" Train classifer to compute state action freq """
self.predicter.train()
output = self.predicter(states)
output = output.squeeze(0)
y = action.type(torch.long).squeeze(1)
loss = nn.CrossEntropyLoss()(output, y)
self.optimizer_pre.zero_grad()
self.encoder_optimizer.zero_grad()
loss.backward()
# torch.nn.utils.clip_grad_norm_(self.predicter.parameters(), 1)
self.optimizer_pre.step()
self.writer.add_scalar('Predict_loss', loss, self.steps)
self.predicter.eval()
def test_predicter(self, memory):
""" Test the classifier """
self.predicter.eval()
same_state_predition = 0
for i in range(memory.idx):
states = memory.obses[i]
actions = memory.actions[i]
states = torch.as_tensor(states, device=self.device).unsqueeze(0)
states = states.type(torch.float32).div_(255)
states = self.encoder.create_vector(states)
actions = torch.as_tensor(actions, device=self.device)
output = self.predicter(states)
output = F.softmax(output, dim=1)
# create one hot encode y from actions
y = actions.type(torch.long).item()
p = torch.argmax(output.data).item()
if y == p:
same_state_predition += 1
text = "Same prediction {} of {} ".format(same_state_predition,
memory.idx)
print(text)
logging.debug(text)
def soft_update(self, local_model, target_model, tau=4):
"""Soft update model parameters.
θ_target = τ*θ_local + (1 - τ)*θ_target
Params
======
local_model (PyTorch model): weights will be copied from
target_model (PyTorch model): weights will be copied to
tau (float): interpolation parameter
"""
# print("use tau", tau)
for target_param, local_param in zip(target_model.parameters(),
local_model.parameters()):
target_param.data.copy_(tau * local_param.data +
(1.0 - tau) * target_param.data)
def load(self, filename):
self.predicter.load_state_dict(torch.load(filename + "_predicter.pth"))
self.optimizer_pre.load_state_dict(
torch.load(filename + "_predicter_optimizer.pth"))
self.R_local.load_state_dict(torch.load(filename + "_r_net.pth"))
self.qnetwork_local.load_state_dict(torch.load(filename +
"_q_net.pth"))
print("Load models to {}".format(filename))
def save(self, filename):
"""
"""
mkdir("", filename)
torch.save(self.predicter.state_dict(), filename + "_predicter.pth")
torch.save(self.optimizer_pre.state_dict(),
filename + "_predicter_optimizer.pth")
torch.save(self.qnetwork_local.state_dict(), filename + "_q_net.pth")
torch.save(self.optimizer.state_dict(),
filename + "_q_net_optimizer.pth")
torch.save(self.R_local.state_dict(), filename + "_r_net.pth")
torch.save(self.q_shift_local.state_dict(),
filename + "_q_shift_net.pth")
print("save models to {}".format(filename))
def test_q_value(self, memory):
test_elements = memory.idx
test_elements = 100
all_diff = 0
error = True
used_elements_r = 0
used_elements_q = 0
r_error = 0
q_error = 0
for i in range(test_elements):
states = memory.obses[i]
actions = memory.actions[i]
states = torch.as_tensor(states, device=self.device).unsqueeze(0)
states = states.type(torch.float32).div_(255)
states = self.encoder.create_vector(states)
actions = torch.as_tensor(actions, device=self.device)
one_hot = torch.Tensor([0 for i in range(self.action_size)],
device="cpu")
one_hot[actions.item()] = 1
with torch.no_grad():
r_values = self.R_local(states.detach()).detach()
q_values = self.qnetwork_local(states.detach()).detach()
soft_r = F.softmax(r_values, dim=1).to("cpu")
soft_q = F.softmax(q_values, dim=1).to("cpu")
actions = actions.type(torch.int64)
kl_q = F.kl_div(soft_q.log(), one_hot, None, None, 'sum')
kl_r = F.kl_div(soft_r.log(), one_hot, None, None, 'sum')
if kl_r == float("inf"):
pass
else:
r_error += kl_r
used_elements_r += 1
if kl_q == float("inf"):
pass
else:
q_error += kl_q
used_elements_q += 1
average_q_kl = q_error / used_elements_q
average_r_kl = r_error / used_elements_r
text = "Kl div of Reward {} of {} elements".format(
average_q_kl, used_elements_r)
print(text)
text = "Kl div of Q_values {} of {} elements".format(
average_r_kl, used_elements_q)
print(text)
self.writer.add_scalar('KL_reward', average_r_kl, self.steps)
self.writer.add_scalar('KL_q_values', average_q_kl, self.steps)
def act(self, states):
states = torch.as_tensor(states, device=self.device).unsqueeze(0)
states = states.type(torch.float32).div_(255)
states = self.encoder.create_vector(states)
q_values = self.qnetwork_local(states.detach()).detach()
action = torch.argmax(q_values).item()
return action
def eval_policy(self, record=False, eval_episodes=2):
if record:
env = wrappers.Monitor(self.env,
str(self.vid_path) +
"/{}".format(self.steps),
video_callable=lambda episode_id: True,
force=True)
else:
env = self.env
average_reward = 0
scores_window = deque(maxlen=100)
s = 0
for i_epiosde in range(eval_episodes):
episode_reward = 0
state = env.reset()
while True:
s += 1
action = self.act(state)
state, reward, done, _ = env.step(action)
episode_reward += reward
if done:
break
scores_window.append(episode_reward)
if record:
return
average_reward = np.mean(scores_window)
print("Eval Episode {} average Reward {} ".format(
eval_episodes, average_reward))
self.writer.add_scalar('Eval_reward', average_reward, self.steps)
class TD3():
def __init__(self, action_size, state_size, config):
self.seed = config["seed"]
print("TD3 seed", self.seed)
torch.manual_seed(self.seed)
np.random.seed(seed=self.seed)
random.seed(self.seed)
self.env = gym.make(config["env_name"])
self.env.seed(self.seed)
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H:%M:%S")
self.env.action_space.seed(self.seed)
self.action_size = action_size
self.state_size = state_size
self.min_action = config["min_action"]
self.max_action = config["max_action"]
self.seed = config["seed"]
self.tau = config["tau"]
self.gamma = config["gamma"]
self.batch_size = config["batch_size"]
if not torch.cuda.is_available():
config["device"] == "cpu"
self.device = config["device"]
self.eval = config["eval"]
self.vid_path = config["vid_path"]
print("actions size ", action_size)
print("actions min ", self.min_action)
print("actions max ", self.max_action)
fc1 = config["fc1_units"]
fc2 = config["fc2_units"]
self.actor = Actor(state_size, action_size, self.seed, fc1, fc2).to(self.device)
self.optimizer_a = torch.optim.Adam(self.actor.parameters(), config["lr_actor"])
self.target_actor = Actor(state_size, action_size, self.seed, fc1, fc2).to(self.device)
self.target_actor.load_state_dict(self.actor.state_dict())
self.critic = QNetwork(state_size, action_size, self.seed, fc1, fc2).to(self.device)
self.optimizer_q = torch.optim.Adam(self.critic.parameters(), config["lr_critic"])
self.target_critic = QNetwork(state_size, action_size, self.seed, fc1, fc2).to(self.device)
self.target_critic.load_state_dict(self.critic.state_dict())
self.max_timesteps = config["max_episodes_steps"]
self.episodes = config["episodes"]
self.memory = ReplayBuffer((state_size, ), (action_size, ), config["buffer_size"], self.seed, self.device)
pathname = str(config["seed"]) + str(dt_string)
tensorboard_name = str(config["res_path"]) + '/runs/'+ "TD3" + str(pathname)
self.writer = SummaryWriter(tensorboard_name)
self.steps= 0
self.actor_freq = config["actor_freq"]
self.policy_noise = config["policy_noise"]
self.noise_clip = config["noise_clip"]
self.expl_noise = config["exp_noise"]
def act(self, state):
state = torch.as_tensor(state, dtype=torch.float32, device=self.device)
action = self.actor(state.unsqueeze(0))
actions = action.detach().cpu().numpy()[0]
actions = actions + + np.random.normal(0, self.max_action * self.expl_noise, size=self.action_size)
actions = np.clip(actions, self.min_action, self.max_action)
return actions
def act_greedy(self, state):
state = torch.as_tensor(state, dtype=torch.float32, device=self.device)
action = self.actor(state.unsqueeze(0))
actions = action.detach().cpu().numpy()[0]
actions = np.clip(actions, self.min_action, self.max_action)
return actions
def train_agent(self):
average_reward = 0
scores_window = deque(maxlen=100)
s = 0
t0 = time.time()
for i_epiosde in range(self.episodes):
episode_reward = 0
state = self.env.reset()
for t in range(self.max_timesteps):
s += 1
action = self.act(state)
next_state, reward, done, _ = self.env.step(action)
episode_reward += reward
if i_epiosde > 10:
self.learn()
self.memory.add(state, reward, action, next_state, done)
state = next_state
if done:
scores_window.append(episode_reward)
break
if i_epiosde % self.eval == 0:
self.eval_policy()
ave_reward = np.mean(scores_window)
print("Epiosde {} Steps {} Reward {} Reward averge{} Time {}".format(i_epiosde, t, episode_reward, np.mean(scores_window), time_format(time.time() - t0)))
self.writer.add_scalar('Aver_reward', ave_reward, self.steps)
self.writer.add_scalar('steps_in_episode', t, self.steps)
def learn(self):
self.steps += 1
states, rewards, actions, next_states, dones = self.memory.sample(self.batch_size)
with torch.no_grad():
next_action = self.target_actor(next_states)
noise = (torch.randn_like(actions) * self.policy_noise).clamp(-self.noise_clip, self.noise_clip)
print(noise)
next_action = (next_action + noise).clamp(-self.max_action, self.max_action)
q1_target, q2_target = self.target_critic(next_states, next_action)
q_target = torch.min(q1_target, q2_target)
q_target = rewards + (self.gamma * q_target * (1 - dones))
q_pre1, q_pre2 = self.critic(states, actions)
loss = F.mse_loss(q_pre1, q_target) + F.mse_loss(q_pre2, q_target)
self.writer.add_scalar('Q_loss', loss, self.steps)
self.optimizer_q.zero_grad()
loss.backward()
self.optimizer_q.step()
# delay actor update
if self.steps % self.actor_freq == 0:
#-------------------------------update-actor-------------------------------------------------
actor_actions = self.actor(states)
q_values = self.critic.Q1(states, actor_actions)
loss_actor = -q_values.mean()
self.optimizer_a.zero_grad()
loss_actor.backward()
self.writer.add_scalar('Actor_loss', loss_actor, self.steps)
self.optimizer_a.step()
#-------------------------------update-networks-------------------------------------------------
self.soft_udapte(self.critic, self.target_critic)
self.soft_udapte(self.actor, self.target_actor)
def soft_udapte(self, online, target):
for param, target_parm in zip(online.parameters(), target.parameters()):
target_parm.data.copy_(self.tau * param.data + (1 - self.tau) * target_parm.data)
def eval_policy(self, eval_episodes=4):
env = wrappers.Monitor(self.env, str(self.vid_path) + "/{}".format(self.steps), video_callable=lambda episode_id: True,force=True)
average_reward = 0
scores_window = deque(maxlen=100)
for i_epiosde in range(eval_episodes):
print("Eval Episode {} of {} ".format(i_epiosde, self.episodes))
episode_reward = 0
state = env.reset()
while True:
action = self.act_greedy(state)
state, reward, done, _ = env.step(action)
episode_reward += reward
if done:
scores_window.append(episode_reward)
break
average_reward = np.mean(scores_window)
self.writer.add_scalar('Eval_reward', average_reward, self.steps)
class SACAgent():
def __init__(self, action_size, state_size, config):
self.seed = config["seed"]
torch.manual_seed(self.seed)
np.random.seed(seed=self.seed)
self.env = gym.make(config["env_name"])
self.env = FrameStack(self.env, config)
self.env.seed(self.seed)
self.action_size = action_size
self.state_size = state_size
self.tau = config["tau"]
self.gamma = config["gamma"]
self.batch_size = config["batch_size"]
self.lr = config["lr"]
self.history_length = config["history_length"]
self.size = config["size"]
if not torch.cuda.is_available():
config["device"] == "cpu"
self.device = config["device"]
self.eval = config["eval"]
self.vid_path = config["vid_path"]
print("actions size ", action_size)
self.critic = QNetwork(state_size, action_size, config["fc1_units"],
config["fc2_units"]).to(self.device)
self.q_optim = torch.optim.Adam(self.critic.parameters(),
config["lr_critic"])
self.target_critic = QNetwork(state_size, action_size,
config["fc1_units"],
config["fc2_units"]).to(self.device)
self.target_critic.load_state_dict(self.critic.state_dict())
self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
self.alpha = self.log_alpha.exp()
self.alpha_optim = Adam([self.log_alpha], lr=config["lr_alpha"])
self.policy = SACActor(state_size, action_size).to(self.device)
self.policy_optim = Adam(self.policy.parameters(),
lr=config["lr_policy"])
self.encoder = Encoder(config).to(self.device)
self.encoder_optimizer = torch.optim.Adam(self.encoder.parameters(),
self.lr)
self.episodes = config["episodes"]
self.memory = ReplayBuffer((self.history_length, self.size, self.size),
(1, ), config["buffer_size"],
config["image_pad"], self.seed, self.device)
pathname = config["seed"]
tensorboard_name = str(config["res_path"]) + '/runs/' + str(pathname)
self.writer = SummaryWriter(tensorboard_name)
self.steps = 0
self.target_entropy = -torch.prod(
torch.Tensor(action_size).to(self.device)).item()
def act(self, state, evaluate=False):
with torch.no_grad():
state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
state = state.type(torch.float32).div_(255)
self.encoder.eval()
state = self.encoder.create_vector(state)
self.encoder.train()
if evaluate is False:
action = self.policy.sample(state)
else:
action_prob, _ = self.policy(state)
action = torch.argmax(action_prob)
action = action.cpu().numpy()
return action
# action = np.clip(action, self.min_action, self.max_action)
action = action.cpu().numpy()[0]
return action
def train_agent(self):
average_reward = 0
scores_window = deque(maxlen=100)
t0 = time.time()
for i_epiosde in range(1, self.episodes):
episode_reward = 0
state = self.env.reset()
t = 0
while True:
t += 1
action = self.act(state)
next_state, reward, done, _ = self.env.step(action)
episode_reward += reward
if i_epiosde > 10:
self.learn()
self.memory.add(state, reward, action, next_state, done)
state = next_state
if done:
scores_window.append(episode_reward)
break
if i_epiosde % self.eval == 0:
self.eval_policy()
ave_reward = np.mean(scores_window)
print("Epiosde {} Steps {} Reward {} Reward averge{:.2f} Time {}".
format(i_epiosde, t, episode_reward, np.mean(scores_window),
time_format(time.time() - t0)))
self.writer.add_scalar('Aver_reward', ave_reward, self.steps)
def learn(self):
self.steps += 1
states, rewards, actions, next_states, dones = self.memory.sample(
self.batch_size)
states = states.type(torch.float32).div_(255)
states = self.encoder.create_vector(states)
states_detached = states.detach()
qf1, qf2 = self.critic(states)
q_value1 = qf1.gather(1, actions)
q_value2 = qf2.gather(1, actions)
with torch.no_grad():
next_states = next_states.type(torch.float32).div_(255)
next_states = self.encoder.create_vector(next_states)
q1_target, q2_target = self.target_critic(next_states)
min_q_target = torch.min(q1_target, q2_target)
next_action_prob, next_action_log_prob = self.policy(next_states)
next_q_target = (
next_action_prob *
(min_q_target - self.alpha * next_action_log_prob)).sum(
dim=1, keepdim=True)
next_q_value = rewards + (1 - dones) * self.gamma * next_q_target
# --------------------------update-q--------------------------------------------------------
loss = F.mse_loss(q_value1, next_q_value) + F.mse_loss(
q_value2, next_q_value)
self.q_optim.zero_grad()
self.encoder_optimizer.zero_grad()
loss.backward()
self.q_optim.step()
self.encoder_optimizer.zero_grad()
self.writer.add_scalar('loss/q', loss, self.steps)
# --------------------------update-policy--------------------------------------------------------
action_prob, log_action_prob = self.policy(states_detached)
with torch.no_grad():
q_pi1, q_pi2 = self.critic(states_detached)
min_q_values = torch.min(q_pi1, q_pi2)
#policy_loss = (action_prob * ((self.alpha * log_action_prob) - min_q_values).detach()).sum(dim=1).mean()
policy_loss = (action_prob *
((self.alpha * log_action_prob) - min_q_values)).sum(
dim=1).mean()
self.policy_optim.zero_grad()
policy_loss.backward()
self.policy_optim.step()
self.writer.add_scalar('loss/policy', policy_loss, self.steps)
# --------------------------update-alpha--------------------------------------------------------
alpha_loss = (action_prob.detach() *
(-self.log_alpha *
(log_action_prob + self.target_entropy).detach())).sum(
dim=1).mean()
self.alpha_optim.zero_grad()
alpha_loss.backward()
self.alpha_optim.step()
self.writer.add_scalar('loss/alpha', alpha_loss, self.steps)
self.soft_udapte(self.critic, self.target_critic)
self.alpha = self.log_alpha.exp()
def soft_udapte(self, online, target):
for param, target_parm in zip(online.parameters(),
target.parameters()):
target_parm.data.copy_(self.tau * param.data +
(1 - self.tau) * target_parm.data)
def eval_policy(self, eval_episodes=4):
env = gym.make(self.env_name)
env = wrappers.Monitor(env,
str(self.vid_path) + "/{}".format(self.steps),
video_callable=lambda episode_id: True,
force=True)
average_reward = 0
scores_window = deque(maxlen=100)
for i_epiosde in range(eval_episodes):
print("Eval Episode {} of {} ".format(i_epiosde, eval_episodes))
episode_reward = 0
state = env.reset()
while True:
action = self.act(state, evaluate=True)
state, reward, done, _ = env.step(action)
episode_reward += reward
if done:
break
scores_window.append(episode_reward)
average_reward = np.mean(scores_window)
self.writer.add_scalar('Eval_reward', average_reward, self.steps)
def train(args):
chrome_driver_path = args.chrome_driver_path
checkpoint_path = args.checkpoint_path
nb_actions = args.nb_actions
initial_epsilon = args.initial_epsilon
epsilon = initial_epsilon
final_epsilon = args.final_epsilon
gamma = args.gamma
nb_memory = args.nb_memory
nb_expolre = args.nb_expolre
is_debug = args.is_debug
batch_size = args.batch_size
nb_observation = args.nb_observation
desired_fps = args.desired_fps
is_cuda = True if args.use_cuda and torch.cuda.is_available() else False
log_frequency = args.log_frequency
save_frequency = args.save_frequency
ratio_of_win = args.ratio_of_win
if args.exploiting:
nb_observation = -1
epsilon = final_epsilon
seed = 22
np.random.seed(seed)
memory = deque()
env = DinoSeleniumEnv(chrome_driver_path, speed=args.game_speed)
agent = Agent(env)
game_state = GameState(agent, debug=is_debug)
qnetwork = QNetwork(nb_actions)
if is_cuda:
qnetwork.cuda()
optimizer = torch.optim.Adam(qnetwork.parameters(), 1e-4)
tmp_param = next(qnetwork.parameters())
try:
m = torch.load(checkpoint_path)
qnetwork.load_state_dict(m["qnetwork"])
optimizer.load_state_dict(m["optimizer"])
except:
logger.warn("No model found in {}".format(checkpoint_path))
loss_fcn = torch.nn.MSELoss()
action_indx = 0 # do nothing as the first action
screen, reward, is_gameover, score = game_state.get_state(action_indx)
current_state = np.expand_dims(screen, 0)
# [IMAGE_CHANNELS,IMAGE_WIDTH,IMAGE_HEIGHT]
current_state = np.tile(current_state, (IMAGE_CHANNELS, 1, 1))
initial_state = current_state
t = 0
last_time = 0
sum_scores = 0
total_loss = 0
max_score = 0
qvalues = np.array([0, 0])
lost_action = []
win_actions = []
action_random = 0
action_greedy = 0
episodes = 0
nb_episodes = 0
if not args.exploiting:
try:
t, memory, epsilon, nb_episodes = pickle.load(open(
"cache.p", "rb"))
except:
logger.warn("Could not load cache file! Starting from scratch.")
try:
while True:
qnetwork.eval()
if np.random.random() < epsilon: # epsilon greedy
action_indx = np.random.randint(nb_actions)
action_random += 1
else:
action_greedy += 1
tensor = torch.from_numpy(current_state).float().unsqueeze(0)
with torch.no_grad():
qvalues = qnetwork(tensor).squeeze()
_, action_indx = qvalues.max(-1)
action_indx = action_indx.item()
if epsilon > final_epsilon and t > nb_observation:
epsilon -= (initial_epsilon - final_epsilon) / nb_expolre
screen, reward, is_gameover, score = game_state.get_state(
action_indx)
if is_gameover:
episodes += 1
nb_episodes += 1
lost_action.append(action_indx)
sum_scores += score
else:
win_actions.append(action_indx)
if score > max_score:
max_score = score
if last_time:
fps = 1 / (time.time() - last_time)
if fps > desired_fps:
time.sleep(1 / desired_fps - 1 / fps)
if last_time and t % log_frequency == 0:
logger.info('fps: {0}'.format(1 / (time.time() - last_time)))
last_time = time.time()
screen = np.expand_dims(screen, 0)
next_state = np.append(screen,
current_state[:IMAGE_CHANNELS - 1, :, :],
axis=0)
if not args.exploiting and (is_gameover
or np.random.random() < ratio_of_win):
memory.append((current_state, action_indx, reward, next_state,
is_gameover))
if len(memory) > nb_memory:
memory.popleft()
if nb_observation > 0 and t > nb_observation:
indxes = np.random.choice(len(memory),
batch_size,
replace=False)
minibatch = [memory[b] for b in indxes]
inputs = tmp_param.new(batch_size, IMAGE_CHANNELS, IMAGE_WIDTH,
IMAGE_HEIGHT).zero_()
targets = tmp_param.new(batch_size, nb_actions).zero_()
for i, (state_t, action_t, reward_t, state_t1,
is_gameover_t1) in enumerate(minibatch):
inputs[i] = torch.from_numpy(state_t).float()
tensor = inputs[i].unsqueeze(0)
with torch.no_grad():
qvalues = qnetwork(tensor).squeeze()
targets[i] = qvalues
if is_gameover_t1:
assert reward_t == -1
targets[i, action_t] = reward_t
else:
tensor = torch.from_numpy(state_t1).float().unsqueeze(
0)
with torch.no_grad():
qvalues = qnetwork(tensor).squeeze()
qvalues = qvalues.cpu().numpy()
targets[i, action_t] = reward_t + gamma * qvalues.max()
qnetwork.train()
qnetwork.zero_grad()
q_values = qnetwork(inputs)
loss = loss_fcn(q_values, targets)
loss.backward()
optimizer.step()
total_loss += loss.item()
current_state = initial_state if is_gameover else next_state
t += 1
if t % log_frequency == 0:
logger.info(
"For t {}: mean score is {} max score is {} mean loss: {} number of episode: {}"
.format(t, sum_scores / (episodes + 0.1), max_score,
total_loss / 1000, episodes))
logger.info(
"t: {} action_index: {} reward: {} max qvalue: {} total number of eposodes so far: {}"
.format(t, action_indx, reward, qvalues.max(),
nb_episodes))
tmp = np.array(lost_action)
dnc = (tmp == 0).sum()
logger.info(
"Lost actions do_nothing: {} jump: {} length of memory {}".
format(dnc,
len(tmp) - dnc, len(memory)))
tmp = np.array(win_actions)
dnc = (tmp == 0).sum()
logger.info("Win actions do_nothing: {} jump: {}".format(
dnc,
len(tmp) - dnc))
logger.info("Greedy action {} Random action {}".format(
action_greedy, action_random))
action_greedy = 0
action_random = 0
lost_action = []
win_actions = []
if episodes != 0:
sum_scores = 0
total_loss = 0
episodes = 0
if t % save_frequency and not args.exploiting:
env.pause_game()
with open("cache.p", "wb") as fh:
pickle.dump((t, memory, epsilon, nb_episodes), fh)
gc.collect()
torch.save(
{
"qnetwork": qnetwork.state_dict(),
"optimizer": optimizer.state_dict()
}, checkpoint_path)
env.resume_game()
except KeyboardInterrupt:
if not args.exploiting:
torch.save(
{
"qnetwork": qnetwork.state_dict(),
"optimizer": optimizer.state_dict()
}, checkpoint_path)
with open("cache.p", "wb") as fh:
pickle.dump((t, memory, epsilon, nb_episodes), fh)
from checkpoints.model_checkpoint_backup_config import config
from models import QNetwork
def get_env_configs(config):
env = gym.make(config["env"])
config["num_actions"] = env.action_space.n
config["observation_shape"] = env.observation_space.shape
return config
if __name__ == '__main__':
config = get_env_configs(config)
env = gym.make('CartPole-v1').unwrapped
net = QNetwork(config)
print(net.net)
net.load_state_dict(torch.load("checkpoints/model_checkpoint_backup"))
high_score = -math.inf
episode = 0
num_samples = 0
while True:
done = False
state = env.reset()
score, frame = 0, 1
while not done:
env.render()
state = torch.tensor(state, dtype=torch.float32)
state = state.unsqueeze(0)
qs = net(state)[0]
class SACAgent():
def __init__(self, action_size, state_size, config):
self.action_size = action_size
self.state_size = state_size
self.min_action = config["min_action"]
self.max_action = config["max_action"]
self.seed = config["seed"]
self.tau = config["tau"]
self.gamma = config["gamma"]
self.batch_size = config["batch_size"]
if not torch.cuda.is_available():
config["device"] == "cpu"
self.device = config["device"]
self.eval = config["eval"]
torch.manual_seed(self.seed)
np.random.seed(self.seed)
self.vid_path = config["vid_path"]
print("actions size ", action_size)
print("actions min ", self.min_action)
print("actions max ", self.max_action)
self.critic = QNetwork(state_size, action_size, config["fc1_units"], config["fc2_units"]).to(self.device)
self.q_optim = torch.optim.Adam(self.critic.parameters(), config["lr_critic"])
self.target_critic = QNetwork(state_size, action_size, config["fc1_units"], config["fc2_units"]).to(self.device)
self.target_critic.load_state_dict(self.critic.state_dict())
self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
self.alpha = self.log_alpha.exp()
self.alpha_optim = Adam([self.log_alpha], lr=config["lr_alpha"])
#self.policy = SACActor(state_size, action_size).to(self.device)
self.policy = GaussianPolicy(state_size, action_size, 256).to(self.device)
self.policy_optim = Adam(self.policy.parameters(), lr=config["lr_policy"])
self.max_timesteps = config["max_episodes_steps"]
self.episodes = config["episodes"]
self.memory = ReplayBuffer((state_size, ), (action_size, ), config["buffer_size"], self.device)
pathname = config["seed"]
tensorboard_name = str(config["res_path"]) + '/runs/' + str(pathname)
self.writer = SummaryWriter(tensorboard_name)
self.steps= 0
self.target_entropy = -torch.prod(torch.Tensor(action_size).to(self.device)).item()
def act(self, state, evaluate=False):
with torch.no_grad():
state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
if evaluate is False:
action, _, _ = self.policy.sample(state)
else:
_, _, action = self.policy.sample(state)
# action = np.clip(action, self.min_action, self.max_action)
action = action.cpu().numpy()[0]
#print(action)
return action
def train_agent(self):
env = gym.make("LunarLanderContinuous-v2")
average_reward = 0
scores_window = deque(maxlen=100)
s = 0
t0 = time.time()
for i_epiosde in range(self.episodes):
episode_reward = 0
state = env.reset()
for t in range(self.max_timesteps):
s += 1
action = self.act(state)
next_state, reward, done, _ = env.step(action)
episode_reward += reward
if i_epiosde > 3:
self.learn()
self.memory.add(state, reward, action, next_state, done)
state = next_state
if done:
scores_window.append(episode_reward)
break
if i_epiosde % self.eval == 0:
self.eval_policy()
ave_reward = np.mean(scores_window)
print("Epiosde {} Steps {} Reward {} Reward averge{} Time {}".format(i_epiosde, t, episode_reward, np.mean(scores_window), time_format(time.time() - t0)))
self.writer.add_scalar('Aver_reward', ave_reward, self.steps)
def learn(self):
self.steps += 1
states, rewards, actions, next_states, dones = self.memory.sample(self.batch_size)
with torch.no_grad():
next_state_action, next_state_log_pi, _ = self.policy.sample(next_states)
target_q1, target_q2 = self.target_critic(next_states, next_state_action)
target_min = torch.min(target_q1, target_q2)
q_target = target_min - (self.alpha * next_state_log_pi)
next_q_value = rewards + (1 - dones) * self.gamma * q_target
qf1, qf2 = self.critic(states, actions)
# --------------------------update-q--------------------------------------------------------
loss = F.mse_loss(qf1, next_q_value) + F.mse_loss(qf2, next_q_value)
self.q_optim.zero_grad()
loss.backward()
self.q_optim.step()
self.writer.add_scalar('loss/q', loss, self.steps)
# --------------------------update-policy--------------------------------------------------------
pi, log_pi, _ = self.policy.sample(states)
q_pi1, q_pi2 = self.critic(states, pi)
min_q_values = torch.min(q_pi1, q_pi2)
policy_loss = ((self.alpha * log_pi) - min_q_values).mean()
self.policy_optim.zero_grad()
policy_loss.backward()
self.policy_optim.step()
self.writer.add_scalar('loss/policy', policy_loss, self.steps)
# --------------------------update-alpha--------------------------------------------------------
alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
self.alpha_optim.zero_grad()
alpha_loss.backward()
self.alpha_optim.step()
self.writer.add_scalar('loss/alpha', alpha_loss, self.steps)
self.soft_udapte(self.critic, self.target_critic)
self.alpha = self.log_alpha.exp()
def soft_udapte(self, online, target):
for param, target_parm in zip(online.parameters(), target.parameters()):
target_parm.data.copy_(self.tau * param.data + (1 - self.tau) * target_parm.data)
def eval_policy(self, eval_episodes=4):
env = gym.make("LunarLanderContinuous-v2")
# env = wrappers.Monitor(env, str(self.vid_path) + "/{}".format(self.steps), video_callable=lambda episode_id: True,force=True)
average_reward = 0
scores_window = deque(maxlen=100)
for i_epiosde in range(eval_episodes):
print("Eval Episode {} of {} ".format(i_epiosde, eval_episodes))
episode_reward = 0
state = env.reset()
while True:
action = self.act(state, evaluate=True)
state, reward, done, _ = env.step(action)
episode_reward += reward
if done:
break
scores_window.append(episode_reward)
average_reward = np.mean(scores_window)
self.writer.add_scalar('Eval_reward', average_reward, self.steps)
class NewSACAgent:
def __init__(self, env, render, config_info):
self.env = env
self.render = render
self._reset_env()
# Create run folder to store parameters, figures, and tensorboard logs
self.path_runs = create_run_folder(config_info)
# Extract training parameters from yaml config file
param = load_training_parameters(config_info["config_param"])
self.train_param = param["training"]
# Define device
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device in use : {self.device}")
# Define state and action dimension spaces
state_dim = env.observation_space.shape[0]
num_actions = env.action_space.shape[0]
# Define models
hidden_size = param["model"]["hidden_size"]
self.q_net = QNetwork(state_dim, num_actions, hidden_size).to(self.device)
self.target_q_net = QNetwork(state_dim, num_actions, hidden_size).to(
self.device
)
self.target_q_net.load_state_dict(self.q_net.state_dict())
self.policy_net = PolicyNetwork(state_dim, num_actions, hidden_size).to(
self.device
)
# Define loss criterion
self.q_criterion = nn.MSELoss()
# Define optimizers
lr = float(param["optimizer"]["learning_rate"])
self.q_opt = optim.Adam(self.q_net.parameters(), lr=lr)
self.policy_opt = optim.Adam(self.policy_net.parameters(), lr=lr)
# Initialize replay buffer
self.replay_buffer = ReplayBuffer(param["training"]["replay_size"])
self.transition = namedtuple(
"transition",
field_names=["state", "action", "reward", "done", "next_state"],
)
# Useful variables
self.batch_size = param["training"]["batch_size"]
self.gamma = param["training"]["gamma"]
self.tau = param["training"]["tau"]
self.start_step = param["training"]["start_step"]
self.max_timesteps = param["training"]["max_timesteps"]
self.alpha = param["training"]["alpha"]
def _reset_env(self):
# Reset the environment and initialize episode reward
self.state, self.done = self.env.reset(), False
self.episode_reward = 0.0
self.episode_step = 0
def train(self):
# Main training loop
total_timestep = 0
all_episode_rewards = []
all_mean_rewards = []
update = 0
# Create tensorboard writer
writer = SummaryWriter(log_dir=self.path_runs, comment="-sac")
for episode in itertools.count(1, 1):
self._reset_env()
while not self.done:
# trick to improve exploration at the start of training
if self.start_step > total_timestep:
action = self.env.action_space.sample() # Sample random action
else:
action = self.policy_net.get_action(
self.state, self.device
) # Sample action from policy
# Fill the replay buffer up with transitions
if len(self.replay_buffer) > self.batch_size:
batch = self.replay_buffer.sample_buffer(self.batch_size)
# Update parameters of all the networks
q1_loss, q2_loss, policy_loss = self.train_on_batch(batch)
writer.add_scalar("loss/q1", q1_loss, update)
writer.add_scalar("loss/q2", q2_loss, update)
writer.add_scalar("loss/policy", policy_loss, update)
update += 1
if self.render:
self.env.render()
# Perform one step in the environment
next_state, reward, self.done, _ = self.env.step(action)
total_timestep += 1
self.episode_step += 1
self.episode_reward += reward
# Create a tuple for the new transition
new_transition = self.transition(
self.state, action, reward, self.done, next_state
)
# Append transition to the replay buffer
self.replay_buffer.store_transition(new_transition)
self.state = next_state
if total_timestep > self.max_timesteps:
break
mean_reward = np.mean(all_episode_rewards[-100:])
all_episode_rewards.append(self.episode_reward)
all_mean_rewards.append(mean_reward)
print(
"Episode n°{} ; total timestep [{}/{}] ; episode steps {} ; "
"reward {} ; mean reward {}".format(
episode,
total_timestep,
self.max_timesteps,
self.episode_step,
round(self.episode_reward, 2),
round(mean_reward, 2)
)
)
writer.add_scalar("reward", self.episode_reward, episode)
writer.add_scalar("mean reward", mean_reward, episode)
# Save networks' weights
path_critic = os.path.join(self.path_runs, "critic.pth")
path_actor = os.path.join(self.path_runs, "actor.pth")
torch.save(self.q_net.state_dict(), path_critic)
torch.save(self.policy_net.state_dict(), path_actor)
# Plot reward
self.plot_reward(all_episode_rewards, all_mean_rewards)
# Close all
writer.close()
self.env.close()
def train_on_batch(self, batch_samples):
# Unpack batch_size of transitions randomly drawn from the replay buffer
(
state_batch,
action_batch,
reward_batch,
done_int_batch,
next_state_batch,
) = batch_samples
# Transform np arrays into tensors and send them to device
state_batch = torch.tensor(state_batch).to(self.device)
next_state_batch = torch.tensor(next_state_batch).to(self.device)
action_batch = torch.tensor(action_batch).to(self.device)
reward_batch = torch.tensor(reward_batch).unsqueeze(1).to(self.device)
done_int_batch = torch.tensor(done_int_batch).unsqueeze(1).to(self.device)
### Update Q
with torch.no_grad():
next_state_action, next_state_log_pi = self.policy_net.sample(
next_state_batch
)
q1_next_target, q2_next_target = self.target_q_net(
next_state_batch, next_state_action
)
min_qf_next_target = (
torch.min(q1_next_target, q2_next_target)
- self.alpha * next_state_log_pi
)
next_q_value = reward_batch + (1 - done_int_batch) * self.gamma * (
min_qf_next_target
)
# Estimated state-action values
q1, q2 = self.q_net(state_batch, action_batch)
q1_loss = self.q_criterion(q1, next_q_value)
q2_loss = self.q_criterion(q2, next_q_value)
### Update policy
pi, log_pi = self.policy_net.sample(state_batch)
q1_pi, q2_pi = self.q_net(state_batch, pi)
# Mitigate positive bias in the policy improvement step
min_q_pi = torch.min(q1_pi, q2_pi)
policy_loss = (self.alpha * log_pi - min_q_pi).mean()
# Losses and optimizers
self.q_opt.zero_grad()
q1_loss.backward()
self.q_opt.step()
self.q_opt.zero_grad()
q2_loss.backward()
self.q_opt.step()
self.policy_opt.zero_grad()
policy_loss.backward()
self.policy_opt.step()
soft_update(self.target_q_net, self.q_net, self.tau)
return q1_loss.item(), q2_loss.item(), policy_loss.item()
def plot_reward(self, data, mean_data):
plt.plot(data, label="reward")
plt.plot(mean_data, label="mean reward")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.title(f"Reward evolution for {self.env.unwrapped.spec.id} Gym environment")
plt.tight_layout()
plt.legend()
path_fig = os.path.join(self.path_runs, "figure.png")
plt.savefig(path_fig)
print(f"Figure saved to {path_fig}")
plt.show()
