[codex] Fix constraint violation logging

robot_payload_id/optimization/nevergrad_augmented_lagrangian.py

@@ -5,6 +5,7 @@
 
 import nevergrad as ng
 import numpy as np
+import wandb
 
 from pydrake.all import (
     AugmentedLagrangianNonsmooth,
@@ -14,8 +15,6 @@
 from tqdm import tqdm
 from wandb.sdk.wandb_summary import SummarySubDict
 
-import wandb
-
 # Create a cache for multiprocessing
 _MAX_NUM_WORKERS = 64
 _CACHED_AUGMENTED_LAGRANGIANS: List[Union[AugmentedLagrangianNonsmooth, None]] = [
@@ -250,36 +249,35 @@ def _compute_and_log_constraint_violations(
         constraint_types = [name.split("_")[0] for name in constraint_names]
 
         # Assumes that constraint name has the form "name_..."
-        constraint_type_violations_map = dict(
-            zip(list(set(constraint_types)), [0] * len(constraint_types))
-        )
+        constraint_type_violations_map = {
+            constraint_type: 0 for constraint_type in constraint_types
+        }
 
         # See Drake augmented_lagrangian.cc::EvalAugmentedLagrangian for context
         lag_idx = 0
         is_equality_mask = nonsmooth_al.is_equality()
-        for constraint, constraint_type, is_equality in zip(
-            constraints, constraint_types, is_equality_mask
-        ):
+        for constraint, constraint_type in zip(constraints, constraint_types):
             if isinstance(constraint, BoundingBoxConstraint):
                 # No residuals exist for these bounds
                 continue
 
-            constraint_tol = (
-                self._equality_constraint_tol
-                if is_equality
-                else self._inequality_constraint_tol
-            )
             for i in range(constraint.num_constraints()):
                 lb = constraint.lower_bound()[i]
                 ub = constraint.upper_bound()[i]
                 if lb == ub:
                     # Constraint adds one Lagrange multiplier
+                    constraint_tol = self._equality_constraint_tol
+                    assert is_equality_mask[lag_idx], "Expected equality constraint."
                     if constraint_residue[lag_idx] ** 2 > constraint_tol:
                         constraint_type_violations_map[constraint_type] += 1
                     lag_idx += 1
                 else:
                     # Constraint adds 0 to 2 Lagrange multipliers
+                    constraint_tol = self._inequality_constraint_tol
                     if not np.isinf(lb):
+                        assert not is_equality_mask[
+                            lag_idx
+                        ], "Expected inequality constraint."
                         if (
                             np.maximum(-constraint_residue[lag_idx], 0) ** 2
                             > constraint_tol
@@ -288,6 +286,9 @@ def _compute_and_log_constraint_violations(
                         lag_idx += 1
 
                     if not np.isinf(ub):
+                        assert not is_equality_mask[
+                            lag_idx
+                        ], "Expected inequality constraint."
                         if (
                             np.maximum(-constraint_residue[lag_idx], 0) ** 2
                             > constraint_tol
@@ -299,6 +300,10 @@ def _compute_and_log_constraint_violations(
             logging.warning(
                 "Skipping to log constraint violations for decision variable bounds."
             )
+        else:
+            assert lag_idx == len(
+                is_equality_mask
+            ), "Constraint metadata did not match augmented Lagrangian size."
 
         num_constraint_violations = sum(list(constraint_type_violations_map.values()))
         wandb.log(

robot_payload_id/optimization/nevergrad_util.py

@@ -2,7 +2,6 @@
 from typing import List, Union
 
 import nevergrad as ng
-
 import wandb
 
 

robot_payload_id/optimization/optimal_experiment_design_b_spline.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
+import wandb
 import yaml
 
 from pydrake.all import (
@@ -20,8 +21,6 @@
     MultibodyPlant,
 )
 
-import wandb
-
 from robot_payload_id.data import (
     compute_base_param_mapping,
     extract_numeric_data_matrix_autodiff,

robot_payload_id/optimization/optimal_experiment_design_base.py

@@ -5,12 +5,11 @@
 from typing import Any, Optional, Tuple
 
 import numpy as np
+import wandb
 
 from numpy import ndarray
 from pydrake.all import MakeVectorVariable, ModelInstanceIndex, MultibodyPlant
 
-import wandb
-
 from robot_payload_id.data import (
     extract_symbolic_data_matrix,
     load_symbolic_data_matrix,

robot_payload_id/optimization/optimal_experiment_design_fourier.py

@@ -10,6 +10,7 @@
 
 import nevergrad as ng
 import numpy as np
+import wandb
 import yaml
 
 from pydrake.all import (
@@ -26,8 +27,6 @@
     SnoptSolver,
 )
 
-import wandb
-
 from robot_payload_id.data import (
     compute_autodiff_joint_data_from_fourier_series_traj_params1,
     compute_base_param_mapping,

robot_payload_id/utils/dataclasses.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import pydrake.symbolic as sym
+import wandb
 
 from pydrake.all import (
     BsplineBasis,
@@ -20,8 +21,6 @@
     VectorLogSink,
 )
 
-import wandb
-
 from .inertia import change_inertia_reference_points_with_parallel_axis_theorem