refactor: 将子模块转为普通目录，移除外部 git 依赖

- 移除 JEPA/lejepa-identifiability 子模块 gitlink
- 移除 research/multiply/MultiPLY 子模块 gitlink
- 删除 .gitmodules（不再有外部 URL 依赖）
- 两个目录内容作为普通文件纳入主仓库追踪
- 删除各自内部 .git 目录，消除嵌套 git 仓库

JEPA/lejepa-identifiability/.gitignore

@@ -0,0 +1,27 @@
+# Python
+__pycache__/
+*.pyc
+*.egg-info/
+.venv/
+venv/
+
+# Jupyter
+.ipynb_checkpoints/
+local/
+
+# Data
+data/
+
+# Results (large tensors)
+results*/
+figures/
+logs/
+
+# Lean
+lean/.lake/
+lean/lake-packages/
+
+# OS / editors
+.DS_Store
+.vscode/
+.idea/

JEPA/lejepa-identifiability/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 David Klindt
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

JEPA/lejepa-identifiability/README.md

@@ -0,0 +1,155 @@
+# LeJEPA Identifiability
+### When Does LeJEPA Learn a World Model?
+
+[David Klindt](https://scholar.google.com/citations?user=EpT-nUAAAAAJ&hl=en), [Yann LeCun](https://scholar.google.com/citations?user=WLN3QrAAAAAJ&hl=en) and [Randall Balestriero](https://scholar.google.com/citations?user=S1x_xqcAAAAJ&hl=en&oi=ao)
+
+**Abstract:** A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization. We prove that LeJEPA (alignment plus Gaussian regularization) linearly recovers the world's latent variables from nonlinear observations, a property known as *linear identifiability*, in a broad class of worlds where latents evolve under stationary, additive-noise transitions. Our main result is that among all such worlds, the Gaussian is the *unique* latent distribution for which this guarantee holds. The forward direction rests on a spectral decomposition in which each degree of nonlinearity is strictly penalized by alignment, making the linear map the optimum; the converse rules out every non-Gaussian alternative. We further prove an *approximate identifiability* result where the guarantee degrades gracefully, and show that linear, orthogonal identifiability enables *optimal latent-space planning*. We validate the theory across 2D examples to 1024-dimensional latents, distributional ablations, and pixel-based robotic control. All theorems are formally verified in Lean 4.
+
+<p align="center">
+   <b>[ <a href="https://arxiv.org/abs/2605.26379">Paper</a> | <a href="https://klindtlab.github.io/lejepa-identifiability/">Website</a> | <a href="https://colab.research.google.com/drive/1ozjRk3FfUIDX7WBqlOKvhNcIamy0JxCH?usp=sharing">Colab</a> | <a href="https://youtu.be/EioGDo67ZDs">Video</a> ]</b>
+</p>
+
+<p align="center">
+   <img width="600" height="338" alt="world_model" src="https://github.com/user-attachments/assets/580f4e9b-cf2e-4b31-88e4-257fa16ea40c" />
+</p>
+
+If you find this work useful, please cite:
+
+```bibtex
+@article{klindt2026lejepa,
+  title={When Does LeJEPA Learn a World Model?},
+  author={Klindt, David and LeCun, Yann and Balestriero, Randall},
+  journal={arXiv preprint arXiv:TODO},
+  year={2026}
+}
+```
+
+## Quick Start
+
+Try the 2D demo in your browser (~30s on a T4 GPU):
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ozjRk3FfUIDX7WBqlOKvhNcIamy0JxCH?usp=sharing)
+
+## Repository Structure
+
+```
+lejepa-identifiability/
+├── lean/                           # Lean 4 formal verification
+│   ├── LeJEPA/
+│   │   ├── Hermite.lean            # Forward direction (Hermite polynomial proof)
+│   │   ├── Uniqueness.lean         # Converse (Gaussian uniqueness)
+│   │   ├── Approx.lean             # Approximate identifiability bound
+│   │   └── Dirichlet.lean          # Alternative proof (Dirichlet energy)
+│   ├── LeJEPA.lean
+│   ├── lakefile.lean
+│   └── lean-toolchain              # Lean 4 v4.28.0
+├── experiments/
+│   ├── lejepa_id/                  # Shared library
+│   │   ├── mixing.py               # Mixing functions (spiral, banana, sinusoid, coupling)
+│   │   ├── models.py               # MLP and matched (inverse-NVP) encoders
+│   │   ├── losses.py               # SIGReg, whitening, alignment, InfoNCE
+│   │   ├── metrics.py              # R², orthogonality, bound quantities
+│   │   ├── data.py                 # Gaussian / generalized-normal sampling, OU augmentation
+│   │   ├── reacher.py              # Reacher pixel data utilities
+│   │   └── engine.py               # Training loop (warmup + cosine LR, online data)
+│   ├── run.py                      # Unified runner for 2D / scaling / gennorm / grid
+│   ├── run_reacher.py              # Reacher pixel-observation runner
+│   ├── prerender.py                # Render Reacher OU and trajectory frames
+│   ├── analysis/                   # Post-hoc plotting and tables
+│   ├── configs/                    # Experiment hyperparameters (YAML)
+│   │   ├── 2d.yaml
+│   │   ├── gennorm.yaml
+│   │   ├── scaling.yaml
+│   │   ├── grid.yaml
+│   │   └── reacher.yaml
+│   └── slurm/                      # SLURM launch scripts (CSHL cluster)
+├── requirements.txt
+└── README.md
+```
+
+## Formal Verification (Lean 4)
+
+All theoretical results are formalized in Lean 4 with Mathlib. The project compiles with **zero `sorry` obligations** — every logical chain from axiomatized premises to conclusions is machine-checked. Axiomatized components are standard results not yet available in Mathlib (Hermite polynomial infrastructure, Mazur–Ulam, AM–GM with uniform weights). See the paper appendix for the full verification inventory.
+
+```bash
+cd lean
+lake build    # requires Lean 4 v4.28.0; fetches Mathlib automatically
+```
+
+## Experiments
+
+All experiments share the same training infrastructure (`lejepa_id/engine.py`) and read parameters from YAML configs. Training uses online data generation, a warmup + cosine LR schedule, and saves results as `.json` (scalars and curves); 2D and ablation runs additionally save `.pt` files with scatter arrays.
+
+```bash
+pip install -r requirements.txt
+cd experiments
+```
+
+### 2D Illustrations
+
+Four mixing functions (spiral, banana, sinusoidal shear, NVP) with MLP or matched encoders.
+
+```bash
+python run.py --config configs/2d.yaml --run spiral_lejepa --seed 1337
+python analysis/plot_2d.py --results_dir results/2d/ --out figures/
+```
+
+### Scaling (N = 2 to 1024)
+
+Matched (inverse-NVP) encoder scaling with latent dimension, swept across SIGReg / VICReg / InfoNCE objectives. Each (N, seed) trains K=3 encoders in parallel for N ≤ 32 and picks the best by final loss.
+
+```bash
+python run.py --config configs/scaling.yaml --N 16 --seed 0
+python run.py --config configs/scaling.yaml --N 16 --seed 0 --mode infonce
+python analysis/plot_scaling.py --results_dir results/scaling/ --out figures/
+```
+
+### Distributional Ablation (Generalized Normal)
+
+Same mixings sweeping the latent shape parameter α (heavy-tailed → Laplace → Gaussian → uniform). Demonstrates that linear identifiability fails away from the Gaussian (α = 2).
+
+```bash
+python run.py --config configs/gennorm.yaml --run spiral_lejepa --alpha 2.0 --seed 1337
+python analysis/plot_gennorm.py --results_dir results/gennorm/ --out figures/
+```
+
+### Grid Search / Bound Verification
+
+Sweep over regularization weight λ and OU correlation ρ on the 2D spiral mixing.
+
+```bash
+python run.py --config configs/grid.yaml --lamb 0.01 --rho 0.9 --seed 0
+python analysis/plot_bound.py --results_dirs results/grid results/2d results/scaling --out figures/
+```
+
+### Reacher (Pixel-Based RL)
+
+CNN encoder on rendered DMC Reacher frames, comparing OU pairs against trajectory pairs from a learned policy.
+
+```bash
+python prerender.py ou --rho 0.95
+python prerender.py traj --delta 16 --h5_path data/reacher.h5
+python run_reacher.py --config configs/reacher.yaml --data_dir data/reacher/ou/rho=0.95
+```
+
+### Cross-Experiment Analysis
+
+```bash
+python analysis/aggregate.py --results_dir results/ --recursive --out results/all.csv
+python analysis/plot_scatter.py --results_dirs results/2d results/gennorm results/scaling results/grid --out figures/
+```
+
+### Regenerate All Figures
+
+```bash
+bash analysis/run_all.sh
+```
+
+## Requirements
+
+- **Lean**: v4.28.0 + Mathlib v4.28.0 (managed by `lake`)
+- **Python**: `pip install -r requirements.txt`
+
+## License
+
+MIT

JEPA/lejepa-identifiability/experiments/analysis/aggregate.py

@@ -0,0 +1,63 @@
+"""
+Aggregate results from any experiment into a flat CSV.
+
+Usage:
+    python analysis/aggregate.py --results_dir results/2d/ --out results/2d/summary.csv
+    python analysis/aggregate.py --results_dir results/ --recursive --out results/all.csv
+"""
+
+import argparse, glob, json, os
+import pandas as pd
+
+
+SCALAR_KEYS = [
+    "experiment", "run_name", "mixing", "encoder", "mode", "source_dist",
+    "seed", "N", "lamb", "rho", "lr", "steps", "batch_size", "n_layers", "hidden",
+    "r2_zx", "r2_xz", "r2_zh", "r2_hz",
+    "orth_err", "orth_err_normalized",
+    "epsilon", "delta", "D_bound", "approx_bound",
+    "procrustes_mse", "L_h", "trace_cov",
+    "final_align", "final_sigreg", "final_whiten", "final_loss",
+]
+
+
+def load_results(results_dir, recursive=False):
+    pattern = os.path.join(results_dir, "**/*.json") if recursive else os.path.join(results_dir, "*.json")
+    files = sorted(glob.glob(pattern, recursive=recursive))
+    print(f"Found {len(files)} .json files")
+
+    rows = []
+    for path in files:
+        try:
+            with open(path) as f:
+                r = json.load(f)
+            row = {k: r.get(k) for k in SCALAR_KEYS}
+            row["file"] = os.path.relpath(path, results_dir)
+            rows.append(row)
+        except Exception as e:
+            print(f"  SKIP {path}: {e}")
+    return pd.DataFrame(rows)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", type=str, required=True)
+    p.add_argument("--out", type=str, default=None)
+    p.add_argument("--recursive", action="store_true")
+    args = p.parse_args()
+
+    df = load_results(args.results_dir, recursive=args.recursive)
+    if len(df) == 0:
+        print("No results found.")
+        return
+
+    print(f"\n{len(df)} runs loaded")
+    print(df.to_string(index=False))
+
+    out = args.out or os.path.join(args.results_dir, "summary.csv")
+    df.to_csv(out, index=False)
+    print(f"\nSaved {out}")
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/make_reacher_distributions.py

@@ -0,0 +1,371 @@
+"""
+Reacher trajectory distribution analysis figures.
+
+Produces two figures for the paper:
+  1. Scatter grid: stationary marginal + per-delta 2D transition differences
+     and per-dim (z_t, z_{t+delta}) scatters, annotated with R² and rho.
+  2. rho-vs-SIGReg scatter: three panels (z_0, z_1, joint), colored by R²,
+     showing the dual constraint that identifiability requires both
+     rho off from 1 and approximately-Gaussian transition shape.
+
+Usage:
+    python -m analysis.make_reacher_distributions \
+        --results_dir results/reacher \
+        --data_path data/reacher.h5 \
+        --out_dir figures/reacher
+"""
+
+import argparse
+import json
+from collections import defaultdict
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from scipy.stats import pearsonr
+
+
+DELTAS = [1, 2, 4, 8, 16, 32, 64]
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+N_MAX = 100_000
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# SIGReg (Epps–Pulley) — matches LeJEPA Algorithm 1
+# ═════════════════════════════════════════════════════════════════════════════
+
+def _sigreg_nd(x, num_slices=64, n_knots=17, seed=0, device=DEVICE):
+    """EP via random slicing on (N, K)."""
+    x = torch.as_tensor(np.asarray(x), dtype=torch.float32, device=device)
+    if x.dim() == 1:
+        x = x[:, None]
+    N, K = x.shape
+    g = torch.Generator(device=device).manual_seed(seed)
+    A = torch.randn(K, num_slices, generator=g, device=device)
+    A = A / A.norm(p=2, dim=0)
+    t = torch.linspace(-5, 5, n_knots, device=device)
+    phi = torch.exp(-0.5 * t ** 2)
+    zt = (x @ A).unsqueeze(2) * t
+    cm, sm = torch.cos(zt).mean(0), torch.sin(zt).mean(0)
+    err = ((cm - phi) ** 2 + sm ** 2) * phi
+    return (torch.trapz(err, t, dim=1) * N).mean().item()
+
+
+def _sigreg_1d(x, n_knots=17, device=DEVICE):
+    """EP directly on 1D (no slicing)."""
+    x = torch.as_tensor(np.asarray(x).reshape(-1), dtype=torch.float32, device=device)
+    N = x.shape[0]
+    t = torch.linspace(-5, 5, n_knots, device=device)
+    phi = torch.exp(-0.5 * t ** 2)
+    zt = x.unsqueeze(1) * t
+    cm, sm = torch.cos(zt).mean(0), torch.sin(zt).mean(0)
+    err = ((cm - phi) ** 2 + sm ** 2) * phi
+    return (torch.trapz(err, t) * N).item()
+
+
+def _zscore(x):
+    return (x - x.mean(0, keepdims=True)) / (x.std(0, keepdims=True) + 1e-8)
+
+
+def measure(x, n_draws=20, N_max=N_MAX, num_slices=64):
+    """
+    SIGReg raw + zscored for joint (K-d) and per-dim marginals.
+    Averages over n_draws random subsamples / projections.
+    Returns dict mapping key -> (mean, std) over draws.
+    """
+    x = np.asarray(x)
+    if x.ndim == 1:
+        x = x[:, None]
+    N, K = x.shape
+    rng = np.random.default_rng(0)
+    keys = ["joint_raw", "joint_zs"]
+    keys += [f"marg_{k}_raw" for k in range(K)]
+    keys += [f"marg_{k}_zs" for k in range(K)]
+    buf = {k: [] for k in keys}
+    for s in range(n_draws):
+        xs = x if N <= N_max else x[rng.choice(N, N_max, replace=False)]
+        xz = _zscore(xs)
+        buf["joint_raw"].append(_sigreg_nd(xs, num_slices, seed=s))
+        buf["joint_zs"].append(_sigreg_nd(xz, num_slices, seed=s))
+        for k in range(K):
+            buf[f"marg_{k}_raw"].append(_sigreg_1d(xs[:, k]))
+            buf[f"marg_{k}_zs"].append(_sigreg_1d(xz[:, k]))
+    return {k: (np.mean(v), np.std(v)) for k, v in buf.items()}
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# R² loading — per-seed best-lambda tuning, median across seeds
+# ═════════════════════════════════════════════════════════════════════════════
+
+def get_best_r2_per_delta(results_dir, agg="median"):
+    """
+    For each (delta, seed), pick the lambda with best R²; aggregate seeds
+    with median (robust to outliers) or mean.
+    """
+    by_dls = defaultdict(list)
+    for p in Path(results_dir).rglob("result.json"):
+        r = json.load(open(p))
+        if "delta" not in r or r.get("rho") is not None:
+            continue
+        by_dls[(r["delta"], r["lamb"], r.get("seed", 0))].append(r)
+
+    deltas = sorted({k[0] for k in by_dls})
+    lambs = sorted({k[1] for k in by_dls})
+    seeds = sorted({k[2] for k in by_dls})
+    agg_fn = np.median if agg == "median" else np.mean
+
+    best = {}
+    for delta in deltas:
+        per_seed = {"r2": [], "d0": [], "d1": [], "lamb": []}
+        for seed in seeds:
+            best_lamb, best_r2 = None, -np.inf
+            for lamb in lambs:
+                runs = by_dls.get((delta, lamb, seed), [])
+                if not runs:
+                    continue
+                r2 = np.mean([r["r2_hz"] for r in runs])
+                if r2 > best_r2:
+                    best_r2, best_lamb = r2, lamb
+            if best_lamb is None:
+                continue
+            runs = by_dls[(delta, best_lamb, seed)]
+            per_seed["r2"].append(np.mean([r["r2_hz"] for r in runs]))
+            per_seed["d0"].append(np.mean([r["r2_hz_per_dim"][0] for r in runs]))
+            per_seed["d1"].append(np.mean([r["r2_hz_per_dim"][1] for r in runs]))
+            per_seed["lamb"].append(best_lamb)
+        best[delta] = {
+            "r2": agg_fn(per_seed["r2"]),
+            "r2_dim0": agg_fn(per_seed["d0"]),
+            "r2_dim1": agg_fn(per_seed["d1"]),
+            "lamb": np.median(per_seed["lamb"]),
+            "n_seeds": len(per_seed["r2"]),
+        }
+    return best
+
+
+def best_ou_rho(results_dir):
+    """Return the rho of the OU run with highest mean R² across seeds/lambdas."""
+    grouped = defaultdict(list)
+    for p in Path(results_dir).rglob("result.json"):
+        r = json.load(open(p))
+        if "rho" not in r or r.get("rho") is None:
+            continue
+        grouped[(r["rho"], r["lamb"])].append(r)
+    if not grouped:
+        return None
+    best_rho, best_mean = None, -np.inf
+    for (rho, lamb), runs in grouped.items():
+        m = np.mean([r["r2_hz"] for r in runs])
+        if m > best_mean:
+            best_mean, best_rho = m, rho
+    return best_rho
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Figure 1: scatter grid
+# ═════════════════════════════════════════════════════════════════════════════
+
+def make_scatter_grid(episodes, r2_dict, save_path, sub=1, s=1e-4):
+    """
+    Left column: stationary marginal scatter of (z_0, z_1).
+    Top row, remaining columns: 2D transition-difference scatter per delta.
+    Bottom row, remaining columns: per-dim (z_t, z_{t+delta}) scatter.
+    Titles show R² and rho.
+    """
+    fig = plt.figure(figsize=0.85 * np.array((1 + 3 * len(DELTAS), 5)))
+    gs = fig.add_gridspec(2, 2 + len(DELTAS))
+
+    # stationary marginal (spans both rows, first two cols)
+    ax = fig.add_subplot(gs[:2, :2])
+    ax.scatter(*episodes.reshape(-1, 2)[::sub].T, s=s * 10)
+    ax.set_title("Marginal")
+    ax.grid()
+    ax.set_xlabel(r"$z_0$ (shoulder)")
+    ax.set_ylabel(r"$z_1$ (wrist)")
+
+    for i, delta in enumerate(DELTAS):
+        # top row: 2D transition differences
+        ax = fig.add_subplot(gs[0, 2 + i])
+        transitions = episodes[:, delta:] - episodes[:, :-delta]
+        ax.scatter(*transitions.reshape(-1, 2)[::sub].T, s=s)
+        r2_d0 = r2_dict[delta]["r2_dim0"]
+        r2_d1 = r2_dict[delta]["r2_dim1"]
+        ax.set_title(r"$\Delta=$" + f"{delta}" + "\n"
+                     r"$R^2=(%.2f, %.2f)$" % (r2_d0, r2_d1))
+        ax.grid()
+
+        # bottom row: per-dim (z_t, z_{t+delta}) with rho
+        ax = fig.add_subplot(gs[1, 2 + i])
+        a = episodes[:, delta:, 0].flatten()[::sub]
+        b = episodes[:, :-delta, 0].flatten()[::sub]
+        rho0 = pearsonr(a, b)[0]
+        ax.scatter(a, b, s=s)
+        c = episodes[:, delta:, 1].flatten()[::sub]
+        d = episodes[:, :-delta, 1].flatten()[::sub]
+        rho1 = pearsonr(c, d)[0]
+        ax.scatter(c, d, s=s)
+        ax.set_title(r"$\rho=(%.2f, %.2f)$" % (rho0, rho1))
+        ax.grid()
+        if i == 0:
+            ax.legend([r"$z_0$ (shoulder)", r"$z_1$ (wrist)"], loc="upper left")
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=500, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Figure 2: rho vs SIGReg scatter (3 panels)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def make_rho_vs_sigreg(trans, r2_dict, save_path, best_ou_rho_val=None,
+                       gaussian_floor=1.2):
+    """
+    Three panels (z_0, z_1, joint) showing rho vs SIGReg(zscored),
+    colored by R². Vertical line marks the best OU rho for reference.
+    """
+    fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))
+    names = [r"$z_0$ (shoulder)", r"$z_1$ (wrist)", "joint (avg across dims)"]
+
+    all_r2 = []
+    for d in DELTAS:
+        all_r2.append(r2_dict[d]["r2_dim0"])
+        all_r2.append(r2_dict[d]["r2_dim1"])
+        all_r2.append(r2_dict[d]["r2"])
+    vmin, vmax = min(all_r2), max(all_r2)
+
+    def _one(ax, rhos, sigs, errs, r2s, xlabel, ylabel, title):
+        if best_ou_rho_val is not None:
+            ax.axvline(best_ou_rho_val, color="crimson", lw=1.4, ls="--",
+                       alpha=0.8, label=fr"best OU $\rho={best_ou_rho_val:.2f}$",
+                       zorder=1)
+        sc = ax.scatter(rhos, sigs, c=r2s, cmap="viridis", s=140,
+                        vmin=vmin, vmax=vmax,
+                        edgecolors="black", linewidths=0.8, zorder=3)
+        ax.errorbar(rhos, sigs, yerr=errs, fmt="none", ecolor="gray",
+                    alpha=0.5, zorder=2)
+        for d, r, s in zip(DELTAS, rhos, sigs):
+            ax.annotate(f"Δ={d}", (r, s), xytext=(6, 6),
+                        textcoords="offset points", fontsize=9)
+        ax.axhline(gaussian_floor, color="red", lw=1, ls=":", alpha=0.6,
+                   label="Gaussian floor")
+        ax.set_yscale("log")
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel(ylabel)
+        ax.set_title(title)
+        ax.grid(alpha=0.3, which="both")
+        ax.legend(loc="lower left", fontsize=8)
+        return sc
+
+    # per-dim panels
+    for k in range(2):
+        rhos = np.array([trans[d]["rho"][k] for d in DELTAS])
+        sigs = np.array([trans[d]["sig"][f"marg_{k}_zs"][0] for d in DELTAS])
+        errs = np.array([trans[d]["sig"][f"marg_{k}_zs"][1] for d in DELTAS])
+        r2s = np.array([r2_dict[d][f"r2_dim{k}"] for d in DELTAS])
+        sc = _one(axes[k], rhos, sigs, errs, r2s,
+                  xlabel=r"auto-correlation $\rho$",
+                  ylabel="SIGReg (zscored, marginal)",
+                  title=names[k])
+        plt.colorbar(sc, ax=axes[k], label=r"$R^2$")
+
+    # joint panel
+    rhos_avg = np.array([np.mean(trans[d]["rho"]) for d in DELTAS])
+    sigs_j = np.array([trans[d]["sig"]["joint_zs"][0] for d in DELTAS])
+    errs_j = np.array([trans[d]["sig"]["joint_zs"][1] for d in DELTAS])
+    r2s_avg = np.array([r2_dict[d]["r2"] for d in DELTAS])
+    sc = _one(axes[2], rhos_avg, sigs_j, errs_j, r2s_avg,
+              xlabel=r"avg auto-correlation $\bar{\rho}$",
+              ylabel="SIGReg (zscored, 2D joint)",
+              title=names[2])
+    plt.colorbar(sc, ax=axes[2], label=r"$R^2$ (avg)")
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Measurement pipeline
+# ═════════════════════════════════════════════════════════════════════════════
+
+def compute_all_transitions(episodes, n_draws=20):
+    """
+    For each delta, compute SIGReg stats and per-dim rho on the
+    transition-difference distribution z(t+delta) - z(t).
+    """
+    trans = {}
+    for d in DELTAS:
+        diffs = (episodes[:, d:] - episodes[:, :-d]).reshape(-1, 2)
+        rho = np.array([
+            pearsonr(episodes[:, d:, k].flatten(),
+                     episodes[:, :-d, k].flatten())[0]
+            for k in range(2)
+        ])
+        trans[d] = {"rho": rho, "sig": measure(diffs, n_draws=n_draws)}
+    return trans
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Main
+# ═════════════════════════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--results_dir", type=str, default="results/reacher",
+                        help="Directory with result.json files")
+    parser.add_argument("--data_path", type=str,
+                        default="data/reacher.h5",
+                        help="HDF5 file with 'qpos' and 'ep_len' datasets "
+                        "(reshaped to (n_episodes, T, 2))")
+    parser.add_argument("--out_dir", type=str, default="figures/reacher")
+    parser.add_argument("--n_draws", type=int, default=20,
+                        help="Random subsamples for SIGReg stats")
+    parser.add_argument("--agg", type=str, default="median",
+                        choices=["median", "mean"],
+                        help="How to aggregate R² across seeds")
+    args = parser.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # load episodes (HDF5 with qpos + ep_len, matching notebook convention)
+    import h5py
+    with h5py.File(args.data_path, "r") as f:
+        qpos = np.array(f["qpos"])
+        ep_len = np.array(f["ep_len"])
+    T = int(ep_len[0])
+    episodes = qpos.reshape(-1, T, 2)
+    print(f"Loaded {len(episodes)} episodes of length {T} from {args.data_path}")
+
+    r2_dict = get_best_r2_per_delta(args.results_dir, agg=args.agg)
+    print(f"Loaded R² for deltas: {sorted(r2_dict.keys())}")
+    for d in DELTAS:
+        v = r2_dict[d]
+        print(f"  Δ={d:2d}  λ={v['lamb']:.0e}  n={v['n_seeds']}  "
+              f"R²={v['r2']:.3f}  dim0={v['r2_dim0']:.3f}  "
+              f"dim1={v['r2_dim1']:.3f}")
+
+    ou_rho = best_ou_rho(args.results_dir)
+    if ou_rho is not None:
+        print(f"Best OU rho: {ou_rho}")
+
+    # measure transitions
+    print("Computing SIGReg on transitions...")
+    trans = compute_all_transitions(episodes, n_draws=args.n_draws)
+
+    # figures
+    make_scatter_grid(episodes, r2_dict,
+                      save_path=out_dir / "distribution.png")
+    make_rho_vs_sigreg(trans, r2_dict,
+                       save_path=out_dir / "rho_vs_sigreg.png",
+                       best_ou_rho_val=ou_rho)
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/make_table_ablation.py

@@ -0,0 +1,98 @@
+"""
+Ablation table: 4 mixings × 3 source distributions × {SIGReg, Whitening}.
+Reports R²(h→z). The winner per (mixing, α) is bolded when the difference
+is statistically significant (Welch's t-test, p < 0.05).
+
+Usage:
+    python analysis/make_table_ablation.py --out figures/tab_ablation.tex
+"""
+import argparse, glob, json, os
+import numpy as np
+from scipy import stats
+from collections import defaultdict
+
+MIXINGS = [("spiral", "Spiral"), ("banana", "Banana"),
+           ("sinusoid", "Sinusoid"), ("nvp", "NVP")]
+ALPHAS = [
+    (None, r"Gaussian ($\alpha = 2$)"),
+    (0.25, r"Heavy tail / sparse ($\alpha = 1/4$)"),
+    (16.0, r"Light tail / uniform ($\alpha = 16$)"),
+]
+# METRIC = "r2_hz" # marginal
+METRIC = "r2_hz_grid" # grid
+P_THRESH = 0.05
+
+
+def alpha_key(r):
+    if r.get("source_dist") == "gennorm":
+        return r.get("source_alpha")
+    return None
+
+
+def fmt(vals, bold=False):
+    if not vals:
+        return r"$-$"
+    s = f"{np.mean(vals):.3f} \\pm {np.std(vals):.3f}"
+    return rf"$\mathbf{{{s}}}$" if bold else f"${s}$"
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dirs", nargs="+",
+                   default=["results/2d",
+                            "results/ablation_alpha_0.25",
+                            "results/ablation_alpha_16"])
+    p.add_argument("--out", default="figures/tab_ablation.tex")
+    args = p.parse_args()
+
+    groups = defaultdict(list)
+    for d in args.dirs:
+        for path in sorted(glob.glob(os.path.join(d, "*.json"))):
+            with open(path) as f:
+                r = json.load(f)
+            if METRIC not in r:
+                continue
+            key = (r["mixing"], alpha_key(r), r["mode"])
+            groups[key].append(r[METRIC])
+
+    col_spec = "l" + " cc" * len(ALPHAS)
+    multicol = " & ".join(rf"\multicolumn{{2}}{{c}}{{{lab}}}" for _, lab in ALPHAS)
+    cmidrules = "".join(rf"\cmidrule(lr){{{2*i+2}-{2*i+3}}}" for i in range(len(ALPHAS)))
+    method_hdr = " & ".join(["SIGReg & Whitening"] * len(ALPHAS))
+
+    lines = [
+        rf"\begin{{tabular}}{{{col_spec}}}",
+        r"\toprule",
+        rf" & {multicol} \\",
+        cmidrules,
+        rf"Mixing & {method_hdr} \\",
+        r"\midrule",
+    ]
+    for mix_key, mix_name in MIXINGS:
+        row = [mix_name]
+        for alpha_k, _ in ALPHAS:
+            v_lej = groups.get((mix_key, alpha_k, "lejepa"), [])
+            v_wht = groups.get((mix_key, alpha_k, "whiten"), [])
+            bold_lej = bold_wht = False
+            if len(v_lej) >= 2 and len(v_wht) >= 2:
+                _, pval = stats.ttest_ind(v_lej, v_wht, equal_var=False)
+                if pval < P_THRESH:
+                    if np.mean(v_lej) > np.mean(v_wht):
+                        bold_lej = True
+                    else:
+                        bold_wht = True
+            row.append(fmt(v_lej, bold=bold_lej))
+            row.append(fmt(v_wht, bold=bold_wht))
+        lines.append(" & ".join(row) + r" \\")
+    lines += [r"\bottomrule", r"\end{tabular}"]
+
+    out = "\n".join(lines)
+    print(out)
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    with open(args.out, "w") as f:
+        f.write(out + "\n")
+    print(f"\nSaved {args.out}")
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/make_table_reacher.py

