【DQN】基于pytorch的强化学习算法Demo

简介

DQN（Deep Q-Network）是一种基于深度神经网络的强化学习算法，于2013年由DeepMind提出。它的目标是解决具有离散动作空间的强化学习问题，并在多个任务中取得了令人瞩目的表现。

DQN的核心思想是使用深度神经网络来逼近状态-动作值函数（Q函数），将当前状态作为输入，输出每个可能动作的Q值估计。通过不断迭代和更新网络参数，DQN能够逐步学习到最优的Q函数，并根据Q值选择具有最大潜在回报的动作。

DQN的训练过程中采用了两个关键技术：经验回放和目标网络。经验回放是一种存储并重复使用智能体经历的经验的方法，它可以破坏数据之间的相关性，提高训练的稳定性。目标网络用于解决训练过程中的估计器冲突问题，通过固定一个与训练网络参数较为独立的目标网络来提供稳定的目标Q值，从而减少训练的不稳定性。

DQN还采用了一种策略称为epsilon-贪心策略来在探索和利用之间进行权衡。初始时，智能体以较高的概率选择随机动作（探索），随着训练的进行，该概率逐渐降低，让智能体更多地依靠Q值选择最佳动作（利用）。

DQN在许多复杂任务中取得了显著的成果，特别是在Atari游戏等需要视觉输入的任务中。它的成功在很大程度上得益于深度神经网络的强大拟合能力和经验回放的效果，使得智能体能够通过与环境的交互进行自主学习。

代码

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-v1',render_mode="human")
#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, 50)
        self.fc1.weight.data.normal_(0, 0.1)   # initialization
        self.out = nn.Linear(50, N_ACTIONS)
        self.out.weight.data.normal_(0, 0.1)   # initialization

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        actions_value = self.out(x)
        return actions_value


class DQN(object):
    def __init__(self):
        self.eval_net, self.target_net = Net(), Net()

        self.learn_step_counter = 0                                     # for target updating
        self.memory_counter = 0                                         # for storing memory
        self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))     # initialize memory
        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)
        # input only one sample
        if np.random.uniform() < EPSILON:   # greedy
            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)  # return the argmax index
        else:   # random
            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_))
        # replace the old memory with new memory
        index = self.memory_counter % MEMORY_CAPACITY
        self.memory[index, :] = transition
        self.memory_counter += 1

    def learn(self):
        # target parameter update
        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

        # sample batch transitions
        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_eval w.r.t the action in experience
        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)
        q_next = self.target_net(b_s_).detach()     # detach from graph, don't backpropagate
        q_target = b_r + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)   # shape (batch, 1)
        loss = self.loss_func(q_eval, q_target)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

dqn = DQN()  # 创建 DQN 对象

print('\nCollecting experience...')
for i_episode in range(400):  # 进行 400 个回合的训练
    s, info = env.reset()  # 环境重置，获取初始状态 s 和其他信息
    ep_r = 0  # 初始化本回合的总奖励 ep_r 为 0
    while True:
        env.render()  # 显示环境,通过调用 render() 方法，可以将当前环境的状态以图形化的方式呈现出来.
        a = dqn.choose_action(s)  # 根据当前状态选择动作 a
        # 下一个状态（nextstate）：返回智能体执行动作a后环境的下一个状态。在示例中，它存储在变量s_中。奖励（reward）：返回智能体执行动作a后在环境中获得的奖励。在示例中，它存储在变中。
        # 完成标志（doneflag）：返回一个布尔值，指示智能体是否已经完成了当前环境。在示例中，它存储在变量done中。
        # 截断标志（truncatedflag）：返回一个布尔值，表示当前状态是否是由于达到了最大时间步骤或其他特定条件而被截断。在示例中，它存储在变量truncated中。
        # 其他信息（info）：返回一个包含其他辅助信息的字典或对象。在示例中，它存储在变量info中。
        # 执行动作，获取下一个状态 s_，奖励 r，done 标志位，以及其他信息
        s_, r, done, truncated, info = env.step(a)

        # 修改奖励值
        #根据智能体在x方向和theta方向上与目标位置的偏离程度，计算两个奖励值r1和r2。具体计算方法是将每个偏离程度除以相应的阈值，然后减去一个常数（0.8和0.5）得到奖励值。这样，如果智能体在这两个方向上的偏离程度越小，奖励值越高。
        x, x_dot, theta, theta_dot = s_  # 从 s_ 中提取参数
        r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8  # 根据 x 的偏离程度计算奖励 r1
        r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5  # 根据 theta 的偏离程度计算奖励 r2
        r = r1 + r2  # 组合两个奖励成为最终的奖励 r

        dqn.store_transition(s, a, r, s_)  # 存储状态转换信息到经验池

        ep_r += r  # 更新本回合的总奖励

        if dqn.memory_counter > MEMORY_CAPACITY:  # 当经验池中的样本数量超过阈值 MEMORY_CAPACITY 时进行学习
            dqn.learn()

            if done:  # 如果本回合结束
                print('Ep: ', i_episode,
                      '| Ep_r: ', round(ep_r, 2))  # 打印本回合的回合数和总奖励

        if done:  # 如果任务结束
            break  # 跳出当前回合的循环

        s = s_  # 更新状态，准备进行下一步动作选择