Skip to content

Commit

Permalink
Add Generic Binding<Cost> and Binding<Constraint> to vertices and edg…
Browse files Browse the repository at this point in the history 
…es within a GcsTrajectoryOptimization::Subgraph (#22179)

This PR includes Python bindings.
  • Loading branch information
@cohnt
cohnt authored Nov 19, 2024
1 parent d44189c commit cbaecba
Show file tree
Hide file tree
Showing 5 changed files with 366 additions and 45 deletions.
66 changes: 58 additions & 8 deletions bindings/pydrake/planning/planning_py_trajectory_optimization.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -442,18 +442,68 @@ void DefinePlanningTrajectoryOptimization(py::module m) {
return self.edge_constituent_vertex_control_points();
},
subgraph_doc.edge_constituent_vertex_control_points.doc)
.def("AddVertexCost", &Class::Subgraph::AddVertexCost, py::arg("e"),
.def("AddVertexCost",
py::overload_cast<const symbolic::Expression&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddVertexCost),
py::arg("e"), py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddVertexCost.doc_2args_e_use_in_transcription)
.def("AddVertexCost",
py::overload_cast<const solvers::Binding<solvers::Cost>&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddVertexCost),
py::arg("binding"),
py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddVertexCost.doc_2args_binding_use_in_transcription)
.def("AddVertexConstraint",
py::overload_cast<const symbolic::Formula&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddVertexConstraint),
py::arg("e"), py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddVertexConstraint.doc_2args_e_use_in_transcription)
.def("AddVertexConstraint",
py::overload_cast<const solvers::Binding<solvers::Constraint>&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddVertexConstraint),
py::arg("binding"),
py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddVertexCost.doc)
.def("AddVertexConstraint", &Class::Subgraph::AddVertexConstraint,
subgraph_doc.AddVertexConstraint
.doc_2args_binding_use_in_transcription)
.def("AddEdgeCost",
py::overload_cast<const symbolic::Expression&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddEdgeCost),
py::arg("e"), py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddVertexConstraint.doc)
.def("AddEdgeCost", &Class::Subgraph::AddEdgeCost, py::arg("e"),
subgraph_doc.AddEdgeCost.doc_2args_e_use_in_transcription)
.def("AddEdgeCost",
py::overload_cast<const solvers::Binding<solvers::Cost>&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddEdgeCost),
py::arg("binding"),
py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddEdgeCost.doc)
.def("AddEdgeConstraint", &Class::Subgraph::AddEdgeConstraint,
subgraph_doc.AddEdgeCost.doc_2args_binding_use_in_transcription)
.def("AddEdgeConstraint",
py::overload_cast<const symbolic::Formula&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddEdgeConstraint),
py::arg("e"), py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddEdgeConstraint.doc);
subgraph_doc.AddEdgeConstraint.doc_2args_e_use_in_transcription)
.def("AddEdgeConstraint",
py::overload_cast<const solvers::Binding<solvers::Constraint>&,
const std::unordered_set<
geometry::optimization::GraphOfConvexSets::Transcription>&>(
&Class::Subgraph::AddEdgeConstraint),
py::arg("binding"),
py::arg("use_in_transcription") = all_transcriptions,
subgraph_doc.AddEdgeConstraint
.doc_2args_binding_use_in_transcription);

// EdgesBetweenSubgraphs
const auto& subgraph_edges_doc =
Expand Down
30 changes: 30 additions & 0 deletions bindings/pydrake/planning/test/trajectory_optimization_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,13 @@
Point,
VPolytope,
)
from pydrake.solvers import (
Binding,
Cost,
Constraint,
ExpressionCost,
ExpressionConstraint
)
from pydrake.multibody.plant import MultibodyPlant
import pydrake.solvers as mp
from pydrake.symbolic import Variable
Expand Down Expand Up @@ -377,6 +384,29 @@ def test_gcs_trajectory_optimization_basic(self):
regions.AddEdgeConstraint(e=edge_constraint,
use_in_transcription=all_transcriptions)

