Blame for scripts/2_cluster_and_visualize.py

"""
VC-ACT 阶段 2: 聚类与可视化

读取阶段 1 输出的 (N, 2048) 特征，执行:
  1. K-Means 聚类 (固定 k=10 和 自动搜索最优 k)
  2. Silhouette score 选最优 k (k=2..max_k)
  3. t-SNE 可视化 (按聚类标签着色 + 按 bag 编号着色)
  4. 每个簇包含哪些 bag 的详细分布

输出:
  - results/cluster_labels_k10.npy        (N,) int, k=10 聚类标签
  - results/cluster_labels_auto.npy       (N,) int, 最优 k 聚类标签
  - results/cluster_centers_k10.npy       (10, 2048) 簇中心
  - results/cluster_centers_auto.npy      (k*, 2048) 簇中心
  - results/optimal_k.json               {"best_k": k*, ...}
  - results/tsne_k10.png                 t-SNE (k=10 聚类着色)
  - results/tsne_auto.png                t-SNE (auto k 聚类着色)
  - results/silhouette_scores.png        Silhouette score 曲线

用法:
  /venv/mult-act/bin/python3 2_cluster_and_visualize.py
  /venv/mult-act/bin/python3 2_cluster_and_visualize.py --max_k 15
"""

import argparse
import json
from pathlib import Path

import numpy as np
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import StandardScaler
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def run_kmeans(features: np.ndarray, k: int, seed: int = 42) -> tuple:
    """Run K-Means and return (labels, centers, inertia)."""
    km = KMeans(n_clusters=k, random_state=seed, n_init=10, max_iter=300)
    labels = km.fit_predict(features)
    return labels, km.cluster_centers_, km.inertia_


def find_optimal_k(features: np.ndarray, max_k: int, seed: int = 42) -> tuple:
    """Compute silhouette scores for k=2..max_k, return results dict and best k."""
    results = {}
    for k in range(2, max_k + 1):
        labels, _, inertia = run_kmeans(features, k, seed)
        score = silhouette_score(features, labels)
        results[k] = {"silhouette": score, "inertia": inertia}
        print(f"  k={k:2d}: silhouette={score:.4f}, inertia={inertia:.1f}")
    best_k = max(results, key=lambda k: results[k]["silhouette"])
    return results, best_k


def plot_silhouette(scores_dict: dict, best_k: int, save_path: str):
    """Plot silhouette scores vs k."""
    ks = sorted(scores_dict.keys())
    sils = [scores_dict[k]["silhouette"] for k in ks]

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(ks, sils, "o-", color="steelblue", linewidth=2, markersize=8)
    ax.axvline(best_k, color="red", linestyle="--", alpha=0.7, label=f"best k={best_k}")
    ax.set_xlabel("Number of clusters (k)", fontsize=12)
    ax.set_ylabel("Silhouette Score", fontsize=12)
    ax.set_title("Silhouette Score vs K", fontsize=14)
    ax.set_xticks(ks)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(save_path, dpi=150)
    plt.close(fig)
    print(f"Saved silhouette plot to {save_path}")