@@ -0,0 +1,236 @@
+"""
+Generate LaTeX tables for the paper, matching scaling table style.
+
+Usage:
+    python analysis/make_table_reacher.py --results_dir results/reacher
+"""
+
+import json
+import argparse
+import numpy as np
+from pathlib import Path
+from collections import defaultdict
+
+
+def load_all_results(results_dir):
+    ou, traj = [], []
+    for p in Path(results_dir).rglob("result.json"):
+        r = json.load(open(p))
+        if "delta" in r and r.get("rho") is None:
+            traj.append(r)
+        elif "rho" in r:
+            ou.append(r)
+    return ou, traj
+
+
+def best_lambda_per_x(results, x_key):
+    grouped = defaultdict(list)
+    for r in results:
+        grouped[(r[x_key], r["lamb"])].append(r)
+
+    best = {}
+    for x in sorted(set(k[0] for k in grouped)):
+        best_mean, best_lamb = -np.inf, None
+        for lamb in set(k[1] for k in grouped if k[0] == x):
+            m = np.mean([r["r2_hz"] for r in grouped[(x, lamb)]])
+            if m > best_mean:
+                best_mean, best_lamb = m, lamb
+        best[x] = {
+            "lamb": best_lamb,
+            "runs": grouped[(x, best_lamb)],
+        }
+    return best
+
+
+def pm(vals, fmt=".2f"):
+    """Format as value\\tiny{±std} matching paper style."""
+    m, s = np.mean(vals), np.std(vals)
+    return f"{m:{fmt}}\\tiny{{$\\pm${s:.0e}}}"
+
+
+def make_combined_table(ou_results, traj_results):
+    ou_best = best_lambda_per_x(ou_results, "rho")
+    traj_best = best_lambda_per_x(traj_results, "delta")
+
+    lines = []
+    lines.append(r"\begin{table}[t]")
+    lines.append(r"\centering")
+    lines.append(r"\caption{")
+    lines.append(r"    \textbf{Pixel-observation identifiability on DMC Reacher} "
+                 r"(mean $\pm$ std over 3 seeds, best $\lambda$ per condition).")
+    lines.append(r"    \textbf{Left:} OU process with Gaussian marginals. "
+                 r"$R^2$ increases monotonically with $\rho$, reaching $0.95$ "
+                 r"at $\rho = 0.99$, confirming linear identifiability from pixels.")
+    lines.append(r"    \textbf{Right:} Real SAC trajectories with non-Gaussian marginals. "
+                 r"The two joints have different autocorrelation timescales ($\rho_0 \neq \rho_1$) "
+                 r"and the wrist has a near-uniform marginal distribution, "
+                 r"leading to anisotropic and reduced identifiability.")
+    lines.append(r"}")
+    lines.append(r"\label{tab:reacher}")
+    lines.append(r"\resizebox{\textwidth}{!}{%")
+    lines.append(r"\begin{tabular}{r cc | r cc ccc}")
+    lines.append(r" \multicolumn{3}{c}{\textbf{OU (Gaussian)}} & "
+                 r"\multicolumn{6}{c}{\textbf{Trajectory (non-Gaussian)}} \\")
+    lines.append(r"\cmidrule(lr){1-3} \cmidrule(lr){4-9}")
+    lines.append(r"$\rho$ & $R^2(z \to h)$ & $R^2(h \to z)$ & "
+                 r"$\delta$ & $\rho_0$ & $\rho_1$ & "
+                 r"$R^2(z \to h)$ & $R^2(h \to z_0)$ & $R^2(h \to z_1)$ \\")
+    lines.append(r"\midrule")
+
+    ou_rhos = sorted(ou_best.keys())
+    traj_deltas = sorted(traj_best.keys())
+    n_rows = max(len(ou_rhos), len(traj_deltas))
+
+    for i in range(n_rows):
+        # OU columns
+        if i < len(ou_rhos):
+            rho = ou_rhos[i]
+            runs = ou_best[rho]["runs"]
+            r2_zh = pm([r["r2_zh"] for r in runs])
+            r2_hz = pm([r["r2_hz"] for r in runs])
+            ou_str = f"  {rho:.2f} & {r2_zh} & {r2_hz}"
+        else:
+            ou_str = r"  & &"
+
+        # Traj columns
+        if i < len(traj_deltas):
+            delta = traj_deltas[i]
+            runs = traj_best[delta]["runs"]
+            rho0 = runs[0].get("rho_shoulder", None)
+            rho1 = runs[0].get("rho_wrist", None)
+            rho0_s = f"{rho0:.3f}" if rho0 is not None else "---"
+            rho1_s = f"{rho1:.3f}" if rho1 is not None else "---"
+            r2_zh = pm([r["r2_zh"] for r in runs])
+            r2_d0 = pm([r["r2_hz_per_dim"][0] for r in runs])
+            r2_d1 = pm([r["r2_hz_per_dim"][1] for r in runs])
+            traj_str = f"{delta} & {rho0_s} & {rho1_s} & {r2_zh} & {r2_d0} & {r2_d1}"
+        else:
+            traj_str = r"& & & & &"
+
+        lines.append(f"{ou_str} & {traj_str} \\\\")
+
+    lines.append(r"\bottomrule")
+    lines.append(r"\end{tabular}}")
+    lines.append(r"\vspace{5pt}")
+    lines.append(r"\vspace{-20pt}")
+    lines.append(r"\end{table}")
+    return "\n".join(lines)
+
+
+def make_ou_table_standalone(ou_results):
+    """Standalone OU table for appendix if needed."""
+    ou_best = best_lambda_per_x(ou_results, "rho")
+
+    lines = []
+    lines.append(r"\begin{table}[t]")
+    lines.append(r"\centering")
+    lines.append(r"\begin{tabular}{r c cc}")
+    lines.append(r"\toprule")
+    lines.append(r" \multicolumn{1}{c}{\textbf{Correlation}} & "
+                 r"\multicolumn{1}{c}{\textbf{Regularizer}} & "
+                 r"\multicolumn{2}{c}{\textbf{Linear identifiability}} \\")
+    lines.append(r"\cmidrule(lr){1-1} \cmidrule(lr){2-2} \cmidrule(lr){3-4}")
+    lines.append(r"$\rho$ & $\lambda$ & $R^2(z \to h)$ & $R^2(h \to z)$ \\")
+    lines.append(r"\midrule")
+
+    for rho in sorted(ou_best.keys()):
+        runs = ou_best[rho]["runs"]
+        lamb = ou_best[rho]["lamb"]
+        r2_zh = pm([r["r2_zh"] for r in runs])
+        r2_hz = pm([r["r2_hz"] for r in runs])
+        lines.append(f"  {rho:.2f} & {lamb:.0e} & {r2_zh} & {r2_hz} \\\\")
+
+    lines.append(r"\bottomrule")
+    lines.append(r"\end{tabular}")
+    lines.append(r"\vspace{5pt}")
+    lines.append(r"\caption{")
+    lines.append(r"    \textbf{OU (Gaussian) identifiability from pixels} "
+                 r"(mean $\pm$ std over 3 seeds).")
+    lines.append(r"    $R^2$ increases monotonically with temporal correlation $\rho$, "
+                 r"reaching $0.95$ at $\rho = 0.99$.")
+    lines.append(r"}")
+    lines.append(r"\label{tab:reacher_ou}")
+    lines.append(r"\end{table}")
+    return "\n".join(lines)
+
+
+def make_traj_table_standalone(traj_results):
+    """Standalone traj table for appendix if needed."""
+    traj_best = best_lambda_per_x(traj_results, "delta")
+
+    lines = []
+    lines.append(r"\begin{table}[t]")
+    lines.append(r"\centering")
+    lines.append(r"\resizebox{\textwidth}{!}{%")
+    lines.append(r"\begin{tabular}{r cc c c cc}")
+    lines.append(r"\toprule")
+    lines.append(r" \multicolumn{1}{c}{\textbf{Stride}} & "
+                 r"\multicolumn{2}{c}{\textbf{Autocorrelation}} & "
+                 r"\multicolumn{1}{c}{\textbf{Regularizer}} & "
+                 r"\multicolumn{1}{c}{\textbf{Identifiability}} & "
+                 r"\multicolumn{2}{c}{\textbf{Per-dimension}} \\")
+    lines.append(r"\cmidrule(lr){1-1} \cmidrule(lr){2-3} \cmidrule(lr){4-4} "
+                 r"\cmidrule(lr){5-5} \cmidrule(lr){6-7}")
+    lines.append(r"$\delta$ & $\rho_0$ & $\rho_1$ & $\lambda$ & "
+                 r"$R^2(z \to h)$ & $R^2(h \to z_0)$ & $R^2(h \to z_1)$ \\")
+    lines.append(r"\midrule")
+
+    for delta in sorted(traj_best.keys()):
+        runs = traj_best[delta]["runs"]
+        lamb = traj_best[delta]["lamb"]
+        rho0 = runs[0].get("rho_shoulder", None)
+        rho1 = runs[0].get("rho_wrist", None)
+        rho0_s = f"{rho0:.3f}" if rho0 is not None else "---"
+        rho1_s = f"{rho1:.3f}" if rho1 is not None else "---"
+        r2_zh = pm([r["r2_zh"] for r in runs])
+        r2_d0 = pm([r["r2_hz_per_dim"][0] for r in runs])
+        r2_d1 = pm([r["r2_hz_per_dim"][1] for r in runs])
+        lines.append(f"  {delta} & {rho0_s} & {rho1_s} & {lamb:.0e} "
+                     f"& {r2_zh} & {r2_d0} & {r2_d1} \\\\")
+
+    lines.append(r"\bottomrule")
+    lines.append(r"\end{tabular}}")
+    lines.append(r"\vspace{5pt}")
+    lines.append(r"\caption{")
+    lines.append(r"    \textbf{Trajectory (non-Gaussian) identifiability from pixels} "
+                 r"(mean $\pm$ std over 3 seeds).")
+    lines.append(r"    The shoulder ($z_0$) and wrist ($z_1$) have different "
+                 r"autocorrelation timescales and marginal distributions, "
+                 r"leading to anisotropic identifiability.")
+    lines.append(r"}")
+    lines.append(r"\label{tab:reacher_traj}")
+    lines.append(r"\end{table}")
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--results_dir", type=str, default="results/reacher")
+    args = parser.parse_args()
+
+    ou, traj = load_all_results(args.results_dir)
+    print(f"Loaded {len(ou)} OU runs, {len(traj)} traj runs\n")
+
+    if ou and traj:
+        print("=" * 70)
+        print("COMBINED TABLE (for main text)")
+        print("=" * 70)
+        print(make_combined_table(ou, traj))
+        print()
+
+    if ou:
+        print("=" * 70)
+        print("OU TABLE (standalone, for appendix)")
+        print("=" * 70)
+        print(make_ou_table_standalone(ou))
+        print()
+
+    if traj:
+        print("=" * 70)
+        print("TRAJ TABLE (standalone, for appendix)")
+        print("=" * 70)
+        print(make_traj_table_standalone(traj))
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/make_table_scaling.py

@@ -0,0 +1,272 @@
+"""
+Generate LaTeX tables from scaling results.
+
+Emits four tables:
+  1. tab:scaling-comparison  (main text)
+       Three-way R^2(h -> z) comparison: SIGReg, VICReg, InfoNCE.
+
+  2. tab:scaling-sigreg      (appendix)
+       Detailed per-method table for SIGReg: mixing difficulty,
+       linear identifiability (both directions), orthogonality error,
+       alignment loss, SIGReg loss.
+
+  3. tab:scaling-vicreg      (appendix)
+       Same structure as SIGReg, but with whitening loss column.
+
+  4. tab:scaling-infonce     (appendix)
+       Same structure as SIGReg, but with InfoNCE loss column.
+
+The three appendix tables let each method tell its own failure-mode story:
+  - SIGReg / VICReg: orthogonality error grows gradually with N
+  - InfoNCE: regularizer loss explodes / fails to converge at high N
+
+Usage:
+    python analysis/make_table_scaling.py --results_dir results/scaling/
+"""
+
+import argparse, glob, json, os
+import numpy as np
+from collections import defaultdict
+import math
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Number formatting helpers (shared)
+# ──────────────────────────────────────────────────────────────────────────
+def column_scale(stds, threshold=0.05):
+    valid = [s for s in stds if s > 0 and not np.isnan(s)]
+    if not valid:
+        return 0
+    m = max(valid)
+    if m >= threshold:
+        return 0
+    return int(np.floor(np.log10(m))) - 1
+
+
+def fmt_std(s, k):
+    if np.isnan(s):
+        return "---"
+    scaled = s if k == 0 else s / (10 ** k)
+    if scaled >= 10:
+        return f"{math.floor(scaled):.0f}"
+    elif scaled >= 1:
+        return f"{math.floor(scaled * 10) / 10:.1f}"
+    else:
+        return f"{math.floor(scaled * 100) / 100:.2f}"
+
+
+def scale_header(k):
+    if k == 0:
+        return r"{\scriptsize $\pm$std}"
+    return rf"{{\scriptsize $\pm$std\,$\times 10^{{{k}}}$}}"
+
+
+def fmt_cell(m, s, dec, k):
+    if np.isnan(m):
+        return "---"
+    factor = 10 ** dec
+    m_floored = math.floor(m * factor) / factor
+    return rf"{m_floored:.{dec}f}\tiny{{$\pm${fmt_std(s, k)}}}"
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Main-text: three-way comparison on R^2(h -> z)
+# ──────────────────────────────────────────────────────────────────────────
+def render_table_comparison(by_key, dims):
+    modes = ("lejepa", "whiten", "infonce")
+
+    def agg(N, mode, key):
+        rs = by_key.get((N, mode), [])
+        vals = [r[key] for r in rs if r.get(key) is not None]
+        if not vals:
+            return float("nan"), float("nan")
+        return float(np.mean(vals)), float(np.std(vals))
+
+    def agg_mixing(N, key):
+        all_rs = sum((by_key.get((N, m), []) for m in modes), [])
+        vals = [r[key] for r in all_rs if r.get(key) is not None]
+        if not vals:
+            return float("nan"), float("nan")
+        return float(np.mean(vals)), float(np.std(vals))
+
+    mix_scale  = column_scale([agg_mixing(N, "r2_xz")[1] for N in dims])
+    r2_scales  = {m: column_scale([agg(N, m, "r2_hz")[1] for N in dims]) for m in modes}
+
+    print(r"\begin{table}[t]")
+    print(r"\centering")
+    print(r"\caption{\textbf{Scaling Comparison Across Regularizers} (mean $\pm$ std, 5 seeds). "
+          r"All three Gaussianity-enforcing objectives are tested on the same RealNVP mixing "
+          r"with matched encoder. SIGReg and VICReg (batch-statistic estimators) maintain "
+          r"$R^2 > 0.999$ up to $N{=}1024$, consistent with Thm.~\ref{thm:approx}. "
+          r"InfoNCE (pair-based) matches at low $N$ but degrades at scale under fixed kernel "
+          r"width $\sigma{=}1$, illustrating the per-dimension tuning required by pair-based estimators. "
+          r"Per-method details (orthogonality, regularizer loss) in App.~\ref{app:scaling}, "
+          r"Tabs.~\ref{tab:scaling-sigreg}--\ref{tab:scaling-infonce}.}")
+    print(r"\label{tab:scaling-comparison}")
+    print(r"\begin{tabular}{r c ccc}")
+    print(r"\toprule")
+    print(r"  & \textbf{Mixing} & \multicolumn{3}{c}{\textbf{Linear identifiability} $R^2(h \to z)$} \\")
+    print(r"\cmidrule(lr){3-5}")
+    print(r"$N$ & $R^2(x \to z)$ & SIGReg & VICReg & InfoNCE \\")
+    sub_cells = [
+        scale_header(mix_scale),
+        scale_header(r2_scales["lejepa"]),
+        scale_header(r2_scales["whiten"]),
+        scale_header(r2_scales["infonce"]),
+    ]
+    print(" & " + " & ".join(sub_cells) + r" \\")
+    print(r"\midrule")
+
+    for N in dims:
+        cells = [rf"{N}"]
+        m, s = agg_mixing(N, "r2_xz")
+        cells.append(fmt_cell(m, s, 3, mix_scale))
+        for mode in modes:
+            m_, s_ = agg(N, mode, "r2_hz")
+            cells.append(fmt_cell(m_, s_, 6, r2_scales[mode]))
+        print("  " + " & ".join(cells) + r" \\")
+
+    print(r"\bottomrule")
+    print(r"\end{tabular}")
+    print(r"\end{table}")
+    print()
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Appendix: detailed per-method table
+# ──────────────────────────────────────────────────────────────────────────
+
+# Each method has its own native regularizer loss key.
+METHOD_SPECS = {
+    "lejepa": {
+        "name":            "SIGReg",
+        "label":           "tab:scaling-sigreg",
+        "reg_loss_key":    "final_sigreg",
+        "reg_loss_label":  "SIGReg",
+        "reg_loss_dec":    2,
+        "caption_tail":    (
+            "The RealNVP mixing is consistently nonlinear across dimensions "
+            r"($R^2(x \to z) < 1$). The learned model nonetheless recovers the "
+            "true latents at all dimensions. Training losses are stable; "
+            r"orthogonality error grows gradually with $N$."
+        ),
+    },
+    "whiten": {
+        "name":            "VICReg",
+        "label":           "tab:scaling-vicreg",
+        "reg_loss_key":    "final_whiten",
+        "reg_loss_label":  "Whitening",
+        "reg_loss_dec":    4,
+        "caption_tail":    (
+            "VICReg matches SIGReg on linear identifiability across all dimensions; "
+            r"orthogonality error grows similarly with $N$. The whitening loss is "
+            "stable and small throughout."
+        ),
+    },
+    "infonce": {
+        "name":            "InfoNCE",
+        "label":           "tab:scaling-infonce",
+        "reg_loss_key":    "final_loss",
+        "reg_loss_label":  "InfoNCE",
+        "reg_loss_dec":    3,
+        "caption_tail":    (
+            r"InfoNCE matches the batch-statistic methods at low $N$ but degrades "
+            "at scale under a fixed Gaussian kernel width. The InfoNCE column "
+            "shows the total contrastive loss (not decomposable into alignment "
+            "plus regularizer), which inflates with $N$ as the per-dimension "
+            "kernel-width assumption breaks down."
+        ),
+    },
+}
+
+
+def render_table_per_method(by_key, dims, mode):
+    spec = METHOD_SPECS[mode]
+    rows = sum((by_key.get((N, mode), []) for N in dims), [])
+    if not rows:
+        print(f"% No rows for mode={mode}, skipping {spec['label']}")
+        return
+
+    cols = [
+        ("r2_zx",                3, r"$R^2(z \to x)$"),
+        ("r2_xz",                3, r"$R^2(x \to z)$"),
+        ("r2_zh",                5, r"$R^2(z \to h)$"),
+        ("r2_hz",                5, r"$R^2(h \to z)$"),
+        ("orth_err_normalized",  3, r"$\|\hat Q^\top \hat Q - I\|_F / \sqrt{N}$"),
+        ("final_align",          4, "Align"),
+        (spec["reg_loss_key"],   spec["reg_loss_dec"], spec["reg_loss_label"]),
+    ]
+
+    agg = {}
+    for N in dims:
+        rs = [r for r in rows if r["N"] == N]
+        for key, _, _ in cols:
+            vals = [r[key] for r in rs if r.get(key) is not None]
+            agg[(N, key)] = ((float(np.mean(vals)), float(np.std(vals)))
+                             if vals else (float("nan"), float("nan")))
+
+    scales = {key: column_scale([agg[(N, key)][1] for N in dims])
+              for key, _, _ in cols}
+
+    print(r"\begin{table}[t]")
+    print(r"\centering")
+    print(rf"\caption{{\textbf{{Scaling Experiment ({spec['name']})}} "
+          r"(mean $\pm$ std, 5 seeds). " + spec["caption_tail"] + "}")
+    print(rf"\label{{{spec['label']}}}")
+    print(r"\resizebox{\textwidth}{!}{%")
+    print(r"\begin{tabular}{r cc cc c cc}")
+    print(r"\toprule")
+    print(r" \multicolumn{1}{c}{\textbf{Latents}} "
+          r"& \multicolumn{2}{c}{\textbf{Mixing difficulty}} "
+          r"& \multicolumn{2}{c}{\textbf{Linear identifiability}} "
+          r"& \multicolumn{1}{c}{\textbf{Orthogonality}} "
+          rf"& \multicolumn{{2}}{{c}}{{\textbf{{{spec['name']} losses}}}} \\")
+    print(r"\cmidrule(lr){1-1} \cmidrule(lr){2-3} \cmidrule(lr){4-5} "
+          r"\cmidrule(lr){6-6} \cmidrule(lr){7-8}")
+    print("$N$ & " + " & ".join(label for _, _, label in cols) + r" \\")
+    print(" & " + " & ".join(scale_header(scales[key]) for key, _, _ in cols) + r" \\")
+    print(r"\midrule")
+
+    for N in dims:
+        log2N = int(np.log2(N))
+        cells = [rf"$2^{{{log2N}}}$"]
+        for key, dec, _ in cols:
+            m, s = agg[(N, key)]
+            cells.append(fmt_cell(m, s, dec, scales[key]))
+        print("  " + " & ".join(cells) + r" \\")
+
+    print(r"\bottomrule")
+    print(r"\end{tabular}}")
+    print(r"\end{table}")
+    print()
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Driver
+# ──────────────────────────────────────────────────────────────────────────
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", default="results/scaling/")
+    args = p.parse_args()
+
+    by_key = defaultdict(list)
+    for path in sorted(glob.glob(os.path.join(args.results_dir, "*.json"))):
+        with open(path) as f:
+            r = json.load(f)
+        mode = r.get("mode", "lejepa")
+        by_key[(r["N"], mode)].append(r)
+
+    dims = sorted({N for (N, _) in by_key})
+    modes_present = sorted({mode for (_, mode) in by_key})
+    print(f"% Found modes: {modes_present}, dims: {dims}\n")
+
+    print(r"% ── Main text: three-way comparison ──")
+    render_table_comparison(by_key, dims)
+
+    print(r"% ── Appendix: detailed per-method tables ──")
+    for mode in ("lejepa", "whiten", "infonce"):
+        render_table_per_method(by_key, dims, mode)
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_2d.py

@@ -0,0 +1,82 @@
+"""
+2D illustration figure: one row per mixing (z, g(z), h(z)).
+Picks best seed per mixing by final_loss.
+
+Usage:
+    python analysis/plot_2d.py --results_dir results/2d/ --out figures/
+"""
+
+import argparse, os, glob
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import colorsys
+
+MIXING_ORDER = ["spiral", "banana", "sinusoid", "nvp"]
+
+
+def make_colors(z):
+    x, y = z[:, 0], z[:, 1]
+    angles = np.arctan2(y, x)
+    radii = np.sqrt(x**2 + y**2)
+    hue = (angles + np.pi) / (2 * np.pi)
+    lightness = 0.3 + 0.4 * (radii / (radii.max() + 1e-8))
+    saturation = np.full_like(hue, 0.85)
+    return [colorsys.hls_to_rgb(h, l, s) for h, l, s in zip(hue, lightness, saturation)]
+
+
+def load_best_per_mixing(results_dir):
+    """Load all results, pick best lejepa run per mixing by final_loss."""
+    files = sorted(glob.glob(os.path.join(results_dir, "*.pt")))
+    by_mix = {}
+    for path in files:
+        r = torch.load(path, map_location="cpu", weights_only=False)
+        if r.get("mode") != "lejepa":
+            continue
+        mix = r["mixing"]
+        if mix not in by_mix or r["final_loss"] < by_mix[mix]["final_loss"]:
+            by_mix[mix] = r
+    return by_mix
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", default="results/2d/")
+    p.add_argument("--out", default="figures/")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    best = load_best_per_mixing(args.results_dir)
+
+    s, lim = 5, 4
+    for mix_name in MIXING_ORDER:
+        if mix_name not in best:
+            print(f"Missing {mix_name}"); continue
+        res = best[mix_name]
+        z, x, h = res["z"], res["x"], res["h"]
+        colors = make_colors(z)
+
+        fig, axes = plt.subplots(1, 3, figsize=(9, 3))
+        for i, (ax, data, labels) in enumerate(zip(
+            axes, [z, x, h],
+            [("True Latent 0", "True Latent 1"),
+             ("Observation 0", "Observation 1"),
+             ("Learned Latent 0", "Learned Latent 1")],
+        )):
+            ax.scatter(data[:, 0], data[:, 1], c=colors, s=s, linewidths=0)
+            ax.set_xlabel(labels[0])
+            ax.set_ylabel(labels[1])
+            ax.grid(alpha=0.3)
+            if i == 0 or i == 2:
+                ax.set_xlim(-4, 4)
+                ax.set_ylim(-4, 4)
+
+        fig.tight_layout()
+        out_path = os.path.join(args.out, f"fig_2d_{mix_name}.jpg")
+        fig.savefig(out_path, bbox_inches="tight", dpi=500)
+        print(f"Saved {out_path}")
+        plt.close()
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_ablation.py

@@ -0,0 +1,91 @@
+"""
+Laplace ablation: 4 rows (mixings) x 4 cols (z, g(z), h_lejepa, h_whiten).
+Picks best seed per (mixing, mode) by final_loss.
+
+Usage:
+    python analysis/plot_ablation.py --results_dir results/ablation_alpha_0.25/ --prefix ablation_alpha_0.25 --out figures/
+    python analysis/plot_ablation.py --results_dir results/ablation_alpha_16/   --prefix ablation_alpha_16   --out figures/
+"""
+
+import argparse, os, glob
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import colorsys
+
+MIXING_ORDER = ["spiral", "banana", "sinusoid", "nvp"]
+
+
+def make_colors(z):
+    x, y = z[:, 0], z[:, 1]
+    angles = np.arctan2(y, x)
+    radii = np.sqrt(x**2 + y**2)
+    hue = (angles + np.pi) / (2 * np.pi)
+    lightness = 0.3 + 0.4 * (radii / (radii.max() + 1e-8))
+    saturation = np.full_like(hue, 0.85)
+    return [colorsys.hls_to_rgb(h, l, s) for h, l, s in zip(hue, lightness, saturation)]
+
+
+def load_best(results_dir):
+    files = sorted(glob.glob(os.path.join(results_dir, "*.pt")))
+    by_key = {}
+    for path in files:
+        r = torch.load(path, map_location="cpu", weights_only=False)
+        key = r["run_name"]  # e.g. "spiral_lejepa"
+        if key not in by_key or r["final_loss"] < by_key[key]["final_loss"]:
+            by_key[key] = r
+    return by_key        
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", default="results/ablation/")
+    p.add_argument("--out", default="figures/")
+    p.add_argument("--prefix", default="ablation",
+                   help="Output filename prefix (e.g. ablation_alpha_0.25)")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    best = load_best(args.results_dir)
+
+    s, lim = 5, 4
+    for mix in MIXING_ORDER:
+        lej_key = f"{mix}_lejepa"
+        wht_key = f"{mix}_whiten"
+        if lej_key not in best or wht_key not in best:
+            print(f"Missing {mix}"); continue
+
+        lej, wht = best[lej_key], best[wht_key]
+        z = lej["z"]
+        colors = make_colors(z)
+
+        fig, axes = plt.subplots(1, 4, figsize=(12, 3))
+        col_labels = [
+            ("True Latent 0", "True Latent 1"),
+            ("Observation 0", "Observation 1"),
+            ("Learned (LeJEPA) 0", "Learned (LeJEPA) 1"),
+            ("Learned (Whiten) 0", "Learned (Whiten) 1"),
+        ]
+        panels = [z, lej["x"], lej["h"], wht["h"]]
+        # r2s = [None, None, lej["r2_hz"], wht["r2_hz"]]
+        r2s = [None, None, lej["r2_hz_grid"], wht["r2_hz_grid"]]
+
+        for i, (ax, data, labels, r2) in enumerate(zip(axes, panels, col_labels, r2s)):
+            ax.scatter(data[:, 0], data[:, 1], c=colors, s=s, linewidths=0)
+            ax.set_xlabel(labels[0])
+            ax.set_ylabel(labels[1])
+            ax.grid(alpha=0.3)
+            if r2 is not None:
+                ax.text(0.95, 0.05, f"$R^2$={r2:.3f}", transform=ax.transAxes,
+                        ha="right", va="bottom", fontsize=9,
+                        bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.8))
+
+        fig.tight_layout()
+        out_path = os.path.join(args.out, f"fig_{args.prefix}_{mix}.jpg")
+        fig.savefig(out_path, bbox_inches="tight", dpi=500)
+        print(f"Saved {out_path}")
+        plt.close()
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_bound.py

@@ -0,0 +1,237 @@
+"""
+Bound verification and grid search plots.
+
+bound_verification.pdf  — pooled across experiments (main paper)
+bound_decomposition.pdf — grid search only (appendix)
+heatmap_*.pdf           — grid search only (appendix)
+
+Usage:
+    python analysis/plot_bound.py \
+    --results_dirs results/grid results/2d results/scaling results/gennorm \
+    --out figures/
+"""
+
+import argparse, os, glob, json
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+
+# mpl.rcParams.update({
+#     "font.size": 10, "axes.titlesize": 11, "axes.labelsize": 10,
+#     "figure.dpi": 200, "font.family": "serif",
+# })
+
+EXPERIMENT_COLORS = {
+    "2d":      "tab:blue",
+    "grid":    "tab:red",
+    "scaling": "tab:green",
+    "reacher": "tab:purple",
+    "gennorm": "tab:orange",
+}
+EXPERIMENT_MARKERS = {
+    "2d":      "o",
+    "grid":    "D",
+    "scaling": "s",
+    "reacher": "v",
+    "gennorm": "^",
+}
+EXPERIMENT_ORDER = ["grid", "scaling", "2d", "reacher", "gennorm"]
+
+
+def is_valid_run(r, path=""):
+    """SIGReg + Gaussian source + non-degenerate (encoder actually learned)."""
+    if r.get("mode") != "lejepa":
+        return False
+    sd = r.get("source_dist", "gaussian")
+    if sd == "gennorm" and abs(r.get("source_alpha", 0) - 2.0) > 1e-6:
+        return False
+    if sd not in ("gaussian", "gennorm"):
+        return False
+    # Drop degenerate runs where the encoder failed to learn
+    if r.get("r2_hz", 0) < 0.5:
+        return False
+    return True
+
+
+def load_all(dirs):
+    data = []
+    for d in dirs:
+        for path in sorted(glob.glob(os.path.join(d, "**", "*.json"), recursive=True)):
+            with open(path) as f:
+                r = json.load(f)
+            if not isinstance(r, dict):
+                continue
+            if r.get("approx_bound") is None or r.get("procrustes_mse") is None:
+                continue
+            if not is_valid_run(r, path):
+                continue
+            data.append(r)
+    print(f"Loaded {len(data)} Gaussian-source runs")
+    return data
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dirs", nargs="+", required=True)
+    p.add_argument("--out", default="figures/")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    data = load_all(args.results_dirs)
+    if not data:
+        print("No results."); return
+
+    # ══════════════════════════════════════════════════════════════════════
+    # Bound verification (main paper, single panel)
+    # ══════════════════════════════════════════════════════════════════════
+    errors = np.array([d["procrustes_mse"] for d in data])
+    bounds = np.array([d["approx_bound"] for d in data])
+    experiments = [d["experiment"] for d in data]
+
+    fig, ax = plt.subplots(figsize=0.8 * np.array((3, 3)))
+
+    pos = (errors > 0) & (bounds > 0)
+    lo = min(errors[pos].min(), bounds[pos].min()) * 0.3
+    hi = max(errors.max(), bounds.max()) * 3
+    pts = np.logspace(np.log10(lo), np.log10(hi), 200)
+    ax.fill_between(pts, lo * 0.1, pts, color='#c8e6c9', alpha=0.35, zorder=0)
+    ax.fill_between(pts, pts, hi * 10, color='#ffcdd2', alpha=0.35, zorder=0)
+    
+    for exp in EXPERIMENT_ORDER:
+        mask = np.array([e == exp for e in experiments])
+        if not mask.any():
+            continue
+        ax.scatter(bounds[mask], errors[mask],
+                   c=EXPERIMENT_COLORS[exp],
+                   marker=EXPERIMENT_MARKERS[exp],
+                   s=32, alpha=0.7, edgecolors='k', linewidths=0.3,
+                   label=exp, zorder=3)
+    
+    ax.plot([lo, hi], [lo, hi], 'k--', alpha=0.5, linewidth=0.8)
+    ax.set_xscale('log'); ax.set_yscale('log')
+    ax.set_xlim(lo, hi)
+    ax.set_ylim(lo, hi)
+    ax.set_xlabel("Recovery error bound")
+    ax.set_ylabel("Recovery error")
+    ax.legend(fontsize=7, loc='upper left', framealpha=0.9)
+    ax.grid(alpha=0.3)
+    ax.set_aspect("equal")
+
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.out, "bound_verification.pdf"), bbox_inches="tight")
+    print("Saved bound_verification.pdf")
+    plt.close()
+
+    # ══════════════════════════════════════════════════════════════════════
+    # Grid-specific plots (appendix)
+    # ══════════════════════════════════════════════════════════════════════
+    grid_dir = None
+    for d in args.results_dirs:
+        if "grid" in d:
+            grid_dir = d
+            break
+    if grid_dir is None:
+        print("No grid dir found, skipping decomposition and heatmaps.")
+        return
+
+    grid_data = []
+    for path in sorted(glob.glob(os.path.join(grid_dir, "*.json"))):
+        with open(path) as f:
+            grid_data.append(json.load(f))
+    if not grid_data:
+        print("No grid results."); return
+
+    errors_g = np.array([d["procrustes_mse"] for d in grid_data])
+    epsilons = np.array([d["epsilon"] for d in grid_data])
+    deltas = np.array([d["delta"] for d in grid_data])
+    rhos = np.array([d["rho"] for d in grid_data])
+    lambs = np.array([d["lamb"] for d in grid_data])
+
+    LAMB_MARKERS = {1e-6: 'h', 1e-5: 'H', 1e-4: 'p',
+                    1e-3: 'o', 5e-3: 's', 1e-2: 'D', 5e-2: '^', 1e-1: 'v', 5e-1: 'P'}
+    RHO_MARKERS = {0.3: 'o', 0.5: 's', 0.7: 'D', 0.8: '^', 0.9: 'v', 0.95: 'P', 0.99: 'X'}
+
+    def scatter_by_lamb(ax, xvals, yvals):
+        cmap = plt.cm.viridis
+        norm = mpl.colors.Normalize(vmin=min(rhos), vmax=max(rhos))
+        for lamb in sorted(set(lambs)):
+            mask = lambs == lamb
+            ax.scatter(xvals[mask], yvals[mask], c=rhos[mask], cmap=cmap, norm=norm,
+                       marker=LAMB_MARKERS.get(lamb, 'o'), s=30, alpha=0.8,
+                       edgecolors='k', linewidths=0.3, label=f"$\\lambda$={lamb:.0e}")
+        return cmap, norm, r"Correlation [$\rho$]"
+
+    def scatter_by_rho(ax, xvals, yvals):
+        log_lambs = np.log10(lambs)
+        cmap = plt.cm.plasma
+        norm = mpl.colors.Normalize(vmin=log_lambs.min(), vmax=log_lambs.max())
+        for rho_val in sorted(set(rhos)):
+            mask = rhos == rho_val
+            ax.scatter(xvals[mask], yvals[mask], c=log_lambs[mask], cmap=cmap, norm=norm,
+                       marker=RHO_MARKERS.get(rho_val, 'o'), s=30, alpha=0.8,
+                       edgecolors='k', linewidths=0.3, label=f"$\\rho$={rho_val:.2f}")
+        return cmap, norm, r"Regularization [$\log_{10}\lambda$]"
+
+    # ── Decomposition ──
+    ylabel = r"$\min_{Q \in O(n)} \mathbb{E}[\|h(z) - Qz\|^2]$"
+    fig, axes = plt.subplots(2, 2, figsize=(8, 6))
+    x_configs = [
+        (epsilons, r"$\varepsilon = \|\mathrm{Cov}(h(z)) - I\|_F$", r"Error vs $\varepsilon$"),
+        (deltas, r"$\delta = \mathcal{L}(h) - 2(1{-}\rho)\,\mathrm{tr}(\Sigma)$", r"Error vs $\delta$"),
+    ]
+    for row, scatter_fn in enumerate([scatter_by_lamb, scatter_by_rho]):
+        for col, (xvals, xlabel, title) in enumerate(x_configs):
+            ax = axes[row, col]
+            cmap, norm, cbar_label = scatter_fn(ax, xvals, errors_g)
+            ax.set_xlabel(xlabel); ax.set_ylabel(ylabel); ax.set_title(title)
+            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm); sm.set_array([])
+            plt.colorbar(sm, ax=ax, label=cbar_label, shrink=0.85)
+            ax.legend(fontsize=5.5, loc='upper right', framealpha=0.9)
+            ax.grid()
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.out, "bound_decomposition.pdf"), bbox_inches="tight")
+    print("Saved bound_decomposition.pdf")
+    plt.close()
+
+    # ── Heatmaps ──
+    unique_lambs = sorted(set(lambs))
+    unique_rhos = sorted(set(rhos))
+    for metric_key, title, cmap_name in [
+        ("r2_hz", "Linear $R^2$ (h -> z)", "viridis"),
+        ("orth_err_normalized", "Orth. error normalized", "viridis_r"),
+    ]:
+        grid = np.full((len(unique_lambs), len(unique_rhos)), np.nan)
+        counts = np.zeros_like(grid)
+        for r in grid_data:
+            li = unique_lambs.index(r["lamb"])
+            ri = unique_rhos.index(r["rho"])
+            val = r.get(metric_key)
+            if val is not None:
+                if np.isnan(grid[li, ri]):
+                    grid[li, ri] = 0
+                grid[li, ri] += val
+                counts[li, ri] += 1
+        grid = np.where(counts > 0, grid / counts, np.nan)
+
+        fig, ax = plt.subplots(figsize=(7, 5))
+        im = ax.imshow(grid, aspect="auto", origin="lower", cmap=cmap_name)
+        plt.colorbar(im, ax=ax, label=title)
+        ax.set_xticks(range(len(unique_rhos)))
+        ax.set_xticklabels([f"{r:.2f}" for r in unique_rhos])
+        ax.set_yticks(range(len(unique_lambs)))
+        ax.set_yticklabels([f"{l:.0e}" for l in unique_lambs])
+        ax.set_xlabel(r"$\rho$"); ax.set_ylabel(r"$\lambda$")
+        ax.set_title(title)
+        for i in range(len(unique_lambs)):
+            for j in range(len(unique_rhos)):
+                if not np.isnan(grid[i, j]):
+                    ax.text(j, i, f"{grid[i,j]:.3f}", ha="center", va="center", fontsize=6)
+        fig.tight_layout()
+        safe = metric_key.replace(".", "_")
+        fig.savefig(os.path.join(args.out, f"heatmap_{safe}.pdf"), bbox_inches="tight")
+        print(f"Saved heatmap_{safe}.pdf")
+        plt.close()
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_gennorm.py

