import argparse
from collections import deque
import os
import random
import numpy as np
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
device = 'cuda' if torch.cuda.is_available() else 'cpu'
class ReplayBuffer:
def __init__(self, max_size):
self.storage = []
self.max_size = max_size
self.ptr = 0
def push(self, data):
if len(self.storage) == self.max_size:
self.storage[int(self.ptr)] = data
self.ptr = (self.ptr + 1) % self.max_size
else:
self.storage.append(data)
def sample(self, batch_size):
ind = np.random.randint(0, len(self.storage), size=batch_size)
x, y, u, r, d = [], [], [], [], []
for i in ind:
X, Y, U, R, D = self.storage[i]
x.append(np.array(X, copy=False))
y.append(np.array(Y, copy=False))
u.append(np.array(U, copy=False))
r.append(np.array(R, copy=False))
d.append(np.array(D, copy=False))
return np.array(x), np.array(y), np.array(u), np.array(r).reshape(-1, 1), np.array(d).reshape(-1, 1)
class Actor(nn.Module):
def __init__(self, state_dim, action_dim, max_action):
super(Actor, self).__init__()
self.fc1 = nn.Linear(state_dim, 400)
self.fc2 = nn.Linear(400, 300)
self.fc3 = nn.Linear(300, action_dim)
self.max_action = max_action
def forward(self, state):
a = F.relu(self.fc1(state))
a = F.relu(self.fc2(a))
a = torch.tanh(self.fc3(a)) * self.max_action
return a
class Critic(nn.Module):
def __init__(self, state_dim, action_dim):
super(Critic, self).__init__()
self.fc1 = nn.Linear(state_dim + action_dim, 400)
self.fc2 = nn.Linear(400, 300)
self.fc3 = nn.Linear(300, 1)
def forward(self, state, action):
state_action = torch.cat([state, action], 1)
q = F.relu(self.fc1(state_action))
q = F.relu(self.fc2(q))
q = self.fc3(q)
return q
class TD3:
def __init__(self, state_dim, action_dim, max_action):
self.actor = Actor(state_dim, action_dim, max_action).to(device)
self.actor_target = Actor(state_dim, action_dim, max_action).to(device)
self.critic_1 = Critic(state_dim, action_dim).to(device)
self.critic_1_target = Critic(state_dim, action_dim).to(device)
self.critic_2 = Critic(state_dim, action_dim).to(device)
self.critic_2_target = Critic(state_dim, action_dim).to(device)
self.actor_optimizer = optim.Adam(self.actor.parameters())
self.critic_1_optimizer = optim.Adam(self.critic_1.parameters())
self.critic_2_optimizer = optim.Adam(self.critic_2.parameters())
self.actor_target.load_state_dict(self.actor.state_dict())
self.critic_1_target.load_state_dict(self.critic_1.state_dict())
self.critic_2_target.load_state_dict(self.critic_2.state_dict())
self.max_action = max_action
self.memory = ReplayBuffer(max_size=50000)
self.num_critic_update_iteration = 0
self.num_actor_update_iteration = 0
self.num_training = 0
def select_action(self, state):
state = torch.tensor(state.reshape(1, -1)).float().to(device)
return self.actor(state).cpu().data.numpy().flatten()
def update(self, num_iteration, args):
for i in range(num_iteration):
x, y, u, r, d = self.memory.sample(args.batch_size)
state = torch.FloatTensor(x).to(device)
action = torch.FloatTensor(u).to(device)
next_state = torch.FloatTensor(y).to(device)
done = torch.FloatTensor(d).to(device)
reward = torch.FloatTensor(r).to(device)
noise = torch.ones_like(action).data.normal_(0, args.policy_noise).to(device)
noise = noise.clamp(-args.noise_clip, args.noise_clip)
next_action = (self.actor_target(next_state) + noise).clamp(-self.max_action, self.max_action)
target_Q1 = self.critic_1_target(next_state, next_action)
target_Q2 = self.critic_2_target(next_state, next_action)
target_Q = torch.min(target_Q1, target_Q2)
target_Q = reward + ((1 - done) * args.gamma * target_Q).detach()
current_Q1 = self.critic_1(state, action)
loss_Q1 = F.mse_loss(current_Q1, target_Q)
self.critic_1_optimizer.zero_grad()
loss_Q1.backward()
self.critic_1_optimizer.step()
current_Q2 = self.critic_2(state, action)
loss_Q2 = F.mse_loss(current_Q2, target_Q)
self.critic_2_optimizer.zero_grad()
loss_Q2.backward()
self.critic_2_optimizer.step()
if i % args.policy_delay == 0:
actor_loss = - self.critic_1(state, self.actor(state)).mean()
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)
for param, target_param in zip(self.critic_1.parameters(), self.critic_1_target.parameters()):
target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)
for param, target_param in zip(self.critic_2.parameters(), self.critic_2_target.parameters()):
target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)
self.num_actor_update_iteration += 1
self.num_critic_update_iteration += 1
self.num_training += 1
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--env_name', default="Pendulum-v0")
parser.add_argument('--tau', default=0.005, type=float)
parser.add_argument('--iteration', default=5, type=int)
parser.add_argument('--learning_rate', default=3e-4, type=float)
parser.add_argument('--gamma', default=0.99, type=float)
parser.add_argument('--capacity', default=50000, type=int)
parser.add_argument('--num_iteration', default=100000, type=int)
parser.add_argument('--batch_size', default=100, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--policy_noise', default=0.2, type=float)
parser.add_argument('--noise_clip', default=0.5, type=float)
parser.add_argument('--policy_delay', default=2, type=int)
parser.add_argument('--exploration_noise', default=0.1, type=float)
parser.add_argument('--max_episode', default=2000, type=int)
args = parser.parse_args()
env = gym.make(args.env_name)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
max_action = float(env.action_space.high[0])
agent = TD3(state_dim, action_dim, max_action)
if args.mode == 'train':
ep_r = 0
for i in range(args.num_iteration):
state = env.reset()
for t in range(2000):
action = agent.select_action(state)
action = action + np.random.normal(0, args.exploration_noise, size=env.action_space.shape[0])
action = action.clip(env.action_space.low, env.action_space.high)
next_state, reward, done, info = env.step(action)
ep_r += reward
agent.memory.push((state, next_state, action, reward, np.float(done)))
if len(agent.memory.storage) >= args.capacity - 1:
agent.update(10, args)
state = next_state
if done or t == args.max_episode - 1:
print(f"Ep_i {i}, ep_r {ep_r:.2f}")
ep_r = 0
break