Python bindings for methods of GcsTrajectoryOptimization::Subgraph fo…

…r adding generic costs and constraints. (RobotLocomotion#22156)

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_<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_<Class::Subgraph>(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<symbolic::Variable> {
+              return self.vertex_control_points();
+            },
+            subgraph_doc.vertex_control_points.doc)
+        .def(
+            "edge_constituent_vertex_control_points",
+            [](const GcsTrajectoryOptimization::Subgraph& self)
+                -> std::pair<MatrixX<symbolic::Variable>,
+                    MatrixX<symbolic::Variable>> {
+              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 =

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())