Add DynamicSimilaritySampler — embedding-kNN batches with mutex constraint

Batches assembled from visually-similar drone queries pressure the model to
learn finer discriminative features. Random mutex batches average ~0.26
pairwise cosine similarity in query embedding space; DSS batches average
~0.71 — confirming the lookalikes grouping works as intended.

Algorithm per batch:
  1. Pick a random seed drone from the remaining pool.
  2. Rank the entire remaining pool by cosine similarity to the seed.
  3. Walk the ranking in descending order; add items whose sat_candidates
     don't collide with the batch's already-claimed set.
  4. Drop the seed if no valid batch can be assembled (rare mutex deadlock).

Inherits MutuallyExclusiveSampler semantics — no false negatives. Degrades
gracefully to mutex-only when no embeddings are set (warmup epochs, or if
`sampler_type="mutex"` is chosen).

Integration in `train_gtauav.py`:
  - New `_embed_drone_queries` helper: model.encode_query forwarded over
    GTAUAVDroneQuery, returns [N, D] CPU tensor. ~13s per 1024 queries on
    a 4090 → ~5 min for the full 25K train set.
  - Epoch loop re-embeds every `dss_reembed_every` epochs after a `dss_warmup_epochs`
    warmup (first epochs use mutex-only since untrained embeddings aren't
    informative for kNN).
  - Config: `sampler_type` ∈ {"mutex", "dss"}. Default flipped to "dss".

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

scripts/smoke_dss.py

@@ -0,0 +1,104 @@
+"""Smoke test for DynamicSimilaritySampler: re-embed + batch composition.
+
+Loads checkpoint, embeds a subset of train drones through model.encode_query,
+feeds embeddings to the sampler, and verifies:
+  - Sampler produces batches
+  - Mutex constraint preserved
+  - Batches differ from plain mutex (sorted by similarity to seed)
+"""
+import time
+import torch
+from torch.utils.data import DataLoader
+
+from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
+from src.datasets.gtauav_dataset import (
+    GTAUAVDataset,
+    GTAUAVDroneQuery,
+    collate_drone_query,
+)
+from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
+
+CKPT = "out/gtauav/with_text/ckpt_epoch005.pt"
+N_SUBSET = 1024  # small subset for smoke speed
+
+
+def main() -> None:
+    model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda")
+    model.eval()
+
+    tf = get_dino_transform(image_size=256)
+    ds = GTAUAVDataset(
+        pair_json="meta/train_80.json",
+        filter_meta="meta/seg_filter.json",
+        image_transform=tf,
+    )
+    # Subset the entries for smoke speed
+    ds.entries = ds.entries[:N_SUBSET]
+    print(f"Using {len(ds)} train entries")
+
+    query_ds = GTAUAVDroneQuery(ds)
+    loader = DataLoader(
+        query_ds, batch_size=32, shuffle=False, num_workers=2,
+        collate_fn=collate_drone_query, pin_memory=True,
+    )
+
+    # Embed
+    t0 = time.time()
+    embs = []
+    with torch.no_grad():
+        for batch in loader:
+            drone_img = batch["drone_img"].to("cuda", non_blocking=True)
+            q = model.encode_query(
+                drone_img,
+                batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+            )
+            embs.append(q.cpu())
+    emb = torch.cat(embs, dim=0)
+    print(f"Embedded {emb.shape[0]} queries in {time.time()-t0:.1f}s, dim={emb.shape[1]}")
+
+    # Build sampler
+    sat_cand = [e["sat_candidates"] for e in ds.entries]
+    sampler = DynamicSimilaritySampler(sat_cand, batch_size=8, seed=42)
+    sampler.set_epoch(0)
+
+    # Mutex fallback
+    mutex_batches = list(iter(sampler))
+    print(f"Mutex-only: {len(mutex_batches)} batches")
+
+    # DSS mode
+    sampler.update_embeddings(emb)
+    t0 = time.time()
+    dss_batches = list(iter(sampler))
+    print(f"DSS: {len(dss_batches)} batches in {time.time()-t0:.2f}s")
+
+    # Verify mutex invariant in DSS output
+    n_ok = 0
+    for b in dss_batches[:20]:
+        all_sats = set()
+        overlap = False
+        for idx in b:
+            s = set(ds.entries[idx]["sat_candidates"])
+            if s & all_sats:
+                overlap = True; break
+            all_sats |= s
+        if not overlap: n_ok += 1
+    print(f"Mutex preserved: {n_ok}/20 DSS batches clean")
+
+    # Similarity check: within a DSS batch, drones should be more similar
+    # to each other than drones in a random mutex batch.
+    def batch_mean_pairwise_sim(batch):
+        e = emb[batch]
+        e = torch.nn.functional.normalize(e, dim=-1)
+        sim = e @ e.t()
+        # Exclude diagonal
+        mask = ~torch.eye(len(batch), dtype=torch.bool)
+        return sim[mask].mean().item()
+
+    mutex_sim = sum(batch_mean_pairwise_sim(b) for b in mutex_batches[:20]) / 20
+    dss_sim = sum(batch_mean_pairwise_sim(b) for b in dss_batches[:20]) / 20
+    print(f"Mean in-batch cosine sim — mutex: {mutex_sim:.4f}  DSS: {dss_sim:.4f}")
+    print(f"  → DSS batches {'MORE' if dss_sim > mutex_sim else 'LESS'} visually similar (expected: MORE)")
+
+
+if __name__ == "__main__":
+    main()