Skip to content

Conversation

@vijk777
Copy link
Collaborator

@vijk777 vijk777 commented Jan 27, 2026

  • Add a new model that stores a bank of initial states for the training dataset z0
  • The model only consists of a decoder, evolver (in latent), stimulus encoder. The stimulus is fed at every time step.
  • We apply a training loss at the staggered time points and simultaneously learn both the initial state and the model
  • Decided to add a new code path latent_stag.py and forked the default config into latent_stag_20step.yaml.
  • Handle SIGUSR2 from the cluster and terminate training so we get all the final figures & logging

vijk777 and others added 10 commits January 26, 2026 08:09
- Add LatentStagModel: decode-only architecture with learned z0
- Add latent_stag.py: training loop with z0 bank optimization
- Add diagnostics_stag.py: validation diagnostics for staggered model
- Add benchmark_rollout.py: compare LatentModel vs LatentStagModel
- Add latent_stag_20step.yaml: config with cross-validation datasets
- Remove simulation_config/training_data_path mutual exclusivity validation

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
use same key patterns as latent.py (Val/, CrossVal/) for consistency.
log diagnostics duration and long rollout metrics.

also fix device mismatch in stimulus_utils_test.py.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
plot_long_rollout_mse returns only a Figure, not a tuple.
call compute_rollout_stability_metrics separately to get metrics.
also fix incorrect type annotation in diagnostics.py.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
only log figures to tensorboard, don't save to disk.
remove unused run_dir parameter.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
match latent.py pattern: run_validation_diagnostics returns
(metrics, figures) and caller logs with proper prefixes.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@vijk777 vijk777 merged commit 0be21f7 into main Jan 27, 2026
1 check passed
@vijk777 vijk777 deleted the vj/learn_mat branch January 27, 2026 16:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants