我想每次做出多个飞机最优选择，编码时，我想使用四个选择的信息，我的代码是否可行，解码编码是否有问题，解码时用到的信息有哪些，得到权重方式是否有不对地方，如何一步使用多个选择信息获取所有选择权重，代码是否还有别的问题class Transformer_encoder_net(nn.Module):
    """
    基于自注意力 Transformer 的编码器网络
    输入：
      h 的大小为 (bsz, nb_nodes+1, dim_emb)    输入城市的批量
    输出：
      h 的大小为 (bsz, nb_nodes+1, dim_emb)    编码后的城市批量
      score 的大小为 (bsz, nb_nodes+1, nb_nodes+1)    注意力分数的批量
    """
    def __init__(self, nb_layers, dim_emb, nb_heads, dim_ff, batchnorm):
        super(Transformer_encoder_net, self).__init__()
        
        # 确保嵌入维度能够被头的数量整除
        assert dim_emb == nb_heads * (dim_emb // nb_heads) 
        
        # 初始化多头自注意力层
        self.MHA_layers = nn.ModuleList([nn.MultiheadAttention(dim_emb, nb_heads) for _ in range(nb_layers)])
        
        # 初始化前馈网络层
        self.linear1_layers = nn.ModuleList([nn.Linear(dim_emb, dim_ff) for _ in range(nb_layers)])
        self.linear2_layers = nn.ModuleList([nn.Linear(dim_ff, dim_emb) for _ in range(nb_layers)])
        
        # 根据输入标志选择 BatchNorm 或 LayerNorm 进行归一化
        if batchnorm:
            self.norm1_layers = nn.ModuleList([nn.BatchNorm1d(dim_emb) for _ in range(nb_layers)])
            self.norm2_layers = nn.ModuleList([nn.BatchNorm1d(dim_emb) for _ in range(nb_layers)])
        else:
            self.norm1_layers = nn.ModuleList([nn.LayerNorm(dim_emb) for _ in range(nb_layers)])
            self.norm2_layers = nn.ModuleList([nn.LayerNorm(dim_emb) for _ in range(nb_layers)])
        
        # 保存层的数量、头的数量和 BatchNorm 标志
        self.nb_layers = nb_layers
        self.nb_heads = nb_heads
        self.batchnorm = batchnorm
        
    def forward(self, h):      
        # PyTorch nn.MultiheadAttention 需要输入大小为 (seq_len, bsz, dim_emb)
        h = h.transpose(0, 1)  # 转置后 size(h)=(nb_nodes, bsz, dim_emb)
        
        # 逐层处理 L 层
        for i in range(self.nb_layers):
            h_rc = h  # 残差连接，size(h_rc)=(nb_nodes, bsz, dim_emb)
            h, score = self.MHA_layers[i](h, h, h)  # size(h)=(nb_nodes, bsz, dim_emb)，size(score)=(bsz, nb_nodes, nb_nodes)
            score = score.mean(dim=1)  # 平均所有头的注意力权重, score: (bsz, nb_nodes, nb_nodes)
            
            # 添加残差连接
            h = h_rc + h  # size(h)=(nb_nodes, bsz, dim_emb)
            
            # 进行归一化
            if self.batchnorm:
                # Pytorch nn.BatchNorm1d 需要输入大小为 (bsz, dim, seq_len)
                h = h.permute(1, 2, 0).contiguous()  # size(h)=(bsz, dim_emb, nb_nodes)
                h = self.norm1_layers[i](h)          # size(h)=(bsz, dim_emb, nb_nodes)
                h = h.permute(2, 0, 1).contiguous()  # size(h)=(nb_nodes, bsz, dim_emb)
            else:
                h = self.norm1_layers[i](h)          # size(h)=(nb_nodes, bsz, dim_emb)
                
            # 前馈网络
            h_rc = h  # 残差连接
            h = self.linear2_layers[i](torch.relu(self.linear1_layers[i](h)))  # 前馈计算
            h = h_rc + h  # size(h)=(nb_nodes, bsz, dim_emb)
            
            # 进行归一化
            if self.batchnorm:
                h = h.permute(1, 2, 0).contiguous()  # size(h)=(bsz, dim_emb, nb_nodes)
                h = self.norm2_layers[i](h)          # size(h)=(bsz, dim_emb, nb_nodes)
                h = h.permute(2, 0, 1).contiguous()  # size(h)=(nb_nodes, bsz, dim_emb)
            else:
                h = self.norm2_layers[i](h)          # size(h)=(nb_nodes, bsz, dim_emb)
        
        # 转置 h 以便于输出
        h = h.transpose(0, 1)  # size(h)=(bsz, nb_nodes, dim_emb)
        return h, score  # 返回编码后的结果和注意力分数
    

def myMHA(Q, K, V, nb_heads, mask=None, clip_value=None):
    bsz, nb_nodes, emd_dim = K.size() #  dim_emb must be divisable by nb_heads
    plane_num = Q.size(1)
    if nb_heads>1:
        # PyTorch view requires contiguous dimensions for correct reshaping
        Q = Q.transpose(1,2).contiguous() # size(Q)=(bsz, dim_emb, 1)
        Q = Q.view(bsz*nb_heads, emd_dim//nb_heads, plane_num) # size(Q)=(bsz*nb_heads, dim_emb//nb_heads, 1)
        Q = Q.transpose(1,2).contiguous() # size(Q)=(bsz*nb_heads, 1, dim_emb//nb_heads)
        K = K.transpose(1,2).contiguous() # size(K)=(bsz, dim_emb, nb_nodes+1)
        K = K.view(bsz*nb_heads, emd_dim//nb_heads, nb_nodes) # size(K)=(bsz*nb_heads, dim_emb//nb_heads, nb_nodes+1)
        K = K.transpose(1,2).contiguous() # size(K)=(bsz*nb_heads, nb_nodes+1, dim_emb//nb_heads)
        V = V.transpose(1,2).contiguous() # size(V)=(bsz, dim_emb, nb_nodes+1)
        V = V.view(bsz*nb_heads, emd_dim//nb_heads, nb_nodes) # size(V)=(bsz*nb_heads, dim_emb//nb_heads, nb_nodes+1)
        V = V.transpose(1,2).contiguous() # size(V)=(bsz*nb_heads, nb_nodes+1, dim_emb//nb_heads)
    attn_weights = torch.bmm(Q, K.transpose(1,2))/ Q.size(-1)**0.5 # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)
    if clip_value is not None:
        attn_weights = clip_value * torch.tanh(attn_weights)
    if mask is not None:
        if nb_heads>1:
            mask = torch.repeat_interleave(mask, repeats=nb_heads, dim=0) # size(mask)=(bsz*nb_heads, nb_nodes+1)
        #attn_weights = attn_weights.masked_fill(mask.unsqueeze(1), float('-inf')) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)
        attn_weights = attn_weights.masked_fill(mask.unsqueeze(1), float('-1e9')) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)
    attn_weights = torch.softmax(attn_weights, dim=-1) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)
    attn_output = torch.bmm(attn_weights, V) # size(attn_output)=(bsz*nb_heads, 1, dim_emb//nb_heads)
    if nb_heads>1:
        attn_output = attn_output.transpose(1,2).contiguous() # size(attn_output)=(bsz*nb_heads, dim_emb//nb_heads, 1)
        attn_output = attn_output.view(bsz, emd_dim, plane_num) # size(attn_output)=(bsz, dim_emb, 1)
        attn_output = attn_output.transpose(1,2).contiguous() # size(attn_output)=(bsz, 1, dim_emb)
        attn_weights = attn_weights.view(bsz, nb_heads, plane_num, nb_nodes) # size(attn_weights)=(bsz, nb_heads, 1, nb_nodes+1)
        attn_weights = attn_weights.mean(dim=1) # mean over the heads, size(attn_weights)=(bsz, 1, nb_nodes+1)
    return attn_output, attn_weights
    
class AutoRegressiveDecoderLayer(nn.Module):
    """
    Single decoder layer based on self-attention and query-attention
    Inputs :  
      h_t of size      (bsz, 1, dim_emb)          batch of input queries
      K_att of size    (bsz, nb_nodes+1, dim_emb) batch of query-attention keys
      V_att of size    (bsz, nb_nodes+1, dim_emb) batch of query-attention values
      mask of size     (bsz, nb_nodes+1)          batch of masks of visited cities
    Output :  
      h_t of size (bsz, nb_nodes+1)               batch of transformed queries
    """
    def __init__(self, dim_emb, nb_heads):
        super(AutoRegressiveDecoderLayer, self).__init__()
        self.dim_emb = dim_emb
        self.nb_heads = nb_heads
        self.Wq_selfatt = nn.Linear(dim_emb, dim_emb)
        self.Wk_selfatt = nn.Linear(dim_emb, dim_emb)
        self.Wv_selfatt = nn.Linear(dim_emb, dim_emb)
        self.W0_selfatt = nn.Linear(dim_emb, dim_emb)
        self.W0_att = nn.Linear(dim_emb, dim_emb)
        self.Wq_att = nn.Linear(dim_emb, dim_emb)
        self.W1_MLP = nn.Linear(dim_emb, dim_emb)
        self.W2_MLP = nn.Linear(dim_emb, dim_emb)
        self.BN_selfatt = nn.LayerNorm(dim_emb)
        self.BN_att = nn.LayerNorm(dim_emb)
        self.BN_MLP = nn.LayerNorm(dim_emb)
        self.K_sa = None
        self.V_sa = None

    def reset_selfatt_keys_values(self):
        self.K_sa = None
        self.V_sa = None
        
    # For beam search
    def reorder_selfatt_keys_values(self, t, idx_top_beams):
        bsz, B = idx_top_beams.size()
        zero_to_B = torch.arange(B, device=idx_top_beams.device) # [0,1,...,B-1]
        B2 = self.K_sa.size(0)// bsz
        self.K_sa = self.K_sa.contiguous().view(bsz, B2, t+1, self.dim_emb) # size(self.K_sa)=(bsz, B2, t+1, dim_emb)
        K_sa_tmp = self.K_sa.clone()
        self.K_sa = torch.zeros(bsz, B, t+1, self.dim_emb, device=idx_top_beams.device)
        for b in range(bsz):
            self.K_sa[b, zero_to_B, :, :] = K_sa_tmp[b, idx_top_beams[b], :, :]
        self.K_sa = self.K_sa.contiguous().view(bsz*B, t+1, self.dim_emb) # size(self.K_sa)=(bsz*B, t+1, dim_emb)
        self.V_sa = self.V_sa.contiguous().view(bsz, B2, t+1, self.dim_emb) # size(self.K_sa)=(bsz, B, t+1, dim_emb)
        V_sa_tmp = self.V_sa.clone()
        self.V_sa = torch.zeros(bsz, B, t+1, self.dim_emb, device=idx_top_beams.device)
        for b in range(bsz):
            self.V_sa[b, zero_to_B, :, :] = V_sa_tmp[b, idx_top_beams[b], :, :]
        self.V_sa = self.V_sa.contiguous().view(bsz*B, t+1, self.dim_emb) # size(self.K_sa)=(bsz*B, t+1, dim_emb)

    # For beam search
    def repeat_selfatt_keys_values(self, B):
        self.K_sa = torch.repeat_interleave(self.K_sa, B, dim=0) # size(self.K_sa)=(bsz.B, t+1, dim_emb)
        self.V_sa = torch.repeat_interleave(self.V_sa, B, dim=0) # size(self.K_sa)=(bsz.B, t+1, dim_emb)
        
    def forward(self, h_t, K_att, V_att, mask):
        bsz = h_t.size(0)
        plane_num = h_t.size(1)
        h_t = h_t.contiguous().view(bsz,plane_num,self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)
        # embed the query for self-attention
        q_sa = self.Wq_selfatt(h_t) # size(q_sa)=(bsz, 1, dim_emb)
        k_sa = self.Wk_selfatt(h_t) # size(k_sa)=(bsz, 1, dim_emb)
        v_sa = self.Wv_selfatt(h_t) # size(v_sa)=(bsz, 1, dim_emb)
        # concatenate the new self-attention key and value to the previous keys and values
        if self.K_sa is None:
            self.K_sa = k_sa # size(self.K_sa)=(bsz, 1, dim_emb)
            self.V_sa = v_sa # size(self.V_sa)=(bsz, 1, dim_emb)
        else:
            self.K_sa = torch.cat([self.K_sa, k_sa], dim=1)
            self.V_sa = torch.cat([self.V_sa, v_sa], dim=1)
        # compute self-attention between nodes in the partial tour
        h_t = h_t + self.W0_selfatt(myMHA(q_sa, self.K_sa, self.V_sa, self.nb_heads)[0] ) # size(h_t)=(bsz, 1, dim_emb)
        h_t = self.BN_selfatt(h_t.squeeze()) # size(h_t)=(bsz, dim_emb)
        h_t = h_t.contiguous().view(bsz, plane_num, self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)
        # compute attention between self-attention nodes and encoding nodes in the partial tour (translation process)
        q_a = self.Wq_att(h_t) # size(q_a)=(bsz, 1, dim_emb)
        h_t = h_t + self.W0_att( myMHA(q_a, K_att, V_att, self.nb_heads, mask)[0] ) # size(h_t)=(bsz, 1, dim_emb)
        h_t = self.BN_att(h_t.squeeze()) # size(h_t)=(bsz, dim_emb)
        h_t = h_t.contiguous().view(bsz, plane_num, self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)
        # MLP
        h_t = h_t + self.W2_MLP(torch.relu(self.W1_MLP(h_t)))
        h_t = self.BN_MLP(h_t.squeeze(1)) # size(h_t)=(bsz, dim_emb)
        return h_t


class Transformer_decoder_net(nn.Module): 
    """
    基于自注意力和查询注意力变换器的解码器网络。
    输入：
      h_t 的大小为 (bsz, 1, dim_emb)                            输入查询的批量
      K_att 的大小为 (bsz, nb_nodes+1, dim_emb*nb_layers_decoder) 所有解码层的查询注意力键的批量
      V_att 的大小为 (bsz, nb_nodes+1, dim_emb*nb_layers_decoder) 所有解码层的查询注意力值的批量
      mask 的大小为 (bsz, nb_nodes+1)                            已访问城市的掩码批量
    输出：
      prob_next_node 的大小为 (bsz, nb_nodes+1)                  下一个节点的概率的批量
      
    """
    def __init__(self, dim_emb, nb_heads, nb_layers_decoder):
        super(Transformer_decoder_net, self).__init__()
        self.dim_emb = dim_emb  # 嵌入维度
        self.nb_heads = nb_heads  # 注意力头的数量
        self.nb_layers_decoder = nb_layers_decoder  # 解码层的数量
        
        # 初始化解码器层
        self.decoder_layers = nn.ModuleList([AutoRegressiveDecoderLayer(dim_emb, nb_heads) for _ in range(nb_layers_decoder - 1)])
        self.Wq_final = nn.Linear(dim_emb, dim_emb)  # 最后输出的线性变换

    # 当解码开始时重置自注意力的键和值为 None
    def reset_selfatt_keys_values(self): 
        for l in range(self.nb_layers_decoder - 1):
            self.decoder_layers[l].reset_selfatt_keys_values()  # 重置每一层的自注意力键和值
            
    def forward(self, h_t, K_att, V_att, mask):
        for l in range(self.nb_layers_decoder):
            # 获取当前解码层的 K_att 和 V_att
            K_att_l = K_att[:, :, l * self.dim_emb:(l + 1) * self.dim_emb].contiguous()  # size(K_att_l)=(bsz, nb_nodes+1, dim_emb)
            V_att_l = V_att[:, :, l * self.dim_emb:(l + 1) * self.dim_emb].contiguous()  # size(V_att_l)=(bsz, nb_nodes+1, dim_emb)
            if l < self.nb_layers_decoder - 1:  # 多头解码层（中间层）
                h_t = self.decoder_layers[l](h_t, K_att_l, V_att_l, mask)  # 逐层传递
            else:  # 单头解码层（最后一层）
                q_final = self.Wq_final(h_t)  # 最后层的线性变换
                bsz = h_t.size(0)  # 批量大小
                plane_num = h_t.size(1)
                q_final = q_final.view(bsz, plane_num, self.dim_emb)  # 调整形状，size(q_final)=(bsz, 1, dim_emb)
                # 计算下一步访问节点的概率
                
                attn_weights = myMHA(q_final, K_att_l, V_att_l, 1, mask, 10)[1]  # 计算注意力权重
                
        prob_next_node = attn_weights.squeeze(1).mean(dim=1)  # 去掉第二维度，得到概率
        return prob_next_node  # 返回下一个节点的概率


def generate_positional_encoding(d_model, max_len):
    """
    创建标准变换器位置编码（PE）。
    输入：
      d_model 是对应于隐藏维度的标量
      max_len 是序列的最大长度
    输出：
      pe 的大小为 (max_len, d_model)，其中 d_model=dim_emb，max_len=1000
    """
    pe = torch.zeros(max_len, d_model)  # 初始化位置编码矩阵
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # 生成位置索引
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))  # 计算缩放因子
    
    # 奇数位置使用正弦函数，偶数位置使用余弦函数
    pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度
    pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度
    return pe  # 返回生成的位置编码

