VAEChannelEstimator#

Note

Estimator constructor defaults mirror the settings used during benchmarking.
generate_training_data() defaults mirror the settings used to generate training data.
train_model() defaults mirror the settings used to train the model.

class sionna.bench.mimo_ofdm_chest.papers.baur2022.VAEChannelEstimator(resource_grid: sionna.phy.ofdm.resource_grid.ResourceGrid, model_states: collections.abc.Sequence[collections.abc.Mapping[str, torch.Tensor]], num_rx_ant: int, model_noise_variances: torch.Tensor | None = None, *, latent_dim: int = 16, err_var_scale: float = 0.541, covariance_floor: float = 1e-05, covariance_mode: str = 'full', interpolation_type: str = 'lin_time_avg', interpolator: sionna.phy.ofdm.channel_estimation.BaseChannelInterpolator | None = None, precision: Literal['single', 'double'] | None = None, device: str | None = None)[source]#

Bases: sionna.phy.ofdm.channel_estimation.BasePilotChannelEstimator, sionna.phy.ofdm.channel_estimation._PilotGroupMixin

Channel estimator that leverages a variational autoencoder (VAE) for denoising.

This implementation follows the method proposed in:

Michael Baur, Benedikt Fesl, Michael Koller, and Wolfgang Utschick. Variational autoencoder leveraged mmse channel estimation. In Asilomar Conf. Signals, Syst., Comput., 527–532. October 2022.

TL;DR

This estimator is based on a VAE that learns a conditionally Gaussian channel model. The VAE operates on the DFT-transformed receive-array channel vectors, separately for each pilot-bearing resource element. Its decoder returns a transformed-domain conditional mean and covariance matrix. The estimator then applies a conditional LMMSE update to obtain the channel estimate at that resource element.

Supported System Geometry

The estimator processes every active pilot resource element independently, while jointly processing the corresponding receive-antenna vector. It therefore supports different resource-grid sizes, pilot layouts, UE counts, OFDM-symbol counts, and positive receive-antenna counts, provided that the following hold:

The resource grid contains at least one active nonzero pilot.
The prepared VAE model state matches num_rx_ant, the latent dimension, covariance parameterization, and exact pilot geometry and values.

Changing the receive-antenna count or pilot configuration requires a separately prepared compatible model state.

Receive arrays that are not ULAs are supported by flattening their antenna dimensions and applying the receive DFT over the resulting num_rx_ant axis.

Deviations from the Paper

The paper assumes a uniform linear array (ULA) at the receive antennas, which motivates the use of a DFT over the receive antennas. This implementation supports rectangular receive antenna arrays and uses a 1D DFT over Sionna’s flattened num_rx_ant axis.
The paper assumes a diagonal covariance matrix in the transform that is motivated by the assumption of a receive-side ULA. This implementation supports a full real-valued covariance matrix in the transform domain.
The paper omits certain architectural details, such as convolutional padding and latent dimension, so the faithfulness of the implementation in this regard cannot be guaranteed.
The paper dos not address interpolation. This implementation has to be combined with a dedicated interpolator.
The paper does not specify detector-facing error variances. Due to the conditionally Gaussian channel model, the conditional variances have an analytic expression, which is used in this implementation.

Parameters:

resource_grid (ResourceGrid) – Sionna resource grid defining the pilot pattern.
model_states (Sequence[Mapping[str, torch.Tensor]]) – Prepared VAE state dictionaries. Every state must use the configured VAE architecture and receive-vector length.
num_rx_ant (int) – Number of receive antennas.
model_noise_variances (Optional[torch.Tensor]) – Noise variance associated with each state, shape [num_models]. These values select SNR-specific states. If None, exactly one mixed-SNR state must be supplied.
latent_dim (int) – Latent dimension of the VAEs.
err_var_scale (float) – Multiplicative scale applied to the detector-facing posterior marginal error variances.
covariance_floor (float) – Positive floor added to decoder covariance output.
covariance_mode (CovarianceMode) – Decoder covariance structure. "diagonal" uses the DFT-domain diagonal covariance from the paper. "full" uses a full real-valued covariance matrix in the transform domain.
interpolation_type (str) – Sionna interpolation policy. This argument is ignored when interpolator is provided.
interpolator (Optional[BaseChannelInterpolator]) – Custom Sionna channel interpolator. If None, the interpolator selected by interpolation_type is used.
precision (Optional[Precision]) – Precision used for Sionna channel estimation. If None, precision is used.
device (Optional[str]) – Device for Sionna channel estimation. If None, device is used.

Methods

estimate_at_pilot_locations(y_pilots: torch.Tensor, no: torch.Tensor) → tuple[torch.Tensor, torch.Tensor][source]#

Estimate denoised channels at Sionna pilot locations.

Parameters:

