地形分类可视化

[LRuYuan]

打扰一下，我想要知道训练后，这个地形分类分出的情况，就用论文中提到的t-SNE可视化encoder网络对于地形的编码结果，发现编码的结果仅仅只有一种地形，下面是我的用于可视化的代码，temperature默认是3,我调整这个参数，便可以分出不同的地形，但是默认是3,请问是不是这个temperature=3的时候这个对地形的编码根本没有实现？还是说我的这个可视化的代码逻辑有问题？
def update(self, obs_history, next_critic_obs, lr=None):
# 用于降采样

    if lr is not None:
        self.learning_rate = lr
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.learning_rate
            
    vel = next_critic_obs[:, self.num_one_step_obs:self.num_one_step_obs+3].detach()
    next_obs = next_critic_obs.detach()[:, 3:self.num_one_step_obs+3]

    z_s = self.encoder(obs_history)
    z_t = self.target(next_obs)
    pred_vel, z_s = z_s[..., :3], z_s[..., 3:]

    # print("obs_history shape:", obs_history.shape)
    # print("next_obs shape:", next_obs.shape)

    z_s = F.normalize(z_s, dim=-1, p=2)
    z_t = F.normalize(z_t, dim=-1, p=2)

    with torch.no_grad():
        w = self.proto.weight.data.clone()
        w = F.normalize(w, dim=-1, p=2)
        self.proto.weight.copy_(w)

    score_s = z_s @ self.proto.weight.T
    score_t = z_t @ self.proto.weight.T

    # print("score_s shape",score_s.shape)

    with torch.no_grad():
        q_s = sinkhorn(score_s)
        q_t = sinkhorn(score_t)
    # print("q_s shape",q_s.shape)

    log_p_s = F.log_softmax(score_s / self.temperature, dim=-1)
    log_p_t = F.log_softmax(score_t / self.temperature, dim=-1)
    print(f"temperature:{self.temperature}")
    
    if self.epoch % 8000 == 0:
        # 1. 提取【全量】10.4万个数据点
        # 这一步是为了让所有的点都参与构建隐式结构（邻域图）
        full_data = z_s.detach().cpu().numpy()
        full_labels = torch.argmax(q_s, dim=1).detach().cpu().numpy()

        # 2. 使用 openTSNE 对全量数据进行快速降维
        # 它底层使用 k-NN 树，速度极快，10万数据通常只需不到 1 分钟
        tsne = FastTSNE(
            n_components=2,
            perplexity=30,
            n_jobs=-1,  # 使用所有 CPU 核心并行计算
            random_state=42
        )
        # full_feature_2d 包含了全部 10.4 万个点的二维坐标
        full_feature_2d = tsne.fit(full_data)

        # 3. 子集采样（对应论文中的“只显示数据子集”）
        # 坐标已经受到 10.4 万点整体流形的影响，此时我们随机抽取 1 万个点用来画图
        n_landmarks = 5000
        indices = np.random.permutation(full_data.shape[0])[:n_landmarks]
        
        sampled_feature_2d = full_feature_2d[indices]
        sampled_labels = full_labels[indices]

        # 4. 绘图（画点太多容易变成一坨黑，所以调小点的大小 s=5）
        plt.figure(figsize=(20, 16))
        plt.scatter(sampled_feature_2d[:, 0], sampled_feature_2d[:, 1], 
                    c=sampled_labels, cmap='viridis', s=5, alpha=0.8)
        plt.title(f"Terrain Classification at Epoch {self.epoch} (10k sampled from 104k)")

        plt.savefig(f"picture/epoch_{self.epoch}.png")
        plt.close()
        
    self.epoch +=1


    swap_loss = -0.5 * (q_s * log_p_t + q_t * log_p_s).mean()

    estimation_loss = F.mse_loss(pred_vel, vel)

    losses = estimation_loss + swap_loss

    self.optimizer.zero_grad()
    losses.backward()
    nn.utils.clip_grad_norm_(self.parameters(), self.max_grad_norm)
    self.optimizer.step()

    return estimation_loss.item(), swap_loss.item()