feat(plate-balancing) : Support plate-balancing example in dairlib

.bazelrc

@@ -51,4 +51,4 @@ build --action_env=LD_LIBRARY_PATH=
 # use python3 by default
 build --python_path=python3
 
-build --define=WITH_GUROBI=ON
+build --define=WITH_GUROBI=ON

BUILD.bazel

@@ -21,6 +21,7 @@ cc_library(
     name = "libc3",
     hdrs = [":c3_headers"],  # Changed from srcs to hdrs for headers
     deps = LIBC3_COMPONENTS + [
+        "//lcmtypes:lcmt_c3",
         "@drake//:drake_shared_library",
     ],
     include_prefix = "c3",

bindings/pyc3/c3_multibody_py.cc

@@ -26,6 +26,20 @@ PYBIND11_MODULE(multibody, m) {
              c3::multibody::ContactModel::kFrictionlessSpring)
       .export_values();
 
+  py::class_<c3::multibody::LCSContactDescription>(m, "LCSContactDescription")
+      .def(py::init<>())
+      .def_readwrite("witness_point_A",
+                     &c3::multibody::LCSContactDescription::witness_point_A)
+      .def_readwrite("witness_point_B",
+                     &c3::multibody::LCSContactDescription::witness_point_B)
+      .def_readwrite("force_basis",
+                     &c3::multibody::LCSContactDescription::force_basis)
+      .def_readwrite("is_slack",
+                     &c3::multibody::LCSContactDescription::is_slack)
+      .def_static("CreateSlackVariableDescription",
+                  &c3::multibody::LCSContactDescription::
+                      CreateSlackVariableDescription);
+
   py::class_<c3::multibody::LCSFactory>(m, "LCSFactory")
       .def(py::init<const drake::multibody::MultibodyPlant<double>&,
                     drake::systems::Context<double>&,
@@ -37,8 +51,8 @@ PYBIND11_MODULE(multibody, m) {
            py::arg("plant"), py::arg("context"), py::arg("plant_ad"),
            py::arg("context_ad"), py::arg("contact_geoms"), py::arg("options"))
       .def("GenerateLCS", &c3::multibody::LCSFactory::GenerateLCS)
-      .def("GetContactJacobianAndPoints",
-           &c3::multibody::LCSFactory::GetContactJacobianAndPoints)
+      .def("GetContactDescriptions",
+           &c3::multibody::LCSFactory::GetContactDescriptions)
       .def("UpdateStateAndInput",
            &c3::multibody::LCSFactory::UpdateStateAndInput, py::arg("state"),
            py::arg("input"))

bindings/pyc3/c3_systems_py.cc

@@ -9,6 +9,9 @@
 #include "systems/c3_controller_options.h"
 #include "systems/framework/c3_output.h"
 #include "systems/framework/timestamped_vector.h"
+#include "systems/lcmt_generators/c3_output_generator.h"
+#include "systems/lcmt_generators/c3_trajectory_generator.h"
+#include "systems/lcmt_generators/contact_force_generator.h"
 #include "systems/lcs_factory_system.h"
 #include "systems/lcs_simulator.h"
 
@@ -100,8 +103,8 @@ PYBIND11_MODULE(systems, m) {
            py::return_value_policy::reference)
       .def("get_output_port_lcs", &LCSFactorySystem::get_output_port_lcs,
            py::return_value_policy::reference)
-      .def("get_output_port_lcs_contact_jacobian",
-           &LCSFactorySystem::get_output_port_lcs_contact_jacobian,
+      .def("get_output_port_lcs_contact_descriptions",
+           &LCSFactorySystem::get_output_port_lcs_contact_descriptions,
            py::return_value_policy::reference);
 
   py::class_<C3Output::C3Solution>(m, "C3Solution")
@@ -178,6 +181,88 @@ PYBIND11_MODULE(systems, m) {
                      &C3ControllerOptions::state_prediction_joints);
 
   m.def("LoadC3ControllerOptions", &LoadC3ControllerOptions);
+
+  // C3OutputGenerator
+  py::class_<lcmt_generators::C3OutputGenerator,
+             drake::systems::LeafSystem<double>>(m, "C3OutputGenerator")
+      .def(py::init<>())
+      .def("get_input_port_c3_solution",
+           &lcmt_generators::C3OutputGenerator::get_input_port_c3_solution,
+           py::return_value_policy::reference)
+      .def("get_input_port_c3_intermediates",
+           &lcmt_generators::C3OutputGenerator::get_input_port_c3_intermediates,
+           py::return_value_policy::reference)
+      .def("get_output_port_c3_output",
+           &lcmt_generators::C3OutputGenerator::get_output_port_c3_output,
+           py::return_value_policy::reference)
+      .def_static("AddLcmPublisherToBuilder",
+                  &lcmt_generators::C3OutputGenerator::AddLcmPublisherToBuilder,
+                  py::arg("builder"), py::arg("solution_port"),
+                  py::arg("intermediates_port"), py::arg("channel"),
+                  py::arg("lcm"), py::arg("publish_triggers"),
+                  py::arg("publish_period"), py::arg("publish_offset"));
+
+  // C3TrajectoryGenerator
+  py::class_<lcmt_generators::TrajectoryDescription::index_range>(
+      m, "TrajectoryDescriptionIndexRange")
+      .def(py::init<>())
+      .def_readwrite(
+          "start", &lcmt_generators::TrajectoryDescription::index_range::start)
+      .def_readwrite("end",
+                     &lcmt_generators::TrajectoryDescription::index_range::end);
+  py::class_<lcmt_generators::TrajectoryDescription>(m, "TrajectoryDescription")
+      .def(py::init<>())
+      .def_readwrite("trajectory_name",
+                     &lcmt_generators::TrajectoryDescription::trajectory_name)
+      .def_readwrite("variable_type",
+                     &lcmt_generators::TrajectoryDescription::variable_type)
+      .def_readwrite("indices",
+                     &lcmt_generators::TrajectoryDescription::indices);
+  py::class_<lcmt_generators::C3TrajectoryGeneratorConfig>(
+      m, "C3TrajectoryGeneratorConfig")
+      .def(py::init<>())
+      .def_readwrite(
+          "trajectories",
+          &lcmt_generators::C3TrajectoryGeneratorConfig::trajectories);
+  py::class_<lcmt_generators::C3TrajectoryGenerator,
+             drake::systems::LeafSystem<double>>(m, "C3TrajectoryGenerator")
+      .def(py::init<lcmt_generators::C3TrajectoryGeneratorConfig>())
+      .def("get_input_port_c3_solution",
+           &lcmt_generators::C3TrajectoryGenerator::get_input_port_c3_solution,
+           py::return_value_policy::reference)
+      .def("get_output_port_lcmt_c3_trajectory",
+           &lcmt_generators::C3TrajectoryGenerator::
+               get_output_port_lcmt_c3_trajectory,
+           py::return_value_policy::reference)
+      .def_static(
+          "AddLcmPublisherToBuilder",
+          &lcmt_generators::C3TrajectoryGenerator::AddLcmPublisherToBuilder,
+          py::arg("builder"), py::arg("config"), py::arg("solution_port"),
+          py::arg("channel"), py::arg("lcm"), py::arg("publish_triggers"),
+          py::arg("publish_period"), py::arg("publish_offset"));
+
+  // ContactForceGenerator
+  py::class_<lcmt_generators::ContactForceGenerator,
+             drake::systems::LeafSystem<double>>(m, "ContactForceGenerator")
+      .def(py::init<>())
+      .def("get_input_port_c3_solution",
+           &lcmt_generators::ContactForceGenerator::get_input_port_c3_solution,
+           py::return_value_policy::reference)
+      .def("get_input_port_lcs_contact_descriptions",
+           &lcmt_generators::ContactForceGenerator::
+               get_input_port_lcs_contact_descriptions,
+           py::return_value_policy::reference)
+      .def("get_output_port_contact_force",
+           &lcmt_generators::ContactForceGenerator::
+               get_output_port_contact_force,
+           py::return_value_policy::reference)
+      .def_static(
+          "AddLcmPublisherToBuilder",
+          &lcmt_generators::ContactForceGenerator::AddLcmPublisherToBuilder,
+          py::arg("builder"), py::arg("solution_port"),
+          py::arg("lcs_contact_info_port"), py::arg("channel"), py::arg("lcm"),
+          py::arg("publish_triggers"), py::arg("publish_period"),
+          py::arg("publish_offset"));
 }
 }  // namespace pyc3
 }  // namespace systems