class TSP_net(nn.Module):
    """
    TSP网络包含两个步骤
      步骤1. 编码器步骤:获取代表plane_num * num_choices个选择,每个选择维度是39维,通过self-transformer对其进行编码。
      步骤2. 解码器步骤：递归/自回归地构建空域搜索集合的选择，每次从代表不同飞机的部分选择一个搜索集合,使用self-transformer和query-transformer完成。
    输入：
      x: 形状为 (bsz, plane_num * num_choices, input_dim) 
      deterministic: 布尔值,如果为True,选择概率最高的选择,如果为False.则通过Categorical采样选择选项
    输出：
      tours: 形状为 (bsz, palne_num,step) 的批量选择
      sumLogProbOfActions: 形状为 (bsz,) 的张量，每个元素为sum_t log prob( pi_t | pi_(t-1),...,pi_0 )。
    """

    def __init__(self, dim_input_nodes, dim_emb, dim_ff, nb_layers_encoder, nb_layers_decoder, nb_heads, max_len_PE,
                 batchnorm=True):
        super(TSP_net, self).__init__()

        self.dim_emb = dim_emb

        # 输入嵌入层：将输入的节点特征映射到嵌入维度
        self.input_emb = nn.Linear(dim_input_nodes, dim_emb)

        # 编码器层：用于对节点进行编码，采用自定义的Transformer编码器
        self.encoder = Transformer_encoder_net(nb_layers_encoder, dim_emb, nb_heads, dim_ff, batchnorm)

        # 解码起始占位符：随机初始化，用于解码起始的向量
        self.start_placeholder = nn.Parameter(torch.randn(dim_emb))

        # 解码器层：用于构建TSP路径，采用自定义的Transformer解码器
        self.decoder = Transformer_decoder_net(dim_emb, nb_heads, nb_layers_decoder)
        self.WK_att_decoder = nn.Linear(dim_emb, nb_layers_decoder * dim_emb)  # 解码器的键
        self.WV_att_decoder = nn.Linear(dim_emb, nb_layers_decoder * dim_emb)  # 解码器的值
        self.PE = generate_positional_encoding(dim_emb, max_len_PE)  # 位置编码，用于引入序列信息

    def forward(self, x, plane_num,num_choices,max_steps,step_now,K_att_decoder, V_att_decoder, mask_visited_nodes,h_encoder,deterministic=False):
        if step_now == 0:
            # 获取批量大小和节点数量
            bsz = x.shape[0]
            nb_nodes = x.shape[1]
            zero_to_bsz = torch.arange(bsz, device=x.device)  # 生成一个[0,1,...,bsz-1]的序列
            # 输入嵌入层：将节点特征映射到嵌入空间
            h = self.input_emb(x)  # h的形状为 (bsz, nb_nodes, dim_emb)
            # 拼接节点和解码起始占位符
            h = torch.cat([h, self.start_placeholder.repeat(bsz, 1, 1)], dim=1)  # 拼接后的形状为 (bsz, nb_nodes+1, dim_emb)
            # 编码器层：对输入进行编码
            h_encoder, _ = self.encoder(h)  # 编码后的形状为 (bsz, nb_nodes+1, dim_emb)
            # 存储每步选择的负对数概率的列表
            sumLogProbOfActions = []
            # 解码器的键和值
            K_att_decoder = self.WK_att_decoder(h_encoder)  # 键的形状为 (bsz, nb_nodes+1, dim_emb*nb_layers_decoder)
            V_att_decoder = self.WV_att_decoder(h_encoder)  # 值的形状为 (bsz, nb_nodes+1, dim_emb*nb_layers_decoder)
            # 解码起始向量：编码后的起始节点加上位置编码
            self.PE = self.PE.to(x.device)
            # 选择最后一个节点作为起始节点的索引
            idx_start_placeholder = torch.tensor([nb_nodes]).long().to(x.device).repeat(bsz, plane_num)
            # 获取最后一个节点的表示，重复 plane_num 次
            # 使用 gather 进行索引
            h_start = h_encoder[zero_to_bsz,  idx_start_placeholder[:,0], :].unsqueeze(1).expand(bsz,plane_num,-1) + self.PE[0].repeat(bsz, plane_num, 1)
            # 更新掩码，标记所有 plane_num 节点为已访问
            mask_visited_nodes[zero_to_bsz, idx_start_placeholder[:,0]] = True  # 标记所有选择的节点为已访问
            # idx_start_placeholder = torch.Tensor([nb_nodes]).long().repeat(bsz).to(x.device)
            # h_start = h_encoder[zero_to_bsz, idx_start_placeholder, :] + self.PE[0].repeat(bsz, 1)  # 起始向量的形状为 (bsz, dim_emb)
            # mask_visited_nodes[zero_to_bsz, idx_start_placeholder] = True  # 标记起始节点为已访问
            # 清除解码器中的存储的键和值
            self.decoder.reset_selfatt_keys_values()
            # 递归构建路径
            h_t = h_start
            tours = torch.zeros((bsz,plane_num),dtype=torch.int,device=device)
        else:
            h_t = x
            bsz = x.shape[0]
            nb_nodes = plane_num * num_choices
            zero_to_bsz = torch.arange(bsz, device=x.device)  # 生成一个[0,1,...,bsz-1]的序列
            sumLogProbOfActions = []
            tours = torch.zeros((bsz,plane_num),dtype=torch.int,device=device)
        prob_next_node = self.decoder(h_t.clone(), K_att_decoder, V_att_decoder, mask_visited_nodes)  # 概率形状为 (bsz, nb_nodes+1) 
        # 初始化选择的概率分布，将其形状转换为 (batch_size, plane_num, plane_num, nb_nodes)
        prob_next_node1 = prob_next_node[:,0:nb_nodes]
        prob_next_node1 = prob_next_node1.contiguous().view(bsz, plane_num, num_choices)
        # 创建一个保存最终选择概率的张量
        final_prob_next_node = torch.zeros(bsz, plane_num, num_choices, device=prob_next_node.device)
        # 对每个 plane 的选择区间应用 softmax，确保每个 plane 只选择属于它的区间
        for i in range(plane_num):
            # 取出每个 plane 的选择区间，即第 i 个区间
            current_prob = prob_next_node1[:, i, :]  # shape = (batch_size, nb_nodes)
            # 对该 plane 的选择区间进行 softmax 归一化
            final_prob_next_node[:, i, :] = torch.softmax(current_prob, dim=-1)
        # 创建一个保存最终选择概率的张量
        # 根据确定性或采样方式选择下一个节点
        # idx_list = torch.zeros(bsz,plane_num,dtype=torch.long).contiguous()
        if deterministic:
            idx_list =  torch.argmax(final_prob_next_node.view(bsz * plane_num, num_choices), dim=1) # shape = (bsz * plane_num, 1)
            idx_list = idx_list.view(bsz, plane_num)  # reshape to (batch_size, plane_num)
        else:
            idx_list =  Categorical(probs=final_prob_next_node.view(bsz * plane_num, num_choices)).sample()  # shape = (bsz * plane_num, 1)
            idx_list = idx_list.view(bsz, plane_num)  # reshape to (batch_size, plane_num)
        
        # 计算当前步骤选择的节点的对数概率，并存储
        # 更新已访问节点的掩码
        for k in range(plane_num):
            #ProbOfChoices = final_prob_next_node[:, k, idx_list[:,k]]
            # 这样索引是有问题的
            # 相当于result = []
            # for i in range(12):
            # result.append(final_prob_next_node[i, k, idx_list[:, k]])
            # result = torch.stack(result)
            ProbOfChoices = final_prob_next_node[torch.arange(bsz), k, idx_list[:, k]]
            sumLogProbOfActions.append(ProbOfChoices)
            # 更新当前访问的节点嵌入
            h_t[:,k,:] = h_encoder[zero_to_bsz, idx_list[:,k], :]  # 更新为当前选择节点的嵌入
        # 更新已访问节点的掩码
        mask_visited_nodes = mask_visited_nodes.clone()
        mask_visited_nodes[zero_to_bsz.unsqueeze(1), idx_list] = True  # 使用高级索引一次性更新
            # 更新路径：将选择的节点索引添加到路径中
        tours= idx_list
        h_t = h_t + self.PE[step_now + 1].expand(bsz, plane_num,self.dim_emb)  # 加上下一位置的编码
        # 将路径的节点索引列表转换为张量：形状为 (bsz, nb_nodes)
        # 计算路径选择的总对数概率：sum_t log prob( pi_t | pi_(t-1),...,pi_0 )
        sumLogProbOfActions = torch.stack(sumLogProbOfActions).sum(dim=0)
        return tours, sumLogProbOfActions,h_t,K_att_decoder, V_att_decoder, mask_visited_nodes,h_encoder # 返回路径和总对数概率
    
