pytorch——通用的低耦合度的DQN模型

使用须知

使用时需要：

1. 向模型提供运行环境（env），env需要包括以下方法：
   reset() 重置环境
   step() 根据动作更新环境得到反馈
   render() 展示环境，可以为空方法
2. 传参：
   state_dim 表示状态的维度
   action_num 表示动作空间中包含动作的数量
3. 调整Q-net结构
4. 调参

模型代码

测试部分我们使用的是gym中平衡车环境：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/5/17 14:56
# @Author  : Liu Lihao
# @File    : DQN.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import collections
from torch import optim
import numpy as np

'''
通用的DQN模型
'''

'''
参数设置
'''
learning_rate = 0.0002 #0.0001
N_episode = 10000  # 训练回合数
max_epsilon = 0.08  # 最大贪婪度
min_epsilon = 0.001  # 最小贪婪度
T = 600  # 一个episode包含的步长
train_threshold = 2000
update_epoch = 10  # 每次调用神经网络更新函数更新次数
copy_time = 50  # 太大模型会陷入过拟合，太小则模型收敛速度过慢
SAVE_PATH = 'DQNmodel/dqn2.pt'  # 模型保存路径
batch_size = 64
buffer_limit = 50000  # 经验池最大容量
gamma = 0.99  # 预测的下一状态的最大动作值折合到reward中的衰减程度
updateStartSize = 2000  # 当经验池大于这个值之后再开始更新
print_interval = 20  # 打印各种参数实时状态的间隔

'''
构建Q网络
使用全连接神经网络
策略采用 epsilon 贪婪策略，有 epsilon 的概率选择随机动作， 1-epsilon 的概率选择贪婪动作。
'''
class Qnet(nn.Module):
    def __init__(self, state_dim, action_num):
        super(Qnet, self).__init__()
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256, action_num)
        self.action_num = action_num

    def forward(self, x):
        x = F.leaky_relu(self.fc1(x))
        x = self.fc2(x)
        return x

    # 选择动作
    def sample_action(self, obs, epsilon):
        out = self.forward(obs)
        coin = random.random()  # 生成0和1之间的随机浮点数
        if coin < epsilon:
            # 随机返回一个动作
            return random.randint(0, self.action_num-1)
        else:
            # 返回最大Q值对应的动作
            return out.argmax().item()


'''
构建经验回放池
经验回放池实际上是一个队列，当经验回放池满时，会抛弃旧的经验值，加入新采样的经验值。
采样时，从经验回放池中随机抽取batch_size个经验值作为一个transition返回给训练机进行学习。
'''
class ReplayBuffer():
    def __init__(self):
        self.buffer = collections.deque(maxlen=buffer_limit)  # 双端队列

    # 插入经验池
    def put(self, transition):
        self.buffer.append(transition)

    # 从经验池中抽样
    def sample(self, n):
        mini_batch = random.sample(self.buffer, n)  # 随机抽取n个样本
        s_lst, a_lst, r_lst, s_prime_lst, done_mask_lst = [], [], [], [], []

        for transition in mini_batch:
            s, a, r, s_prime, done_mask = transition
            s_lst.append(s)
            a_lst.append([a])
            r_lst.append([r])
            s_prime_lst.append(s_prime)
            done_mask_lst.append([done_mask])

        '''
        s_lst: [batch_size,input_size]
        a_lst: [batch_size,1]
        r_lst: [batch_size,1]
        '''

        return torch.tensor(s_lst, dtype=torch.float),\
               torch.tensor(a_lst),\
               torch.tensor(r_lst, dtype=torch.float),\
               torch.tensor(s_prime_lst, dtype=torch.float),\
               torch.tensor(done_mask_lst)

    # 返回buffer长度
    def size(self):
        return len(self.buffer)


class DQNModel:
    def __init__(self, env, state_dim, action_num):
        self.env = env
        self.q = Qnet(state_dim, action_num)
        self.q_target = Qnet(state_dim, action_num)
        self.memory = ReplayBuffer()
        self.optimizer = optim.Adam(self.q.parameters(), lr=learning_rate)

    '''
    梯度下降更新参数
    同时对batch_size个经验样本 进行更新
    '''
    def update(self):
        LOSS = []
        for i in range(update_epoch):
            s, a, r, s_prime, done_mask = self.memory.sample(batch_size)
            q_out = self.q(s)  # 将状态s输入神经网络得到输出
            q_a = q_out.gather(1, a)  # 在状态s处执行a动作对应的q值
            max_q_prime = self.q_target(s_prime).max(1)[0].unsqueeze(1)  # 获取s_prime状态（即下一个状态）可能会得到的回报
            target = r + gamma * max_q_prime * done_mask  # 计算真实回报
            loss = F.smooth_l1_loss(q_a, target)  # 计算loss
            LOSS.append(loss.item())
            '''
            shape：
            q_out:[batch_size, action_number]
            q_a, max_q_prime, target:[batch_size, 1]
            loss:[]
            '''

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
        return np.mean(LOSS)


    '''
    构建训练模型
    此代码设置的是每个 episode 更新一次。也可以选择每个时间步长更新一次
    '''
    def train(self):

        # checkpoint = torch.load(SAVE_PATH)
        # self.q.load_state_dict(checkpoint['q_state_dict'])

        self.q_target.load_state_dict(self.q.state_dict())
        score = 0.0

        for n_epi in range(N_episode):

            epsilon = max(min_epsilon, max_epsilon - (max_epsilon - min_epsilon) * (n_epi / N_episode))  # 贪婪度
            s = self.env.reset()  # 重新设置环境

            '''探索'''
            for t in range(T):
                a = self.q.sample_action(torch.from_numpy(s).float(), epsilon)  # （状态, 贪婪度）
                s_prime, r, done, info = self.env.step(a)
                done_mask = 0.0 if done else 1.0
                self.memory.put((s, a, r, s_prime, done_mask))
                s = s_prime

                score += r
                if done:
                    break

            '''更新'''
            loss = 0
            # 经验池长度大于updateStartSize以后开始更新
            if self.memory.size() > updateStartSize:
                loss = self.update()

            # 将q网络中的参数拷贝到q_target网络中
            if n_epi % copy_time == 0:
                self.q_target.load_state_dict(self.q.state_dict())

            if n_epi % print_interval == 0 and n_epi != 0:
                print("# of episode :{}, avg score : {:.1f}, loss : {}, buffer size : {}, epsilon : {:.1f}%".format(
                    n_epi, score / print_interval, loss, self.memory.size(), epsilon * 100))
                score = 0.0


        # 保存q网络、q_target网络、优化器梯度的参数
        torch.save(
            {
                'q_state_dict': self.q.state_dict(),
                'q_target_state_dict': self.q_target.state_dict(),
                'optim_state_dict': self.optimizer.state_dict()
            }, SAVE_PATH)


    '''
    运行DDPG
    '''
    def evaluate(self):
        checkpoint = torch.load(SAVE_PATH)
        self.q.load_state_dict(checkpoint['q_state_dict'])
        # q.eval()
        while(1):
            s = self.env.reset()
            done = False
            while not done:
                a = self.q.sample_action(torch.from_numpy(s).float(), min_epsilon)
                s_prime, r, done, info = self.env.step(a)
                self.env.render()  # 图形化显示
                s = s_prime




'''
测试
'''
# import gym
# test = DQNModel(env = gym.make('CartPole-v1'),state_dim = 4, action_num = 2)
# test.train()
# test.evaluate()