ILQRVAE

Optimization-based iLQR-VAE with posterior-control inference and ELBO training.

Source

  • Registry name: ilqr_vae
  • Model class: ladys.models.ilqr_vae.ILQRVAE
  • Config class: ladys.models.ilqr_vae.ILQRVAEConfig
  • Source file: src/ladys/models/ilqr_vae.py

Method

iLQR-VAE is a latent dynamical model with no amortized recognition network for the posterior mean. For each trial and current generative parameter setting, the model solves an inner optimal-control problem over a sequence of latent inputs u. Those inputs drive a recurrent dynamical system to produce latent states z, and a likelihood readout decodes observations from z.

The LaDyS implementation ports the tutorial Student input prior, Mini-GRU-IO dynamics, Poisson spike likelihood, and iLQR posterior-control solver into PyTorch. The posterior covariance is shared across trials as a Kronecker product of learned time and input-space factors, matching the structure used by the original tutorial code.

Inference-only checkpoint mode

With objective="posterior_control" and initialization="pretrained", the model loads the supplied params_path and uses iLQR only to infer posterior controls for each evaluation trial. It can register checkpoint tensors as buffers (trainable_parameters=false) or as nn.Parameters (trainable_parameters=true) for regression checks; with zero training epochs both paths are numerically identical. The corrected solver includes the terminal observation likelihood and is not an exact reproduction of historical predictions made before that correction.

For NLB-style co-smoothing, the inner solve can be restricted to held-in neurons by setting held_in_neurons, and the returned rates can be sliced to held-out neurons with output_neuron_start and output_neurons. Returned rates are expected spike counts per bin, not Hz, so they can be consumed directly by LaDyS/NLB bits-per-spike metrics.

ELBO training mode

With objective="ilqr_vae_elbo" and a gradient optimizer, training follows the original iLQR-VAE outer objective:

ELBO = H[q(u | o)] + E_q[log p(u) + log p(o | z(u))]
loss = -ELBO / num_observations + regularizer

The inner iLQR solve provides the posterior mean controls. By default the outer backward pass differentiates through the executed solver updates as well as the sampled ELBO, including the Student prior, dynamics, likelihood, and shared posterior covariance. This unrolled finite-iteration derivative differs from the upstream implicit adjoint at an optimum. Adam fallback also preserves its control gradient; a detached-control approximation is available only by explicitly setting differentiate_controls=false. Bounded parameters such as prior scales, degrees of freedom, gains, and covariance diagonals are projected back into their valid domains after optimizer steps.

Current scope

The trainable path is designed for spike-count datasets and currently uses the Poisson/Mini-GRU-IO variant that was translated for MC_Maze. The original repository's standalone Lorenz example uses MGU2 dynamics with a 3D Gaussian observation model; exact parity with that example would require adding that dynamics/likelihood variant as a separate model option. The provided configs/experiment/synthetic/lorenz/ilqr_vae/ilqr_vae_lorenz_100.yaml trains the Poisson/Mini-GRU-IO variant on the LaDyS Lorenz-100 spike population for comparison with LFADS and NDT.

Configuration

Config for the PyTorch iLQR-VAE adapter.

The default ELBO objective trains a randomly initialized model. Loading a fixed tutorial checkpoint is opt-in through initialization="pretrained", an explicit params_path, objective="posterior_control", and optimization.name="inference_only".

latent_dim is the recurrent latent state dimension and input_dim is the dimensionality of the inferred control input. The current trainable LaDyS path uses the translated Student prior, Mini-GRU-IO dynamics, Poisson likelihood, and shared Kronecker posterior covariance. For co-smoothing datasets, held_in_neurons selects the neurons used by the inner posterior solve, while output_neuron_start and output_neurons select the decoded prediction slice returned to the benchmark metrics. build_from_data derives those slices and dt from the dataset; explicit inconsistent values are rejected. Direct build uses unit-width bins if dt is omitted.

Field Type Default
name Literal['ilqr_vae'] 'ilqr_vae'
objective Literal['posterior_control', 'ilqr_vae_elbo'] 'ilqr_vae_elbo'
params_path Optional[str] None
initialization Literal['pretrained', 'random', 'checkpoint_transfer'] 'random'
template_params_path Optional[str] None
random_init_profile Literal['default', 'tutorial_mc_maze'] 'default'
readout_bias_initialization Literal['none', 'empirical_rates'] 'none'
empirical_rate_floor_hz float 0.001
latent_dim int 20
input_dim int 5
init_seed int 0
solver Literal['ilqr', 'lbfgs', 'adam'] 'ilqr'
max_iter int 5
lr Optional[float] None
control_hessian_mode Literal['true', 'fisher', 'clamped'] 'true'
ilqr_failure_fallback Literal['none', 'adam', 'lbfgs'] 'adam'
ilqr_fallback_max_iter int 25
ilqr_fallback_lr Optional[float] None
differentiate_controls bool True
trainable_parameters bool True
n_posterior_samples int 1
include_elbo_constants bool True
dynamics_regularizer float 0.0
held_in_neurons Optional[int] None
output_neuron_start Optional[int] None
output_neurons Optional[int] None
rate_mode Literal['likelihood', 'pre_sample'] 'likelihood'
dt Optional[float] None
optimization OptimizationConfig OptimizationConfig(name='gradient', lr=0.001)

Contracts