Add DiG (DINO-embedded Gaussians) feature training

[connorsoohoo]

Summary

Implements the DiG model from Robot See Robot Do (reference implementation: kerrj/dig, MIT) natively in Brush: per-Gaussian DINO feature embeddings trained alongside RGB, useful for downstream part segmentation and tracking. The reference stack is CUDA-only (nerfstudio/gsplat/tiny-cuda-nn); this port runs on every Brush backend, including Metal and the web.

Everything is explicit opt-in via --dino — with the flag off there is no change to training, and the feature pass adds no cost. No new Rust dependencies.

What's in it

• Feature rasterization (brush-render / brush-render-bwd): a new forward kernel reusing the existing projection/sort pipeline; the 64-d per-Gaussian feature vector is read from global memory only for contributing splats, so the feature dimension isn't bound by the 32 KiB threadgroup budget. The backward is a forward-order replay scattering gradients via the existing AtomicAddF32 (CAS fallback on Metal). Geometry is detached in the feature pass — matching the reference — so this is a single-input autodiff custom op with no changes to the RGB backward.
• Dataset (brush-dataset): per-view .npy feature-map loading (--features-dir-name, default dino_features), with a minimal pure-Rust npy reader.
• Training (brush-train): per-Gaussian [N, 64] feature table + shared no-bias MLP decoder (64→64→64→96), DINO MSE at 5× the GT feature resolution with bilinear-upsampled GT, 3-NN feature-variance regularizer after step 1000 (CPU grid-hash kNN, recomputed on refine), the reference LR schedule (1e-2→1e-3 over 6k steps), and prune/split remapping of the feature table + its Adam state in refine. All reference constants are CLI flags with the reference values as defaults (--dino-feature-dim, --dino-rescale-factor, --dino-lr/--dino-lr-end, --dino-nn-reg-weight).
• Export (brush-process): <name>_dig_features.npy (rows match PLY order) + <name>_dig_mlp.json alongside every exported PLY, so features survive round-trips out of Brush.
• Viewer: a live "DINO feature view" toggle in the scene controls recolors splats by their learned features (MLP-decoded, top-3 PCA channels → RGB), refreshed every 50 training steps; --dino-view starts with it enabled.
• Preprocessing: scripts/extract_dino_features.py — DINOv2 → PCA cache, numerically matching the reference format; runs on MPS/CUDA/CPU; self-contained via PEP 723 (uv run needs no env setup).
• Docs: docs/dig-port/dig-port-usage.md — end-to-end usage; a short README section.

Tests

crates/brush-bench-test/tests/dig_features.rs:

• feature render reproduces the RGB rasterizer bit-for-bit (max diff < 1e-3) when features are set to the degree-0 colors
• analytic feature gradients match central finite differences
• multi-step training with feature supervision + refine (prune/split) smoke test

Full workspace test suite, clippy --all-targets -D warnings, cargo fmt --check, and a wasm32 cargo check pass locally (macOS/Metal).

Known limitations

• LOD generation (--lod-levels > 0) resets the trainer between levels and isn't supported with feature training (--dino + --lod-levels errors out with a clear message).
• GARField affinity training, interactive click-to-segment, and camera-pose optimization are out of scope here.

Happy to split this up or adjust conventions (flag names, doc placement) however you prefer.

🤖 Generated with Claude Code

[connorsoohoo]