大模型- 强化学习-Deep-Q-learning-pytorch实现--83
参考
https://newfacade.github.io/notes-on-reinforcement-learning/10-dqn-torch.html
https://g.co/gemini/share/834917e5f7e3
内容
这是一个在强化学习领域非常经典和著名的任务，叫做 “车杆” (CartPole)，有时也被称为 “倒立摆” (Inverted Pendulum)。
可以把它看作一个简单的小游戏，是很多强化学习算法入门的“Hello, World!”。
游戏目标 (Goal)
这个任务的目标非常明确：
控制下方的小车（cart）左右移动，来保持竖立在它上面的杆子（pole）不会倒下，并让这个平衡状态维持尽可能长的时间。
游戏规则 (Rules)
动作 (Actions)：在任何时刻，控制小车的智能体（agent）只有两种选择：向左推小车或向右推小车。
奖励 (Reward)：只要杆子保持在直立状态，每经过一个时间步（timestep），智能体就会获得 +1 的奖励。这意味着，坚持的时间越长，累积的奖励就越高。
结束条件 (Termination)：游戏（或称为一个“回合”，episode）在以下两种情况之一发生时结束：
杆子倾斜的角度过大（倒了）。
小车移动超出了预设的边界范围（离中心太远）。
这个任务很有名：
经典基准：它是测试和验证各种强化学习算法（如 Q-learning, DQN, Policy Gradients 等）性能的一个标准环境。
简单直观：它的规则简单，目标明确，很容易让人理解强化学习的基本概念：状态（state）、动作（action）、奖励（reward）。
恰到好处的难度：它虽然简单，但又不是一个能轻易解决的问题，需要智能体学会“预判”和“精细控制”，是学习和实践强化学习理论的绝佳起点。
代码
背景回顾
任务: 控制小车左右移动，让杆子保持平衡。
状态 (State): 一个包含4个值的向量：[小车位置, 小车速度, 杆子角度, 杆子角速度]。
动作 (Action): 2个离散动作：0 (向左推) 和 1 (向右推)。
目标: 训练一个神经网络（Q-Network），输入是状态，输出是每个动作的Q值，从而根据Q值选择最优动作。
准备工作：环境、参数与经验回放
这部分代码负责初始化环境、设定超参数，并定义用于经验回放 (Experience Replay) 的数据结构。
import gym
import math
import random
import collections
import torch
import torch.nn as nn
import torch.optim as optim
# --- 超参数 ---
BATCH_SIZE = 128 # 每次从Replay Buffer中采样的样本数量
GAMMA = 0.999 # 折扣因子γ
EPS_START = 0.9 # ε-greedy策略的初始探索率
EPS_END = 0.05 # 最终探索率
EPS_DECAY = 200 # 探索率衰减速率
TARGET_UPDATE = 10 # 目标网络更新频率（每10个episodes）
# --- Replay Memory ---
Transition = collections.namedtuple('Transition', ('state', 'action', 'next_state', 'reward'))
class ReplayMemory(object):
def __init__(self, capacity):
self.memory = collections.deque([], maxlen=capacity) # 使用deque实现固定容量的缓冲区
def push(self, *args):
self.memory.append(Transition(*args)) # 存入一个transition
def sample(self, batch_size):
return random.sample(self.memory, batch_size) # 随机采样一个batch
def __len__(self):
return len(self.memory)****
超参数 (Hyperparameters)：
BATCH_SIZE: 在每次训练时，我们从经验池中随机抽取多少条经验数据。这是为了进行小批量梯度下降（mini-batch gradient descent）。
GAMMA (γ): 折扣因子。它出现在Q值计算公式中，决定了未来奖励的重要性。γ 接近1表示我们更看重长期回报。
EPS_...: 这三个参数用于 ε-greedy 策略。探索率 ε 会从 EPS_START (0.9) 开始，随着训练的进行逐渐衰减到 EPS_END (0.05)。这意味着在早期，智能体会有很大几率随机探索（90%）；在后期，它会更多地利用已学到的知识（95%）。
TARGET_UPDATE: 这个参数对应DQN的第二个关键技巧——固定Q目标。它规定了每隔多少个完整的游戏回合（episodes），我们才把主网络（Policy Network）的权重复制给目标网络（Target Network）
经验回放 (ReplayMemory):(人类的大脑 睡眠中 也会对一天中发生的记忆进行整理)
这个类完美地实现了我们在理论中讲到的经验回放池。它使用 collections.deque 结构，这是一个双端队列，当存满后，再有新数据加入时，会自动从另一端丢弃最老的数据。
push 方法：智能体每执行一步，就会得到一个经验元组 (state, action, next_state, reward)，这个方法将该元组存入 memory 中。
sample 方法：这是经验回放的核心。它从 memory 中随机抽取 BATCH_SIZE 个经验，用于后续的训练。这种随机性打破了数据之间的时间相关性，使训练更稳定。
Q-Network 模型定义
这部分定义了作为函数逼近器的神经网络结构。
class DQN(nn.Module):
def __init__(self, h, w, outputs):
super(DQN, self).__init__()
# 在CartPole中，我们直接用向量，所以这里会是一个简单的全连接网络
# 假设输入是状态向量的4个特征
self.fc1 = nn.Linear(4, 128) # 输入层 -> 隐藏层
self.fc2 = nn.Linear(128, 128) # 隐藏层
self.fc3 = nn.Linear(128, outputs) # 隐藏层 -> 输出层
def forward(self, x):
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
return self.fc3(x)
(注意: 原始网页中可能使用卷积层来处理图像输入，但对于CartPole，通常使用上述的全连接网络(Linear layers)。这里的代码是为了解读CartPole这个背景下的实现。)
这是一个非常简单的全连接神经网络（也叫多层感知机, MLP）。
输入: 状态 s，在CartPole中是一个长度为4的向量。[小车位置, 小车速度, 杆子角度, 杆子角速度]
输出: 一个长度为2的向量，分别代表 Q(s, action=0) 和 Q(s, action=1) 的值。2个离散动作：0 (向左推) 和 1 (向右推)
数学关联: 这个网络 DQN(s) 就是我们理论中的函数逼近器 Q(s, a; θ)。输入 s，它会一次性计算出所有可能动作的Q值。参数 θ 就是这个网络中所有全连接层的权重（weights）和偏置（biases）。
整个算法的核心，包含了动作选择和模型优化两大部分。
动作选择 (select_action)
def select_action(state):
global steps_done
sample = random.random()
eps_threshold = EPS_END + (EPS_START - EPS_END) * \
math.exp(-1. * steps_done / EPS_DECAY)
steps_done += 1
if sample > eps_threshold:
with torch.no_grad():
# 利用：选择Q值最大的动作
# policy_net(state) -> [Q(s, a0), Q(s, a1)]
# .max(1)[1] -> 找到最大值的索引（即动作0或1）
return policy_net(state).max(1)[1].view(1, 1)
else:
# 探索：随机选择一个动作
return torch.tensor([[random.randrange(n_actions)]], device=device, dtype=torch.long)
该函数实现了 ε-greedy 策略。
首先计算当前步数下 ε 的值 (eps_threshold)，这个值会随 steps_done 的增加而衰减。
- 利用 (Exploitation): 如果随机数大于 ε，我们就相信网络。(或者是table 记录了 s a对应的Q值)
policy_net(state) 计算出当前状态下两个动作的Q值,
.max(1)[1] 找出其中Q值较大者的索引（0或1）,这个索引就是我们要执行的动作。
torch.no_grad()表示这里不需要计算梯度，可以节省计算资源。 - 探索 (Exploration): 如果随机数小于等于 ε，我们就随机选择一个动作。（随机的往左或者往右）
模型优化 (optimize_model)
最关键的一步，它实现了DQN的损失计算和参数更新。
def optimize_model():
if len(memory) < BATCH_SIZE:
return
transitions = memory.sample(BATCH_SIZE)
batch = Transition(*zip(*transitions)) # 将Transition的batch转换成batch的Transition
# 1. 准备数据
state_batch = torch.cat(batch.state)
action_batch = torch.cat(batch.action)
reward_batch = torch.cat(batch.reward)
next_state_batch = torch.cat(batch.next_state) # 假设非终止状态
# 2. 计算预测Q值: Q(s, a; θ)
# policy_net计算出所有动作的Q值，gather根据action_batch选出实际执行动作的Q值
state_action_values = policy_net(state_batch).gather(1, action_batch)
# 3. 计算目标Q值: y = r + γ * max_{a'} Q(s', a'; θ⁻)
# 使用 target_net 计算 V(s_{t+1})
# detach()使其不参与梯度计算
next_state_values = target_net(next_state_batch).max(1)[0].detach()
expected_state_action_values = (next_state_values * GAMMA) + reward_batch
# 4. 计算损失
criterion = nn.SmoothL1Loss() # 使用Huber Loss
loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))
# 5. 优化
optimizer.zero_grad()
loss.backward()
for param in policy_net.parameters():
param.grad.data.clamp_(-1, 1) # 梯度裁剪，防止梯度爆炸
optimizer.step()
采样与数据整理: 从 memory 中随机采样一个批次的数据，并把 (s, a, r, s') 分别整理成一个 batch
计算预测Q值 (Prediction):
state_action_values = policy_net(state_batch).gather(1, action_batch)
policy_net(state_batch) 输出了一个 [BATCH_SIZE, 2] 的张量,代表了批次中每个状态下两个动作的Q值,
gather 函数则根据 action_batch（记录了当时实际执行的动作）从中精确地挑选出我们需要的那个Q值。
这一步计算的就是损失函数中的预测项 Q(s_j, a_j; θ)
计算目标Q值 (Target):
next_state_values = target_net(next_state_batch).max(1)[0].detach()
这对应了 max_{a'} Q(s'_j, a'; θ⁻)
注意这里用的是 target_net！这是为了保证目标稳定。
.max(1)[0] 得到了下一状态的最大Q值
.detach() 则切断了梯度流,确保我们只更新 policy_net 而不更新 target_net
expected_state_action_values = (next_state_values * GAMMA) + reward_batch
这一步完整地构造出了 TD Target：y_j = r_j + γ * max_{a'} Q(s'_j, a'; θ⁻)。
计算损失
loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))
这就是计算损失函数 L(θ) = E[(y_j - Q(s_j, a_j; θ))^2]。代码中用了 SmoothL1Loss (Huber Loss)，它在误差较小时像均方误差(MSE)，在误差较大时像绝对值误差(MAE)，比纯MSE更稳健，能防止由异常值引起的巨大梯度。
optimizer.zero_grad(), loss.backward(), optimizer.step()
这是标准的PyTorch模型训练流程：清空旧梯度、反向传播计算新梯度、根据梯度更新 policy_net 的参数 θ。
主训练循环
将所有组件串联起来，执行完整的训练流程。
num_episodes = 500
for i_episode in range(num_episodes):
# 初始化环境和状态
state = env.reset()
for t in count():
# 1. 选择并执行动作
action = select_action(state)
next_state, reward, done, _ = env.step(action.item())
# 2. 存入Replay Memory
memory.push(state, action, next_state, reward)
# 3. 状态转移
state = next_state
# 4. 执行一步优化
optimize_model()
if done:
break # 回合结束
# 5. 更新目标网络
if i_episode % TARGET_UPDATE == 0:
target_net.load_state_dict(policy_net.state_dict())
外层循环遍历所有游戏回合（episodes）。
内层循环是单个回合中的每一步。
在每一步，智能体选择动作、与环境交互、获得反馈，然后将经验存入 memory。
接着，立刻调用 optimize_model() 从 memory 中采样并训练网络一步。
当一个回合结束后 (done=True)，内层循环中断。
在外层循环中，每隔 TARGET_UPDATE 个回合，就将 policy_net 的最新权重复制给 target_net，实现 Fixed Q-Targets 的更新。
这份代码是DQN理论非常标准和清晰的工程实现，完美地体现了经验回放和固定Q目标两大核心思想，希望能帮助您更好地理解DQN的运作方式。
全部代码
import math
import random
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count
import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
env = gym.make("CartPole-v1")
# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
from IPython import display
plt.ion()
# if GPU is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward'))
class ReplayMemory(object):
def __init__(self, capacity):
self.memory = deque([], maxlen=capacity)
def push(self, *args):
"""Save a transition"""
self.memory.append(Transition(*args))
def sample(self, batch_size):
return random.sample(self.memory, batch_size)
def __len__(self):
return len(self.memory)
class DQN(nn.Module):
def __init__(self, n_observations, n_actions):
super(DQN, self).__init__()
self.layer1 = nn.Linear(n_observations, 128)
self.layer2 = nn.Linear(128, 128)
self.layer3 = nn.Linear(128, n_actions)
# Called with either one element to determine next action, or a batch
# during optimization. Returns tensor([[left0exp,right0exp]...]).
def forward(self, x):
x = F.relu(self.layer1(x))
x = F.relu(self.layer2(x))
return self.layer3(x)
# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the ``AdamW`` optimizer
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 1000
TAU = 0.005
LR = 1e-4
# Get number of actions from gym action space
n_actions = env.action_space.n
# Get the number of state observations
state, info = env.reset()
n_observations = len(state)
policy_net = DQN(n_observations, n_actions).to(device)
target_net = DQN(n_observations, n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())
optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
memory = ReplayMemory(10000)
steps_done = 0
def select_action(state):
global steps_done
sample = random.random()
eps_threshold = EPS_END + (EPS_START - EPS_END) * \
math.exp(-1. * steps_done / EPS_DECAY)
steps_done += 1
if sample > eps_threshold:
with torch.no_grad():
# t.max(1) will return the largest column value of each row.
# second column on max result is index of where max element was
# found, so we pick action with the larger expected reward.
return policy_net(state).max(1).indices.view(1, 1) # exploitation
else:
return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long) # explore
episode_durations = []
def plot_durations(show_result=False):
plt.figure(1)
durations_t = torch.tensor(episode_durations, dtype=torch.float)
if show_result:
plt.title('Result')
else:
plt.clf()
plt.title('Training...')
plt.xlabel('Episode')
plt.ylabel('Duration')
plt.plot(durations_t.numpy())
# Take 100 episode averages and plot them too
if len(durations_t) >= 100:
means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
means = torch.cat((torch.zeros(99), means))
plt.plot(means.numpy())
plt.pause(0.001) # pause a bit so that plots are updated
if is_ipython:
if not show_result:
display.display(plt.gcf())
display.clear_output(wait=True)
else:
display.display(plt.gcf())
def optimize_model():
if len(memory) < BATCH_SIZE:
return
transitions = memory.sample(BATCH_SIZE)
# Converts batch-array of Transitions to Transition of batch-arrays.
batch = Transition(*zip(*transitions))
# Compute a mask of non-final states and concatenate the batch elements
# (a final state would've been the one after which simulation ended)
non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
batch.next_state)), device=device, dtype=torch.bool)
non_final_next_states = torch.cat([s for s in batch.next_state
if s is not None])
state_batch = torch.cat(batch.state)
action_batch = torch.cat(batch.action)
reward_batch = torch.cat(batch.reward)
# Compute Q(s_t, a) - the model computes Q(s_t), then we select the
# columns of actions taken. These are the actions which would've been taken
# for each batch state according to policy_net
state_action_values = policy_net(state_batch).gather(1, action_batch)
# Compute V(s_{t+1}) for all next states.
# Expected values of actions for non_final_next_states are computed based
# on the "older" target_net; selecting their best reward with max(1).values
# This is merged based on the mask, such that we'll have either the expected
# state value or 0 in case the state was final.
next_state_values = torch.zeros(BATCH_SIZE, device=device)
with torch.no_grad():
next_state_values[non_final_mask] = target_net(non_final_next_states).max(1).values
# Compute the expected Q values
expected_state_action_values = (next_state_values * GAMMA) + reward_batch
# Compute Huber loss
criterion = nn.SmoothL1Loss()
loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))
# Optimize the model
optimizer.zero_grad()
loss.backward()
# In-place gradient clipping
torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)
optimizer.step()
def train(num_episodes=600):
for i_episode in range(num_episodes):
# Initialize the environment and get it's state
state, info = env.reset()
state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
for t in count():
action = select_action(state)
observation, reward, terminated, truncated, _ = env.step(action.item())
reward = torch.tensor([reward], device=device)
done = terminated or truncated
if terminated:
next_state = None
else:
next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)
# Store the transition in memory
memory.push(state, action, next_state, reward)
# Move to the next state
state = next_state
# Perform one step of the optimization (on the policy network)
optimize_model()
# Soft update of the target network's weights
# θ′ ← τ θ + (1 −τ )θ′
target_net_state_dict = target_net.state_dict()
policy_net_state_dict = policy_net.state_dict()
for key in policy_net_state_dict:
target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)
target_net.load_state_dict(target_net_state_dict)
if done:
episode_durations.append(t + 1)
plot_durations()
break
# 保存模型
torch.save(policy_net.state_dict(), "policy_dqn.pth")
print('Complete')
# plot_durations(show_result=True)
# plt.ioff()
# plt.show()
def eval_model():
policy_net.load_state_dict(torch.load("policy_dqn.pth"))
policy_net.eval()
env = gym.make('CartPole-v1', render_mode='rgb_array') # 改为rgb_array以便保存视频
# 设置视频保存
from gymnasium.wrappers import RecordVideo
env = RecordVideo(env, "videos", episode_trigger=lambda x: True)
for i in range(5): # 跑5个回合看看效果
state, _ = env.reset()
for t in range(500): # CartPole-v1最多500步
env.render() # 渲染画面
with torch.no_grad(): # 在评估时不需要计算梯度
q_values = policy_net(torch.from_numpy(state).float().unsqueeze(0))
action = q_values.max(1)[1] # 直接选择Q值最大的动作
state, reward, done, _, _ = env.step(action.item())
if done:
print(f"Episode {i+1} finished after {t+1} timesteps")
break
env.close()
if __name__ == "__main__":
# train()
eval_model()
浙公网安备 33010602011771号