Merge pull request #31 from Simple-Robotics/topic/get_contact_state

edantec · web-flow · commit 1997888f0b2a · 2024-12-20T13:44:42.000+01:00
Add getters for contact vector and state derivative
diff --git a/bindings/expose-mpc.cpp b/bindings/expose-mpc.cpp
@@ -90,6 +90,7 @@ namespace simple_mpc
         .def("switchToStand", &MPC::switchToStand, "self"_a)
         .def("getFootTakeoffCycle", &MPC::getFootTakeoffCycle, ("self"_a, "ee_name"))
         .def("getFootLandCycle", &MPC::getFootLandCycle, ("self"_a, "ee_name"))
+        .def("getStateDerivative", &MPC::getStateDerivative, ("self"_a, "t"))
         .def("getCyclingContactState", &MPC::getCyclingContactState, ("self"_a, "t", "ee_name"))
         .def(
           "getModelHandler", &MPC::getModelHandler, "self"_a, bp::return_internal_reference<>(),
diff --git a/bindings/expose-problem.cpp b/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"))
diff --git a/examples/go2_fulldynamics.py b/examples/go2_fulldynamics.py
@@ -221,18 +221,11 @@
     end = time.time()
     solve_time.append(end - start)
 
-    a0 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[0]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
-    a1 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[1]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
+    a0 = mpc.getStateDerivative(0)[nv:]
+    a1 = mpc.getStateDerivative(1)[nv:]
 
-    FL_f, FR_f, RL_f, RR_f, contact_states = extract_forces(mpc.getTrajOptProblem(), mpc.solver.workspace, 0)
+    FL_f, FR_f, RL_f, RR_f = extract_forces(mpc.getTrajOptProblem(), mpc.solver.workspace, 0)
+    contact_states = mpc.ocp_handler.getContactState(0)
     total_forces = np.concatenate((FL_f, FR_f, RL_f, RR_f))
     force_FL.append(FL_f)
     force_FR.append(FR_f)
diff --git a/examples/go2_kinodynamics.py b/examples/go2_kinodynamics.py
@@ -225,23 +225,14 @@
     com_measured.append(mpc.getDataHandler().getData().com[0].copy())
     L_measured.append(mpc.getDataHandler().getData().hg.angular.copy())
 
-    a0 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[0]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
-    a1 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[1]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
+    a0 = mpc.getStateDerivative(0)[nv:]
+    a1 = mpc.getStateDerivative(1)[nv:]
+
     a0[6:] = mpc.us[0][nk * force_size :]
     a1[6:] = mpc.us[1][nk * force_size :]
     forces0 = mpc.us[0][: nk * force_size]
     forces1 = mpc.us[1][: nk * force_size]
-    contact_states = (
-        mpc.getTrajOptProblem().stages[0].dynamics.differential_dynamics.contact_states
-    )
+    contact_states = mpc.ocp_handler.getContactState(0)
 
     device.moveQuadrupedFeet(
         mpc.getReferencePose(0, "FL_foot").translation,
diff --git a/examples/talos_centroidal.py b/examples/talos_centroidal.py
@@ -201,18 +201,10 @@
         mpc.getReferencePose(0, "right_sole_link").translation,
     )
 
-    contact_states = (
-        mpc.getTrajOptProblem()
-        .stages[0]
-        .dynamics.differential_dynamics.contact_map.contact_states.tolist()
-    )
+    contact_states = mpc.ocp_handler.getContactState(0)
     foot_ref = [mpc.getReferencePose(0, name) for name in model_handler.getFeetNames()]
     foot_ref_next = [mpc.getReferencePose(1, name) for name in model_handler.getFeetNames()]
-    dH = (
-        mpc.solver
-        .workspace.problem_data.stage_data[0]
-        .dynamics_data.continuous_data.xdot[3:9]
-    )
+    dH = mpc.getStateDerivative(0)[3:9]
     qp.computeDifferences(
         mpc.getDataHandler().getData(), x_measured, foot_ref, foot_ref_next
     )
diff --git a/examples/talos_kinodynamics.py b/examples/talos_kinodynamics.py
@@ -210,23 +210,14 @@
     mpc.iterate(x_measured)
     end = time.time()
     print("MPC iterate = " + str(end - start))
-    a0 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[0]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
-    a1 = (
-        mpc.solver
-        .workspace.problem_data.stage_data[1]
-        .dynamics_data.continuous_data.xdot[nv:]
-    )
+
+    a0 = mpc.getStateDerivative(0)[nv:]
+    a1 = mpc.getStateDerivative(1)[nv:]
     a0[6:] = mpc.us[0][nk * force_size :]
     a1[6:] = mpc.us[1][nk * force_size :]
     forces0 = mpc.us[0][: nk * force_size]
     forces1 = mpc.us[1][: nk * force_size]
-    contact_states = (
-        mpc.getTrajOptProblem().stages[0].dynamics.differential_dynamics.contact_states
-    )
+    contact_states = mpc.ocp_handler.getContactState(0)
 
     device.moveMarkers(
         mpc.getReferencePose(0, "left_sole_link").translation,
diff --git a/examples/utils.py b/examples/utils.py
@@ -71,21 +71,15 @@ def extract_forces(problem, workspace, id):
     force_FR = np.zeros(3)
     force_RL = np.zeros(3)
     force_RR = np.zeros(3)
-    contact_states = [False, False, False, False]
     in_contact = problem.stages[id].dynamics.differential_dynamics.constraint_models.__len__()
-
     for i in range(in_contact):
         if problem.stages[id].dynamics.differential_dynamics.constraint_models[i].name == 'FL_foot':
-            contact_states[0] = True
             force_FL = workspace.problem_data.stage_data[id].dynamics_data.continuous_data.constraint_datas[i].contact_force.linear
         elif problem.stages[id].dynamics.differential_dynamics.constraint_models[i].name == 'FR_foot':
-            contact_states[1] = True
             force_FR = workspace.problem_data.stage_data[id].dynamics_data.continuous_data.constraint_datas[i].contact_force.linear
         elif problem.stages[id].dynamics.differential_dynamics.constraint_models[i].name == 'RL_foot':
-            contact_states[2] = True
             force_RL = workspace.problem_data.stage_data[id].dynamics_data.continuous_data.constraint_datas[i].contact_force.linear
         elif problem.stages[id].dynamics.differential_dynamics.constraint_models[i].name == 'RR_foot':
-            contact_states[3] = True
             force_RR = workspace.problem_data.stage_data[id].dynamics_data.continuous_data.constraint_datas[i].contact_force.linear
 
-    return force_FL, force_FR, force_RL, force_RR, contact_states
+    return force_FL, force_FR, force_RL, force_RR
diff --git a/include/simple-mpc/centroidal-dynamics.hpp b/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()
     {
diff --git a/include/simple-mpc/fulldynamics.hpp b/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_;
diff --git a/include/simple-mpc/kinodynamics.hpp b/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);
 
diff --git a/include/simple-mpc/mpc.hpp b/include/simple-mpc/mpc.hpp
@@ -7,6 +7,10 @@
 ///////////////////////////////////////////////////////////////////////////////
 #pragma once
 
+#include <aligator/core/stage-data.hpp>
+#include <aligator/fwd.hpp>
+#include <aligator/modelling/dynamics/fwd.hpp>
+#include <aligator/modelling/dynamics/integrator-explicit.hpp>
 #include <aligator/solvers/proxddp/solver-proxddp.hpp>
 
 #include "simple-mpc/deprecated.hpp"
@@ -17,6 +21,7 @@
 
 namespace simple_mpc
 {
+  using ExplicitIntegratorData = dynamics::ExplicitIntegratorDataTpl<double>;
 
   struct MPCSettings
   {
@@ -171,6 +176,14 @@ namespace simple_mpc
       }
     }
 
+    const ConstVectorRef getStateDerivative(const std::size_t t)
+    {
+      ExplicitIntegratorData * int_data =
+        dynamic_cast<ExplicitIntegratorData *>(&*solver_->workspace_.problem_data.stage_data[t]->dynamics_data);
+      assert(int_data != nullptr);
+      return int_data->continuous_data->xdot_;
+    }
+
     void switchToWalk(const Vector6d & velocity_base);
 
     void switchToStand();
diff --git a/include/simple-mpc/ocp-handler.hpp b/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
 
diff --git a/include/simple-mpc/python/py-ocp-handler.hpp b/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);
diff --git a/src/centroidal-dynamics.cpp b/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>();
+    assert(ode != nullptr);
+    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_);
diff --git a/src/fulldynamics.cpp b/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>();
+    assert(ode != nullptr);
+    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());
diff --git a/src/kinodynamics.cpp b/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>();
+    assert(ode != nullptr);
+    return ode->contact_states_;
+  }
+
   CostStack KinodynamicsOCP::createTerminalCost()
   {
     auto ter_space = MultibodyPhaseSpace(model_handler_.getModel());
diff --git a/tests/mpc.cpp b/tests/mpc.cpp
@@ -88,6 +88,8 @@ BOOST_AUTO_TEST_CASE(mpc_fulldynamics)
   BOOST_CHECK_EQUAL(mpc.foot_takeoff_times_.at("right_sole_link")[0], 100);
   BOOST_CHECK_EQUAL(mpc.foot_land_times_.at("left_sole_link")[0], 209);
   BOOST_CHECK_EQUAL(mpc.foot_land_times_.at("right_sole_link")[0], 150);
+
+  Eigen::VectorXd xdot = mpc.getStateDerivative(0);
 }
 
 BOOST_AUTO_TEST_CASE(mpc_kinodynamics)
