1 lightning.one_step.module

lightning.one_step.module

Train the runnable one-step VIN candidate scorer with Lightning.

This module owns Lightning lifecycle, fitted-binner state, CORAL ordinal loss, coverage reweighting, candidate-table metrics, optimizer construction, and checkpoint persistence for aria_nbv.vin.types.VinPrediction. The scorer forward pass receives only actor-visible snippet evidence, world-frame candidate poses, and optional cached EVL fields. Oracle RRI labels and point-mesh diagnostics remain loss/evaluation targets on :class:aria_nbv.data_handling.VinOracleBatch and are never scorer inputs.

For a collated candidate table of shape B × N_q and K ordinal classes, the scorer emits logits Tensor["B N_q K-1", float32], probabilities Tensor["B N_q K", float32], and expected scores Tensor["B N_q", float32]. Training flattens only rows selected by the hard batch prefix mask, finite oracle labels, and finite logits. Scorer-produced candidate_valid is a diagnostic/coverage signal, not the authoritative action mask.

Full-shell Q_H replay is intentionally outside this owner. Rollout action and training masks are Tensor["S N_shell", bool] and reward tables are Tensor["S N_shell", float32]; invalid rows retain false masks and NaN labels. That objective requires masked bootstrap action selection, so the contract guard rejects the scorer family until a dedicated Lightning module exists.

1.1 Classes

Name	Description
VinLightningModule	Own one-step CORAL training state for a VIN candidate scorer.
VinLightningModuleConfig	Configure one-step ordinal candidate scoring and optimization.

1.1.1 VinLightningModule

lightning.one_step.module.VinLightningModule(config)

Own one-step CORAL training state for a VIN candidate scorer.

The module registers the scorer under vin to preserve historical checkpoint keys, loads or restores the fitted :class:RriOrdinalBinner, and accumulates stage metrics. For ordinal class label \(y \in [0,K-1]\), CORAL predicts threshold levels \(l_k = 1[y > k]\) for K-1 logits per candidate. Only the flattened M rows satisfying the batch hard mask, finite oracle label, and finite-logit checks contribute to loss.

VinPrediction.candidate_valid is intentionally diagnostic. The hard action/training mask is derived from VinOracleBatch.candidate_count; invalid rollout full-shell rows never enter this compact one-step batch.

1.1.1.1 Attributes

Name	Type	Description
config		Composed scorer, loss, optimizer, scheduler, binner, and logging policy.
vin		Registered scorer module that owns the vin.* checkpoint namespace.

1.1.1.2 Methods

Name	Description
setup	Initialize binner-derived scorer state for a Lightning stage.
on_save_checkpoint	Persist fitted ordinal-binner state inside a Lightning checkpoint.
on_load_checkpoint	Restore embedded ordinal-binner state before stage setup runs.
load_for_inference	Load a checkpoint and return an evaluation-ready scorer module.
prepare_for_inference	Initialize mandatory binner-derived state before direct inference.
training_step	Compute one masked CORAL update for an oracle-labelled training batch.
validation_step	Evaluate masked ordinal loss and candidate-table metrics without updates.
test_step	Evaluate the test split using the validation loss/mask contract.
on_train_epoch_end	Emit and reset accumulated training metrics at the epoch boundary.
on_validation_epoch_end	Emit/reset validation metrics and expected-RRI error statistics.
on_test_epoch_end	Emit and reset accumulated test metrics at the epoch boundary.
on_after_backward	Log configured scorer gradient norms after Lightning backpropagation.
configure_optimizers	Construct optimizer and optional scheduler from trainable scorer state.
summarize_vin	Summarize VIN inputs/outputs for a single oracle-labeled batch.
plot_vin_encodings_batch	Generate VIN encoding plots for a single oracle-labeled batch.

1.1.1.2.1 setup

lightning.one_step.module.VinLightningModule.setup(stage)

Initialize binner-derived scorer state for a Lightning stage.

Lightning owns invocation timing. This hook integrates logging, loads the configured binner when a checkpoint did not restore one, initializes ordinal class representatives and optional prior biases, then records the effective scorer config once per module instance.

1.1.1.2.2 on_save_checkpoint

lightning.one_step.module.VinLightningModule.on_save_checkpoint(checkpoint)

Persist fitted ordinal-binner state inside a Lightning checkpoint.

1.1.1.2.3 on_load_checkpoint

lightning.one_step.module.VinLightningModule.on_load_checkpoint(checkpoint)

Restore embedded ordinal-binner state before stage setup runs.

1.1.1.2.4 load_for_inference

lightning.one_step.module.VinLightningModule.load_for_inference(
    checkpoint_path,
    *,
    device='cpu',
    fallback_binner_path=None,
)

Load a checkpoint and return an evaluation-ready scorer module.

1.1.1.2.4.1 Parameters

Name	Type	Description	Default
checkpoint_path	Path | str	Lightning checkpoint containing config and strict vin.* state-dict entries.	required
device	torch.device | str	Destination device for the restored module.	'cpu'
fallback_binner_path	Path | str | None	Binner JSON used only when the checkpoint and restored config do not provide fitted binner state.	None

1.1.1.2.4.2 Returns

Name	Type	Description
	'VinLightningModule'	Evaluation-mode VinLightningModule with ordinal
	'VinLightningModule'	representatives initialized and gradients still owned by the
	'VinLightningModule'	caller’s surrounding context.

1.1.1.2.5 prepare_for_inference

lightning.one_step.module.VinLightningModule.prepare_for_inference(
    fallback_binner_path=None,
)

Initialize mandatory binner-derived state before direct inference.

This is the non-Trainer lifecycle path used by diagnostics and checkpoint loading. It never fits a binner; missing persisted state is an error because ordinal probabilities cannot be mapped to stable RRI representatives otherwise.

1.1.1.2.6 training_step

lightning.one_step.module.VinLightningModule.training_step(batch, batch_idx)

Compute one masked CORAL update for an oracle-labelled training batch.

Returns None when no row has a hard-valid candidate, finite oracle RRI label, and finite K-1 logits; Lightning then skips the update.

1.1.1.2.7 validation_step

lightning.one_step.module.VinLightningModule.validation_step(batch, batch_idx)

Evaluate masked ordinal loss and candidate-table metrics without updates.

1.1.1.2.8 test_step

lightning.one_step.module.VinLightningModule.test_step(batch, batch_idx)

Evaluate the test split using the validation loss/mask contract.

1.1.1.2.9 on_train_epoch_end

lightning.one_step.module.VinLightningModule.on_train_epoch_end()

Emit and reset accumulated training metrics at the epoch boundary.

1.1.1.2.10 on_validation_epoch_end

lightning.one_step.module.VinLightningModule.on_validation_epoch_end()

Emit/reset validation metrics and expected-RRI error statistics.

1.1.1.2.11 on_test_epoch_end

lightning.one_step.module.VinLightningModule.on_test_epoch_end()

Emit and reset accumulated test metrics at the epoch boundary.

1.1.1.2.12 on_after_backward

lightning.one_step.module.VinLightningModule.on_after_backward()

Log configured scorer gradient norms after Lightning backpropagation.

Logging is disabled during sanity validation, evaluation mode, or when GradNormLoggingConfig is disabled. This hook observes gradients only and does not clip or mutate them.

1.1.1.2.13 configure_optimizers

lightning.one_step.module.VinLightningModule.configure_optimizers()

Construct optimizer and optional scheduler from trainable scorer state.

1.1.1.2.13.1 Returns

Name	Type	Description
	dict[str, Any]	Lightning optimizer mapping. One-cycle scheduling resolves its step
	dict[str, Any]	budget from the attached Trainer; plateau scheduling includes its
	dict[str, Any]	configured monitor, interval, and frequency.

1.1.1.2.13.2 Raises

Name	Type	Description
	RuntimeError	If the scorer exposes no trainable parameters.

1.1.1.2.14 summarize_vin

lightning.one_step.module.VinLightningModule.summarize_vin(
    batch,
    *,
    include_torchsummary=True,
    torchsummary_depth=3,
)

Summarize VIN inputs/outputs for a single oracle-labeled batch.

1.1.1.2.14.1 Parameters

Name	Type	Description	Default
batch	VinOracleBatch	Oracle-labeled VIN batch from VinDataModule.	required
include_torchsummary	bool	Whether to append torchsummary module summaries.	True
torchsummary_depth	int	Max depth for torchsummary module traversal.	3

1.1.1.2.14.2 Returns

Name	Type	Description
	str	Multiline string with VIN summary information.

1.1.1.2.15 plot_vin_encodings_batch

lightning.one_step.module.VinLightningModule.plot_vin_encodings_batch(
    batch,
    *,
    out_dir,
    lmax,
    sh_normalization,
    radius_freqs,
    file_stem_prefix,
)

Generate VIN encoding plots for a single oracle-labeled batch.

1.1.1.2.15.1 Parameters

Name	Type	Description	Default
batch	VinOracleBatch	Oracle-labeled VIN batch from VinDataModule.	required
out_dir	Path	Output directory for plots.	required
lmax	int	Max SH degree for visualization.	required
sh_normalization	str	Spherical harmonics normalization mode.	required
radius_freqs	list[float]	Fourier frequencies for radius plot.	required
file_stem_prefix	str	Filename prefix for the plots.	required

1.1.1.2.15.2 Returns

Name	Type	Description
	dict[str, Path]	Mapping of plot names to saved paths.

Additional attribute identifiers: candidate_scorer

1.1.2 VinLightningModuleConfig

lightning.one_step.module.VinLightningModuleConfig()

Configure one-step ordinal candidate scoring and optimization.

The config composes the scorer, optimizer, optional scheduler, fitted RRI binner, loss variants, and logging policy. It does not configure dataset construction or rollout Q_H targets; those remain owned by aria_nbv.lightning.VinDataModuleConfig and aria_nbv.rollouts, respectively.

1.1.2.1 Attributes

Name	Description
target_type	Return the VinLightningModule factory target.
vin	Candidate scorer configuration.
optimizer	AdamW config instantiated from scorer parameters during trainer setup.
lr_scheduler	Optional scheduler created after Lightning knows the optimizer and step budget.
num_classes	Ordinal class count K; scorer logits have final width K-1.
coral_bias_init	CORAL threshold-bias initialization, optionally derived from binner priors.
coral_loss_variant	Per-valid-row ordinal loss applied to Tensor["M K-1", float32] logits.
coral_balance_source	Source of K-1 threshold priors for balanced BCE or focal loss.
coral_balance_eps	Dimensionless clamp keeping threshold priors strictly inside (0, 1).
coral_focal_gamma	Non-negative focusing exponent for the CORAL focal variant.
coral_focal_alpha	Optional positive-class focal weight in [0, 1]; None uses priors.
binner_fit_snippets	Oracle-labelled snippets used to fit ordinal thresholds; None uses the source extent.
binner_max_attempts	Maximum skipped/invalid oracle batches tolerated during binner fitting.
save_binner	Persist fitted ordinal thresholds and class representatives with the run.
binner_path	Fitted binner JSON used for class labels and expected-RRI representatives.
aux_regression_loss	Optional regression loss from expected ordinal RRI to finite oracle RRI.
aux_regression_weight	Initial dimensionless multiplier for the auxiliary expected-RRI loss.
aux_regression_weight_gamma	Per-interval multiplicative decay in (0, 1] for auxiliary weight.
aux_regression_weight_min	Non-negative floor for the decayed auxiliary-loss multiplier.
aux_regression_weight_interval	Lightning counter used as the auxiliary-weight decay exponent.
log_interval_steps	Optional training-step cadence for buffered rank and ordinal diagnostics.
log_spearman	Enable Spearman rank-correlation metrics.
grad_norms	Post-backward gradient-norm targets and logging policy.
coverage_weight_mode	Unitless coverage proxy used to weight valid candidate rows during training.
coverage_weight_floor	Minimum per-row coverage weight in [0, 1] before schedule blending.
coverage_weight_power	Non-negative exponent applied to coverage fractions in [0, 1].
coverage_weight_strength_start	Initial blend from uniform weights (0) to coverage weights (1).
coverage_weight_strength_end	Final uniform-to-coverage blend after the configured anneal window.
coverage_weight_schedule	Interpolation schedule between starting and ending coverage strength.
coverage_weight_interval	Lightning counter used to advance coverage-strength annealing.
coverage_weight_anneal_epochs	Positive epoch horizon for annealing; None keeps the start strength.
coverage_weight_anneal_steps	Positive global-step horizon for annealing; None keeps the start strength.
coverage_weight_apply_aux	Apply the same valid-row coverage weights to expected-RRI regression.