Add Expression Costs and Formula Constraints to Subgraphs in GcsTrajectoryOptimization

planning/trajectory_optimization/gcs_trajectory_optimization.cc

@@ -29,9 +29,6 @@ namespace trajectory_optimization {
 using Subgraph = GcsTrajectoryOptimization::Subgraph;
 using EdgesBetweenSubgraphs = GcsTrajectoryOptimization::EdgesBetweenSubgraphs;
 
-using drake::solvers::MathematicalProgram;
-using drake::solvers::Solve;
-using drake::solvers::VectorXDecisionVariable;
 using Eigen::MatrixXd;
 using Eigen::SparseMatrix;
 using Eigen::VectorXd;
@@ -65,11 +62,17 @@ using solvers::L2NormCost;
 using solvers::LinearConstraint;
 using solvers::LinearCost;
 using solvers::LinearEqualityConstraint;
+using solvers::MathematicalProgram;
+using solvers::MatrixXDecisionVariable;
 using solvers::QuadraticCost;
+using solvers::Solve;
+using solvers::VectorXDecisionVariable;
 using symbolic::DecomposeLinearExpressions;
 using symbolic::Expression;
+using symbolic::Formula;
 using symbolic::MakeMatrixContinuousVariable;
 using symbolic::MakeVectorContinuousVariable;
+using symbolic::Variable;
 using trajectories::BezierCurve;
 using trajectories::CompositeTrajectory;
 using trajectories::PiecewiseTrajectory;
@@ -316,6 +319,16 @@ Subgraph::Subgraph(
         path_continuity_constraint,
         {GetControlPoints(*u).col(order), GetControlPoints(*v).col(0)}));
   }
+
+  // Construct placeholder variables.
+  placeholder_vertex_duration_var_ = symbolic::Variable("t");
+  placeholder_vertex_control_points_var_ =
+      MakeMatrixContinuousVariable(num_positions(), order + 1, "x");
+  placeholder_edge_durations_var_ =
+      std::make_pair(symbolic::Variable("tu"), symbolic::Variable("tv"));
+  placeholder_edge_control_points_var_ = std::make_pair(
+      MakeMatrixContinuousVariable(num_positions(), order + 1, "xu"),
+      MakeMatrixContinuousVariable(num_positions(), order + 1, "xv"));
 }
 
 Subgraph::~Subgraph() = default;