y_pilots (torch.Tensor) – [batch_size, num_rx, num_rx_ant, num_tx, num_streams, num_pilot_symbols], torch.complex. Observed signals for the pilot-carrying resource elements.
no (torch.Tensor) – [batch_size, num_rx, num_rx_ant] or only the first n>=0 dims, torch.float. Variance of the AWGN.

Outputs:

h_hat – [batch_size, num_rx, num_rx_ant, num_tx, num_streams, num_pilot_symbols], torch.complex. Channel estimates for the pilot-carrying resource elements.
err_var – Same shape as h_hat, torch.float. Channel estimation error variance for the pilot-carrying resource elements.

Neural Network Training#

sionna.bench.mimo_ofdm_chest.papers.baur2022.asset.generate_training_data(resource_grid: sionna.phy.ofdm.resource_grid.ResourceGrid, cfg: sionna.bench.mimo_ofdm_chest.configs.system_config.SystemConfig, *, train_batches: int = 400, train_batch_size: int = 16, validation_batches: int = 16, validation_batch_size: int = 16, training_snr_dbs: list[float] | None = None, seed: int = 42, asset_root: str | pathlib.Path, overwrite: bool = False) → pathlib.Path[source]#

Generate and persist a deterministic VAE-noisy training-data bank.

The asset contains clean propagated channels. The same channel realizations are reused across the requested SNR points, while deterministic SNR-specific observation noise is generated during model training. The completed asset is strictly validated before its manifest path is returned.

Parameters:

resource_grid (sionna.phy.ofdm.resource_grid.ResourceGrid) – Sionna resource grid whose pilot geometry and values define the derived VAE training observations.
cfg (sionna.bench.mimo_ofdm_chest.configs.system_config.SystemConfig) – System and channel configuration used to generate the channel realizations.
train_batches (int) – Number of independent training-channel batches.
train_batch_size (int) – Number of complete system-channel realizations in each training batch. The total number of training channels is train_batches * train_batch_size.
validation_batches (int) – Number of independent validation-channel batches.
validation_batch_size (int) – Number of complete system-channel realizations in each validation batch. The total number of validation channels is validation_batches * validation_batch_size.
training_snr_dbs (list[float] | None) – SNR points for which the data bank is prepared. The clean channel realizations are shared across these SNR points. If None, use the SNR grid defined by cfg.
seed (int) – Base random seed for deterministic channel generation. Training and validation use separate deterministic random streams.
asset_root (str | pathlib.Path) – Writable external root beneath which the generated payloads and manifest are stored.
overwrite (bool) – If True, replace an existing asset with the same identity. If False, raise an error when that asset already exists.

Returns:

Path to the strictly validated manifest.json of the generated training-data asset.

sionna.bench.mimo_ofdm_chest.papers.baur2022.asset.train_model(resource_grid: sionna.phy.ofdm.resource_grid.ResourceGrid, cfg: sionna.bench.mimo_ofdm_chest.configs.system_config.SystemConfig, *, method: str = 'noisy', latent_dim: int = 16, train_batches: int = 400, train_batch_size: int = 16, validation_batches: int = 16, validation_batch_size: int | None = 16, epochs: int = 200, optimizer_batch_size: int = 512, learning_rate: float = 0.00023, free_bits: float = 0.875, min_train_seconds: float = 0.0, patience_epochs: int = 20, relative_min_delta: float = 0.0001, snr_specific_weights: bool = True, training_snr_dbs: list[float] | None = None, covariance_floor: float = 1e-05, covariance_mode: str = 'full', seed: int = 42, training_data_manifest: str | pathlib.Path | None = None, asset_root: str | pathlib.Path, overwrite: bool = False) → tuple[list[dict[str, torch.Tensor]], torch.Tensor | None, list[pathlib.Path]][source]#

Train and persist an ordered bank of VAE channel-estimation models.

For SNR-specific training, one independent model is trained for every requested SNR. The supplied training-data asset provides fixed clean channel realizations; fresh deterministic observation noise is generated for every training epoch. Each completed model asset is persisted atomically and strictly reloaded before being returned.

Deviations from the Paper

The paper trains a separate VAE at each SNR until the ELBO saturates and retains the model with the highest ELBO, but does not specify a dataset size, latent dimension, free-bits threshold, optimizer batch size, stopping patience, or minimum improvement. This implementation uses finite deterministic training and validation channel banks, a maximum epoch count, and validation-loss early stopping with restoration of the best validation state.
For VAE-noisy, the clean channel bank is reused across epochs while fresh deterministic observation noise is generated every epoch. This retains the paper’s observation-noise resampling while making the channel sample and training procedure reproducible. Shared mixed-SNR training is an additional configurable mode; per-SNR models follow the paper.
The paper uses a diagonal complex Gaussian decoder covariance. The "full" mode is an adaptation that learns a full covariance over the real-stacked transform-domain channel vector and uses the corresponding multivariate Gaussian likelihood.

