lightning_policy

lightning_policy

PyTorch Lightning wrapper for Policy.

LightningPolicy

LightningPolicy(policy, training_config, total_training_steps=None)

Bases: LightningModule

PyTorch Lightning wrapper around Policy.

This wrapper handles: - Training and validation steps - Optimizer configuration with parameter groups - Learning rate scheduling - Metric accumulation and logging - Gradient clipping (via Lightning Trainer)

Initialize LightningPolicy.

Parameters:

Name	Type	Description	Default
policy	Policy	The policy to train	required
training_config	TrainingConfig	Training configuration	required
total_training_steps	int | None	Total number of training steps for LR scheduling. Calculated as ceil(len(train_loader) / gradient_accumulate_every) * num_epochs, because Lightning flushes the final partial accumulation window each epoch. If None, will use trainer.estimated_stepping_batches as fallback.	None

Source code in src/versatil/training/lightning_policy.py

def __init__(
    self,
    policy: Policy,
    training_config: TrainingConfig,
    total_training_steps: int | None = None,
):
    """Initialize LightningPolicy.

    Args:
        policy: The policy to train
        training_config: Training configuration
        total_training_steps: Total number of training steps for LR scheduling.
            Calculated as ceil(len(train_loader) / gradient_accumulate_every)
            * num_epochs, because Lightning flushes the final partial
            accumulation window each epoch.
            If None, will use trainer.estimated_stepping_batches as fallback.
    """
    super().__init__()
    self.policy = policy
    self.training_config = training_config
    self.total_training_steps = total_training_steps
    self.train_metrics = MetricsAccumulator()
    self.val_metrics = MetricsAccumulator()
    self.save_hyperparameters(ignore=["policy"])
    self._train_dataloader = None
    self._val_dataloader = None
    self.lr = None

on_train_epoch_start

on_train_epoch_start()

Record epoch start time for duration tracking.

Source code in src/versatil/training/lightning_policy.py

def on_train_epoch_start(self) -> None:
    """Record epoch start time for duration tracking."""
    self._epoch_start_time = time.monotonic()

training_step

training_step(batch, batch_idx)

Training step.

Parameters:

Name	Type	Description	Default
batch	dict[str, dict[str, Tensor]]	Batch dictionary with observations and actions	required
batch_idx	int	Batch index	required

Returns:

Type	Description
Tensor	Total loss tensor

Source code in src/versatil/training/lightning_policy.py

def training_step(
    self, batch: dict[str, dict[str, torch.Tensor]], batch_idx: int
) -> torch.Tensor:
    """Training step.

    Args:
        batch: Batch dictionary with observations and actions
        batch_idx: Batch index

    Returns:
        Total loss tensor
    """
    loss_output: LossOutput = self.policy.compute_loss(batch)
    batch_size = _batch_size_of(batch)
    self.train_metrics.add_loss_output(loss_output, batch_size=batch_size)
    # Log only on epoch to avoid batch size dependency in plots
    self.log(
        "train_loss",
        loss_output.total_loss,
        on_step=False,
        on_epoch=True,
        prog_bar=True,
        batch_size=batch_size,
    )
    return loss_output.total_loss

on_train_epoch_end

on_train_epoch_end()

Called at the end of training epoch to log accumulated metrics.

Source code in src/versatil/training/lightning_policy.py

def on_train_epoch_end(self) -> None:
    """Called at the end of training epoch to log accumulated metrics."""
    metrics = self.train_metrics.to_dict()
    self.log_dict(
        {f"train/{k}": v for k, v in metrics.items()},
        on_epoch=True,
        sync_dist=True,
    )
    self.train_metrics.reset()

    # Log epoch duration in seconds
    if hasattr(self, "_epoch_start_time"):
        epoch_duration = time.monotonic() - self._epoch_start_time
        self.log("train/epoch_time_seconds", epoch_duration, on_epoch=True)

    # Log peak GPU memory usage in GB, then reset for next epoch
    if torch.cuda.is_available() and self.device.type == "cuda":
        peak_memory_gb = torch.cuda.max_memory_allocated(device=self.device) / (
            1024**3
        )
        self.log("train/gpu_memory_peak_gb", peak_memory_gb, on_epoch=True)
        torch.cuda.reset_peak_memory_stats(device=self.device)

validation_step

validation_step(batch, batch_idx)

Validation step.

Parameters:

Name	Type	Description	Default
batch	dict[str, dict[str, Tensor]]	Batch dictionary with observations and actions	required
batch_idx	int	Batch index	required

Returns:

Type	Description
Tensor	Total loss tensor

Source code in src/versatil/training/lightning_policy.py

