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 final_params.bin and uses iLQR only to infer posterior controls for each evaluation trial. This is the mode used to reproduce the MC_Maze tutorial result. It can register the checkpoint tensors either as buffers (trainable_parameters=false) or as nn.Parameters (trainable_parameters=true) for regression checks; with zero training epochs both paths are numerically identical.

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 and is treated as an implicit inference step; the outer PyTorch backward pass differentiates the sampled ELBO through the Student prior, dynamics, likelihood readout, and shared posterior covariance parameters. 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.

Use objective="posterior_control" with initialization="pretrained" and optimization.name="inference_only" to reproduce a fixed checkpoint, such as the MC_Maze tutorial model. Use objective="ilqr_vae_elbo", initialization="random", trainable_parameters=true, and a gradient optimizer to train a new model from scratch.

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.

Field Type Default
name Literal['ilqr_vae'] 'ilqr_vae'
objective Literal['posterior_control', 'ilqr_vae_elbo'] 'posterior_control'
params_path Optional[str] 'data/real/ilqr_vae/final_params.bin'
initialization Literal['pretrained', 'random', 'checkpoint_transfer'] 'pretrained'
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 100
lr Optional[float] None
trainable_parameters bool False
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 float 0.005
optimization OptimizationConfig OptimizationConfig(name='inference_only')

Contracts