@@ -638,6 +651,9 @@ void Subgraph::AddPathContinuityConstraints(int continuity_order) {
 }
 
 void Subgraph::AddContinuityConstraints(int continuity_order) {
+  // TODO(cohnt): Rewrite to use the generic AddCost and AddConstraint
+  // interfaces.
+
   if (continuity_order == 0) {
     throw std::runtime_error(
         "Path continuity is enforced by default. Choose a higher order.");
@@ -739,6 +755,137 @@ symbolic::Variable Subgraph::GetTimeScaling(
   return v.x()(v.x().size() - 1);
 }
 
+namespace {
+
+std::vector<Variable> FlattenVariables(
+    const std::vector<Variable>& scalar_variables,
+    const std::vector<VectorXDecisionVariable>& vector_variables,
+    const std::vector<MatrixXDecisionVariable>& matrix_variables) {
+  std::vector<Variable> all_variables;
+
+  // Append scalar variables.
+  all_variables.insert(all_variables.end(), scalar_variables.begin(),
+                       scalar_variables.end());
+
+  // Flatten vector variables.
+  for (const auto& vector_variable : vector_variables) {
+    for (int i = 0; i < vector_variable.size(); ++i) {
+      all_variables.push_back(vector_variable(i));
+    }
+  }
+
+  // Flatten matrix variables.
+  for (const auto& matrix_variable : matrix_variables) {
+    for (int i = 0; i < matrix_variable.rows(); ++i) {
+      for (int j = 0; j < matrix_variable.cols(); ++j) {
+        all_variables.push_back(matrix_variable(i, j));
+      }
+    }
+  }
+
+  return all_variables;
+}
+
+template <typename T>
+T SubstituteAllVariables(
+    T e, std::vector<Variable> old_scalar_variables = {},
+    std::vector<Variable> new_scalar_variables = {},
+    std::vector<VectorXDecisionVariable> old_vector_variables = {},
+    std::vector<VectorXDecisionVariable> new_vector_variables = {},
+    std::vector<MatrixXDecisionVariable> old_matrix_variables = {},
+    std::vector<MatrixXDecisionVariable> new_matrix_variables = {}) {
+  DRAKE_DEMAND(old_scalar_variables.size() == new_scalar_variables.size());
+  DRAKE_DEMAND(old_vector_variables.size() == new_vector_variables.size());
+  DRAKE_DEMAND(old_matrix_variables.size() == new_matrix_variables.size());
+
+  std::vector<Variable> old_variables = FlattenVariables(
+      old_scalar_variables, old_vector_variables, old_matrix_variables);
+  std::vector<Variable> new_variables = FlattenVariables(
+      new_scalar_variables, new_vector_variables, new_matrix_variables);
+  DRAKE_DEMAND(old_variables.size() == new_variables.size());
+
+  // Note: the logic is identical for Expression and Formula, except for one
+  // differing method name -- Expression::GetVariables versus
+  // Formula::GetFreeVariables.
+  symbolic::Variables e_vars;
+  if constexpr (std::is_same_v<T, Expression>) {
+    e_vars = e.GetVariables();
+  } else {  // std::is_same_v<T, Formula>
+    e_vars = e.GetFreeVariables();
+  }
+
+  for (int i = 0; i < ssize(old_variables); ++i) {
+    if (e_vars.include(old_variables[i])) {
+      e = e.Substitute(old_variables[i], new_variables[i]);
+    }
+  }
+
+  return e;
+}
+
+}  // namespace
+
+template <typename T>
+T Subgraph::SubstituteVertexPlaceholderVariables(T e,
+                                                 const Vertex& vertex) const {
+  return SubstituteAllVariables(
+      e, {placeholder_vertex_duration_var_}, {GetTimeScaling(vertex)}, {}, {},
+      {placeholder_vertex_control_points_var_}, {GetControlPoints(vertex)});
+}
+
+template <typename T>
+T Subgraph::SubstituteEdgePlaceholderVariables(T e, const Edge& edge) const {
+  const Vertex& v1 = edge.u();
+  const Vertex& v2 = edge.v();
+
+  return SubstituteAllVariables(e,
+                                {placeholder_edge_durations_var_.first,
+                                 placeholder_edge_durations_var_.second},
+                                {GetTimeScaling(v1), GetTimeScaling(v2)}, {},
+                                {},
+                                {placeholder_edge_control_points_var_.first,
+                                 placeholder_edge_control_points_var_.second},
+                                {GetControlPoints(v1), GetControlPoints(v2)});
+}
+
+void Subgraph::AddVertexCost(
+    const Expression& e,
+    const std::unordered_set<Transcription>& used_in_transcription) {
+  for (Vertex*& vertex : vertices_) {
+    Expression post_substitution =
+        SubstituteVertexPlaceholderVariables(e, *vertex);
+    vertex->AddCost(post_substitution, used_in_transcription);
+  }
+}
+
+void Subgraph::AddVertexConstraint(
+    const Formula& e,
+    const std::unordered_set<Transcription>& used_in_transcription) {
+  for (Vertex*& vertex : vertices_) {
+    Formula post_substitution =
+        SubstituteVertexPlaceholderVariables(e, *vertex);
+    vertex->AddConstraint(post_substitution, used_in_transcription);
+  }
+}
+
+void Subgraph::AddEdgeCost(
+    const Expression& e,
+    const std::unordered_set<Transcription>& used_in_transcription) {
+  for (Edge*& edge : edges_) {
+    Expression post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
+    edge->AddCost(post_substitution, used_in_transcription);
+  }
+}
+
+void Subgraph::AddEdgeConstraint(
+    const symbolic::Formula& e,
+    const std::unordered_set<Transcription>& used_in_transcription) {
+  for (Edge*& edge : edges_) {
+    Formula post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
+    edge->AddConstraint(post_substitution, used_in_transcription);
+  }
+}
+
 EdgesBetweenSubgraphs::EdgesBetweenSubgraphs(
     const Subgraph& from_subgraph, const Subgraph& to_subgraph,
     const ConvexSet* subspace, GcsTrajectoryOptimization* traj_opt,

planning/trajectory_optimization/gcs_trajectory_optimization.h

@@ -5,7 +5,9 @@
 #include <optional>
 #include <string>
 #include <tuple>
+#include <unordered_set>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "drake/common/trajectories/bezier_curve.h"
@@ -273,6 +275,155 @@ class GcsTrajectoryOptimization final {
     */
     void AddContinuityConstraints(int continuity_order);
 
+    /** Returns a placeholder decision variable (not actually declared as a
+    decision variable in the MathematicalProgram) associated with the time
+    scaling of the trajectory in a set within this subgraph. This variable will
+    be substituted for real decision variables in methods like AddVertexCost and
+    AddVertexConstraint. Passing this variable directly into
+    objectives/constraints will result in an error. */
+    const symbolic::Variable& vertex_duration() const {
+      return placeholder_vertex_duration_var_;
+    }
+
+    /** Returns a placeholder decision variable (not actually declared as a
+    decision variable in the MathematicalProgram) associated with the control
+    points of the trajectory in a set within this subgraph. The variable will be
+    of shape (num_positions(), order+1), where the ith column is the ith control
+    point. This variable will be substituted for real decision variables in
+    methods like AddVertexCost and AddVertexConstraint. Passing this variable
+    directly into objectives/constraints will result in an error. */
+    const solvers::MatrixXDecisionVariable& vertex_control_points() const {
+      return placeholder_vertex_control_points_var_;
+    }
+
+    /** Returns a pair of placeholder decision variables (not actually declared
+    as decision variables in the MathematicalProgram) associated with the time
+    scaling of the trajectory in two sets within this subgraph that are
+    connected by an internal edge. This variable will be substituted for real
+    decision variables in methods like AddEdgeCost and AddEdgeConstraint.
+    Passing this variable directly into objectives/constraints will result in an
+    error. */
+    const std::pair<symbolic::Variable, symbolic::Variable>&
+    edge_constituent_vertex_durations() const {
+      return placeholder_edge_durations_var_;
+    }
+
+    /** Returns a pair of placeholder decision variables (not actually declared
+    as decision variables in the MathematicalProgram) associated with the
+    control points of the trajectory in two sets within this subgraph that are
+    connected by an internal edge. Each variable will be of shape
+    (num_positions(), order+1), where the ith column is the ith control point.
+    This variable will be substituted for real decision variables in methods
+    like AddEdgeCost and AddEdgeConstraint. Passing this variable directly into
+    objectives/constraints will result in an error. */
+    const std::pair<solvers::MatrixXDecisionVariable,
+                    solvers::MatrixXDecisionVariable>&
+    edge_constituent_vertex_control_points() const {
+      return placeholder_edge_control_points_var_;
+    }
+
+    /** Adds an arbitrary user-defined cost to every vertex in the subgraph. The
+    cost should be defined using the placeholder control point variables
+    (obtained from vertex_control_points()) and the placeholder time scaling
+    variable (obtained from vertex_time()). This enables greater modeling
+    freedom, but we cannot guarantee a feasible solution for all possible costs.
+
+    @throws std::exception if any variables besides those from @ref
+    vertex_duration and @ref vertex_control_points are used.
+
+    Costs which do not support the perspective operation cannot be used with
+    Transcription::kMIP or Transcription::kRelaxation. Consider providing an
+    appropriate "convex surrogate" that is supported within GraphOfConvexSets,
+    or exclusively using the SolveConvexRestriction method. */
+    void AddVertexCost(
+        const symbolic::Expression& e,
+        const std::unordered_set<
+            geometry::optimization::GraphOfConvexSets::Transcription>&
+            used_in_transcription = {
+                geometry::optimization::GraphOfConvexSets::Transcription::kMIP,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRelaxation,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRestriction});
+
+    /** Adds an arbitrary user-defined constraint to every vertex in the
+    subgraph. The constraint should be defined using the placeholder control
+    point variables (obtained from vertex_control_points()) and the placeholder
+    time scaling variable (obtained from vertex_time()). This enables greater
+    modeling freedom, but we cannot guarantee a feasible solution for all
+    possible costs.
+
+    @throws std::exception if any variables besides those from @ref
+    vertex_duration and @ref vertex_control_points are used.
+
+    Constraints which do not support the perspective operation cannot be used
+    with Transcription::kMIP or Transcription::kRelaxation. Consider providing
+    an appropriate "convex surrogate" that is supported within
+    GraphOfConvexSets, or exclusively using the SolveConvexRestriction method.
+    */
+    void AddVertexConstraint(
+        const symbolic::Formula& e,
+        const std::unordered_set<
+            geometry::optimization::GraphOfConvexSets::Transcription>&
+            used_in_transcription = {
+                geometry::optimization::GraphOfConvexSets::Transcription::kMIP,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRelaxation,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRestriction});
+
+    /** Adds an arbitrary user-defined cost to every internal edge within the
+    subgraph. The cost should be defined using the placeholder control point
+    variables (obtained from edge_control_points()) and the placeholder time
+    scaling variables (obtained from edge_time()). This enables greater modeling
+    freedom, but we cannot guarantee a feasible solution for all possible costs.
+
+    @throws std::exception if any variables besides those from @ref
+    edge_constituent_vertex_durations and @ref
+    edge_constituent_vertex_control_points are used.
+
+    Costs which do not support the perspective operation cannot be used with
+    Transcription::kMIP or Transcription::kRelaxation. Consider providing an
+    appropriate "convex surrogate" that is supported within GraphOfConvexSets,
+    or exclusively using the SolveConvexRestriction method. */
+    void AddEdgeCost(
+        const symbolic::Expression& e,
+        const std::unordered_set<
+            geometry::optimization::GraphOfConvexSets::Transcription>&
+            used_in_transcription = {
+                geometry::optimization::GraphOfConvexSets::Transcription::kMIP,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRelaxation,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRestriction});
+
+    /** Adds an arbitrary user-defined constraint to every internal edge within
+    the subgraph. The constraint should be defined using the placeholder control
+    point variables (obtained from edge_control_points()) and the placeholder
+    time scaling variables (obtained from edge_time()). This enables greater
+    modeling freedom, but we cannot guarantee a feasible solution for all
+    possible costs.
+
+    @throws std::exception if any variables besides those from @ref
+    edge_constituent_vertex_durations and @ref
+    edge_constituent_vertex_control_points are used.
+
+    Constraints which do not support the perspective operation cannot be used
+    with Transcription::kMIP or Transcription::kRelaxation. Consider providing
+    an appropriate "convex surrogate" that is supported within
+    GraphOfConvexSets, or exclusively using the SolveConvexRestriction method.
+    */
+    void AddEdgeConstraint(
+        const symbolic::Formula& e,
+        const std::unordered_set<
+            geometry::optimization::GraphOfConvexSets::Transcription>&
+            used_in_transcription = {
+                geometry::optimization::GraphOfConvexSets::Transcription::kMIP,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRelaxation,
+                geometry::optimization::GraphOfConvexSets::Transcription::
+                    kRestriction});
+
    private:
     /* Constructs a new subgraph and copies the regions. */
     Subgraph(const geometry::optimization::ConvexSets& regions,
@@ -301,6 +452,19 @@ class GcsTrajectoryOptimization final {
     symbolic::Variable GetTimeScaling(
         const geometry::optimization::GraphOfConvexSets::Vertex& v) const;
 
+    /* Substitute any placeholder variables with the versions corresponding to
+    a specific vertex. The return type will match the argument type. */
+    template <typename T>
+    T SubstituteVertexPlaceholderVariables(
+        T e,
+        const geometry::optimization::GraphOfConvexSets::Vertex& vertex) const;
+
+    /* Substitute any placeholder variables with the versions corresponding to
+    a specific internal edge. The return type will match the argument type. */
+    template <typename T>
+    T SubstituteEdgePlaceholderVariables(
+        T e, const geometry::optimization::GraphOfConvexSets::Edge& edge) const;
+
     const geometry::optimization::ConvexSets regions_;
     const int order_;
     const double h_min_;
@@ -314,6 +478,15 @@ class GcsTrajectoryOptimization final {
     // design costs and constraints for the underlying vertices and edges.
     trajectories::BezierCurve<double> r_trajectory_;
 
+    symbolic::Variable placeholder_vertex_duration_var_;
+    solvers::MatrixXDecisionVariable placeholder_vertex_control_points_var_;
+
+    std::pair<symbolic::Variable, symbolic::Variable>
+        placeholder_edge_durations_var_;
+    std::pair<solvers::MatrixXDecisionVariable,
+              solvers::MatrixXDecisionVariable>
+        placeholder_edge_control_points_var_;
+
     friend class GcsTrajectoryOptimization;
   };