Skip to content

Commit 9f470d4

Browse files
ManifoldFRManifoldFR
committed
[modelling/costs] CostStack : change storage from std::vector to boost::unordered_map
+ add casting getter for components
1 parent ac8f8d9 commit 9f470d4

File tree

4 files changed

+109
-54
lines changed

4 files changed

+109
-54
lines changed

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

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33

44
#include "aligator/modelling/costs/sum-of-costs.hpp"
55

6+
#include <eigenpy/std-map.hpp>
7+
#include <eigenpy/std-pair.hpp>
8+
69
namespace aligator {
710
namespace python {
811
using context::CostAbstract;
@@ -13,20 +16,35 @@ using context::Scalar;
1316
void exposeCostStack() {
1417
using CostStack = CostStackTpl<Scalar>;
1518
using CostStackData = CostStackDataTpl<Scalar>;
19+
using CostKey = CostStack::CostKey;
20+
using PolyCost = CostStack::PolyCost;
21+
using CostItem = CostStack::CostItem;
22+
using CostMap = CostStack::CostMap;
23+
eigenpy::StdPairConverter<CostItem>::registration();
1624

1725
bp::class_<CostStack, bp::bases<CostAbstract>>(
18-
"CostStack", "A weighted sum of other cost functions.",
19-
bp::init<xyz::polymorphic<Manifold>, int,
20-
const std::vector<xyz::polymorphic<CostAbstract>> &,
21-
const std::vector<Scalar> &>(("self"_a, "space", "nu",
22-
"components"_a = bp::list(),
23-
"weights"_a = bp::list())))
24-
.def_readwrite("components", &CostStack::components_,
25-
"Components of this cost stack.")
26-
.def_readonly("weights", &CostStack::weights_,
27-
"Weights of this cost stack.")
28-
.def("addCost", &CostStack::addCost, ("self"_a, "cost", "weight"_a = 1.),
29-
"Add a cost to the stack of costs.")
26+
"CostStack", "A weighted sum of other cost functions.", bp::no_init)
27+
.def(bp::init<xyz::polymorphic<Manifold>, const int,
28+
const std::vector<PolyCost> &, const std::vector<Scalar> &>(
29+
("self"_a, "space", "nu", "components"_a = bp::list(),
30+
"weights"_a = bp::list())))
31+
.def(bp::init<const PolyCost &>(("self"_a, "cost")))
32+
// .def_readwrite("components", &CostStack::components_,
33+
// "Components of this cost stack.")
34+
.def(
35+
"addCost",
36+
+[](CostStack &self, const PolyCost &cost, const Scalar weight)
37+
-> CostItem & { return self.addCost(cost, weight); },
38+
("self"_a, "cost", "weight"_a = 1.),
39+
bp::return_internal_reference<>())
40+
.def(
41+
"addCost",
42+
+[](CostStack &self, CostKey key, const PolyCost &cost,
43+
const Scalar weight) -> CostItem & {
44+
return self.addCost(key, cost, weight);
45+
},
46+
("self"_a, "key", "cost", "weight"_a = 1.),
47+
bp::return_internal_reference<>())
3048
.def("size", &CostStack::size, "Get the number of cost components.")
3149
.def(CopyableVisitor<CostStack>())
3250
.def(PolymorphicMultiBaseVisitor<CostAbstract>());

‎examples/talos-walk-utils.cpp‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,8 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
187187

188188
auto rcost = CostStack(stage_space, nu);
189189

190-
rcost.addCost(QuadraticStateCost(stage_space, nu, x0, w_x));
191-
rcost.addCost(QuadraticControlCost(stage_space, u0, w_u));
190+
rcost.addCost("quad_state", QuadraticStateCost(stage_space, nu, x0, w_x));
191+
rcost.addCost("quad_control", QuadraticControlCost(stage_space, u0, w_u));
192192
pin::SE3 LF_placement = rdata.oMf[foot_frame_ids[0]];
193193
pin::SE3 RF_placement = rdata.oMf[foot_frame_ids[1]];
194194
std::shared_ptr<FramePlacementResidual> frame_fn_RF;
@@ -198,13 +198,15 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
198198
foot_traj(RF_placement.translation(), T_ss, ts);
199199
frame_fn_RF = std::make_shared<FramePlacementResidual>(
200200
stage_space.ndx(), nu, rmodel, RF_placement, foot_frame_ids[1]);
201-
rcost.addCost(QuadraticResidualCost(stage_space, *frame_fn_RF, w_LFRF));
201+
rcost.addCost("frame_fn_RF",
202+
QuadraticResidualCost(stage_space, *frame_fn_RF, w_LFRF));
202203
break;
203204
case RIGHT:
204205
foot_traj(LF_placement.translation(), T_ss, ts);
205206
frame_fn_LF = std::make_shared<FramePlacementResidual>(
206207
stage_space.ndx(), nu, rmodel, LF_placement, foot_frame_ids[0]);
207-
rcost.addCost(QuadraticResidualCost(stage_space, *frame_fn_LF, w_LFRF));
208+
rcost.addCost("frame_fn_LF",
209+
QuadraticResidualCost(stage_space, *frame_fn_LF, w_LFRF));
208210
break;
209211
case DOUBLE:
210212
ts = 0;
@@ -216,7 +218,7 @@ TrajOptProblem defineLocomotionProblem(const std::size_t T_ss,
216218
}
217219
auto ter_space = MultibodyPhaseSpace(rmodel);
218220
auto term_cost = CostStack(ter_space, nu);
219-
term_cost.addCost(QuadraticStateCost(ter_space, nu, x0, w_x));
221+
term_cost.addCost("quad_state", QuadraticStateCost(ter_space, nu, x0, w_x));
220222

221223
return TrajOptProblem(x0, stage_models, term_cost);
222224
}

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

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
#pragma once
22

33
#include "aligator/core/cost-abstract.hpp"
4+
// Faster than std::unordered_map with Bost 1.80
5+
// https://martin.ankerl.com/2022/08/27/hashmap-bench-01/#boost__unordered_map
6+
#include <boost/unordered_map.hpp>
47

58
namespace aligator {
69

@@ -18,12 +21,15 @@ template <typename _Scalar> struct CostStackTpl : CostAbstractTpl<_Scalar> {
1821
ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
1922
using CostBase = CostAbstractTpl<Scalar>;
2023
using CostData = CostDataAbstractTpl<Scalar>;
21-
using CostPtr = xyz::polymorphic<CostBase>;
24+
using PolyCost = xyz::polymorphic<CostBase>;
2225
using SumCostData = CostStackDataTpl<Scalar>;
2326
using Manifold = ManifoldAbstractTpl<Scalar>;
27+
using CostItem = std::pair<PolyCost, Scalar>;
28+
using CostKey = std::variant<std::size_t, std::string>;
29+
using CostMap = boost::unordered::unordered_map<CostKey, CostItem>;
30+
using CostIterator = typename CostMap::iterator;
2431

25-
std::vector<CostPtr> components_;
26-
std::vector<Scalar> weights_;
32+
CostMap components_;
2733

2834
/// @brief Check the dimension of a component.
2935
/// @returns A bool value indicating whether the component is OK to be added
@@ -33,15 +39,41 @@ template <typename _Scalar> struct CostStackTpl : CostAbstractTpl<_Scalar> {
3339
/// @brief Constructor with a specified dimension, and optional vector of
3440
/// components and weights.
3541
CostStackTpl(xyz::polymorphic<Manifold> space, const int nu,
36-
const std::vector<CostPtr> &comps = {},
42+
const std::vector<PolyCost> &comps = {},
3743
const std::vector<Scalar> &weights = {});
3844

45+
CostStackTpl(xyz::polymorphic<Manifold> space, const int nu,
46+
const CostMap &comps)
47+
: CostBase(space, nu), components_(comps) {
48+
for (const auto &[key, item] : comps) {
49+
this->checkDimension(item.first);
50+
}
51+
}
52+
3953
/// @brief Constructor from a single CostBase instance.
40-
CostStackTpl(const CostPtr &cost);
54+
CostStackTpl(const PolyCost &cost);
55+
56+
inline CostItem &addCost(const PolyCost &cost, const Scalar weight = 1.) {
57+
const std::size_t size = components_.size();
58+
return this->addCost(size, cost, weight);
59+
}
60+
61+
CostItem &addCost(const CostKey &key, const PolyCost &cost,
62+
const Scalar weight = 1.);
4163

42-
void addCost(const CostPtr &cost, const Scalar weight = 1.);
64+
inline std::size_t size() const { return components_.size(); }
4365

44-
std::size_t size() const;
66+
/// @brief Get component, cast down to the specified type.
67+
template <typename Derived> Derived *getComponent(const CostKey &key) {
68+
CostItem &item = components_.at(key);
69+
return dynamic_cast<Derived *>(&*item.first);
70+
}
71+
72+
template <typename Derived>
73+
const Derived *getComponent(const CostKey &key) const {
74+
CostItem &item = components_.at(key);
75+
return dynamic_cast<const Derived *>(&*item.first);
76+
}
4577

4678
void evaluate(const ConstVectorRef &x, const ConstVectorRef &u,
4779
CostData &data) const;
@@ -89,16 +121,11 @@ operator+(const xyz::polymorphic<CostStackTpl<T>> &c1, CostPtr<T> &&c2) {
89121
return c1;
90122
}
91123

92-
template <typename T>
93-
xyz::polymorphic<CostStackTpl<T>> operator*(T u, const CostPtr<T> &c1) {
94-
return std::make_shared<CostStackTpl<T>>({c1}, {u});
95-
}
96-
97124
template <typename T>
98125
xyz::polymorphic<CostStackTpl<T>>
99126
operator*(T u, xyz::polymorphic<CostStackTpl<T>> &&c1) {
100-
for (std::size_t i = 0; i < c1->size(); i++) {
101-
c1->weights_[i] *= u;
127+
for (auto &[key, item] : c1->components_) {
128+
item.second *= u;
102129
}
103130
return c1;
104131
}
@@ -107,7 +134,11 @@ template <typename _Scalar>
107134
struct CostStackDataTpl : CostDataAbstractTpl<_Scalar> {
108135
using Scalar = _Scalar;
109136
using CostData = CostDataAbstractTpl<Scalar>;
110-
std::vector<shared_ptr<CostData>> sub_cost_data;
137+
using CostStack = CostStackTpl<Scalar>;
138+
using CostKey = typename CostStack::CostKey;
139+
using DataMap =
140+
boost::unordered::unordered_map<CostKey, shared_ptr<CostData>>;
141+
DataMap sub_cost_data;
111142
CostStackDataTpl(const CostStackTpl<Scalar> &obj);
112143
};
113144
} // namespace aligator

‎include/aligator/modelling/costs/sum-of-costs.hxx‎

Lines changed: 26 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@ namespace aligator {
66
template <typename Scalar>
77
CostStackTpl<Scalar>::CostStackTpl(xyz::polymorphic<Manifold> space,
88
const int nu,
9-
const std::vector<CostPtr> &comps,
9+
const std::vector<PolyCost> &comps,
1010
const std::vector<Scalar> &weights)
11-
: CostBase(space, nu), components_(comps), weights_(weights) {
11+
: CostBase(space, nu) {
1212
if (comps.size() != weights.size()) {
1313
auto msg = fmt::format(
1414
"Inconsistent number of components ({:d}) and weights ({:d}).",
@@ -25,12 +25,19 @@ CostStackTpl<Scalar>::CostStackTpl(xyz::polymorphic<Manifold> space,
2525
ALIGATOR_RUNTIME_ERROR(msg);
2626
}
2727
}
28+
29+
for (std::size_t i = 0; i < comps.size(); i++) {
30+
components_.emplace(
31+
std::make_pair(i, std::make_pair(comps[i], weights[i])));
32+
}
2833
}
2934
}
3035

3136
template <typename Scalar>
32-
CostStackTpl<Scalar>::CostStackTpl(const CostPtr &cost)
33-
: CostStackTpl(cost->space, cost->nu, {cost}, {1.}) {}
37+
CostStackTpl<Scalar>::CostStackTpl(const PolyCost &cost)
38+
: CostBase(cost->space, cost->nu) {
39+
components_.emplace(0UL, std::make_pair(cost, 1.0));
40+
}
3441

3542
template <typename Scalar>
3643
bool CostStackTpl<Scalar>::checkDimension(
@@ -39,20 +46,17 @@ bool CostStackTpl<Scalar>::checkDimension(
3946
(comp->nu == this->nu);
4047
}
4148

42-
template <typename Scalar> std::size_t CostStackTpl<Scalar>::size() const {
43-
return components_.size();
44-
}
45-
4649
template <typename Scalar>
47-
void CostStackTpl<Scalar>::addCost(const CostPtr &cost, const Scalar weight) {
50+
auto CostStackTpl<Scalar>::addCost(const CostKey &key, const PolyCost &cost,
51+
const Scalar weight) -> CostItem & {
4852
if (!this->checkDimension(cost)) {
4953
ALIGATOR_DOMAIN_ERROR(fmt::format(
5054
"Cannot add new component due to inconsistent input dimensions "
5155
"(got ({:d}, {:d}), expected ({:d}, {:d}))",
5256
cost->ndx(), cost->nu, this->ndx(), this->nu));
5357
}
54-
components_.push_back(cost);
55-
weights_.push_back(weight);
58+
components_.emplace(key, std::make_pair(cost, weight));
59+
return components_.at(key);
5660
}
5761

5862
template <typename Scalar>
@@ -61,9 +65,9 @@ void CostStackTpl<Scalar>::evaluate(const ConstVectorRef &x,
6165
CostData &data) const {
6266
SumCostData &d = static_cast<SumCostData &>(data);
6367
d.value_ = 0.;
64-
for (std::size_t i = 0; i < components_.size(); i++) {
65-
components_[i]->evaluate(x, u, *d.sub_cost_data[i]);
66-
d.value_ += this->weights_[i] * d.sub_cost_data[i]->value_;
68+
for (const auto &[key, item] : components_) {
69+
item.first->evaluate(x, u, *d.sub_cost_data[key]);
70+
d.value_ += item.second * d.sub_cost_data[key]->value_;
6771
}
6872
}
6973

@@ -73,9 +77,9 @@ void CostStackTpl<Scalar>::computeGradients(const ConstVectorRef &x,
7377
CostData &data) const {
7478
SumCostData &d = static_cast<SumCostData &>(data);
7579
d.grad_.setZero();
76-
for (std::size_t i = 0; i < components_.size(); i++) {
77-
components_[i]->computeGradients(x, u, *d.sub_cost_data[i]);
78-
d.grad_.noalias() += this->weights_[i] * d.sub_cost_data[i]->grad_;
80+
for (const auto &[key, item] : components_) {
81+
item.first->computeGradients(x, u, *d.sub_cost_data[key]);
82+
d.grad_.noalias() += item.second * d.sub_cost_data[key]->grad_;
7983
}
8084
}
8185

@@ -85,9 +89,9 @@ void CostStackTpl<Scalar>::computeHessians(const ConstVectorRef &x,
8589
CostData &data) const {
8690
SumCostData &d = static_cast<SumCostData &>(data);
8791
d.hess_.setZero();
88-
for (std::size_t i = 0; i < components_.size(); i++) {
89-
components_[i]->computeHessians(x, u, *d.sub_cost_data[i]);
90-
d.hess_.noalias() += this->weights_[i] * d.sub_cost_data[i]->hess_;
92+
for (const auto &[key, item] : components_) {
93+
item.first->computeHessians(x, u, *d.sub_cost_data[key]);
94+
d.hess_.noalias() += item.second * d.sub_cost_data[key]->hess_;
9195
}
9296
}
9397

@@ -102,8 +106,8 @@ CostStackTpl<Scalar>::createData() const {
102106
template <typename Scalar>
103107
CostStackDataTpl<Scalar>::CostStackDataTpl(const CostStackTpl<Scalar> &obj)
104108
: CostData(obj.ndx(), obj.nu) {
105-
for (std::size_t i = 0; i < obj.size(); i++) {
106-
sub_cost_data.push_back(obj.components_[i]->createData());
109+
for (const auto &[key, item] : obj.components_) {
110+
sub_cost_data[key] = item.first->createData();
107111
}
108112
}
109113

0 commit comments

Comments
 (0)