dim_input_nodes = 39
# target_space 8 * 3 detect_space 5 * 3
dim_emb = 128
dim_ff = 512
nb_layers_encoder = 6
nb_layers_decoder = 2
nb_heads = 16
max_len_PE = 1000
batch = 12
plane_num = 4
num_choices = 121 # 11 * 11
batchnorm = True # if batchnorm=True  than batch norm is used
max_steps = 15

# 初始化训练模型
model_train = TSP_net(dim_input_nodes, dim_emb, dim_ff, 
              nb_layers_encoder, nb_layers_decoder, nb_heads, max_len_PE,
              batchnorm=batchnorm).to(device)
# 初始化基准模型
model_baseline = TSP_net(dim_input_nodes, dim_emb, dim_ff, 
              nb_layers_encoder, nb_layers_decoder, nb_heads, max_len_PE,
              batchnorm=batchnorm).to(device)
# 初始化优化器，Adam优化器用于更新训练模型的参数
optimizer = torch.optim.Adam(model_train.parameters(), lr=0.0001)
model_baseline.eval()  # 设置基准模型为评估模式
epochs = 200
epislion = 0.9
max_mean_train =30000
for epoch in range(epochs):
    mean_epoch_train = 0
    mean_epoch_baseline = 0
    model_train.train()  # 设置训练模型为训练模式
    max_step_per_epoch = 30
    epislion = 0.9 * epislion
    d_train = False
    # if random.uniform(0,1) > epislion:
    #     d_train = True
    # else:
    #     d_train = False
    model_train.train()  # 设置训练模型为训练模式
    max_step_per_epoch = 30
    # current_dir = os.getcwd()
    # log_dir = os.path.join(current_dir, "logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
    # writer = SummaryWriter(log_dir=log_dir)
    # if epoch == 0:
    #     model_1 = torch.load('8.4.4.20_31179.0.pkl')
    #     # 加载已经训练好的模型
    #     model_train.load_state_dict(model_1.state_dict())
    #     model_baseline.load_state_dict(model_1.state_dict())
    for step_per_epoch in tqdm(range(max_step_per_epoch), desc='step_per_epochs'):
        t1,t2,start_x = get_train_text(batch,num_choices,plane_num)
        com = combine_features(t1,t2)
        # 初始化访问节点的掩码：表示节点是否已被访问，起始占位符节点标记为已访问
        mask_visited_nodes_train = torch.zeros((batch, plane_num * num_choices + 1), device=device).bool()  # 全部初始化为False
        mask_visited_nodes_baseline = torch.zeros((batch, plane_num * num_choices + 1), device=device).bool()  # 全部初始化为False
        K_att_decoder = torch.zeros((batch,plane_num * num_choices + 1,dim_emb*nb_layers_decoder),device=device)
        V_att_decoder = torch.zeros((batch,plane_num * num_choices + 1,dim_emb*nb_layers_decoder),device=device)
        h_encoder = torch.zeros((batch, plane_num * num_choices + 1, dim_emb),device=device)
        reward_baseline_old = torch.zeros(batch,device=device)
        reward_train_old = torch.zeros(batch,device=device)
        names = locals()
        voxel_flags_train = torch.zeros(batch,10 * 5 * 60 * 5 *60 * 5, dtype=torch.int, device=device)
        voxel_flags_baseline = torch.zeros(batch,10 * 5 * 60 * 5 *60 * 5, dtype=torch.int, device=device)
        all_voxel_flags = torch.zeros((batch,10 * 5 * 60 * 5 *60 * 5, 3), device = device ,dtype=torch.float16)
        total_loss = 0
        for num_step in range(max_steps): 
            if num_step == 0: 
                tour_train,names['sumLogProbOfAction'+str(num_step)],h_t_train,K_att_decoder_train, V_att_decoder_train, mask_visited_nodes_train,h_encoder_train = model_train(com,
                                                                                                                                                                   h_encoder=h_encoder,step_now=num_step,
                                                                                                                                                                   plane_num=plane_num,num_choices=num_choices,max_steps=max_steps,
                                                                                                                                                                   mask_visited_nodes=mask_visited_nodes_train,K_att_decoder=K_att_decoder,
                                                                                                                                                                   V_att_decoder=V_att_decoder,deterministic=d_train)
            else:
                tour_train,names['sumLogProbOfAction'+str(num_step)],h_t_train,_,_, mask_visited_nodes_train,_  = model_train(h_t_train,
                                                                                                                                        h_encoder=h_encoder_train,
                                                                                                                                        step_now=num_step,plane_num=plane_num,num_choices=num_choices,max_steps = max_steps,
                                                                                                                                        mask_visited_nodes=mask_visited_nodes_train,K_att_decoder=K_att_decoder_train,
                                                                                                                                        V_att_decoder=V_att_decoder_train,deterministic=d_train)
            
            # 计算基准模型的路径
            with torch.no_grad():  # 基准模型不计算梯度
                if num_step == 0: 
                    tour_baseline,_,h_t_baseline,K_att_decoder_baseline, V_att_decoder_baseline, mask_visited_nodes_baseline,h_encoder_baseline  = model_baseline(com,h_encoder=h_encoder,
                                                                                                                                                               step_now=num_step,plane_num=plane_num,num_choices=num_choices
                                                                                                                                                               ,max_steps = max_steps,mask_visited_nodes=mask_visited_nodes_baseline,K_att_decoder=K_att_decoder,
                                                                                                                                                               V_att_decoder=V_att_decoder,deterministic=True) 
                else:
                    tour_baseline,_,h_t_baseline,K_att_decoder_baseline, V_att_decoder_baseline, mask_visited_nodes_baseline,h_encoder_baseline  = model_baseline(h_t_baseline,h_encoder=h_encoder_baseline,
                                                                                                                                                               step_now=num_step,plane_num=plane_num,num_choices=num_choices,max_steps = max_steps,
                                                                                                                                                               mask_visited_nodes=mask_visited_nodes_baseline,K_att_decoder=K_att_decoder_baseline,
                                                                                                                                                               V_att_decoder=V_att_decoder_baseline,deterministic=True)
            reward_train = torch.zeros(batch)
            grid_size = (10, 60, 60)  # 区域大小
            resolution = 0.2 # 体素分辨率
            polyhedra_vertices_train = [[]for k in range(batch)]
            polyhedra_vertices_baseline = [[]for k in range(batch)]
            all_voxel_flags = create_voxel_grid_torch(all_voxel_flags,start_x,grid_size,resolution,y_start=20,z_start=20,device=device)
            for k in range(batch):
                for p_num in range(plane_num):
                    polyhedra_vertices_train[k].append(t1[k][p_num][tour_train[k][p_num]])
            # 体素化并标记
            with torch.no_grad():
                voxel_flags_train = voxelize_with_convex_polyhedra_torch(all_voxel_flags, voxel_flags_train,grid_size, resolution, polyhedra_vertices_train, device=device)
            reward_train = torch.sum(voxel_flags_train>0,dim=1) * resolution ** 3
            reward_train1 = reward_train.clone() - reward_train_old.clone()
            reward_train_old = reward_train.clone()
            reward_baseline = torch.zeros(batch)
            for k in range(batch):
                for p_num in range(plane_num):
                    polyhedra_vertices_baseline[k].append(t1[k][p_num][tour_baseline[k][p_num]])
            # 体素化并标记
            with torch.no_grad():
                voxel_flags_baseline = voxelize_with_convex_polyhedra_torch(all_voxel_flags, voxel_flags_baseline,grid_size, resolution, polyhedra_vertices_baseline, device=device)
            reward_baseline = torch.sum(voxel_flags_baseline>0,dim=1) * resolution ** 3
            # print(len(polyhedra_vertices_train[0]),len(polyhedra_vertices_baseline[0]))
            
            reward_baseline1 = reward_baseline.clone() - reward_baseline_old.clone()
            # print(reward_train1,reward_baseline1)
            reward_baseline_old = reward_baseline.clone()
            # 反向传播：损失函数为路径差异的加权和
            loss = torch.mean((- reward_train1.detach() + reward_baseline1.detach()) * names['sumLogProbOfAction'+str(num_step)])
            total_loss = total_loss + loss
        mean_epoch_train = mean_epoch_train + reward_train
        mean_epoch_baseline = mean_epoch_baseline + reward_baseline
        optimizer.zero_grad()  # 梯度清零
        total_loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新模型参数
    print(f'mean_train: {torch.mean(mean_epoch_train)/max_step_per_epoch}')
    print(f'mean_base: {torch.mean(mean_epoch_baseline)/max_step_per_epoch}')
    model_train.eval()  # 设置训练模型为评估模式
    mean_tour_length_train = 0
    mean_tour_length_baseline = 0
    for step_per_epoch1 in range(5):
        t1,t2,start_x = get_train_text(batch,num_choices,plane_num)
        com = combine_features(t1,t2)
        # 初始化访问节点的掩码：表示节点是否已被访问，起始占位符节点标记为已访问
        mask_visited_nodes_train = torch.zeros((batch, plane_num * num_choices + 1), device=device).bool()  # 全部初始化为False
        mask_visited_nodes_baseline = torch.zeros((batch, plane_num * num_choices + 1), device=device).bool()  # 全部初始化为False
        K_att_decoder = torch.zeros((batch,plane_num * num_choices + 1,dim_emb*nb_layers_decoder),device=device)
        V_att_decoder = torch.zeros((batch,plane_num * num_choices + 1,dim_emb*nb_layers_decoder),device=device)
        h_encoder = torch.zeros((batch, plane_num * num_choices + 1, dim_emb),device=device)
        reward_baseline_old = torch.zeros(batch,device=device)
        reward_train_old = torch.zeros(batch,device=device)
        voxel_flags_train = torch.zeros(batch,10 * 5 * 60 * 5 *60 * 5, dtype=torch.int, device=device)
        voxel_flags_baseline = torch.zeros(batch,10 * 5 * 60 * 5 *60 * 5, dtype=torch.int, device=device)
        all_voxel_flags = torch.zeros((batch,10 * 5 * 60 * 5 *60 * 5, 3), device = device ,dtype=torch.float16)
        for num_step in range(max_steps): 
        # 前二十步进行验证
            with torch.no_grad():  # 基准模型不计算梯度
                if num_step == 0: 
                    tour_train,_,h_t_train,K_att_decoder_train, V_att_decoder_train, mask_visited_nodes_train,h_encoder_train = model_train(com,h_encoder=h_encoder,step_now=num_step,
                                                                                                                                                                   plane_num=plane_num,num_choices=num_choices,max_steps=max_steps,
                                                                                                                                                                   mask_visited_nodes=mask_visited_nodes_train,K_att_decoder=K_att_decoder,
                                                                                                                                                                   V_att_decoder=V_att_decoder,deterministic=True)
                else:
                    tour_train,_,h_t_train,K_att_decoder_train, V_att_decoder_train, mask_visited_nodes_train,h_encoder_train  = model_train(h_t_train,
                                                                                                                                        h_encoder=h_encoder_train,
                                                                                                                                        step_now=num_step,plane_num=plane_num,num_choices=num_choices,max_steps = max_steps,
                                                                                                                                        mask_visited_nodes=mask_visited_nodes_train,K_att_decoder=K_att_decoder_train,
                                                                                                                                        V_att_decoder=V_att_decoder_train,deterministic=True)
            # 计算基准模型的路径
            with torch.no_grad():  # 基准模型不计算梯度
                if num_step == 0: 
                    tour_baseline,_,h_t_baseline,K_att_decoder_baseline, V_att_decoder_baseline, mask_visited_nodes_baseline,h_encoder_baseline  = model_baseline(com,h_encoder=h_encoder,
                                                                                                                                                               step_now=num_step,plane_num=plane_num,num_choices=num_choices
                                                                                                                                                               ,max_steps = max_steps,mask_visited_nodes=mask_visited_nodes_baseline,K_att_decoder=K_att_decoder,
                                                                                                                                                               V_att_decoder=V_att_decoder,deterministic=True) 
                else:
                    tour_baseline,_,h_t_baseline,K_att_decoder_baseline, V_att_decoder_baseline, mask_visited_nodes_baseline,h_encoder_baseline  = model_baseline(h_t_baseline,h_encoder=h_encoder_baseline,
                                                                                                                                                               step_now=num_step,plane_num=plane_num,num_choices=num_choices,max_steps = max_steps,
                                                                                                                                                               mask_visited_nodes=mask_visited_nodes_baseline,K_att_decoder=K_att_decoder_baseline,
                                                                                                                                                               V_att_decoder=V_att_decoder_baseline,deterministic=True)
            reward_train = torch.zeros(batch)
            grid_size = (10, 60, 60)  # 区域大小
            resolution = 0.2 # 体素分辨率
            polyhedra_vertices_train = [[]for k in range(batch)]
            polyhedra_vertices_baseline = [[]for k in range(batch)]
            all_voxel_flags = create_voxel_grid_torch(all_voxel_flags,start_x,grid_size,resolution,y_start=20,z_start=20,device=device)
            for k in range(batch):
                for p_num in range(plane_num):
                    polyhedra_vertices_train[k].append(t1[k][p_num][tour_train[k][p_num]])
            # 体素化并标记
            with torch.no_grad():
                voxel_flags_train = voxelize_with_convex_polyhedra_torch(all_voxel_flags, voxel_flags_train,grid_size, resolution, polyhedra_vertices_train, device=device)
            reward_train = torch.sum(voxel_flags_train>0,dim=1) * resolution ** 3
            reward_train1 = reward_train.clone() - reward_train_old.clone()
            reward_train_old = reward_train.clone()
            reward_baseline = torch.zeros(batch)
            for k in range(batch):
                for p_num in range(plane_num):
                    polyhedra_vertices_baseline[k].append(t1[k][p_num][tour_baseline[k][p_num]])
            # 体素化并标记
            with torch.no_grad():
                voxel_flags_baseline = voxelize_with_convex_polyhedra_torch(all_voxel_flags, voxel_flags_baseline,grid_size, resolution, polyhedra_vertices_baseline, device=device)
            reward_baseline = torch.sum(voxel_flags_baseline>0,dim=1) * resolution ** 3
            # print(len(polyhedra_vertices_train[0]),len(polyhedra_vertices_baseline[0]))
            # print(reward_train,reward_baseline)
            reward_baseline1 = reward_baseline.clone() - reward_baseline_old.clone()
            reward_baseline_old = reward_baseline.clone()
            mean_tour_length_train += reward_train1.mean().item()
            mean_tour_length_baseline += reward_baseline1.mean().item()
    mean_tour_length_train = mean_tour_length_train /5  # 训练模型的平均路径长度
    mean_tour_length_baseline = mean_tour_length_baseline /5  # 基准模型的平均路径长度
    print(f"train:  {mean_tour_length_train}")
    print(f'base_line: {mean_tour_length_baseline}')
    # 评估并更新基准模型
    update_baseline = (mean_tour_length_train > 0.1 + mean_tour_length_baseline) or ( mean_tour_length_train  > max_mean_train)  # 如果训练模型表现更好，则更新基准模型
    if update_baseline:
        torch.save(model_train,  '8.4.4.20_' + str(round(mean_tour_length_train,0)) + '.pkl')
        if (mean_tour_length_train > 0.1 + mean_tour_length_baseline):
            model_baseline.load_state_dict(model_train.state_dict())
        elif ( mean_tour_length_train  > max_mean_train):
            max_mean_train = mean_tour_length_train