Merge pull request #58 from issp-center-dev/fix_57

fix `policy.write` method to write the last action
issp-center-dev · Aug 20, 2024 · e9ab4e3 · e9ab4e3
2 parents 606c76b + 0a180ed
commit e9ab4e3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/physbo/search/discrete/policy.py b/physbo/search/discrete/policy.py
@@ -138,11 +138,15 @@ def write(
             time_run_simulator=time_run_simulator,
         )
         self.training.add(X=X, t=t, Z=Z)
-        local_index = np.searchsorted(self.actions, action)
-        local_index = local_index[
-            np.take(self.actions, local_index, mode="clip") == action
-        ]
-        self.actions = self._delete_actions(local_index)
+
+        # remove the selected actions from the list of candidates if exists
+        if len(self.actions) > 0:
+            local_index = np.searchsorted(self.actions, action)
+            local_index = local_index[
+                np.take(self.actions, local_index, mode="clip") == action
+            ]
+            self.actions = self._delete_actions(local_index)
+
         if self.new_data is None:
             self.new_data = variable(X=X, t=t, Z=Z)
         else:

diff --git a/physbo/search/discrete_multi/policy.py b/physbo/search/discrete_multi/policy.py
@@ -106,11 +106,14 @@ def write(
             else:
                 self.new_data_list[i].add(X=X, t=t[:, i], Z=Z)
             self.training_list[i].add(X=X, t=t[:, i], Z=Z)
-        local_index = np.searchsorted(self.actions, action)
-        local_index = local_index[
-            np.take(self.actions, local_index, mode="clip") == action
-        ]
-        self.actions = self._delete_actions(local_index)
+
+        # remove action from candidates if exists
+        if len(self.actions) > 0:
+            local_index = np.searchsorted(self.actions, action)
+            local_index = local_index[
+                np.take(self.actions, local_index, mode="clip") == action
+            ]
+            self.actions = self._delete_actions(local_index)
 
     def _model(self, i):
         training = self.training_list[i]