1 lightning

lightning

Lightning orchestration for VIN candidate-scorer training.

This package provides the config-as-factory experiment, data-module, training module, trainer, callback, and optimizer surfaces used by the runnable one-step CORAL scorer. Its finite-horizon leaf modules additionally own QH stage admission, masked Double-Q optimization, exact-Q_2 certification, and immutable experiment/bundle publication. Actor-visible feature construction belongs to aria_nbv.vin; oracle labels belong to aria_nbv.rri_metrics and aria_nbv.oracle.

The package root intentionally retains the established eight-name one-step interface. Import task-specific leaves from .one_step, .finite_horizon, or .training.

1.1 Package guide

aria_nbv.lightning owns training-stage admission, optimization, metric lifecycle, checkpoint selection and experiment publication. Scorer architecture belongs to aria_nbv.vin; actor/supervision DTOs belong to aria_nbv.data_handling; factual replay and stored masks belong to aria_nbv.rollouts. Lightning composes those interfaces and preserves their scientific contracts.

1.2 Responsibility map

Package Conceptual owner and caller interface
one_step One-step CORAL training and checkpoint recovery through AriaNBVExperimentConfig
finite_horizon QH admission, fitting, evaluation and immutable publication through QhExperiment
finite_horizon/studies Frozen study protocols, schema-specific receipt admission and scientific reductions
training Shared trainer, callback, optimizer and gradient-observation machinery

Task packages expose composition interfaces and retain their data/objective semantics. Shared training machinery handles framework construction and observation. Study readers preserve historical scientific identities; reporting presents admitted results. Configuration, receipts, tests and logs provide cross-cutting evidence of those owners’ behavior.

Code 
flowchart LR
  D["data_handling DTOs"] --> O["one_step lifecycle"]
  D --> Q["finite_horizon lifecycle"]
  F["rollouts replay contracts"] --> Q
  V["vin scorer interfaces"] --> O
  V --> Q
  T["training machinery"] --> O
  T --> Q
  Q --> E["Receipts and immutable artifacts"]
  E --> S["studies admission and reductions"]
  S --> R["reporting snapshots and figures"]

flowchart LR
  D["data_handling DTOs"] --> O["one_step lifecycle"]
  D --> Q["finite_horizon lifecycle"]
  F["rollouts replay contracts"] --> Q
  V["vin scorer interfaces"] --> O
  V --> Q
  T["training machinery"] --> O
  T --> Q
  Q --> E["Receipts and immutable artifacts"]
  E --> S["studies admission and reductions"]
  S --> R["reporting snapshots and figures"]

1.3 Entry points and retained contracts

The package root retains the established eight one-step names. QH callers use aria_nbv.lightning.finite_horizon.*; study callers import their owning module under finite_horizon.studies.*. The package-owned CLI adapters dispatch through their corresponding experiment lifecycle:

nbv-train --help
nbv-summary --help
nbv-qh --help

One-step supports online/offline VIN data, binner fitting, Lightning resume and Optuna. QH distinguishes diagnostic fitting from deployment-bundle publication; verified inference weights are different from resumed training state. The task READMEs describe admission, support, ordering, failure and verification contracts.

The legacy qh_datamodule, qh_experiment and qh_module import paths are retired; qh_q2_certification remains available. Public tensor/callable contracts are documented with their source and tests.

1.4 Classes

Name Description
AdamWConfig Configure AdamW for a Lightning module’s trainable parameters.
AriaNBVExperimentConfig Own one reproducible VIN Lightning run and its nested factories.
TrainerCallbacksConfig Configure one mutually compatible Lightning callback collection.
TrainerFactoryConfig Configure one external Lightning trainer and its attached services.
VinDataModule Own stage datasets and loaders for oracle-labelled VIN batches.
VinDataModuleConfig Configure stage datasets and candidate-table collation.
VinLightningModule Own one-step CORAL training state for a VIN candidate scorer.
VinLightningModuleConfig Configure one-step ordinal candidate scoring and optimization.

