From 544bd62aad0984fce957929ae3d8a733229a6899 Mon Sep 17 00:00:00 2001 From: cohnt Date: Mon, 11 Nov 2024 19:21:11 -0500 Subject: [PATCH] Python bindings for methods of GcsTrajectoryOptimization::Subgraph for adding generic costs and constraints. --- .../planning_py_trajectory_optimization.cc | 46 ++++++++++++++++- .../test/trajectory_optimization_test.py | 49 +++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/bindings/pydrake/planning/planning_py_trajectory_optimization.cc b/bindings/pydrake/planning/planning_py_trajectory_optimization.cc index c8b899d545db..f65363979535 100644 --- a/bindings/pydrake/planning/planning_py_trajectory_optimization.cc +++ b/bindings/pydrake/planning/planning_py_trajectory_optimization.cc @@ -346,6 +346,15 @@ void DefinePlanningTrajectoryOptimization(py::module m) { constexpr auto& cls_doc = doc.GcsTrajectoryOptimization; py::class_ gcs_traj_opt(m, "GcsTrajectoryOptimization", cls_doc.doc); + const std::unordered_set< + geometry::optimization::GraphOfConvexSets::Transcription> + all_transcriptions = { + geometry::optimization::GraphOfConvexSets::Transcription::kMIP, + geometry::optimization::GraphOfConvexSets::Transcription:: + kRelaxation, + geometry::optimization::GraphOfConvexSets::Transcription:: + kRestriction}; + // Subgraph const auto& subgraph_doc = doc.GcsTrajectoryOptimization.Subgraph; py::class_(gcs_traj_opt, "Subgraph", subgraph_doc.doc) @@ -409,7 +418,42 @@ void DefinePlanningTrajectoryOptimization(py::module m) { .def("AddContinuityConstraints", &Class::Subgraph::AddContinuityConstraints, py::arg("continuity_order"), - subgraph_doc.AddContinuityConstraints.doc); + subgraph_doc.AddContinuityConstraints.doc) + .def("vertex_duration", &Class::Subgraph::vertex_duration, + subgraph_doc.vertex_duration.doc) + .def("edge_constituent_vertex_durations", + &Class::Subgraph::edge_constituent_vertex_durations, + subgraph_doc.edge_constituent_vertex_durations.doc) + // As in trajectory_optimization_py.cc, we use a lambda to *copy* + // the decision variables; otherwise we get dtype=object arrays + // cannot be referenced. + .def( + "vertex_control_points", + [](const GcsTrajectoryOptimization::Subgraph& self) + -> MatrixX { + return self.vertex_control_points(); + }, + subgraph_doc.vertex_control_points.doc) + .def( + "edge_constituent_vertex_control_points", + [](const GcsTrajectoryOptimization::Subgraph& self) + -> std::pair, + MatrixX> { + return self.edge_constituent_vertex_control_points(); + }, + subgraph_doc.edge_constituent_vertex_control_points.doc) + .def("AddVertexCost", &Class::Subgraph::AddVertexCost, py::arg("e"), + py::arg("use_in_transcription") = all_transcriptions, + subgraph_doc.AddVertexCost.doc) + .def("AddVertexConstraint", &Class::Subgraph::AddVertexConstraint, + py::arg("e"), py::arg("use_in_transcription") = all_transcriptions, + subgraph_doc.AddVertexConstraint.doc) + .def("AddEdgeCost", &Class::Subgraph::AddEdgeCost, py::arg("e"), + py::arg("use_in_transcription") = all_transcriptions, + subgraph_doc.AddEdgeCost.doc) + .def("AddEdgeConstraint", &Class::Subgraph::AddEdgeConstraint, + py::arg("e"), py::arg("use_in_transcription") = all_transcriptions, + subgraph_doc.AddEdgeConstraint.doc); // EdgesBetweenSubgraphs const auto& subgraph_edges_doc = diff --git a/bindings/pydrake/planning/test/trajectory_optimization_test.py b/bindings/pydrake/planning/test/trajectory_optimization_test.py index dbd0fd28c01f..dfcc0891c231 100644 --- a/bindings/pydrake/planning/test/trajectory_optimization_test.py +++ b/bindings/pydrake/planning/test/trajectory_optimization_test.py @@ -331,6 +331,55 @@ def test_gcs_trajectory_optimization_basic(self): np.testing.assert_allclose(restricted_traj_start, start, atol=1e-6) np.testing.assert_allclose(restricted_traj_end, end, atol=1e-6) + # We can add additional costs and constraints with the placeholder + # variables. + vertex_duration = regions.vertex_duration() + vertex_control_points = regions.vertex_control_points() + edge_constituent_vertex_durations = ( + regions.edge_constituent_vertex_durations() + ) + edge_constituent_vertex_control_points = ( + regions.edge_constituent_vertex_control_points() + ) + + vertex_cost = sum([ + vertex_duration, + vertex_control_points[0, 0], + vertex_control_points[0, 1], + ]) + edge_cost = sum([ + edge_constituent_vertex_durations[0], + edge_constituent_vertex_durations[1], + edge_constituent_vertex_control_points[0][0, 0], + edge_constituent_vertex_control_points[0][0, 1], + edge_constituent_vertex_control_points[1][0, 0], + edge_constituent_vertex_control_points[1][0, 1] + ]) + vertex_constraint = vertex_cost >= 0 + edge_constraint = edge_cost >= 0 + + regions.AddVertexCost(e=vertex_cost) + regions.AddVertexConstraint(e=vertex_constraint) + regions.AddEdgeCost(e=edge_cost) + regions.AddEdgeConstraint(e=edge_constraint) + + all_transcriptions = { + GraphOfConvexSets.Transcription.kMIP, + GraphOfConvexSets.Transcription.kRelaxation, + GraphOfConvexSets.Transcription.kRestriction + } + regions.AddVertexCost(e=vertex_cost, + use_in_transcription=all_transcriptions) + regions.AddVertexConstraint(e=vertex_constraint, + use_in_transcription=all_transcriptions) + regions.AddEdgeCost(e=edge_cost, + use_in_transcription=all_transcriptions) + regions.AddEdgeConstraint(e=edge_constraint, + use_in_transcription=all_transcriptions) + + traj, result = gcs.SolvePath(source, target) + self.assertTrue(result.is_success()) + # Test that removing all subgraphs, removes all vertices and edges. self.assertEqual(len(gcs.graph_of_convex_sets().Vertices()), len(source.Vertices())