Add convergence check

Saurav Agarwal · Saurav Agarwal · commit af89c0fee6a1 · 2024-09-10T17:31:10.000-04:00
diff --git a/python/coverage_control/algorithms/controllers.py b/python/coverage_control/algorithms/controllers.py
@@ -150,4 +150,7 @@ def step(self, env):
         point_vector_actions = PointVector(actions.cpu().numpy())
         env.StepActions(point_vector_actions)
 
+        # Check if actions are all zeros (1e-12)
+        if torch.allclose(actions, torch.zeros_like(actions), atol=1e-5):
+            return env.GetObjectiveValue(), True
         return env.GetObjectiveValue(), False
diff --git a/python/scripts/evaluators/constrained_learning.py b/python/scripts/evaluators/constrained_learning.py
@@ -155,17 +155,19 @@ def update_idf(self, coefficients, normalize=False):
             )
 
     def advance_state(self):
-        self.controller.step(self.env_main)
+        obj_val, is_converged = self.controller.step(self.env_main)
         self.step_counter = self.step_counter + 1
 
         if self.generate_video and self.step_counter % 1 == 0:
             self.env_main.RecordPlotData()
             # self.env_main.PlotRobotLocalMap("./robot_maps/", 0, self.step_counter)
             # self.env_main.PlotRobotSensorView("./robot_maps/", 0, self.step_counter)
-        robot_positions = self.env_main.GetRobotPositions()
-
-        for env in self.envs:
-            env.SetGlobalRobotPositions(robot_positions)
+        if is_converged == False:
+            robot_positions = self.env_main.GetRobotPositions()
+            for env in self.envs:
+                env.SetGlobalRobotPositions(robot_positions)
+        is_state_updated = not is_converged
+        return is_state_updated
 
     def compute_obj_values(self):
         obj_values = np.array(
@@ -188,33 +190,37 @@ def evaluate(self):
         K = self.num_steps // self.T_0
         self.lambda_duals = self.fun_dual_updater(self.dual_updater, self.lambda_duals)
 
+        self.update_idf(self.lambda_duals, normalize=self.normalize)
         obj_values = self.compute_obj_values()
         print(
             f"{0} Objective values: {obj_values} Lambda duals: {self.lambda_duals}, self alphas: {self.alphas}"
         )
 
         for k in range(K):
-            obj_values = np.zeros(self.num_idfs)
-
-            self.update_idf(self.lambda_duals, normalize=self.normalize)
 
+            is_state_updated = False
             for _ in range(self.T_0):
-                self.advance_state()
+                is_state_updated = is_state_updated or self.advance_state()
                 # obj_values += self.compute_obj_values()  # This is a vector
 
-            # obj_values /= self.T_0
-            obj_values = self.compute_obj_values()
-            obj_max = np.max(obj_values)
-            self.lambda_duals = np.maximum(
-                self.lambda_duals
-                + self.eta_dual * (obj_values - self.alphas) / obj_max,
-                0,
-            )
-            if self.dual_updater == "max_one" or self.dual_updater == "malencia":
-                self.lambda_duals = self.compute_obj_values()
-            self.lambda_duals = self.fun_dual_updater(
-                self.dual_updater, self.lambda_duals
-            )
+            if is_state_updated == True:
+                # obj_values /= self.T_0
+                obj_values = self.compute_obj_values()
+                obj_max = np.max(obj_values)
+                self.lambda_duals = np.maximum(
+                    self.lambda_duals
+                    + self.eta_dual * (obj_values - self.alphas) / obj_max,
+                    0,
+                )
+                if self.dual_updater == "max_one" or self.dual_updater == "malencia":
+                    self.lambda_duals = self.compute_obj_values()
+                self.lambda_duals = self.fun_dual_updater(
+                    self.dual_updater, self.lambda_duals
+                )
+                self.update_idf(self.lambda_duals, normalize=self.normalize)
+            else:
+                obj_values = self.all_obj_values[:, k]
+                self.lambda_duals = self.all_lambda_duals[:, k]
             
             self.all_obj_values[:, k + 1] = obj_values
             self.all_lambda_duals[:, k + 1] = self.lambda_duals
@@ -296,7 +302,7 @@ def fun_dual_updater(self, configs, lambdas):
     envs = list(range(100))
     # T_0s = [25, 50, 75, 100]
     # envs = [72]
-    T_0s = [25]
+    T_0s = [1]
     eta_duals = [1]
     eval_dir = sys.argv[2]
 
@@ -309,7 +315,7 @@ def fun_dual_updater(self, configs, lambdas):
                     env_id,
                     eta_dual,
                     T_0,
-                    dual_updater="proj_1",
+                    dual_updater="malencia",
                     alpha=0.0,
                     normalize=True,
                     obj_normalize_factor=1e10,
