🐛 bug fixes

scripts/hpo.py

@@ -157,6 +157,14 @@ def _create_base_task(
         add_task_init_call=False,
     )
 
+    # Explicitly set Hydra config-group choices so cloned tasks
+    # pick up the correct env / runner / training groups.
+    # Task.create() does not populate the Hydra parameter section
+    # because Hydra never actually runs during creation.
+    base_task.set_parameter("Hydra/env", env)
+    base_task.set_parameter("Hydra/runner", runner)
+    base_task.set_parameter("Hydra/training", training)
+
     # Set docker config
     base_task.set_base_docker(
         "registry.kube.optimize/worker-image:latest",
@@ -198,12 +206,12 @@ def main() -> None:
         help="Total HPO trial budget",
     )
     parser.add_argument(
-        "--min-budget", type=int, default=3,
-        help="Minimum budget (epochs) per trial",
+        "--min-budget", type=int, default=50_000,
+        help="Minimum budget (total_timesteps) per trial",
     )
     parser.add_argument(
-        "--max-budget", type=int, default=81,
-        help="Maximum budget (epochs) for promoted trials",
+        "--max-budget", type=int, default=500_000,
+        help="Maximum budget (total_timesteps) for promoted trials",
     )
     parser.add_argument("--eta", type=int, default=3, help="Successive halving reduction factor")
     parser.add_argument(
@@ -303,6 +311,7 @@ def main() -> None:
         pool_period_min=1,
         time_limit_per_job=240,  # 4 hours per trial max
         eta=args.eta,
+        budget_param_name="Hydra/training.total_timesteps",
     )
 
     # Send this HPO controller to a remote services worker