import random
from collections import namedtuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
import matplotlib.pyplot as plt
# 使用gpu加速
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Hyper-parameters
batch_size = 32
learning_rate = 0.01
epsilon = 0.9
gamma = 0.9
target_replace_iter = 100
memory_size = 2000
env = gym.make('CartPole-v0')
env = env.unwrapped
action_size = env.action_space.n
state_size = env.observation_space.shape[0]
env_a_shape = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample().shape
# replay_buffer对应的元素
Transition = namedtuple('Transition', ('state', 'action', 'reward', 'next_state'))
class ReplayBuffer:
def __init__(self):
self.capacity = memory_size
self.memory = []
self.position = 0
self.counter = 0
def push(self, *args):
if len(self.memory) < self.capacity:
self.memory.append(None)
self.memory[self.position] = Transition(*args)
self.position = (self.position + 1) % self.capacity
self.counter += 1
def sample(self, batch_size):
return random.sample(self.memory, batch_size)
def __len__(self):
return len(self.memory)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(state_size, 50)
self.fc1.weight.data.normal_(0, 0.1) # initialization
self.out = nn.Linear(50, action_size)
self.out.weight.data.normal_(0, 0.1)
def forward(self, x):
x = self.fc1(x)
x = F.relu(x)
action_value = self.out(x)
return action_value
class DQN:
def __init__(self):
super(DQN, self).__init__()
self.eval_net, self.target_net = Net().to(device), Net().to(device) # 1.GPU加速部分,将网络部分放到GPU中
self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=learning_rate)
self.loss_func = nn.MSELoss()
self.learn_step_counter = 0
def choose_action(self, x):
if np.random.uniform() < epsilon:
action_value = self.eval_net(x)
action = action_value.max(1)[1].unsqueeze(0)
action = action[0] if env_a_shape == 0 else action.reshape(env_a_shape)
else:
action = np.random.randint(0, action_size)
action = torch.IntTensor([action]).to(device) # 2.将action放到GPU中
return action
def learn_dqn(self):
if self.learn_step_counter % target_replace_iter == 0:
self.target_net.load_state_dict(self.eval_net.state_dict())
self.learn_step_counter += 1
# 抽样本
transitions = memory.sample(batch_size=batch_size)
batch = Transition(*zip(*transitions))
b_s = torch.cat(batch.state)
b_a = torch.cat(batch.action).unsqueeze(1)
b_r = torch.cat(batch.reward).unsqueeze(1)
b_s_ = torch.cat(batch.next_state)
# DQN
q_eval = self.eval_net(b_s).gather(1, b_a)
q_next = self.target_net(b_s_).detach()
q_target = b_r + gamma * q_next.max(1)[0].view(batch_size, 1)
loss = self.loss_func(q_eval, q_target)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if __name__=="__main__":
dqn = DQN()
print('DQN start collecting...')
result_dqn = []
memory = ReplayBuffer()
for i_episode in range(400):
s = env.reset()
# 将state转换成tensor,并复制张量到gpu
s = torch.Tensor([s]).to(device) # 3.将状态放到device中
ep_r = 0
while True:
# env.render()
a = dqn.choose_action(s)
s_, reward, done, info = env.step(a.data.item())
# 计算奖励
x, x_dot, theta, theta_dot = s_
r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5
r = r1 + r2
# 将变量转为tensor,同时将设备设置为cuda # 4.将奖励放到GPU中
r = torch.Tensor([r]).to(device)
s_ = torch.Tensor([s_]).to(device)
# 保存交互信息
memory.push(s, a, r, s_)
# dqn.store_transition(s, a, r, s_)
ep_r += reward
# print(memory.counter)
if memory.counter > memory_size:
dqn.learn_dqn()
if done:
result_dqn.append(ep_r)
print('Episode: ', i_episode, '|ep_r: ', ep_r)
if done:
break
s = s_
plt.plot(result_dqn)
plt.show()
强化学习:GPU加速DQN玩倒立摆代码
最新推荐文章于 2024-07-03 16:17:09 发布