From ea8ed2cd2490ec26641719d0bd9709942d8e9249 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Tue, 1 Sep 2026 09:05:55 +0200
Subject: [PATCH 01/22] feat: Add HEALPix curriculum

---
 packages/common/src/weathergen/common/config.py | 11 +++++++++++
 src/weathergen/train/trainer.py                 | 17 +++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 1f506fffb..937f6ad40 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -463,6 +463,17 @@ def load_merge_configs(
     assert isinstance(c, Config)
     c = _sanitize_time_keys(c)
 
+    if c.get("healpix_curriculum"):
+        istep = c.get("general", {}).get("istep", 0)
+        cumulative = 0
+        current_hl = None
+        for hl, steps in c.healpix_curriculum.items():
+            cumulative += steps
+            current_hl = int(hl)
+            if istep < cumulative:
+                break
+        c.healpix_level = current_hl
+
     return c
 
 
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 276da0bd6..4c5e3f491 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -407,6 +407,11 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
                 )
             self.save_model(mini_epoch)
 
+            if getattr(self.cf, "_curriculum_exit", False):
+                if is_root():
+                    logger.info("Curriculum stage completed. Exiting training loop.")
+                break
+
         # log final model
         self.save_model(self.training_cfg.num_mini_epochs)
 
@@ -568,6 +573,18 @@ def train(self, mini_epoch):
 
             self.cf.general.istep += 1
 
+            if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
+                cumulative = 0
+                for hl, steps in self.cf.healpix_curriculum.items():
+                    cumulative += steps
+                    if self.cf.healpix_level == int(hl):
+                        break
+                if self.cf.general.istep >= cumulative:
+                    if is_root():
+                        logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")
+                    self.cf._curriculum_exit = True
+                    break
+
         self.dataset.advance()
 
     def validate(self, mini_epoch, mode_cfg, batch_size):

From a422f9ebcd4ddfdbf721eaef6a21cf219b12c60c Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Tue, 1 Sep 2026 09:41:55 +0200
Subject: [PATCH 02/22] add HL curriculum config

---
 config/config_curriculum.yml    | 13 ++++++++++
 src/weathergen/run_train.py     | 43 +++++++++++++++++++++++++++++++--
 src/weathergen/train/trainer.py |  6 +----
 3 files changed, 55 insertions(+), 7 deletions(-)
 create mode 100644 config/config_curriculum.yml

diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml
new file mode 100644
index 000000000..ebbca19e7
--- /dev/null
+++ b/config/config_curriculum.yml
@@ -0,0 +1,13 @@
+healpix_curriculum:
+  3: 1000
+  4: 1000
+  5: 1500
+  6: 3000
+
+curriculum_streams:
+  3: "./config/streams/era5_8deg/"
+  4: "./config/streams/era5_4deg/"
+  5: "./config/streams/era5_2deg/"
+  6: "./config/streams/era5_1deg/"
+
+streams_directory: "${curriculum_streams.${healpix_level}}"
diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 7995b5864..1e61971d7 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -147,7 +147,27 @@ def run_continue(args):
     trainer = Trainer(cf.train_logging)
 
     try:
-        trainer.run(cf, devices, args.from_run_id, args.mini_epoch)
+        from_run_id_iter = args.from_run_id
+        mini_epoch_iter = args.mini_epoch
+        first_run = True
+        while True:
+            if not first_run:
+                cf = config.load_merge_configs(
+                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, {}, cli_overwrite
+                )
+                cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf.streams = config.load_streams(Path(cf.streams_directory))
+                trainer = Trainer(cf.train_logging)
+                
+            trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
+            first_run = False
+            
+            if not getattr(cf, "_curriculum_exit", False):
+                break
+                
+            logger.info("Restarting training for next curriculum stage...")
+            from_run_id_iter = cf.general.run_id
+            mini_epoch_iter = -1
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()
@@ -188,7 +208,26 @@ def run_train(args):
     trainer = Trainer(cf.train_logging)
 
     try:
-        trainer.run(cf, devices)
+        from_run_id_iter = None
+        mini_epoch_iter = None
+        while True:
+            if from_run_id_iter is not None:
+                cf = config.load_merge_configs(
+                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, cli_overwrite
+                )
+                cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf.streams = config.load_streams(Path(cf.streams_directory))
+                trainer = Trainer(cf.train_logging)
+                trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
+            else:
+                trainer.run(cf, devices)
+                
+            if not getattr(cf, "_curriculum_exit", False):
+                break
+                
+            logger.info("Restarting training for next curriculum stage...")
+            from_run_id_iter = cf.general.run_id
+            mini_epoch_iter = -1
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 4c5e3f491..422142e1b 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -574,11 +574,7 @@ def train(self, mini_epoch):
             self.cf.general.istep += 1
 
             if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
-                cumulative = 0
-                for hl, steps in self.cf.healpix_curriculum.items():
-                    cumulative += steps
-                    if self.cf.healpix_level == int(hl):
-                        break
+                cumulative = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
                 if self.cf.general.istep >= cumulative:
                     if is_root():
                         logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")

From 6e4572edf77c83fe355696ca5df47e2a13a8cff7 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 09:11:40 +0200
Subject: [PATCH 03/22] change config for testing

---
 config/config_curriculum.yml | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml
index ebbca19e7..5479bbfbb 100644
--- a/config/config_curriculum.yml
+++ b/config/config_curriculum.yml
@@ -1,13 +1,13 @@
 healpix_curriculum:
-  3: 1000
-  4: 1000
-  5: 1500
-  6: 3000
+  3: 250
+  4: 250
+  5: 250
+  6: 250
 
 curriculum_streams:
-  3: "./config/streams/era5_8deg/"
-  4: "./config/streams/era5_4deg/"
-  5: "./config/streams/era5_2deg/"
+  3: "./config/streams/era5_1deg/"
+  4: "./config/streams/era5_1deg/"
+  5: "./config/streams/era5_1deg/"
   6: "./config/streams/era5_1deg/"
 
 streams_directory: "${curriculum_streams.${healpix_level}}"

From 08e65fc589aace203b38d564957c579ca9b82a03 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 09:15:57 +0200
Subject: [PATCH 04/22] fix: catch interpolation error during config load