to_bind_vertex_cost = ExpressionCost(vertex_cost)
to_bind_vertex_constraint = ExpressionConstraint([vertex_cost],
[-1],
[1])
to_bind_edge_cost = ExpressionCost(edge_cost)
to_bind_edge_constraint = ExpressionConstraint([edge_cost], [-1], [1])

restriction_only = {
GraphOfConvexSets.Transcription.kRestriction
}
regions.AddVertexCost(binding=Binding[Cost](
to_bind_vertex_cost, to_bind_vertex_cost.vars()),
use_in_transcription=restriction_only)
regions.AddVertexConstraint(binding=Binding[Constraint](
to_bind_vertex_constraint, to_bind_vertex_constraint.vars()),
use_in_transcription=restriction_only)
regions.AddEdgeCost(binding=Binding[Cost](
to_bind_edge_cost, to_bind_edge_cost.vars()),
use_in_transcription=restriction_only)
regions.AddEdgeConstraint(binding=Binding[Constraint](
to_bind_edge_constraint, to_bind_edge_constraint.vars()),
use_in_transcription=restriction_only)

# We can add additional costs and constraints to EdgesBetweenSubgraphs
# with the placeholder variables.
edge_constituent_vertex_durations = (
Expand Down
140 changes: 119 additions & 21 deletions planning/trajectory_optimization/gcs_trajectory_optimization.cc
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include "drake/planning/trajectory_optimization/gcs_trajectory_optimization.h"

#include <algorithm>
#include <limits>
#include <tuple>
#include <unordered_map>
Expand Down Expand Up @@ -786,6 +787,7 @@ std::vector<Variable> FlattenVariables(
return all_variables;
}

// Compatible with Expression and Formula.
template <typename T>
T SubstituteAllVariables(
T e, std::vector<Variable> old_scalar_variables = {},
Expand Down Expand Up @@ -813,10 +815,57 @@ T SubstituteAllVariables(
return e;
}

// Compatible with Binding<Cost> and Binding<Constraint>.
template <typename T>
Binding<T> SubstituteAllVariables(
const Binding<T>& binding, 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());

VectorXDecisionVariable new_binding_variables(binding.variables());
for (int i = 0; i < new_binding_variables.size(); ++i) {
// For the current placeholder variable, find its index in old_variables,
// and grab the corresponding entry in new_variables.
auto iterator = std::find_if(
old_variables.begin(), old_variables.end(), [&](const Variable& var) {
return var.equal_to(new_binding_variables[i]);
});
if (iterator == old_variables.end()) {
// We throw an error if the user gave an unknown variable.
throw std::runtime_error(
fmt::format("Unknown variable with name {} provided.",
new_binding_variables[i].get_name()));
}
size_t index = std::distance(std::begin(old_variables), iterator);
new_binding_variables[i] = new_variables[index];
}

return Binding<T>{binding.evaluator(), new_binding_variables};
}

// Compatible with Expression, Formula, Binding<Cost>, and Binding<Constraint>.
template <typename T>
void ThrowIfContainsVariables(const T& e, const std::vector<Variable>& vars,
const std::string& error_message) {
symbolic::Variables e_vars = e.GetFreeVariables();
symbolic::Variables e_vars;
if constexpr (std::disjunction_v<std::is_same<T, Formula>,
std::is_same<T, Expression>>) {
e_vars = e.GetFreeVariables();
} else {
e_vars = symbolic::Variables(e.variables());
}
for (const auto& var : vars) {
if (e_vars.include(var)) {
throw(std::runtime_error(error_message));
Expand All @@ -826,6 +875,7 @@ void ThrowIfContainsVariables(const T& e, const std::vector<Variable>& vars,

} // namespace

// Compatible with Expression, Formula, Binding<Cost>, and Binding<Constraint>.
template <typename T>
T Subgraph::SubstituteVertexPlaceholderVariables(T e,
const Vertex& vertex) const {
Expand All @@ -847,6 +897,7 @@ T Subgraph::SubstituteVertexPlaceholderVariables(T e,
{placeholder_vertex_control_points_var_}, {GetControlPoints(vertex)});
}

// Compatible with Expression, Formula, Binding<Cost>, and Binding<Constraint>.
template <typename T>
T Subgraph::SubstituteEdgePlaceholderVariables(T e, const Edge& edge) const {
// Check that a user hasn't used the vertex placeholder variables.
Expand Down Expand Up @@ -874,39 +925,86 @@ T Subgraph::SubstituteEdgePlaceholderVariables(T e, const Edge& edge) const {

void Subgraph::AddVertexCost(
const Expression& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddVertexCost(e, use_in_transcription);
}
void Subgraph::AddVertexCost(
const Binding<Cost>& binding,
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddVertexCost(binding, use_in_transcription);
}

// Compatible with Expression and Binding<Cost>.
template <typename T>
void Subgraph::DoAddVertexCost(
const T& e, const std::unordered_set<Transcription>& use_in_transcription) {
for (Vertex*& vertex : vertices_) {
Expression post_substitution =
SubstituteVertexPlaceholderVariables(e, *vertex);
vertex->AddCost(post_substitution, used_in_transcription);
T post_substitution = SubstituteVertexPlaceholderVariables(e, *vertex);
vertex->AddCost(post_substitution, use_in_transcription);
}
}

void Subgraph::AddVertexConstraint(
const Formula& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddVertexConstraint(e, use_in_transcription);
}
void Subgraph::AddVertexConstraint(
const Binding<Constraint>& binding,
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddVertexConstraint(binding, use_in_transcription);
}

// Compatible with Formula and Binding<Constraint>.
template <typename T>
void Subgraph::DoAddVertexConstraint(
const T& e, const std::unordered_set<Transcription>& use_in_transcription) {
for (Vertex*& vertex : vertices_) {
Formula post_substitution =
SubstituteVertexPlaceholderVariables(e, *vertex);
vertex->AddConstraint(post_substitution, used_in_transcription);
T post_substitution = SubstituteVertexPlaceholderVariables(e, *vertex);
vertex->AddConstraint(post_substitution, use_in_transcription);
}
}

void Subgraph::AddEdgeCost(
const Expression& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddEdgeCost(e, use_in_transcription);
}

void Subgraph::AddEdgeCost(
const Binding<Cost>& binding,
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddEdgeCost(binding, use_in_transcription);
}

// Compatible with Expression and Binding<Cost>.
template <typename T>
void Subgraph::DoAddEdgeCost(
const T& e, const std::unordered_set<Transcription>& use_in_transcription) {
for (Edge*& edge : edges_) {
Expression post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddCost(post_substitution, used_in_transcription);
T post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddCost(post_substitution, use_in_transcription);
}
}

void Subgraph::AddEdgeConstraint(
const symbolic::Formula& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const Formula& e,
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddEdgeConstraint(e, use_in_transcription);
}

void Subgraph::AddEdgeConstraint(
const Binding<Constraint>& binding,
const std::unordered_set<Transcription>& use_in_transcription) {
DoAddEdgeConstraint(binding, use_in_transcription);
}

// Compatible with Formula and Binding<Constraint>.
template <typename T>
void Subgraph::DoAddEdgeConstraint(
const T& e, const std::unordered_set<Transcription>& use_in_transcription) {
for (Edge*& edge : edges_) {
Formula post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddConstraint(post_substitution, used_in_transcription);
T post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddConstraint(post_substitution, use_in_transcription);
}
}

Expand Down Expand Up @@ -1600,19 +1698,19 @@ T EdgesBetweenSubgraphs::SubstituteEdgePlaceholderVariables(

void EdgesBetweenSubgraphs::AddEdgeCost(
const Expression& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const std::unordered_set<Transcription>& use_in_transcription) {
for (Edge*& edge : edges_) {
Expression post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddCost(post_substitution, used_in_transcription);
edge->AddCost(post_substitution, use_in_transcription);
}
}

void EdgesBetweenSubgraphs::AddEdgeConstraint(
const symbolic::Formula& e,
const std::unordered_set<Transcription>& used_in_transcription) {
const std::unordered_set<Transcription>& use_in_transcription) {
for (Edge*& edge : edges_) {
Formula post_substitution = SubstituteEdgePlaceholderVariables(e, *edge);
edge->AddConstraint(post_substitution, used_in_transcription);
edge->AddConstraint(post_substitution, use_in_transcription);
}
}

Expand Down
Loading

0 comments on commit cbaecba

Please sign in to comment.