1.4.1 AdamWConfig

lightning.AdamWConfig()

Configure AdamW for a Lightning module’s trainable parameters.

1.4.1.1 Attributes

Name Description
target_type Factory target for aria_nbv.configs.base.BaseConfig.setup_target.
learning_rate Learning rate for AdamW.
weight_decay Weight decay for AdamW.
foreach Explicit foreach policy; None preserves PyTorch’s historical backend selection.
fused Explicit fused policy; None preserves PyTorch’s historical backend selection.

1.4.1.2 Methods

Name Description
setup_target Create AdamW over the exact trainable parameter list.
1.4.1.2.1 setup_target
lightning.AdamWConfig.setup_target(params)

Create AdamW over the exact trainable parameter list.

1.4.1.2.1.1 Parameters
Name Type Description Default
params list[Tensor] Materialized parameters selected by the consuming aria_nbv.lightning.VinLightningModule.configure_optimizers or aria_nbv.lightning.finite_horizon.module.QhLightningModule.configure_optimizers. required
1.4.1.2.1.2 Returns
Name Type Description
Optimizer Optimizer owned and stepped by Lightning.

1.4.2 AriaNBVExperimentConfig

lightning.AriaNBVExperimentConfig()

Own one reproducible VIN Lightning run and its nested factories.

Calling :meth:setup_target constructs but does not execute the three runtime owners. :meth:setup_target_and_run then delegates loop ownership to Lightning. Run artifacts, checkpoint paths, and fitted binner state are resolved through :class:PathConfig before construction.

1.4.2.1 Attributes

Name Description
run_mode Requested experiment action, normalized from kebab or snake case.
run_name Human-readable identifier used by loggers, checkpoints, and snapshots.
stage Lightning stage constructed and executed by :meth:setup_target_and_run.
summary_stage Dataset stage sampled by the non-training VIN summary action.
summary_num_batches Positive number of real oracle-labelled batches included in a summary.
summary_include_torchsummary Append external torchsummary module traversal to scorer diagnostics.
summary_torchsummary_depth Positive maximum module-tree depth for optional torchsummary output.
plot_out_dir Path for non-training VIN encoding figures, resolved by the plot action.
ckpt_path Optional checkpoint for strict validation/test loading or training resume.
seed Torch CPU/CUDA seed set before dataset, scorer, and trainer construction.
out_dir Run-artifact root for config snapshots, logs, checkpoints, and binner state.
fit_binner_only Compatibility flag that rewrites run_mode to "fit_binner".
fit_binner_overwrite Overwrite module_config.binner_path when run_mode="fit_binner".
inspect_config Print rich datamodule and LightningModule config summaries before running.
paths Repository-aware resolver for configs, run roots, checkpoints, and artifacts.
trainer_config Trainer factory including callback ownership, logger, devices, and loop limits.
datamodule_config VIN source, stage-reuse, worker, and candidate-collation configuration.
module_config One-step scorer, CORAL loss, binner, optimizer, scheduler, and metric config.
optuna_config Optional study and search-space config used only by the Optuna action.
verbosity Experiment-orchestration logging verbosity.
verbose Legacy Optuna-facing verbosity value retained for config compatibility.
resolved_out_dir Return out_dir resolved under the configured repository run root.
default_config_path Default path for saving the experiment TOML.

1.4.2.2 Methods

Name Description
save_config Save this config (and nested configs) as TOML.
setup_target Construct the trainer, scorer module, and stage-aware datamodule.
setup_target_and_run Construct runtime owners and delegate one stage loop to Lightning.
run_optuna_study Run an Optuna hyperparameter sweep for this experiment config.
summarize_vin Run VIN summary on a real oracle batch from the datamodule.
fit_binner_and_save Fit RriOrdinalBinner from oracle batches and persist it.
run Execute the configured action.
1.4.2.2.1 save_config
lightning.AriaNBVExperimentConfig.save_config(
    path=None,
    *,
    include_comments=True,
    include_type_hints=True,
)

