Update train_kaggle.py

train_kaggle.py

@@ -217,12 +217,25 @@ def build_model(env, args: argparse.Namespace):
         raise RuntimeError(f"ML stack missing: {_ML_IMPORT_ERROR}")
 
     if _GRU_AVAILABLE and create_gru_weather_policy_kwargs is not None:
+        # GRUWeatherFeaturesExtractor's forward() reshapes per-zone features
+        # into (batch, n_zones * hidden_size * 2) -- features_dim MUST match
+        # that exactly, or SB3's MlpExtractor (built from the declared
+        # features_dim) gets a differently-shaped tensor and crashes at the
+        # first policy.forward() call. Confirmed empirically: the actual
+        # working checkpoints (final_normal/drought/humidity.zip) all have
+        # features_dim == n_zones * hidden_size * 2 (256/384/512 for
+        # n_zones=2/3/4, hidden_size=64) -- NOT hidden_size * 2, which only
+        # coincides with the correct value when n_zones == 1.
+        features_dim = args.n_zones * args.hidden_size * 2
         policy_kwargs = create_gru_weather_policy_kwargs(
             hidden_size=args.hidden_size,
-            features_dim=args.hidden_size * 2,
+            features_dim=features_dim,
         )
         policy = get_equivariant_policy_class() if get_equivariant_policy_class is not None else "MultiInputPolicy"
-        logger.info("Using GRU policy (hidden_size=%d)", args.hidden_size)
+        logger.info(
+            "Using GRU policy (hidden_size=%d, n_zones=%d, features_dim=%d)",
+            args.hidden_size, args.n_zones, features_dim,
+        )
     else:
         policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64]))
         policy = "MultiInputPolicy"