Parameters:

resource_grid (sionna.phy.ofdm.resource_grid.ResourceGrid) – Resource grid defining the pilot geometry and pilot values used to construct the VAE observations.
cfg (sionna.bench.mimo_ofdm_chest.configs.system_config.SystemConfig) – System and channel configuration defining the channel distribution, SNR grid, numerical precision, and device.
method (str) – VAE training method. "noisy" uses noisy encoder inputs and clean channel targets. "real" uses noisy inputs and noisy targets.
latent_dim (int) – Dimension of the Gaussian latent representation.
train_batches (int) – Number of independent training-channel batches.
train_batch_size (int) – Number of complete channel realizations per training batch. Their product is the number of training channels.
validation_batches (int) – Number of independent validation-channel batches.
validation_batch_size (int | None) – Number of channel realizations per validation batch. If None, use train_batch_size.
epochs (int) – Maximum number of complete training epochs.
optimizer_batch_size (int) – Number of derived VAE examples in each optimizer update.
learning_rate (float) – Adam learning rate.
free_bits (float) – Minimum KL contribution assigned to each latent dimension.
min_train_seconds (float) – Minimum wall-clock duration before early stopping is permitted.
patience_epochs (int) – Stop after this many epochs without the required validation-loss improvement. Zero disables early stopping.
relative_min_delta (float) – Minimum relative validation-loss reduction counted as an improvement.
snr_specific_weights (bool) – If True, train one model per requested SNR. Otherwise, train one model shared across all requested SNRs.
training_snr_dbs (list[float] | None) – SNR values used for training. If None, use the SNR grid defined by cfg.
covariance_floor (float) – Positive floor applied to decoder covariance outputs.
covariance_mode (str) – Decoder covariance parameterization, either "diagonal" or "full".
seed (int) – Base seed controlling model initialization, data ordering, and deterministic observation-noise generation.
training_data_manifest (str | pathlib.Path | None) – Manifest of the fixed training-data asset. This is required for method="noisy".
asset_root (str | pathlib.Path) – Writable external root beneath which model assets are stored.
overwrite (bool) – If True, replace assets with the same identity.

Returns:

A tuple (model_states, model_noise_variances, manifest_paths). model_states contains the strictly reloaded state dictionaries. model_noise_variances contains the ordered inference noise variances for an SNR-specific model bank, or None for a shared model. manifest_paths contains the corresponding validated manifest paths.

Hyperparameter Tuning#

Training Hyperparameter Tuning

The "training" stage minimizes weighted NMSE at -10, -5, 0, 5, and 10 dB. It runs 10 seeded TPE trials. Every trial trains a separate VAE model for each of these five SNRs and evaluates three batches of four channels at each SNR.

All trials reuse the same training-data bank. Every SNR-specific model uses 6,400 training channels and 256 validation channels, at most 200 epochs, patience of 20 epochs, and a relative minimum validation improvement of 1e-4. The best checkpoint is restored.

Parameter	Search space
`latent_dim`	categorical set {4, 8, 16}
`learning_rate`	logarithmic from 1e-5 through 1e-3
`free_bits`	uniform from 0 through 1
`optimizer_batch_size`	categorical set {256, 512, 1024, 2048}

The "training-validation" stage retrains and compares the rounded winner with the existing training defaults using 63 batches of four channels at each SNR.

Error-Variance Calibration

The "variance" stage tunes err_var_scale logarithmically from 1e-3 through 1 using 10 TPE trials. Codeword BLER is evaluated at 10 dB over 63 batches of four channels using the frozen canonical model bank.

The "variance-validation" stage compares the rounded winner with the existing scale.

Parameters:

study (str) – Stage to execute: "training", "training-validation", "variance", or "variance-validation".
n_trials (int | None) – Trial-count override. If None, use 10 for both training tuning and variance calibration. Validation stages do not run Optuna trials.
num_mc_batches (int | None) – Monte Carlo budget override. If None, use three batches per SNR for training tuning and 63 for validation and variance calibration.
batch_size (int | None) – Channels per evaluation batch. If None, use the canonical configuration value of four.
output_dir (str | pathlib.Path | None) – Output directory override for the selected stage.
algorithm_output_dir (str | pathlib.Path | None) – Directory containing the training study and paired-validation evidence used to select the canonical training recipe.
asset_root (str | pathlib.Path) – External root containing or receiving trained VAE model assets.
training_data_manifest (str | pathlib.Path | None) – Manifest of the fixed training and validation data bank. This is required by the training and training-validation stages.

Returns:

An Optuna study for tuning stages or a dictionary containing paired-validation results for validation stages.