mirror of
https://github.com/wassname/Deep-reinforcement-learning-with-pytorch.git
synced 2026-09-21 12:30:14 +08:00
189 lines
6.6 KiB
Python
189 lines
6.6 KiB
Python
|
|
import argparse
|
|
import pickle
|
|
from collections import namedtuple
|
|
|
|
import os
|
|
import numpy as np
|
|
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
|
|
|
|
# Parameters
|
|
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', default=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()
|
|
|
|
env = gym.make('Pendulum-v0').unwrapped
|
|
num_state = env.observation_space.shape[0]
|
|
num_action = env.action_space.shape[0]
|
|
torch.manual_seed(args.seed)
|
|
env.seed(args.seed)
|
|
|
|
Transition = namedtuple('Transition',['state', 'aciton', 'reward', 'a_log_prob', 'next_state'])
|
|
TrainRecord = namedtuple('TrainRecord',['episode', 'reward'])
|
|
|
|
class Actor(nn.Module):
|
|
def __init__(self):
|
|
super(Actor, self).__init__()
|
|
self.fc1 = nn.Linear(num_state, 64)
|
|
self.fc2 = nn.Linear(64,8)
|
|
self.mu_head = nn.Linear(8, 1)
|
|
self.sigma_head = nn.Linear(8, 1)
|
|
|
|
def forward(self, x):
|
|
x = F.leaky_relu(self.fc1(x))
|
|
x = F.leaky_relu(self.fc2(x))
|
|
|
|
mu = self.mu_head(x)
|
|
sigma = self.sigma_head(x)
|
|
|
|
return mu, sigma
|
|
|
|
class Critic(nn.Module):
|
|
def __init__(self):
|
|
super(Critic, self).__init__()
|
|
self.fc1 = nn.Linear(num_state, 64)
|
|
self.fc2 = nn.Linear(64, 8)
|
|
self.state_value= nn.Linear(8, 1)
|
|
|
|
def forward(self, x):
|
|
x = F.leaky_relu(self.fc1(x))
|
|
x = F.leaky_relu(self.fc2(x))
|
|
value = self.state_value(x)
|
|
return value
|
|
|
|
class PPO():
|
|
clip_param = 0.2
|
|
max_grad_norm = 0.5
|
|
ppo_epoch = 10
|
|
buffer_capacity = 1000
|
|
batch_size = 8
|
|
|
|
def __init__(self):
|
|
super(PPO, self).__init__()
|
|
self.actor_net = Actor().float()
|
|
self.critic_net = Critic().float()
|
|
self.buffer = []
|
|
self.counter = 0
|
|
self.training_step = 0
|
|
|
|
self.actor_optimizer = optim.Adam(self.actor_net.parameters(), 1e-3)
|
|
self.critic_net_optimizer = optim.Adam(self.critic_net.parameters(), 4e-3)
|
|
if not os.path.exists('../param'):
|
|
os.makedirs('../param/net_param')
|
|
os.makedirs('../param/img')
|
|
|
|
def select_action(self, state):
|
|
state = torch.from_numpy(state).float().unsqueeze(0)
|
|
with torch.no_grad():
|
|
mu, sigma = self.actor_net(state)
|
|
dist = Normal(mu, sigma)
|
|
action = dist.sample()
|
|
action_log_prob = dist.log_prob(action)
|
|
action = action.clamp(-2, 2)
|
|
return action.item(), action_log_prob.item()
|
|
|
|
|
|
def get_value(self, state):
|
|
state = torch.from_numpy(state)
|
|
with torch.no_grad():
|
|
value = self.critic_net(state)
|
|
return value.item()
|
|
|
|
def save_param(self):
|
|
torch.save(self.actor_net.state_dict(), '../param/net_param/actor_net'+str(time.time())[:10],+'.pkl')
|
|
torch.save(self.critic_net.state_dict(), '../param/net_param/critic_net'+str(time.time())[:10],+'.pkl')
|
|
|
|
def store_transition(self, transition):
|
|
self.buffer.append(transition)
|
|
self.counter+=1
|
|
return counter % self.buffer_capacity == 0
|
|
|
|
def update(self):
|
|
self.training_step +=1
|
|
|
|
state = torch.tensor([t.state for t in self.buffer ], dtype=torch.float)
|
|
action = torch.tensor([t.action for t in self.buffer], dtype=torch.float).view(-1, 1)
|
|
reward = torch.tensor([t.reward for t in self.buffer], dtype=torch.float).view(-1, 1)
|
|
next_state = torch.tensor([t.next_state for t in self.buffer], dtype=torch.float)
|
|
old_action_log_prob = torch.tensor([t.a_log_prob for t in self.buffer], dtype=torch.float).view(-1, 1)
|
|
|
|
reward = (reward - reward.mean())/(reward.std() + 1e-10)
|
|
with torch.no_grad():
|
|
target_v = reward + args.gamma * self.critic_net(next_state)
|
|
|
|
advantage = (target_v - self.critic_net(state)).detach()
|
|
for _ in range(self.ppo_epoch): # iteration ppo_epoch
|
|
for index in BatchSampler(SubsetRandomSampler(range(self.buffer_capacity), self.batch_size, True)):
|
|
# epoch iteration, PPO core!!!
|
|
mu, sigma = self.actor_net(state[index])
|
|
n = Normal(mu, sigma)
|
|
action_log_prob = n.log_prob(action[index])
|
|
ratio = torch.exp(action_log_prob - old_action_log_prob)
|
|
|
|
L1 = ratio * advantage[index]
|
|
L2 = torch.clamp(ratio, 1-self.clip_param, 1+self.clip_param) * advantage[index]
|
|
action_loss = -torch.min(L1, L2).mean() # MAX->MIN desent
|
|
self.actor_optimizer.zero_grad()
|
|
action_loss.backward()
|
|
nn.utils.clip_grad_norm_(self.actor_net.parameters(), self.max_grad_norm)
|
|
self.actor_optimizer.step()
|
|
|
|
value_loss = F.smooth_l1_loss(self.critic_net(state[index]), target_v[index])
|
|
self.critic_net_optimizer.zero_grad()
|
|
value_loss.backward()
|
|
nn.utils.clip_grad_norm_(self.critic_net.parameters(), self.max_grad_norm)
|
|
self.critic_net_optimizer.step()
|
|
|
|
del self.buffer[:]
|
|
|
|
def main():
|
|
|
|
agent = PPO()
|
|
|
|
training_records = []
|
|
running_reward = -1000
|
|
|
|
for i_epoch in range(1000):
|
|
score = 0
|
|
state = env.reset()
|
|
if args.render: env.render()
|
|
for t in range(200):
|
|
action, action_log_prob = agent.select_action(state)
|
|
next_state, reward, done, info = env.step(action)
|
|
trans = Transition(state, action, reward, action_log_prob, next_state)
|
|
if args.render: env.render()
|
|
if agent.store_transition(trans):
|
|
agent.update()
|
|
score += reward
|
|
state = next_state
|
|
|
|
running_reward = running_reward * 0.9 + score * 0.1
|
|
training_records.append(TrainingRecord(i_epoch, running_reward))
|
|
if i_epoch % 10 ==0:
|
|
print("Epoch {}, Moving average score is: {:.2f} ".format(i_epoch, running_reward))
|
|
if running_reward > -200:
|
|
print("Solved! Moving average score is now {}!".format(running_reward))
|
|
env.close()
|
|
agent.save_param()
|
|
break
|
|
|
|
if __name__ == '__main__':
|
|
main()
|