小车立杆之DQN实现详解（PyTorch）
前面的话
对于DQN的原理网上很多，故不再赘述，主要针对小车立杆这一场景的DQN实现代码进行详细说明，自我学习。
代码from莫烦老师.
完整代码
主要部分：
- 给出程序所需超参数；主要是与算法相关的参数
- 神经网络类；DQN中所使用的的神经网络
- DQN类；DQN算法的实现
- 主循环，训练过程
# -*- coding: utf-8 -*-
# @Time : 2019/12/8 14:05
# @Author : Chen
# @File : DQN_CartPole.py
# @Software: PyCharm
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
# Hyper Parameters
BATCH_SIZE = 32
LR = 0.01 # learning rate
EPSILON = 0.9 # greedy policy
GAMMA = 0.9 # reward discount
TARGET_REPLACE_ITER = 100 # target update frequency
MEMORY_CAPACITY = 2000
env = gym.make('CartPole-v0') # 小车立杆
env = env.unwrapped
N_ACTIONS = env.action_space.n
N_STATES = env.observation_space.shape[0]
ENV_A_SHAPE = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample().shape # to confirm the shape
# 定义神经网络
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 一个隐层，一个输出层
self.fc1 = nn.Linear(N_STATES, 100)
self.fc1.weight.data.normal_(0, 0.1)
self.out = nn.Linear(100, N_ACTIONS)
self.out.weight.data.normal_(0, 0.1)
def forward(self, x):
# Net的执行逻辑 Linear_fc1 --> relu --> out --> actions_value
x = self.fc1(x)
x = F.relu(x)
actions_value = self.out(x)
return actions_value
# 定义DQN模型
class DQN(object):
def __init__(self):
# 模型初始化。初始化main net和target net
# 设置target更新计数器、存储计数器、记忆库初始化、优化器和loss定义。
self.eval_net, self.target_net = Net(), Net()
self.learn_step_counter = 0
self.memory_counter = 0
self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))
self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
self.loss_func = nn.MSELoss()
def choose_action(self, x):
x = torch.unsqueeze(torch.FloatTensor(x), 0)
if np.random.uniform() < EPSILON:
actions_value = self.eval_net.forward(x)
action = torch.max(actions_value, 1)[1].data.numpy()
action = action[0] if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
else:
action = np.random.randint(0, N_ACTIONS)
action = action if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
return action
def store_transition(self, s, a, r, s_):
transition = np.hstack((s, [a, r], s_))
index = self.memory_counter % MEMORY_CAPACITY
self.memory[index,:] = transition
self.memory_counter += 1
def learn(self):
if self.learn_step_counter % TARGET_REPLACE_ITER == 0:
self.target_net.load_state_dict(self.eval_net.state_dict())
print('target update....')
self.learn_step_counter += 1
sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)
b_memory = self.memory[sample_index, :]
b_s = torch.FloatTensor(b_memory[:, :N_STATES])
b_a = torch.LongTensor(b_memory[:, N_STATES:N_STATES + 1].astype(int))
b_r = torch.FloatTensor(b_memory[:, N_STATES + 1:N_STATES + 2])
b_s_ = torch.FloatTensor(b_memory[:, -N_STATES:])
# 计算q_target,loss，反向传递
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()
dqn = DQN()
# training...
print('\nCollecting experience....')
for i_episode in range(400):
s = env.reset()
ep_r = 0
while True:
env.render()
a = dqn.choose_action(s)
s_, r, done, info = env.step(a)
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
dqn.store_transition(s, a, r, s_)
ep_r += r
print('memory_counter:', dqn.memory_counter)
if dqn.memory_counter > MEMORY_CAPACITY:
print('\nStart learning....', i_episode, '--')
dqn.learn()
if done:
print('Ep:', i_episode, '| Ep_r:', round(ep_r, 2))
if done:
break
s = s_
神经网络类
# 定义神经网络
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 一个隐层，一个输出层
self.fc1 = nn.Linear(N_STATES, 100)
self.fc1.weight.data.normal_(0, 0.1)
self.out = nn.Linear(100, N_ACTIONS)
self.out.weight.data.normal_(0, 0.1)
def forward(self, x):
# Net的执行逻辑 Linear_fc1 --> relu --> out --> actions_value
x = self.fc1(x)
x = F.relu(x)
actions_value = self.out(x)
return actions_value继承Module，实现__init__,forward两个方法。init主要是定义Net结构，forward主要给出Net的执行逻辑即流程。
这里的网络十分简单，一个线形层和一个输出层。
Net的执行逻辑： input state --> Linear_fc1 --> relu激活函数 --> out --> actions_value
输入为一个状态，通过神经网络后输出该状态下所有动作值。
DQN定义
class DQN(object):
def __init__(self):
# 模型初始化。初始化main net和target net
# 设置target更新计数器、存储计数器、记忆库初始化、优化器和loss定义。
self.eval_net, self.target_net = Net(), Net()
self.learn_step_counter = 0
self.memory_counter = 0
self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))
self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
self.loss_func = nn.MSELoss()
def choose_action(self, x):
x = torch.unsqueeze(torch.FloatTensor(x), 0)
if np.random.uniform() < EPSILON:
actions_value = self.eval_net.forward(x)
action = torch.max(actions_value, 1)[1].data.numpy()
action = action[0] if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
else:
action = np.random.randint(0, N_ACTIONS)
action = action if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
return action
def store_transition(self, s, a, r, s_):
transition = np.hstack((s, [a, r], s_))
index = self.memory_counter % MEMORY_CAPACITY
self.memory[index,:] = transition
self.memory_counter += 1
def learn(self):
if self.learn_step_counter % TARGET_REPLACE_ITER == 0:
self.target_net.load_state_dict(self.eval_net.state_dict())
print('target update....')
self.learn_step_counter += 1
sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)
b_memory = self.memory[sample_index, :]
b_s = torch.FloatTensor(b_memory[:, :N_STATES])
b_a = torch.LongTensor(b_memory[:, N_STATES:N_STATES + 1].astype(int))
b_r = torch.FloatTensor(b_memory[:, N_STATES + 1:N_STATES + 2])
b_s_ = torch.FloatTensor(b_memory[:, -N_STATES:])
# 计算q_target,loss，反向传递
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()包括：__init__、choose_action、store_transition、learn四个方法的实现；
主循环
for i_episode in range(400):
s = env.reset()
ep_r = 0
while True:
env.render()
a = dqn.choose_action(s)
s_, r, done, info = env.step(a)
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
dqn.store_transition(s, a, r, s_)
ep_r += r
print('memory_counter:', dqn.memory_counter)
if dqn.memory_counter > MEMORY_CAPACITY:
print('\nStart learning....', i_episode, '--')
dqn.learn()
if done:
print('Ep:', i_episode, '| Ep_r:', round(ep_r, 2))
if done:
break
s = s_设置进行400个episode，首先reset初始化环境对象env，ep_r为累计期望reward；
while循环：
每次循环开始都重新生成图像，然后让dqn在状态s下选择动作，这里得到当前状态s下action值最大的动作，我们再在env环境中执行这一动作a，得到执行动作a后的下一个状态s_和奖赏reward；由于模拟的是小车立杆，所以状态中包含四个变量 x, x_dot, theta, theta_dot。这里莫烦对r进行了修正，然后将（s,a,r,s_）存储到记忆库中，并累计期望奖赏。
这里的两个if语句，第一个是要在dqn收集了一定量的经验后才开始学习；第二个是这一轮结束开始下一个episode；
否则就接着下一状态继续训练