core/BUILD.bazel

@@ -29,6 +29,11 @@ filegroup(
     ],
 )
 
+cc_library(
+    name = "logging",
+    hdrs = ["common/logging_utils.hpp"],
+)
+
 cc_library(
     name = "c3",
     srcs = [
@@ -57,6 +62,7 @@ cc_library(
     deps = [
         ":lcs",
         ":options",
+        ":logging",
         "@drake//:drake_shared_library",
     ] + select({
          "//tools:with_gurobi": ["@gurobi//:gurobi_cxx"],

core/c3.cc

@@ -7,6 +7,7 @@
 #include <drake/common/find_runfiles.h>
 #include <omp.h>
 
+#include "common/logging_utils.hpp"
 #include "lcs.h"
 #include "solver_options_io.h"
 
@@ -89,7 +90,6 @@ C3::C3(const LCS& lcs, const CostMatrices& costs,
   w_sol_ = std::make_unique<std::vector<VectorXd>>();
   delta_sol_ = std::make_unique<std::vector<VectorXd>>();
   for (int i = 0; i < N_; ++i) {
-    z_sol_->push_back(Eigen::VectorXd::Zero(n_z_));
     x_sol_->push_back(Eigen::VectorXd::Zero(n_x_));
     lambda_sol_->push_back(Eigen::VectorXd::Zero(n_lambda_));
     u_sol_->push_back(Eigen::VectorXd::Zero(n_u_));
@@ -267,6 +267,7 @@ const std::vector<drake::solvers::QuadraticCost*>& C3::GetTargetCost() {
 }
 
 void C3::Solve(const VectorXd& x0) {
+  drake::log()->debug("C3::Solve called");
   auto start = std::chrono::high_resolution_clock::now();
   // Set the initial state constraint
   if (initial_state_constraint_) {
@@ -317,6 +318,8 @@ void C3::Solve(const VectorXd& x0) {
   std::vector<VectorXd> w(N_, VectorXd::Zero(n_z_));
   vector<MatrixXd> G = cost_matrices_.G;
 
+  drake::log()->debug("C3::Solve starting ADMM iterations.");
+
   for (int iter = 0; iter < options_.admm_iter; iter++) {
     ADMMStep(x0, &delta, &w, &G, iter);
   }
@@ -332,6 +335,7 @@ void C3::Solve(const VectorXd& x0) {
   *delta_sol_ = delta;
 
   if (!options_.end_on_qp_step) {
+    drake::log()->debug("C3::Solve compute a half step.");
     *z_sol_ = delta;
     z_sol_->at(0).segment(0, n_x_) = x0;
     x_sol_->at(0) = x0;
@@ -355,6 +359,7 @@ void C3::Solve(const VectorXd& x0) {
   solve_time_ =
       std::chrono::duration_cast<std::chrono::microseconds>(elapsed).count() /
       1e6;
+  drake::log()->debug("C3::Solve completed in {} seconds.", solve_time_);
 }
 
 void C3::ADMMStep(const VectorXd& x0, vector<VectorXd>* delta,
@@ -373,6 +378,7 @@ void C3::ADMMStep(const VectorXd& x0, vector<VectorXd>* delta,
     ZW[i] = w->at(i) + z[i];
   }
 
+  drake::log()->debug("C3::ADMMStep SolveProjection step.");
   if (cost_matrices_.U[0].isZero(0)) {
     *delta = SolveProjection(*G, ZW, admm_iteration);
   } else {
@@ -417,10 +423,17 @@ void C3::StoreQPResults(const MathematicalProgramResult& result,
     z_sol_->at(i).segment(0, n_x_) = result.GetSolution(x_[i]);
     z_sol_->at(i).segment(n_x_, n_lambda_) = result.GetSolution(lambda_[i]);
     z_sol_->at(i).segment(n_x_ + n_lambda_, n_u_) = result.GetSolution(u_[i]);
+
+    drake::log()->trace(
+        "C3::StoreQPResults storing solution for time step {}:  "
+        "lambda = {}",
+        i, EigenToString(lambda_sol_->at(i).transpose()));
   }
 
   if (!warm_start_)
     return;  // No warm start, so no need to update warm start parameters
+
+  drake::log()->trace("C3::StoreQPResults storing warm start values.");
   for (int i = 0; i < N_ + 1; ++i) {
     if (i < N_) {
       warm_start_x_[admm_iteration][i] = result.GetSolution(x_[i]);
@@ -438,15 +451,18 @@ vector<VectorXd> C3::SolveQP(const VectorXd& x0, const vector<MatrixXd>& G,
   AddAugmentedCost(G, WD, delta, is_final_solve);
   SetInitialGuessQP(x0, admm_iteration);
 
-  MathematicalProgramResult result = osqp_.Solve(prog_);
-
-  if (!result.is_success()) {
-    drake::log()->warn("C3::SolveQP failed to solve the QP with status: {}",
-                       result.get_solution_result());
+  drake::log()->trace("C3::SolveQP calling solver.");
+  try {
+    MathematicalProgramResult result = osqp_.Solve(prog_);
+    if (!result.is_success()) {
+      drake::log()->warn("C3::SolveQP failed to solve the QP with status: {}",
+                         result.get_solution_result());
+    }
+    StoreQPResults(result, admm_iteration, is_final_solve);
+  } catch (const std::exception& e) {
+    drake::log()->error("C3::SolveQP failed with exception: {}", e.what());
   }
 
-  StoreQPResults(result, admm_iteration, is_final_solve);
-
   return *z_sol_;
 }
 
@@ -460,7 +476,7 @@ void C3::AddAugmentedCost(const vector<MatrixXd>& G, const vector<VectorXd>& WD,
 
   // If this is the final solve, we may want to scale the augmented cost
   if (is_final_solve) {
-    auto Gfinal = G;
+    vector<MatrixXd> Gfinal = G;
     for (int i = 0; i < N_; ++i) {
       Gfinal.at(i) *= options_.final_augmented_cost_scaling.value_or(1.0);
       DRAKE_ASSERT(Gfinal.at(i).cols() == Gfinal.at(i).rows());

core/c3_options.h

@@ -71,6 +71,7 @@ struct C3Options {
   std::vector<double> g_lambda_t;
   std::vector<double> g_lambda;
   std::vector<double> g_u;
+  std::vector<double> g_eta_vector;
   std::optional<std::vector<double>> g_eta_slack;
   std::optional<std::vector<double>> g_eta_n;
   std::optional<std::vector<double>> g_eta_t;
@@ -83,6 +84,7 @@ struct C3Options {
   std::vector<double> u_lambda_t;
   std::vector<double> u_lambda;
   std::vector<double> u_u;
+  std::vector<double> u_eta_vector;
   std::optional<std::vector<double>> u_eta_slack;
   std::optional<std::vector<double>> u_eta_n;
   std::optional<std::vector<double>> u_eta_t;
@@ -144,16 +146,14 @@ struct C3Options {
     }
     g_vector.insert(g_vector.end(), g_lambda.begin(), g_lambda.end());
     g_vector.insert(g_vector.end(), g_u.begin(), g_u.end());
-    if (g_eta != std::nullopt || g_eta_slack != std::nullopt) {
-      if (g_eta == std::nullopt || g_eta->empty()) {
-        g_vector.insert(g_vector.end(), g_eta_slack->begin(),
-                        g_eta_slack->end());
-        g_vector.insert(g_vector.end(), g_eta_n->begin(), g_eta_n->end());
-        g_vector.insert(g_vector.end(), g_eta_t->begin(), g_eta_t->end());
-      } else {
-        g_vector.insert(g_vector.end(), g_eta->begin(), g_eta->end());
-      }
+    g_eta_vector = g_eta.value_or(std::vector<double>());
+    if (g_eta_vector.empty() && g_eta_slack.has_value()) {
+      g_eta_vector.insert(g_eta_vector.end(), g_eta_slack->begin(),
+                          g_eta_slack->end());
+      g_eta_vector.insert(g_eta_vector.end(), g_eta_n->begin(), g_eta_n->end());
+      g_eta_vector.insert(g_eta_vector.end(), g_eta_t->begin(), g_eta_t->end());
     }
+    g_vector.insert(g_vector.end(), g_eta_vector.begin(), g_eta_vector.end());
 
     u_vector = std::vector<double>();
     u_vector.insert(u_vector.end(), u_x.begin(), u_x.end());
@@ -164,16 +164,14 @@ struct C3Options {
     }
     u_vector.insert(u_vector.end(), u_lambda.begin(), u_lambda.end());
     u_vector.insert(u_vector.end(), u_u.begin(), u_u.end());
-    if (u_eta != std::nullopt || u_eta_slack != std::nullopt) {
-      if (u_eta == std::nullopt || u_eta->empty()) {
-        u_vector.insert(u_vector.end(), u_eta_slack->begin(),
-                        u_eta_slack->end());
-        u_vector.insert(u_vector.end(), u_eta_n->begin(), u_eta_n->end());
-        u_vector.insert(u_vector.end(), u_eta_t->begin(), u_eta_t->end());
-      } else {
-        u_vector.insert(u_vector.end(), u_eta->begin(), u_eta->end());
-      }
+    u_eta_vector = u_eta.value_or(std::vector<double>());
+    if (u_eta_vector.empty() && u_eta_slack.has_value()) {
+      u_eta_vector.insert(u_eta_vector.end(), u_eta_slack->begin(),
+                          u_eta_slack->end());
+      u_eta_vector.insert(u_eta_vector.end(), u_eta_n->begin(), u_eta_n->end());
+      u_eta_vector.insert(u_eta_vector.end(), u_eta_t->begin(), u_eta_t->end());
     }
+    u_vector.insert(u_vector.end(), u_eta_vector.begin(), u_eta_vector.end());
 
     Eigen::VectorXd q = Eigen::Map<Eigen::VectorXd, Eigen::Unaligned>(
         this->q_vector.data(), this->q_vector.size());
@@ -184,8 +182,6 @@ struct C3Options {
     Eigen::VectorXd u = Eigen::Map<Eigen::VectorXd, Eigen::Unaligned>(
         this->u_vector.data(), this->u_vector.size());
 
-    DRAKE_DEMAND(g.size() == u.size());
-
     Q = w_Q * q.asDiagonal();
     R = w_R * r.asDiagonal();
     G = w_G * g.asDiagonal();