LangSplatV2: Training pipeline for sparse language feature fields

[swahtz]

Summary

Implements the LangSplatV2 training pipeline for learning per-Gaussian sparse language feature fields using fVDB and fvdb-reality-capture. This builds on the preprocessing transforms already merged in #42, adding the model, loss, training loop, and supporting utilities needed to train language-aware Gaussian splats. closes #32

Key components:

• Model (model.py): LangSplatV2Model wraps a frozen GaussianSplat3d with learnable per-Gaussian logits and codebooks. Renders sparse coefficient weight maps via splatting and decodes them into dense CLIP feature maps through codebook lookup.
• Vector quantization (vq_utils.py): Implements softmax_to_topk_soft_code for efficient sparse coefficient generation and ResidualVectorQuantization for K-means codebook initialization from ground-truth CLIP features.
• Loss (loss.py): Cosine similarity and L1 losses with per-pixel masking for regions without valid language features.
• Dataset (training/dataset.py): LangSplatV2Dataset loads pre-computed CLIP features and segmentation maps in compact form. Dense ground-truth feature maps are materialized on-device after transfer using build_feature_map, avoiding large CPU-to-GPU transfers.
• Training runner (training/trainer.py): LangSplatV2Training handles the full workflow — dataset construction, K-means codebook initialization, optimizer setup, training/eval loops with gradient accumulation, and checkpointing.
• Config (config.py): Extended with LangSplatV2ModelConfig and LangSplatV2TrainingConfig dataclasses.
• Entry point (train_langsplatv2.py): CLI script using tyro for launching training.

Performance optimizations

• Compact feature storage with JaggedTensor for variable-length per-image features, avoiding padding overhead
• GPU-side dense feature map construction (build_feature_map) using torch.empty to eliminate costly zero-fill of ~4 GB tensors

Test plan

• Run training on a preprocessed scene: python train_langsplatv2.py --scene-dir <path> --checkpoint-dir <path>
• Verify K-means codebook initialization completes and logs cluster info
• Confirm training loop progresses without OOM on a single GPU (tested with 1080p images)
• Check that checkpoints are saved and can be resumed
• Profile with Nsight Systems to verify no unexpected data transfer bottlenecks

[Copilot]

Pull request overview

Adds the LangSplatV2 training pipeline to the langsplatv2 package, enabling learning per-Gaussian sparse language feature fields (decoded via shared codebooks) on top of a frozen fVDB Gaussian splat reconstruction.

Changes:

• Introduces LangSplatV2Model, VQ/codebook initialization utilities, and masked CLIP-feature losses.
• Adds dataset + dataloader utilities for compact CLIP feature storage and GPU-side dense feature map materialization.
• Implements a full training runner (init, train/eval loops, checkpointing) and a tyro-based CLI entrypoint.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
open_vocabulary_segmentation/langsplatv2/train_langsplatv2.py	Adds CLI entrypoint for launching LangSplatV2 training.
open_vocabulary_segmentation/langsplatv2/testers/dataset_tester.py	Adds a local dataset smoke-test script (currently non-portable).
open_vocabulary_segmentation/langsplatv2/testers/data_tester.py	Adds a local preprocessing smoke-test script (currently non-portable).
open_vocabulary_segmentation/langsplatv2/pyproject.toml	Adds runtime deps needed for training (e.g., scikit-learn, tyro).
open_vocabulary_segmentation/langsplatv2/langsplatv2/vq_utils.py	Adds top-k soft code helper + RVQ K-means codebook initialization.
open_vocabulary_segmentation/langsplatv2/langsplatv2/util.py	Adds helpers for loading SfM scenes and GS reconstructions.
open_vocabulary_segmentation/langsplatv2/langsplatv2/training/trainer.py	Implements the end-to-end training runner (init/train/eval/checkpoint).
open_vocabulary_segmentation/langsplatv2/langsplatv2/training/dataset.py	Adds dataset + collate + GPU dense feature map builder + infinite sampler.
open_vocabulary_segmentation/langsplatv2/langsplatv2/training/init.py	Exposes training APIs from the package.
open_vocabulary_segmentation/langsplatv2/langsplatv2/model.py	Implements the LangSplatV2 model (render sparse weights, decode via codebooks).
open_vocabulary_segmentation/langsplatv2/langsplatv2/loss.py	Adds masked cosine/L1 losses for CLIP feature supervision.
open_vocabulary_segmentation/langsplatv2/langsplatv2/config.py	Extends configs with model + training dataclasses.
open_vocabulary_segmentation/langsplatv2/langsplatv2/init.py	Re-exports key configs/model/loss for package users.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.