def surrogate_loss(actor, advants, states, old_policy, actions, batch_index): mu, std = actor(states) new_policy = log_prob_density(actions, mu, std) old_policy = old_policy[batch_index] ratio = torch.exp(new_policy - old_policy) surrogate_loss = ratio * advants return surrogate_loss, ratio
def surrogate_loss(actor, advants, states, old_policy, actions, batch_index): mu, std = actor(states) new_policy = log_prob_density(actions, mu, std) old_policy = old_policy[batch_index] ratio = torch.exp(new_policy - old_policy) surrogate_loss = ratio * advants return surrogate_loss, ratio
def train_actor_critic(actor, critic, memory, actor_optim, critic_optim, args): memory = np.array(memory) states = np.vstack(memory[:, 0]) actions = list(memory[:, 1]) rewards = list(memory[:, 2]) masks = list(memory[:, 3]) old_values = critic(torch.Tensor(states)) returns, advants = get_gae(rewards, masks, old_values, args) mu, std = actor(torch.Tensor(states)) old_policy = log_prob_density(torch.Tensor(actions), mu, std) criterion = torch.nn.MSELoss() n = len(states) arr = np.arange(n) for _ in range(args.actor_critic_update_num): np.random.shuffle(arr) for i in range(n // args.batch_size): batch_index = arr[args.batch_size * i : args.batch_size * (i + 1)] batch_index = torch.LongTensor(batch_index) inputs = torch.Tensor(states)[batch_index] actions_samples = torch.Tensor(actions)[batch_index] returns_samples = returns.unsqueeze(1)[batch_index] advants_samples = advants.unsqueeze(1)[batch_index] oldvalue_samples = old_values[batch_index].detach() values = critic(inputs) clipped_values = oldvalue_samples + \ torch.clamp(values - oldvalue_samples, -args.clip_param, args.clip_param) critic_loss1 = criterion(clipped_values, returns_samples) critic_loss2 = criterion(values, returns_samples) critic_loss = torch.max(critic_loss1, critic_loss2).mean() loss, ratio, entropy = surrogate_loss(actor, advants_samples, inputs, old_policy.detach(), actions_samples, batch_index) clipped_ratio = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) clipped_loss = clipped_ratio * advants_samples actor_loss = -torch.min(loss, clipped_loss).mean() loss = actor_loss + 0.5 * critic_loss - 0.001 * entropy critic_optim.zero_grad() loss.backward(retain_graph=True) critic_optim.step() actor_optim.zero_grad() loss.backward() actor_optim.step()
def train_model(actor, critic, memory, actor_optim, critic_optim, args): memory = np.array(memory) states = np.vstack(memory[:, 0]) actions = list(memory[:, 1]) rewards = list(memory[:, 2]) masks = list(memory[:, 3]) old_values = critic(torch.Tensor(states)) returns, advants = get_gae(rewards, masks, old_values, args) mu, std = actor(torch.Tensor(states)) old_policy = log_prob_density(torch.Tensor(actions), mu, std) criterion = torch.nn.MSELoss() n = len(states) arr = np.arange(n) for _ in range(args.model_update_num): np.random.shuffle(arr) for i in range(n // args.batch_size): batch_index = arr[args.batch_size * i : args.batch_size * (i + 1)] batch_index = torch.LongTensor(batch_index) inputs = torch.Tensor(states)[batch_index] actions_samples = torch.Tensor(actions)[batch_index] returns_samples = returns.unsqueeze(1)[batch_index] advants_samples = advants.unsqueeze(1)[batch_index] oldvalue_samples = old_values[batch_index].detach() values = critic(inputs) clipped_values = oldvalue_samples + \ torch.clamp(values - oldvalue_samples, -args.clip_param, args.clip_param) critic_loss1 = criterion(clipped_values, returns_samples) critic_loss2 = criterion(values, returns_samples) critic_loss = torch.max(critic_loss1, critic_loss2).mean() loss, ratio = surrogate_loss(actor, advants_samples, inputs, old_policy.detach(), actions_samples, batch_index) clipped_ratio = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) clipped_loss = clipped_ratio * advants_samples actor_loss = -torch.min(loss, clipped_loss).mean() loss = actor_loss + 0.5 * critic_loss critic_optim.zero_grad() loss.backward(retain_graph=True) critic_optim.step() actor_optim.zero_grad() loss.backward() actor_optim.step()