@@ -0,0 +1,202 @@
+"""
+Generalized normal sweep across 4 mixings.
+
+Emits three figures:
+  1. fig_gennorm.pdf       — 4-panel R^2(h -> z) vs alpha for SIGReg/VICReg/InfoNCE
+  2. fig_gennorm_orth.pdf  — 4-panel orthogonality error vs alpha (unconstrained ylim
+                             to show InfoNCE excursions off the chart)
+  3. fig_gennorm_main.pdf  — single-panel spiral-only headline figure for main text,
+                             matching the Fig.~4b style of the paper
+
+Usage:
+    python analysis/plot_gennorm.py --results_dir results/gennorm/ --out figures/
+"""
+import argparse, glob, json, os, re
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import defaultdict
+
+MIXINGS = [("spiral", "Spiral"), ("banana", "Banana"),
+           ("sinusoid", "Sinusoid"), ("nvp", "NVP")]
+
+MODES = ("lejepa", "whiten", "infonce")
+COLORS = {"lejepa": "#d62728", "whiten": "#1f77b4", "infonce": "#2ca02c"}
+LABELS = {"lejepa": "SIGReg", "whiten": "VICReg", "infonce": "InfoNCE"}
+
+YLABELS = {
+    "r2_hz_grid":               r"Linear identifiability  $R^2(h \to z)$",
+    "r2_hz":                    r"Linear identifiability  $R^2(h \to z)$",
+    "orth_err_normalized_grid": r"Orthogonality error  $\|\hat Q^\top \hat Q - I\|_F / \sqrt{n}$",
+    "orth_err_normalized":      r"Orthogonality error  $\|\hat Q^\top \hat Q - I\|_F / \sqrt{n}$",
+}
+
+# Sensible y-limits per metric. R^2 is bounded in [0,1] so we clip there.
+# Orthogonality error is unbounded above (Whitening/InfoNCE off-Gaussian can spike
+# into the tens), so we use log scale and let matplotlib autoscale.
+YLIMS = {
+    "r2_hz_grid":               (-0.05, 1.05),
+    "r2_hz":                    (-0.05, 1.05),
+    "orth_err_normalized_grid": None,  # autoscale; log scale (see YSCALES) handles outliers
+    "orth_err_normalized":      None,
+}
+
+# Y-axis scale per metric. Linear by default; log for orth error to compress
+# off-Gaussian excursions while still showing structure near zero.
+YSCALES = {
+    "r2_hz_grid":               "linear",
+    "r2_hz":                    "linear",
+    "orth_err_normalized_grid": "log",
+    "orth_err_normalized":      "log",
+}
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Data loading
+# ──────────────────────────────────────────────────────────────────────────
+def load_groups(results_dir, metric):
+    """groups[(mixing, mode, alpha)] -> list of seed values for `metric`."""
+    groups = defaultdict(list)
+    for path in sorted(glob.glob(os.path.join(results_dir, "*.json"))):
+        with open(path) as f:
+            r = json.load(f)
+        if r.get("experiment") != "gennorm":
+            continue
+        alpha = r.get("source_alpha")
+        if alpha is None:
+            m = re.search(r"alpha=([\d.]+)", r.get("run_name", ""))
+            if m:
+                alpha = float(m.group(1))
+        if alpha is None or metric not in r:
+            continue
+        groups[(r["mixing"], r["mode"], alpha)].append(r[metric])
+    return groups
+
+
+def curve(groups, mixing, mode):
+    alphas = sorted({a for (mx, m, a) in groups if mx == mixing and m == mode})
+    mu = np.array([np.mean(groups[(mixing, mode, a)]) for a in alphas])
+    sd = np.array([np.std (groups[(mixing, mode, a)]) for a in alphas])
+    return np.array(alphas), mu, sd
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Figure 1 & 2: 4-panel grids (one per metric)
+# ──────────────────────────────────────────────────────────────────────────
+def plot_grid(groups, metric, out_path):
+    fig, axes = plt.subplots(1, 4, figsize=(13, 3.0), sharey=True)
+    use_log = YSCALES[metric] == "log"
+
+    for ax, (mix_key, mix_name) in zip(axes, MIXINGS):
+        for mode in MODES:
+            alphas, mu, sd = curve(groups, mix_key, mode)
+            if len(alphas) == 0:
+                continue
+            ax.plot(alphas, mu, marker="o", ms=5, lw=1.8,
+                    color=COLORS[mode], label=LABELS[mode], zorder=3)
+            # On log axes, clip the lower edge of the band away from zero
+            # so fill_between doesn't disappear / warn.
+            lower = mu - sd
+            if use_log:
+                lower = np.maximum(lower, 1e-3)
+            ax.fill_between(alphas, lower, mu + sd,
+                            color=COLORS[mode], alpha=0.2, zorder=2)
+        ax.set_xscale("log", base=2)
+        if use_log:
+            ax.set_yscale("log")
+        ax.axvline(2.0, color="black", lw=0.7, ls="--", alpha=0.6, zorder=1)
+        ax.set_xlabel(r"Source shape $\alpha$")
+        ax.set_title(mix_name)
+        if YLIMS[metric] is not None:
+            ax.set_ylim(*YLIMS[metric])
+        ax.grid(alpha=0.3, which="both" if use_log else "major")
+
+    axes[0].set_ylabel(YLABELS[metric])
+    axes[-1].legend(frameon=False, loc="best")
+
+    fig.tight_layout()
+    fig.savefig(out_path, bbox_inches="tight")
+    print(f"Saved {out_path}")
+    plt.close(fig)
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Figure 3: main-text single-panel headline (spiral mixing, all three methods)
+# ──────────────────────────────────────────────────────────────────────────
+def plot_main_panel(groups, out_path):
+    """Single-panel spiral-only figure to sit next to the bound-verification panel
+    in the main-text composite figure (Fig.~4b in the paper)."""
+    FIGSIZE  = 0.8 * np.array((4.0, 3.0))
+    LW       = 2.0
+    MS       = 6
+    FONTSIZE = 11
+
+    rc_saved = plt.rcParams.copy()
+    plt.rcParams.update({
+        "font.size":         FONTSIZE,
+        "axes.labelsize":    FONTSIZE,
+        "xtick.labelsize":   FONTSIZE - 1,
+        "ytick.labelsize":   FONTSIZE - 1,
+        "legend.fontsize":   FONTSIZE - 1,
+        "axes.spines.top":   False,
+        "axes.spines.right": False,
+    })
+
+    fig, ax = plt.subplots(figsize=FIGSIZE)
+
+    for mode in MODES:
+        alphas, mu, sd = curve(groups, "spiral", mode)
+        if len(alphas) == 0:
+            continue
+        ax.plot(alphas, mu, marker="o", ms=MS, lw=LW,
+                color=COLORS[mode], label=LABELS[mode], zorder=3)
+        ax.fill_between(alphas, mu - sd, mu + sd,
+                        color=COLORS[mode], alpha=0.2, zorder=2)
+
+    # Reference lines for canonical distributions
+    ax.axvline(1.0, 0, 0.95, color="gray",  lw=0.8, ls=":",  alpha=0.7, zorder=1)
+    ax.axvline(2.0, 0, 0.90, color="black", lw=0.8, ls="--", alpha=0.7, zorder=1)
+    ax.text(1.0, 1.04, "Laplace", ha="center", va="bottom",
+            fontsize=FONTSIZE - 1, color="gray")
+
+    ax.set_xscale("log", base=2)
+    ax.set_ylabel(r"Linearity")
+    ax.set_ylim(-0.05, 1.12)
+    ax.grid(alpha=0.3, which="both")
+    ax.legend(frameon=False, loc="lower right")
+    ax.set_xticks([2**(-2), 2, 16])
+    ax.set_xticklabels([r"$\leftarrow$ sparse", "Gaussian", r"uniform $\rightarrow$"])
+
+    fig.tight_layout()
+    fig.savefig(out_path, bbox_inches="tight", dpi=500)
+    print(f"Saved {out_path}")
+    plt.close(fig)
+    plt.rcParams.update(rc_saved)
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Driver
+# ──────────────────────────────────────────────────────────────────────────
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", default="results/gennorm/")
+    p.add_argument("--out",         default="figures/")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    # 4-panel R^2 grid (appendix)
+    groups_r2 = load_groups(args.results_dir, "r2_hz_grid")
+    plot_grid(groups_r2, "r2_hz_grid",
+              os.path.join(args.out, "fig_gennorm.pdf"))
+
+    # 4-panel orthogonality grid (appendix, autoscaled)
+    groups_orth = load_groups(args.results_dir, "orth_err_normalized_grid")
+    plot_grid(groups_orth, "orth_err_normalized_grid",
+              os.path.join(args.out, "fig_gennorm_orth.pdf"))
+
+    # Single-panel main-text headline (spiral, all three methods)
+    plot_main_panel(groups_r2,
+                    os.path.join(args.out, "fig_gennorm_main.pdf"))
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_reacher.py

@@ -0,0 +1,503 @@
+"""
+Aggregate and plot Reacher sweep results.
+
+Generates:
+  1. OU: R² vs ρ (per lambda)
+  2. OU: R² vs ρ (per dimension)
+  3. Traj: per-dim R² vs δ (with per-dim ρ annotations)
+  4. OU vs Traj on same axes (ρ on x-axis, traj uses measured ρ_mean)
+  5. Lambda robustness panel (OU, R² vs λ for each ρ)
+  6. Traj: R² vs measured ρ
+  7. Orthogonality error plots
+
+Usage:
+    python analysis/plot_reacher.py --results_dir results/reacher
+"""
+
+import os
+import json
+import argparse
+import numpy as np
+from pathlib import Path
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from collections import defaultdict
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# DATA LOADING
+# ═════════════════════════════════════════════════════════════════════════════
+
+def load_summaries(results_dir):
+    """Load all summary_*.json files, split into OU and traj."""
+    ou_results, traj_results = [], []
+    for p in sorted(Path(results_dir).glob("summary_*.json")):
+        with open(p) as f:
+            summary = json.load(f)
+        for run_name, r in summary.items():
+            if r.get("r2_hz", -999) <= -1:
+                continue
+            if "delta" in r:
+                traj_results.append(r)
+            else:
+                ou_results.append(r)
+    return ou_results, traj_results
+
+
+def _group_by(results, x_key):
+    """Group results by (x_key, lamb) → list of result dicts."""
+    grouped = defaultdict(list)
+    for r in results:
+        grouped[(r[x_key], r["lamb"])].append(r)
+    return grouped
+
+
+def _get_sorted(results, key):
+    return sorted(set(r[key] for r in results))
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 1. OU: R² vs ρ (one curve per lambda)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_ou_r2_vs_rho(results, save_path):
+    grouped = _group_by(results, "rho")
+    lambs = _get_sorted(results, "lamb")
+    rhos = _get_sorted(results, "rho")
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    for lamb in lambs:
+        medians, q25, q75, xs = [], [], [], []
+        for rho in rhos:
+            vals = [r["r2_hz"] for r in grouped.get((rho, lamb), [])]
+            if vals:
+                medians.append(np.median(vals))
+                q25.append(np.percentile(vals, 25))
+                q75.append(np.percentile(vals, 75))
+                xs.append(rho)
+        medians, q25, q75 = np.array(medians), np.array(q25), np.array(q75)
+        ax.plot(xs, medians, "o-", label=f"λ={lamb:.0e}", markersize=5)
+        ax.fill_between(xs, q25, q75, alpha=0.15)
+
+    ax.set_xlabel("ρ (OU autocorrelation)", fontsize=12)
+    ax.set_ylabel("R² (embed → true state)", fontsize=12)
+    ax.set_title("OU: Linear identifiability vs. ρ", fontsize=13)
+    ax.legend(fontsize=9)
+    ax.set_ylim(-0.05, 1.05)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 2. OU: per-dimension R² vs ρ
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_ou_perdim_r2(results, save_path):
+    """Two curves: shoulder vs wrist, averaged over lambda and seeds."""
+    by_rho = defaultdict(list)
+    for r in results:
+        if "r2_hz_per_dim" in r:
+            by_rho[r["rho"]].append(r["r2_hz_per_dim"])
+
+    rhos = sorted(by_rho.keys())
+    dim0_med, dim1_med = [], []
+    for rho in rhos:
+        vals = np.array(by_rho[rho])
+        dim0_med.append(np.median(vals[:, 0]))
+        dim1_med.append(np.median(vals[:, 1]))
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    ax.plot(rhos, dim0_med, "o-", label="Shoulder (dim 0)", markersize=5)
+    ax.plot(rhos, dim1_med, "s-", label="Wrist (dim 1)", markersize=5)
+
+    ax.set_xlabel("ρ (OU autocorrelation)", fontsize=12)
+    ax.set_ylabel("R² per dimension", fontsize=12)
+    ax.set_title("OU: Per-dimension identifiability", fontsize=13)
+    ax.legend(fontsize=10)
+    ax.set_ylim(-0.05, 1.05)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 3. Traj: per-dimension R² vs δ (with ρ annotations)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_traj_perdim_r2(results, save_path):
+    """Two curves: shoulder vs wrist, annotated with per-dim ρ."""
+    by_delta = defaultdict(list)
+    for r in results:
+        if "r2_hz_per_dim" in r:
+            by_delta[r["delta"]].append(r)
+
+    deltas = sorted(by_delta.keys())
+    dim0_med, dim1_med = [], []
+    rho0_vals, rho1_vals = [], []
+
+    for delta in deltas:
+        runs = by_delta[delta]
+        vals = np.array([r["r2_hz_per_dim"] for r in runs])
+        dim0_med.append(np.median(vals[:, 0]))
+        dim1_med.append(np.median(vals[:, 1]))
+        rho0 = [r.get("rho_shoulder", None) for r in runs]
+        rho1 = [r.get("rho_wrist", None) for r in runs]
+        rho0_vals.append(rho0[0] if rho0[0] is not None else None)
+        rho1_vals.append(rho1[0] if rho1[0] is not None else None)
+
+    fig, ax = plt.subplots(figsize=(8, 5))
+    ax.plot(deltas, dim0_med, "o-", label="Shoulder (dim 0)",
+            markersize=5, color="C0")
+    ax.plot(deltas, dim1_med, "s-", label="Wrist (dim 1)",
+            markersize=5, color="C1")
+
+    # Annotate with ρ values
+    for i, delta in enumerate(deltas):
+        if rho0_vals[i] is not None:
+            y_pos = max(dim0_med[i], dim1_med[i]) + 0.03
+            ax.annotate(f"ρ₀={rho0_vals[i]:.3f}\nρ₁={rho1_vals[i]:.3f}",
+                        (delta, y_pos),
+                        fontsize=7, ha="center", alpha=0.7)
+
+    ax.set_xlabel("δ (temporal stride)", fontsize=12)
+    ax.set_ylabel("R² per dimension", fontsize=12)
+    ax.set_title("Trajectory: Per-dimension identifiability vs. δ", fontsize=13)
+    ax.legend(fontsize=10)
+    ax.set_ylim(-0.15, 1.15)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 4. OU vs Traj on same axes (ρ on x-axis)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_ou_vs_traj(ou_results, traj_results, save_path):
+    """Both conditions on one plot, using ρ as common x-axis."""
+    # OU: group by rho, average over lambda and seeds
+    ou_by_rho = defaultdict(list)
+    for r in ou_results:
+        ou_by_rho[r["rho"]].append(r["r2_hz"])
+
+    # Traj: use rho_mean, group by delta
+    traj_by_rho = defaultdict(list)
+    for r in traj_results:
+        rho_mean = r.get("rho_mean", None)
+        if rho_mean is not None:
+            traj_by_rho[rho_mean].append(r["r2_hz"])
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+
+    # OU
+    rhos_ou = sorted(ou_by_rho.keys())
+    med_ou = [np.median(ou_by_rho[rho]) for rho in rhos_ou]
+    q25_ou = [np.percentile(ou_by_rho[rho], 25) for rho in rhos_ou]
+    q75_ou = [np.percentile(ou_by_rho[rho], 75) for rho in rhos_ou]
+    ax.plot(rhos_ou, med_ou, "o-", label="OU (Gaussian)", markersize=6,
+            color="C0", linewidth=2)
+    ax.fill_between(rhos_ou, q25_ou, q75_ou, alpha=0.15, color="C0")
+
+    # Traj
+    rhos_traj = sorted(traj_by_rho.keys())
+    med_traj = [np.median(traj_by_rho[rho]) for rho in rhos_traj]
+    q25_traj = [np.percentile(traj_by_rho[rho], 25) for rho in rhos_traj]
+    q75_traj = [np.percentile(traj_by_rho[rho], 75) for rho in rhos_traj]
+    ax.plot(rhos_traj, med_traj, "s-", label="Trajectory (non-Gaussian)",
+            markersize=6, color="C3", linewidth=2)
+    ax.fill_between(rhos_traj, q25_traj, q75_traj, alpha=0.15, color="C3")
+
+    ax.set_xlabel("ρ (autocorrelation)", fontsize=12)
+    ax.set_ylabel("R² (embed → true state)", fontsize=12)
+    ax.set_title("Gaussian vs. non-Gaussian latents", fontsize=13)
+    ax.legend(fontsize=10)
+    ax.set_ylim(-0.05, 1.05)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 5. Lambda robustness (OU)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_lambda_robustness(results, save_path):
+    """R² vs lambda for each rho."""
+    grouped = _group_by(results, "rho")
+    rhos = _get_sorted(results, "rho")
+    lambs = _get_sorted(results, "lamb")
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    for rho in rhos:
+        medians, xs = [], []
+        for lamb in lambs:
+            vals = [r["r2_hz"] for r in grouped.get((rho, lamb), [])]
+            if vals:
+                medians.append(np.median(vals))
+                xs.append(lamb)
+        if medians:
+            ax.plot(xs, medians, "o-", label=f"ρ={rho}", markersize=4)
+
+    ax.set_xscale("log")
+    ax.set_xlabel("λ (SIGReg weight)", fontsize=12)
+    ax.set_ylabel("R² (embed → true state)", fontsize=12)
+    ax.set_title("OU: Robustness to λ", fontsize=13)
+    ax.legend(fontsize=8, ncol=2)
+    ax.set_ylim(-0.05, 1.05)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 6. Traj: R² vs measured ρ (per lambda)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_traj_r2_vs_rho(results, save_path):
+    """Traj R² plotted against measured autocorrelation, per lambda."""
+    grouped = defaultdict(list)
+    for r in results:
+        rho_mean = r.get("rho_mean", None)
+        if rho_mean is not None:
+            grouped[(rho_mean, r["lamb"])].append(r["r2_hz"])
+
+    lambs = _get_sorted(results, "lamb")
+    rhos = sorted(set(k[0] for k in grouped.keys()))
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    for lamb in lambs:
+        medians, xs = [], []
+        for rho in rhos:
+            vals = grouped.get((rho, lamb), [])
+            if vals:
+                medians.append(np.median(vals))
+                xs.append(rho)
+        if medians:
+            ax.plot(xs, medians, "o-", label=f"λ={lamb:.0e}", markersize=5)
+
+    ax.set_xlabel("ρ (measured autocorrelation)", fontsize=12)
+    ax.set_ylabel("R² (embed → true state)", fontsize=12)
+    ax.set_title("Trajectory: Identifiability vs. measured ρ", fontsize=13)
+    ax.legend(fontsize=9)
+    ax.set_ylim(-0.05, 1.05)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 7. Orthogonality error
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_orth_err(results, x_key, x_label, title, save_path):
+    grouped = _group_by(results, x_key)
+    lambs = _get_sorted(results, "lamb")
+    x_vals = _get_sorted(results, x_key)
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    for lamb in lambs:
+        medians, xs = [], []
+        for x in x_vals:
+            vals = [r.get("orth_error", r.get("procrustes_error", 1.0))
+                    for r in grouped.get((x, lamb), [])]
+            if vals:
+                medians.append(np.median(vals))
+                xs.append(x)
+        ax.plot(xs, medians, "o-", label=f"λ={lamb:.0e}", markersize=4)
+
+    ax.set_xlabel(x_label, fontsize=12)
+    ax.set_ylabel("Orthogonality error", fontsize=12)
+    ax.set_title(title, fontsize=13)
+    ax.legend(fontsize=9)
+    ax.grid(alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {save_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# TABLE
+# ═════════════════════════════════════════════════════════════════════════════
+
+def print_table(results, x_key, label):
+    grouped = defaultdict(list)
+    for r in results:
+        grouped[(r[x_key], r["lamb"])].append(r)
+
+    print(f"\n=== {label} ===")
+    print(f"{x_key:>6s} {'lamb':>8s} {'R²(h→z)':>10s} "
+          f"{'R²(dim0)':>10s} {'R²(dim1)':>10s} "
+          f"{'R²(sincos)':>10s} {'orth_err':>10s} {'n':>4s}")
+    print("-" * 76)
+    for (x, lamb), runs in sorted(grouped.items()):
+        r2s = [r["r2_hz"] for r in runs]
+        errs = [r.get("orth_error", r.get("procrustes_error", 1.0))
+                for r in runs]
+
+        dim0 = [r["r2_hz_per_dim"][0] for r in runs
+                if "r2_hz_per_dim" in r]
+        dim1 = [r["r2_hz_per_dim"][1] for r in runs
+                if "r2_hz_per_dim" in r]
+        sincos = [r["r2_sincos"] for r in runs if "r2_sincos" in r]
+
+        d0 = f"{np.median(dim0):10.4f}" if dim0 else f"{'n/a':>10s}"
+        d1 = f"{np.median(dim1):10.4f}" if dim1 else f"{'n/a':>10s}"
+        sc = f"{np.median(sincos):10.4f}" if sincos else f"{'n/a':>10s}"
+
+        print(f"{x:6g} {lamb:8.1e} "
+              f"{np.median(r2s):10.4f} "
+              f"{d0} {d1} "
+              f"{sc} "
+              f"{np.median(errs):10.4f} "
+              f"{len(runs):4d}")
+
+
+def plot_traj_distributions(h5_path, results_dir, save_path):
+    """Marginal + transition distributions for trajectory data, annotated with R²."""
+    import h5py
+    from scipy.stats import pearsonr
+
+    with h5py.File(h5_path, "r") as f:
+        qpos = np.array(f["qpos"])
+        ep_len = np.array(f["ep_len"])
+    T = ep_len[0]
+    episodes = qpos.reshape(-1, T, 2)
+
+    # Load R² per delta
+    grouped = defaultdict(list)
+    for p in Path(results_dir).rglob("result.json"):
+        r = json.load(open(p))
+        if "delta" not in r or r.get("rho") is not None:
+            continue
+        grouped[(r["delta"], r["lamb"])].append(r)
+
+    r2_dict = {}
+    for delta in set(k[0] for k in grouped):
+        best_mean, best_lamb = -np.inf, None
+        for lamb in set(k[1] for k in grouped if k[0] == delta):
+            m = np.mean([r["r2_hz"] for r in grouped[(delta, lamb)]])
+            if m > best_mean:
+                best_mean, best_lamb = m, lamb
+        runs = grouped[(delta, best_lamb)]
+        r2_dict[delta] = {
+            "r2_dim0": np.mean([r["r2_hz_per_dim"][0] for r in runs]),
+            "r2_dim1": np.mean([r["r2_hz_per_dim"][1] for r in runs]),
+        }
+
+    DELTAS = [1, 2, 4, 8, 16, 32, 64]
+    s = 0.0001
+    sub = 1
+
+    fig = plt.figure(figsize=0.85 * np.array((1 + 3 * len(DELTAS), 5)))
+    gs = fig.add_gridspec(2, 2 + len(DELTAS))
+
+    ax = fig.add_subplot(gs[:2, :2])
+    ax.scatter(*episodes.reshape(-1, 2)[::sub].T, s=s * 10)
+    ax.set_title("Marginal")
+    ax.grid()
+    ax.set_xlabel(r"$z_0$ (shoulder)")
+    ax.set_ylabel(r"$z_1$ (wrist)")
+
+    for i, delta in enumerate(DELTAS):
+        # Transition scatter
+        ax = fig.add_subplot(gs[0, 2 + i])
+        transitions = (episodes[:, delta:] - episodes[:, :-delta]).reshape(-1, 2)
+        ax.scatter(*transitions[::sub].T, s=s)
+        ax.set_title(
+            r"$\Delta=%d$" % delta + "\n"
+            + r"$R^2=(%.2f,\,%.2f)$"
+            % (r2_dict[delta]["r2_dim0"], r2_dict[delta]["r2_dim1"])
+        )
+        ax.grid()
+
+        # Autocorrelation scatter
+        ax = fig.add_subplot(gs[1, 2 + i])
+        a = episodes[:, delta:, 0].flatten()[::sub]
+        b = episodes[:, :-delta, 0].flatten()[::sub]
+        rho0 = pearsonr(a, b)[0]
+        ax.scatter(a, b, s=s)
+
+        c = episodes[:, delta:, 1].flatten()[::sub]
+        d = episodes[:, :-delta, 1].flatten()[::sub]
+        rho1 = pearsonr(c, d)[0]
+        ax.scatter(c, d, s=s)
+
+        ax.set_title(r"$\rho=(%.2f,\,%.2f)$" % (rho0, rho1))
+        ax.grid()
+        if i == 0:
+            ax.legend([r"$z_0$ (shoulder)", r"$z_1$ (wrist)"], loc="upper left")
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=500, bbox_inches="tight", format="jpg")
+    plt.close()
+    print(f"Saved {save_path}")
+    
+
+# ═════════════════════════════════════════════════════════════════════════════
+# MAIN
+# ═════════════════════════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--results_dir", type=str, default="results/reacher")
+    parser.add_argument("--out_dir", type=str, default="figures/reacher")
+    parser.add_argument("--h5_path", type=str, default="data/reacher.h5")
+    args = parser.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    ou_results, traj_results = load_summaries(args.results_dir)
+    print(f"Loaded {len(ou_results)} OU runs, {len(traj_results)} traj runs")
+
+    # Tables
+    if ou_results:
+        print_table(ou_results, "rho", "OU")
+    if traj_results:
+        print_table(traj_results, "delta", "Trajectory")
+
+    # OU plots
+    if ou_results:
+        plot_ou_r2_vs_rho(ou_results, out_dir / "ou_r2_vs_rho.png")
+        plot_ou_perdim_r2(ou_results, out_dir / "ou_perdim_r2.png")
+        plot_lambda_robustness(ou_results, out_dir / "ou_lambda_robustness.png")
+        plot_orth_err(ou_results, "rho", "ρ",
+                      "OU: Orthogonality error vs. ρ",
+                      out_dir / "ou_orth_err.png")
+
+    # Traj plots
+    if traj_results:
+        plot_traj_perdim_r2(traj_results, out_dir / "traj_perdim_r2.png")
+        plot_traj_r2_vs_rho(traj_results, out_dir / "traj_r2_vs_rho.png")
+        plot_orth_err(traj_results, "delta", "δ",
+                      "Trajectory: Orthogonality error vs. δ",
+                      out_dir / "traj_orth_err.png")
+        
+        h5_path = os.path.join(args.h5_path)
+        if os.path.exists(h5_path):
+            plot_traj_distributions(h5_path, args.results_dir,
+                                     out_dir / "traj_distributions.png")
+
+    # Combined
+    if ou_results and traj_results:
+        plot_ou_vs_traj(ou_results, traj_results,
+                        out_dir / "ou_vs_traj.png")
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_scaling.py

@@ -0,0 +1,73 @@
+"""
+Scaling plots: R² and orthogonality vs latent dimension N.
+
+Usage:
+    python analysis/plot_scaling.py --results_dir results/scaling/ --out figures/
+"""
+
+import argparse, os, glob, json
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+
+def load_results(results_dir):
+    rows = []
+    for path in sorted(glob.glob(os.path.join(results_dir, "*.json"))):
+        with open(path) as f:
+            r = json.load(f)
+        rows.append({k: r.get(k) for k in [
+            "N", "seed", "r2_zx", "r2_xz", "r2_zh", "r2_hz",
+            "orth_err", "orth_err_normalized", "final_loss",
+            "final_align", "final_sigreg", "final_whiten",
+            "epsilon", "delta", "approx_bound", "procrustes_mse",
+        ]})
+    return pd.DataFrame(rows)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dir", default="results/scaling/")
+    p.add_argument("--out", default="figures/")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    df = load_results(args.results_dir)
+    if len(df) == 0:
+        print("No results."); return
+
+    summary = df.groupby("N").agg(
+        r2_xz_mean=("r2_xz", "mean"), r2_xz_std=("r2_xz", "std"),
+        r2_hz_mean=("r2_hz", "mean"), r2_hz_std=("r2_hz", "std"),
+        orth_mean=("orth_err_normalized", "mean"), orth_std=("orth_err_normalized", "std"),
+    ).reset_index()
+    dims = summary["N"].values
+
+    fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
+
+    ax = axes[0]
+    ax.errorbar(dims, summary["r2_xz_mean"], yerr=summary["r2_xz_std"],
+                fmt="o-", capsize=3, color="gray", label=r"Probe: $g(z) \to z$")
+    ax.errorbar(dims, summary["r2_hz_mean"], yerr=summary["r2_hz_std"],
+                fmt="s-", capsize=3, label=r"Probe: $f \circ g(z) \to z$")
+    ax.set_xscale("log", base=2)
+    ax.set_xlabel("Latent dimension $N$"); ax.set_ylabel(r"Linearity [$R^2$]")
+    ax.set_title("Latent Recovery"); ax.set_ylim(-0.05, 1.05)
+    ax.set_xticks(dims); ax.legend(); ax.grid(alpha=0.3)
+
+    ax = axes[1]
+    ax.errorbar(dims, summary["orth_mean"], yerr=summary["orth_std"],
+                fmt="D-", capsize=3, color="tab:green")
+    ax.set_xscale("log", base=2)
+    ax.set_xlabel("Latent dimension $N$")
+    ax.set_ylabel(r"$\|A^\top A - I\|_F / \sqrt{N}$")
+    ax.set_title("Orthogonality Error"); ax.set_xticks(dims); ax.grid(alpha=0.3)
+
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.out, "fig_scaling.pdf"), bbox_inches="tight")
+    print("Saved fig_scaling.pdf")
+    plt.close()
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/plot_scatter.py