---
 packages/common/src/weathergen/common/config.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 937f6ad40..26f40ae7c 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -480,7 +480,12 @@ def load_merge_configs(
 def _load_streams_in_config(config: Config) -> Config:
     """If the config contains a streams_directory, loads the streams and returns the config with
     the streams set."""
-    streams_directory = config.get("streams_directory", None)
+    import omegaconf
+    try:
+        streams_directory = config.get("streams_directory", None)
+    except (omegaconf.errors.InterpolationKeyError, omegaconf.errors.InterpolationResolutionError):
+        streams_directory = None
+        
     config = config.copy()
     if streams_directory is not None:
         streams_directory = Path(streams_directory)

From fe0310791ab14b1124e7ec407e84171081f976f9 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 09:24:36 +0200
Subject: [PATCH 05/22] fix: avoid evaluating streams_directory early and fix
 imports

---
 packages/common/src/weathergen/common/config.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 26f40ae7c..d2deff0b9 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -24,6 +24,7 @@
 import yaml.scanner
 from omegaconf import DictConfig, ListConfig, OmegaConf
 from omegaconf.omegaconf import open_dict
+from omegaconf.errors import InterpolationKeyError, InterpolationResolutionError
 
 from weathergen.common.io import StoreType
 from weathergen.common.paths import _REPO_ROOT, get_wg_private_path
@@ -456,7 +457,7 @@ def load_merge_configs(
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
         # streams from an overwrite's streams_directory replace inherited streams
-        if any(o.get("streams_directory") is not None for o in overwrite_configs):
+        if any("streams_directory" in o for o in overwrite_configs):
             base_config.streams = None
     # use OmegaConf.unsafe_merge if too slow
     c = OmegaConf.merge(base_config, private_config, *overwrite_configs)
@@ -480,10 +481,9 @@ def load_merge_configs(
 def _load_streams_in_config(config: Config) -> Config:
     """If the config contains a streams_directory, loads the streams and returns the config with
     the streams set."""
-    import omegaconf
     try:
         streams_directory = config.get("streams_directory", None)
-    except (omegaconf.errors.InterpolationKeyError, omegaconf.errors.InterpolationResolutionError):
+    except (InterpolationKeyError, InterpolationResolutionError):
         streams_directory = None
         
     config = config.copy()

From 9f45c08481de519f1278aec75154d92d48a2344c Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 09:36:38 +0200
Subject: [PATCH 06/22] fix: map curriculum streams natively to avoid omegaconf
 int key error

---
 packages/common/src/weathergen/common/config.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index d2deff0b9..2ffc326d3 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -474,6 +474,10 @@ def load_merge_configs(
             if istep < cumulative:
                 break
         c.healpix_level = current_hl
+        
+        if c.get("curriculum_streams"):
+            # Support both integer and string keys in the yaml
+            c.streams_directory = c.curriculum_streams.get(current_hl) or c.curriculum_streams.get(str(current_hl))
 
     return c
 

From 828f8688cb9791b2de756b49af48bbf3a1064028 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 09:38:35 +0200
Subject: [PATCH 07/22] fix: remove buggy streams_directory interpolation from
 config

---
 config/config_curriculum.yml | 2 --
 1 file changed, 2 deletions(-)

diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml
index 5479bbfbb..27aeb1c52 100644
--- a/config/config_curriculum.yml
+++ b/config/config_curriculum.yml
@@ -9,5 +9,3 @@ curriculum_streams:
   4: "./config/streams/era5_1deg/"
   5: "./config/streams/era5_1deg/"
   6: "./config/streams/era5_1deg/"
-
-streams_directory: "${curriculum_streams.${healpix_level}}"

From 2413625eb14ad7e423484ee4e9eb5694808957f7 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 10:03:26 +0200
Subject: [PATCH 08/22] fix: safely handle integer keys when stripping
 interpolations from OmegaConf dicts

---
 .../common/src/weathergen/common/config.py    | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 2ffc326d3..fd4ea043e 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -140,25 +140,25 @@ def _strip_interpolation(conf: Config) -> Config:
     """Recursively convert interpolated timedelta/datetime objects to strings."""
     stripped = {}
     if OmegaConf.is_dict(conf):
-        for key in list(conf.keys()):
-            key = str(key)
-            if OmegaConf.is_missing(conf, key):
+        for orig_key in list(conf.keys()):
+            str_key = str(orig_key)
+            if OmegaConf.is_missing(conf, orig_key):
                 val = "???"
-            elif OmegaConf.is_config(conf[key]):
-                val = _strip_interpolation(conf[key])
-            elif key.startswith("_"):
+            elif OmegaConf.is_config(conf[orig_key]):
+                val = _strip_interpolation(conf[orig_key])
+            elif str_key.startswith("_"):
                 continue  # Skip hidden/backup keys
-            elif OmegaConf.is_interpolation(conf, key):
-                raw_key = f"_{key}"
+            elif OmegaConf.is_interpolation(conf, orig_key):
+                raw_key = f"_{str_key}"
                 assert raw_key in conf, (
-                    f"Backup key: {raw_key} expected for interpolated key: {key}"
+                    f"Backup key: {raw_key} expected for interpolated key: {orig_key}"
                 )
                 # Retrieve the value from the backup key (resolves interpolation)
                 val = conf[raw_key]
             else:
-                val = conf[key]
+                val = conf[orig_key]
 
-            stripped[key] = val
+            stripped[str_key] = val
     elif OmegaConf.is_list(conf):
         stripped = [
             _strip_interpolation(item) if OmegaConf.is_config(item) else item for item in conf

From 8d329bd26d8a7a7939b2fd185e52ae720e43b59b Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 10:18:21 +0200
Subject: [PATCH 09/22] fix: read _curriculum_exit from trainer.cf, not the
 local cf copy

---
 src/weathergen/run_train.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 1e61971d7..b433dad56 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -162,11 +162,11 @@ def run_continue(args):
             trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
             first_run = False
             
-            if not getattr(cf, "_curriculum_exit", False):
+            if not getattr(trainer.cf, "_curriculum_exit", False):
                 break
                 
             logger.info("Restarting training for next curriculum stage...")
-            from_run_id_iter = cf.general.run_id
+            from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
     except Exception:
         extype, value, tb = sys.exc_info()
@@ -222,11 +222,11 @@ def run_train(args):
             else:
                 trainer.run(cf, devices)
                 
-            if not getattr(cf, "_curriculum_exit", False):
+            if not getattr(trainer.cf, "_curriculum_exit", False):
                 break
                 
             logger.info("Restarting training for next curriculum stage...")
-            from_run_id_iter = cf.general.run_id
+            from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
     except Exception:
         extype, value, tb = sys.exc_info()

From 5329eaaa38f8b27381fdf0123f5636001e7ea836 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 10:30:17 +0200
Subject: [PATCH 10/22] fix: save as _latest checkpoint on curriculum exit so
 restart can find it

---
 src/weathergen/train/trainer.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 422142e1b..531291c10 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -413,7 +413,10 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
                 break
 
         # log final model
-        self.save_model(self.training_cfg.num_mini_epochs)
+        if getattr(self.cf, "_curriculum_exit", False):
+            self.save_model(-1)
+        else:
+            self.save_model(self.training_cfg.num_mini_epochs)
 
     def validate_before_training(self):
         """

From 02d1532fe50863341e34358cc6f234108e221568 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 10:42:16 +0200
Subject: [PATCH 11/22] fix: re-apply init_ddp on curriculum restart to fix
 device placement

---
 src/weathergen/run_train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index b433dad56..8db3152ad 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -156,6 +156,7 @@ def run_continue(args):
                     args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, {}, cli_overwrite
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf = Trainer.init_ddp(cf)
                 cf.streams = config.load_streams(Path(cf.streams_directory))
                 trainer = Trainer(cf.train_logging)
                 
@@ -216,6 +217,7 @@ def run_train(args):
                     args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, cli_overwrite
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf = Trainer.init_ddp(cf)
                 cf.streams = config.load_streams(Path(cf.streams_directory))
                 trainer = Trainer(cf.train_logging)
                 trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)

From f7d068e5e7086d8ac80fd682a8bdc767630fa1a0 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 10:53:09 +0200
Subject: [PATCH 12/22] fix: correctly read rank/local_rank from dist when
 already initialized on curriculum restart

---
 src/weathergen/train/trainer_base.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/weathergen/train/trainer_base.py b/src/weathergen/train/trainer_base.py
index d0359892a..b142d60d2 100644
--- a/src/weathergen/train/trainer_base.py
+++ b/src/weathergen/train/trainer_base.py
@@ -136,6 +136,11 @@ def init_ddp(cf):
                     dist.all_reduce(l_seed, op=torch.distributed.ReduceOp.SUM)
                     cf.data_loader_rng_seed = l_seed.item()
 
+        if dist.is_initialized():
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+            local_rank = int(os.environ.get("LOCAL_RANK", os.environ.get("SLURM_LOCALID", "0")))
+
         cf.world_size = world_size
         cf.rank = rank
         cf.local_rank = local_rank

From ac9cd0c82c3f14f8f3b3ac6073c6c12399c38536 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 11:13:05 +0200
Subject: [PATCH 13/22] fix: preserve istep across curriculum restarts so
 healpix_level is computed correctly

---
 src/weathergen/run_train.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 8db3152ad..e84840606 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -150,10 +150,11 @@ def run_continue(args):
         from_run_id_iter = args.from_run_id
         mini_epoch_iter = args.mini_epoch
         first_run = True
+        istep_override = {}
         while True:
             if not first_run:
                 cf = config.load_merge_configs(
-                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, {}, cli_overwrite
+                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, istep_override, cli_overwrite
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
                 cf = Trainer.init_ddp(cf)
@@ -169,6 +170,7 @@ def run_continue(args):
             logger.info("Restarting training for next curriculum stage...")
             from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
+            istep_override = {"general": {"istep": trainer.cf.general.istep}}
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()
@@ -211,10 +213,11 @@ def run_train(args):
     try:
         from_run_id_iter = None
         mini_epoch_iter = None
+        istep_override = {}
         while True:
             if from_run_id_iter is not None:
                 cf = config.load_merge_configs(
-                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, cli_overwrite
+                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, istep_override, cli_overwrite
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
                 cf = Trainer.init_ddp(cf)
@@ -230,6 +233,7 @@ def run_train(args):
             logger.info("Restarting training for next curriculum stage...")
             from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
+            istep_override = {"general": {"istep": trainer.cf.general.istep}}
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()

From 61d18046cd295aaff5ae2ab9195ddf01d2498a10 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 11:27:35 +0200
Subject: [PATCH 14/22] debug: log curriculum state at first batch to diagnose
 exit not firing

---
 src/weathergen/train/trainer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 531291c10..dc7f3109e 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -577,6 +577,9 @@ def train(self, mini_epoch):
             self.cf.general.istep += 1
 
             if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
+                if bidx == 0 and is_root():
+                    cumulative_debug = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
+                    logger.info(f"[Curriculum] healpix_level={self.cf.healpix_level}, istep={self.cf.general.istep}, cumulative_exit_at={cumulative_debug}, curriculum={dict(self.cf.healpix_curriculum)}")
                 cumulative = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
                 if self.cf.general.istep >= cumulative:
                     if is_root():

From ed67dc6dedd1cf68687ba4a139cb93d6be5cef2e Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 11:28:09 +0200
Subject: [PATCH 15/22] change isteps for faster testing

---
 config/config_curriculum.yml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml
index 27aeb1c52..4945ed80e 100644
--- a/config/config_curriculum.yml
+++ b/config/config_curriculum.yml
@@ -1,8 +1,8 @@
 healpix_curriculum:
-  3: 250
-  4: 250
-  5: 250
-  6: 250
+  3: 100
+  4: 100
+  5: 100
+  6: 100
 
 curriculum_streams:
   3: "./config/streams/era5_1deg/"

From a1c51301f340fa7c4c0d0743151ae93159370911 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 11:37:01 +0200
Subject: [PATCH 16/22] chore: remove curriculum debug log

---
 src/weathergen/train/trainer.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index dc7f3109e..531291c10 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -577,9 +577,6 @@ def train(self, mini_epoch):
             self.cf.general.istep += 1
 
             if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
-                if bidx == 0 and is_root():
-                    cumulative_debug = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
-                    logger.info(f"[Curriculum] healpix_level={self.cf.healpix_level}, istep={self.cf.general.istep}, cumulative_exit_at={cumulative_debug}, curriculum={dict(self.cf.healpix_curriculum)}")
                 cumulative = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
                 if self.cf.general.istep >= cumulative:
                     if is_root():

From f573b84f2432726ac05950bf5c4dbf42fb9c6da2 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 11:40:17 +0200
Subject: [PATCH 17/22] fix: deduplicate string/int keys in healpix_curriculum
 dictionary on restart

---
 packages/common/src/weathergen/common/config.py | 7 ++++---
 src/weathergen/train/trainer.py                 | 3 ++-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index fd4ea043e..193bb635c 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -468,9 +468,10 @@ def load_merge_configs(
         istep = c.get("general", {}).get("istep", 0)
         cumulative = 0
         current_hl = None
-        for hl, steps in c.healpix_curriculum.items():
-            cumulative += steps
-            current_hl = int(hl)
+        unique_curr = {int(hl): steps for hl, steps in c.healpix_curriculum.items()}
+        for hl in sorted(unique_curr.keys()):
+            cumulative += unique_curr[hl]
+            current_hl = hl
             if istep < cumulative:
                 break
         c.healpix_level = current_hl
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 531291c10..08dacf885 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -577,7 +577,8 @@ def train(self, mini_epoch):
             self.cf.general.istep += 1
 
             if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
-                cumulative = sum(steps for hl, steps in self.cf.healpix_curriculum.items() if int(hl) <= self.cf.healpix_level)
+                unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
+                cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
                 if self.cf.general.istep >= cumulative:
                     if is_root():
                         logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")

From 498c4c650a148bc31e738a9e0961cbe55a0d20c0 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 13:12:41 +0200
Subject: [PATCH 18/22] fix: prevent infinite restart loops when curriculum
 reaches its maximum level

---
 src/weathergen/train/trainer.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 08dacf885..ee433d367 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -578,12 +578,14 @@ def train(self, mini_epoch):
 
             if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
                 unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
-                cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
-                if self.cf.general.istep >= cumulative:
-                    if is_root():
-                        logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")
-                    self.cf._curriculum_exit = True
-                    break
+                max_hl = max(unique_curr.keys())
+                if self.cf.healpix_level < max_hl:
+                    cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
+                    if self.cf.general.istep >= cumulative:
+                        if is_root():
+                            logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")
+                        self.cf._curriculum_exit = True
+                        break
 
         self.dataset.advance()
 

From 66d763a99f45828132c5a726a0c5f799e8c9d3a3 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 13:13:32 +0200
Subject: [PATCH 19/22] increase isteps for meaningful plots

---
 config/config_curriculum.yml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml
index 4945ed80e..27aeb1c52 100644
--- a/config/config_curriculum.yml
+++ b/config/config_curriculum.yml
@@ -1,8 +1,8 @@
 healpix_curriculum:
-  3: 100
-  4: 100
-  5: 100
-  6: 100
+  3: 250
+  4: 250
+  5: 250
+  6: 250
 
 curriculum_streams:
   3: "./config/streams/era5_1deg/"

From 28a245f866ceee7e2887302dac32277aa9c08349 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 14:18:06 +0200
Subject: [PATCH 20/22] feat: add clear logging to explain curriculum vs
 standard training run length

---
 src/weathergen/train/trainer.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index ee433d367..07486842d 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -358,6 +358,22 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         if self.cf.general.istep > 0 and is_root():
             logger.info(f"Continuing run with learning rate: {self.lr_scheduler.get_lr()}")
 
+        if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum and is_root():
+            unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
+            max_hl = max(unique_curr.keys())
+            if self.cf.healpix_level < max_hl:
+                cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
+                logger.info(
+                    f"Curriculum active: Training HEALPix level {self.cf.healpix_level}. "
+                    f"Next stage will begin at istep {cumulative}. "
+                    f"(Note: Total run length is dictated by num_mini_epochs)"
+                )
+            else:
+                logger.info(
+                    f"Curriculum max level ({max_hl}) reached. "
+                    f"Continuing standard training until num_mini_epochs limit."
+                )
+
         # Instantiate loss calculator modules to compute losses
         self.loss_calculator = LossCalculator(cf, self.training_cfg, TRAIN, device=self.device)
         val_cfg = self.validation_cfg

From d9fb931e0a279fb8609a42bc1c7e53e63de05d27 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 14:33:51 +0200
Subject: [PATCH 21/22] ruff

---
 .../common/src/weathergen/common/config.py    | 10 ++++---
 src/weathergen/run_train.py                   | 26 ++++++++++++++-----
 src/weathergen/train/trainer.py               | 13 +++++++---
 3 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 193bb635c..2ca738f11 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -23,8 +23,8 @@
 import yaml.constructor
 import yaml.scanner
 from omegaconf import DictConfig, ListConfig, OmegaConf
-from omegaconf.omegaconf import open_dict
 from omegaconf.errors import InterpolationKeyError, InterpolationResolutionError
+from omegaconf.omegaconf import open_dict
 
 from weathergen.common.io import StoreType
 from weathergen.common.paths import _REPO_ROOT, get_wg_private_path
@@ -475,10 +475,12 @@ def load_merge_configs(
             if istep < cumulative:
                 break
         c.healpix_level = current_hl
-        
+
         if c.get("curriculum_streams"):
             # Support both integer and string keys in the yaml
-            c.streams_directory = c.curriculum_streams.get(current_hl) or c.curriculum_streams.get(str(current_hl))
+            c.streams_directory = c.curriculum_streams.get(current_hl) or c.curriculum_streams.get(
+                str(current_hl)
+            )
 
     return c
 
@@ -490,7 +492,7 @@ def _load_streams_in_config(config: Config) -> Config:
         streams_directory = config.get("streams_directory", None)
     except (InterpolationKeyError, InterpolationResolutionError):
         streams_directory = None
-        
+
     config = config.copy()
     if streams_directory is not None:
         streams_directory = Path(streams_directory)
diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index e84840606..a670887fa 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -154,19 +154,25 @@ def run_continue(args):
         while True:
             if not first_run:
                 cf = config.load_merge_configs(
-                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, istep_override, cli_overwrite
+                    args.private_config,
+                    from_run_id_iter,
+                    mini_epoch_iter,
+                    args.base_config,
+                    *args.config,
+                    istep_override,
+                    cli_overwrite,
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
                 cf = Trainer.init_ddp(cf)
                 cf.streams = config.load_streams(Path(cf.streams_directory))
                 trainer = Trainer(cf.train_logging)
-                
+
             trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
             first_run = False
-            
+
             if not getattr(trainer.cf, "_curriculum_exit", False):
                 break
-                
+
             logger.info("Restarting training for next curriculum stage...")
             from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
@@ -217,7 +223,13 @@ def run_train(args):
         while True:
             if from_run_id_iter is not None:
                 cf = config.load_merge_configs(
-                    args.private_config, from_run_id_iter, mini_epoch_iter, args.base_config, *args.config, istep_override, cli_overwrite
+                    args.private_config,
+                    from_run_id_iter,
+                    mini_epoch_iter,
+                    args.base_config,
+                    *args.config,
+                    istep_override,
+                    cli_overwrite,
                 )
                 cf = config.set_run_id(cf, cf.general.run_id, True)
                 cf = Trainer.init_ddp(cf)
@@ -226,10 +238,10 @@ def run_train(args):
                 trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
             else:
                 trainer.run(cf, devices)
-                
+
             if not getattr(trainer.cf, "_curriculum_exit", False):
                 break
-                
+
             logger.info("Restarting training for next curriculum stage...")
             from_run_id_iter = trainer.cf.general.run_id
             mini_epoch_iter = -1
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 07486842d..80da137e0 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -362,7 +362,9 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
             unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
             max_hl = max(unique_curr.keys())
             if self.cf.healpix_level < max_hl:
-                cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
+                cumulative = sum(
+                    steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level
+                )
                 logger.info(
                     f"Curriculum active: Training HEALPix level {self.cf.healpix_level}. "
                     f"Next stage will begin at istep {cumulative}. "
@@ -596,10 +598,15 @@ def train(self, mini_epoch):
                 unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
                 max_hl = max(unique_curr.keys())
                 if self.cf.healpix_level < max_hl:
-                    cumulative = sum(steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level)
+                    cumulative = sum(
+                        steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level
+                    )
                     if self.cf.general.istep >= cumulative:
                         if is_root():
-                            logger.info(f"Curriculum stage for HEALPix level {self.cf.healpix_level} finished at istep {self.cf.general.istep}. Exiting mini_epoch early.")
+                            logger.info(
+                                f"Curriculum stage for HEALPix level {self.cf.healpix_level} "
+                                f"finished at istep {self.cf.general.istep}. Exiting early."
+                            )
                         self.cf._curriculum_exit = True
                         break
 

From fcbc70698a6497bb71ac3c4e1545bfb6d48e54a4 Mon Sep 17 00:00:00 2001
From: TillHae <till@web-hauer.de>
Date: Wed, 2 Sep 2026 14:42:09 +0200
Subject: [PATCH 22/22] fix: replace getattr with .get to satisfy pylint W0141

---
 src/weathergen/run_train.py     | 4 ++--
 src/weathergen/train/trainer.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index a670887fa..23132f8bc 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -170,7 +170,7 @@ def run_continue(args):
             trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
             first_run = False
 
-            if not getattr(trainer.cf, "_curriculum_exit", False):
+            if not trainer.cf.get("_curriculum_exit", False):
                 break
 
             logger.info("Restarting training for next curriculum stage...")
@@ -239,7 +239,7 @@ def run_train(args):
             else:
                 trainer.run(cf, devices)
 
-            if not getattr(trainer.cf, "_curriculum_exit", False):
+            if not trainer.cf.get("_curriculum_exit", False):
                 break
 
             logger.info("Restarting training for next curriculum stage...")
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 80da137e0..13457fbb3 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -425,13 +425,13 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
                 )
             self.save_model(mini_epoch)
 
-            if getattr(self.cf, "_curriculum_exit", False):
+            if self.cf.get("_curriculum_exit", False):
                 if is_root():
                     logger.info("Curriculum stage completed. Exiting training loop.")
                 break
 
         # log final model
-        if getattr(self.cf, "_curriculum_exit", False):
+        if self.cf.get("_curriculum_exit", False):
             self.save_model(-1)
         else:
             self.save_model(self.training_cfg.num_mini_epochs)