pbt tests

Author: TheEimer

examples/configs/sac_pb2.yaml

@@ -7,7 +7,7 @@ defaults:
 
 hydra:
   sweeper:
-    budget: ${algorithm.total_timesteps}
+    budget: ${multiply:${algorithm.total_timesteps},${hydra.sweeper.sweeper_kwargs.optimizer_kwargs.population_size}}
     budget_variable: algorithm.total_timesteps
     loading_variable: load
     saving_variable: save

examples/configs/sac_pbt.yaml

@@ -7,7 +7,7 @@ defaults:
 
 hydra:
   sweeper:
-    budget: ${algorithm.total_timesteps}
+    budget: ${multiply:${algorithm.total_timesteps},${hydra.sweeper.sweeper_kwargs.optimizer_kwargs.population_size}}
     budget_variable: algorithm.total_timesteps
     loading_variable: load
     saving_variable: save

examples/sb3_rl_agent.py

@@ -9,6 +9,8 @@
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor
 
+OmegaConf.register_new_resolver("multiply", lambda x, y: x * y, replace=True)
+
 __copyright__ = "Copyright 2022, Theresa Eimer"
 __license__ = "3-clause BSD"
 

hydra_plugins/hyper_pbt/hyper_pbt.py

@@ -87,7 +87,7 @@ def ask(self):
             self.population_id += 1
             if iteration_end:
                 self.iteration += 1
-            return Info(config=config, budget=self.budget_per_run, load_path=None, seed=None), iteration_end
+            return Info(config=config, budget=self.budget_per_run, load_path=None, seed=None), iteration_end, False
         config, load_path = self.perturb_config(self.population_id)
         self.population_id += 1
         if iteration_end:
@@ -172,8 +172,8 @@ def tell(self, info, value):
             # Now that we have finished the iteration,
             # we can safely remove all checkpoints from the previous iteration
             print(f"Finished iteration {self.iteration}")
-            print("Remove checkpoints")
             if self.self_destruct and self.iteration > 1:
+                print("Remove checkpoints")
                 self.remove_checkpoints(self.iteration - 2)
 
     def finish_run(self, output_path):

hydra_plugins/hyper_pbt/pb2_utils.py

@@ -187,9 +187,8 @@ def K2(self, x, x2):
     def K(self, x, x2):
         """Compute the kernel."""
         # clip epsilons
-        self.epsilon_1 = min(self.epsilon_1, 0.5)
-
-        self.epsilon_2 = min(self.epsilon_2, 0.5)
+        self.epsilon_1 = GPy.core.Param("epsilon_1", min(self.epsilon_1, 0.5))
+        self.epsilon_2 = GPy.core.Param("epsilon_1", min(self.epsilon_2, 0.5))
 
         # format data
         if x2 is None:

hydra_plugins/hypersweeper/hypersweeper_sweeper.py

@@ -456,13 +456,12 @@ def run(self, verbose=False):
             infos = []
             t = 0
             terminate = False
-            while (
-                t < self.max_parallel
-                and not terminate
-                and not trial_termination
-                and not budget_termination
-                and not optimizer_termination
-            ):
+            while t < self.max_parallel and not terminate and not trial_termination and not optimizer_termination:
+                if not any(b is None for b in self.history["budget"]) and self.budget is not None:
+                    budget_termination = sum(self.history["budget"]) >= self.budget
+                    if budget_termination:
+                        break
+
                 try:
                     info, terminate, optimizer_termination = self.optimizer.ask()
                 except Exception as e:  # noqa: BLE001
@@ -484,8 +483,7 @@ def run(self, verbose=False):
                 if info.load_path is not None:
                     loading_paths.append(info.load_path)
                 infos.append(info)
-                if not any(b is None for b in self.history["budget"]) and self.budget is not None:
-                    budget_termination = sum(self.history["budget"]) >= self.budget
+
                 if self.n_trials is not None:
                     trial_termination = self.trials_run + len(configs) >= self.n_trials
 

pyproject.toml

@@ -70,7 +70,6 @@ log_level = "DEBUG"
 xfail_strict = true
 addopts = "--durations=10 -vv"
 markers = ["example: An example"]
-adopts = "--ignore=tests/test_runs/test_neps.py"  
 
 [tool.coverage.run]
 branch = true

tests/test_runs/test_neps.py

@@ -29,4 +29,4 @@
 #     runhistory = pd.read_csv("./tmp/mlp_neps_priorband/runhistory.csv")
 #     assert not runhistory.empty, "Run history is empty"
 #     assert len(runhistory) == 5, "Run history should contain 5 entries"
-#     shutil.rmtree(Path("./tmp/mlp_neps_priorband"))
+#     shutil.rmtree(Path("./tmp/mlp_neps_priorband"))

tests/test_runs/test_pbt.py

@@ -0,0 +1,47 @@
+import shutil
+import subprocess
+from pathlib import Path
+
+import pandas as pd
+
+
+def test_pbt_sac_example():
+    if Path("./tmp/sac_pbt").exists():
+        shutil.rmtree(Path("./tmp/sac_pbt"))
+    subprocess.call(
+        [
+            "python",
+            "examples/sb3_rl_agent.py",
+            "--config-name=sac_pbt",
+            "-m",
+            "algorithm.total_timesteps=1000",
+            "hydra.sweeper.sweeper_kwargs.optimizer_kwargs.config_interval=500",
+        ]
+    )
+    assert Path("./tmp/sac_pbt").exists(), "Run directory not created"
+    assert Path("./tmp/sac_pbt/runhistory.csv").exists(), "Run history file not created"
+    runhistory = pd.read_csv("./tmp/sac_pbt/runhistory.csv")
+    assert not runhistory.empty, "Run history is empty"
+    assert len(runhistory) == 4, "Run history should contain 4 entries"
+    shutil.rmtree(Path("./tmp/sac_pbt"))
+
+
+def test_pb2_sac_example():
+    if Path("./tmp/sac_pb2").exists():
+        shutil.rmtree(Path("./tmp/sac_pb2"))
+    subprocess.call(
+        [
+            "python",
+            "examples/sb3_rl_agent.py",
+            "--config-name=sac_pb2",
+            "-m",
+            "algorithm.total_timesteps=1000",
+            "hydra.sweeper.sweeper_kwargs.optimizer_kwargs.config_interval=500",
+        ]
+    )
+    assert Path("./tmp/sac_pb2").exists(), "Run directory not created"
+    assert Path("./tmp/sac_pb2/runhistory.csv").exists(), "Run history file not created"
+    runhistory = pd.read_csv("./tmp/sac_pb2/runhistory.csv")
+    assert not runhistory.empty, "Run history is empty"
+    assert len(runhistory) == 4, "Run history should contain 4 entries"
+    shutil.rmtree(Path("./tmp/sac_pb2"))