Save this config (and nested configs) as TOML.

1.4.2.2.1.1 Parameters
Name Type Description Default
path Path | str | None Destination TOML path. If None, uses default_config_path. None
include_comments bool Include docstring comments in the TOML. True
include_type_hints bool Include type hints in the TOML comments. True
1.4.2.2.1.2 Returns
Name Type Description
Path Resolved path that was written.
1.4.2.2.2 setup_target
lightning.AriaNBVExperimentConfig.setup_target(setup_stage=None, *, trial=None)

Construct the trainer, scorer module, and stage-aware datamodule.

1.4.2.2.2.1 Parameters
Name Type Description Default
setup_stage Stage | str | None Dataset stage to initialize; defaults to stage. None
trial 'optuna.Trial | None' Optional Optuna trial forwarded to trainer callback setup. None
1.4.2.2.2.2 Returns
Name Type Description
pl.Trainer Tuple of external Lightning trainer, VinLightningModule,
and VinLightningModule class:VinDataModule. No fit/validation/test loop has run.
1.4.2.2.2.3 Notes

Torch seeds and artifact directories are established before any runtime owner is constructed. Checkpoint weights remain Lightning’s responsibility when the returned trainer starts its loop.

1.4.2.2.3 setup_target_and_run
lightning.AriaNBVExperimentConfig.setup_target_and_run(stage=None)

Construct runtime owners and delegate one stage loop to Lightning.

Training verifies the fitted binner class count before fit; validation and test pass the resolved checkpoint directly to Lightning. A globally ranked interrupted trainer writes one explicit recovery checkpoint.

1.4.2.2.3.1 Returns
Name Type Description
pl.Trainer Trainer after the requested loop finishes.
1.4.2.2.4 run_optuna_study
lightning.AriaNBVExperimentConfig.run_optuna_study()

Run an Optuna hyperparameter sweep for this experiment config.

1.4.2.2.5 summarize_vin
lightning.AriaNBVExperimentConfig.summarize_vin()

Run VIN summary on a real oracle batch from the datamodule.

1.4.2.2.6 fit_binner_and_save
lightning.AriaNBVExperimentConfig.fit_binner_and_save(
    datamodule=None,
    *,
    overwrite=False,
)

Fit RriOrdinalBinner from oracle batches and persist it.

1.4.2.2.6.1 Parameters
Name Type Description Default
datamodule VinDataModule | None Optional pre-built datamodule (avoids re-instantiation). None
overwrite bool Overwrite existing binner JSON when True. False
1.4.2.2.6.2 Returns
Name Type Description
Path Path to the saved rri_binner.json.
1.4.2.2.7 run
lightning.AriaNBVExperimentConfig.run()

Execute the configured action.

This method is intended to be called from CLI entry points.

1.4.3 TrainerCallbacksConfig

lightning.TrainerCallbacksConfig()

Configure one mutually compatible Lightning callback collection.

:meth:setup_target may adapt checkpoint, early-stopping, and pruning flags for an Optuna trial. The returned callback objects are then owned and invoked by the external Lightning trainer.

1.4.3.1 Attributes

