latent_visualization

latent_visualization

Latent-space visualization callback for variational policies.

LatentVisualizationCallback

LatentVisualizationCallback(log_every_n_epochs=5, max_samples=5000, label_keys=None)

Bases: Callback

Visualize latent spaces with optional metadata coloring.

Creates t-SNE/PCA projections of latent spaces. Label metadata is configured explicitly so synthetic modes, phase labels, task ids, or other categorical annotations can be used without task-specific callback logic.

Initialize latent visualization callback.

Parameters:

Name	Type	Description	Default
log_every_n_epochs	int	Log visualization every N epochs.	5
max_samples	int	Maximum samples for t-SNE (subsamples if exceeded).	5000
label_keys	list[str] | None	Metadata keys to use as categorical color labels.	None

Source code in src/versatil/training/callbacks/latent_visualization.py

def __init__(
    self,
    log_every_n_epochs: int = 5,
    max_samples: int = 5000,
    label_keys: list[str] | None = None,
) -> None:
    """Initialize latent visualization callback.

    Args:
        log_every_n_epochs: Log visualization every N epochs.
        max_samples: Maximum samples for t-SNE (subsamples if exceeded).
        label_keys: Metadata keys to use as categorical color labels.
    """
    super().__init__()
    self.log_every_n_epochs = log_every_n_epochs
    self.max_samples = max_samples
    self.label_keys = (
        label_keys
        if label_keys is not None
        else [
            MetadataKey.LATENT_COLOR_LABEL.value,
            MetadataKey.PHASE_LABEL.value,
        ]
    )

on_train_epoch_end

on_train_epoch_end(trainer, pl_module)

Create and log latent space visualization at end of training epoch.

Source code in src/versatil/training/callbacks/latent_visualization.py

def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Create and log latent space visualization at end of training epoch."""
    self._log_latent(
        trainer=trainer,
        metrics_accumulator=pl_module.train_metrics,
        split="train",
    )

on_validation_epoch_end

on_validation_epoch_end(trainer, pl_module)

Create and log latent space visualization at end of validation epoch.

Source code in src/versatil/training/callbacks/latent_visualization.py

def on_validation_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Create and log latent space visualization at end of validation epoch."""
    self._log_latent(
        trainer=trainer,
        metrics_accumulator=pl_module.val_metrics,
        split="val",
    )