Mh/diffusion full pipeline latent rmse curve

[moritzhauschulz]

Description

This PR enables plotting of latent RMSE during inference of the diffusion model. Simply enable via test_config.latent_rollout_rmse=True.

NB: This involved adding a new batch to the sampler (see changes). I dropped the idea of reusing the num_steps_input as this would mess with the indexing of the entire pipeline, and there are many ways that it could fail silently. This implementation is a bit more intrusive, but all changes are gated with a flag, so that the standard path is minimally disrupted. If this is to become a permanent feature on develop, I think a broader discussion should be had about what is the best way.

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[clessig]

Why do we need this? How do you distinguish between latent (R)MSE and physical (R)MSE in the training config?