[modelling] Change CostStack backend to an unordered_map type

.github/workflows/macos-linux-conda.yml

@@ -79,7 +79,7 @@ jobs:
           -DBUILD_CROCODDYL_COMPAT=ON \
           -DBUILD_BENCHMARKS=ON \
           -DBUILD_EXAMPLES=ON
-        cmake --build .
+        cmake --build . -j2
         ctest --output-on-failure
         cmake --install .
 

CHANGELOG.md

@@ -10,9 +10,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Getter `getResidual<Derived>()` for composite cost functions ([#198](https://github.com/Simple-Robotics/aligator/pull/198))
+- Getter `getComponent<Derived>()` for `CostStack` ([#199](https://github.com/Simple-Robotics/aligator/pull/199))
 
 ### Changed
 
+- Change storage of `CostStack` to `boost::unordered::unordered_map`, pointing to pair of cost function and weight ([#199](https://github.com/Simple-Robotics/aligator/pull/199))
 - Change storage for `ConstraintStack` to using two `std::vector<polymorphic<>>` the struct `StageConstraintTpl` is now merely a convenient API shortcut for the end-user.
 - Remove `StageConstraintTpl::nr()` (in C++ only)
 - Update minimum required version of eigenpy to 3.7.0

bindings/python/src/modelling/expose-cost-stack.cpp

@@ -3,6 +3,18 @@
 
 #include "aligator/modelling/costs/sum-of-costs.hpp"
 
+#include <eigenpy/std-pair.hpp>
+#include <eigenpy/variant.hpp>
+#if EIGENPY_VERSION_AT_LEAST(3, 9, 1)
+#define ALIGATOR_EIGENPY_HAS_MAP_SUPPORT 1
+#else
+#define ALIGATOR_EIGENPY_HAS_MAP_SUPPORT 0
+#endif
+
+#if ALIGATOR_EIGENPY_HAS_MAP_SUPPORT
+#include <eigenpy/std-map.hpp>
+#endif
+
 namespace aligator {
 namespace python {
 using context::CostAbstract;
@@ -13,28 +25,63 @@ using context::Scalar;
 void exposeCostStack() {
   using CostStack = CostStackTpl<Scalar>;
   using CostStackData = CostStackDataTpl<Scalar>;
+  using CostKey = CostStack::CostKey;
+  using PolyCost = CostStack::PolyCost;
+  using PolyManifold = xyz::polymorphic<Manifold>;
+  using CostItem = CostStack::CostItem;
+  using CostMap = CostStack::CostMap;
+  eigenpy::StdPairConverter<CostItem>::registration();
+  eigenpy::VariantConverter<CostKey>::registration();
+
+  {
+    bp::scope scope =
+        bp::class_<CostStack, bp::bases<CostAbstract>>(
+            "CostStack", "A weighted sum of other cost functions.", bp::no_init)
+            .def(
+                bp::init<PolyManifold, const int, const std::vector<PolyCost> &,
+                         const std::vector<Scalar> &>(
+                    ("self"_a, "space", "nu", "components"_a = bp::list(),
+                     "weights"_a = bp::list())))
+            .def(bp::init<const PolyCost &>(("self"_a, "cost")))
+#if ALIGATOR_EIGENPY_HAS_MAP_SUPPORT
+            .def(bp::init<PolyManifold, int, const CostMap &>(
+                ("self"_a, "components"),
+                "Construct the CostStack from a CostMap object."))
+            .def_readonly("components", &CostStack::components_,
+                          "Components of this cost stack.")
+#endif
+            .def(
+                "addCost",
+                +[](CostStack &self, const PolyCost &cost, const Scalar weight)
+                    -> PolyCost & { return self.addCost(cost, weight).first; },
+                bp::return_internal_reference<>(),
+                ("self"_a, "cost", "weight"_a = 1.))
+            .def(
+                "addCost",
+                +[](CostStack &self, CostKey key, const PolyCost &cost,
+                    const Scalar weight) -> PolyCost & {
+                  return self.addCost(key, cost, weight).first;
+                },
+                bp::return_internal_reference<>(),
+                ("self"_a, "key", "cost", "weight"_a = 1.))
+            .def("size", &CostStack::size, "Get the number of cost components.")
+            .def(CopyableVisitor<CostStack>())
+            .def(PolymorphicMultiBaseVisitor<CostAbstract>());
+#if ALIGATOR_EIGENPY_HAS_MAP_SUPPORT
+    eigenpy::GenericMapVisitor<CostMap, true>::expose("CostMap");
+#endif
+  }
 
-  bp::class_<CostStack, bp::bases<CostAbstract>>(
-      "CostStack", "A weighted sum of other cost functions.",
-      bp::init<xyz::polymorphic<Manifold>, int,
-               const std::vector<xyz::polymorphic<CostAbstract>> &,
-               const std::vector<Scalar> &>(("self"_a, "space", "nu",
-                                             "components"_a = bp::list(),
-                                             "weights"_a = bp::list())))
-      .def_readwrite("components", &CostStack::components_,
-                     "Components of this cost stack.")
-      .def_readonly("weights", &CostStack::weights_,
-                    "Weights of this cost stack.")
-      .def("addCost", &CostStack::addCost, ("self"_a, "cost", "weight"_a = 1.),
-           "Add a cost to the stack of costs.")
-      .def("size", &CostStack::size, "Get the number of cost components.")
-      .def(CopyableVisitor<CostStack>())
-      .def(PolymorphicMultiBaseVisitor<CostAbstract>());
-
-  bp::register_ptr_to_python<shared_ptr<CostStackData>>();
-  bp::class_<CostStackData, bp::bases<CostData>>(
-      "CostStackData", "Data struct for CostStack.", bp::no_init)
-      .def_readonly("sub_cost_data", &CostStackData::sub_cost_data);
+  {
+    bp::register_ptr_to_python<shared_ptr<CostStackData>>();
+    bp::scope scope =
+        bp::class_<CostStackData, bp::bases<CostData>>(
+            "CostStackData", "Data struct for CostStack.", bp::no_init)
+            .def_readonly("sub_cost_data", &CostStackData::sub_cost_data);
+#if ALIGATOR_EIGENPY_HAS_MAP_SUPPORT
+    eigenpy::GenericMapVisitor<CostStackData::DataMap, true>::expose("DataMap");
+#endif
+  }
 }
 
 } // namespace python

examples/talos-walk-utils.cpp

@@ -187,8 +187,8 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
 
     auto rcost = CostStack(stage_space, nu);
 
-    rcost.addCost(QuadraticStateCost(stage_space, nu, x0, w_x));
-    rcost.addCost(QuadraticControlCost(stage_space, u0, w_u));
+    rcost.addCost("quad_state", QuadraticStateCost(stage_space, nu, x0, w_x));
+    rcost.addCost("quad_control", QuadraticControlCost(stage_space, u0, w_u));
     pin::SE3 LF_placement = rdata.oMf[foot_frame_ids[0]];
     pin::SE3 RF_placement = rdata.oMf[foot_frame_ids[1]];
     std::shared_ptr<FramePlacementResidual> frame_fn_RF;
@@ -198,13 +198,15 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
       foot_traj(RF_placement.translation(), T_ss, ts);
       frame_fn_RF = std::make_shared<FramePlacementResidual>(
           stage_space.ndx(), nu, rmodel, RF_placement, foot_frame_ids[1]);
-      rcost.addCost(QuadraticResidualCost(stage_space, *frame_fn_RF, w_LFRF));
+      rcost.addCost("frame_fn_RF",
+                    QuadraticResidualCost(stage_space, *frame_fn_RF, w_LFRF));
       break;
     case RIGHT:
       foot_traj(LF_placement.translation(), T_ss, ts);
       frame_fn_LF = std::make_shared<FramePlacementResidual>(
           stage_space.ndx(), nu, rmodel, LF_placement, foot_frame_ids[0]);
-      rcost.addCost(QuadraticResidualCost(stage_space, *frame_fn_LF, w_LFRF));
+      rcost.addCost("frame_fn_LF",
+                    QuadraticResidualCost(stage_space, *frame_fn_LF, w_LFRF));
       break;
     case DOUBLE:
       ts = 0;
@@ -216,7 +218,7 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
   }
   auto ter_space = MultibodyPhaseSpace(rmodel);
   auto term_cost = CostStack(ter_space, nu);
-  term_cost.addCost(QuadraticStateCost(ter_space, nu, x0, w_x));
+  term_cost.addCost("quad_state", QuadraticStateCost(ter_space, nu, x0, w_x));
 
   return TrajOptProblem(x0, stage_models, term_cost);
 }

include/aligator/modelling/costs/sum-of-costs.hpp

@@ -1,6 +1,9 @@
 #pragma once
 
 #include "aligator/core/cost-abstract.hpp"
+// Faster than std::unordered_map with Bost 1.80
+// https://martin.ankerl.com/2022/08/27/hashmap-bench-01/#boost__unordered_map
+#include <boost/unordered_map.hpp>
 
 namespace aligator {
 
@@ -18,30 +21,65 @@ template <typename _Scalar> struct CostStackTpl : CostAbstractTpl<_Scalar> {
   ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
   using CostBase = CostAbstractTpl<Scalar>;
   using CostData = CostDataAbstractTpl<Scalar>;
-  using CostPtr = xyz::polymorphic<CostBase>;
+  using PolyCost = xyz::polymorphic<CostBase>;
   using SumCostData = CostStackDataTpl<Scalar>;
   using Manifold = ManifoldAbstractTpl<Scalar>;
+  using CostItem = std::pair<PolyCost, Scalar>;
+  using CostKey = std::variant<std::size_t, std::string>;
+  using CostMap = boost::unordered::unordered_map<CostKey, CostItem>;
+  using CostIterator = typename CostMap::iterator;
 
-  std::vector<CostPtr> components_;
-  std::vector<Scalar> weights_;
+  CostMap components_;
 
   /// @brief    Check the dimension of a component.
   /// @returns  A bool value indicating whether the component is OK to be added
   /// to this instance.
-  bool checkDimension(const xyz::polymorphic<CostBase> comp) const;
+  bool checkDimension(const CostBase &comp) const;
 
   /// @brief  Constructor with a specified dimension, and optional vector of
   /// components and weights.
   CostStackTpl(xyz::polymorphic<Manifold> space, const int nu,
-               const std::vector<CostPtr> &comps = {},
+               const std::vector<PolyCost> &comps = {},
                const std::vector<Scalar> &weights = {});
 
+  CostStackTpl(xyz::polymorphic<Manifold> space, const int nu,
+               const CostMap &comps)
+      : CostBase(space, nu), components_(comps) {
+    for (const auto &[key, item] : comps) {
+      auto &cost = *item.first;
+      if (!this->checkDimension(cost)) {
+        ALIGATOR_DOMAIN_ERROR(fmt::format(
+            "Cannot add new component due to inconsistent input dimensions "
+            "(got ({:d}, {:d}), expected ({:d}, {:d}))",
+            cost.ndx(), cost.nu, this->ndx(), this->nu));
+      }
+    }
+  }
+
   /// @brief  Constructor from a single CostBase instance.
-  CostStackTpl(const CostPtr &cost);
+  CostStackTpl(const PolyCost &cost);
 
-  void addCost(const CostPtr &cost, const Scalar weight = 1.);
+  inline CostItem &addCost(const PolyCost &cost, const Scalar weight = 1.) {
+    const std::size_t size = components_.size();
+    return this->addCost(size, cost, weight);
+  }
 
-  std::size_t size() const;
+  CostItem &addCost(const CostKey &key, const PolyCost &cost,
+                    const Scalar weight = 1.);
+
+  inline std::size_t size() const { return components_.size(); }
+
+  /// @brief Get component, cast down to the specified type.
+  template <typename Derived> Derived *getComponent(const CostKey &key) {
+    CostItem &item = components_.at(key);
+    return dynamic_cast<Derived *>(&*item.first);
+  }
+
+  template <typename Derived>
+  const Derived *getComponent(const CostKey &key) const {
+    CostItem &item = components_.at(key);
+    return dynamic_cast<const Derived *>(&*item.first);
+  }
 
   void evaluate(const ConstVectorRef &x, const ConstVectorRef &u,
                 CostData &data) const;
@@ -89,16 +127,11 @@ operator+(const xyz::polymorphic<CostStackTpl<T>> &c1, CostPtr<T> &&c2) {
   return c1;
 }
 
-template <typename T>
-xyz::polymorphic<CostStackTpl<T>> operator*(T u, const CostPtr<T> &c1) {
-  return std::make_shared<CostStackTpl<T>>({c1}, {u});
-}
-
 template <typename T>
 xyz::polymorphic<CostStackTpl<T>>
 operator*(T u, xyz::polymorphic<CostStackTpl<T>> &&c1) {
-  for (std::size_t i = 0; i < c1->size(); i++) {
-    c1->weights_[i] *= u;
+  for (auto &[key, item] : c1->components_) {
+    item.second *= u;
   }
   return c1;
 }
@@ -107,13 +140,15 @@ template <typename _Scalar>
 struct CostStackDataTpl : CostDataAbstractTpl<_Scalar> {
   using Scalar = _Scalar;
   using CostData = CostDataAbstractTpl<Scalar>;
-  std::vector<shared_ptr<CostData>> sub_cost_data;
+  using CostStack = CostStackTpl<Scalar>;
+  using CostKey = typename CostStack::CostKey;
+  using DataMap =
+      boost::unordered::unordered_map<CostKey, shared_ptr<CostData>>;
+  DataMap sub_cost_data;
   CostStackDataTpl(const CostStackTpl<Scalar> &obj);
 };
 } // namespace aligator
 
-#include "aligator/modelling/costs/sum-of-costs.hxx"
-
 #ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
 #include "aligator/modelling/costs/sum-of-costs.txx"
 #endif