Upload code/gate_success_fail.py with huggingface_hub

code/gate_success_fail.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Definitive RLT encoder gate: do SUCCESS vs FAILURE z_rl separate?
+
+The RL Token paper validates the token by showing success frames form a smooth
+manifold while failures (failed grasps / tower collapses) cluster apart. This is
+the real test (stronger than the success-only smoothness gate in tsne_gate.py).
+
+  success = the 44 teleop demos        -> encoder_cache_prefix/
+  failure = frozen-baseline rollouts   -> encoder_cache_fail/   (SR~0)
+
+Reports: linear separability (LogisticRegression 5-fold CV accuracy on z_rl) and
+silhouette score, plus a t-SNE colored by class. High CV-acc + clear visual
+split = z_rl encodes task-success information => good RL state.
+
+CPU-only (won't fight the GPU policy server). Run AFTER collecting failures:
+  ./lerobot/.venv/bin/python gate_success_fail.py --ckpt checkpoints/rl_token_encoder_FINAL.pt
+"""
+from __future__ import annotations
+import argparse, glob, os
+import numpy as np, torch
+import matplotlib; matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
+from sklearn.manifold import TSNE
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import cross_val_score
+from sklearn.metrics import silhouette_score
+from rl_token_encoder import RLTokenAutoencoder, RLTokenConfig
+
+
+def encode_dir(ae, d, cap):
+    fs = sorted(glob.glob(os.path.join(d, "*.npz")))
+    if not fs:
+        return np.zeros((0, 2560), np.float32)
+    if len(fs) > cap:
+        fs = fs[:: len(fs) // cap]
+    Z = []
+    with torch.no_grad():
+        for f in fs:
+            e = torch.from_numpy(np.load(f)["embeddings"].astype(np.float32))[None]
+            m = torch.ones(1, e.shape[1], dtype=torch.bool)
+            Z.append(ae.encode(e, m)[0].numpy())
+    return np.stack(Z)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="checkpoints/rl_token_encoder_FINAL.pt")
+    p.add_argument("--success-dir", default="./encoder_cache_prefix")
+    p.add_argument("--fail-dir", default="./encoder_cache_fail")
+    p.add_argument("--weights", default="ema", choices=["ema", "model"])
+    p.add_argument("--cap", type=int, default=800, help="max points per class (balanced)")
+    p.add_argument("--out", default="outputs/gate_success_fail.png")
+    args = p.parse_args()
+
+    ck = torch.load(args.ckpt, map_location="cpu")
+    ae = RLTokenAutoencoder(RLTokenConfig(dim=2560)); ae.load_state_dict(ck[args.weights]); ae.eval()
+    print(f"loaded {args.ckpt} ({args.weights})")
+
+    Zs = encode_dir(ae, args.success_dir, args.cap)
+    Zf = encode_dir(ae, args.fail_dir, args.cap)
+    print(f"success z_rl: {len(Zs)}   failure z_rl: {len(Zf)}")
+    if len(Zf) < 10:
+        print("⚠️ <10 failure shards — collect baseline rollouts into", args.fail_dir, "first."); return
+
+    n = min(len(Zs), len(Zf))  # balance
+    Zs, Zf = Zs[:n], Zf[:n]
+    X = np.concatenate([Zs, Zf]); y = np.array([0] * n + [1] * n)
+
+    # 1) linear separability (the quantitative verdict)
+    clf = LogisticRegression(max_iter=2000, C=1.0)
+    acc = cross_val_score(clf, X, y, cv=5, scoring="accuracy").mean()
+    sil = silhouette_score(X, y)
+    print(f"\nSEPARATION:")
+    print(f"  LogReg 5-fold CV accuracy = {acc:.1%}   (50%=indistinguishable, >80%=clearly separable)")
+    print(f"  silhouette (by class)     = {sil:.3f}   (>0 = classes form distinct groups)")
+
+    # 2) t-SNE colored by class
+    Xp = PCA(n_components=min(50, X.shape[0])).fit_transform(X)
+    emb = TSNE(n_components=2, perplexity=30, init="pca", random_state=0).fit_transform(Xp)
+    plt.figure(figsize=(7, 6))
+    plt.scatter(emb[:n, 0], emb[:n, 1], c="tab:green", s=10, label="success (demos)", alpha=0.6)
+    plt.scatter(emb[n:, 0], emb[n:, 1], c="tab:red", s=10, label="failure (baseline)", alpha=0.6)
+    plt.legend(); plt.title(f"z_rl: success vs failure  (LogReg CV acc {acc:.0%}, silhouette {sil:.2f})")
+    os.makedirs(os.path.dirname(args.out), exist_ok=True)
+    plt.tight_layout(); plt.savefig(args.out, dpi=120)
+    print(f"saved {args.out}")
+    print("\nGATE:", "✅ z_rl separates success from failure — good RL state" if acc > 0.8
+          else "⚠️ weak separation — z_rl may not capture task success cleanly")
+
+
+if __name__ == "__main__":
+    main()