@@ -0,0 +1,103 @@
+"""
+Cross-experiment scatter plots (2x2 panel).
+
+Usage:
+    python analysis/plot_scatter.py --results_dirs results/2d results/scaling results/grid results/ablation --out figures/
+"""
+
+import argparse, os, glob, json
+import numpy as np
+import matplotlib.pyplot as plt
+
+EXPERIMENT_COLORS = {"2d": "tab:blue", "grid": "tab:red", "scaling": "tab:green", "ablation": "tab:orange"}
+EXPERIMENT_ORDER = ["grid", "scaling", "2d", "ablation"]
+
+
+def load_all(dirs):
+    rows = []
+    for d in dirs:
+        for path in sorted(glob.glob(os.path.join(d, "*.json"))):
+            try:
+                with open(path) as f:
+                    rows.append(json.load(f))
+            except Exception:
+                pass
+    return rows
+
+
+def scatter_by_experiment(ax, rows, x_key, y_key):
+    for exp in EXPERIMENT_ORDER:
+        pts = [r for r in rows if r.get("experiment") == exp
+               and r.get(x_key) is not None and r.get(y_key) is not None]
+        if not pts:
+            continue
+        ax.scatter([r[x_key] for r in pts],
+                   [r[y_key] for r in pts],
+                   c=EXPERIMENT_COLORS[exp],
+                   s=25, alpha=0.7, edgecolors='k', linewidths=0.3,
+                   label=exp, zorder=3)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--results_dirs", nargs="+", required=True)
+    p.add_argument("--out", default="figures/")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    rows = load_all(args.results_dirs)
+    if not rows:
+        print("No results."); return
+
+    xlim = (5e-3, 1e0)
+    ylim = (0.9, 1.01)
+
+    fig = plt.figure(figsize=0.65 * np.array((8, 7)))
+
+    # ── Total loss vs R² ──
+    ax = plt.subplot(2, 2, 1)
+    scatter_by_experiment(ax, rows, "final_loss", "r2_hz")
+    ax.set_xlabel("Total loss")
+    ax.set_ylabel("Linear Identifiability")
+    ax.legend(fontsize=8)
+    ax.grid(alpha=0.3)
+    ax.set_xlim(*xlim)
+    ax.set_ylim(*ylim)
+    ax.set_xscale("log")
+
+    # ── Alignment vs R² ──
+    ax = plt.subplot(2, 2, 2)
+    scatter_by_experiment(ax, rows, "final_align", "r2_hz")
+    ax.set_xlabel("Alignment loss")
+    ax.set_ylabel("Linear Identifiability")
+    ax.grid(alpha=0.3)
+    ax.set_xlim(*xlim)
+    ax.set_ylim(*ylim)
+    ax.set_xscale("log")
+
+    # ── SIGReg vs R² ──
+    ax = plt.subplot(2, 2, 3)
+    scatter_by_experiment(ax, rows, "final_sigreg", "r2_hz")
+    ax.set_xlabel("SIGReg loss")
+    ax.set_ylabel("Linear Identifiability")
+    ax.grid(alpha=0.3)
+    ax.set_ylim(*ylim)
+    ax.set_xscale("log")
+
+    # ── SIGReg vs whitening ──
+    ax = plt.subplot(2, 2, 4)
+    scatter_by_experiment(ax, rows, "final_sigreg", "final_whiten")
+    ax.set_xlabel("SIGReg loss")
+    ax.set_ylabel("Whitening loss")
+    ax.grid(alpha=0.3)
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.out, "scatter_plots.pdf"), bbox_inches="tight")
+    print("Saved scatter_plots.pdf")
+    plt.close()
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/analysis/run_all.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+python analysis/aggregate.py --results_dir results/2d/ --out results/2d/summary.csv
+python analysis/aggregate.py --results_dir results/ --recursive --out results/all.csv
+python analysis/plot_2d.py --results_dir results/2d/ --out figures/
+python analysis/plot_bound.py --results_dirs results/grid results/2d results/scaling results/ablation --out figures/
+python analysis/plot_ablation.py --results_dir results/ablation/ --out figures/
+python analysis/plot_scaling.py --results_dir results/scaling/ --out figures/
+python analysis/plot_scatter.py --results_dirs results/2d results/scaling results/grid results/ablation --out figures/
+python analysis/make_table_scaling.py --results_dir results/scaling/

JEPA/lejepa-identifiability/experiments/configs/2d.yaml

@@ -0,0 +1,30 @@
+# 2D illustration + Gaussian half of the regularizer ablation
+# python run.py --config configs/2d.yaml --run spiral_lejepa --seed 1337
+
+experiment: 2d
+out: results/2d
+
+# Data
+N: 2
+source_dist: gaussian
+num_eval: 10000
+
+# Training (shared)
+steps: 20000
+lr: 3.0e-3
+batch_size: 256
+rho: 0.95
+log_every: 500
+
+# Per-run specs: 4 mixings x 2 objectives
+runs:
+  spiral_lejepa:   {mixing: spiral,   encoder: mlp,     hidden: 256, mode: lejepa, lamb: 1.0e-3}
+  spiral_whiten:   {mixing: spiral,   encoder: mlp,     hidden: 256, mode: whiten, lamb: 0.5}
+  banana_lejepa:   {mixing: banana,   encoder: mlp,     hidden: 256, mode: lejepa, lamb: 1.0e-3}
+  banana_whiten:   {mixing: banana,   encoder: mlp,     hidden: 256, mode: whiten, lamb: 0.5}
+  sinusoid_lejepa: {mixing: sinusoid, encoder: mlp,     hidden: 256, mode: lejepa, lamb: 1.0e-3}
+  sinusoid_whiten: {mixing: sinusoid, encoder: mlp,     hidden: 256, mode: whiten, lamb: 0.5}
+  nvp_lejepa:      {mixing: nvp,      encoder: matched, n_layers: 8, mode: lejepa, lamb: 1.0e-3}
+  nvp_whiten:      {mixing: nvp,      encoder: matched, n_layers: 8, mode: whiten, lamb: 0.5}
+
+seeds: [1337, 1338, 1339]

JEPA/lejepa-identifiability/experiments/configs/gennorm.yaml

@@ -0,0 +1,32 @@
+# Generalized normal sweep across mixings (main-text figure)
+# python run.py --config configs/gennorm.yaml --run spiral_lejepa --alpha 2.0 --seed 1337
+
+experiment: gennorm
+out: results/gennorm
+
+N: 2
+source_dist: gennorm   # alpha provided per-run via CLI
+num_eval: 10000
+
+steps: 20000
+lr: 3.0e-3
+batch_size: 256
+rho: 0.95
+log_every: 500
+
+runs:
+  spiral_lejepa:    {mixing: spiral,   encoder: mlp,     hidden: 256, mode: lejepa,  lamb: 1.0e-3}
+  spiral_whiten:    {mixing: spiral,   encoder: mlp,     hidden: 256, mode: whiten,  lamb: 0.5}
+  spiral_infonce:   {mixing: spiral,   encoder: mlp,     hidden: 256, mode: infonce, sigma: 1.0}
+  banana_lejepa:    {mixing: banana,   encoder: mlp,     hidden: 256, mode: lejepa,  lamb: 1.0e-3}
+  banana_whiten:    {mixing: banana,   encoder: mlp,     hidden: 256, mode: whiten,  lamb: 0.5}
+  banana_infonce:   {mixing: banana,   encoder: mlp,     hidden: 256, mode: infonce, sigma: 1.0}
+  sinusoid_lejepa:  {mixing: sinusoid, encoder: mlp,     hidden: 256, mode: lejepa,  lamb: 1.0e-3}
+  sinusoid_whiten:  {mixing: sinusoid, encoder: mlp,     hidden: 256, mode: whiten,  lamb: 0.5}
+  sinusoid_infonce: {mixing: sinusoid, encoder: mlp,     hidden: 256, mode: infonce, sigma: 1.0}
+  nvp_lejepa:       {mixing: nvp,      encoder: matched, n_layers: 8, mode: lejepa,  lamb: 1.0e-3}
+  nvp_whiten:       {mixing: nvp,      encoder: matched, n_layers: 8, mode: whiten,  lamb: 0.5}
+  nvp_infonce:      {mixing: nvp,      encoder: matched, n_layers: 8, mode: infonce, sigma: 1.0}
+
+alphas: [0.125, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]
+seeds: [1337, 1338, 1339]

JEPA/lejepa-identifiability/experiments/configs/grid.yaml

@@ -0,0 +1,27 @@
+# Grid search over lambda and rho (bound verification figure)
+# python run.py --config configs/grid.yaml --lamb 0.01 --rho 0.9 --seed 0
+
+experiment: grid
+out: results/grid
+
+# Data
+N: 2
+source_dist: gaussian
+num_eval: 10000
+
+# Training (shared)
+steps: 20000
+lr: 3.0e-3
+batch_size: 256
+log_every: 500
+
+# Encoder
+encoder: mlp
+hidden: 256
+mixing: spiral
+mode: lejepa
+
+# Sweep dimensions
+lambs: [1.0e-6, 1.0e-5, 1.0e-4, 1.0e-3, 5.0e-3, 1.0e-2, 5.0e-2, 1.0e-1, 5.0e-1]
+rhos: [0.3, 0.5, 0.7, 0.8, 0.9, 0.95, 0.99]
+seeds: [0, 1, 2]

JEPA/lejepa-identifiability/experiments/configs/reacher.yaml

@@ -0,0 +1,29 @@
+# Reacher pixel-observation identifiability experiment
+# Works for both OU and trajectory data — just point --data_dir at the right place.
+#
+# Prerender:
+#   python prerender.py eval
+#   python prerender.py ou --rho 0.95
+#   python prerender.py traj --delta 16 --h5_path data/reacher.h5
+#
+# Train:
+#   python run_reacher.py --config configs/reacher.yaml \
+#       --data_dir data/reacher/ou/rho=0.95
+
+experiment: reacher
+out: results/reacher
+data_root: data/reacher
+
+# Model
+d_latent: 2
+
+# Training
+epochs: 100
+batch_size: 256
+lr: 3.0e-3
+n_slices: 256
+n_eval_fast: 2000
+
+# Sweep dimensions
+lambs: [1.0e-3, 5.0e-3, 1.0e-2, 5.0e-2]
+seeds: [0, 1, 2]

JEPA/lejepa-identifiability/experiments/configs/scaling.yaml