Name Description
target_type Factory target for aria_nbv.configs.base.BaseConfig.setup_target.
use_model_checkpoint Attach a model-checkpoint callback using the fields below.
checkpoint_monitor Logged metric key used to rank checkpoints.
checkpoint_mode Checkpoint ranking direction, conventionally "min" or "max".
checkpoint_dir Checkpoint directory; None resolves through :class:PathConfig.
checkpoint_filename Lightning format template evaluated from epoch, step, and logged metrics.
checkpoint_save_top_k Number of monitor-ranked checkpoints retained by Lightning.
checkpoint_auto_insert_metric_name Let Lightning inject monitor names into the checkpoint filename.
checkpoint_save_last Lightning tri-state policy for maintaining a last.ckpt alias/copy.
checkpoint_every_n_train_steps Optional positive optimizer-step checkpoint cadence.
checkpoint_train_time_interval Optional positive wall-clock cadence accepted by datetime.timedelta.
checkpoint_every_n_epochs Optional positive epoch checkpoint cadence.
checkpoint_save_on_train_epoch_end Optional override for Lightning’s train-epoch-end checkpoint timing.
use_early_stopping Attach an early-stopping callback for the configured monitor.
early_stopping_monitor Logged metric key observed by early stopping.
early_stopping_mode Early-stopping improvement direction, "min" or "max".
early_stopping_patience Validation checks without improvement allowed before stopping.
optimizer_diagnostics Opt-in bounded gradient and optimizer scalars at the unscaled pre-step hook.
use_lr_monitor Attach learning-rate logging when the trainer has a logger.
lr_logging_interval Learning-rate logging cadence passed to Lightning.
use_optuna_pruning Enable Optuna pruning callback for hyperparameter optimisation runs.
use_rich_progress_bar Enable the Rich progress bar; mutually exclusive with TQDM.
use_tqdm_progress_bar Enable the TQDM progress bar; mutually exclusive with Rich.
tqdm_refresh_rate TQDM refresh cadence in training batches.
use_rich_model_summary Attach a Rich model-tree summary callback.
rich_summary_max_depth Maximum module-tree depth rendered by the Rich summary.
use_backbone_finetuning Attach Lightning’s staged backbone-unfreezing callback.
backbone_unfreeze_at_epoch Zero-based epoch at which the callback unfreezes the backbone.
backbone_lambda_func Optional serialized learning-rate lambda consumed by backbone finetuning.
backbone_train_bn Keep backbone batch-normalization layers trainable after unfreezing.
use_timer Attach a Lightning wall-clock timer callback.
timer_duration Optional duration mapping accepted by datetime.timedelta.
timer_interval Timer check cadence passed to Lightning.

1.4.3.2 Methods

Name Description
setup_target Build callbacks and adapt trial-specific ownership flags.
1.4.3.2.1 setup_target
lightning.TrainerCallbacksConfig.setup_target(
    model_name=None,
    *,
    has_logger=True,
    trial=None,
    optuna_config=None,
)

Build callbacks and adapt trial-specific ownership flags.

1.4.3.2.1.1 Parameters
Name Type Description Default
model_name str | None Optional prefix added to checkpoint filenames. None
has_logger bool Whether a trainer logger exists; LR monitoring is omitted when false. True
trial 'Trial | None' Optional Optuna trial that owns pruning decisions. None
optuna_config 'OptunaConfig | None' Required pruning config when trial enables the pruning callback. None
1.4.3.2.1.2 Returns
Name Type Description
list[Callback] Fresh callback list owned by one external Lightning trainer.
1.4.3.2.1.3 Notes

Trial runs disable checkpointing and early stopping to avoid competing persistence/termination owners. Outside a trial, stale pruning flags are disabled explicitly.

1.4.4 TrainerFactoryConfig

lightning.TrainerFactoryConfig()

Configure one external Lightning trainer and its attached services.

1.4.4.1 Attributes

Name Description
target_type Return the external Lightning trainer factory target.
is_debug Force CPU fast-dev execution, anomaly detection, and no checkpoint callback.
fast_dev_run Ask Lightning to run one batch per loop as an integration smoke test.
accelerator Lightning accelerator selector such as "auto", "cpu", or "gpu".
devices Device count, selector, or explicit device-index sequence passed to Lightning.
strategy Optional distributed-execution strategy; "auto" delegates selection.
max_epochs Maximum training epochs; None delegates the unbounded policy to Lightning.
max_steps Exact optimizer-step ceiling; -1 leaves the existing epoch-only policy.
precision Lightning numeric-precision policy, including mixed-precision strings.
default_root_dir Root directory used by Lightning for logger-free run artifacts.
tf32_matmul_precision Optional process-wide float32 matmul precision set before trainer construction.
gradient_clip_val Optional gradient clipping threshold applied by Lightning.
accumulate_grad_batches Positive number of batches accumulated per optimizer step.
log_every_n_steps Trainer logging cadence in optimizer/training steps.
deterministic Lightning deterministic-algorithm policy; None preserves framework defaults.
limit_train_batches Optional absolute count or fraction limiting each training epoch.
limit_val_batches Optional absolute count or fraction limiting each validation epoch.
check_val_every_n_epoch Epoch cadence for validation when validation is enabled.
num_sanity_val_steps Validation batches run before training; -1 checks the entire loader.
use_distributed_sampler Allow Lightning to replace data-loader samplers during distributed execution.
enable_validation Enable validation; false zeroes validation limits/cadence during validation.
enable_model_summary Enable Lightning’s default model summary callback.
callbacks Callback collection factory attached to the constructed trainer.
use_wandb Use the configured W&B logger outside debug mode.
wandb_config W&B logger config instantiated only when use_wandb is true.

