chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mju_mulVecMatVec, mutiplies a square matrix M by a vector x on both sides. Returns x^T * M * x.

Browse Source 
PiperOrigin-RevId: 474292806
Change-Id: I3432469dbe1f02ccf5a13241c7aa12d824cbe034
This commit is contained in:
Yuval Tassa
2022-09-14 08:04:40 -07:00
committed by Copybara-Service
parent ee9eccc992
commit f4e7fa97af
12 changed files with 109 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference.rst
+11
View File
@@ -6040,6 +6040,17 @@ mju_mulMatTVec
Multiply transposed matrix and vector: res = mat' \* vec.
.. _mju_mulVecMatVec:
mju_mulVecMatVec
~~~~~~~~~~~~~~~~
.. code-block:: C
mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
Multiply square matrix with vectors on both sides: return vec1' \* mat \* vec2.
mju_transpose
~~~~~~~~~~~~~
doc/changelog.rst
+2
View File
@@ -25,6 +25,8 @@ General
- The algorithm, introduced in `Tassa et al. 2014 <https://doi.org/10.1109/ICRA.2014.6907001>`_,
converges after 2-5 Cholesky factorisations, independent of problem size.
- Added :ref:`mju_mulVecMatVec` to multiply a square matrix :math:`M` with vectors :math:`x` and :math:`y` on both
sides. The function returns :math:`x^TMy`.
Version 2.2.2 (September 7, 2022)
---------------------------------
include/mujoco/mujoco.h
+3
View File
@@ -911,6 +911,9 @@ MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// Multiply transposed matrix and vector: res = mat' * vec.
MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc);
// Multiply square matrix with vectors on both sides: returns vec1'*mat*vec2.
MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
// Transpose matrix: res = mat'.
MJAPI void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc);
introspect/functions.py
+30
View File
@@ -5644,6 +5644,36 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc="Multiply transposed matrix and vector: res = mat' * vec.",
)),
('mju_mulVecMatVec',
FunctionDecl(
name='mju_mulVecMatVec',
return_type=ValueType(name='mjtNum'),
parameters=(
FunctionParameterDecl(
name='vec1',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='vec2',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='n',
type=ValueType(name='int'),
),
),
doc="Multiply square matrix with vectors on both sides: returns vec1'*mat*vec2.", # pylint: disable=line-too-long
)),
('mju_transpose',
FunctionDecl(
name='mju_transpose',
python/mujoco/bindings_test.py
+6
View File
@@ -1006,6 +1006,12 @@ Euler integrator, semi-implicit in velocity.
rank = mujoco.mju_boxQP(res, r, index, h, g, lower, upper)
self.assertGreater(rank, -1)
def test_mju_mul_vec_mat_vec(self):
vec1 = np.array([1., 2., 3.])
vec2 = np.array([3., 2., 1.])
mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)
@parameterized.product(flg_html=(False, True), flg_pad=(False, True))
def test_mj_printSchema(self, flg_html, flg_pad): # pylint: disable=invalid-name
# Make sure that mj_printSchema doesn't raise an exception
python/mujoco/functions.cc
+20
View File
@@ -784,6 +784,26 @@ PYBIND11_MODULE(_functions, pymodule) {
return InterceptMjErrors(::mju_mulMatTVec)(
res.data(), mat.data(), vec.data(), mat.rows(), mat.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulVecMatVec, "n")(
pymodule,
[](Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenVectorX> vec2) {
if (vec1.size() != vec2.size()) {
throw py::type_error(
"size of vec1 should equal the size of vec2");
}
if (vec1.size() != mat.cols()) {
throw py::type_error(
"size of vectors should equal the number of columns in mat");
}
if (vec1.size() != mat.rows()) {
throw py::type_error(
"size of vectors should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_mulVecMatVec)(
vec1.data(), mat.data(), vec2.data(), vec1.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_transpose, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
src/engine/engine_solver.c
+2 -3
View File
@@ -190,7 +190,7 @@ static void residual(const mjModel* m, mjData* d, mjtNum* res, int i, int dim, i
// compute cost change
static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
const mjtNum* res, int dim) {
mjtNum delta[6], v[6], change;
mjtNum delta[6], change;
// compute change
if (dim==1) {
@@ -198,8 +198,7 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
change = 0.5*delta[0]*delta[0]*A[0] + delta[0]*res[0];
} else {
mju_sub(delta, force, oldforce, dim);
mju_mulMatVec(v, A, delta, dim, dim);
change = 0.5*mju_dot(delta, v, dim) + mju_dot(delta, res, dim);
change = 0.5*mju_mulVecMatVec(delta, A, delta, dim) + mju_dot(delta, res, dim);
}
// positive change: restore
src/engine/engine_util_blas.c
+13 -4
View File
@@ -685,8 +685,7 @@ mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, const int n) {
//------------------------------ matrix-vector operations ------------------------------------------
// multiply matrix and vector
void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int nr, int nc) {
void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
for (int r=0; r<nr; r++) {
res[r] = mju_dot(mat + r*nc, vec, nc);
}
@@ -695,8 +694,7 @@ void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
// multiply transposed matrix and vector
void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int nr, int nc) {
void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
mjtNum tmp;
mju_zero(res, nc);
@@ -709,6 +707,17 @@ void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
// multiply square matrix with vectors on both sides: return vec1'*mat*vec2
mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n) {
mjtNum res = 0;
for (int i=0; i<n; i++) {
res += vec1[i] * mju_dot(mat + i*n, vec2, n);
}
return res;
}
//------------------------------ matrix-matrix operations ------------------------------------------
// transpose matrix
src/engine/engine_util_blas.h
+3
View File
@@ -180,6 +180,9 @@ MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int nr, int nc);
// multiply square matrix with vectors on both sides: return vec1'*mat*vec2
MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
//------------------------------ matrix-matrix operations ------------------------------------------
src/engine/engine_util_solve.c
+2 -4
View File
@@ -934,8 +934,7 @@ int mju_boxQPoption(mjtNum* res, mjtNum* R, int* index, // outputs
}
// compute objective: value = 0.5*res'*H*res + res'*g
mju_mulMatVec(temp, H, res, n, n); // TODO(b/246267542): do this in one call
value = 0.5 * mju_dot(res, temp, n) + mju_dot(res, g, n);
value = 0.5 * mju_mulVecMatVec(res, H, res, n) + mju_dot(res, g, n);
// save last value
oldvalue = value;
@@ -1056,8 +1055,7 @@ int mju_boxQPoption(mjtNum* res, mjtNum* R, int* index, // outputs
}
// new objective value
mju_mulMatVec(temp, H, candidate, n, n);
value = 0.5 * mju_dot(candidate, temp, n) + mju_dot(candidate, g, n);
value = 0.5 * mju_mulVecMatVec(candidate, H, candidate, n) + mju_dot(candidate, g, n);
// increment and break if step is too small
nstep++;
test/engine/engine_util_blas_test.cc
+12
View File
@@ -41,5 +41,17 @@ TEST_F(EngineUtilBlasTest, MjuDot) {
EXPECT_EQ(mju_dot(a, b, 7), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
}
TEST_F(EngineUtilBlasTest, MjuMulVecMatVec) {
mjtNum vec1[] = {1, 2, 3};
mjtNum vec2[] = {3, 2, 1};
mjtNum mat[] = {
1, 2, 3,
4, 5, 6,
7, 8, 9
};
EXPECT_EQ(mju_mulVecMatVec(vec1, mat, vec2, 3), 204);
}
} // namespace
} // namespace mujoco
test/engine/engine_util_solve_test.cc
+5 -9
View File
@@ -75,10 +75,8 @@ TEST_F(QCQP3Test, DegenerateAMatrix) {
using BoxQPTest = MujocoTest;
// utility: compute QP objective = 0.5*x'*H*x + x'*g
mjtNum objective(const mjtNum* x, const mjtNum* H, const mjtNum* g, int n,
mjtNum* temp) {
mju_mulMatVec(temp, H, x, n, n);
return 0.5 * mju_dot(x, temp, n) + mju_dot(x, g, n);
mjtNum objective(const mjtNum* x, const mjtNum* H, const mjtNum* g, int n) {
return 0.5 * mju_mulVecMatVec(x, H, x, n) + mju_dot(x, g, n);
}
// utility: test if res is the minimum of a given box-QP problem
@@ -86,11 +84,10 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
const mjtNum* lower, const mjtNum* upper) {
static const mjtNum eps = 1e-4; // epsilon used for nudging
bool is_minimum = true;
mjtNum* temp = (mjtNum*) mju_malloc(sizeof(mjtNum)*n);
mjtNum* res_nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*n);
// get solution value
mjtNum value = objective(res, H, g, n, temp);
mjtNum value = objective(res, H, g, n);
mjtNum value_nudge;
// compare to nudged solution
@@ -102,7 +99,7 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
if (lower) {
res_nudge[i] = mju_max(lower[i], res_nudge[i]);
}
value_nudge = objective(res_nudge, H, g, n, temp);
value_nudge = objective(res_nudge, H, g, n);
if (value_nudge - value < 0) {
is_minimum = false;
break;
@@ -113,7 +110,7 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
if (upper) {
res_nudge[i] = mju_min(upper[i], res_nudge[i]);
}
value_nudge = objective(res_nudge, H, g, n, temp);
value_nudge = objective(res_nudge, H, g, n);
if (value_nudge - value < 0) {
is_minimum = false;
break;
@@ -124,7 +121,6 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
}
mju_free(res_nudge);
mju_free(temp);
return is_minimum;
}