简介:这份资源围绕指针网络求解旅行商问题(TSP)的强化学习实现展开,面向具备一定Python与深度学习基础、希望动手复现组合优化算法的开发者与研究者。代码以最佳路径长度作为critic值,省去额外critic网络,训练样本由[0,1]×[0,1]网格均匀采样生成,最优解借助Concorde求解,需将其加入系统PATH。压缩包共14个文件,约4.01MB,包含8个py源码(模型、训练器、数据加载、层定义与配置等)、2个npz测试数据集、2张结果图、1份readme说明及gitignore,结构清晰便于按模块阅读。资源中给出TSP10在10万步训练后的测试结果,并以diff指标衡量强化学习解与最优解的差距,读者可据此评估模型收敛与泛化表现。目前已有867人学习下载,适合作为组合优化与强化学习交叉方向的入门实践参考。
1. 指针网络做 TSP 强化学习:为什么它比传统启发式更值得投入
如果你做过 TSP(旅行商问题)求解,大概率经历过这样的场景:城市规模从 20 涨到 100,原本跑得好好的遗传算法突然变得不稳定,换 CPLEX 又嫌太重,手写 2-opt 邻域搜索调参调到怀疑人生。指针网络(Pointer Network)配合强化学习,恰好切入了这个痛点——它不需要标注好的最优解,直接用奖励信号驱动模型学会输出城市访问序列,推理时一次前向传播就能给出解,速度比迭代式启发式快几个数量级。
这个方向适合两类人:一是想入门深度强化学习但不想碰游戏环境的 Python 开发者,TSP 的奖励函数天然清晰,没有稀疏奖励的玄学问题;二是做组合优化落地的工程师,需要一套能泛化到不同城市规模的求解框架。本文从零搭建一套可运行的指针网络 + REINFORCE 训练流程,覆盖数据生成、模型定义、训练循环、贪心与采样解码、避坑排查,最后给出规模化验证的实用技巧。代码基于 PyTorch,Python 3.8 以上即可跑通。
2. 指针网络与 REINFORCE 的配合逻辑:为什么不用交叉熵
2.1 指针网络解决的是「输出字典随输入变化」的问题
标准 Seq2Seq 模型在解码时,输出词表是固定的。TSP 不一样:输入 10 个城市,输出就是 10 个位置的排列;输入 50 个城市,输出就是 50 个位置的排列。词表大小随输入变化,固定 softmax 层没法处理。
指针网络的核心改动是:解码每一步,不再从固定词表选 token,而是用注意力机制计算当前解码状态与所有编码器隐藏状态的匹配分数,归一化后作为指向输入位置的指针概率分布。数学上,第 $t$ 步指向城市 $i$ 的概率为:
$$p(i \mid \text{context}) = \text{softmax}(u_i)$$
其中 $u_i = v^T \tanh(W_1 e_i + W_2 d_t)$,$e_i$ 是编码器对城市 $i$ 的输出,$d_t$ 是解码器当前状态。这个设计让模型天然支持变长输入,且输出必然是输入的一个排列(配合 mask 机制)。
2.2 为什么用 REINFORCE 而不是监督学习
监督学习需要标注最优解,而 TSP 最优解在 50 城市以上就极难获取。REINFORCE 属于策略梯度方法,直接用路径长度的负值作为奖励:
$$\nabla_\theta J(\theta) \approx \frac{1}{B} \sum_{b=1}^{B} (L(\tau_b) - b) \nabla_\theta \log p_\theta(\tau_b)$$
其中 $L(\tau_b)$ 是第 $b$ 条采样路径的总长度,$b$ 是基线(baseline),用于降低方差。常见做法是用贪心解码的路径长度作为基线,这样不需要额外训练 Critic 网络,实现简单且效果稳定。
注意:基线不参与梯度回传,只做数值减法。如果用 Critic 网络做基线,需要 detach 后再减。
2.3 最小可运行代码:数据生成与模型定义
先解决数据。TSP 实例生成很简单:在单位正方形内均匀采样 N 个点。
import torch import torch.nn as nn import torch.nn.functional as F def generate_tsp_data(batch_size, num_cities, device): """生成 batch_size 个 TSP 实例,每个实例 num_cities 个城市坐标""" # 坐标范围 [0, 1],形状 (batch_size, num_cities, 2) return torch.rand(batch_size, num_cities, 2, device=device)指针网络模型分编码器和解码器两部分。编码器用 LSTM 或 Transformer 均可,这里用 LSTM 做最小实现。
class PointerNetwork(nn.Module): def __init__(self, input_dim=2, hidden_dim=128): super().__init__() self.hidden_dim = hidden_dim # 编码器:将城市坐标序列编码为隐藏状态 self.encoder = nn.LSTM(input_dim, hidden_dim, batch_first=True) # 解码器:输入是上一步选中的城市坐标 + 上一步隐藏状态 self.decoder = nn.LSTMCell(input_dim, hidden_dim) # 注意力参数:将编码器输出和解码器状态映射为指针分数 self.W1 = nn.Linear(hidden_dim, hidden_dim, bias=False) self.W2 = nn.Linear(hidden_dim, hidden_dim, bias=False) self.v = nn.Linear(hidden_dim, 1, bias=False) def forward(self, x, decode_type='sampling'): """ x: (batch, num_cities, 2) decode_type: 'sampling' 用于训练,'greedy' 用于基线和推理 """ batch_size, num_cities, _ = x.shape # 编码 encoder_out, (h, c) = self.encoder(x) # encoder_out: (B, N, H) # 解码器初始状态用编码器最后一步的隐藏状态 decoder_h = h.squeeze(0) # (B, H) decoder_c = c.squeeze(0) # 初始输入:一个可学习的起始向量,这里简化为全零 decoder_input = torch.zeros(batch_size, 2, device=x.device) # 记录已访问城市,防止重复选择 mask = torch.zeros(batch_size, num_cities, device=x.device) # 记录路径 pointers = [] for _ in range(num_cities): decoder_h, decoder_c = self.decoder(decoder_input, (decoder_h, decoder_c)) # 注意力分数计算 query = self.W2(decoder_h).unsqueeze(1) # (B, 1, H) keys = self.W1(encoder_out) # (B, N, H) scores = self.v(torch.tanh(query + keys)).squeeze(-1) # (B, N) # 已访问城市分数置为极小值 scores = scores.masked_fill(mask == 1, -1e9) probs = F.softmax(scores, dim=-1) if decode_type == 'greedy': idx = probs.argmax(dim=-1) else: idx = torch.multinomial(probs, 1).squeeze(-1) pointers.append(idx) mask = mask.scatter(1, idx.unsqueeze(1), 1) # 下一步输入是当前选中城市的坐标 decoder_input = x[torch.arange(batch_size), idx] return torch.stack(pointers, dim=1) # (B, N)这段代码里几个关键点:mask保证每个城市只被访问一次;decode_type控制训练时用采样、评估时用贪心;注意力分数计算采用加性注意力,比点积注意力更适合小规模 TSP。参数hidden_dim建议从 128 起步,城市数超过 50 时加到 256。
3. 训练循环与奖励设计:从随机路径到稳定收敛
3.1 奖励函数与损失计算
TSP 的奖励就是路径长度的负值。给定指针序列,计算总距离:
def compute_tour_length(x, pointers): """ x: (B, N, 2) 城市坐标 pointers: (B, N) 访问顺序 返回: (B,) 每条路径的总长度 """ batch_size, num_cities, _ = x.shape # 按 pointers 重排城市坐标 ordered = x[torch.arange(batch_size).unsqueeze(1), pointers] # (B, N, 2) # 计算相邻城市距离,包括首尾闭合 diff = ordered - ordered.roll(-1, dims=1) distances = diff.norm(dim=-1) # (B, N) return distances.sum(dim=-1)REINFORCE 损失需要 log 概率。上面的 forward 只返回了 pointers,需要额外记录 log_probs。修改 forward 增加返回值:
# 在 forward 循环内,选择 idx 后追加: log_prob = torch.log(probs.gather(1, idx.unsqueeze(1)) + 1e-9) log_probs.append(log_prob.squeeze(1)) # 循环结束后返回 return torch.stack(pointers, dim=1), torch.stack(log_probs, dim=1).sum(dim=1)训练循环:
def train_step(model, optimizer, x): model.train() # 采样解码,获得一条路径及其 log 概率 pointers, log_prob = model(x, decode_type='sampling') tour_length = compute_tour_length(x, pointers) # 贪心解码作为基线,不计算梯度 with torch.no_grad(): greedy_pointers, _ = model(x, decode_type='greedy') baseline = compute_tour_length(x, greedy_pointers) # REINFORCE 损失:最大化 (baseline - tour_length) * log_prob advantage = (baseline - tour_length).detach() loss = (advantage * log_prob).mean() # 等价于最小化负的期望奖励 optimizer.zero_grad() loss.backward() # 梯度裁剪,防止 LSTM 梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() return tour_length.mean().item(), baseline.mean().item()参数说明:advantage用baseline - tour_length,因为路径越短奖励越高,取负号后作为损失方向;clip_grad_norm_的max_norm=1.0是经验值,TSP 训练中梯度范数经常冲到 10 以上,不裁剪容易发散。
3.2 训练超参与收敛判断
完整训练脚本:
def train(num_epochs=100, batch_size=256, num_cities=20, lr=1e-3): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = PointerNetwork(hidden_dim=128).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=lr) for epoch in range(num_epochs): x = generate_tsp_data(batch_size, num_cities, device) sample_len, greedy_len = train_step(model, optimizer, x) if epoch % 10 == 0: print(f"Epoch {epoch:3d} | Sample: {sample_len:.4f} | Greedy: {greedy_len:.4f}") return model超参建议表:
| 参数 | 推荐值 | 调整方向 |
|---|---|---|
| hidden_dim | 128(N≤20)/ 256(N>20) | 太小欠拟合,太大过拟合且慢 |
| learning_rate | 1e-3 | 不收敛降到 1e-4 |
| batch_size | 256 | 显存不够降到 64,但梯度噪声变大 |
| num_epochs | 100~200 | 看 greedy 长度是否平稳 |
| clip max_norm | 1.0 | 梯度爆炸时降到 0.5 |
收敛判断看贪心解码的路径长度:前 20 个 epoch 快速下降,之后缓慢改善。如果 50 个 epoch 后还在震荡,检查学习率是否过大或基线是否失效。
提示:训练时采样解码的路径长度通常比贪心差 5%~15%,这是正常的探索代价。如果两者差距超过 30%,说明策略方差太大,可以增大 batch_size 或改用 rollout baseline。
4. 避坑与排查:指针网络训练 TSP 最常见的 5 个翻车现场
4.1 损失不降反升,路径长度爆炸
现象:训练几个 epoch 后,采样路径长度从 4.0 涨到 20 以上,贪心解码也同步恶化。
原因:REINFORCE 的梯度方差过大,加上 LSTM 梯度爆炸,参数更新方向完全随机。常见触发条件是学习率设成 1e-2 或没有梯度裁剪。
解决:学习率降到 1e-3 或 1e-4,加clip_grad_norm_(max_norm=1.0)。如果还不行,把基线从贪心解码改成指数移动平均(EMA)基线,平滑效果更好。
4.2 模型学会「摆烂」:所有路径都指向同一个城市
现象:指针序列出现大量重复索引,mask 机制似乎失效。
原因:mask 的scatter操作写错了维度,或者masked_fill的值不够小(比如用了 -1e4 但 logits 量级到了 1e5)。另一个可能是解码器初始输入全零导致第一步注意力均匀分布,后续陷入局部循环。
解决:检查mask.scatter(1, idx.unsqueeze(1), 1)中 idx 的形状必须是(B, 1);masked_fill用-1e9而不是-1e4;解码器初始输入改用可学习的参数向量,不要用全零。
4.3 训练集表现好,换一组随机城市就崩
现象:在固定随机种子的 TSP 实例上路径长度 3.8,换一组种子变成 6.5。
原因:模型过拟合到了特定坐标分布。指针网络本身有泛化能力,但训练时如果 batch 内城市分布太集中(比如都挤在角落),编码器学到的特征没有覆盖全空间。
解决:确保torch.rand生成的坐标覆盖[0,1]×[0,1]全空间;每轮重新生成数据,不要固定一个 batch 反复训练;如果城市数可变,训练时混合不同 N 的实例(如 15、20、25 交替)。
4.4 GPU 显存溢出,batch_size 降到 1 才能跑
现象:N=50 时 batch_size=256 直接 OOM,降到 64 还是不够。
原因:指针网络的注意力分数矩阵是(B, N, H),解码 N 步后中间变量累积。加上 LSTM 的隐藏状态,显存占用是 $O(B \cdot N \cdot H)$。
解决:用梯度累积模拟大 batch:batch_size=32,累积 8 次梯度再更新,等效 batch_size=256。或者把编码器换成 Transformer,注意力计算可以分块。
4.5 贪心解码结果比采样还差
现象:训练日志里 greedy 长度始终高于 sample 长度。
原因:模型还没收敛时,贪心解码容易陷入局部最优,而采样有随机性反而能跳出。这不是 bug,是训练早期的正常现象。
解决:继续训练,通常 30 个 epoch 后贪心会反超。如果 100 个 epoch 后仍然如此,说明模型容量不够,增大hidden_dim或加一层编码器 LSTM。
5. 规模化验证与推理加速:从 N=20 到 N=100 的实用技巧
训练完 N=20 的模型,直接拿去做 N=50 或 N=100 的推理,路径长度会明显劣化。指针网络对城市规模有一定泛化能力,但需要配合几个技巧。
技巧一:课程学习(Curriculum Learning)。不要一上来就训 N=50,先从 N=10 开始,每 20 个 epoch 增加 5 个城市,直到目标规模。这样编码器逐步适应更长的序列,最终 N=100 的路径长度比直接训练低 8%~12%。
技巧二:推理时用 beam search 替代贪心。贪心每步只选概率最大的城市,beam search 保留 top-k 条候选路径。k=3 时推理时间增加不到 2 倍,但路径长度平均改善 3%~5%。实现上维护 k 个解码状态,每步扩展后按累积 log 概率排序取前 k。
技巧三:坐标归一化与尺度不变性。训练时坐标在[0,1],推理时如果输入坐标范围是[0,100],路径长度会放大 100 倍,但模型输出的指针序列不变。所以推理前务必把坐标归一化到[0,1],否则注意力分数的数值范围偏移会导致选择错误。
验证方法:用随机生成的 1000 个 TSP 实例,分别跑贪心解码和 beam search,统计平均路径长度和标准差。如果标准差超过均值的 15%,说明模型在某些分布上不稳定,需要增加训练数据多样性。
def evaluate(model, num_cities=50, num_instances=1000, beam_width=3): """在随机实例上评估模型""" device = next(model.parameters()).device model.eval() total_length = 0.0 with torch.no_grad(): for _ in range(num_instances // 100): x = generate_tsp_data(100, num_cities, device) pointers, _ = model(x, decode_type='greedy') lengths = compute_tour_length(x, pointers) total_length += lengths.sum().item() return total_length / num_instances我自己的习惯是:每次改完模型结构,先跑 N=20 的 100 个 epoch 看收敛曲线,确认没有震荡后再上课程学习训 N=100。血泪经验是不要跳过小规模验证直接怼大规模,否则调参调到最后都不知道是模型问题还是数据问题。希望帮到你。
本文还有配套的精品资源,点击获取