1.4.4.2 Methods

Name Description
setup_target Instantiate a Trainer after logger and callback ownership is resolved.
1.4.4.2.1 setup_target
lightning.TrainerFactoryConfig.setup_target(
    experiment=None,
    *,
    trial=None,
    optuna_config=None,
)

Instantiate a Trainer after logger and callback ownership is resolved.

1.4.4.2.1.1 Parameters
Name Type Description Default
experiment Any | None Optional experiment providing an optuna_config fallback. None
trial 'Trial | None' Optional Optuna trial forwarded to callback construction. None
optuna_config 'OptunaConfig | None' Explicit pruning configuration, preferred over the experiment fallback. None
1.4.4.2.1.2 Returns
Name Type Description
pl.Trainer External Lightning trainer owning the fresh logger and callbacks.
1.4.4.2.1.3 Notes

TF32 precision is process-wide. Debug mode uses Lightning’s default logger; normal mode uses W&B only when configured. Checkpointing is enabled precisely when callback construction produced a checkpoint callback.

1.4.5 VinDataModule

lightning.VinDataModule(config)

Own stage datasets and loaders for oracle-labelled VIN batches.

Lightning may call :meth:setup more than once. Dataset instances are therefore initialized lazily and retained by stage. Scorer-device transfer remains the module’s responsibility; this class yields CPU/source-device VinOracleBatch objects whose candidate width is compact for a single sample or right-padded after offline collation.

1.4.5.1 Attributes

Name Type Description
config Immutable construction policy for sources, batching, and worker lifecycle.

1.4.5.2 Methods

Name Description
setup Initialize and retain the datasets required by one Lightning stage.
train_dataloader Build the training loader with optional row and candidate shuffling.
val_dataloader Build the deterministic validation loader for the retained source.
test_dataloader Reuse the validation-loader contract for Lightning test evaluation.
iter_oracle_batches Return a direct iterator over one stage’s oracle-labelled samples.
1.4.5.2.1 setup
lightning.VinDataModule.setup(stage=None)

Initialize and retain the datasets required by one Lightning stage.

1.4.5.2.1.1 Parameters
Name Type Description Default
stage Stage | str | None train, val, test, or None. None prepares both training and validation ownership paths. None
1.4.5.2.1.2 Notes

Repeated calls are idempotent for already-created sources. Validation and test share _val_source; when use_train_as_val is enabled they deliberately share _train_source as well.

1.4.5.2.2 train_dataloader
lightning.VinDataModule.train_dataloader()

Build the training loader with optional row and candidate shuffling.

1.4.5.2.2.1 Returns
Name Type Description
DataLoader DataLoader yielding VinOracleBatch. Map-style batching
DataLoader produces poses PoseTW["B N_q 12"], RRI labels
DataLoader Tensor["B N_q", float32], and a hard prefix mask derived from
DataLoader candidate_count; iterable/online sources remain unbatched.
1.4.5.2.3 val_dataloader
lightning.VinDataModule.val_dataloader()

Build the deterministic validation loader for the retained source.