@@ -161,6 +163,8 @@ BOOST_AUTO_TEST_CASE(mpc_kinodynamics)
   {
     mpc.iterate(model_handler.getReferenceState());
   }
+
+  Eigen::VectorXd xdot = mpc.getStateDerivative(0);
 }
 
 BOOST_AUTO_TEST_CASE(mpc_centroidal)
@@ -236,6 +240,8 @@ BOOST_AUTO_TEST_CASE(mpc_centroidal)
   {
     mpc.iterate(x_multibody);
   }
+
+  Eigen::VectorXd xdot = mpc.getStateDerivative(0);
 }
 
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/tests/problem.cpp b/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);

Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,7 @@ namespace simple_mpc`
`96`	`96`	`void setPoseBase(const std::size_t t, const ConstVectorRef & pose_base) override;`
`97`	`97`	`const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) override;`
`98`	`98`	`size_t getContactSupport(const std::size_t t) override;`
	`99`	`+ std::vector<bool> getContactState(const std::size_t t) override;`
`99`	`100`
`100`	`101`	`CentroidalSettings getSettings()`
`101`	`102`	`{`
Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,7 @@ namespace simple_mpc`
`101`	`101`	`void setPoseBase(const std::size_t t, const ConstVectorRef & pose_base) override;`
`102`	`102`	`const Eigen::VectorXd getProblemState(const RobotDataHandler & data_handler) override;`
`103`	`103`	`size_t getContactSupport(const std::size_t t) override;`
	`104`	`+ std::vector<bool> getContactState(const std::size_t t) override;`
`104`	`105`	`FullDynamicsSettings getSettings()`
`105`	`106`	`{`
`106`	`107`	`return settings_;`