[Feat] +automatic multistart

ai4co · Nov 30, 2023 · ba7b8c4 · ba7b8c4
1 parent 0aaed80
commit ba7b8c4
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/rl4co/models/zoo/active_search/search.py b/rl4co/models/zoo/active_search/search.py
@@ -112,7 +112,7 @@ def training_step(self, batch, batch_idx):
         td_init = self.env.reset(batch)
         n_aug, n_start, n_runs = (
             self.augmentation.num_augment,
-            get_num_starts(td_init),
+            get_num_starts(td_init, self.env.name),
             self.hparams.num_parallel_runs,
         )
         td_init = self.augmentation(td_init)

diff --git a/rl4co/models/zoo/eas/search.py b/rl4co/models/zoo/eas/search.py
@@ -142,7 +142,7 @@ def training_step(self, batch, batch_idx):
         td_init = self.env.reset(batch)
         n_aug, n_start, n_runs = (
             self.augmentation.num_augment,
-            get_num_starts(td_init),
+            get_num_starts(td_init, self.env.name),
             self.hparams.num_parallel_runs,
         )
         td_init = self.augmentation(td_init)

diff --git a/rl4co/models/zoo/pomo/model.py b/rl4co/models/zoo/pomo/model.py
@@ -81,7 +81,7 @@ def shared_step(
     ):
         td = self.env.reset(batch)
         n_aug, n_start = self.num_augment, self.num_starts
-        n_start = get_num_starts(td) if n_start is None else n_start
+        n_start = get_num_starts(td, self.env.name) if n_start is None else n_start
 
         # During training, we do not augment the data
         if phase == "train":

diff --git a/rl4co/models/zoo/symnco/model.py b/rl4co/models/zoo/symnco/model.py
@@ -71,7 +71,7 @@ def shared_step(
     ):
         td = self.env.reset(batch)
         n_aug, n_start = self.num_augment, self.num_starts
-        n_start = get_num_starts(td) if n_start is None else n_start
+        n_start = get_num_starts(td, self.env.name) if n_start is None else n_start
 
         # Symmetric augmentation
         if n_aug > 1: