Add getContactState function to ocp_handler

Author: edantec

bindings/expose-problem.cpp

@@ -71,6 +71,7 @@ namespace simple_mpc
         .def("getPoseBase", bp::pure_virtual(&OCPHandler::getPoseBase), bp::args("self", "t"))
         .def("getProblemState", bp::pure_virtual(&OCPHandler::getProblemState), bp::args("self", "data_handler"))
         .def("getContactSupport", bp::pure_virtual(&OCPHandler::getContactSupport), bp::args("self", "t"))
+        .def("getContactState", bp::pure_virtual(&OCPHandler::getContactState), bp::args("self", "t"))
         .def(
           "createProblem", &OCPHandler::createProblem,
           ("self"_a, "x0", "horizon", "force_size", "gravity", "terminal_constraint"))

include/simple-mpc/centroidal-dynamics.hpp

@@ -96,6 +96,7 @@ namespace simple_mpc
     void setPoseBase(const std::size_t t, const ConstVectorRef & pose_base) override;
     const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) override;
     size_t getContactSupport(const std::size_t t) override;
+    std::vector<bool> getContactState(const std::size_t t) override;
 
     CentroidalSettings getSettings()
     {

include/simple-mpc/fulldynamics.hpp

@@ -101,6 +101,7 @@ namespace simple_mpc
     void setPoseBase(const std::size_t t, const ConstVectorRef & pose_base) override;
     const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) override;
     size_t getContactSupport(const std::size_t t) override;
+    std::vector<bool> getContactState(const std::size_t t) override;
     FullDynamicsSettings getSettings()
     {
       return settings_;

include/simple-mpc/kinodynamics.hpp

@@ -89,6 +89,7 @@ namespace simple_mpc
     void setVelocityBase(const std::size_t t, const ConstVectorRef & velocity_base) override;
     const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) override;
     size_t getContactSupport(const std::size_t t) override;
+    std::vector<bool> getContactState(const std::size_t t) override;
 
     void computeControlFromForces(const std::map<std::string, Eigen::VectorXd> & force_refs);
 

include/simple-mpc/ocp-handler.hpp

@@ -95,6 +95,7 @@ namespace simple_mpc
     virtual const Eigen::VectorXd getReferenceForce(const std::size_t t, const std::string & ee_name) = 0;
     virtual const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) = 0;
     virtual size_t getContactSupport(const std::size_t t) = 0;
+    virtual std::vector<bool> getContactState(const std::size_t t) = 0;
 
     /// Common functions for all problems
 

include/simple-mpc/python/py-ocp-handler.hpp

@@ -165,6 +165,11 @@ namespace simple_mpc
         SIMPLE_MPC_PYTHON_OVERRIDE_PURE(std::size_t, "getContactSupport", t);
       }
 
+      std::vector<bool> getContactState(const std::size_t t) override
+      {
+        SIMPLE_MPC_PYTHON_OVERRIDE_PURE(std::vector<bool>, "getContactState", t);
+      }
+
       void setReferenceControl(const std::size_t t, const ConstVectorRef & u_ref)
       {
         SIMPLE_MPC_PYTHON_OVERRIDE(void, OCPHandler, setReferenceControl, t, u_ref);

src/centroidal-dynamics.cpp

@@ -278,6 +278,18 @@ namespace simple_mpc
     return active_contacts;
   }
 
+  std::vector<bool> CentroidalOCP::getContactState(const std::size_t t)
+  {
+    std::vector<bool> contact_state;
+    CentroidalFwdDynamics * ode =
+      problem_->stages_[t]->getDynamics<IntegratorEuler>()->getDynamics<CentroidalFwdDynamics>();
+
+    for (auto name : model_handler_.getFeetNames())
+      contact_state.push_back(ode->contact_map_.getContactState(name));
+
+    return contact_state;
+  }
+
   CostStack CentroidalOCP::createTerminalCost()
   {
     auto ter_space = VectorSpace(nx_);

src/fulldynamics.cpp

@@ -373,6 +373,31 @@ namespace simple_mpc
     return ode->constraint_models_.size();
   }
 