1.4.5.2.3.1 Returns
Name Type Description
DataLoader DataLoader yielding unshuffled VinOracleBatch objects with
DataLoader the same compact/right-padded shape contract as the training loader.
1.4.5.2.4 test_dataloader
lightning.VinDataModule.test_dataloader()

Reuse the validation-loader contract for Lightning test evaluation.

1.4.5.2.5 iter_oracle_batches
lightning.VinDataModule.iter_oracle_batches(stage)

Return a direct iterator over one stage’s oracle-labelled samples.

This bypasses worker processes and VinOracleBatch.collate, so yielded candidate payloads retain their per-sample N_q width.

1.4.6 VinDataModuleConfig

lightning.VinDataModuleConfig()

Configure stage datasets and candidate-table collation.

Online sources yield one oracle-labelled snippet at a time and therefore require num_workers=0. Map-style offline sources may collate a batch of size B; candidate rows are right-padded to N_q and accompanied by per-sample candidate_count values used to construct the hard prefix mask.

1.4.6.1 Attributes

Name Description
target_type Return the VinDataModule factory target.
paths Project path resolver used by nested dataset/store configurations.
source Config-as-factory source for online labels or an immutable offline store.
shuffle Shuffle map-style training samples each epoch; ignored for iterable sources.
shuffle_candidates Permute compact valid candidate rows and aligned labels per training sample.
num_workers DataLoader worker count; must be zero for the non-multiprocess-safe online labeler.
batch_size Optional offline batch size B using VinOracleBatch.collate.
persistent_workers Keep worker processes between epochs; forced off when num_workers is zero.
verbosity Dataset and oracle-labeler diagnostic verbosity.
is_debug Serialized compatibility flag; worker and trainer debug policy is configured elsewhere.
use_train_as_val Reuse the exact train dataset instance for validation and testing.

1.4.7 VinLightningModule

lightning.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.4.7.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.4.7.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.4.7.2.1 setup
lightning.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.4.7.2.2 on_save_checkpoint
lightning.VinLightningModule.on_save_checkpoint(checkpoint)

Persist fitted ordinal-binner state inside a Lightning checkpoint.

1.4.7.2.3 on_load_checkpoint
lightning.VinLightningModule.on_load_checkpoint(checkpoint)

Restore embedded ordinal-binner state before stage setup runs.

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

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

1.4.7.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.4.7.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.4.7.2.5 prepare_for_inference
lightning.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.4.7.2.6 training_step
lightning.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.4.7.2.7 validation_step
lightning.VinLightningModule.validation_step(batch, batch_idx)

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

1.4.7.2.8 test_step
lightning.VinLightningModule.test_step(batch, batch_idx)

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

1.4.7.2.9 on_train_epoch_end
lightning.VinLightningModule.on_train_epoch_end()

Emit and reset accumulated training metrics at the epoch boundary.

1.4.7.2.10 on_validation_epoch_end
lightning.VinLightningModule.on_validation_epoch_end()

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

1.4.7.2.11 on_test_epoch_end
lightning.VinLightningModule.on_test_epoch_end()

Emit and reset accumulated test metrics at the epoch boundary.

1.4.7.2.12 on_after_backward
lightning.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.4.7.2.13 configure_optimizers
lightning.VinLightningModule.configure_optimizers()

Construct optimizer and optional scheduler from trainable scorer state.

1.4.7.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.4.7.2.13.2 Raises
Name Type Description
RuntimeError If the scorer exposes no trainable parameters.
1.4.7.2.14 summarize_vin
lightning.VinLightningModule.summarize_vin(
    batch,
    *,
    include_torchsummary=True,
    torchsummary_depth=3,
)

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

1.4.7.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.4.7.2.14.2 Returns
Name Type Description
str Multiline string with VIN summary information.
1.4.7.2.15 plot_vin_encodings_batch
lightning.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.4.7.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.4.7.2.15.2 Returns
Name Type Description
dict[str, Path] Mapping of plot names to saved paths.

Additional attribute identifiers: candidate_scorer

1.4.8 VinLightningModuleConfig

lightning.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.4.8.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.
Back to top