From 7dcfbb3e5de7767ed8907ea8db1bc6db5e83b6ba Mon Sep 17 00:00:00 2001 From: Johnny He <269401927@qq.com> Date: Mon, 5 Nov 2018 20:55:03 +0800 Subject: [PATCH] Create PPO_pendulum.py --- PPO/PPO_pendulum.py | 195 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 195 insertions(+) create mode 100644 PPO/PPO_pendulum.py diff --git a/PPO/PPO_pendulum.py b/PPO/PPO_pendulum.py new file mode 100644 index 0000000..5f5fa5d --- /dev/null +++ b/PPO/PPO_pendulum.py @@ -0,0 +1,195 @@ +import argparse +import pickle +from collections import namedtuple + +import matplotlib.pyplot as plt + +import gym +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.distributions import Normal +from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler + +parser = argparse.ArgumentParser(description='Solve the Pendulum-v0 with PPO') +parser.add_argument( + '--gamma', type=float, default=0.9, metavar='G', help='discount factor (default: 0.9)') +parser.add_argument('--seed', type=int, default=0, metavar='N', help='random seed (default: 0)') +parser.add_argument('--render', action='store_true', help='render the environment') +parser.add_argument( + '--log-interval', + type=int, + default=10, + metavar='N', + help='interval between training status logs (default: 10)') +args = parser.parse_args() + +torch.manual_seed(args.seed) + +TrainingRecord = namedtuple('TrainingRecord', ['ep', 'reward']) +Transition = namedtuple('Transition', ['s', 'a', 'a_log_p', 'r', 's_']) + + +class ActorNet(nn.Module): + + def __init__(self): + super(ActorNet, self).__init__() + self.fc = nn.Linear(3, 100) + self.mu_head = nn.Linear(100, 1) + self.sigma_head = nn.Linear(100, 1) + + def forward(self, x): + x = F.relu(self.fc(x)) + mu = 2.0 * F.tanh(self.mu_head(x)) + sigma = F.softplus(self.sigma_head(x)) + return (mu, sigma) + + +class CriticNet(nn.Module): + + def __init__(self): + super(CriticNet, self).__init__() + self.fc = nn.Linear(3, 100) + self.v_head = nn.Linear(100, 1) + + def forward(self, x): + x = F.relu(self.fc(x)) + state_value = self.v_head(x) + return state_value + + +class Agent(): + + clip_param = 0.2 + max_grad_norm = 0.5 + ppo_epoch = 10 + buffer_capacity, batch_size = 1000, 32 + + def __init__(self): + self.training_step = 0 + self.anet = ActorNet().float() + self.cnet = CriticNet().float() + self.buffer = [] + self.counter = 0 + + self.optimizer_a = optim.Adam(self.anet.parameters(), lr=1e-4) + self.optimizer_c = optim.Adam(self.cnet.parameters(), lr=3e-4) + + def select_action(self, state): + state = torch.from_numpy(state).float().unsqueeze(0) + with torch.no_grad(): + (mu, sigma) = self.anet(state) + dist = Normal(mu, sigma) + action = dist.sample() + action_log_prob = dist.log_prob(action) + action.clamp(-2.0, 2.0) + return action.item(), action_log_prob.item() + + def get_value(self, state): + + state = torch.from_numpy(state).float().unsqueeze(0) + with torch.no_grad(): + state_value = self.cnet(state) + return state_value.item() + + def save_param(self): + torch.save(self.anet.state_dict(), 'param/ppo_anet_params.pkl') + torch.save(self.cnet.state_dict(), 'param/ppo_cnet_params.pkl') + + def store(self, transition): + self.buffer.append(transition) + self.counter += 1 + return self.counter % self.buffer_capacity == 0 + + def update(self): + self.training_step += 1 + + s = torch.tensor([t.s for t in self.buffer], dtype=torch.float) + a = torch.tensor([t.a for t in self.buffer], dtype=torch.float).view(-1, 1) + r = torch.tensor([t.r for t in self.buffer], dtype=torch.float).view(-1, 1) + s_ = torch.tensor([t.s_ for t in self.buffer], dtype=torch.float) + + old_action_log_probs = torch.tensor( + [t.a_log_p for t in self.buffer], dtype=torch.float).view(-1, 1) + + r = (r - r.mean()) / (r.std() + 1e-5) + with torch.no_grad(): + target_v = r + args.gamma * self.cnet(s_) + + adv = (target_v - self.cnet(s)).detach() + + for _ in range(self.ppo_epoch): + for index in BatchSampler( + SubsetRandomSampler(range(self.buffer_capacity)), self.batch_size, False): + + (mu, sigma) = self.anet(s[index]) + dist = Normal(mu, sigma) + action_log_probs = dist.log_prob(a[index]) + ratio = torch.exp(action_log_probs - old_action_log_probs[index]) + + surr1 = ratio * adv[index] + surr2 = torch.clamp(ratio, 1.0 - self.clip_param, + 1.0 + self.clip_param) * adv[index] + action_loss = -torch.min(surr1, surr2).mean() + + self.optimizer_a.zero_grad() + action_loss.backward() + nn.utils.clip_grad_norm_(self.anet.parameters(), self.max_grad_norm) + self.optimizer_a.step() + + value_loss = F.smooth_l1_loss(self.cnet(s[index]), target_v[index]) + self.optimizer_c.zero_grad() + value_loss.backward() + nn.utils.clip_grad_norm_(self.cnet.parameters(), self.max_grad_norm) + self.optimizer_c.step() + + del self.buffer[:] + + +def main(): + env = gym.make('Pendulum-v0') + env.seed(args.seed) + + agent = Agent() + + training_records = [] + running_reward = -1000 + state = env.reset() + for i_ep in range(1000): + score = 0 + state = env.reset() + + for t in range(200): + action, action_log_prob = agent.select_action(state) + state_, reward, done, _ = env.step([action]) + if args.render: + env.render() + if agent.store(Transition(state, action, action_log_prob, (reward + 8) / 8, state_)): + agent.update() + score += reward + state = state_ + + running_reward = running_reward * 0.9 + score * 0.1 + training_records.append(TrainingRecord(i_ep, running_reward)) + + if i_ep % args.log_interval == 0: + print('Ep {}\tMoving average score: {:.2f}\t'.format(i_ep, running_reward)) + if running_reward > -200: + print("Solved! Moving average score is now {}!".format(running_reward)) + env.close() + agent.save_param() + with open('log/ppo_training_records.pkl', 'wb') as f: + pickle.dump(training_records, f) + break + + plt.plot([r.ep for r in training_records], [r.reward for r in training_records]) + plt.title('PPO') + plt.xlabel('Episode') + plt.ylabel('Moving averaged episode reward') + plt.savefig("img/ppo.png") + plt.show() + + +if __name__ == '__main__': + main()