From d56c5f20a892d20a7ea3bcd02947acbdd30e905a Mon Sep 17 00:00:00 2001 From: nepfaff Date: Wed, 13 May 2026 11:53:14 -0700 Subject: [PATCH] Fix constraint violation logging --- .../nevergrad_augmented_lagrangian.py | 31 +++++++++++-------- .../optimization/nevergrad_util.py | 1 - .../optimal_experiment_design_b_spline.py | 3 +- .../optimal_experiment_design_base.py | 3 +- .../optimal_experiment_design_fourier.py | 3 +- robot_payload_id/utils/dataclasses.py | 3 +- 6 files changed, 22 insertions(+), 22 deletions(-) diff --git a/robot_payload_id/optimization/nevergrad_augmented_lagrangian.py b/robot_payload_id/optimization/nevergrad_augmented_lagrangian.py index c3b9207..3e521c0 100644 --- a/robot_payload_id/optimization/nevergrad_augmented_lagrangian.py +++ b/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( diff --git a/robot_payload_id/optimization/nevergrad_util.py b/robot_payload_id/optimization/nevergrad_util.py index 7d5c86e..573d61d 100644 --- a/robot_payload_id/optimization/nevergrad_util.py +++ b/robot_payload_id/optimization/nevergrad_util.py @@ -2,7 +2,6 @@ from typing import List, Union import nevergrad as ng - import wandb diff --git a/robot_payload_id/optimization/optimal_experiment_design_b_spline.py b/robot_payload_id/optimization/optimal_experiment_design_b_spline.py index 5a6b8b1..8ce0053 100644 --- a/robot_payload_id/optimization/optimal_experiment_design_b_spline.py +++ b/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, diff --git a/robot_payload_id/optimization/optimal_experiment_design_base.py b/robot_payload_id/optimization/optimal_experiment_design_base.py index 67c3f8c..5290e7c 100644 --- a/robot_payload_id/optimization/optimal_experiment_design_base.py +++ b/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, diff --git a/robot_payload_id/optimization/optimal_experiment_design_fourier.py b/robot_payload_id/optimization/optimal_experiment_design_fourier.py index 4f56e4b..fd7b1be 100644 --- a/robot_payload_id/optimization/optimal_experiment_design_fourier.py +++ b/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, diff --git a/robot_payload_id/utils/dataclasses.py b/robot_payload_id/utils/dataclasses.py index 3f6e108..9fa44e7 100644 --- a/robot_payload_id/utils/dataclasses.py +++ b/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