def validation_step(
    self, batch: dict[str, dict[str, torch.Tensor]], batch_idx: int
) -> torch.Tensor:
    """Validation step.

    Args:
        batch: Batch dictionary with observations and actions
        batch_idx: Batch index

    Returns:
        Total loss tensor
    """
    loss_output: LossOutput = self.policy.compute_loss(batch)
    batch_size = _batch_size_of(batch)
    self.val_metrics.add_loss_output(loss_output, batch_size=batch_size)
    self.log(
        "val_loss",
        loss_output.total_loss,
        on_step=False,
        on_epoch=True,
        prog_bar=True,
        batch_size=batch_size,
    )
    return loss_output.total_loss

on_validation_epoch_end

on_validation_epoch_end()

Called at the end of validation epoch to log accumulated metrics.

Source code in src/versatil/training/lightning_policy.py

def on_validation_epoch_end(self) -> None:
    """Called at the end of validation epoch to log accumulated metrics."""
    metrics = self.val_metrics.to_dict()
    self.log_dict(
        {f"val/{k}": v for k, v in metrics.items()},
        on_epoch=True,
        sync_dist=True,
    )
    self.val_metrics.reset()

configure_optimizers

configure_optimizers()

Configure optimizers and learning rate schedulers.

Uses Hydra's get_class() to resolve optimizer from config.target_class. Supports parameter groups with different learning rates.

Returns:

Type	Description
dict[str, Any]	Dictionary with optimizer and optional lr_scheduler

Source code in src/versatil/training/lightning_policy.py

def configure_optimizers(self) -> dict[str, Any]:
    """Configure optimizers and learning rate schedulers.

    Uses Hydra's get_class() to resolve optimizer from config.target_class.
    Supports parameter groups with different learning rates.

    Returns:
        Dictionary with optimizer and optional lr_scheduler
    """
    optimizer_config = self.training_config.optimizer
    param_groups = self._create_parameter_groups(optimizer_config)

    optimizer_config_omega = OmegaConf.structured(optimizer_config)
    optimizer_config_dict = OmegaConf.to_container(
        optimizer_config_omega, resolve=True
    )
    target = optimizer_config_dict.pop("target_class")
    optimizer_cls = get_class(target)
    # Remove custom field that's not passed to torch.optim
    optimizer_config_dict.pop("param_groups", None)
    optimizer = optimizer_cls(param_groups, **optimizer_config_dict)
    if self.training_config.lr_schedule is None:
        return {"optimizer": optimizer}
    scheduler_config = self._create_scheduler_config(optimizer)
    return {
        "optimizer": optimizer,
        "lr_scheduler": scheduler_config,
    }

on_load_checkpoint

on_load_checkpoint(checkpoint)

Called when loading a checkpoint.

Ensures that observation_space and action_space are converted from OmegaConf DictConfig to proper dataclass instances.

Parameters:

Name	Type	Description	Default
checkpoint	dict[str, Any]	The loaded checkpoint dictionary	required

Source code in src/versatil/training/lightning_policy.py

def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
    """Called when loading a checkpoint.

    Ensures that observation_space and action_space are converted from
    OmegaConf DictConfig to proper dataclass instances.

    Args:
        checkpoint: The loaded checkpoint dictionary
    """
    super().on_load_checkpoint(checkpoint)

forward

forward(obs_dict)

Forward pass for inference.

Parameters:

Name	Type	Description	Default
obs_dict	dict[str, Tensor]	Observation dictionary	required

Returns:

Type	Description
dict[str, Tensor]	Predicted actions

Source code in src/versatil/training/lightning_policy.py

def forward(self, obs_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Forward pass for inference.

    Args:
        obs_dict: Observation dictionary

    Returns:
        Predicted actions
    """
    return self.policy.predict_action(obs_dict)

set_dataloaders

set_dataloaders(train_loader, val_loader)

Attach the dataloaders returned by the Lightning dataloader hooks.

Parameters:

Name	Type	Description	Default
train_loader	DataLoader	Training dataloader.	required
val_loader	DataLoader | None	Optional validation dataloader.	required

Source code in src/versatil/training/lightning_policy.py

def set_dataloaders(
    self,
    train_loader: torch.utils.data.DataLoader,
    val_loader: torch.utils.data.DataLoader | None,
) -> None:
    """Attach the dataloaders returned by the Lightning dataloader hooks.

    Args:
        train_loader: Training dataloader.
        val_loader: Optional validation dataloader.
    """
    self._train_dataloader = train_loader
    self._val_dataloader = val_loader

train_dataloader

train_dataloader()

Return training dataloader for Lightning.

Source code in src/versatil/training/lightning_policy.py

def train_dataloader(self) -> torch.utils.data.DataLoader:
    """Return training dataloader for Lightning."""
    return self._train_dataloader

val_dataloader

val_dataloader()

Return validation dataloader for Lightning, or None if validation is disabled.

Source code in src/versatil/training/lightning_policy.py

def val_dataloader(self) -> torch.utils.data.DataLoader | None:
    """Return validation dataloader for Lightning, or None if validation is disabled."""
    return self._val_dataloader