+  std::vector<bool> FullDynamicsOCP::getContactState(const std::size_t t)
+  {
+    std::vector<bool> contact_state;
+    MultibodyConstraintFwdDynamics * ode =
+      problem_->stages_[t]->getDynamics<IntegratorSemiImplEuler>()->getDynamics<MultibodyConstraintFwdDynamics>();
+
+    for (auto name : model_handler_.getFeetNames())
+    {
+      std::size_t i;
+      for (i = 0; i < ode->constraint_models_.size(); i++)
+      {
+        if (ode->constraint_models_[i].name == name)
+        {
+          contact_state.push_back(true);
+          break;
+        }
+      }
+      if (i == ode->constraint_models_.size())
+      {
+        contact_state.push_back(false);
+      }
+    }
+    return contact_state;
+  }
+
   CostStack FullDynamicsOCP::createTerminalCost()
   {
     auto ter_space = MultibodyPhaseSpace(model_handler_.getModel());

src/kinodynamics.cpp

@@ -322,6 +322,14 @@ namespace simple_mpc
     return active_contacts;
   }
 
+  std::vector<bool> KinodynamicsOCP::getContactState(const std::size_t t)
+  {
+    KinodynamicsFwdDynamics * ode =
+      problem_->stages_[t]->getDynamics<IntegratorSemiImplEuler>()->getDynamics<KinodynamicsFwdDynamics>();
+
+    return ode->contact_states_;
+  }
+
   CostStack KinodynamicsOCP::createTerminalCost()
   {
     auto ter_space = MultibodyPhaseSpace(model_handler_.getModel());

tests/problem.cpp

@@ -55,8 +55,10 @@ BOOST_AUTO_TEST_CASE(fulldynamics)
   QuadraticResidualCost * crc = csp->getComponent<QuadraticResidualCost>("centroidal_cost");
   QuadraticResidualCost * cpc = csp->getComponent<QuadraticResidualCost>("left_sole_link_pose_cost");
 
+  std::vector<bool> cs2 = {true, true};
   BOOST_CHECK_EQUAL(fdproblem.getSize(), 100);
   BOOST_CHECK_EQUAL(fdproblem.getContactSupport(2), 2);
+  BOOST_TEST(fdproblem.getContactState(2) == cs2);
   BOOST_CHECK_EQUAL(cc->weights_, settings.w_u);
   BOOST_CHECK_EQUAL(crc->weights_, settings.w_cent);
   BOOST_CHECK_EQUAL(cpc->weights_, settings.w_frame);
@@ -137,7 +139,9 @@ BOOST_AUTO_TEST_CASE(kinodynamics)
   QuadraticResidualCost * crc = csp->getComponent<QuadraticResidualCost>("centroidal_cost");
   QuadraticResidualCost * cpc = csp->getComponent<QuadraticResidualCost>("left_sole_link_pose_cost");
 
+  std::vector<bool> cs2 = {true, true};
   BOOST_CHECK_EQUAL(knproblem.getContactSupport(2), 2);
+  BOOST_TEST(knproblem.getContactState(2) == cs2);
   BOOST_CHECK_EQUAL(cc->weights_, settings.w_u);
   BOOST_CHECK_EQUAL(crc->weights_, settings.w_cent);
   BOOST_CHECK_EQUAL(cpc->weights_, settings.w_frame);
@@ -220,7 +224,9 @@ BOOST_AUTO_TEST_CASE(centroidal)
   QuadraticResidualCost * crc = csp->getComponent<QuadraticResidualCost>("linear_mom_cost");
   QuadraticResidualCost * cpc = csp->getComponent<QuadraticResidualCost>("angular_acc_cost");
 
+  std::vector<bool> cs2 = {true, true};
   BOOST_CHECK_EQUAL(cproblem.getContactSupport(2), 2);
+  BOOST_TEST(cproblem.getContactState(2) == cs2);
   BOOST_CHECK_EQUAL(cc->weights_, settings.w_u);
   BOOST_CHECK_EQUAL(crc->weights_, settings.w_linear_mom);
   BOOST_CHECK_EQUAL(cpc->weights_, settings.w_angular_acc);
@@ -308,7 +314,9 @@ BOOST_AUTO_TEST_CASE(centroidal_solo)
   QuadraticResidualCost * crc = csp->getComponent<QuadraticResidualCost>("linear_mom_cost");
   QuadraticResidualCost * cpc = csp->getComponent<QuadraticResidualCost>("angular_acc_cost");
 
+  std::vector<bool> cs2 = {true, true, true, true};
   BOOST_CHECK_EQUAL(cproblem.getContactSupport(2), 4);
+  BOOST_TEST(cproblem.getContactState(2) == cs2);
   BOOST_CHECK_EQUAL(cc->weights_, settings.w_u);
   BOOST_CHECK_EQUAL(crc->weights_, settings.w_linear_mom);
   BOOST_CHECK_EQUAL(cpc->weights_, settings.w_angular_acc);