def plot_tsne(features: np.ndarray, cluster_labels: np.ndarray, bag_names: list,
              k_value: int, save_path: str, seed: int = 42):
    """Plot t-SNE colored by cluster labels, annotated with bag names."""
    tsne = TSNE(n_components=2, random_state=seed, perplexity=min(30, len(features) - 1))
    coords = tsne.fit_transform(features)

    fig, ax = plt.subplots(figsize=(12, 9))
    scatter = ax.scatter(coords[:, 0], coords[:, 1], c=cluster_labels,
                         cmap="tab10", s=60, alpha=0.8, edgecolors="k", linewidths=0.3)
    ax.set_title(f"t-SNE Visualization (k={k_value})", fontsize=14)
    ax.legend(*scatter.legend_elements(), title="Cluster", loc="best", fontsize=9)
    ax.set_xlabel("t-SNE dim 1")
    ax.set_ylabel("t-SNE dim 2")
    ax.grid(True, alpha=0.2)

    # Annotate a subset of points with bag names to avoid clutter
    # Only annotate every N-th point if too many
    step = max(1, len(bag_names) // 30)
    for i in range(0, len(bag_names), step):
        short_name = bag_names[i].replace("pick_hanger_", "").replace(".bag", "")
        ax.annotate(short_name, (coords[i, 0], coords[i, 1]),
                    fontsize=6, alpha=0.6, ha="center", va="bottom")

    fig.tight_layout()
    fig.savefig(save_path, dpi=150)
    plt.close(fig)
    print(f"Saved t-SNE plot to {save_path}")


def main(args):
    results_dir = Path(args.results_dir)

    # Load features and metadata
    features = np.load(results_dir / "features_2048.npy")
    with open(results_dir / "episode_meta.json") as f:
        meta = json.load(f)
    bag_names = [m["bag_name"] for m in meta]
    print(f"Loaded features: {features.shape}, {len(bag_names)} bags")

    # Standardize features for better clustering
    scaler = StandardScaler()
    features_scaled = scaler.fit_transform(features)

    # === 1. Fixed k=10 ===
    print("\n=== K-Means k=10 ===")
    labels_k10, centers_k10, _ = run_kmeans(features_scaled, k=10, seed=args.seed)
    np.save(results_dir / "cluster_labels_k10.npy", labels_k10)
    np.save(results_dir / "cluster_centers_k10.npy", centers_k10)
    sil_k10 = silhouette_score(features_scaled, labels_k10)
    print(f"k=10 silhouette score: {sil_k10:.4f}")

    # === 2. Auto k selection ===
    print(f"\n=== Silhouette search k=2..{args.max_k} ===")
    scores_dict, best_k = find_optimal_k(features_scaled, args.max_k, args.seed)
    print(f"\nOptimal k = {best_k} (silhouette={scores_dict[best_k]['silhouette']:.4f})")

    labels_auto, centers_auto, _ = run_kmeans(features_scaled, k=best_k, seed=args.seed)
    np.save(results_dir / "cluster_labels_auto.npy", labels_auto)
    np.save(results_dir / "cluster_centers_auto.npy", centers_auto)

    with open(results_dir / "optimal_k.json", "w") as f:
        json.dump({
            "best_k": best_k,
            "silhouette_score": scores_dict[best_k]["silhouette"],
            "all_scores": {str(k): v for k, v in scores_dict.items()},
        }, f, indent=2)

    # === 3. Visualizations ===
    print("\n=== Generating plots ===")

    # Silhouette scores
    plot_silhouette(scores_dict, best_k, str(results_dir / "silhouette_scores.png"))

    # t-SNE for k=10
    plot_tsne(features_scaled, labels_k10, bag_names, 10,
              str(results_dir / "tsne_k10.png"), seed=args.seed)

    # t-SNE for auto k
    if best_k != 10:
        plot_tsne(features_scaled, labels_auto, bag_names, best_k,
                  str(results_dir / "tsne_auto.png"), seed=args.seed)

    # === 4. Summary ===
    print("\n=== Cluster distribution (k=10) ===")
    for c in range(10):
        bag_indices = np.where(labels_k10 == c)[0]
        names = [bag_names[i] for i in bag_indices]
        print(f"  cluster {c}: {len(bag_indices)} bags")
        for n in names:
            print(f"    - {n}")

    print(f"\n=== Cluster distribution (k={best_k}, auto) ===")
    for c in range(best_k):
        bag_indices = np.where(labels_auto == c)[0]
        names = [bag_names[i] for i in bag_indices]
        print(f"  cluster {c}: {len(bag_indices)} bags")
        for n in names:
            print(f"    - {n}")

    print("\nDone!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VC-ACT: Cluster and visualize")
    parser.add_argument("--results_dir", type=str,
                        default="/workspace/Mult-skill ACT/VC-ACT/results")
    parser.add_argument("--max_k", type=int, default=10,
                        help="Maximum k to search (silhouette search from 2 to max_k)")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    main(args)