简介面向需要上手强化学习与TSP组合优化研究的开发者这份Python代码包提供了一套基于指针网络的完整训练与测试流程。实现中未单独搭建critic网络而是直接用最优路径长度作为critic值进行策略梯度更新逻辑简洁易懂。训练样本在[0,1]×[0,1]网格内均匀采样2D点最优解由Concorde求解器生成并需将其加入系统PATH环境变量便于复现实验测试覆盖TSP10与TSP50场景通过gap指标对比强化学习解与最优解差距。资源包共14个文件以8个py源码文件为核心涵盖模型、训练器、数据加载、配置与工具函数另有2张结果对比图、2个npz测试数据集规模8192、1份README说明文档及gitignore等整体压缩后4.01MB。目前已有866人学习下载适合作为入门指针网络求解TSP的参考实现通过对照代码与README可快速掌握数据生成、模型训练、基准测试及结果评估方法。1. 指针网络 强化学习为什么是 TSP 最值得入门的解法组合优化问题里TSP 是个绕不开的靶子。城市数量一上去精确解就指数级爆炸传统启发式又要为每个新实例重新设计邻域搜索策略。指针网络Pointer Networks走的是一条完全不同的路它不靠人工设计的邻域结构而是让模型从数据里学会“下一步该指向哪个城市”并用强化学习来训练这个指向策略。这个思路最早能追溯到 Vinyals 等人的工作但其真正的工程价值在于你不再需要为每个 TSP 实例调用外部的求解器模型前向一次就能给出一个解。更反直觉的是完全不用监督标签也能训练出不错的解——这正是 REINFORCE 这类策略梯度方法的用武之地。对于想进入深度强化学习与组合优化交叉领域的工程师来说这几乎是门槛最低、闭环最完整的实践项目。本文基于 PyTorch 实现一套可在本地跑通的完整代码覆盖从模型设计、数据生成到训练与可视化验证的全部环节。读者需要具备基本的 Python、PyTorch 和深度学习概念基础读完可以用同一套框架迁移到车辆路径规划、序列排序等更广泛的离散决策问题上。2. 指针网络的核心把“选择城市”建模为注意力上的分布2.1 从 Sequence-to-Sequence 到 Pointer输出不再是词表而是输入序列本身传统 Seq2Seq 模型比如机器翻译在解码阶段会从一个固定的词表里产生 softmax 分布然后选词。但 TSP 的解是输入城市的一个排列城市集合随实例变化而变化。如果硬套 Seq2Seq就得为每种城市数量训练一个不同输出维度的模型泛化性极差。指针网络的关键改动在解码器的最后一步不把隐状态映射到词表而是把它和编码器输出的所有城市嵌入做点积注意力得到一个覆盖所有输入位置的分布。这个分布表示“下一个访问哪个城市”的概率。换句话说模型每一步输出的是指向输入序列某个元素的“指针”。import torch import torch.nn as nn import torch.nn.functional as F class PointerDecoder(nn.Module): def __init__(self, hidden_dim): super().__init__() self.W1 nn.Linear(hidden_dim, hidden_dim, biasFalse) self.W2 nn.Linear(hidden_dim, hidden_dim, biasFalse) self.v nn.Linear(hidden_dim, 1, biasFalse) def forward(self, encoder_outputs, decoder_hidden, maskNone): # encoder_outputs: (batch, seq_len, hidden_dim) # decoder_hidden: (batch, hidden_dim) energy self.v(torch.tanh(self.W1(encoder_outputs) self.W2(decoder_hidden).unsqueeze(1))) # 去掉最后一维得到 (batch, seq_len) logits energy.squeeze(-1) if mask is not None: logits logits.masked_fill(mask 0, float(-inf)) return F.softmax(logits, dim-1)这段代码实现了注意力打分与前向传播。W1对编码器的每个城市嵌入做线性变换W2对解码器当前隐状态做变换后加在每座城市上形成逐位置的注意力能量v把能量压缩成标量。mask参数用于屏蔽已访问过的城市确保模型不会重复选同一个点。对 TSP 这种每个城市只能访问一次的约束来说mask 是保证可行性最关键的部分。2.2 静态欧式 TSP 的编码器选型与位置嵌入静态欧式 TSP 的输入是二维平面上的点集没有顺序概念。这决定了编码器不能用 LSTM——LSTM 本身对顺序敏感而 TSP 输入本质上是集合。更稳妥的做法是使用 Transformer Encoder通过自注意力让每个城市感知全局信息。class TSPEncoder(nn.Module): def __init__(self, input_dim, hidden_dim, num_layers3): super().__init__() self.embed nn.Linear(input_dim, hidden_dim) encoder_layer nn.TransformerEncoderLayer(d_modelhidden_dim, nhead8, batch_firstTrue) self.encoder nn.TransformerEncoder(encoder_layer, num_layersnum_layers) def forward(self, points): # points: (batch, seq_len, 2) x self.embed(points) # 将二维坐标映射到高维空间 return self.encoder(x)这里用了一个简单的线性层把二维坐标映射到hidden_dim维。之所以不用正弦位置编码是因为 TSP 的输入没有天然顺序加入位置反而不利于模型理解拓扑结构。Transformer 的自注意力层能让每座城市在编码阶段就收到其他所有城市的信息后续解码时的指针选择才能综合考虑全局分布。num_layers3是经验值层数太少城市间的交互建模不充分层数太多训练成本上升快收益却变小。2.3 自回归解码与确定性转化mask 机制解码过程是典型的自回归给定已访问的城市序列预测下一个城市的分布把预测结果作为下一步的输入重复直到所有城市都被选完。def decode_tour(model, points, deterministicTrue): batch_size, seq_len, _ points.shape encoder_outputs model.encoder(points) # 初始解码状态用编码器输出的平均池化 decoder_hidden encoder_outputs.mean(dim1) visited torch.zeros(batch_size, seq_len, dtypetorch.bool) tour [] for _ in range(seq_len): logits model.pointer(encoder_outputs, decoder_hidden, mask~visited) if deterministic: next_node logits.argmax(dim-1) else: next_node torch.multinomial(logits, 1).squeeze(-1) tour.append(next_node) visited visited.scatter(1, next_node.unsqueeze(-1), True) # 用选中的城市嵌入更新解码器状态 decoder_hidden model.decoder_cell(decoder_hidden, encoder_outputs[torch.arange(batch_size), next_node]) return torch.stack(tour, dim1)decoder_cell通常是一个 GRU 或线性层负责把“上一步选中城市的嵌入”融合进当前隐状态。deterministicTrue时走贪心解码测试阶段用False时采样训练阶段用。mask~visited表示只允许在未访问的城市上产生概率。有些实现会用加性注意力而不是点积注意力来算 logits两者在 TSP 上差距不大但加性注意力在维度较大时数值更稳定推荐在hidden_dim 256时切换到加性实现。3. 用 REINFORCE 训练没有标签也能让模型学会规划路径3.1 为什么监督学习在这里不适用如果走监督学习路线需要为每个训练实例准备最优解标注。用精确算法比如 Concorde生成标注不仅慢而且只适用于小规模。更麻烦的是TSP 的最优解不唯一对称性和旋转不变性会让模型学到大量冲突信号。反观强化学习只需要定义奖励——路径总长度的倒数——就能让模型自己在探索中发现策略。对于 TSP 这种组合优化问题奖励稀疏性并不严重因为每个完整 tour 的路径长度都能计算且长度本身就是比 baseline 好坏的信号。3.2 策略梯度目标与 baseline 的两种选择REINFORCE 的目标函数是最大化期望奖励用对数似然技巧转化为梯度估计∇θ J(θ) ≈ E[∇θ log πθ(tour | points) * (L(tour) - baseline)]这里的L(tour)是路径总长度baseline的作用是减小方差。如果去掉 baseline梯度的方差会大到模型几乎不收敛。常见做法有两种方法思路优点缺点Critic 网络用一个回归网络预测当前实例的期望路径长度训练快方差小需要额外训练一个 Critic且 Critic 的误差会传导到策略Greedy Rollout baseline用一个固定周期性更新的贪心策略作为 baseline无需单独训练网络偏差小每步都要额外解码一次训练耗时约增加一倍选择哪种 baseline 直接决定了代码的复杂度和训练时间。对第一次跑通管道的团队我建议用 Critic 网络代码量少调参也相对简单。class Critic(nn.Module): def __init__(self, hidden_dim): super().__init__() self.inp nn.Linear(hidden_dim, hidden_dim) self.out nn.Linear(hidden_dim, 1) def forward(self, encoder_outputs): pooled encoder_outputs.mean(dim1) return self.out(F.relu(self.inp(pooled)))Critic 把编码器输出的所有城市嵌入做平均池化然后回归出一个标量代表模型对该实例路径长度的估计。3.3 损失函数、梯度裁剪与训练超参数设定def reinforce_loss(tour_log_probs, tour_lengths, baseline_values): # tour_log_probs: (batch, seq_len) 每一步选节点的对数概率 # tour_lengths: (batch,) # baseline_values: (batch,) log_probs tour_log_probs.sum(dim1) advantage (tour_lengths - baseline_values).detach() # 关键advantage 不参与梯度传播 return (log_probs * advantage).mean()detach()是这段代码最关键的一步。advantage需要被视为常数否则梯度会同时经过策略和 baseline导致训练不稳。路径长度本身作为标量奖励数值大建议除以seq_len做归一化或者用tour_lengths / (seq_len * scale_factor)的形式避免梯度尺度过大。训练时通常搭配 Adam 优化器学习率设在1e-4到3e-4之间并配合梯度裁剪clip_grad_norm_设为max_norm1.0防止 RNN 或 Transformer 的梯度爆炸。4. 从数据生成到训练循环一套完整的 Python 实现4.1 动态生成训练数据静态欧式 TSP 的批处理采样TSP 训练集不适合固定生成一次——每个 epoch 都重新采样可以让模型在每个 batch 看到全新的城市分布天然起到正则化作用也能避免模型死记硬背固定坐标。这里统一采用静态欧式 TSP城市坐标在[0, 1]方形区域内均匀采样。def generate_tsp_batch(batch_size, num_cities, seedNone): if seed is not None: torch.manual_seed(seed) # 生成二维坐标范围 [0, 1] points torch.rand(batch_size, num_cities, 2) return pointsnum_cities决定了问题难度。模型在训练时最好从较小的规模如 20 个城市起步验证时再逐步提升到 50 或 100。这种策略叫课程学习Curriculum Learning可以显著加速收敛。4.2 训练循环优势计算、梯度更新与模型保存def train(model, critic, optimizer, critic_optimizer, num_epochs, batch_size512, num_cities20): for epoch in range(num_epochs): points generate_tsp_batch(batch_size, num_cities) tours, log_probs model(points, decode_typestochastic) tour_lengths calculate_tour_length(points, tours) baseline_values critic(points).squeeze(-1) loss reinforce_loss(log_probs, tour_lengths, baseline_values) critic_loss F.mse_loss(baseline_values, tour_lengths.detach()) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() critic_optimizer.zero_grad() critic_loss.backward() critic_optimizer.step() if epoch % 10 0: avg_len tour_lengths.mean().item() print(fEpoch {epoch}, avg tour length: {avg_len:.4f})calculate_tour_length的计算逻辑是tours存储的是城市索引用torch.gather从points中取出对应的坐标然后计算相邻点之间的欧氏距离再加上最后一个回程到起点的距离。def calculate_tour_length(points, tours): batch_idx torch.arange(points.size(0)).unsqueeze(1) tour_coords points[batch_idx, tours] # (batch, seq_len, 2) # 让序列首尾闭合 next_coords torch.cat([tour_coords[:, 1:, :], tour_coords[:, :1, :]], dim1) return torch.norm(tour_coords - next_coords, dim-1).sum(dim-1)路径长度计算是训练循环里每个 batch 都要执行的算子性能很关键。上面的向量化实现避免了对每个 batch 逐一循环GPU 利用率更高。对变长 TSP比如训练实例的城市数量不固定则需要在数据生成时按长度分组或加 padding mask这会显著增加代码复杂度建议第一版先固定城市数量。4.3 训练效果与经典基线对比训练 20 个城市的静态欧式 TSP拿贪心最近邻算法做对照通常几分钟到几十分钟内就能看到明显差距。方法平均路径长度20 城市计算耗时备注最近邻贪心4.2~4.5毫秒级不需要训练随机策略7~8毫秒级基线指针网络 REINFORCE3.8~4.0训练后单次推理亚毫秒级训练约 15 分钟这里的绝对数值取决于随机种子与训练轮数不同机器上的结果会略有浮动但指针网络的优势在于它不需要每个实例单独计算——训练完成后推理成本极低这是它相比传统启发式方法的根本性差异。5. 推理阶段提升解质量采样、集束搜索与验证技巧训练完成后模型本身就是一个可部署的求解器。但贪心解码往往只给出平均水平的解想要在固定预算内拿到更好的 tour可以在推理时多花一点计算量。对规模在 50 城市以下的静态欧式 TSP最简单有效的手段是采用集束搜索Beam Search。用 PyTorch 实现 Top-K 采样时需要注意每一步要有 K 个不同的部分序列同时维护各自的 mask 状态和概率累加值还要考虑路径闭合时的边界处理。def beam_search_decode(model, points, beam_width5, deterministicTrue): batch_size, seq_len, _ points.shape encoder_outputs model.encoder(points) decoder_hidden encoder_outputs.mean(dim1) visited torch.zeros(batch_size, seq_len, dtypetorch.bool) # 以 (log_prob, tour, hidden, mask) 作为堆叠变量 sequences torch.zeros(batch_size, 1, dtypetorch.long) log_probs torch.zeros(batch_size, 1) visited visited.unsqueeze(1).expand(batch_size, beam_width, -1).contiguous() for step in range(seq_len): # 展开 beam解码概率 flat_logits model.pointer(encoder_outputs, decoder_hidden, mask~visited) # 取 Top-K 候选 top_k torch.topk(flat_logits, kbeam_width, dim-1) # 更新序列、visited、hidden # ... return sequences, log_probs这段代码只给出骨架完整实现比贪心解码多出至少三个细节一是 hidden 状态的复制与重排二是每步从选中的top_k节点更新 beam 内的所有序列后需要重新排序合并三是最终从beam_width条路径中选择总长度最短的那条。实际经验是 beam width 从 1 增加到 10路径长度大约能降低 1%~3%计算量却线性增加所以推理预算有限时优先选 beam search 而不是更大的模型。最后一个推荐技巧是验证模型是否过拟合训练集分布在测试时把城市坐标做随机旋转和缩放保持问题的欧氏特性不变如果模型的解质量显著下降说明模型学到了坐标空间的某种偏置而不是拓扑结构的不变特征。这类验证对强化学习模型尤为重要因为 REINFORCE 的训练信号弱模型容易走捷径。本文还有配套的精品资源点击获取