@@ -0,0 +1,36 @@
+# Scaling experiment (paper figure)
+# python run.py --config configs/scaling.yaml --N 16 --seed 0
+# python run.py --config configs/scaling.yaml --N 16 --seed 0 --mode infonce
+# python run.py --config configs/scaling.yaml --N 16 --seed 0 --mode whiten
+
+experiment: scaling
+out: results/scaling
+
+# Data
+source_dist: gaussian
+num_eval: 10000
+
+# Training (shared)
+steps: 20000
+lr: 3.0e-3
+batch_size: 256
+rho: 0.95
+log_every: 500
+
+# Encoder
+encoder: matched
+n_layers: 4
+mode: lejepa     # default; override with --mode
+
+# Mode-specific defaults (used based on --mode)
+lamb: 1.0e-6     # for lejepa
+lamb_whiten: 0.5  # used when mode=whiten
+sigma: 1.0       # for infonce
+
+# Mixing
+mixing: coupling
+
+# Sweep dimensions
+dims: [2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
+seeds: [0, 1, 2, 3, 4]
+K: 3  # parallel encoder runs per (N, seed); pick lowest loss

JEPA/lejepa-identifiability/experiments/lejepa_id/data.py

@@ -0,0 +1,36 @@
+"""Data generation: latent sources and OU augmentation."""
+
+import math
+import torch
+
+
+def _gennorm_unit_var_scale(alpha):
+    """Scale β so gennorm(α, β) has unit variance: β = sqrt(Γ(1/α) / Γ(3/α))."""
+    return math.exp(0.5 * (math.lgamma(1.0 / alpha) - math.lgamma(3.0 / alpha)))
+
+
+def sample_latents(D, N, dist="gaussian", device="cuda", alpha=None):
+    """Sample D points in R^N (unit variance)."""
+    if dist == "gaussian":
+        return torch.randn(D, N, device=device)
+    elif dist == "laplace":
+        return torch.distributions.Laplace(0, 1 / (2 ** 0.5)).sample((D, N)).to(device)
+    elif dist == "gennorm":
+        if alpha is None:
+            raise ValueError("gennorm requires alpha")
+        scale = _gennorm_unit_var_scale(alpha)
+        u = torch.distributions.Gamma(1.0 / alpha, 1.0).sample((D, N)).to(device)
+        sign = torch.randint(0, 2, (D, N), device=device).float() * 2 - 1
+        return scale * sign * u.pow(1.0 / alpha)
+    else:
+        raise ValueError(f"Unknown distribution: {dist}")
+
+
+def ou_augment(z, rho, n_views=2, dist="gaussian", alpha=None):
+    """OU channel: z' = ρz + √(1-ρ²)η, η drawn from same dist as source.
+    Returns (V, B, N)."""
+    fac = (1 - rho ** 2) ** 0.5
+    D, N = z.shape
+    eta = sample_latents(n_views * D, N, dist=dist, device=z.device, alpha=alpha)
+    eta = eta.reshape(n_views, D, N)
+    return rho * z.unsqueeze(0) + fac * eta

JEPA/lejepa-identifiability/experiments/lejepa_id/engine.py

@@ -0,0 +1,143 @@
+"""
+Core training engine — single function used by all experiments.
+
+Handles: LR schedule (warmup + cosine), online data generation,
+loss computation (lejepa or whiten), periodic evaluation of ALL metrics
+on a fixed eval set, standardized output schema.
+"""
+
+import torch
+import numpy as np
+
+from .losses import SIGReg, whitening_loss, alignment_loss, infonce_loss
+from .data import sample_latents, ou_augment
+from .metrics import compute_all_metrics
+
+
+def warmup_cosine_lr(step, total_steps, base_lr):
+    """Constant for first half, cosine decay for second half."""
+    warmup = total_steps // 2
+    if step < warmup:
+        return base_lr
+    t = (step - warmup) / (total_steps - warmup)
+    return base_lr * 0.5 * (1 + np.cos(np.pi * t))
+
+
+def train_and_evaluate(
+    encoder,
+    mix_fn,
+    *,
+    N,
+    rho,
+    lamb,
+    sigma=1.0,
+    mode="lejepa",
+    source_dist="gaussian",
+    source_alpha=None,
+    steps=20000,
+    batch_size=256,
+    lr=3e-3,
+    z_eval,
+    log_every=100,
+    device="cuda",
+):
+    """Train encoder and evaluate periodically.
+
+    Args:
+        encoder: nn.Module, x -> h
+        mix_fn: callable, z -> x
+        N: latent dimension
+        rho: OU correlation
+        lamb: regularization weight
+        mode: "lejepa" or "whiten"
+        source_dist: "gaussian", "laplace", or "gennorm"
+        steps: total training steps
+        batch_size: batch size (online data)
+        lr: peak learning rate
+        z_eval: (num_eval, N) fixed eval tensor
+        log_every: eval frequency
+        device: torch device string
+
+    Returns:
+        encoder: trained encoder
+        log: dict of lists — training curves and periodic eval metrics
+    """
+    sigreg = SIGReg().to(device)
+    opt = torch.optim.AdamW(encoder.parameters(), lr=lr)
+
+    # Precompute eval mixing (constant across training)
+    x_eval = mix_fn(z_eval)
+
+    log_keys = [
+        "step", "lr",
+        # training losses
+        "align", "sigreg", "whiten", "total",
+        # eval metrics
+        "r2_zx", "r2_xz", "r2_zh", "r2_hz",
+        "orth_err", "orth_err_normalized",
+        "epsilon", "delta", "D_bound", "approx_bound",
+        "procrustes_mse", "L_h", "trace_cov",
+    ]
+    log = {k: [] for k in log_keys}
+
+    for step in range(steps + 1):
+        # LR schedule
+        current_lr = warmup_cosine_lr(step, steps, lr)
+        for pg in opt.param_groups:
+            pg["lr"] = current_lr
+
+        # Online data
+        z_batch = sample_latents(batch_size, N, dist=source_dist,
+                                 device=device, alpha=source_alpha)
+        z_aug = ou_augment(z_batch, rho, dist=source_dist, alpha=source_alpha)  # (2, B, N)
+        h = encoder(mix_fn(z_aug).flatten(0, 1)).reshape(2, batch_size, N)
+
+        align = alignment_loss(h)
+        sig = sigreg(h)
+        wht = whitening_loss(h)
+
+        if mode == "lejepa":
+            loss = lamb * sig + (1 - lamb) * align
+        elif mode == "whiten":
+            loss = lamb * wht + (1 - lamb) * align
+        elif mode == "infonce":
+            loss = infonce_loss(h, sigma)
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+
+        if step % log_every == 0 or (step < 1000 and step % 100 == 0):
+            log["step"].append(step)
+            log["lr"].append(current_lr)
+            log["align"].append(align.item())
+            log["sigreg"].append(sig.item())
+            log["whiten"].append(wht.item())
+            log["total"].append(loss.item())
+
+            # Full eval on fixed set
+            encoder.eval()
+            with torch.no_grad():
+                h_eval = encoder(x_eval)
+                z_prime = ou_augment(
+                    z_eval, rho, n_views=1, 
+                    dist=source_dist, alpha=source_alpha
+                ).squeeze(0)
+                h_prime = encoder(mix_fn(z_prime))
+
+            metrics = compute_all_metrics(z_eval, x_eval, h_eval, h_prime, rho, N)
+            
+            for k, v in metrics.items():
+                log[k].append(v)
+
+            encoder.train()
+
+            if step % (log_every * 10) == 0:
+                print(f"  step {step:5d} | lr={current_lr:.1e} "
+                      f"align={align.item():.2e} sig={sig.item():.1f} "
+                      f"R²(h->z)={metrics['r2_hz']:.4f} "
+                      f"orth={metrics['orth_err']:.4f}")
+
+    return encoder, log

JEPA/lejepa-identifiability/experiments/lejepa_id/losses.py

@@ -0,0 +1,53 @@
+"""Loss functions."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SIGReg(nn.Module):
+    """Sliced characteristic function regularizer (Balestriero & LeCun 2025)."""
+
+    def __init__(self, knots=17, n_slices=256, t_max=3.0):
+        super().__init__()
+        self.n_slices = n_slices
+        t = torch.linspace(0, t_max, knots)
+        dt = t_max / (knots - 1)
+        w = torch.full((knots,), 2 * dt)
+        w[[0, -1]] = dt
+        self.register_buffer("t", t)
+        self.register_buffer("phi", torch.exp(-t**2 / 2))
+        self.register_buffer("weights", w * torch.exp(-t**2 / 2))
+
+    def forward(self, h):
+        """h: (V, B, N) -> scalar."""
+        flat = h.flatten(0, 1)
+        A = F.normalize(torch.randn(flat.size(-1), self.n_slices, device=flat.device), dim=0)
+        xt = (flat @ A).unsqueeze(-1) * self.t
+        err = (xt.cos().mean(0) - self.phi) ** 2 + xt.sin().mean(0) ** 2
+        return (err @ self.weights).mean() * flat.size(0)
+
+
+def whitening_loss(h):
+    """||Cov(h) - I||²_F.  h: (V, B, N) -> scalar."""
+    flat = h.flatten(0, 1)
+    flat = flat - flat.mean(dim=0)
+    cov = (flat.T @ flat) / (flat.shape[0] - 1)
+    return (cov - torch.eye(flat.shape[1], device=h.device)).square().mean()
+
+
+def alignment_loss(h):
+    """Pull positive-pair views together. h: (V, B, N) -> scalar."""
+    return (h.mean(0) - h).square().mean()
+
+
+def infonce_loss(h, sigma):
+    """Symmetric Gaussian-kernel InfoNCE: sim(u, v) = -||u - v||² / (2σ²).
+    h: (V, B, N) with V=2 views. Negatives are other batch elements.
+    """
+    h1, h2 = h[0], h[1]                                              # (B, N) each
+    d12 = ((h1.unsqueeze(1) - h2.unsqueeze(0)) ** 2).sum(-1)         # (B, B)
+    sim = -d12 / (2 * sigma ** 2)
+    loss_a = -(sim.diag() - torch.logsumexp(sim, dim=1)).mean()
+    loss_b = -(sim.diag() - torch.logsumexp(sim, dim=0)).mean()
+    return 0.5 * (loss_a + loss_b)

JEPA/lejepa-identifiability/experiments/lejepa_id/metrics.py

@@ -0,0 +1,65 @@
+"""Evaluation metrics — standardized across all experiments."""
+
+import torch
+
+def bidirectional_r2(a, b):
+    """R²(a->b) and R²(b->a) via torch lstsq on GPU. a, b are tensors."""
+    def _r2(x, y):
+        x1 = torch.cat([x, torch.ones(len(x), 1, device=x.device)], dim=1)
+        W = torch.linalg.lstsq(x1, y).solution
+        ss_res = ((y - x1 @ W) ** 2).sum()
+        ss_tot = ((y - y.mean(0)) ** 2).sum()
+        return (1 - ss_res / ss_tot).item()
+    return _r2(a, b), _r2(b, a)
+
+
+def compute_all_metrics(z, x, h, h_prime, rho, N):
+    """All metrics on GPU. z, x, h, h_prime are torch tensors."""
+    r2_zx, r2_xz = bidirectional_r2(z, x)
+    r2_zh, r2_hz = bidirectional_r2(z, h)
+
+    # Orthogonality
+    z1 = torch.cat([z, torch.ones(len(z), 1, device=z.device)], dim=1)
+    W = torch.linalg.lstsq(z1, h).solution
+    A = W[:N].T
+    orth_err = torch.linalg.norm(A.T @ A - torch.eye(N, device=A.device), 'fro').item()
+    orth_err_normalized = orth_err / (N ** 0.5)
+
+    # Bound quantities
+    cov_h = torch.cov(h.T)
+    epsilon = torch.linalg.norm(cov_h - torch.eye(N, device=h.device), 'fro').item()
+    trace_cov = torch.trace(cov_h).item()
+    L_h = ((h_prime - h) ** 2).sum(dim=1).mean().item()
+    delta = max(L_h - 2 * (1 - rho) * trace_cov, 0.0)
+    spectral_gap = 2 * rho * (1 - rho)
+    D_bound = delta / spectral_gap if spectral_gap > 0 else float("inf")
+    approx_bound = D_bound + (epsilon + D_bound) ** 2
+
+    # Procrustes
+    M = (h.T @ z) / len(z)
+    U, S, Vt = torch.linalg.svd(M)
+    Q = U @ Vt
+    procrustes_mse = ((h - z @ Q.T) ** 2).sum(dim=1).mean().item()
+
+    return {
+        "r2_zx": r2_zx, "r2_xz": r2_xz,
+        "r2_zh": r2_zh, "r2_hz": r2_hz,
+        "orth_err": orth_err, "orth_err_normalized": orth_err_normalized,
+        "epsilon": epsilon, "delta": delta, "D_bound": D_bound,
+        "approx_bound": approx_bound, "procrustes_mse": procrustes_mse,
+        "L_h": L_h, "trace_cov": trace_cov,
+    }
+
+
+def compute_recovery_metrics(z, h, N, suffix=""):
+    """R² in both directions + orthogonality. Suffix appended to keys."""
+    r2_zh, r2_hz = bidirectional_r2(z, h)
+    z1 = torch.cat([z, torch.ones(len(z), 1, device=z.device)], dim=1)
+    W = torch.linalg.lstsq(z1, h).solution
+    A = W[:N].T
+    orth_err = torch.linalg.norm(A.T @ A - torch.eye(N, device=A.device), 'fro').item()
+    return {
+        f"r2_zh{suffix}": r2_zh, f"r2_hz{suffix}": r2_hz,
+        f"orth_err{suffix}": orth_err,
+        f"orth_err_normalized{suffix}": orth_err / (N ** 0.5),
+    }

JEPA/lejepa-identifiability/experiments/lejepa_id/mixing.py

@@ -0,0 +1,56 @@
+"""Nonlinear mixing functions."""
+
+import torch
+
+
+# ── 2D mixing functions ──────────────────────────────────────────────────────
+
+def mix_spiral(z):
+    """g(z) = R(π‖z‖) z — measure-preserving spiral diffeomorphism."""
+    norms = z.norm(dim=-1) * torch.pi
+    c, s = norms.cos(), norms.sin()
+    R = torch.stack([torch.stack([c, -s], dim=-1),
+                     torch.stack([s,  c], dim=-1)], dim=-2)
+    return (R @ z.unsqueeze(-1)).squeeze(-1)
+
+
+def mix_banana(z):
+    """Banana: x0 = z0, x1 = z1 + z0²."""
+    return torch.stack([z[..., 0], z[..., 1] + z[..., 0] ** 2], dim=-1)
+
+
+def mix_sinusoid(z):
+    """Sinusoidal shear: x0 = z0 + sin(1.5 z1), x1 = z1."""
+    return torch.stack([z[..., 0] + torch.sin(1.5 * z[..., 1]), z[..., 1]], dim=-1)
+
+
+MIXINGS_2D = {
+    "spiral": mix_spiral,
+    "banana": mix_banana,
+    "sinusoid": mix_sinusoid,
+    # "nvp" handled via make_coupling_mixing(N=2, n_layers=...)
+}
+
+
+# ── Coupling-layer mixing (any dimension) ────────────────────────────────────
+
+def make_coupling_mixing(N, n_layers=4, seed=1337, device="cuda"):
+    """RealNVP-style coupling layers. Works for any even N (including N=2)."""
+    half = N // 2
+    torch.manual_seed(seed)
+    Ws = []
+    for _ in range(n_layers):
+        W, _ = torch.linalg.qr(torch.randn(half, half, device=device))
+        Ws.append(W * 2.0)
+
+    def mix(z):
+        for i, W in enumerate(Ws):
+            z1, z2 = z[..., :half], z[..., half:]
+            if i % 2 == 0:
+                z2 = z2 + torch.tanh(z1 @ W)
+            else:
+                z1 = z1 + torch.tanh(z2 @ W)
+            z = torch.cat([z1, z2], dim=-1)
+        return z
+
+    return mix

JEPA/lejepa-identifiability/experiments/lejepa_id/models.py

@@ -0,0 +1,55 @@
+"""Encoder architectures."""
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def make_mlp_encoder(N, hidden=256, n_layers=4, device="cuda"):
+    """MLP encoder."""
+    layers = [nn.Linear(N, hidden), nn.GELU()]
+    for _ in range(n_layers - 1):
+        layers += [nn.Linear(hidden, hidden), nn.GELU()]
+    layers.append(nn.Linear(hidden, N))
+    return nn.Sequential(*layers).to(device)
+
+
+class MatchedEncoder(nn.Module):
+    """Inverse coupling-layer encoder matched to NVP mixing architecture."""
+
+    def __init__(self, N, n_layers=4, device="cuda"):
+        super().__init__()
+        half = N // 2
+        self.half = half
+        self.n_layers = n_layers
+        self.Ws = nn.ParameterList([
+            nn.Parameter(torch.randn(half, half, device=device) / np.sqrt(half))
+            for _ in range(n_layers)
+        ])
+
+    def forward(self, x):
+        for i, W in reversed(list(enumerate(self.Ws))):
+            z1, z2 = x[..., :self.half], x[..., self.half:]
+            if i % 2 == 0:
+                z2 = z2 - torch.tanh(z1 @ W)
+            else:
+                z1 = z1 - torch.tanh(z2 @ W)
+            x = torch.cat([z1, z2], dim=-1)
+        return x
+
+
+def make_matched_encoder(N, n_layers=4, seed=42, device="cuda"):
+    torch.manual_seed(seed)
+    return MatchedEncoder(N, n_layers=n_layers, device=device).to(device)
+
+
+def make_cnn_encoder(d_latent=2, device="cuda"):
+    return nn.Sequential(
+        nn.Conv2d(3, 32, 4, 2, 1), nn.BatchNorm2d(32), nn.GELU(),
+        nn.Conv2d(32, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.GELU(),
+        nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.GELU(),
+        nn.Conv2d(128, 256, 4, 2, 1), nn.BatchNorm2d(256), nn.GELU(),
+        torch.nn.AvgPool2d(4), nn.Flatten(),
+        nn.Linear(256, 256), nn.BatchNorm1d(256), nn.GELU(),
+        nn.Linear(256, d_latent),
+    ).to(device)

JEPA/lejepa-identifiability/experiments/lejepa_id/reacher.py

@@ -0,0 +1,95 @@
+"""DMC Reacher rendering and dataset utilities."""
+
+import os
+os.environ.setdefault("MUJOCO_GL", "egl")
+
+import numpy as np
+import torch
+from dm_control import suite
+from tqdm import tqdm
+
+from .data import ou_augment
+
+
+def make_env():
+    return suite.load(domain_name="reacher", task_name="hard")
+
+
+def render_at(env, qpos, target, height=64, width=64):
+    """Set joint angles and render → (3, H, W) float32 in [0, 1]."""
+    env.physics.data.qpos[:2] = qpos
+    env.physics.data.qvel[:] = 0
+    env.physics.named.model.geom_pos['target', :2] = target
+    env.physics.forward()
+    rgb = env.physics.render(height=height, width=width, camera_id=0)
+    return rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
+
+
+def render_batch(env, qpos_batch, target, height=64, width=64):
+    """Render batch → (N, 3, H, W)."""
+    N = len(qpos_batch)
+    imgs = np.empty((N, 3, height, width), dtype=np.float32)
+    for i in tqdm(range(N), desc="Rendering"):
+        imgs[i] = render_at(env, qpos_batch[i], target, height, width)
+    return imgs
+
+
+def generate_ou_image_pairs(env, N, rho, target, seed=9999):
+    """
+    Sample OU latent pairs, render both → (img_t, img_tp1, z_t, z_tp1).
+
+    Uses the same OU process as the rest of the repo but renders through MuJoCo.
+    """
+    rng = np.random.default_rng(seed)
+    z_t = rng.standard_normal((N, 2)).astype(np.float32)
+    eps = rng.standard_normal((N, 2)).astype(np.float32)
+    z_tp1 = rho * z_t + np.sqrt(1 - rho**2) * eps
+
+    print(f"Rendering {2 * N} images (rho={rho})...")
+    img_t = render_batch(env, z_t, target)
+    img_tp1 = render_batch(env, z_tp1, target)
+    return img_t, img_tp1, z_t, z_tp1
+
+
+def normalize_images(img_t, img_tp1, img_eval=None):
+    """Per-channel mean/std normalization. Returns normalized arrays + stats."""
+    mean = img_t.mean(axis=(0, 2, 3), keepdims=True)
+    std = img_t.std(axis=(0, 2, 3), keepdims=True) + 1e-6
+    img_t = (img_t - mean) / std
+    img_tp1 = (img_tp1 - mean) / std
+    if img_eval is not None:
+        img_eval = (img_eval - mean) / std
+        return img_t, img_tp1, img_eval, mean, std
+    return img_t, img_tp1, mean, std
+
+
+def solve_ik_grid(env, target, n_grid=200):
+    """Find joint angles that place fingertip at target via grid search."""
+    best_dist, best_qpos = np.inf, None
+    for q0 in np.linspace(-np.pi, np.pi, n_grid):
+        for q1 in np.linspace(-np.pi, np.pi, n_grid):
+            env.physics.data.qpos[:2] = [q0, q1]
+            env.physics.named.model.geom_pos['target', :2] = target
+            env.physics.forward()
+            tip = env.physics.named.data.geom_xpos['finger'][:2]
+            d = np.linalg.norm(tip - target)
+            if d < best_dist:
+                best_dist = d
+                best_qpos = np.array([q0, q1])
+    return best_qpos, best_dist
+
+
+class ReacherOUDataset(torch.utils.data.Dataset):
+    """Prerendered OU image pairs with ground-truth latents."""
+
+    def __init__(self, img_t, img_tp1, z_t, z_tp1):
+        self.img_t = torch.from_numpy(img_t)
+        self.img_tp1 = torch.from_numpy(img_tp1)
+        self.z_t = torch.from_numpy(z_t)
+        self.z_tp1 = torch.from_numpy(z_tp1)
+
+    def __len__(self):
+        return len(self.img_t)
+
+    def __getitem__(self, i):
+        return self.img_t[i], self.img_tp1[i], self.z_t[i], self.z_tp1[i]

JEPA/lejepa-identifiability/experiments/prerender.py

@@ -0,0 +1,279 @@
+"""
+Pre-render all Reacher datasets to disk.
+
+Three modes:
+  eval  — 10k Gaussian samples, rendered once, shared by all runs
+  ou    — 100k OU pairs for a given rho
+  traj  — 100k pairs subsampled from LeWM trajectories at a given delta
+
+Usage:
+    python prerender.py eval
+    python prerender.py ou   --rho 0.95
+    python prerender.py traj --delta 16 --h5_path data/reacher.h5
+
+Saves images as uint8 (3, 64, 64) to keep disk usage ~1.2 GB per 100k images.
+Normalization stats computed and saved; applied at training time.
+"""
+
+import os
+os.environ.setdefault("MUJOCO_GL", "egl")
+
+import argparse
+import json
+import numpy as np
+from pathlib import Path
+from scipy.stats import pearsonr, shapiro, skew, kurtosis
+from tqdm import tqdm
+from dm_control import suite
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# RENDERING
+# ═════════════════════════════════════════════════════════════════════════════
+
+TARGET = np.array([0.1, 0.1])
+IMG_SIZE = 64
+
+
+def make_env():
+    return suite.load(domain_name="reacher", task_name="hard")
+
+
+def render_at(env, qpos, height=IMG_SIZE, width=IMG_SIZE):
+    """Render → (3, H, W) uint8."""
+    env.physics.data.qpos[:2] = qpos
+    env.physics.data.qvel[:] = 0
+    env.physics.named.model.geom_pos['target', :2] = TARGET
+    env.physics.forward()
+    rgb = env.physics.render(height=height, width=width, camera_id=0)
+    return rgb.transpose(2, 0, 1)  # uint8, (3, H, W)
+
+
+def render_batch(env, qpos_batch):
+    """Render → (N, 3, H, W) uint8."""
+    N = len(qpos_batch)
+    imgs = np.empty((N, 3, IMG_SIZE, IMG_SIZE), dtype=np.uint8)
+    for i in tqdm(range(N), desc="Rendering"):
+        imgs[i] = render_at(env, qpos_batch[i])
+    return imgs
+
+
+def compute_norm_stats(imgs_uint8):
+    """Compute per-channel mean/std from uint8 images. Returns float32 arrays."""
+    imgs = imgs_uint8.astype(np.float32) / 255.0
+    mean = imgs.mean(axis=(0, 2, 3))   # (3,)
+    std = imgs.std(axis=(0, 2, 3))     # (3,)
+    return mean.astype(np.float32), std.astype(np.float32)
+
+
+def save_dataset(out_dir, z_t, z_tp1, img_t, img_tp1, meta):
+    """Save arrays + metadata to directory."""
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    np.save(out_dir / "z_t.npy", z_t)
+    np.save(out_dir / "z_tp1.npy", z_tp1)
+    np.save(out_dir / "img_t.npy", img_t)
+    np.save(out_dir / "img_tp1.npy", img_tp1)
+
+    # Norm stats from img_t
+    mean, std = compute_norm_stats(img_t)
+    np.save(out_dir / "img_mean.npy", mean)
+    np.save(out_dir / "img_std.npy", std)
+
+    meta["img_mean"] = mean.tolist()
+    meta["img_std"] = std.tolist()
+    with open(out_dir / "meta.json", "w") as f:
+        json.dump(meta, f, indent=2)
+
+    size_gb = sum(
+        (out_dir / fn).stat().st_size for fn in
+        ["img_t.npy", "img_tp1.npy", "z_t.npy", "z_tp1.npy"]
+    ) / 1e9
+    print(f"  Saved to {out_dir} ({size_gb:.2f} GB)")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# EVAL
+# ═════════════════════════════════════════════════════════════════════════════
+
+def prerender_eval(args):
+    """10k i.i.d. Gaussian samples + rendered images."""
+    out_dir = Path(args.data_root) / "eval"
+    if (out_dir / "img.npy").exists() and not args.force:
+        print(f"Eval data already exists at {out_dir}, skipping (use --force)")
+        return
+
+    rng = np.random.default_rng(args.eval_seed)
+    z = rng.standard_normal((args.n_eval, 2)).astype(np.float32)
+
+    env = make_env()
+    print(f"Rendering {args.n_eval} eval images...")
+    imgs = render_batch(env, z)
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    np.save(out_dir / "z.npy", z)
+    np.save(out_dir / "img.npy", imgs)
+
+    mean, std = compute_norm_stats(imgs)
+    np.save(out_dir / "img_mean.npy", mean)
+    np.save(out_dir / "img_std.npy", std)
+
+    with open(out_dir / "meta.json", "w") as f:
+        json.dump({"n_eval": args.n_eval, "seed": args.eval_seed,
+                    "img_mean": mean.tolist(), "img_std": std.tolist()}, f, indent=2)
+    print(f"  Saved to {out_dir}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# OU
+# ═════════════════════════════════════════════════════════════════════════════
+
+def prerender_ou(args):
+    """100k OU pairs for a given rho."""
+    rho = args.rho
+    out_dir = Path(args.data_root) / "ou" / f"rho={rho:.2f}"
+    if (out_dir / "img_t.npy").exists() and not args.force:
+        print(f"OU data for rho={rho} already exists, skipping (use --force)")
+        return
+
+    N = args.n_train
+    rng = np.random.default_rng(args.render_seed)
+    z_t = rng.standard_normal((N, 2)).astype(np.float32)
+    eps = rng.standard_normal((N, 2)).astype(np.float32)
+    z_tp1 = rho * z_t + np.sqrt(1 - rho**2) * eps
+
+    env = make_env()
+    print(f"OU rho={rho}: rendering {2 * N} images...")
+    img_t = render_batch(env, z_t)
+    img_tp1 = render_batch(env, z_tp1)
+
+    meta = {"type": "ou", "rho": rho, "n": N, "seed": args.render_seed}
+    save_dataset(out_dir, z_t, z_tp1, img_t, img_tp1, meta)
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# TRAJECTORY
+# ═════════════════════════════════════════════════════════════════════════════
+
+def load_episodes(h5_path):
+    """Load qpos grouped by episode → (n_episodes, T, 2)."""
+    import h5py
+    with h5py.File(h5_path, "r") as f:
+        qpos = np.array(f["qpos"])
+        ep_len = np.array(f["ep_len"])
+    T = ep_len[0]
+    assert (ep_len == T).all(), f"Non-uniform episode lengths"
+    episodes = qpos.reshape(len(ep_len), T, 2)
+    print(f"Loaded {len(episodes)} episodes, {T} steps each")
+    return episodes
+
+
+def subsample_pairs(episodes, delta, n_per_episode, seed):
+    """Sample n_per_episode (t, t+delta) pairs from each episode."""
+    rng = np.random.default_rng(seed)
+    n_ep, T, d = episodes.shape
+    max_start = T - delta
+    z_t_list, z_tp1_list = [], []
+    for ep in episodes:
+        starts = rng.choice(max_start, size=n_per_episode, replace=False)
+        z_t_list.append(ep[starts])
+        z_tp1_list.append(ep[starts + delta])
+    return (np.concatenate(z_t_list).astype(np.float32),
+            np.concatenate(z_tp1_list).astype(np.float32))
+
+
+def traj_diagnostics(episodes, delta):
+    """Compute autocorrelation + normality stats."""
+    n_ep, T, d = episodes.shape
+    ms = T - delta
+    z_t = episodes[:, :ms].reshape(-1, d)
+    z_tp1 = episodes[:, delta:delta+ms].reshape(-1, d)
+
+    diag = {"delta": delta}
+    for i, name in enumerate(["shoulder", "wrist"]):
+        r, _ = pearsonr(z_t[:, i], z_tp1[:, i])
+        diag[f"rho_{name}"] = float(r)
+        diag[f"skew_{name}"] = float(skew(z_t[:, i]))
+        diag[f"kurtosis_{name}"] = float(kurtosis(z_t[:, i]))
+        sub = z_t[np.random.choice(len(z_t), 5000, replace=False), i]
+        _, p = shapiro(sub)
+        diag[f"shapiro_p_{name}"] = float(p)
+    diag["rho_mean"] = (diag["rho_shoulder"] + diag["rho_wrist"]) / 2
+    return diag
+
+
+def prerender_traj(args):
+    """100k pairs subsampled from LeWM trajectories at a given delta."""
+    delta = args.delta
+    out_dir = Path(args.data_root) / "traj" / f"delta={delta}"
+    if (out_dir / "img_t.npy").exists() and not args.force:
+        print(f"Traj data for delta={delta} already exists, skipping")
+        return
+
+    episodes = load_episodes(args.h5_path)
+    n_episodes = len(episodes)
+    n_per_episode = args.n_train // n_episodes
+    N_actual = n_per_episode * n_episodes
+    print(f"delta={delta}: {n_per_episode} pairs/episode × {n_episodes} = {N_actual}")
+
+    # Diagnostics
+    diag = traj_diagnostics(episodes, delta)
+    print(f"  rho: shoulder={diag['rho_shoulder']:.4f}, "
+          f"wrist={diag['rho_wrist']:.4f}")
+    print(f"  skew: {diag['skew_shoulder']:.3f}, {diag['skew_wrist']:.3f}")
+
+    # Subsample
+    z_t, z_tp1 = subsample_pairs(episodes, delta, n_per_episode, args.render_seed)
+
+    # Render
+    env = make_env()
+    print(f"  Rendering {2 * len(z_t)} images...")
+    img_t = render_batch(env, z_t)
+    img_tp1 = render_batch(env, z_tp1)
+
+    meta = {"type": "traj", "delta": delta, "n": len(z_t),
+            "n_per_episode": n_per_episode, "seed": args.render_seed,
+            **diag}
+    save_dataset(out_dir, z_t, z_tp1, img_t, img_tp1, meta)
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# MAIN
+# ═════════════════════════════════════════════════════════════════════════════
+
+def main():
+    p = argparse.ArgumentParser()
+    sub = p.add_subparsers(dest="mode", required=True)
+
+    # Shared
+    for name in ["eval", "ou", "traj"]:
+        sp = sub.add_parser(name)
+        sp.add_argument("--data_root", type=str, default="data/reacher")
+        sp.add_argument("--force", action="store_true")
+        sp.add_argument("--render_seed", type=int, default=9999)
+
+    # eval
+    sub.choices["eval"].add_argument("--n_eval", type=int, default=10000)
+    sub.choices["eval"].add_argument("--eval_seed", type=int, default=8888)
+
+    # ou
+    sub.choices["ou"].add_argument("--rho", type=float, required=True)
+    sub.choices["ou"].add_argument("--n_train", type=int, default=100000)
+
+    # traj
+    sub.choices["traj"].add_argument("--delta", type=int, required=True)
+    sub.choices["traj"].add_argument("--h5_path", type=str, required=True)
+    sub.choices["traj"].add_argument("--n_train", type=int, default=100000)
+
+    args = p.parse_args()
+
+    if args.mode == "eval":
+        prerender_eval(args)
+    elif args.mode == "ou":
+        prerender_ou(args)
+    elif args.mode == "traj":
+        prerender_traj(args)
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/run.py

@@ -0,0 +1,324 @@
+"""
+Unified experiment runner. Loads config YAML, builds mixing + encoder,
+calls engine.train_and_evaluate, saves standardized .pt output.
+
+Usage:
+    python run.py --config configs/2d.yaml --run spiral --seed 1337
+    python run.py --config configs/ablation.yaml --run spiral_lejepa --seed 1337
+    python run.py --config configs/scaling.yaml --N 16 --seed 0
+    python run.py --config configs/grid.yaml --lamb 0.01 --rho 0.9 --seed 0
+"""
+
+import argparse, os, json, yaml
+import torch
+import numpy as np
+
+from lejepa_id.mixing import MIXINGS_2D, make_coupling_mixing
+from lejepa_id.models import make_mlp_encoder, make_matched_encoder
+from lejepa_id.data import sample_latents, ou_augment
+from lejepa_id.metrics import compute_all_metrics, compute_recovery_metrics
+from lejepa_id.engine import train_and_evaluate
+
+
+def _jsonify(obj):
+    """Convert numpy types to Python natives for JSON serialization."""
+    if isinstance(obj, dict):
+        return {k: _jsonify(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_jsonify(v) for v in obj]
+    elif isinstance(obj, (np.floating,)):
+        return float(obj)
+    elif isinstance(obj, (np.integer,)):
+        return int(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    return obj
+
+def build_mixing(mixing_name, N, n_layers=4, seed=1337, device="cuda"):
+    """Build mixing function from name."""
+    if mixing_name in MIXINGS_2D:
+        return MIXINGS_2D[mixing_name]
+    elif mixing_name in ("nvp", "coupling"):
+        return make_coupling_mixing(N, n_layers=n_layers, seed=seed, device=device)
+    else:
+        raise ValueError(f"Unknown mixing: {mixing_name}")
+
+
+def build_encoder(encoder_type, N, hidden=256, n_layers=4, seed=42, device="cuda"):
+    """Build encoder from type string."""
+    if encoder_type == "mlp":
+        return make_mlp_encoder(N, hidden=hidden, device=device)
+    elif encoder_type == "matched":
+        return make_matched_encoder(N, n_layers=n_layers, seed=seed, device=device)
+    else:
+        raise ValueError(f"Unknown encoder: {encoder_type}")
+
+
+def resolve_run_spec(cfg, args):
+    """Resolve the full run specification from config + CLI args.
+    Returns a dict with all parameters needed for one training run."""
+    experiment = cfg["experiment"]
+
+    # Start with config-level defaults
+    spec = {
+        "experiment": experiment,
+        "N": cfg.get("N", 2),
+        "source_dist": cfg.get("source_dist", "gaussian"),
+        "num_eval": cfg.get("num_eval", 10000),
+        "steps": cfg.get("steps", 10000),
+        "lr": cfg.get("lr", 3e-3),
+        "batch_size": cfg.get("batch_size", 256),
+        "rho": cfg.get("rho", 0.95),
+        "lamb": cfg.get("lamb"),
+        "sigma": cfg.get("sigma", 1.0),
+        "source_alpha": cfg.get("source_alpha"),   # NEW
+        "log_every": cfg.get("log_every", 100),
+        "encoder": cfg.get("encoder", "mlp"),
+        "hidden": cfg.get("hidden", 256),
+        "n_layers": cfg.get("n_layers", 4),
+        "mixing": cfg.get("mixing", "spiral"),
+        "mode": cfg.get("mode", "lejepa"),
+        "seed": args.seed,
+    }
+
+    if experiment in ("2d", "ablation"):
+        # Look up run-specific overrides
+        run_name = args.run
+        run_cfg = cfg["runs"][run_name]
+        spec["run_name"] = run_name
+        for k in ("mixing", "encoder", "hidden", "n_layers", "mode", "lamb", "sigma"):
+            if k in run_cfg:
+                spec[k] = run_cfg[k]
+
+    elif experiment == "scaling":
+        N = args.N
+        spec["N"] = N
+        spec["mixing"] = "coupling"
+        if args.mode is not None:
+            spec["mode"] = args.mode
+        # Mode-specific lamb (whiten uses different default)
+        if spec["mode"] == "whiten":
+            spec["lamb"] = cfg.get("lamb_whiten", 0.5)
+        spec["run_name"] = f"N={N}_{spec['mode']}"
+
+    elif experiment == "grid":
+        spec["lamb"] = args.lamb
+        spec["rho"] = args.rho
+        spec["run_name"] = f"lamb={args.lamb:.0e}_rho={args.rho:.2f}"
+
+    elif experiment == "gennorm":
+        if args.alpha is None:
+            raise ValueError("--alpha required for gennorm experiment")
+        spec["source_dist"] = "gennorm"
+        spec["source_alpha"] = args.alpha
+        run_name = args.run
+        run_cfg = cfg["runs"][run_name]
+        for k in ("mixing", "encoder", "hidden", "n_layers", "mode", "lamb", "sigma"):
+            if k in run_cfg:
+                spec[k] = run_cfg[k]
+        spec["run_name"] = f"{run_name}_alpha={args.alpha:g}"
+
+    return spec
+
+
+def run_single(spec, device):
+    """Execute one training run from a resolved spec. Returns result dict."""
+    N = spec["N"]
+    seed = spec["seed"]
+
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    # Build mixing
+    mix_seed = seed
+    n_layers = spec.get("n_layers", 4)
+    mix_fn = build_mixing(spec["mixing"], N, n_layers=n_layers,
+                          seed=mix_seed, device=device)
+
+    # Build encoder (different seed from mixing)
+    enc_seed = seed + 77777
+    encoder = build_encoder(spec["encoder"], N, hidden=spec.get("hidden", 256),
+                            n_layers=n_layers, seed=enc_seed, device=device)
+
+    # Fixed eval set
+    z_eval = sample_latents(spec["num_eval"], N, dist=spec["source_dist"],
+                            device=device, alpha=spec.get("source_alpha"))
+
+    # Train
+    encoder, log = train_and_evaluate(
+        encoder, mix_fn,
+        N=N, rho=spec["rho"], lamb=spec["lamb"], mode=spec["mode"],
+        source_dist=spec["source_dist"],
+        source_alpha=spec.get("source_alpha"),
+        sigma=spec["sigma"],
+        steps=spec["steps"], batch_size=spec["batch_size"], lr=spec["lr"],
+        z_eval=z_eval, log_every=spec["log_every"], device=device,
+    )
+    
+    # Final metrics from 10k eval set
+    encoder.eval()
+    with torch.no_grad():
+        x_eval = mix_fn(z_eval)
+        h_eval = encoder(x_eval)
+        # z_prime = ou_augment(z_eval, spec["rho"], n_views=1).squeeze(0)
+        z_prime = ou_augment(
+            z_eval, spec["rho"], n_views=1,
+            dist=spec["source_dist"],
+            alpha=spec.get("source_alpha")
+        ).squeeze(0)
+        h_prime = encoder(mix_fn(z_prime))
+
+    final_metrics = compute_all_metrics(
+        z_eval, mix_fn(z_eval), h_eval, h_prime, spec["rho"], N,
+    )
+
+    # Fixed-grid evaluation (cross-distribution comparable, only for 2D)
+    if N == 2:
+        with torch.no_grad():
+            g = torch.linspace(-3.0, 3.0, 100, device=device)
+            z_grid = torch.stack(torch.meshgrid(g, g, indexing='ij'), dim=-1).reshape(-1, N)
+            h_grid = encoder(mix_fn(z_grid))
+        final_metrics.update(compute_recovery_metrics(z_grid, h_grid, N, suffix="_grid"))
+
+    # Large scatter data for plotting (only for 2d/ablation)
+    if spec["experiment"] in ("2d", "ablation"):
+        with torch.no_grad():
+            z_plot = sample_latents(100000, N, dist=spec["source_dist"],
+                                    device=device, alpha=spec.get("source_alpha"))
+            x_plot = mix_fn(z_plot)
+            h_chunks = []
+            for i in range(0, len(z_plot), 10000):
+                h_chunks.append(encoder(x_plot[i:i+10000]))
+            h_plot = torch.cat(h_chunks, dim=0)
+        z_np = z_plot.cpu().numpy()
+        x_np = x_plot.cpu().numpy()
+        h_np = h_plot.cpu().numpy()
+    else:
+        z_np, x_np, h_np = None, None, None
+    
+    # JSON-serializable result (scalars + training curves)
+    result = {
+        # Identity
+        "experiment": spec["experiment"],
+        "run_name": spec["run_name"],
+        "mixing": spec["mixing"],
+        "encoder": spec["encoder"],
+        "mode": spec["mode"],
+        "source_dist": spec["source_dist"],
+        "source_alpha": spec.get("source_alpha"),
+        "seed": seed,
+        "N": N,
+        # Hyperparameters
+        "lamb": spec["lamb"],
+        "rho": spec["rho"],
+        "lr": spec["lr"],
+        "steps": spec["steps"],
+        "batch_size": spec["batch_size"],
+        "n_layers": n_layers,
+        "hidden": spec.get("hidden", None),
+        # Final metrics
+        **final_metrics,
+        "final_align": log["align"][-1],
+        "final_sigreg": log["sigreg"][-1],
+        "final_whiten": log["whiten"][-1],
+        "final_loss": log["total"][-1],
+        # Training curves
+        "log": log,
+    }
+
+    # Heavy data (arrays + model) — only saved as .pt for 2d/ablation
+    arrays = {
+        "z": z_np, "x": x_np, "h": h_np,
+        "model_state_dict": encoder.state_dict(),
+    }
+
+    return result, arrays
+
+
+def save_result(result, arrays, out_dir, fname_base, save_pt=False):
+    """Save JSON always; .pt with arrays/model only when requested."""
+    # JSON
+    json_path = os.path.join(out_dir, fname_base + ".json")
+    with open(json_path, "w") as f:
+        json.dump(_jsonify(result), f, indent=2)
+    print(f"Saved {fname_base}.json")
+
+    # .pt (arrays + model) for 2d/ablation scatter plots
+    if save_pt:
+        pt_path = os.path.join(out_dir, fname_base + ".pt")
+        torch.save({**result, **arrays}, pt_path)
+        print(f"Saved {fname_base}.pt")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", type=str, required=True)
+    # Sweep variables (CLI overrides)
+    p.add_argument("--run", type=str, default=None, help="Run name (2d/ablation)")
+    p.add_argument("--seed", type=int, required=True)
+    p.add_argument("--N", type=int, default=None, help="Latent dim (scaling)")
+    p.add_argument("--lamb", type=float, default=None, help="Lambda (grid)")
+    p.add_argument("--rho", type=float, default=None, help="Rho (grid)")
+    p.add_argument("--alpha", type=float, default=None, help="Gennorm shape (gennorm)")
+    p.add_argument("--mode", type=str, default=None,
+               help="Override mode (lejepa/whiten/infonce)")
+    args = p.parse_args()
+
+    with open(args.config) as f:
+        cfg = yaml.safe_load(f)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    spec = resolve_run_spec(cfg, args)
+    out_dir = cfg["out"]
+    os.makedirs(out_dir, exist_ok=True)
+
+    experiment = cfg["experiment"]
+    save_pt = experiment in ("2d", "ablation")
+
+    if experiment == "scaling":
+        # For small N, train K encoders, pick best
+        K = cfg.get("K", 1)
+        # For large N, all converge
+        if spec["N"] > 32:
+            K = 1
+        best_result = None
+        best_arrays = None
+        best_loss = float("inf")
+
+        for k in range(K):
+            spec_k = dict(spec)
+            spec_k["seed"] = spec["seed"] + k * 1000
+            print(f"\n  Encoder {k+1}/{K} (seed={spec_k['seed']})")
+            result, arrays = run_single(spec_k, device)
+            print(f"    R²(h->z)={result['r2_hz']:.4f} "
+                  f"orth={result['orth_err']:.4f} "
+                  f"loss={result['final_loss']:.6f}")
+
+            if result["final_loss"] < best_loss:
+                best_loss = result["final_loss"]
+                best_result = result
+                best_arrays = arrays
+
+        best_result["K"] = K
+        best_result["seed"] = spec["seed"]  # original seed
+        fname = f"{spec['run_name']}_seed={spec['seed']}"
+        save_result(best_result, best_arrays, out_dir, fname, save_pt=False)
+        print(f"  R²(h->z)={best_result['r2_hz']:.4f} orth={best_result['orth_err']:.4f}")
+
+    else:
+        # Single run
+        print(f"\n{'='*50}")
+        print(f"{spec['run_name']} seed={spec['seed']}")
+        print(f"{'='*50}")
+        result, arrays = run_single(spec, device)
+
+        fname = f"{spec['run_name']}_seed={spec['seed']}"
+        save_result(result, arrays, out_dir, fname, save_pt=save_pt)
+        print(f"  R²(z->h)={result['r2_zh']:.4f} R²(h->z)={result['r2_hz']:.4f} "
+              f"orth={result['orth_err']:.4f}")
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/run_reacher.py

@@ -0,0 +1,361 @@
+"""
+Train LeJEPA on prerendered Reacher data.
+
+Loads images + latents from disk (output of prerender.py),
+sweeps lambda × seed. Final model is always used (no K inits, no loss selection).
+
+Works identically for OU and trajectory data — just point --data_dir
+at the right directory.
+
+Usage:
+    python run_reacher.py --config configs/reacher.yaml \
+        --data_dir data/reacher/ou/rho=0.95
+    python run_reacher.py --config configs/reacher.yaml \
+        --data_dir data/reacher/traj/delta=16
+"""
+
+import argparse, os, json, yaml
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from lejepa_id.losses import SIGReg, alignment_loss
+from lejepa_id.models import make_cnn_encoder
+from lejepa_id.metrics import bidirectional_r2
+
+from sklearn.linear_model import LinearRegression
+from scipy.linalg import orthogonal_procrustes
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# DATA LOADING
+# ═════════════════════════════════════════════════════════════════════════════
+
+def load_dataset(data_dir):
+    """Load prerendered (img, z) pairs. Images stored as uint8."""
+    data_dir = str(data_dir)
+    z_t = np.load(os.path.join(data_dir, "z_t.npy"))
+    z_tp1 = np.load(os.path.join(data_dir, "z_tp1.npy"))
+    img_t = np.load(os.path.join(data_dir, "img_t.npy"))
+    img_tp1 = np.load(os.path.join(data_dir, "img_tp1.npy"))
+    mean = np.load(os.path.join(data_dir, "img_mean.npy"))
+    std = np.load(os.path.join(data_dir, "img_std.npy"))
+    with open(os.path.join(data_dir, "meta.json")) as f:
+        meta = json.load(f)
+    print(f"Loaded {len(z_t)} pairs from {data_dir}")
+    return z_t, z_tp1, img_t, img_tp1, mean, std, meta
+
+
+def load_eval(eval_dir):
+    """Load prerendered eval set."""
+    eval_dir = str(eval_dir)
+    z = np.load(os.path.join(eval_dir, "z.npy"))
+    img = np.load(os.path.join(eval_dir, "img.npy"))
+    mean = np.load(os.path.join(eval_dir, "img_mean.npy"))
+    std = np.load(os.path.join(eval_dir, "img_std.npy"))
+    print(f"Loaded {len(z)} eval samples")
+    return z, img, mean, std
+
+
+def normalize_uint8(img_uint8, mean, std):
+    """Convert uint8 → float32 normalized. mean/std are (3,) arrays."""
+    img = img_uint8.astype(np.float32) / 255.0
+    img = (img - mean[None, :, None, None]) / (std[None, :, None, None] + 1e-6)
+    return img
+
+
+class ImageDataset(torch.utils.data.Dataset):
+    """Normalized float32 image pairs + latents."""
+    def __init__(self, img_t, img_tp1, z_t, z_tp1, mean, std):
+        self.img_t = torch.from_numpy(normalize_uint8(img_t, mean, std))
+        self.img_tp1 = torch.from_numpy(normalize_uint8(img_tp1, mean, std))
+        self.z_t = torch.from_numpy(z_t)
+        self.z_tp1 = torch.from_numpy(z_tp1)
+
+    def __len__(self):
+        return len(self.img_t)
+
+    def __getitem__(self, i):
+        return self.img_t[i], self.img_tp1[i], self.z_t[i], self.z_tp1[i]
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# TRAINING
+# ═════════════════════════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def extract_embeddings(encoder, images, device, batch_size=512):
+    encoder.eval()
+    embeds = []
+    for i in range(0, len(images), batch_size):
+        batch = images[i:i+batch_size].to(device)
+        embeds.append(encoder(batch).cpu())
+    return torch.cat(embeds)
+
+
+def train_one(encoder, loader, eval_data, lamb, cfg, device):
+    """Train one encoder. Returns final model (no selection)."""
+    sigreg = SIGReg(n_slices=cfg["n_slices"]).to(device)
+    opt = torch.optim.AdamW(encoder.parameters(), lr=cfg["lr"], weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=cfg["epochs"])
+
+    eval_imgs, eval_z = eval_data
+
+    log = {"align": [], "sigreg": [], "total": [], "z_std": [], "r2_hz": []}
+
+    for epoch in range(cfg["epochs"]):
+        encoder.train()
+        ep = {k: [] for k in ["align", "sigreg", "total", "z_std"]}
+
+        for img_t, img_tp1, _, _ in loader:
+            img_t, img_tp1 = img_t.to(device), img_tp1.to(device)
+            z_t = encoder(img_t)
+            z_tp1 = encoder(img_tp1)
+
+            h = torch.stack([z_t, z_tp1], dim=0)
+            L_align = alignment_loss(h)
+            L_sig = sigreg(h)
+            loss = lamb * L_sig + (1 - lamb) * L_align
+
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(encoder.parameters(), 1.0)
+            opt.step()
+
+            ep["total"].append(loss.item())
+            ep["align"].append(L_align.item())
+            ep["sigreg"].append(L_sig.item())
+            with torch.no_grad():
+                ep["z_std"].append(z_t.std(0).mean().item())
+
+        scheduler.step()
+
+        # Quick eval (on eval subset, for logging only)
+        encoder.eval()
+        h_eval = extract_embeddings(encoder, eval_imgs, device)
+        _, r2_hz = bidirectional_r2(eval_z, h_eval)
+
+        for k in ep:
+            log[k].append(float(np.mean(ep[k])))
+        log["r2_hz"].append(r2_hz)
+
+        if (epoch + 1) % 10 == 0 or epoch == 0:
+            print(f"      epoch {epoch+1:3d}/{cfg['epochs']}  "
+                  f"align={log['align'][-1]:.5f}  "
+                  f"sig={log['sigreg'][-1]:.1f}  "
+                  f"z_std={log['z_std'][-1]:.3f}  "
+                  f"R²={r2_hz:.4f}")
+
+    return log
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# EVALUATION
+# ═════════════════════════════════════════════════════════════════════════════
+
+def final_eval(encoder, train_imgs, train_z, eval_imgs, eval_z, device):
+    """
+    Full eval with proper train/test split.
+    Fit linear regression on train embeddings, score on eval embeddings.
+    """
+    h_train = extract_embeddings(encoder, train_imgs, device).numpy()
+    h_eval = extract_embeddings(encoder, eval_imgs, device).numpy()
+    z_train = train_z.numpy() if isinstance(train_z, torch.Tensor) else train_z
+    z_eval = eval_z.numpy() if isinstance(eval_z, torch.Tensor) else eval_z
+
+    # Overall R² (fit on train, score on test)
+    reg_hz = LinearRegression().fit(h_train, z_train)
+    r2_hz = reg_hz.score(h_eval, z_eval)
+
+    reg_zh = LinearRegression().fit(z_train, h_train)
+    r2_zh = reg_zh.score(z_eval, h_eval)
+
+    # Per-dimension R² (fit on train, score on test)
+    r2_hz_per = []
+    for i in range(z_train.shape[1]):
+        reg_i = LinearRegression().fit(h_train, z_train[:, i])
+        r2_hz_per.append(reg_i.score(h_eval, z_eval[:, i]))
+
+    # Sin/cos diagnostic (fit on train, score on test)
+    z_train_sc = np.column_stack([np.sin(z_train), np.cos(z_train)])
+    z_eval_sc = np.column_stack([np.sin(z_eval), np.cos(z_eval)])
+    reg_sc = LinearRegression().fit(h_train, z_train_sc)
+    r2_sincos = reg_sc.score(h_eval, z_eval_sc)
+
+    # Per-component sin/cos R²
+    sincos_names = ["sin_shoulder", "cos_shoulder", "sin_wrist", "cos_wrist"]
+    r2_sincos_per = {}
+    for i, name in enumerate(sincos_names):
+        reg_i = LinearRegression().fit(h_train, z_train_sc[:, i])
+        r2_sincos_per[name] = reg_i.score(h_eval, z_eval_sc[:, i])
+
+    # Orthogonality error
+    d = min(z_eval.shape[1], h_eval.shape[1])
+    Zt = (z_eval[:, :d] - z_eval[:, :d].mean(0)).copy()
+    Zl = (h_eval[:, :d] - h_eval[:, :d].mean(0)).copy()
+    for Z in [Zt, Zl]:
+        cov = np.cov(Z, rowvar=False)
+        evals, evecs = np.linalg.eigh(cov)
+        evals = np.maximum(evals, 1e-8)
+        W = evecs @ np.diag(1 / np.sqrt(evals)) @ evecs.T
+        Z[:] = Z @ W
+    R, _ = orthogonal_procrustes(Zl, Zt)
+    orth_err = float(np.linalg.norm(Zl @ R - Zt) / np.linalg.norm(Zt))
+
+    return {
+        "r2_zh": r2_zh,
+        "r2_hz": r2_hz,
+        "r2_hz_per_dim": r2_hz_per,
+        "r2_sincos": r2_sincos,
+        "r2_sincos_per": r2_sincos_per,
+        "orth_error": orth_err,
+        "linear_map_W": reg_hz.coef_.T,
+        "linear_map_b": reg_hz.intercept_,
+    }
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# MAIN
+# ═════════════════════════════════════════════════════════════════════════════
+
+def _jsonify(obj):
+    if isinstance(obj, dict):
+        return {k: _jsonify(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_jsonify(v) for v in obj]
+    elif isinstance(obj, (np.floating,)):
+        return float(obj)
+    elif isinstance(obj, (np.integer,)):
+        return int(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    return obj
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", type=str, required=True)
+    p.add_argument("--data_dir", type=str, required=True,
+                   help="Path to prerendered dataset (ou/rho=X or traj/delta=X)")
+    args = p.parse_args()
+
+    with open(args.config) as f:
+        cfg = yaml.safe_load(f)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    out_dir = cfg["out"]
+    os.makedirs(out_dir, exist_ok=True)
+
+    # ── Load data ────────────────────────────────────────────────────────
+    z_t, z_tp1, img_t, img_tp1, train_mean, train_std, data_meta = \
+        load_dataset(args.data_dir)
+
+    eval_dir = os.path.join(cfg["data_root"], "eval")
+    z_eval, img_eval_u8, eval_mean, eval_std = load_eval(eval_dir)
+
+    # Normalize with training stats
+    dataset = ImageDataset(img_t, img_tp1, z_t, z_tp1, train_mean, train_std)
+    eval_imgs = torch.from_numpy(
+        normalize_uint8(img_eval_u8, train_mean, train_std))
+    eval_z = torch.from_numpy(z_eval)
+
+    # Train embeddings for linreg fitting (subsample for speed)
+    n_fit = min(10000, len(dataset))
+    fit_imgs = dataset.img_t[:n_fit]
+    fit_z = dataset.z_t[:n_fit]
+
+    # Fast eval subset for in-training monitoring
+    n_fast = cfg.get("n_eval_fast", 2000)
+    eval_data_fast = (eval_imgs[:n_fast], eval_z[:n_fast])
+
+    loader = torch.utils.data.DataLoader(
+        dataset, batch_size=cfg["batch_size"], shuffle=True,
+        num_workers=4, pin_memory=True, drop_last=True)
+
+    # Dataset label for output paths
+    data_label = os.path.basename(args.data_dir)
+
+    # ── Sweep lambda × seed ──────────────────────────────────────────────
+    all_results = []
+
+    for lamb in cfg["lambs"]:
+        for seed in cfg["seeds"]:
+            run_name = f"{data_label}_lamb={lamb:.0e}_seed={seed}"
+            print(f"\n{'='*60}")
+            print(f"  {run_name}")
+            print(f"{'='*60}")
+
+            torch.manual_seed(seed)
+            np.random.seed(seed)
+
+            encoder = make_cnn_encoder(
+                d_latent=cfg["d_latent"], device=device)
+
+            log = train_one(
+                encoder, loader, eval_data_fast,
+                lamb=lamb, cfg=cfg, device=device)
+
+            # Full eval with train/test split
+            metrics = final_eval(encoder, fit_imgs, fit_z,
+                                 eval_imgs, eval_z, device)
+
+            result = {
+                "experiment": "reacher",
+                "run_name": run_name,
+                "data_dir": args.data_dir,
+                "lamb": lamb,
+                "seed": seed,
+                "d_latent": cfg["d_latent"],
+                # Data meta (exclude 'seed' key to avoid overwriting training seed)
+                **{k: v for k, v in data_meta.items() if k != "seed"},
+                "render_seed": data_meta.get("seed", None),
+                # Metrics
+                **{k: v for k, v in metrics.items()
+                   if not isinstance(v, np.ndarray)},
+                "best_r2_during_training": max(log["r2_hz"]),
+                "final_r2_during_training": log["r2_hz"][-1],
+                "final_loss": log["total"][-1],
+                "final_align": log["align"][-1],
+                "final_sigreg": log["sigreg"][-1],
+                "log": log,
+            }
+            all_results.append(result)
+
+            print(f"  → R²(h→z)={metrics['r2_hz']:.4f}  "
+                  f"orth_err={metrics['orth_error']:.4f}  "
+                  f"R²(sincos)={metrics['r2_sincos']:.4f}")
+            print(f"    per-dim R²: {['%.4f' % r for r in metrics['r2_hz_per_dim']]}")
+            print(f"    sincos:  {metrics['r2_sincos_per']}")
+
+            # Save checkpoint + result
+            run_dir = os.path.join(out_dir, run_name)
+            os.makedirs(run_dir, exist_ok=True)
+            torch.save({
+                "encoder_state_dict": encoder.state_dict(),
+                "train_mean": train_mean,
+                "train_std": train_std,
+                "d_latent": cfg["d_latent"],
+            }, os.path.join(run_dir, "checkpoint.pt"))
+
+            with open(os.path.join(run_dir, "result.json"), "w") as f:
+                json.dump(_jsonify(result), f, indent=2)
+
+    # ── Summary ──────────────────────────────────────────────────────────
+    summary = {r["run_name"]: {k: v for k, v in r.items() if k != "log"}
+               for r in all_results}
+    with open(os.path.join(out_dir, f"summary_{data_label}.json"), "w") as f:
+        json.dump(_jsonify(summary), f, indent=2)
+
+    print(f"\n{'data':>12s} {'lamb':>8s} {'seed':>4s} "
+          f"{'R²(h→z)':>8s} {'R²(sc)':>8s} {'orth_err':>8s}")
+    print("-" * 56)
+    for r in all_results:
+        print(f"{data_label:>12s} {r['lamb']:8.1e} {r['seed']:4d} "
+              f"{r['r2_hz']:8.4f} {r['r2_sincos']:8.4f} "
+              f"{r['orth_error']:8.4f}")
+
+
+if __name__ == "__main__":
+    main()

JEPA/lejepa-identifiability/experiments/slurm/launch_2d.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_2d
+#SBATCH --output=logs/2d_%A_%a.out
+#SBATCH --error=logs/2d_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=16G
+#SBATCH --time=02:00:00
+#SBATCH --array=0-7   # 8 runs; seeds loop inside
+
+RUNS=(spiral_lejepa spiral_whiten banana_lejepa banana_whiten \
+      sinusoid_lejepa sinusoid_whiten nvp_lejepa nvp_whiten)
+SEEDS=(1337 1338 1339)
+
+RUN=${RUNS[$SLURM_ARRAY_TASK_ID]}
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+
+mkdir -p logs
+for SEED in "${SEEDS[@]}"; do
+    echo "${RUN} seed=${SEED}"
+    python run.py --config configs/2d.yaml \
+        --run "${RUN}" --seed "${SEED}"
+done

JEPA/lejepa-identifiability/experiments/slurm/launch_gennorm.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_gennorm
+#SBATCH --output=logs/gennorm_%A_%a.out
+#SBATCH --error=logs/gennorm_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=16G
+#SBATCH --time=03:00:00
+#SBATCH --array=0-71   # 8 runs x 9 alphas; seeds loop inside
+
+RUNS=(spiral_lejepa spiral_whiten banana_lejepa banana_whiten \
+      sinusoid_lejepa sinusoid_whiten nvp_lejepa nvp_whiten)
+ALPHAS=(0.125 0.25 0.5 1.0 2.0 4.0 8.0 16.0 32.0)
+SEEDS=(1337 1338 1339)
+
+N_ALPHAS=${#ALPHAS[@]}
+RUN_IDX=$(( SLURM_ARRAY_TASK_ID / N_ALPHAS ))
+ALPHA_IDX=$(( SLURM_ARRAY_TASK_ID % N_ALPHAS ))
+RUN=${RUNS[$RUN_IDX]}
+ALPHA=${ALPHAS[$ALPHA_IDX]}
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+
+mkdir -p logs
+for SEED in "${SEEDS[@]}"; do
+    echo "${RUN} alpha=${ALPHA} seed=${SEED}"
+    python run.py --config configs/gennorm.yaml \
+        --run "${RUN}" --alpha "${ALPHA}" --seed "${SEED}"
+done

JEPA/lejepa-identifiability/experiments/slurm/launch_grid.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_grid
+#SBATCH --output=logs/grid_%A_%a.out
+#SBATCH --error=logs/grid_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=16G
+#SBATCH --time=12:00:00
+#SBATCH --array=0-8   # 9 lambda values; rhos and seeds loop inside
+
+LAMBS=(1e-6 1e-5 1e-4 1e-3 5e-3 1e-2 5e-2 1e-1 5e-1)
+RHOS=(0.3 0.5 0.7 0.8 0.9 0.95 0.99)
+SEEDS=(0 1 2)
+
+LAMB=${LAMBS[$SLURM_ARRAY_TASK_ID]}
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+
+mkdir -p logs
+for RHO in "${RHOS[@]}"; do
+    for SEED in "${SEEDS[@]}"; do
+        echo "lamb=${LAMB} rho=${RHO} seed=${SEED}"
+        python run.py --config configs/grid.yaml \
+            --lamb "${LAMB}" --rho "${RHO}" --seed "${SEED}"
+    done
+done

JEPA/lejepa-identifiability/experiments/slurm/launch_reacher_ou.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_ou
+#SBATCH --output=logs/ou_%A_%a.out
+#SBATCH --error=logs/ou_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=32G
+#SBATCH --time=24:00:00
+#SBATCH --array=0-6
+
+# Each task: prerender eval (skipped if exists) + 200k images for one rho,
+# then train 4 lambdas × 3 seeds × 3 inits = 36 training runs
+
+RHOS=(0.3 0.5 0.7 0.8 0.9 0.95 0.99)
+RHO_RAW=${RHOS[$SLURM_ARRAY_TASK_ID]}
+RHO=$(printf "%.2f" $RHO_RAW)
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+export MUJOCO_GL=egl
+mkdir -p logs
+
+echo "Node: $(hostname) | rho=${RHO} | Start: $(date)"
+
+# Step 1: Prerender (eval + this rho)
+python prerender.py eval
+python prerender.py ou --rho "${RHO}"
+
+# Step 2: Train
+python run_reacher.py --config configs/reacher.yaml \
+    --data_dir "data/reacher/ou/rho=${RHO}"
+
+echo "Done: $(date)"

JEPA/lejepa-identifiability/experiments/slurm/launch_reacher_traj.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_traj
+#SBATCH --output=logs/traj_%A_%a.out
+#SBATCH --error=logs/traj_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=32G
+#SBATCH --time=24:00:00
+#SBATCH --array=0-6
+
+# Each task: prerender eval (skipped if exists) + 200k images for one delta,
+# then train 4 lambdas × 3 seeds × 3 inits = 36 training runs
+
+DELTAS=(1 2 4 8 16 32 64)
+DELTA=${DELTAS[$SLURM_ARRAY_TASK_ID]}
+
+H5_PATH="data/reacher.h5"
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+export MUJOCO_GL=egl
+mkdir -p logs
+
+echo "Node: $(hostname) | delta=${DELTA} | Start: $(date)"
+
+# Step 1: Prerender (eval + this delta)
+python prerender.py eval
+python prerender.py traj --delta "${DELTA}" --h5_path "${H5_PATH}"
+
+# Step 2: Train
+python run_reacher.py --config configs/reacher.yaml \
+    --data_dir "data/reacher/traj/delta=${DELTA}"
+
+echo "Done: $(date)"

JEPA/lejepa-identifiability/experiments/slurm/launch_scaling.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+#SBATCH --job-name=lejepa_scale
+#SBATCH --output=logs/scale_%A_%a.out
+#SBATCH --error=logs/scale_%A_%a.err
+#SBATCH --partition=gpuq
+#SBATCH --qos=slow_nice
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=16G
+#SBATCH --time=12:00:00     # ← bumped from 6h: 3 modes × 5 seeds = 15 runs per N
+#SBATCH --array=0-9   # 10 dims
+
+DIMS=(2 4 8 16 32 64 128 256 512 1024)
+SEEDS=(0 1 2 3 4)
+MODES=(lejepa whiten infonce)
+
+N=${DIMS[$SLURM_ARRAY_TASK_ID]}
+
+eval "$(conda shell.bash hook)"
+conda activate pytorch
+
+mkdir -p logs
+
+for MODE in "${MODES[@]}"; do
+    for SEED in "${SEEDS[@]}"; do
+        echo "N=${N} seed=${SEED} mode=${MODE}"
+        python -u run.py --config configs/scaling.yaml \
+            --N "${N}" --seed "${SEED}" --mode "${MODE}"
+    done
+done

JEPA/lejepa-identifiability/lean/LeJEPA.lean

@@ -0,0 +1,65 @@
+import LeJEPA.Hermite
+import LeJEPA.Uniqueness
+import LeJEPA.Approx
+import LeJEPA.Dirichlet
+import LeJEPA.Planning
+
+/-!
+# LeJEPA Identifiability: Formal Verification in Lean 4
+
+Comprehensive formalization of the theoretical results in the paper.
+
+## Files
+
+- **`LeJEPA.Hermite`**: Main theorem (Theorem 4.1) via Hermite
+  polynomial spectral decomposition and the correlation bound.
+
+- **`LeJEPA.Uniqueness`**: Converse direction (Theorem 4.2), that
+  the Gaussian is the unique latent distribution yielding linear
+  identifiability under the Sturm–Liouville operator.
+
+- **`LeJEPA.Approx`**: Approximate identifiability bound
+  (Theorem 4.3) with D + (ε + D)² recovery error.
+
+- **`LeJEPA.Dirichlet`**: Alternative proof (Appendix D) via
+  Dirichlet energy, AM-GM / Jensen, and Mazur–Ulam.
+
+- **`LeJEPA.Planning`**: Planning equivalence corollary
+  (Corollary 4.5): under orthogonal identifiability, expected
+  costs, optimal values, and optimal plans coincide between the
+  learned latent and the true latent for any O(n)-invariant cost.
+
+## Verification Summary
+
+  | Component                            | Status      |
+  |--------------------------------------|-------------|
+  | Hermite basis & completeness         | axiomatized |
+  | Contraction lemma (ρᵈ decay)         | axiomatized |
+  | Mehler's formula                     | axiomatized |
+  | Correlation bound ≤ ρ                | VERIFIED    |
+  | Equality ⟺ w₁ = 1 (linearity)       | VERIFIED    |
+  | Loss lower bound 2(1−ρ)n             | VERIFIED    |
+  | Hermite theorem assembly h = Qz      | VERIFIED    |
+  | Affine eigenfunction → affine score  | VERIFIED    |
+  | Affine score → Gaussian density      | axiomatized |
+  | Gaussian → Hermite eigenfunctions    | axiomatized |
+  | Gaussian uniqueness biconditional    | VERIFIED    |
+  | Polar decomposition                  | axiomatized |
+  | Cross-degree Hermite orthogonality   | axiomatized |
+  | Spectral gap → W_nl ≤ D              | VERIFIED    |
+  | ‖M − Q‖²_F bound                     | VERIFIED    |
+  | Pythagorean decomposition            | axiomatized |
+  | Bound monotonicity                   | VERIFIED    |
+  | Approximate bound assembly           | VERIFIED    |
+  | Exact recovery (δ = ε = 0)           | VERIFIED    |
+  | AM-GM / Jensen                       | axiomatized |
+  | Mazur–Ulam                           | axiomatized |
+  | Orthogonal Jacobian → Lipschitz      | VERIFIED    |
+  | Bilipschitz → global isometry        | VERIFIED    |
+  | Dirichlet theorem assembly h = Qz    | VERIFIED    |
+  | Trajectory pushforward (stage/term)  | axiomatized |
+  | Per-step stage / terminal equiv.     | VERIFIED    |
+  | Planning equivalence (main step)     | VERIFIED    |
+  | Minimizer equivalence                | VERIFIED    |
+  | Value equivalence                    | VERIFIED    |
+-/

JEPA/lejepa-identifiability/lean/LeJEPA/Approx.lean

@@ -0,0 +1,187 @@
+import Mathlib
+
+/-!
+# Part C — Approximate Identifiability (Proposition 4.3)
+
+  Under approximate alignment (gap δ) and approximate covariance
+  (error ε), the recovery error satisfies:
+
+    𝔼[‖h(z) − Qz‖²] ≤ D + (ε + D)²
+
+  where D = δ/(2ρ(1−ρ)) is the alignment gap normalized by the
+  spectral gap between Hermite degrees 1 and 2.
+
+  When δ = ε = 0 this recovers Theorem 4.1: h(z) = Qz a.e.
+
+  ## Verification status
+
+  | Component                          | Status      |
+  |------------------------------------|-------------|
+  | Spectral gap positivity            | VERIFIED    |
+  | W_nl ≤ D from gap inequality       | VERIFIED    |
+  | Polar decomposition ‖M−Q‖ bound    | axiomatized |
+  | Cross-degree Hermite orthogonality  | axiomatized |
+  | Linear deviation ‖M−Q‖² bound      | VERIFIED    |
+  | Pythagorean decomposition           | axiomatized |
+  | Bound monotonicity in W_nl          | VERIFIED    |
+  | Full bound assembly                 | VERIFIED    |
+  | Exact recovery (δ=ε=0 ⟹ error=0)  | VERIFIED    |
+-/
+
+noncomputable section
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 1: SPECTRAL GAP CONTROLS NONLINEAR ENERGY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- The spectral gap ρ(1−ρ) is positive for 0 < ρ < 1. -/
+theorem spectral_gap_pos (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    0 < ρ * (1 - ρ) := by
+  apply mul_pos hρ0; linarith
+
+/-- 2ρ(1−ρ) is positive. -/
+theorem two_spectral_gap_pos (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    0 < 2 * ρ * (1 - ρ) := by
+  have : 0 < ρ * (1 - ρ) := spectral_gap_pos ρ hρ0 hρ1
+  linarith
+
+/-- **Nonlinear energy bound** (VERIFIED): from the spectral gap
+    inequality δ ≥ 2ρ(1−ρ) W_nl, we get W_nl ≤ D = δ/(2ρ(1−ρ)). -/
+theorem nonlinear_energy_le_D
+    (ρ δ W_nl : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (_hδ_nonneg : 0 ≤ δ) (_hW_nonneg : 0 ≤ W_nl)
+    (hgap : δ ≥ 2 * ρ * (1 - ρ) * W_nl) :
+    W_nl ≤ δ / (2 * ρ * (1 - ρ)) := by
+  have hsgap : (0 : ℝ) < 2 * ρ * (1 - ρ) := two_spectral_gap_pos ρ hρ0 hρ1
+  rw [le_div_iff₀ hsgap]
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 2: LINEAR PART DEVIATION
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Polar decomposition bound** (axiomatized): ‖M − Q‖_F ≤ ε + W_nl.
+    Combines polar decomposition, |σᵢ−1| ≤ |σᵢ²−1|, covariance
+    decomposition Cov(h) = MM^T + N, and triangle inequality. -/
+axiom polar_bound_axiom
+    (M_Q_norm ε W_nl : ℝ)
+    (hε : 0 ≤ ε) (hW : 0 ≤ W_nl) :
+    M_Q_norm ≤ ε + W_nl →
+    M_Q_norm ≤ ε + W_nl
+
+/-- **Linear deviation squared** (VERIFIED): ‖M−Q‖ ≤ ε+W_nl implies
+    ‖M−Q‖² ≤ (ε+W_nl)². -/
+theorem linear_deviation_sq_bound
+    (M_Q_norm ε W_nl : ℝ)
+    (hMQ_nonneg : 0 ≤ M_Q_norm)
+    (hε : 0 ≤ ε) (hW : 0 ≤ W_nl)
+    (hbound : M_Q_norm ≤ ε + W_nl) :
+    M_Q_norm ^ 2 ≤ (ε + W_nl) ^ 2 := by
+  have h1 : 0 ≤ ε + W_nl := by linarith
+  nlinarith [sq_nonneg (ε + W_nl - M_Q_norm)]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 3: PYTHAGOREAN DECOMPOSITION
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Pythagorean decomposition** (axiomatized): the recovery error
+    splits into linear deviation and nonlinear energy.
+    Requires Hermite orthogonality and z ~ N(0,I). -/
+axiom pythagorean_axiom
+    (total_error M_Q_norm_sq W_nl : ℝ) :
+    total_error = M_Q_norm_sq + W_nl →
+    total_error = M_Q_norm_sq + W_nl
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 4: MONOTONICITY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Monotonicity** (VERIFIED): f(t) = (ε + t)² + t is increasing
+    for t ≥ 0. So W_nl ≤ D implies (ε+W_nl)²+W_nl ≤ (ε+D)²+D. -/
+theorem bound_monotone (ε W_nl D : ℝ)
+    (_hε : 0 ≤ ε) (_hW : 0 ≤ W_nl) (_hD : 0 ≤ D)
+    (hle : W_nl ≤ D) :
+    (ε + W_nl) ^ 2 + W_nl ≤ (ε + D) ^ 2 + D := by
+  have h1 : ε + W_nl ≤ ε + D := by linarith
+  nlinarith [sq_nonneg (ε + D - ε - W_nl)]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- MAIN BOUND ASSEMBLY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Approximate identifiability** (Proposition 4.3, VERIFIED assembly):
+
+    𝔼[‖h(z) − Qz‖²] ≤ D + (ε + D)²
+
+    where D = δ/(2ρ(1−ρ)). -/
+theorem approximate_identifiability
+    (ρ δ ε W_nl M_Q_norm total_error : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hδ : 0 ≤ δ) (hε : 0 ≤ ε)
+    (hW : 0 ≤ W_nl) (hMQ : 0 ≤ M_Q_norm)
+    (hgap : δ ≥ 2 * ρ * (1 - ρ) * W_nl)
+    (hpolar : M_Q_norm ≤ ε + W_nl)
+    (hpythag : total_error = M_Q_norm ^ 2 + W_nl) :
+    total_error ≤ δ / (2 * ρ * (1 - ρ))
+              + (ε + δ / (2 * ρ * (1 - ρ))) ^ 2 := by
+  set D := δ / (2 * ρ * (1 - ρ)) with hD_def
+  have hsgap := two_spectral_gap_pos ρ hρ0 hρ1
+  have hD_nonneg : 0 ≤ D := div_nonneg hδ (le_of_lt hsgap)
+  -- Step 1: W_nl ≤ D
+  have hW_le_D : W_nl ≤ D := nonlinear_energy_le_D ρ δ W_nl hρ0 hρ1 hδ hW hgap
+  -- Step 4: ‖M−Q‖² ≤ (ε + W_nl)²
+  have hMQ_sq : M_Q_norm ^ 2 ≤ (ε + W_nl) ^ 2 :=
+    linear_deviation_sq_bound M_Q_norm ε W_nl hMQ hε hW hpolar
+  -- Step 3 + 4: total_error ≤ (ε + W_nl)² + W_nl
+  have h_inter : total_error ≤ (ε + W_nl) ^ 2 + W_nl := by
+    rw [hpythag]; linarith
+  -- Step 5: monotonicity
+  have h_mono := bound_monotone ε W_nl D hε hW hD_nonneg hW_le_D
+  -- Combine
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- EXACT RECOVERY AS SPECIAL CASE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Exact recovery** (VERIFIED): setting δ = ε = 0 gives error = 0,
+    recovering Theorem 4.1: h(z) = Qz almost everywhere. -/
+theorem exact_recovery_special_case
+    (ρ W_nl M_Q_norm total_error : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hW : 0 ≤ W_nl) (hMQ : 0 ≤ M_Q_norm)
+    (hgap : (0 : ℝ) ≥ 2 * ρ * (1 - ρ) * W_nl)
+    (hpolar : M_Q_norm ≤ 0 + W_nl)
+    (hpythag : total_error = M_Q_norm ^ 2 + W_nl)
+    (_htotal_nonneg : 0 ≤ total_error) :
+    total_error = 0 := by
+  -- δ = 0 forces W_nl = 0
+  have hsgap := two_spectral_gap_pos ρ hρ0 hρ1
+  have hW_zero : W_nl = 0 := by nlinarith
+  -- W_nl = 0 and ε = 0 force ‖M − Q‖ = 0
+  have hMQ_zero : M_Q_norm = 0 := by
+    have : M_Q_norm ≤ 0 := by linarith [hpolar, hW_zero]
+    linarith
+  -- Total error = 0² + 0 = 0
+  rw [hpythag, hMQ_zero, hW_zero]; ring
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- BOUND STRUCTURE ANALYSIS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **First-order approximation** (VERIFIED): when ε + D ≤ 1,
+    the quadratic term (ε+D)² ≤ ε+D, so the bound ≤ 2D + ε. -/
+theorem bound_small_perturbation (ε D : ℝ)
+    (hε : 0 ≤ ε) (hD : 0 ≤ D) (hsmall : ε + D ≤ 1) :
+    D + (ε + D) ^ 2 ≤ D + ε + D := by
+  have h1 : 0 ≤ ε + D := by linarith
+  nlinarith [sq_nonneg (1 - (ε + D))]
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/Dirichlet.lean

@@ -0,0 +1,227 @@
+import Mathlib.Analysis.InnerProductSpace.Basic
+import Mathlib.Analysis.InnerProductSpace.PiL2
+import Mathlib.Analysis.Normed.Module.Basic
+import Mathlib.Analysis.Calculus.MeanValue
+import Mathlib.Analysis.SpecialFunctions.Pow.Real
+import Mathlib.Analysis.SpecialFunctions.ExpDeriv
+import Mathlib.LinearAlgebra.Matrix.NonsingularInverse
+import Mathlib.LinearAlgebra.Matrix.Determinant.Basic
+import Mathlib.Topology.MetricSpace.Isometry
+import Mathlib.Topology.MetricSpace.Lipschitz
+
+/-!
+# Part B — Alternative Proof via Dirichlet Energy (Appendix C)
+
+  Any C¹ diffeomorphism h : ℝⁿ → ℝⁿ that preserves the standard
+  Gaussian measure and minimizes the Dirichlet energy 𝔼[‖Jₕ‖²_F]
+  must be a linear orthogonal map h(z) = Uz.
+
+  ## Proof sketch
+
+  Steps 1–2 (reduction to Dirichlet energy and the log-determinant
+  lemma) involve measure-theoretic integration. We axiomatize their
+  conclusions.
+
+  Steps 3–5 are verified:
+    Step 3: AM-GM + Jensen → 𝓙(h) ≥ n                 (axiomatized)
+    Step 4: Equality forces Jₕ orthogonal everywhere   (axiomatized)
+    Step 5: Orthogonal Jacobian → global isometry →
+            Mazur–Ulam → linear                        (VERIFIED)
+
+  ## Verification status
+
+  | Component                        | Status      |
+  |----------------------------------|-------------|
+  | AM-GM for singular values        | axiomatized |
+  | Jensen for log-determinant       | axiomatized |
+  | Mazur–Ulam theorem               | axiomatized |
+  | Norm-preserving CLM → isometry   | VERIFIED    |
+  | Orthogonal Jacobian → Lipschitz  | VERIFIED    |
+  | Bilipschitz → global isometry    | VERIFIED    |
+  | h(0)=0 → b=0 → linear isometry  | VERIFIED    |
+  | Full theorem assembly            | VERIFIED    |
+-/
+
+open scoped Matrix BigOperators
+open Matrix
+
+noncomputable section
+
+variable {n : ℕ}
+
+/-- The type we work with: ℝⁿ as a Euclidean space. -/
+private abbrev E (n : ℕ) := EuclideanSpace ℝ (Fin n)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED KNOWN RESULTS
+-- ═══════════════════════════════════════════════════════════════
+
+/-!
+These are standard results available in Mathlib but requiring
+nontrivial plumbing to connect to our specific statement forms.
+-/
+
+/-- **AM-GM inequality**: arithmetic mean of nonneg reals ≥ geometric
+    mean. Special case of `Real.geom_mean_le_arith_mean_weighted`
+    in `Mathlib.Analysis.MeanInequalities` with uniform weights. -/
+axiom amgm_sum_ge_prod_pow {m : ℕ} (a : Fin m → ℝ)
+    (ha : ∀ i, 0 ≤ a i) :
+    (∑ i : Fin m, a i) / m ≥ (∏ i : Fin m, a i) ^ ((1 : ℝ) / m)
+
+/-- **Jensen's inequality** applied to strictly convex exp:
+    mean of exp(cxᵢ) ≥ 1 when xᵢ sum to zero. Follows from
+    `StrictConvexOn` of `Real.exp` and the weighted AM-GM. -/
+axiom exp_mean_ge_mean_exp {m : ℕ}
+    (f : Fin m → ℝ) (hsum : ∑ i : Fin m, f i = 0) :
+    (∑ i : Fin m, Real.exp ((2 : ℝ) / m * f i)) / m ≥ 1
+
+/-- **Mazur–Ulam theorem**: every surjective isometry of a real normed
+    space is affine. Available in Mathlib as the combination of
+    `Isometry.right_inv` and affine isometry machinery in
+    `Mathlib.Analysis.Normed.Affine.Isometry`. -/
+axiom mazur_ulam
+    {V : Type*} [NormedAddCommGroup V] [NormedSpace ℝ V]
+    {f : V → V} (hiso : Isometry f) (hsurj : Function.Surjective f) :
+    ∃ (A : V →ₗ[ℝ] V) (b : V), ∀ x, f x = A x + b
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- DIFFEOMORPHISM STRUCTURE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A smooth map h : ℝⁿ → ℝⁿ with its Jacobian, modeling a C¹
+    diffeomorphism that preserves the standard Gaussian. -/
+structure GaussianDiffeo (n : ℕ) where
+  /-- The map itself -/
+  toFun : E n → E n
+  /-- The Jacobian at each point, as a continuous linear map -/
+  jacobian : E n → (E n →L[ℝ] E n)
+  /-- h is differentiable with the given Jacobian -/
+  hasFDeriv : ∀ z, HasFDerivAt toFun (jacobian z) z
+  /-- h is a homeomorphism (hence bijective) -/
+  isHomeo : (E n) ≃ₜ (E n)
+  /-- The homeomorphism agrees with toFun -/
+  homeo_eq : ∀ z, isHomeo z = toFun z
+  /-- Inverse differentiability from the **inverse function theorem**
+      (`HasStrictFDerivAt.toOpenPartialHomeomorph` in
+      `Mathlib.Analysis.Calculus.InverseFunctionTheorem.FDeriv`). -/
+  hasFDeriv_inv : ∀ y, HasFDerivAt isHomeo.symm
+    (ContinuousLinearMap.inverse (jacobian (isHomeo.symm y))) y
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: ORTHOGONAL JACOBIAN → GLOBAL ISOMETRY → LINEAR
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A norm-preserving continuous linear map is an isometry. -/
+theorem clm_isometry_of_norm_preserving
+    (L : E n →L[ℝ] E n)
+    (hL : ∀ v, ‖L v‖ = ‖v‖) :
+    Isometry L := by
+  rw [isometry_iff_dist_eq]
+  intro x y
+  simp only [dist_eq_norm, ← map_sub L x y]
+  exact hL (x - y)
+
+/-- **Mean value theorem** (VERIFIED): orthogonal Jacobian everywhere
+    ⟹ h is 1-Lipschitz. By the MVT, ‖h(x)-h(y)‖ ≤ sup ‖Jₕ‖_op · ‖x-y‖,
+    and the operator norm of a norm-preserving map is 1. -/
+theorem lipschitz_of_orthogonal_jacobian
+    (h : GaussianDiffeo n)
+    (horth : ∀ z v, ‖h.jacobian z v‖ = ‖v‖) :
+    LipschitzWith 1 h.toFun := by
+  apply lipschitzWith_of_nnnorm_fderiv_le (𝕜 := ℝ)
+  · intro x; exact (h.hasFDeriv x).differentiableAt
+  · intro x
+    have hfderiv : fderiv ℝ h.toFun x = h.jacobian x :=
+      (h.hasFDeriv x).fderiv
+    rw [hfderiv, ContinuousLinearMap.opNNNorm_le_iff]
+    intro y; simp only [one_mul]
+    exact_mod_cast le_of_eq (horth x y)
+
+/-- **Bilipschitz → isometry** (VERIFIED): if both h and h⁻¹ are
+    1-Lipschitz, h is a global isometry. Forward Lipschitz gives
+    dist(hx,hy) ≤ dist(x,y); applying to h⁻¹ gives ≥. -/
+theorem isometry_of_bilipschitz
+    (h : GaussianDiffeo n)
+    (hlip : LipschitzWith 1 h.toFun)
+    (hinvlip : LipschitzWith 1 h.isHomeo.symm) :
+    Isometry h.toFun := by
+  rw [isometry_iff_dist_eq]
+  intro x y
+  apply le_antisymm
+  · -- Forward: dist(hx, hy) ≤ dist(x, y)
+    have hfwd := hlip.dist_le_mul x y
+    simp only [NNReal.coe_one, one_mul] at hfwd; exact hfwd
+  · -- Backward: apply Lipschitz to h⁻¹
+    have hbwd := hinvlip.dist_le_mul (h.toFun x) (h.toFun y)
+    simp only [NNReal.coe_one, one_mul] at hbwd
+    have hx : h.isHomeo.symm (h.toFun x) = x := by
+      rw [← h.homeo_eq]; exact h.isHomeo.symm_apply_apply x
+    have hy : h.isHomeo.symm (h.toFun y) = y := by
+      rw [← h.homeo_eq]; exact h.isHomeo.symm_apply_apply y
+    rw [hx, hy] at hbwd; exact hbwd
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: MAIN THEOREM (APPENDIX C)
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **LeJEPA identifiability via Dirichlet energy** (VERIFIED):
+
+    C¹ diffeomorphism + Gaussian-preserving + orthogonal Jacobian
+    ⟹ h(z) = Uz for a linear isometry U ∈ O(n).
+
+    Verified chain:
+    1. Orth. Jacobian → h is 1-Lipschitz    (MVT)
+    2. Orth. inverse  → h⁻¹ is 1-Lipschitz  (IFT + MVT)
+    3. Bilipschitz → global isometry
+    4. Mazur–Ulam → h is affine: h(z) = Az + b
+    5. h(0) = 0 → b = 0
+    6. A preserves norms → A is a LinearIsometry -/
+theorem dirichlet_identifiability
+    (h : GaussianDiffeo n)
+    (horth : ∀ z v, ‖h.jacobian z v‖ = ‖v‖)
+    (horth_inv : ∀ z v,
+      ‖(ContinuousLinearMap.inverse (h.jacobian z)) v‖ = ‖v‖)
+    (hmean : h.toFun 0 = 0) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, h.toFun z = U z := by
+  -- Step 1: h is 1-Lipschitz
+  have hlip := lipschitz_of_orthogonal_jacobian h horth
+  -- Step 2: h⁻¹ is 1-Lipschitz (IFT gives derivative = J⁻¹, also orth.)
+  have hinvlip : LipschitzWith 1 h.isHomeo.symm := by
+    apply lipschitzWith_of_nnnorm_fderiv_le (𝕜 := ℝ)
+    · intro x; exact (h.hasFDeriv_inv x).differentiableAt
+    · intro x
+      have hfderiv : fderiv ℝ h.isHomeo.symm x =
+        (h.jacobian (h.isHomeo.symm x)).inverse :=
+        (h.hasFDeriv_inv x).fderiv
+      rw [hfderiv, ContinuousLinearMap.opNNNorm_le_iff]
+      intro y; simp only [one_mul]
+      exact_mod_cast le_of_eq (horth_inv (h.isHomeo.symm x) y)
+  -- Step 3: h is a global isometry
+  have hiso := isometry_of_bilipschitz h hlip hinvlip
+  -- Step 4: Mazur–Ulam → h(z) = Az + b
+  have hsurj : Function.Surjective h.toFun := by
+    intro y
+    exact ⟨h.isHomeo.symm y,
+      by rw [← h.homeo_eq]; exact h.isHomeo.apply_symm_apply y⟩
+  obtain ⟨A, b, hab⟩ := mazur_ulam hiso hsurj
+  -- Step 5: b = 0 from h(0) = 0
+  have hb : b = 0 := by
+    have h0 := hab 0; simp [map_zero] at h0
+    rw [hmean] at h0; exact h0.symm
+  -- h(z) = Az for all z
+  have hab' : ∀ z, h.toFun z = A z := by
+    intro z; have := hab z; rw [hb, add_zero] at this; exact this
+  -- Step 6: A preserves norms → LinearIsometry
+  have hA_norm : ∀ v, ‖A v‖ = ‖v‖ := by
+    intro v
+    have hv := hiso.dist_eq v 0
+    simp [dist_eq_norm] at hv
+    rw [hab' v, hab' 0, map_zero] at hv
+    simpa using hv
+  exact ⟨⟨A, hA_norm⟩, hab'⟩
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/Hermite.lean

@@ -0,0 +1,270 @@
+import Mathlib.Analysis.InnerProductSpace.PiL2
+import Mathlib.Topology.Algebra.InfiniteSum.Order
+import Mathlib.Topology.Algebra.InfiniteSum.Ring
+
+/-!
+# Part A — Main Theorem via Hermite Polynomials (Theorem 4.1)
+
+  Any measurable h : ℝⁿ → ℝⁿ satisfying Gaussianity h(z) ~ N(0,Iₙ)
+  and minimizing the alignment loss must be h(z) = Uz for U ∈ O(n).
+
+  ## Verification status
+
+  | Component                        | Status      |
+  |----------------------------------|-------------|
+  | Hermite basis & completeness     | axiomatized |
+  | Contraction lemma (ρᵈ decay)     | axiomatized |
+  | Mehler's formula                 | axiomatized |
+  | ρᵈ ≤ ρ for d ≥ 1                | VERIFIED    |
+  | ρᵈ < ρ for d ≥ 2                | VERIFIED    |
+  | Pointwise term bound w_d·ρᵈ≤w_d·ρ| VERIFIED    |
+  | Correlation bound ≤ ρ            | VERIFIED    |
+  | Equality ⟺ w₁ = 1 (linearity)  | VERIFIED    |
+  | Loss lower bound 2(1-ρ)n        | VERIFIED    |
+  | Theorem assembly h = Uz         | VERIFIED    |
+-/
+
+set_option maxHeartbeats 400000
+
+open scoped BigOperators
+
+noncomputable section
+
+abbrev E (n : ℕ) := EuclideanSpace ℝ (Fin n)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- SPECTRAL WEIGHTS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- Spectral weights of a single encoder component in its Hermite
+    expansion. `w d` is the fraction of L²(γₙ) variance at degree d. -/
+structure SpectralWeights where
+  w : ℕ → ℝ
+  nonneg : ∀ d, 0 ≤ w d
+  zero_degree : w 0 = 0
+  summable : Summable w
+  total_variance : ∑' d, w d = 1
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED: HERMITE BASIS & MEHLER
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Mehler's formula** (axiomatized): the spectral correlation
+    series Σ_d w_d · ρᵈ is summable. -/
+axiom mehler_summability
+    (sw : SpectralWeights) (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    Summable (fun d => sw.w d * ρ ^ d)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: POINTWISE BOUNDS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- For 0 < ρ ≤ 1 and d ≥ 1, ρᵈ ≤ ρ. -/
+theorem pow_le_self_of_pos_lt_one (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1)
+    (d : ℕ) (hd : 1 ≤ d) : ρ ^ d ≤ ρ := by
+  calc ρ ^ d ≤ ρ ^ 1 := pow_le_pow_of_le_one (le_of_lt hρ0) hρ1 hd
+       _ = ρ := pow_one ρ
+
+/-- Each term w_d · ρᵈ ≤ w_d · ρ. -/
+theorem spectral_term_le (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1) (d : ℕ) :
+    sw.w d * ρ ^ d ≤ sw.w d * ρ := by
+  match d with
+  | 0 => simp [sw.zero_degree]
+  | d + 1 =>
+    exact mul_le_mul_of_nonneg_left
+      (pow_le_self_of_pos_lt_one ρ hρ0 hρ1 (d + 1)
+        (Nat.succ_le_succ (Nat.zero_le d)))
+      (sw.nonneg (d + 1))
+
+/-- For 0 < ρ < 1 and d ≥ 2, ρᵈ < ρ (strict). -/
+theorem pow_lt_self_of_ge_two (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (d : ℕ) (hd : 2 ≤ d) : ρ ^ d < ρ := by
+  calc ρ ^ d ≤ ρ ^ 2 := pow_le_pow_of_le_one (le_of_lt hρ0) (le_of_lt hρ1) hd
+       _ = ρ * ρ      := by ring
+       _ < ρ * 1      := mul_lt_mul_of_pos_left hρ1 hρ0
+       _ = ρ           := mul_one ρ
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: SUMMABILITY AND TSUM OF UPPER BOUND
+-- ═══════════════════════════════════════════════════════════════
+
+/-- The constant-ρ series fun d ↦ w d * ρ is summable
+    (via Summable.mul_right from Ring.lean). -/
+theorem summable_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
+    Summable (fun d => sw.w d * ρ) :=
+  sw.summable.mul_right ρ
+
+/-- Σ w_d · ρ = (Σ w_d) · ρ = 1 · ρ = ρ
+    (via tsum_mul_right from Ring.lean). -/
+theorem tsum_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
+    ∑' d, sw.w d * ρ = ρ := by
+  rw [tsum_mul_right, sw.total_variance, one_mul]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: CORRELATION BOUND (Lemma 3.3)
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Correlation bound** (VERIFIED): Σ_d w_d ρᵈ ≤ ρ.
+    Uses Summable.tsum_le_tsum (from Order.lean via @[to_additive]). -/
+theorem correlation_le_rho (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hsum : Summable (fun d => sw.w d * ρ ^ d)) :
+    ∑' d, sw.w d * ρ ^ d ≤ ρ := by
+  calc ∑' d, sw.w d * ρ ^ d
+      ≤ ∑' d, sw.w d * ρ :=
+        hsum.tsum_le_tsum
+          (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
+          (summable_spectral_upper sw ρ)
+    _ = ρ := tsum_spectral_upper sw ρ
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: EQUALITY FORCES LINEARITY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Equality characterization** (VERIFIED): if Σ w_d ρᵈ = ρ, then
+    w_d = 0 for all d ≥ 2.
+
+    Strategy: by contradiction. If w_{d₀} > 0 for some d₀ ≥ 2, then
+    w_{d₀}·ρ^{d₀} < w_{d₀}·ρ strictly, while all other terms satisfy ≤.
+    By Summable.tsum_lt_tsum (from Order.lean via @[to_additive]),
+    Σ w_d·ρᵈ < Σ w_d·ρ = ρ, contradicting Σ w_d·ρᵈ = ρ. -/
+theorem equality_forces_degree_one (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hsum : Summable (fun d => sw.w d * ρ ^ d))
+    (heq : ∑' d, sw.w d * ρ ^ d = ρ) :
+    ∀ d, 2 ≤ d → sw.w d = 0 := by
+  by_contra h
+  push_neg at h
+  obtain ⟨d₀, hd₀_ge, hd₀_ne⟩ := h
+  -- w_{d₀} > 0
+  have hwd₀_pos : 0 < sw.w d₀ :=
+    lt_of_le_of_ne (sw.nonneg d₀) (Ne.symm hd₀_ne)
+  -- Strict inequality at d₀: w_{d₀} · ρ^{d₀} < w_{d₀} · ρ
+  have hstrict : sw.w d₀ * ρ ^ d₀ < sw.w d₀ * ρ :=
+    mul_lt_mul_of_pos_left (pow_lt_self_of_ge_two ρ hρ0 hρ1 d₀ hd₀_ge) hwd₀_pos
+  -- By tsum_lt_tsum: one strict + rest ≤ ⟹ strict on tsums
+  have hlt : ∑' d, sw.w d * ρ ^ d < ∑' d, sw.w d * ρ :=
+    hsum.tsum_lt_tsum
+      (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
+      hstrict
+      (summable_spectral_upper sw ρ)
+  -- But Σ w_d·ρᵈ = ρ = Σ w_d·ρ
+  rw [tsum_spectral_upper, heq] at hlt
+  exact lt_irrefl ρ hlt
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- ENCODER STRUCTURE & LOSS
+-- ═══════════════════════════════════════════════════════════════
+
+variable {n : ℕ}
+
+/-- An encoder h : ℝⁿ → ℝⁿ with its Hermite spectral decomposition. -/
+structure HermiteEncoder (n : ℕ) where
+  toFun : E n → E n
+  spectrum : Fin n → SpectralWeights
+  correlation : Fin n → ℝ
+
+/-- The alignment loss: 𝓛(h) = 2n − 2 Σᵢ corr_i. -/
+def alignmentLoss (enc : HermiteEncoder n) : ℝ :=
+  2 * n - 2 * ∑ i : Fin n, enc.correlation i
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED: BRIDGE LEMMAS
+-- ═══════════════════════════════════════════════════════════════
+
+axiom correlation_eq_spectral_sum (enc : HermiteEncoder n) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1) (i : Fin n) :
+    enc.correlation i = ∑' d, (enc.spectrum i).w d * ρ ^ d
+
+axiom linear_of_degree_one (enc : HermiteEncoder n)
+    (hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0) :
+    ∃ (M : E n →ₗ[ℝ] E n), ∀ z, enc.toFun z = M z
+
+axiom orthogonal_of_gaussian_linear (M : E n →ₗ[ℝ] E n)
+    (hiso : ∀ v, ‖M v‖ = ‖v‖) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, M z = U z
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: LOSS LOWER BOUND
+-- ═══════════════════════════════════════════════════════════════
+
+theorem loss_lower_bound (enc : HermiteEncoder n) (ρ : ℝ)
+    (_hρ0 : 0 < ρ) (_hρ1 : ρ < 1)
+    (hcorr : ∀ i, enc.correlation i ≤ ρ) :
+    alignmentLoss enc ≥ 2 * (1 - ρ) * n := by
+  unfold alignmentLoss
+  have hsum_le : ∑ i : Fin n, enc.correlation i ≤ ∑ _i : Fin n, ρ :=
+    Finset.sum_le_sum (fun i _ => hcorr i)
+  simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_le
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: MAIN THEOREM ASSEMBLY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Main Theorem** (Theorem 4.1, VERIFIED assembly):
+
+    Any measurable h : ℝⁿ → ℝⁿ with h(z) ~ 𝒩(0, Iₙ) that
+    achieves 𝓛(h) = 2(1−ρ)n must satisfy h(z) = Uz for U ∈ O(n).
+
+    Verified chain:
+    1. Mehler → correlation = Σ w_d ρᵈ     (axiomatized)
+    2. Weighted average → corr_i ≤ ρ        (VERIFIED: correlation_le_rho)
+    3. Loss sum → 𝓛 ≥ 2(1−ρ)n              (VERIFIED: loss_lower_bound)
+    4. 𝓛 = 2(1−ρ)n → each corr_i = ρ       (VERIFIED: Finset.sum_lt_sum)
+    5. corr_i = ρ → w₁ = 1 for all i        (VERIFIED: equality_forces_degree_one)
+    6. w₁ = 1 → h linear                    (axiomatized: linear_of_degree_one)
+    7. Gaussianity + linear → U orthogonal   (axiomatized: orthogonal_of_gaussian_linear)
+-/
+theorem hermite_identifiability
+    (enc : HermiteEncoder n)
+    (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hMehler : ∀ i, Summable (fun d => (enc.spectrum i).w d * ρ ^ d))
+    (hcorr_eq : ∀ i, enc.correlation i =
+      ∑' d, (enc.spectrum i).w d * ρ ^ d)
+    (hopt : alignmentLoss enc = 2 * (1 - ρ) * ↑n)
+    (hnorm : ∀ v, ‖enc.toFun v - enc.toFun 0‖ = ‖v - 0‖) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, enc.toFun z = U z := by
+  -- Step 1: Each correlation ≤ ρ
+  have hcorr_le : ∀ i, enc.correlation i ≤ ρ := by
+    intro i; rw [hcorr_eq i]
+    exact correlation_le_rho (enc.spectrum i) ρ hρ0 hρ1 (hMehler i)
+  -- Step 2: At optimality, each correlation = ρ exactly
+  have hcorr_eq_rho : ∀ i, enc.correlation i = ρ := by
+    by_contra hne; push_neg at hne
+    obtain ⟨i₀, hi₀⟩ := hne
+    have hi₀_lt : enc.correlation i₀ < ρ :=
+      lt_of_le_of_ne (hcorr_le i₀) hi₀
+    have hsum_lt : ∑ i : Fin n, enc.correlation i < ∑ _i : Fin n, ρ :=
+      Finset.sum_lt_sum (fun i _ => hcorr_le i) ⟨i₀, Finset.mem_univ _, hi₀_lt⟩
+    simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_lt
+    unfold alignmentLoss at hopt; linarith
+  -- Step 3: corr_i = ρ forces degree-1 concentration
+  have hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0 := by
+    intro i d hd
+    have hci : ∑' d, (enc.spectrum i).w d * ρ ^ d = ρ := by
+      rw [← hcorr_eq i]; exact hcorr_eq_rho i
+    exact equality_forces_degree_one
+      (enc.spectrum i) ρ hρ0 hρ1 (hMehler i) hci d hd
+  -- Step 4: Linearity
+  obtain ⟨M, hM⟩ := linear_of_degree_one enc hdeg
+  -- Step 5: Orthogonality
+  have hnorm_M : ∀ v, ‖M v‖ = ‖v‖ := by
+    intro v; have hv := hnorm v
+    simp only [sub_zero] at hv
+    rwa [hM v, hM 0, map_zero, sub_zero] at hv
+  obtain ⟨U, hU⟩ := orthogonal_of_gaussian_linear M hnorm_M
+  exact ⟨U, fun z => by rw [hM z, hU z]⟩
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/Planning.lean

@@ -0,0 +1,245 @@
+import Mathlib
+
+/-!
+# Part D — Planning Equivalence (Corollary)
+
+  Let h(z) = Qz with Q ∈ O(n) be the encoder at the optimum of Theorem 4.1.
+  For any finite-horizon optimal control problem whose stage and terminal
+  costs are O(n)-invariant in the state argument, the optimal value function
+  and the set of optimal action sequences agree between the learned latent
+  and the true latent.
+
+  The proof reduces — via the rotation-invariance hypothesis and the
+  pushforward property of expected costs — to the trivial fact that pointwise
+  equal real-valued functions share minimizers.
+
+  ## Verification status
+
+  | Component                             | Status      |
+  |---------------------------------------|-------------|
+  | ControlProblem structure              | structural  |
+  | Orthogonal invariance definition      | structural  |
+  | ExpectedCosts abstraction             | structural  |
+  | Total-cost definition                 | structural  |
+  | Trajectory pushforward (stage)        | axiomatized |
+  | Trajectory pushforward (terminal)     | axiomatized |
+  | Per-step stage cost equivalence       | VERIFIED    |
+  | Terminal cost equivalence             | VERIFIED    |
+  | Total cost equivalence (main step)    | VERIFIED    |
+  | Minimizer equivalence (plan agreement)| VERIFIED    |
+  | Value equivalence                     | VERIFIED    |
+-/
+
+set_option maxHeartbeats 400000
+
+open scoped BigOperators
+
+noncomputable section
+
+abbrev Latent (n : ℕ) := Fin n → ℝ
+abbrev Plan (Action : Type*) (T : ℕ) := Fin T → Action
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STRUCTURE: CONTROL PROBLEM AND ROTATION INVARIANCE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A finite-horizon optimal control problem with stage cost ℓ(z,a) and
+    terminal cost ℓ_T(z). -/
+structure ControlProblem (n : ℕ) (Action : Type*) where
+  stage_cost : Latent n → Action → ℝ
+  terminal_cost : Latent n → ℝ
+
+/-- The costs of the control problem are O(n)-invariant in the state argument
+    under a map Q:  ℓ(Q z, a) = ℓ(z, a) for all z, a, and ℓ_T(Q z) = ℓ_T(z)
+    for all z. In the corollary, Q is the orthogonal recovery matrix from
+    Theorem 4.1; the definition does not itself require Q to be linear or
+    orthogonal — only the invariance property is used. -/
+def IsOrthogonalInvariant {n : ℕ} {Action : Type*}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n) : Prop :=
+  (∀ z a, cp.stage_cost (Q z) a = cp.stage_cost z a) ∧
+  (∀ z, cp.terminal_cost (Q z) = cp.terminal_cost z)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STRUCTURE: EXPECTED COSTS UNDER SOME DYNAMICS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- Expected costs along a trajectory under a specific (stochastic) dynamics.
+
+    `stage_exp a z₀ t c` is the expected value of `c(z_t, a_t)` at time `t`
+    along the trajectory starting from `z₀` and following the action sequence
+    `a`. `term_exp a z₀ c` is the expected value of `c(z_T)` at the final
+    time. Parameterizing over the cost function `c` lets the same dynamics
+    object be reused for different costs, and makes the pushforward relation
+    (below) statable without explicit measure theory. -/
+structure ExpectedCosts (n : ℕ) (Action : Type*) (T : ℕ) where
+  stage_exp :
+    Plan Action T → Latent n → Fin T → (Latent n → Action → ℝ) → ℝ
+  term_exp :
+    Plan Action T → Latent n → (Latent n → ℝ) → ℝ
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- TOTAL EXPECTED COST
+-- ═══════════════════════════════════════════════════════════════
+
+/-- Total expected cost for a plan `a` from initial state `z₀`: the sum of
+    per-step stage costs plus the terminal cost. -/
+def totalCost {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (E : ExpectedCosts n Action T)
+    (a : Plan Action T) (z₀ : Latent n) : ℝ :=
+  (∑ t : Fin T, E.stage_exp a z₀ t cp.stage_cost)
+    + E.term_exp a z₀ cp.terminal_cost
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED: TRAJECTORY PUSHFORWARD
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Stage pushforward** (axiomatized): under the pushforward dynamics
+    `E_hat`, the expected value of any cost `c` at time `t` starting from
+    `Q z` equals the expected value under the original dynamics `E` starting
+    from `z` of the pre-composed cost `c ∘ (Q × id)`.
+
+    Mathematically this is the content of "the joint law of (ẑ_0, …, ẑ_T)
+    under the pushforward dynamics starting from ẑ_0 = Q z equals the joint
+    law of (Q z_0, …, Q z_T) under the original dynamics starting from
+    z_0 = z", restricted to per-time-step marginals and evaluated against
+    arbitrary test functions. -/
+axiom stage_pushforward
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (E_hat E : ExpectedCosts n Action T) (Q : Latent n → Latent n)
+    (a : Plan Action T) (z : Latent n) (t : Fin T)
+    (c : Latent n → Action → ℝ) :
+    E_hat.stage_exp a (Q z) t c
+      = E.stage_exp a z t (fun z' act => c (Q z') act)
+
+/-- **Terminal pushforward** (axiomatized): the same relation at the
+    terminal time. -/
+axiom terminal_pushforward
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (E_hat E : ExpectedCosts n Action T) (Q : Latent n → Latent n)
+    (a : Plan Action T) (z : Latent n) (c : Latent n → ℝ) :
+    E_hat.term_exp a (Q z) c = E.term_exp a z (fun z' => c (Q z'))
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: PER-STEP COST EQUIVALENCE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Stage-cost equivalence** (VERIFIED): the per-step expected stage cost
+    at `Q z` under the pushforward dynamics equals the per-step expected
+    stage cost at `z` under the original dynamics, when the stage cost is
+    O(n)-invariant. This is the point where orthogonal invariance of the
+    cost (hypothesis) meets trajectory pushforward (axiom). -/
+theorem stage_cost_equiv
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n)
+    (E_hat E : ExpectedCosts n Action T)
+    (hinv : IsOrthogonalInvariant cp Q)
+    (a : Plan Action T) (z : Latent n) (t : Fin T) :
+    E_hat.stage_exp a (Q z) t cp.stage_cost
+      = E.stage_exp a z t cp.stage_cost := by
+  rw [stage_pushforward E_hat E Q a z t cp.stage_cost]
+  have hfun : (fun z' act => cp.stage_cost (Q z') act) = cp.stage_cost := by
+    funext z'
+    funext act
+    exact hinv.1 z' act
+  rw [hfun]
+
+/-- **Terminal-cost equivalence** (VERIFIED). -/
+theorem terminal_cost_equiv
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n)
+    (E_hat E : ExpectedCosts n Action T)
+    (hinv : IsOrthogonalInvariant cp Q)
+    (a : Plan Action T) (z : Latent n) :
+    E_hat.term_exp a (Q z) cp.terminal_cost
+      = E.term_exp a z cp.terminal_cost := by
+  rw [terminal_pushforward E_hat E Q a z cp.terminal_cost]
+  have hfun : (fun z' => cp.terminal_cost (Q z')) = cp.terminal_cost := by
+    funext z'
+    exact hinv.2 z'
+  rw [hfun]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: TOTAL COST EQUIVALENCE (PLANNING EQUIVALENCE)
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Planning equivalence** (VERIFIED, main step): for any action sequence,
+    the total expected cost under the pushforward dynamics at `Q z₀` equals
+    the total expected cost under the original dynamics at `z₀`.
+
+    This is the central computational content of the corollary; everything
+    that follows (value and minimizer equivalence) is a consequence. -/
+theorem planning_equivalence
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n)
+    (E_hat E : ExpectedCosts n Action T)
+    (hinv : IsOrthogonalInvariant cp Q)
+    (a : Plan Action T) (z : Latent n) :
+    totalCost cp E_hat a (Q z) = totalCost cp E a z := by
+  unfold totalCost
+  have hstage :
+      (∑ t : Fin T, E_hat.stage_exp a (Q z) t cp.stage_cost)
+        = ∑ t : Fin T, E.stage_exp a z t cp.stage_cost := by
+    apply Finset.sum_congr rfl
+    intro t _
+    exact stage_cost_equiv cp Q E_hat E hinv a z t
+  have hterm :
+      E_hat.term_exp a (Q z) cp.terminal_cost
+        = E.term_exp a z cp.terminal_cost :=
+    terminal_cost_equiv cp Q E_hat E hinv a z
+  rw [hstage, hterm]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: MINIMIZER AND VALUE EQUIVALENCE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Minimizer equivalence** (VERIFIED): an action sequence minimizes the
+    expected cost under the pushforward dynamics at `Q z` iff it minimizes
+    the expected cost under the original dynamics at `z`.
+
+    Consequence: the optimal plan is the same whether it is computed in the
+    learned latent or the true latent. -/
+theorem minimizer_equivalence
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n)
+    (E_hat E : ExpectedCosts n Action T)
+    (hinv : IsOrthogonalInvariant cp Q)
+    (a : Plan Action T) (z : Latent n) :
+    (∀ a', totalCost cp E_hat a (Q z) ≤ totalCost cp E_hat a' (Q z)) ↔
+    (∀ a', totalCost cp E a z ≤ totalCost cp E a' z) := by
+  have h : ∀ a', totalCost cp E_hat a' (Q z) = totalCost cp E a' z :=
+    fun a' => planning_equivalence cp Q E_hat E hinv a' z
+  constructor
+  · intro hmin a'
+    have ha := h a
+    have ha' := h a'
+    have := hmin a'
+    linarith
+  · intro hmin a'
+    have ha := h a
+    have ha' := h a'
+    have := hmin a'
+    linarith
+
+/-- **Value equivalence** (VERIFIED): if `a` achieves total cost `V` under
+    the original dynamics at `z`, it achieves the same `V` under the
+    pushforward dynamics at `Q z`. Combined with `minimizer_equivalence`,
+    this gives the corollary's `V̂*(Q z) = V*(z)` statement. -/
+theorem value_equivalence
+    {n : ℕ} {Action : Type*} {T : ℕ}
+    (cp : ControlProblem n Action) (Q : Latent n → Latent n)
+    (E_hat E : ExpectedCosts n Action T)
+    (hinv : IsOrthogonalInvariant cp Q)
+    (a : Plan Action T) (z : Latent n) (V : ℝ)
+    (hV : totalCost cp E a z = V) :
+    totalCost cp E_hat a (Q z) = V := by
+  rw [planning_equivalence cp Q E_hat E hinv a z, hV]
+
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/PropApprox.lean

@@ -0,0 +1,187 @@
+import Mathlib
+
+/-!
+# Part C — Approximate Identifiability (Proposition 4.3)
+
+  Under approximate alignment (gap δ) and approximate covariance
+  (error ε), the recovery error satisfies:
+
+    𝔼[‖h(z) − Qz‖²] ≤ D + (ε + D)²
+
+  where D = δ/(2ρ(1−ρ)) is the alignment gap normalized by the
+  spectral gap between Hermite degrees 1 and 2.
+
+  When δ = ε = 0 this recovers Theorem 4.1: h(z) = Qz a.e.
+
+  ## Verification status
+
+  | Component                          | Status      |
+  |------------------------------------|-------------|
+  | Spectral gap positivity            | VERIFIED    |
+  | W_nl ≤ D from gap inequality       | VERIFIED    |
+  | Polar decomposition ‖M−Q‖ bound    | axiomatized |
+  | Cross-degree Hermite orthogonality  | axiomatized |
+  | Linear deviation ‖M−Q‖² bound      | VERIFIED    |
+  | Pythagorean decomposition           | axiomatized |
+  | Bound monotonicity in W_nl          | VERIFIED    |
+  | Full bound assembly                 | VERIFIED    |
+  | Exact recovery (δ=ε=0 ⟹ error=0)  | VERIFIED    |
+-/
+
+noncomputable section
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 1: SPECTRAL GAP CONTROLS NONLINEAR ENERGY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- The spectral gap ρ(1−ρ) is positive for 0 < ρ < 1. -/
+theorem spectral_gap_pos (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    0 < ρ * (1 - ρ) := by
+  apply mul_pos hρ0; linarith
+
+/-- 2ρ(1−ρ) is positive. -/
+theorem two_spectral_gap_pos (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    0 < 2 * ρ * (1 - ρ) := by
+  have : 0 < ρ * (1 - ρ) := spectral_gap_pos ρ hρ0 hρ1
+  linarith
+
+/-- **Nonlinear energy bound** (VERIFIED): from the spectral gap
+    inequality δ ≥ 2ρ(1−ρ) W_nl, we get W_nl ≤ D = δ/(2ρ(1−ρ)). -/
+theorem nonlinear_energy_le_D
+    (ρ δ W_nl : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (_hδ_nonneg : 0 ≤ δ) (_hW_nonneg : 0 ≤ W_nl)
+    (hgap : δ ≥ 2 * ρ * (1 - ρ) * W_nl) :
+    W_nl ≤ δ / (2 * ρ * (1 - ρ)) := by
+  have hsgap : (0 : ℝ) < 2 * ρ * (1 - ρ) := two_spectral_gap_pos ρ hρ0 hρ1
+  rw [le_div_iff₀ hsgap]
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 2: LINEAR PART DEVIATION
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Polar decomposition bound** (axiomatized): ‖M − Q‖_F ≤ ε + W_nl.
+    Combines polar decomposition, |σᵢ−1| ≤ |σᵢ²−1|, covariance
+    decomposition Cov(h) = MM^T + N, and triangle inequality. -/
+axiom polar_bound_axiom
+    (M_Q_norm ε W_nl : ℝ)
+    (hε : 0 ≤ ε) (hW : 0 ≤ W_nl) :
+    M_Q_norm ≤ ε + W_nl →
+    M_Q_norm ≤ ε + W_nl
+
+/-- **Linear deviation squared** (VERIFIED): ‖M−Q‖ ≤ ε+W_nl implies
+    ‖M−Q‖² ≤ (ε+W_nl)². -/
+theorem linear_deviation_sq_bound
+    (M_Q_norm ε W_nl : ℝ)
+    (hMQ_nonneg : 0 ≤ M_Q_norm)
+    (hε : 0 ≤ ε) (hW : 0 ≤ W_nl)
+    (hbound : M_Q_norm ≤ ε + W_nl) :
+    M_Q_norm ^ 2 ≤ (ε + W_nl) ^ 2 := by
+  have h1 : 0 ≤ ε + W_nl := by linarith
+  nlinarith [sq_nonneg (ε + W_nl - M_Q_norm)]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 3: PYTHAGOREAN DECOMPOSITION
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Pythagorean decomposition** (axiomatized): the recovery error
+    splits into linear deviation and nonlinear energy.
+    Requires Hermite orthogonality and z ~ N(0,I). -/
+axiom pythagorean_axiom
+    (total_error M_Q_norm_sq W_nl : ℝ) :
+    total_error = M_Q_norm_sq + W_nl →
+    total_error = M_Q_norm_sq + W_nl
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STEP 4: MONOTONICITY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Monotonicity** (VERIFIED): f(t) = (ε + t)² + t is increasing
+    for t ≥ 0. So W_nl ≤ D implies (ε+W_nl)²+W_nl ≤ (ε+D)²+D. -/
+theorem bound_monotone (ε W_nl D : ℝ)
+    (_hε : 0 ≤ ε) (_hW : 0 ≤ W_nl) (_hD : 0 ≤ D)
+    (hle : W_nl ≤ D) :
+    (ε + W_nl) ^ 2 + W_nl ≤ (ε + D) ^ 2 + D := by
+  have h1 : ε + W_nl ≤ ε + D := by linarith
+  nlinarith [sq_nonneg (ε + D - ε - W_nl)]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- MAIN BOUND ASSEMBLY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Approximate identifiability** (Proposition 4.3, VERIFIED assembly):
+
+    𝔼[‖h(z) − Qz‖²] ≤ D + (ε + D)²
+
+    where D = δ/(2ρ(1−ρ)). -/
+theorem approximate_identifiability
+    (ρ δ ε W_nl M_Q_norm total_error : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hδ : 0 ≤ δ) (hε : 0 ≤ ε)
+    (hW : 0 ≤ W_nl) (hMQ : 0 ≤ M_Q_norm)
+    (hgap : δ ≥ 2 * ρ * (1 - ρ) * W_nl)
+    (hpolar : M_Q_norm ≤ ε + W_nl)
+    (hpythag : total_error = M_Q_norm ^ 2 + W_nl) :
+    total_error ≤ δ / (2 * ρ * (1 - ρ))
+              + (ε + δ / (2 * ρ * (1 - ρ))) ^ 2 := by
+  set D := δ / (2 * ρ * (1 - ρ)) with hD_def
+  have hsgap := two_spectral_gap_pos ρ hρ0 hρ1
+  have hD_nonneg : 0 ≤ D := div_nonneg hδ (le_of_lt hsgap)
+  -- Step 1: W_nl ≤ D
+  have hW_le_D : W_nl ≤ D := nonlinear_energy_le_D ρ δ W_nl hρ0 hρ1 hδ hW hgap
+  -- Step 4: ‖M−Q‖² ≤ (ε + W_nl)²
+  have hMQ_sq : M_Q_norm ^ 2 ≤ (ε + W_nl) ^ 2 :=
+    linear_deviation_sq_bound M_Q_norm ε W_nl hMQ hε hW hpolar
+  -- Step 3 + 4: total_error ≤ (ε + W_nl)² + W_nl
+  have h_inter : total_error ≤ (ε + W_nl) ^ 2 + W_nl := by
+    rw [hpythag]; linarith
+  -- Step 5: monotonicity
+  have h_mono := bound_monotone ε W_nl D hε hW hD_nonneg hW_le_D
+  -- Combine
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- EXACT RECOVERY AS SPECIAL CASE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Exact recovery** (VERIFIED): setting δ = ε = 0 gives error = 0,
+    recovering Theorem 4.1: h(z) = Qz almost everywhere. -/
+theorem exact_recovery_special_case
+    (ρ W_nl M_Q_norm total_error : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hW : 0 ≤ W_nl) (hMQ : 0 ≤ M_Q_norm)
+    (hgap : (0 : ℝ) ≥ 2 * ρ * (1 - ρ) * W_nl)
+    (hpolar : M_Q_norm ≤ 0 + W_nl)
+    (hpythag : total_error = M_Q_norm ^ 2 + W_nl)
+    (_htotal_nonneg : 0 ≤ total_error) :
+    total_error = 0 := by
+  -- δ = 0 forces W_nl = 0
+  have hsgap := two_spectral_gap_pos ρ hρ0 hρ1
+  have hW_zero : W_nl = 0 := by nlinarith
+  -- W_nl = 0 and ε = 0 force ‖M − Q‖ = 0
+  have hMQ_zero : M_Q_norm = 0 := by
+    have : M_Q_norm ≤ 0 := by linarith [hpolar, hW_zero]
+    linarith
+  -- Total error = 0² + 0 = 0
+  rw [hpythag, hMQ_zero, hW_zero]; ring
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- BOUND STRUCTURE ANALYSIS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **First-order approximation** (VERIFIED): when ε + D ≤ 1,
+    the quadratic term (ε+D)² ≤ ε+D, so the bound ≤ 2D + ε. -/
+theorem bound_small_perturbation (ε D : ℝ)
+    (hε : 0 ≤ ε) (hD : 0 ≤ D) (hsmall : ε + D ≤ 1) :
+    D + (ε + D) ^ 2 ≤ D + ε + D := by
+  have h1 : 0 ≤ ε + D := by linarith
+  nlinarith [sq_nonneg (1 - (ε + D))]
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/ThmDirichlet.lean

@@ -0,0 +1,227 @@
+import Mathlib.Analysis.InnerProductSpace.Basic
+import Mathlib.Analysis.InnerProductSpace.PiL2
+import Mathlib.Analysis.Normed.Module.Basic
+import Mathlib.Analysis.Calculus.MeanValue
+import Mathlib.Analysis.SpecialFunctions.Pow.Real
+import Mathlib.Analysis.SpecialFunctions.ExpDeriv
+import Mathlib.LinearAlgebra.Matrix.NonsingularInverse
+import Mathlib.LinearAlgebra.Matrix.Determinant.Basic
+import Mathlib.Topology.MetricSpace.Isometry
+import Mathlib.Topology.MetricSpace.Lipschitz
+
+/-!
+# Part B — Alternative Proof via Dirichlet Energy (Appendix C)
+
+  Any C¹ diffeomorphism h : ℝⁿ → ℝⁿ that preserves the standard
+  Gaussian measure and minimizes the Dirichlet energy 𝔼[‖Jₕ‖²_F]
+  must be a linear orthogonal map h(z) = Uz.
+
+  ## Proof sketch
+
+  Steps 1–2 (reduction to Dirichlet energy and the log-determinant
+  lemma) involve measure-theoretic integration. We axiomatize their
+  conclusions.
+
+  Steps 3–5 are verified:
+    Step 3: AM-GM + Jensen → 𝓙(h) ≥ n                 (axiomatized)
+    Step 4: Equality forces Jₕ orthogonal everywhere   (axiomatized)
+    Step 5: Orthogonal Jacobian → global isometry →
+            Mazur–Ulam → linear                        (VERIFIED)
+
+  ## Verification status
+
+  | Component                        | Status      |
+  |----------------------------------|-------------|
+  | AM-GM for singular values        | axiomatized |
+  | Jensen for log-determinant       | axiomatized |
+  | Mazur–Ulam theorem               | axiomatized |
+  | Norm-preserving CLM → isometry   | VERIFIED    |
+  | Orthogonal Jacobian → Lipschitz  | VERIFIED    |
+  | Bilipschitz → global isometry    | VERIFIED    |
+  | h(0)=0 → b=0 → linear isometry  | VERIFIED    |
+  | Full theorem assembly            | VERIFIED    |
+-/
+
+open scoped Matrix BigOperators
+open Matrix
+
+noncomputable section
+
+variable {n : ℕ}
+
+/-- The type we work with: ℝⁿ as a Euclidean space. -/
+private abbrev E (n : ℕ) := EuclideanSpace ℝ (Fin n)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED KNOWN RESULTS
+-- ═══════════════════════════════════════════════════════════════
+
+/-!
+These are standard results available in Mathlib but requiring
+nontrivial plumbing to connect to our specific statement forms.
+-/
+
+/-- **AM-GM inequality**: arithmetic mean of nonneg reals ≥ geometric
+    mean. Special case of `Real.geom_mean_le_arith_mean_weighted`
+    in `Mathlib.Analysis.MeanInequalities` with uniform weights. -/
+axiom amgm_sum_ge_prod_pow {m : ℕ} (a : Fin m → ℝ)
+    (ha : ∀ i, 0 ≤ a i) :
+    (∑ i : Fin m, a i) / m ≥ (∏ i : Fin m, a i) ^ ((1 : ℝ) / m)
+
+/-- **Jensen's inequality** applied to strictly convex exp:
+    mean of exp(cxᵢ) ≥ 1 when xᵢ sum to zero. Follows from
+    `StrictConvexOn` of `Real.exp` and the weighted AM-GM. -/
+axiom exp_mean_ge_mean_exp {m : ℕ}
+    (f : Fin m → ℝ) (hsum : ∑ i : Fin m, f i = 0) :
+    (∑ i : Fin m, Real.exp ((2 : ℝ) / m * f i)) / m ≥ 1
+
+/-- **Mazur–Ulam theorem**: every surjective isometry of a real normed
+    space is affine. Available in Mathlib as the combination of
+    `Isometry.right_inv` and affine isometry machinery in
+    `Mathlib.Analysis.Normed.Affine.Isometry`. -/
+axiom mazur_ulam
+    {V : Type*} [NormedAddCommGroup V] [NormedSpace ℝ V]
+    {f : V → V} (hiso : Isometry f) (hsurj : Function.Surjective f) :
+    ∃ (A : V →ₗ[ℝ] V) (b : V), ∀ x, f x = A x + b
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- DIFFEOMORPHISM STRUCTURE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A smooth map h : ℝⁿ → ℝⁿ with its Jacobian, modeling a C¹
+    diffeomorphism that preserves the standard Gaussian. -/
+structure GaussianDiffeo (n : ℕ) where
+  /-- The map itself -/
+  toFun : E n → E n
+  /-- The Jacobian at each point, as a continuous linear map -/
+  jacobian : E n → (E n →L[ℝ] E n)
+  /-- h is differentiable with the given Jacobian -/
+  hasFDeriv : ∀ z, HasFDerivAt toFun (jacobian z) z
+  /-- h is a homeomorphism (hence bijective) -/
+  isHomeo : (E n) ≃ₜ (E n)
+  /-- The homeomorphism agrees with toFun -/
+  homeo_eq : ∀ z, isHomeo z = toFun z
+  /-- Inverse differentiability from the **inverse function theorem**
+      (`HasStrictFDerivAt.toOpenPartialHomeomorph` in
+      `Mathlib.Analysis.Calculus.InverseFunctionTheorem.FDeriv`). -/
+  hasFDeriv_inv : ∀ y, HasFDerivAt isHomeo.symm
+    (ContinuousLinearMap.inverse (jacobian (isHomeo.symm y))) y
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: ORTHOGONAL JACOBIAN → GLOBAL ISOMETRY → LINEAR
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A norm-preserving continuous linear map is an isometry. -/
+theorem clm_isometry_of_norm_preserving
+    (L : E n →L[ℝ] E n)
+    (hL : ∀ v, ‖L v‖ = ‖v‖) :
+    Isometry L := by
+  rw [isometry_iff_dist_eq]
+  intro x y
+  simp only [dist_eq_norm, ← map_sub L x y]
+  exact hL (x - y)
+
+/-- **Mean value theorem** (VERIFIED): orthogonal Jacobian everywhere
+    ⟹ h is 1-Lipschitz. By the MVT, ‖h(x)-h(y)‖ ≤ sup ‖Jₕ‖_op · ‖x-y‖,
+    and the operator norm of a norm-preserving map is 1. -/
+theorem lipschitz_of_orthogonal_jacobian
+    (h : GaussianDiffeo n)
+    (horth : ∀ z v, ‖h.jacobian z v‖ = ‖v‖) :
+    LipschitzWith 1 h.toFun := by
+  apply lipschitzWith_of_nnnorm_fderiv_le (𝕜 := ℝ)
+  · intro x; exact (h.hasFDeriv x).differentiableAt
+  · intro x
+    have hfderiv : fderiv ℝ h.toFun x = h.jacobian x :=
+      (h.hasFDeriv x).fderiv
+    rw [hfderiv, ContinuousLinearMap.opNNNorm_le_iff]
+    intro y; simp only [one_mul]
+    exact_mod_cast le_of_eq (horth x y)
+
+/-- **Bilipschitz → isometry** (VERIFIED): if both h and h⁻¹ are
+    1-Lipschitz, h is a global isometry. Forward Lipschitz gives
+    dist(hx,hy) ≤ dist(x,y); applying to h⁻¹ gives ≥. -/
+theorem isometry_of_bilipschitz
+    (h : GaussianDiffeo n)
+    (hlip : LipschitzWith 1 h.toFun)
+    (hinvlip : LipschitzWith 1 h.isHomeo.symm) :
+    Isometry h.toFun := by
+  rw [isometry_iff_dist_eq]
+  intro x y
+  apply le_antisymm
+  · -- Forward: dist(hx, hy) ≤ dist(x, y)
+    have hfwd := hlip.dist_le_mul x y
+    simp only [NNReal.coe_one, one_mul] at hfwd; exact hfwd
+  · -- Backward: apply Lipschitz to h⁻¹
+    have hbwd := hinvlip.dist_le_mul (h.toFun x) (h.toFun y)
+    simp only [NNReal.coe_one, one_mul] at hbwd
+    have hx : h.isHomeo.symm (h.toFun x) = x := by
+      rw [← h.homeo_eq]; exact h.isHomeo.symm_apply_apply x
+    have hy : h.isHomeo.symm (h.toFun y) = y := by
+      rw [← h.homeo_eq]; exact h.isHomeo.symm_apply_apply y
+    rw [hx, hy] at hbwd; exact hbwd
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: MAIN THEOREM (APPENDIX C)
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **LeJEPA identifiability via Dirichlet energy** (VERIFIED):
+
+    C¹ diffeomorphism + Gaussian-preserving + orthogonal Jacobian
+    ⟹ h(z) = Uz for a linear isometry U ∈ O(n).
+
+    Verified chain:
+    1. Orth. Jacobian → h is 1-Lipschitz    (MVT)
+    2. Orth. inverse  → h⁻¹ is 1-Lipschitz  (IFT + MVT)
+    3. Bilipschitz → global isometry
+    4. Mazur–Ulam → h is affine: h(z) = Az + b
+    5. h(0) = 0 → b = 0
+    6. A preserves norms → A is a LinearIsometry -/
+theorem dirichlet_identifiability
+    (h : GaussianDiffeo n)
+    (horth : ∀ z v, ‖h.jacobian z v‖ = ‖v‖)
+    (horth_inv : ∀ z v,
+      ‖(ContinuousLinearMap.inverse (h.jacobian z)) v‖ = ‖v‖)
+    (hmean : h.toFun 0 = 0) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, h.toFun z = U z := by
+  -- Step 1: h is 1-Lipschitz
+  have hlip := lipschitz_of_orthogonal_jacobian h horth
+  -- Step 2: h⁻¹ is 1-Lipschitz (IFT gives derivative = J⁻¹, also orth.)
+  have hinvlip : LipschitzWith 1 h.isHomeo.symm := by
+    apply lipschitzWith_of_nnnorm_fderiv_le (𝕜 := ℝ)
+    · intro x; exact (h.hasFDeriv_inv x).differentiableAt
+    · intro x
+      have hfderiv : fderiv ℝ h.isHomeo.symm x =
+        (h.jacobian (h.isHomeo.symm x)).inverse :=
+        (h.hasFDeriv_inv x).fderiv
+      rw [hfderiv, ContinuousLinearMap.opNNNorm_le_iff]
+      intro y; simp only [one_mul]
+      exact_mod_cast le_of_eq (horth_inv (h.isHomeo.symm x) y)
+  -- Step 3: h is a global isometry
+  have hiso := isometry_of_bilipschitz h hlip hinvlip
+  -- Step 4: Mazur–Ulam → h(z) = Az + b
+  have hsurj : Function.Surjective h.toFun := by
+    intro y
+    exact ⟨h.isHomeo.symm y,
+      by rw [← h.homeo_eq]; exact h.isHomeo.apply_symm_apply y⟩
+  obtain ⟨A, b, hab⟩ := mazur_ulam hiso hsurj
+  -- Step 5: b = 0 from h(0) = 0
+  have hb : b = 0 := by
+    have h0 := hab 0; simp [map_zero] at h0
+    rw [hmean] at h0; exact h0.symm
+  -- h(z) = Az for all z
+  have hab' : ∀ z, h.toFun z = A z := by
+    intro z; have := hab z; rw [hb, add_zero] at this; exact this
+  -- Step 6: A preserves norms → LinearIsometry
+  have hA_norm : ∀ v, ‖A v‖ = ‖v‖ := by
+    intro v
+    have hv := hiso.dist_eq v 0
+    simp [dist_eq_norm] at hv
+    rw [hab' v, hab' 0, map_zero] at hv
+    simpa using hv
+  exact ⟨⟨A, hA_norm⟩, hab'⟩
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/ThmHermite.lean

@@ -0,0 +1,271 @@
+-- import Mathlib
+import Mathlib.Analysis.InnerProductSpace.PiL2
+import Mathlib.Topology.Algebra.InfiniteSum.Order
+import Mathlib.Topology.Algebra.InfiniteSum.Ring
+
+/-!
+# Part A — Main Theorem via Hermite Polynomials (Theorem 4.1)
+
+  Any measurable h : ℝⁿ → ℝⁿ satisfying Gaussianity h(z) ~ N(0,Iₙ)
+  and minimizing the alignment loss must be h(z) = Uz for U ∈ O(n).
+
+  ## Verification status
+
+  | Component                        | Status      |
+  |----------------------------------|-------------|
+  | Hermite basis & completeness     | axiomatized |
+  | Contraction lemma (ρᵈ decay)     | axiomatized |
+  | Mehler's formula                 | axiomatized |
+  | ρᵈ ≤ ρ for d ≥ 1                | VERIFIED    |
+  | ρᵈ < ρ for d ≥ 2                | VERIFIED    |
+  | Pointwise term bound w_d·ρᵈ≤w_d·ρ| VERIFIED    |
+  | Correlation bound ≤ ρ            | VERIFIED    |
+  | Equality ⟺ w₁ = 1 (linearity)  | VERIFIED    |
+  | Loss lower bound 2(1-ρ)n        | VERIFIED    |
+  | Theorem assembly h = Uz         | VERIFIED    |
+-/
+
+set_option maxHeartbeats 400000
+
+open scoped BigOperators
+
+noncomputable section
+
+abbrev E (n : ℕ) := EuclideanSpace ℝ (Fin n)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- SPECTRAL WEIGHTS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- Spectral weights of a single encoder component in its Hermite
+    expansion. `w d` is the fraction of L²(γₙ) variance at degree d. -/
+structure SpectralWeights where
+  w : ℕ → ℝ
+  nonneg : ∀ d, 0 ≤ w d
+  zero_degree : w 0 = 0
+  summable : Summable w
+  total_variance : ∑' d, w d = 1
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED: HERMITE BASIS & MEHLER
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Mehler's formula** (axiomatized): the spectral correlation
+    series Σ_d w_d · ρᵈ is summable. -/
+axiom mehler_summability
+    (sw : SpectralWeights) (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
+    Summable (fun d => sw.w d * ρ ^ d)
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: POINTWISE BOUNDS
+-- ═══════════════════════════════════════════════════════════════
+
+/-- For 0 < ρ ≤ 1 and d ≥ 1, ρᵈ ≤ ρ. -/
+theorem pow_le_self_of_pos_lt_one (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1)
+    (d : ℕ) (hd : 1 ≤ d) : ρ ^ d ≤ ρ := by
+  calc ρ ^ d ≤ ρ ^ 1 := pow_le_pow_of_le_one (le_of_lt hρ0) hρ1 hd
+       _ = ρ := pow_one ρ
+
+/-- Each term w_d · ρᵈ ≤ w_d · ρ. -/
+theorem spectral_term_le (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1) (d : ℕ) :
+    sw.w d * ρ ^ d ≤ sw.w d * ρ := by
+  match d with
+  | 0 => simp [sw.zero_degree]
+  | d + 1 =>
+    exact mul_le_mul_of_nonneg_left
+      (pow_le_self_of_pos_lt_one ρ hρ0 hρ1 (d + 1)
+        (Nat.succ_le_succ (Nat.zero_le d)))
+      (sw.nonneg (d + 1))
+
+/-- For 0 < ρ < 1 and d ≥ 2, ρᵈ < ρ (strict). -/
+theorem pow_lt_self_of_ge_two (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (d : ℕ) (hd : 2 ≤ d) : ρ ^ d < ρ := by
+  calc ρ ^ d ≤ ρ ^ 2 := pow_le_pow_of_le_one (le_of_lt hρ0) (le_of_lt hρ1) hd
+       _ = ρ * ρ      := by ring
+       _ < ρ * 1      := mul_lt_mul_of_pos_left hρ1 hρ0
+       _ = ρ           := mul_one ρ
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: SUMMABILITY AND TSUM OF UPPER BOUND
+-- ═══════════════════════════════════════════════════════════════
+
+/-- The constant-ρ series fun d ↦ w d * ρ is summable
+    (via Summable.mul_right from Ring.lean). -/
+theorem summable_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
+    Summable (fun d => sw.w d * ρ) :=
+  sw.summable.mul_right ρ
+
+/-- Σ w_d · ρ = (Σ w_d) · ρ = 1 · ρ = ρ
+    (via tsum_mul_right from Ring.lean). -/
+theorem tsum_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
+    ∑' d, sw.w d * ρ = ρ := by
+  rw [tsum_mul_right, sw.total_variance, one_mul]
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: CORRELATION BOUND (Lemma 3.3)
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Correlation bound** (VERIFIED): Σ_d w_d ρᵈ ≤ ρ.
+    Uses Summable.tsum_le_tsum (from Order.lean via @[to_additive]). -/
+theorem correlation_le_rho (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hsum : Summable (fun d => sw.w d * ρ ^ d)) :
+    ∑' d, sw.w d * ρ ^ d ≤ ρ := by
+  calc ∑' d, sw.w d * ρ ^ d
+      ≤ ∑' d, sw.w d * ρ :=
+        hsum.tsum_le_tsum
+          (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
+          (summable_spectral_upper sw ρ)
+    _ = ρ := tsum_spectral_upper sw ρ
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: EQUALITY FORCES LINEARITY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Equality characterization** (VERIFIED): if Σ w_d ρᵈ = ρ, then
+    w_d = 0 for all d ≥ 2.
+
+    Strategy: by contradiction. If w_{d₀} > 0 for some d₀ ≥ 2, then
+    w_{d₀}·ρ^{d₀} < w_{d₀}·ρ strictly, while all other terms satisfy ≤.
+    By Summable.tsum_lt_tsum (from Order.lean via @[to_additive]),
+    Σ w_d·ρᵈ < Σ w_d·ρ = ρ, contradicting Σ w_d·ρᵈ = ρ. -/
+theorem equality_forces_degree_one (sw : SpectralWeights) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hsum : Summable (fun d => sw.w d * ρ ^ d))
+    (heq : ∑' d, sw.w d * ρ ^ d = ρ) :
+    ∀ d, 2 ≤ d → sw.w d = 0 := by
+  by_contra h
+  push_neg at h
+  obtain ⟨d₀, hd₀_ge, hd₀_ne⟩ := h
+  -- w_{d₀} > 0
+  have hwd₀_pos : 0 < sw.w d₀ :=
+    lt_of_le_of_ne (sw.nonneg d₀) (Ne.symm hd₀_ne)
+  -- Strict inequality at d₀: w_{d₀} · ρ^{d₀} < w_{d₀} · ρ
+  have hstrict : sw.w d₀ * ρ ^ d₀ < sw.w d₀ * ρ :=
+    mul_lt_mul_of_pos_left (pow_lt_self_of_ge_two ρ hρ0 hρ1 d₀ hd₀_ge) hwd₀_pos
+  -- By tsum_lt_tsum: one strict + rest ≤ ⟹ strict on tsums
+  have hlt : ∑' d, sw.w d * ρ ^ d < ∑' d, sw.w d * ρ :=
+    hsum.tsum_lt_tsum
+      (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
+      hstrict
+      (summable_spectral_upper sw ρ)
+  -- But Σ w_d·ρᵈ = ρ = Σ w_d·ρ
+  rw [tsum_spectral_upper, heq] at hlt
+  exact lt_irrefl ρ hlt
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- ENCODER STRUCTURE & LOSS
+-- ═══════════════════════════════════════════════════════════════
+
+variable {n : ℕ}
+
+/-- An encoder h : ℝⁿ → ℝⁿ with its Hermite spectral decomposition. -/
+structure HermiteEncoder (n : ℕ) where
+  toFun : E n → E n
+  spectrum : Fin n → SpectralWeights
+  correlation : Fin n → ℝ
+
+/-- The alignment loss: 𝓛(h) = 2n − 2 Σᵢ corr_i. -/
+def alignmentLoss (enc : HermiteEncoder n) : ℝ :=
+  2 * n - 2 * ∑ i : Fin n, enc.correlation i
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED: BRIDGE LEMMAS
+-- ═══════════════════════════════════════════════════════════════
+
+axiom correlation_eq_spectral_sum (enc : HermiteEncoder n) (ρ : ℝ)
+    (hρ0 : 0 < ρ) (hρ1 : ρ < 1) (i : Fin n) :
+    enc.correlation i = ∑' d, (enc.spectrum i).w d * ρ ^ d
+
+axiom linear_of_degree_one (enc : HermiteEncoder n)
+    (hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0) :
+    ∃ (M : E n →ₗ[ℝ] E n), ∀ z, enc.toFun z = M z
+
+axiom orthogonal_of_gaussian_linear (M : E n →ₗ[ℝ] E n)
+    (hiso : ∀ v, ‖M v‖ = ‖v‖) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, M z = U z
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: LOSS LOWER BOUND
+-- ═══════════════════════════════════════════════════════════════
+
+theorem loss_lower_bound (enc : HermiteEncoder n) (ρ : ℝ)
+    (_hρ0 : 0 < ρ) (_hρ1 : ρ < 1)
+    (hcorr : ∀ i, enc.correlation i ≤ ρ) :
+    alignmentLoss enc ≥ 2 * (1 - ρ) * n := by
+  unfold alignmentLoss
+  have hsum_le : ∑ i : Fin n, enc.correlation i ≤ ∑ _i : Fin n, ρ :=
+    Finset.sum_le_sum (fun i _ => hcorr i)
+  simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_le
+  linarith
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: MAIN THEOREM ASSEMBLY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Main Theorem** (Theorem 4.1, VERIFIED assembly):
+
+    Any measurable h : ℝⁿ → ℝⁿ with h(z) ~ 𝒩(0, Iₙ) that
+    achieves 𝓛(h) = 2(1−ρ)n must satisfy h(z) = Uz for U ∈ O(n).
+
+    Verified chain:
+    1. Mehler → correlation = Σ w_d ρᵈ     (axiomatized)
+    2. Weighted average → corr_i ≤ ρ        (VERIFIED: correlation_le_rho)
+    3. Loss sum → 𝓛 ≥ 2(1−ρ)n              (VERIFIED: loss_lower_bound)
+    4. 𝓛 = 2(1−ρ)n → each corr_i = ρ       (VERIFIED: Finset.sum_lt_sum)
+    5. corr_i = ρ → w₁ = 1 for all i        (VERIFIED: equality_forces_degree_one)
+    6. w₁ = 1 → h linear                    (axiomatized: linear_of_degree_one)
+    7. Gaussianity + linear → U orthogonal   (axiomatized: orthogonal_of_gaussian_linear)
+-/
+theorem hermite_identifiability
+    (enc : HermiteEncoder n)
+    (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
+    (hMehler : ∀ i, Summable (fun d => (enc.spectrum i).w d * ρ ^ d))
+    (hcorr_eq : ∀ i, enc.correlation i =
+      ∑' d, (enc.spectrum i).w d * ρ ^ d)
+    (hopt : alignmentLoss enc = 2 * (1 - ρ) * ↑n)
+    (hnorm : ∀ v, ‖enc.toFun v - enc.toFun 0‖ = ‖v - 0‖) :
+    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, enc.toFun z = U z := by
+  -- Step 1: Each correlation ≤ ρ
+  have hcorr_le : ∀ i, enc.correlation i ≤ ρ := by
+    intro i; rw [hcorr_eq i]
+    exact correlation_le_rho (enc.spectrum i) ρ hρ0 hρ1 (hMehler i)
+  -- Step 2: At optimality, each correlation = ρ exactly
+  have hcorr_eq_rho : ∀ i, enc.correlation i = ρ := by
+    by_contra hne; push_neg at hne
+    obtain ⟨i₀, hi₀⟩ := hne
+    have hi₀_lt : enc.correlation i₀ < ρ :=
+      lt_of_le_of_ne (hcorr_le i₀) hi₀
+    have hsum_lt : ∑ i : Fin n, enc.correlation i < ∑ _i : Fin n, ρ :=
+      Finset.sum_lt_sum (fun i _ => hcorr_le i) ⟨i₀, Finset.mem_univ _, hi₀_lt⟩
+    simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_lt
+    unfold alignmentLoss at hopt; linarith
+  -- Step 3: corr_i = ρ forces degree-1 concentration
+  have hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0 := by
+    intro i d hd
+    have hci : ∑' d, (enc.spectrum i).w d * ρ ^ d = ρ := by
+      rw [← hcorr_eq i]; exact hcorr_eq_rho i
+    exact equality_forces_degree_one
+      (enc.spectrum i) ρ hρ0 hρ1 (hMehler i) hci d hd
+  -- Step 4: Linearity
+  obtain ⟨M, hM⟩ := linear_of_degree_one enc hdeg
+  -- Step 5: Orthogonality
+  have hnorm_M : ∀ v, ‖M v‖ = ‖v‖ := by
+    intro v; have hv := hnorm v
+    simp only [sub_zero] at hv
+    rwa [hM v, hM 0, map_zero, sub_zero] at hv
+  obtain ⟨U, hU⟩ := orthogonal_of_gaussian_linear M hnorm_M
+  exact ⟨U, fun z => by rw [hM z, hU z]⟩
+
+end

JEPA/lejepa-identifiability/lean/LeJEPA/Uniqueness.lean

@@ -0,0 +1,139 @@
+import Mathlib.Analysis.SpecialFunctions.Log.Basic
+import Mathlib.Analysis.SpecialFunctions.Pow.Real
+
+/-!
+# Gaussian Uniqueness (Proposition: Converse Direction)
+
+  The first non-constant eigenfunction of the transition operator
+  is affine **if and only if** p is Gaussian.
+
+  ## Verification status
+
+  | Component                              | Status      |
+  |----------------------------------------|-------------|
+  | SL eigenfunction equation              | structural  |
+  | Score slope negativity (−ev/K < 0)     | VERIFIED    |
+  | Affine eigenfunction → affine score    | VERIFIED    |
+  | Affine score → Gaussian density        | axiomatized |
+  | Only-if assembly                       | VERIFIED    |
+  | Gaussian → Hermite eigenfunctions      | axiomatized |
+  | If assembly                            | VERIFIED    |
+  | Full biconditional                     | VERIFIED    |
+  | Zero-mean specialization               | VERIFIED    |
+-/
+
+set_option maxHeartbeats 400000
+
+noncomputable section
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- STURM–LIOUVILLE STRUCTURE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- A scalar latent component under constant diffusion K > 0. -/
+structure LatentComponent where
+  K : ℝ
+  hK : 0 < K
+  score : ℝ → ℝ        -- (log p)'
+  ev : ℝ               -- first non-constant eigenvalue λ₁
+  hev : 0 < ev
+
+/-- Score corresponds to a Gaussian: ∃ α < 0, β, score(z) = αz + β. -/
+def IsGaussianScore (score : ℝ → ℝ) : Prop :=
+  ∃ α β : ℝ, α < 0 ∧ ∀ z, score z = α * z + β
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- AXIOMATIZED
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Affine score → Gaussian** (axiomatized): integrating
+    score(z) = αz + β gives log p = (α/2)z² + βz + C. -/
+axiom gaussian_of_affine_score (score : ℝ → ℝ) (α β : ℝ)
+    (hα : α < 0) (hscore : ∀ z, score z = α * z + β) :
+    IsGaussianScore score
+
+/-- **Gaussian → affine eigenfunction** (axiomatized): Gaussian
+    density ⟹ SL eigenfunctions are Hermite polynomials ⟹
+    first non-constant eigenfunction is He₁(z) = z. -/
+axiom hermite_first_eigenfunction_of_gaussian
+    (lc : LatentComponent) (hgauss : IsGaussianScore lc.score) :
+    ∃ (a b : ℝ), a ≠ 0 ∧
+      ∀ z, lc.K * lc.score z * a = -(lc.ev * (a * z + b))
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: AFFINE EIGENFUNCTION → AFFINE SCORE
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Core algebraic step** (VERIFIED):
+    K · score(z) · a = −ev·(az + b)  with a ≠ 0
+    ⟹ score(z) = (−ev/K)z + (−ev·b/(Ka)),  slope < 0. -/
+theorem score_affine_of_eigenfunction
+    (lc : LatentComponent) (a b : ℝ) (ha : a ≠ 0)
+    (heigen : ∀ z, lc.K * lc.score z * a = -(lc.ev * (a * z + b))) :
+    ∃ (α β : ℝ), α < 0 ∧ (∀ z, lc.score z = α * z + β) := by
+  refine ⟨-(lc.ev / lc.K), -(lc.ev * b / (lc.K * a)), ?_, ?_⟩
+  · -- −ev/K < 0 since ev > 0 and K > 0
+    have := div_pos lc.hev lc.hK
+    linarith
+  · intro z
+    have hK_ne : lc.K ≠ 0 := ne_of_gt lc.hK
+    have hKa_ne : lc.K * a ≠ 0 := mul_ne_zero hK_ne ha
+    have h := heigen z
+    -- Isolate score(z): divide by K·a
+    have h1 : lc.score z = -(lc.ev * (a * z + b)) / (lc.K * a) := by
+      field_simp at h ⊢; linarith
+    rw [h1]; field_simp; ring
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: ONLY-IF ASSEMBLY
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Only-if** (VERIFIED): affine eigenfunction ⟹ Gaussian. -/
+theorem gaussian_of_affine_eigenfunction
+    (lc : LatentComponent) (a b : ℝ) (ha : a ≠ 0)
+    (heigen : ∀ z, lc.K * lc.score z * a = -(lc.ev * (a * z + b))) :
+    IsGaussianScore lc.score := by
+  obtain ⟨α, β, hα_neg, hscore⟩ :=
+    score_affine_of_eigenfunction lc a b ha heigen
+  exact gaussian_of_affine_score lc.score α β hα_neg hscore
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: FULL BICONDITIONAL
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Gaussian uniqueness** (VERIFIED):
+    First eigenfunction is affine ⟺ p is Gaussian. -/
+theorem gaussian_uniqueness (lc : LatentComponent) :
+    (IsGaussianScore lc.score →
+      ∃ (a b : ℝ), a ≠ 0 ∧
+        ∀ z, lc.K * lc.score z * a = -(lc.ev * (a * z + b)))
+    ∧
+    (∀ (a b : ℝ), a ≠ 0 →
+      (∀ z, lc.K * lc.score z * a = -(lc.ev * (a * z + b))) →
+      IsGaussianScore lc.score) :=
+  ⟨hermite_first_eigenfunction_of_gaussian lc,
+   fun a b ha heigen => gaussian_of_affine_eigenfunction lc a b ha heigen⟩
+
+
+-- ═══════════════════════════════════════════════════════════════
+-- VERIFIED: ZERO-MEAN SPECIALIZATION
+-- ═══════════════════════════════════════════════════════════════
+
+/-- **Zero mean** (VERIFIED): with b = 0, a = 1,
+    score(z) = −(ev/K)·z. -/
+theorem score_pure_linear_zero_mean
+    (lc : LatentComponent)
+    (heigen : ∀ z, lc.K * lc.score z * 1 = -(lc.ev * (1 * z + 0))) :
+    ∀ z, lc.score z = -(lc.ev / lc.K) * z := by
+  intro z
+  have hK_ne : lc.K ≠ 0 := ne_of_gt lc.hK
+  have h := heigen z
+  simp only [mul_one, add_zero] at h
+  field_simp; linarith
+
+end

JEPA/lejepa-identifiability/lean/lake-manifest.json

@@ -0,0 +1,95 @@
+{"version": "1.1.0",
+ "packagesDir": ".lake/packages",
+ "packages":
+ [{"url": "https://github.com/leanprover-community/mathlib4",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "8f9d9cff6bd728b17a24e163c9402775d9e6a365",
+   "name": "mathlib",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "v4.28.0",
+   "inherited": false,
+   "configFile": "lakefile.lean"},
+  {"url": "https://github.com/leanprover-community/plausible",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "55c8532eb21ec9f6d565d51d96b8ca50bd1fbef3",
+   "name": "plausible",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "main",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/LeanSearchClient",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "c5d5b8fe6e5158def25cd28eb94e4141ad97c843",
+   "name": "LeanSearchClient",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "main",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/import-graph",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "85b59af46828c029a9168f2f9c35119bd0721e6e",
+   "name": "importGraph",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "main",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/ProofWidgets4",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "be3b2e63b1bbf496c478cef98b86972a37c1417d",
+   "name": "proofwidgets",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "v0.0.87",
+   "inherited": true,
+   "configFile": "lakefile.lean"},
+  {"url": "https://github.com/leanprover-community/aesop",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "f642a64c76df8ba9cb53dba3b919425a0c2aeaf1",
+   "name": "aesop",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "master",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/quote4",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "b8f98e9087e02c8553945a2c5abf07cec8e798c3",
+   "name": "Qq",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "master",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/batteries",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "495c008c3e3f4fb4256ff5582ddb3abf3198026f",
+   "name": "batteries",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "main",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover/lean4-cli",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover",
+   "rev": "4f10f47646cb7d5748d6f423f4a07f98f7bbcc9e",
+   "name": "Cli",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "v4.28.0",
+   "inherited": true,
+   "configFile": "lakefile.toml"}],
+ "name": "lejepa",
+ "lakeDir": ".lake"}

JEPA/lejepa-identifiability/lean/lakefile.lean

@@ -0,0 +1,12 @@
+import Lake
+open Lake DSL
+
+package lejepa where
+  leanOptions := #[
+    ⟨`autoImplicit, false⟩
+  ]
+
+@[default_target]
+lean_lib LeJEPA where
+
+require "leanprover-community" / "mathlib" @ git "v4.28.0"

JEPA/lejepa-identifiability/lean/lean-toolchain

@@ -0,0 +1 @@
+leanprover/lean4:v4.28.0

JEPA/lejepa-identifiability/requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0
+numpy
+scipy
+scikit-learn
+matplotlib
+pandas
+pyyaml