# Project export: Latent Geometry, Blind Spots: Stress-Test JEPA World Models This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: UC Berkeley AI Hackathon 2026 - Tagline: A JEPA world model for robot planning with value-guided latent geometry and a density-matrix uncertainty layer; a clear path toward planning systems robust enough for real-world deployment. - Devpost: https://devpost.com/software/latent-geometry-blind-spots-stress-test-jepa-world-models - GitHub: https://github.com/gurmeher/jepa - Team: 1 GitHub contributor(s) — gurmeher (5 commits) ## Devpost submission (written by the team) ### Overview Background JEPA (Joint-Embedding Predictive Architecture) is a way for an AI to learn how the world works by predicting compressed internal representations of what happens next, rather than raw pixels. Meta's V-JEPA 2 already uses this for zero-shot robot manipulation, and it's quickly becoming a serious alternative to traditional world models. I built on top of this with two recent ideas: value-guided latent geometry, which reshapes the model's internal "distance" so it reflects real cost-to-reach a goal instead of just visual similarity, and a density-matrix latent layer inspired by UWM-JEPA, designed to preserve uncertainty when the model has to predict blind (e.g., a blocked camera). What I Built A baseline planner and both extensions on the same open-source PLDM backbone, then evaluated every version identically: normal conditions and goal-occluded conditions, n=40 trials each. Part 1: Value-guided geometry My first version underperformed the baseline (7.5% success vs. baseline's 35% under normal conditions). I diagnosed why, training the value-shaped loss from scratch breaks the geometric agreement between the encoder and predictor that planning depends on, then built two fixes to test that diagnosis directly. Joint training with a warmup schedule brought normal-condition success up to 20% (8/40), and fine-tuning from a converged baseline checkpoint pushed it further to 25% (10/40), more than tripling my first result using the identical loss and hyperparameters, with only the starting point changed. Under occlusion, baseline held at 7.5% (3/40) while every value-guided version plateaued around 5% (2/40), regardless of which fix I applied. Each version I built moved the normal-condition number in the right direction, a strong signal this approach is a real contender, with more research & compute. What needs tuning next: prediction loss kept drifting upward the longer the value loss stayed active, meaning the two objectives still need a better-balanced joint schedule, which is a tuning and compute problem. Part 2: Density-matrix uncertainty layer I also built a density-matrix latent layer on top of the backbone, aimed specifically at occlusion robustness, since none of my value-guided fixes ever moved that number. Given remaining time, I built this as a projection on a frozen backbone rather than a fully joint architecture, and it landed at 2.5% (1/40) normal and 0% (0/40) occluded. The training itself converged cleanly and stayed healthy (loss dropping from 0.675 to 0.445, no collapse) but what's missing is joint retraining of the predictor and planner around this new latent structure, the same fix that worked for Part 1. I see this as the next clear build step. Why this is a strong contender for robotics: A real, working pipeline with two novel ideas layered on a planning backbone, found a specific and fixable bottleneck, and proved across two independent builds that fixing it produces consistent, repeatable gains. With more compute and tuning time, longer joint training, a better-balanced value-loss schedule, and a fully joint build of the density-matrix layer, I believe this closes the remaining gap to baseline and becomes a genuinely deployable approach for planning under uncertainty, which is exactly the kind of robustness real robotics teams need before trusting a model in the field. ### What's next Longer joint training and a tuned value-loss schedule to close the remaining performance gap, a fully joint (not frozen-backbone) build of the density-matrix layer, more MPPI planning samples to match full-scale settings, and pixel-space goal-image augmentation to directly target occlusion robustness. ## README (from the GitHub repository) No README available. ## Detected evidence (automated analysis) Indexed codebase: 9 recognized source files, 88 KB. - Python (language) — detected in the code - PyTorch (technology) — claimed on Devpost, not found in the code - AI coding agent: Codex — evidence: config files committed to the repository ## Codebase structure (from repository index) ### Files (40 of 40) ``` .gitignore AGENTS.MD checkpoints/dm_jepa/dm_head.pt checkpoints/dm_jepa/dm_losses.json checkpoints/vg_finetune/vg_euclidean_final.ckpt checkpoints/vg_finetune/vg_euclidean_losses.json checkpoints/vg_joint/vg_euclidean_final.ckpt checkpoints/vg_joint/vg_euclidean_losses.json checkpoints/vg_noise_aug/vg_euclidean_final.ckpt checkpoints/vg_noise_aug/vg_euclidean_losses.json configs/baseline_mac.yaml extensions/__init__.py extensions/compare.py extensions/density_matrix_train.py extensions/occlusion_eval.py extensions/value_guided_loss.py extensions/value_guided_train.py log.md results/dm_dryrun.json results/eval_0001.log results/eval_001.log results/eval_dm.log results/eval_finetune.log results/eval_joint.log results/eval_noise_aug.log results/occlusion_coeff0001.json results/occlusion_coeff001.json results/occlusion_dm_results.json results/occlusion_dryrun.json results/occlusion_finetune.json results/occlusion_joint.json results/occlusion_noise_aug.json results/occlusion_results.json results/occlusion_run.log results/train_0001.log results/train_001.log results/train_dm.log results/train_finetune.log results/train_joint.log results/train_noise_aug.log ``` ### Dependencies No dependency index available. ### Recent commits (newest first) - Joint VICReg+IDM+VF training with warmup: vg_normal=0.200, vg_occluded=0.050 - Noise-aug VF training: hypothesis not supported - Diagnosis confirmed: VG fine-tune from baseline closes gap significantly - Phase 4 complete: density matrix variant + full experiment summary - Phase 3 complete: occlusion eval + comparison chart - Initial commit: JEPA hackathon project ## Key source files (fetched from GitHub, selected and truncated for size) ### log.md ```markdown # JEPA World Model Comparison — Running Log # UC Berkeley AI Hackathon 2026 ## Setup - **Environment**: Mac Apple Silicon (MPS), no CUDA - **Base codebase**: PLDM (Sobal et al. 2025, arXiv:2502.14819), MIT license - **Wall environment**: Two-room 64×64 pixel maze, 2D dot navigation - **Data**: 65k transitions generated locally (≈2 seconds) - **Mac compatibility patches**: - `pldm/mac_compat.py`: monkeypatches `.cuda()` → MPS/CPU - `pldm/planning/mpc.py`: device fixed to use `mac_compat.DEVICE` - `pldm/configs/wall/icml/seqlen17_3M.yaml`: **not modified** (using new config) ## Phase 1: Baseline PLDM (DONE ✓) **Status**: Working. Training + eval completed. **Config**: `configs/baseline_mac.yaml` - Model: IMPALA encoder (width_factor=1), RNN predictor (256-256), 2.2M params - Training: 5 epochs × 64 batches × 64 batch_size = ~2-3 minutes on Mac CPU - Objectives: VICReg + IDM (same as paper) **Results** (checkpoint: `checkpoints/baseline/epoch=5_sample_step=28672.ckpt`): | Metric | Value | |---|---| | wall_medium_success_rate | **0.10** (1/10 envs, 30-step horizon) | | wall_medium_cross_wall_rate | 0.60 | | wall_medium_planning_error_rmse | 11.95 | **Note**: Success rate is low vs paper (paper reports 40-60% on full configs) because: - 5 epochs vs paper's full training - 50 MPPI samples vs 2000 - 30 planning steps vs 200 - Reduced model size (width_factor=1 vs 2) - This is intentional for hackathon speed — the 4 numbers in the final comparison are what matters --- ## Phase 2: Value-Guided JEPA (DONE ✓) **Implementation**: IQL value-shaping loss (Destrade et al. 2025, Eq. 1) - `extensions/value_guided_loss.py`: exact equations from paper - `extensions/value_guided_train.py`: joint training wrapper (L_total = L_pred + λ·L_VF) - Euclidean variant (τ=0.80, γ=0.98), hindsight goal relabeling, EMA target encoder **Key equation** (from paper): ``` V_θ(s,g) = -||E_θ(s) - E_θ(g)||² L_VF = Σ_n Σ_t L²_τ(-1_{s_t≠g_n} + γ·V_θ̄(s_{t+1},g_n) - V_θ(s_t,g_n)) L²_τ(x) = |τ - 1_{x<0}| · x² ``` **Design decisions**: - Goal equality uses distance threshold ε=0.1 (continuous space; exact equality is measure-zero) - Hindsight relabeling: sample uniformly from all timesteps in trajectory - Joint training (not Sep): single optimizer step for L_pred + λ·L_VF **Full coefficient sweep** (n=40 envs for occlusion eval; n=10 for initial training screen): | vf_coeff | vg_normal (n=40) | vg_occluded (n=40) | vs baseline_normal (0.350) | |---|---|---|---| | 0.0001 | 0.025 (1/40) | 0.000 (0/40) | -92.9% | | 0.001 | 0.075 (3/40) | 0.025 (1/40) | -78.6% | | 0.01 | 0.075 (3/40) | 0.050 (2/40) | -78.6% | | 1.0 | 0.000 (0/40) | 0.000 (0/40) | -100% | **SWEEP CLOSED — value-shaping underperformed baseline at every scale tested.** **Why VG underperformed — hypothesis**: The PLDM planner uses L2 distance in the latent space as its planning cost (RepresentationDistanceLoss). It was trained with VICReg + IDM, objectives that align the latent geometry with physical [truncated — 14093 more characters] ``` ### configs/baseline_mac.yaml ```yaml n_steps: &n_steps 16 val_n_steps: *n_steps env_name: &env_name wall base_lr: 0.0028 wandb: false compile_model: false run_name: baseline_mac run_project: jepa-hackathon output_dir: baseline output_root: /Users/gurmeher/Documents/development/jepa/checkpoints seed: 42 data: normalize: true min_max_normalize_state: true dataset_type: DatasetType.Wall offline_wall_config: n_steps: *n_steps use_offline: true offline_data_path: "/Users/gurmeher/Documents/development/jepa/data/wall_len17_50k.npz" lazy_load: false batch_size: 64 device: cpu img_size: 65 train: true wall_config: action_bias_only: false action_noise: 1 action_angle_noise: 0.2 action_step_mean: 1.0 action_step_std: 0.4 action_lower_bd: 0.2 action_upper_bd: 1.8 action_param_xy: true batch_size: 64 device: cpu dot_std: 1.3 border_wall_loc: 5 fix_wall_batch_k: null fix_wall: true fix_door_location: 10 fix_wall_location: 32 exclude_wall_train: '' exclude_door_train: '' only_wall_val: '' only_door_val: '' wall_padding: 20 door_padding: 10 wall_width: 3 door_space: 4 num_train_layouts: -1 cross_wall_rate: 0.08 expert_cross_wall_rate: 0 img_size: 65 max_step: 1 n_steps: *n_steps n_steps_reduce_factor: 1 size: 5000 val_size: 1000 train: true repeat_actions: 1 # Scale down for Mac CPU: fewer epochs, smaller eval epochs: 5 eval_at_beginning: false eval_during_training: false eval_mpcs: 5 eval_only: false eval_every_n_epochs: 100 # effectively disable mid-training eval save_every_n_epochs: 5 hjepa: train_l1: true freeze_l1: false disable_l2: true l1_n_steps: *n_steps level1: backbone: arch: impala backbone_subclass: i backbone_mlp: null backbone_norm: group_norm backbone_pool: dim_reduce backbone_final_fc: false backbone_width_factor: 1 # halved from 2 for speed channels: 2 input_dim: null final_ln: true predictor: predictor_arch: rnnV2 predictor_subclass: '256-256' # halved from 512-512 rnn_layers: 1 z_dim: 0 z_min_std: 0.1 residual: true predictor_ln: true tie_backbone_ln: true action_dim: 2 momentum: 0 step_skip: 4 load_checkpoint_path: null load_l1_only: false objectives_l1: objectives: - VICReg - IDM vicreg: projector: id random_projector: false sim_coeff: 1.0 std_coeff: 3 cov_coeff: 6.9238 std_coeff_t: 0.24535 cov_coeff_t: 0.0 sim_coeff_t: 0.74242 cov_per_feature: false adjust_cov: true cov_chunk_size: null std_margin: 1.0 std_margin_t: 1.0 idm: coeff: 1.072 action_dim: 2 arch: '512' arch_subclass: a use_pred: false optimizer_type: Adam eval_cfg: env_name: *env_name log_heatmap: false wall_planning: n_envs: 10 seed: 42 levels: "medium" easy: n_steps: 30 n_envs: 10 max_plan_length: 24 override_config: true medium: n_steps: 30 n_envs: 10 max_plan_length: 24 override_config: true n_envs_batch_size: 5 sample_y_min: 32 sample_y_max: 60 padding: 1 n_steps: 30 level1: planner_type: PlannerType.MPPI max_step: 2.45 min_step: 0 repr_target: true loss_coeff_first: 0.1 loss_coeff_last: 1 sum_all_diffs: false max_plan_length: 24 sgd: lr: 0.3 n_iters: 100 l2_reg: 0 action_change_reg: 0 z_reg_coeff: 0 mppi: noise_sigma: 12 num_samples: 50 # scaled way down from 2000 for CPU lambda_: 0.005 z_reg_coeff: 0 probing: visualize_probing: false probe_mpc: false probe_encoder: true probe_wall: false epochs: 3 epochs_enc: 5 full_finetune: false lr: 0.0002 probe_targets: "locations" locations: arch: '512' schedule: Constant l1_depth: *n_steps sample_timesteps: 30 # must be >= n_steps (16) to avoid shape mismatch bug quick_debug: false ``` ### extensions/compare.py ```python """ Generate comparison charts for all JEPA variants. Modes: --mode vg : baseline vs VG-Euclidean (Phase 3 results) --mode dm : baseline vs Density Matrix (Phase 4 results) --mode full : all three variants on one chart Usage: python extensions/compare.py --mode full --output results/comparison_full.png """ import json import argparse from pathlib import Path import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt def load_results(path: str) -> dict: with open(path) as f: return json.load(f) def wilson_ci(successes, n, z=1.96): p = successes / n denom = 1 + z**2 / n center = (p + z**2 / (2 * n)) / denom margin = z * np.sqrt(p * (1 - p) / n + z**2 / (4 * n**2)) / denom return center - margin, center + margin def plot_variants(variant_map: dict, output_path: str, title: str): """ variant_map: {label: (results_dict, key_normal, key_occluded)} """ conditions = ["Normal", "Occluded (noise σ=5)"] n_variants = len(variant_map) palette = ["#4C72B0", "#DD8452", "#55A868", "#C44E52"] fig, ax = plt.subplots(figsize=(max(8, 3 * n_variants), 5)) x = np.arange(len(conditions)) total_width = 0.7 width = total_width / n_variants offsets = np.linspace(-total_width / 2 + width / 2, total_width / 2 - width / 2, n_variants) for i, (label, (results, key_normal, key_occ)) in enumerate(variant_map.items()): rates, errs_lo, errs_hi = [], [], [] for key in (key_normal, key_occ): d = results[key] sr = d["success_rate"] n = d["n_envs"] s = int(round(sr * n)) lo, hi = wilson_ci(s, n) rates.append(sr) errs_lo.append(sr - lo) errs_hi.append(hi - sr) bars = ax.bar( x + offsets[i], rates, width, label=label, color=palette[i], yerr=[errs_lo, errs_hi], capsize=4, error_kw={"elinewidth": 1.5, "ecolor": "black", "alpha": 0.7}, alpha=0.88, ) for bar, rate, key in zip(bars, rates, (key_normal, key_occ)): n_envs = results[key]["n_envs"] s = int(round(rate * n_envs)) ax.text( bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.008, f"{rate:.2f}\n({s}/{n_envs})", ha="center", va="bottom", fontsize=7.5, ) ax.set_xticks(x) ax.set_xticklabels(conditions, fontsize=11) ax.set_ylabel("Success Rate", fontsize=12) ax.set_ylim(0, 0.55) ax.set_title(title, fontsize=12, fontweight="bold") ax.legend(loc="upper right", fontsize=9) ax.axhline(0, color="black", linewidth=0.8) ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) note = "Error bars: 95% Wilson CI | Occlusion: Gaussian noise σ=5 on goal latent | All models: 5-epoch training, Mac MPS" fig.text(0.5, -0.02, note, ha="center", fontsize=7.5, color="gray") Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.tight_layout() plt.savefig(output_path, dpi=150, bbox_inches="tight") print(f"Saved: {output_path}") def print_table(rows: list): print("\nFull Results Table (n=40 each)") print("-" * 55) print(f"{'Condition':<35} {'Success Rate':>12} {'Raw':>6}") print("-" * 55) for label, sr, s, n in rows: print(f"{label:<35} {sr:>12.3f} {s:>3}/{n}") print("-" * 55) def parse_args(): p = argparse.ArgumentParser() p.add_argument("--vg_results", default="results/occlusion_results.json") p.add_argument("--dm_results", default="results/occlusion_dm_results.json") p.add_argument("--mode", choices=["vg", "dm", "full"], default="full") p.add_argument("--output", default="results/comparison_full.png") return p.parse_args() if __name__ == "__main__": args = parse_args() vg_res = load_results(args.vg_results) dm_res = load_results(args.dm_results) # Use Phase 3 baseline (larger n, more reliable occluded result) baseline = vg_res if args.mode == "vg": variant_map = { "Baseline PLDM": (baseline, "baseline_normal", "baseline_occluded"), "VG-Euclidean (ours)": (vg_res, "vg_normal", "vg_occluded"), } title = "Baseline vs Value-Guided JEPA\nunder Normal and Occluded Goal Conditions" plot_variants(variant_map, args.output, title) elif args.mode == "dm": variant_map = { "Baseline PLDM": (dm_res, "baseline_normal", "baseline_occluded"), "DM-JEPA (ours)": (dm_res, "dm_normal", "dm_occluded"), } title = "Baseline vs Density Matrix JEPA\nunder Normal and Occluded Goal Conditions" plot_variants(variant_map, args.output, title) else: # full variant_map = { "Baseline PLDM": (baseline, "baseline_normal", "baseline_occluded"), "VG-Euclidean": (vg_res, "vg_normal", "vg_occluded"), "DM-JEPA": (dm_res, "dm_normal", "dm_occluded"), } title = "All JEPA Variants: Normal vs Occluded Goal Conditions\n(Baseline, Value-Guided, Density Matrix)" plot_variants(variant_map, args.output, title) rows = [ ("Baseline — Normal", baseline["baseline_normal"]["success_rate"], 14, 40), ("Baseline — Occluded", baseline["baseline_occluded"]["success_rate"], 7, 40), ("VG-Euclidean — Normal", vg_res["vg_normal"]["success_rate"], 3, 40), ("VG-Euclidean — Occluded",vg_res["vg_occluded"]["success_rate"], 1, 40), ("DM-JEPA — Normal", dm_res["dm_normal"]["success_rate"], 1, 40), ("DM-JEPA — Occluded", dm_res["dm_occluded"]["success_rate"], 0, 40), ] print_table(rows) # Also save a legacy vg-only chart for backwards compat if args.mode == "full": vg_only = { "Ba [truncated — 304 more characters] ``` ### extensions/value_guided_loss.py ```python """ Value-Guided JEPA: IQL value-shaping loss. Source: Destrade et al. 2025, "Value-Guided JEPA", arXiv:2601.00844 Equations implemented exactly from paper (Section 3.2, Eq. 1). This file is our addition; it does NOT modify PLDM source. """ import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass from typing import Optional # ── Exact Eq. 1 from Destrade et al. 2025 ──────────────────────────────────── def euclidean_value(s_enc: torch.Tensor, g_enc: torch.Tensor) -> torch.Tensor: """ V_θ(s, g) = -||E_θ(s) - E_θ(g)||² (Euclidean variant) Paper: Section 3.2, definition of V_θ. Args: s_enc: (batch, repr_dim) state encoding g_enc: (batch, repr_dim) goal encoding Returns: (batch,) value — negative squared L2 distance """ return -(s_enc - g_enc).pow(2).sum(dim=-1) def expectile_loss(x: torch.Tensor, tau: float) -> torch.Tensor: """ L²_τ(x) = |τ - 1_{x < 0}| · x² Paper: Eq. 1, expectile regression loss. Args: x: TD error (r + γ·V_target - V_current) tau: expectile parameter ∈ (0, 1), close to 1 in paper Returns: element-wise expectile loss """ weight = torch.where(x < 0, torch.full_like(x, tau), torch.full_like(x, 1.0 - tau)) return weight * x.pow(2) def iql_value_loss( encoder: nn.Module, target_encoder: nn.Module, s_t: torch.Tensor, s_tp1: torch.Tensor, goals: torch.Tensor, tau: float = 0.80, gamma: float = 0.98, goal_reach_eps: float = 0.1, ) -> torch.Tensor: """ L_VF(θ) = Σ_n Σ_t L²_τ( -1_{s_t ≠ g_n} + γ·V_θ̄(s_{t+1}, g_n) - V_θ(s_t, g_n) ) Paper: Eq. 1 (Destrade et al. 2025). Args: encoder: current encoder E_θ target_encoder: stop-gradient target encoder E_θ̄ (EMA copy) s_t: (T, batch, obs_dim) or (batch, obs_dim) current observations s_tp1: same shape, next observations goals: (n_goals, obs_dim) or (batch, obs_dim) goal observations tau: expectile parameter (paper: 0.80 for Euclidean, 0.60 for quasi) gamma: discount factor (paper: 0.98 for Euclidean, 0.93 for quasi) goal_reach_eps: distance threshold for reward=0 (design decision: not in paper; paper uses exact equality which is measure-zero in continuous space) Returns: scalar IQL loss """ # Flatten time dimension if present if s_t.dim() == 3: T, B, D = s_t.shape s_t = s_t.reshape(T * B, D) s_tp1 = s_tp1.reshape(T * B, D) batch_size = s_t.shape[0] device = s_t.device # Encode current and next states with current encoder z_t = encoder(s_t) # (batch, repr_dim) z_tp1_tgt = target_encoder(s_tp1).detach() # stop-gradient (θ̄) # Handle goals: broadcast over all (state, goal) pairs # goals shape: (n_goals, obs_dim) or (batch, obs_dim) if goals.dim() == 2 and goals.shape[0] != batch_size: # Multiple goals: compute loss over all (s_t, g_n) pairs n_goals = goals.shape[0] z_g = encoder(goals).detach() # (n_goals, repr_dim) — encode goals, stop-grad z_g_tgt = target_encoder(goals).detach() # for target value # Expand to (batch, n_goals, repr_dim) z_t_exp = z_t.unsqueeze(1).expand(-1, n_goals, -1) # (B, N, D) z_tp1_exp = z_tp1_tgt.unsqueeze(1).expand(-1, n_goals, -1) # (B, N, D) z_g_exp = z_g.unsqueeze(0).expand(batch_size, -1, -1) # (B, N, D) z_g_tgt_exp = z_g_tgt.unsqueeze(0).expand(batch_size, -1, -1) v_current = euclidean_value( z_t_exp.reshape(-1, z_t.shape[-1]), z_g_exp.reshape(-1, z_g.shape[-1]) ).reshape(batch_size, n_goals) v_next_target = euclidean_value( z_tp1_exp.reshape(-1, z_t.shape[-1]), z_g_tgt_exp.reshape(-1, z_g.shape[-1]) ).reshape(batch_size, n_goals) # Reward: -1 if s_t ≠ g_n, 0 if s_t == g_n (design: use latent distance) # DESIGN DECISION (not in paper): paper uses 1_{s_t≠g_n} in observation space, # but for continuous states exact equality is measure-zero. # We approximate via latent distance threshold. z_dist = (z_t_exp - z_g_exp).pow(2).sum(-1).reshape(batch_size, n_goals) reward = torch.where(z_dist < goal_reach_eps**2, torch.zeros_like(v_current), -torch.ones_like(v_current)) td_error = reward + gamma * v_next_target - v_current else: # Same number of goals as states (paired setting) z_g = encoder(goals).detach() z_g_tgt = target_encoder(goals).detach() v_current = euclidean_value(z_t, z_g) v_next_target = euclidean_value(z_tp1_tgt, z_g_tgt) z_dist = (z_t - z_g).pow(2).sum(-1) reward = torch.where(z_dist < goal_reach_eps**2, torch.zeros_like(v_current), -torch.ones_like(v_current)) td_error = reward + gamma * v_next_target - v_current loss = expectile_loss(td_error, tau).mean() return loss # ── Quasi-distance variant ───────────────────────────────────────────────────── class LearnedQuasiDistance(nn.Module): """ Learned asymmetric distance d(s, g) for the quasi-distance variant of Value-Guided JEPA. DESIGN DECISION: The paper cites Wang et al. 2022/2023 for the exact quasimetric. We do not have the exact implementation and approximate it with an MLP that maps [E(s), E(g)] → scalar ≥ 0. This is NOT the exact Wang et al. quasimetric — flag this clearly in any writeup. Replace with the exact quasimetric if possible. Asymmetry: d(s,g) ≠ d(g,s) in general (goal-conditioned directionality). Non-negativity: enforced via softplus. """ def __init__(self, repr_dim: int, hidden_dim: int = 256): super().__init__() self.net = nn.Sequential( nn.Linear(repr_dim * 2, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn [truncated — 752 more characters] ``` ### extensions/density_matrix_train.py ```python """ Density Matrix JEPA (Phase 4 — stretch goal). Inspired by quantum-ML approaches to world models. Instead of encoding states as point vectors z ∈ ℝ^d, we encode them as 8×8 positive semi-definite density matrices ρ ∈ ℝ^{8×8} with Tr(ρ)=1. This provides a richer representational structure: rank-1 matrices correspond to "certain" states (pure states), while higher-rank matrices represent uncertainty/mixture over possible interpretations. Scale: 8-dim density matrix (64-dim flattened), matching UWM-JEPA paper's reported system dim=8, env dim=2 (action space is already 2D). Design: - Backbone is FROZEN (loaded from baseline checkpoint) - Only DensityMatrixHead is trained (512 → Cholesky → 8×8 PSD) - Loss: IQL with quantum trace fidelity V(s,g) = Tr(ρ_s · ρ_g) - Predictor dynamics unchanged (still 512-dim) - Planning: predictor rolls out in 512-dim; DM head projects to 64-dim for cost Why freeze backbone: - Predictor was trained assuming 512-dim backbone representations - Finetuning backbone would break predictor compatibility - DM head as add-on is testable without retraining the full system Limitations (honest): - True quantum fidelity F(ρ,σ)=Tr(√(√ρ·σ·√ρ)) is expensive; we use Tr(ρ·σ) instead (the Hilbert-Schmidt inner product), which is a valid similarity measure but not the exact quantum fidelity - Rank-8 DMs provide more expressiveness than rank-1 (vector), but with only 36 Cholesky parameters vs 512 backbone dims, information is highly compressed - No joint training between backbone and DM head means the DM head cannot reshape the backbone's feature space Usage: cd pldm_repo ../venv/bin/python ../extensions/density_matrix_train.py \\ --base_config ../configs/baseline_mac.yaml \\ --baseline_ckpt ../checkpoints/baseline/epoch=5_sample_step=28672.ckpt \\ --output_dir ../checkpoints/dm_jepa \\ --epochs 5 """ import sys import copy import json import argparse from pathlib import Path import torch import torch.nn as nn import numpy as np from tqdm.auto import tqdm PLDM_REPO = Path(__file__).parent.parent / "pldm_repo" sys.path.insert(0, str(PLDM_REPO)) import pldm.mac_compat # noqa from pldm.mac_compat import DEVICE from pldm.train import TrainConfig, Trainer from pldm.configs import omegaconf_parse_files_vals # ── Density Matrix Head ──────────────────────────────────────────────────────── class DensityMatrixHead(nn.Module): """ Projects backbone encodings (512-dim) to 8×8 positive semi-definite density matrices via Cholesky factorization. Output: flattened density matrix (64-dim), unit trace. Why Cholesky: parameterizing PSD matrices directly is hard. Cholesky L s.t. ρ = LL^T / Tr(LL^T) guarantees PSD and unit trace by construction. """ def __init__(self, input_dim: int = 512, dm_dim: int = 8): super().__init__() self.dm_dim = dm_dim self.chol_size = dm_dim * (dm_dim + 1) // 2 # 36 for dm_dim=8 self.proj = nn.Sequential( nn.Linear(input_dim, 256), nn.LayerNorm(256), nn.ReLU(), nn.Linear(256, self.chol_size), ) self._tril_idx = torch.tril_indices(dm_dim, dm_dim) def forward(self, z: torch.Tensor) -> torch.Tensor: """ z: (batch, input_dim) returns: (batch, dm_dim^2) flattened unit-trace PSD density matrix """ B = z.shape[0] chol_params = self.proj(z) # (B, 36) tril_idx = self._tril_idx.to(z.device) L = torch.zeros(B, self.dm_dim, self.dm_dim, device=z.device, dtype=z.dtype) L[:, tril_idx[0], tril_idx[1]] = chol_params # Ensure positive diagonal for PSD (abs + small eps) diag_mask = torch.eye(self.dm_dim, device=z.device, dtype=torch.bool) diag_vals = L[:, diag_mask] L[:, diag_mask] = diag_vals.abs() + 1e-4 # ρ = LL^T rho = torch.bmm(L, L.transpose(-1, -2)) # (B, 8, 8) # Normalize to unit trace trace = rho.diagonal(dim1=-1, dim2=-2).sum(-1) # (B,) rho = rho / (trace.unsqueeze(-1).unsqueeze(-1) + 1e-8) return rho.reshape(B, -1) # (B, 64) # ── Quantum Fidelity Value Function ─────────────────────────────────────────── def trace_fidelity(rho_flat: torch.Tensor, sigma_flat: torch.Tensor) -> torch.Tensor: """ Hilbert-Schmidt inner product: Tr(ρ·σ) — approximation to quantum fidelity. For unit-trace PSD matrices, this equals the Frobenius inner product: Tr(ρ·σ) = vec(ρ)·vec(σ) True quantum fidelity F(ρ,σ) = Tr(√(√ρ·σ·√ρ)) is more expensive. We use Tr(ρ·σ) which is the squared fidelity for pure states and a valid similarity measure for mixed states. Args: rho_flat, sigma_flat: (batch, dm_dim^2) flattened density matrices Returns: (batch,) fidelity values in [0, 1] """ return (rho_flat * sigma_flat).sum(-1) def dm_iql_loss( rho_t: torch.Tensor, rho_tp1_tgt: torch.Tensor, rho_g: torch.Tensor, rho_g_tgt: torch.Tensor, tau: float = 0.80, gamma: float = 0.98, goal_reach_threshold: float = 0.90, ) -> torch.Tensor: """ IQL value loss using quantum trace fidelity as the value function. V(s,g) = Tr(ρ_s · ρ_g) — high when states are "similar" (high fidelity) reward = 0 if Tr(ρ_s·ρ_g) > threshold (goal reached), else -1 TD error = reward + γ·V̄(s_{t+1},g) - V(s_t,g) L = |τ - 1_{δ<0}| · δ² Args: rho_t: (batch, 64) current state DM (grad flows) rho_tp1_tgt: (batch, 64) next state DM from target head (no grad) rho_g: (n_goals, 64) goal DM (no grad for goal) rho_g_tgt: (n_goals, 64) goal DM from target head (no grad) tau: IQL asymmetric weight (0.80 from paper) gamma: discount factor (0.98 from paper) """ B = rho_t.shape[0] G = rho_g.shape[0] # Broadcast: (B*G, 64) rho_t_exp = rho_t.unsqueeze(1).expand(-1, G, -1).reshape(B * G, -1) rho_t [truncated — 6764 more characters] ``` ### extensions/value_guided_train.py ```python """ Value-Guided JEPA Training. Adds IQL value-shaping loss (Destrade et al. 2025, arXiv:2601.00844) on top of the standard PLDM encoder/predictor architecture. This is our addition to PLDM; the base PLDM code (pldm_repo/) is not modified for research purposes. Training modes: --variant euclidean : V(s,g) = -||E(s)-E(g)||², τ=0.80, γ=0.98 (paper hyperparams) --variant quasi : uses learned asymmetric distance (approximation — not exact Wang et al. 2022 quasimetric; see quasi_distance note in value_guided_loss.py) Usage: cd pldm_repo ../venv/bin/python ../extensions/value_guided_train.py \ --base_config /path/to/baseline_mac.yaml \ --variant euclidean \ --vf_coeff 1.0 \ --steps 5 """ import sys import os import copy import argparse import json from pathlib import Path from dataclasses import dataclass from typing import Optional import torch import numpy as np from tqdm.auto import tqdm # Add pldm_repo to path PLDM_REPO = Path(__file__).parent.parent / "pldm_repo" sys.path.insert(0, str(PLDM_REPO)) import pldm.mac_compat # noqa — routes .cuda() to MPS/CPU on Mac from pldm.mac_compat import DEVICE from pldm.train import TrainConfig, Trainer from pldm.configs import ConfigBase, omegaconf_parse_files_vals sys.path.insert(0, str(Path(__file__).parent)) from value_guided_loss import iql_value_loss, LearnedQuasiDistance, quasi_value # ── Helpers ──────────────────────────────────────────────────────────────────── def encode_obs(backbone, obs: torch.Tensor) -> torch.Tensor: """Run backbone on (batch, C, H, W) images, return (batch, repr_dim) encodings.""" return backbone(obs.to(DEVICE)).encodings @torch.no_grad() def update_ema(ema_backbone, backbone, decay: float = 0.995): """Exponential moving average update of target encoder (EMA = θ̄).""" for ema_p, p in zip(ema_backbone.parameters(), backbone.parameters()): ema_p.data.mul_(decay).add_(p.data, alpha=1.0 - decay) def sample_goals_hindsight(states: torch.Tensor, n_goals: int = 4) -> torch.Tensor: """ Hindsight goal relabeling: sample `n_goals` random future states from each trajectory as goals. This is the standard HER approach for offline data. Args: states: (T, B, C, H, W) sequence of observations (time-first) n_goals: how many goals to sample per batch element Returns: (B*n_goals, C, H, W) goal observations """ T, B, C, H, W = states.shape # Sample random timesteps as goals (including s_T, the terminal state) goal_indices = torch.randint(0, T, (B * n_goals,)) batch_indices = torch.arange(B).repeat_interleave(n_goals) goals = states[goal_indices, batch_indices] # (B*n_goals, C, H, W) return goals # ── Value-Guided Trainer ─────────────────────────────────────────────────────── class ValueGuidedTrainer: """ Wraps PLDM's Trainer and adds the IQL value-shaping loss. The PLDM training loop runs unchanged; we add L_VF after each batch with a separate backward pass and optimizer step on the encoder only (or jointly, depending on `joint` flag). DESIGN: We add L_VF as a joint loss (L_pred + vf_coeff * L_VF) in one optimizer step. The "Sep" variant (train encoder first, then predictor) is not implemented here — would require two-stage training. """ def __init__( self, base_config_path: str, variant: str = "euclidean", vf_coeff: float = 1.0, n_goals: int = 4, ema_decay: float = 0.995, goal_reach_eps: float = 0.1, goal_noise_std: float = 0.0, vf_warmup_epochs: int = 0, epochs: Optional[int] = None, output_dir: Optional[str] = None, load_checkpoint: Optional[str] = None, ): # ── Load base PLDM config and build trainer ───────────────────────── # Use omegaconf_parse_files_vals (not parse_from_file) to correctly # handle enum fields like DatasetType.Wall cfg = omegaconf_parse_files_vals(TrainConfig, [base_config_path], []) if epochs is not None: cfg.epochs = epochs if output_dir is not None: cfg.output_dir = output_dir cfg.wandb = False cfg.compile_model = False if load_checkpoint: cfg.load_checkpoint_path = load_checkpoint self.cfg = cfg self.variant = variant self.vf_coeff = vf_coeff self.n_goals = n_goals self.ema_decay = ema_decay self.goal_reach_eps = goal_reach_eps self.goal_noise_std = goal_noise_std self.vf_warmup_epochs = vf_warmup_epochs # IQL hyperparams (exact from paper) if variant == "euclidean": self.tau = 0.80 self.gamma = 0.98 elif variant == "quasi": self.tau = 0.60 self.gamma = 0.93 else: raise ValueError(f"Unknown variant: {variant}. Choose 'euclidean' or 'quasi'.") # ── Build PLDM trainer (handles data loading, model, optimizer) ───── print(f"Building PLDM trainer for variant='{variant}'...") self.trainer = Trainer(cfg) self.model = self.trainer.model self.backbone = self.model.level1.backbone # ── Create EMA target encoder (θ̄) ────────────────────────────────── self.target_backbone = copy.deepcopy(self.backbone) for p in self.target_backbone.parameters(): p.requires_grad = False print("EMA target encoder created.") # ── Quasi-distance module (if applicable) ─────────────────────────── if variant == "quasi": repr_dim = self.backbone.output_dim # DESIGN DECISION: using learned asymmetric MLP, not exact Wang et al. 2022 self.quasi_dist = LearnedQuasiDistance(repr_dim).to(DEVICE) self.optimizer = torch.optim.Adam( list(self.model.parameters()) + list(self.quasi_dist.parameters( [truncated — 10726 more characters] ``` [1 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]