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

Add utility functions for banded-then-dense symmetric matrices.

Browse Source 
PiperOrigin-RevId: 528757637
Change-Id: I916d843140322c28e857100a6b305a041f704d53
This commit is contained in:
Yuval Tassa
2023-05-02 05:57:22 -07:00
committed by Copybara-Service
parent 2ad82d5998
commit b0bc330b54
13 changed files with 1147 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+96 -1
View File
@@ -2748,7 +2748,7 @@ mju_cholSolve
.. mujoco-include:: mju_cholSolve
Solve mat * res = vec, where mat is Cholesky-factorized
Solve (mat*mat') * res = vec, where mat is a Cholesky factor.
.. _mju_cholUpdate:
@@ -2759,6 +2759,101 @@ mju_cholUpdate
Cholesky rank-one update: L*L' +/- x*x'; return rank.
.. _mju_cholFactorBand:
mju_cholFactorBand
~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_cholFactorBand
Band-dense Cholesky decomposition.
|br| Add ``diagadd + diagmul*mat_ii`` to diagonal before decomposition.
|br| Returns the minimum value of the factorized diagonal or 0 if rank-defficient.
**Symmetric band-dense matrices**
:ref:`mju_cholFactorBand` and subsequent functions containing the substring "band" operate on matrices which are a
generalization of symmetric `band matrices <https://en.wikipedia.org/wiki/Band_matrix>`_. *Symmetric band-dense* or
"arrowhead" matrices have non-zeros along proximal diagonal bands and dense blocks on the bottom rows and right
columns. These matrices have the property that Cholesky factorization creates no fill-in and can therefore be
performed efficiently in-place. Matrix structure is defined by three integers:
- ``ntotal``: the number of rows (columns) of the symmetric matrix.
- ``nband``: the number of bands under (over) the diagonal, inclusive of the diagonal.
- ``ndense``: the number of dense rows (columns) at the bottom (right).
The non-zeros are stored in memory as two contiguous row-major blocks, colored green and blue in the illustration
below. The first block has size ``nband x (ntotal-ndense)`` and contains the diagonal and the bands below it. The
second block has size ``ndense x ntotal`` and contains the dense part. Total required memory is the sum of the block
sizes.
.. figure:: /images/APIreference/arrowhead.svg
:width: 750px
:align: left
For example, consider an arrowhead matrix with ``nband = 3``, ``ndense = 2`` and ``ntotal = 8``. In this example, the
total memory required is ``3*(8-2) + 2*8 = 34`` mjtNum's, laid out as follows:
.. code-block::
0 1 2
3 4 5
6 7 8
9 10 11
12 13 14
15 16 17
18 19 20 21 22 23 24 25
26 27 28 29 30 31 32 33
The diagonal elements are ``2, 5, 8, 11, 14, 17, 24, 33``.
|br| Elements ``0, 1, 3, 25`` are present in memory but never touched.
.. _mju_cholSolveBand:
mju_cholSolveBand
~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_cholSolveBand
Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor.
.. _mju_band2Dense:
mju_band2Dense
~~~~~~~~~~~~~~
.. mujoco-include:: mju_band2Dense
Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.
.. _mju_dense2Band:
mju_dense2Band
~~~~~~~~~~~~~~
.. mujoco-include:: mju_dense2Band
Convert dense matrix to banded matrix.
.. _mju_bandMulMatVec:
mju_bandMulMatVec
~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_bandMulMatVec
Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.
.. _mju_bandDiag:
mju_bandDiag
~~~~~~~~~~~~
.. mujoco-include:: mju_bandDiag
Address of diagonal element i in band-dense matrix representation.
.. _mju_eig3:
mju_eig3
doc/APIreference/functions_override.rst
+45
View File
@@ -321,6 +321,51 @@ mju_ceil
.. _Decompositions:
.. _mju_cholFactorBand:
Band-dense Cholesky decomposition.
|br| Add ``diagadd + diagmul*mat_ii`` to diagonal before decomposition.
|br| Returns the minimum value of the factorized diagonal or 0 if rank-defficient.
**Symmetric band-dense matrices**
:ref:`mju_cholFactorBand` and subsequent functions containing the substring "band" operate on matrices which are a
generalization of symmetric `band matrices <https://en.wikipedia.org/wiki/Band_matrix>`_. *Symmetric band-dense* or
"arrowhead" matrices have non-zeros along proximal diagonal bands and dense blocks on the bottom rows and right
columns. These matrices have the property that Cholesky factorization creates no fill-in and can therefore be
performed efficiently in-place. Matrix structure is defined by three integers:
- ``ntotal``: the number of rows (columns) of the symmetric matrix.
- ``nband``: the number of bands under (over) the diagonal, inclusive of the diagonal.
- ``ndense``: the number of dense rows (columns) at the bottom (right).
The non-zeros are stored in memory as two contiguous row-major blocks, colored green and blue in the illustration
below. The first block has size ``nband x (ntotal-ndense)`` and contains the diagonal and the bands below it. The
second block has size ``ndense x ntotal`` and contains the dense part. Total required memory is the sum of the block
sizes.
.. figure:: /images/APIreference/arrowhead.svg
:width: 750px
:align: left
For example, consider an arrowhead matrix with ``nband = 3``, ``ndense = 2`` and ``ntotal = 8``. In this example, the
total memory required is ``3*(8-2) + 2*8 = 34`` mjtNum's, laid out as follows:
.. code-block::
0 1 2
3 4 5
6 7 8
9 10 11
12 13 14
15 16 17
18 19 20 21 22 23 24 25
26 27 28 29 30 31 32 33
The diagonal elements are ``2, 5, 8, 11, 14, 17, 24, 33``.
|br| Elements ``0, 1, 3, 25`` are present in memory but never touched.
.. _mju_boxQP:
Minimize :math:`\tfrac{1}{2} x^T H x + x^T g \quad \text{s.t.} \quad l \le x \le u`, return rank or -1 if failed.
doc/changelog.rst
+13 -5
View File
@@ -5,18 +5,26 @@ Changelog
Upcoming version (not yet released)
-----------------------------------
Plugins
^^^^^^^
.. youtube:: hqIMTNGaLF4
:align: right
:width: 240px
Plugins
^^^^^^^
- Added touch-grid sensor plugin. See `documentation <https://github.com/deepmind/mujoco/blob/main/plugin/sensor/README.md>`_
for details, and associated `touch_grid.xml <https://github.com/deepmind/mujoco/blob/main/model/plugin/touch_grid.xml>`_
example model. The plugin includes `in-scene visualisation <https://youtu.be/0LOJ3WMnqeA>`_.
- Add ``mj_multiRay`` function for intersecting multiple rays emanating from a single point. This is significantly
faster than calling ``mj_ray`` multiple times.
General
^^^^^^^
- Added :ref:`mjd_inverseFD` for finite-differenced inverse-dynamics derivatives.
- Added functions for operations on banded-then-dense "arrowhead" matrices. Such matrices are
common when doing direct trajectory optimization.
See :ref:`mju_cholFactorBand` documentation for details.
- Added :ref:`mj_multiRay` function for intersecting multiple rays emanating from a single point.
This is significantly faster than calling :ref:`mj_ray` multiple times.
Version 2.3.5 (April 25, 2023)
------------------------------
doc/images/APIreference/arrowhead.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 17 KiB

doc/includes/references.h
+10
View File
@@ -2408,6 +2408,16 @@ void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[4],
int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);
int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense,
mjtNum diagadd, mjtNum diagmul);
void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense);
void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
mjtByte flg_sym);
void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
int mju_bandDiag(int i, int ntotal, int nband, int ndense);
int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
int mju_boxQP(mjtNum* res, mjtNum* R, int* index, const mjtNum* H, const mjtNum* g, int n,
const mjtNum* lower, const mjtNum* upper);
include/mujoco/mujoco.h
+28 -1
View File
@@ -1068,12 +1068,39 @@ MJAPI void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[
// Cholesky decomposition: mat = L*L'; return rank, decomposition performed in-place into mat.
MJAPI int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
// Solve mat * res = vec, where mat is Cholesky-factorized
// Solve (mat*mat') * res = vec, where mat is a Cholesky factor.
MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);
// Cholesky rank-one update: L*L' +/- x*x'; return rank.
MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
// Band-dense Cholesky decomposition.
// Returns minimum value in the factorized diagonal, or 0 if rank-deficient.
// mat has (ntotal-ndense) x nband + ndense x ntotal elements.
// The first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive.
// The second ndense x ntotal store the band part as entire dense rows.
// Add diagadd+diagmul*mat_ii to diagonal before factorization.
MJAPI mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense,
mjtNum diagadd, mjtNum diagmul);
// Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor.
MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense);
// Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.
MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
mjtByte flg_sym);
// Convert dense matrix to banded matrix.
MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
// Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.
MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
// Address of diagonal element i in band-dense matrix representation.
MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
// Eigenvalue decomposition of symmetric 3x3 matrix.
MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
introspect/functions.py
+211 -1
View File
@@ -6815,7 +6815,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
type=ValueType(name='int'),
),
),
doc='Solve mat * res = vec, where mat is Cholesky-factorized',
doc="Solve (mat*mat') * res = vec, where mat is a Cholesky factor.",
)),
('mju_cholUpdate',
FunctionDecl(
@@ -6845,6 +6845,216 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc="Cholesky rank-one update: L*L' +/- x*x'; return rank.",
)),
('mju_cholFactorBand',
FunctionDecl(
name='mju_cholFactorBand',
return_type=ValueType(name='mjtNum'),
parameters=(
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='diagadd',
type=ValueType(name='mjtNum'),
),
FunctionParameterDecl(
name='diagmul',
type=ValueType(name='mjtNum'),
),
),
doc='Band-dense Cholesky decomposition. Returns minimum value in the factorized diagonal, or 0 if rank-deficient. mat has (ntotal-ndense) x nband + ndense x ntotal elements. The first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive. The second ndense x ntotal store the band part as entire dense rows. Add diagadd+diagmul*mat_ii to diagonal before factorization.', # pylint: disable=line-too-long
)),
('mju_cholSolveBand',
FunctionDecl(
name='mju_cholSolveBand',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='vec',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
),
doc="Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor.", # pylint: disable=line-too-long
)),
('mju_band2Dense',
FunctionDecl(
name='mju_band2Dense',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='flg_sym',
type=ValueType(name='mjtByte'),
),
),
doc='Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.', # pylint: disable=line-too-long
)),
('mju_dense2Band',
FunctionDecl(
name='mju_dense2Band',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
),
doc='Convert dense matrix to banded matrix.',
)),
('mju_bandMulMatVec',
FunctionDecl(
name='mju_bandMulMatVec',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='vec',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nvec',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='flg_sym',
type=ValueType(name='mjtByte'),
),
),
doc='Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.', # pylint: disable=line-too-long
)),
('mju_bandDiag',
FunctionDecl(
name='mju_bandDiag',
return_type=ValueType(name='int'),
parameters=(
FunctionParameterDecl(
name='i',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ntotal',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nband',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='ndense',
type=ValueType(name='int'),
),
),
doc='Address of diagonal element i in band-dense matrix representation.', # pylint: disable=line-too-long
)),
('mju_eig3',
FunctionDecl(
name='mju_eig3',
python/mujoco/bindings_test.py
+30
View File
@@ -1075,6 +1075,36 @@ Euler integrator, semi-implicit in velocity.
self.assertGreater(np.linalg.norm(ds_dv), eps)
self.assertGreater(np.linalg.norm(ds_da), eps)
def test_banded(self):
n_total = 4
n_band = 1
n_dense = 1
dense = np.array([[1.0, 0, 0, 0.1],
[0, 2.0, 0, 0.2],
[0, 0, 3.0, 0.3],
[0.1, 0.2, 0.3, 4.0]])
band = np.zeros(n_band*(n_total-n_dense) + n_dense*n_total)
mujoco.mju_dense2Band(band, dense, n_total, n_band, n_dense)
for i in range(4):
index = mujoco.mju_bandDiag(i, n_total, n_band, n_dense)
self.assertEqual(band[index], i+1)
dense2 = np.zeros((n_total, n_total))
flg_sym = 1
mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym)
np.testing.assert_array_equal(dense, dense2)
vec = np.array([[2.0], [2.0], [3.0], [4.0]])
res = np.zeros_like(vec)
n_vec = 1
mujoco.mju_bandMulMatVec(res, band, vec,
n_total, n_band, n_dense, n_vec, flg_sym)
np.testing.assert_array_equal(res, dense @ vec)
diag_add = 0
diag_mul = 0
mujoco.mju_cholFactorBand(band, n_total, n_band, n_dense,
diag_add, diag_mul)
mujoco.mju_cholSolveBand(res, band, vec, n_total, n_band, n_dense)
np.testing.assert_almost_equal(res, np.linalg.solve(dense, vec))
def test_mju_box_qp(self):
n = 5
res = np.zeros(n)
python/mujoco/functions.cc
+93
View File
@@ -991,6 +991,99 @@ PYBIND11_MODULE(_functions, pymodule) {
return InterceptMjErrors(::mju_cholUpdate)(
mat.data(), x.data(), mat.rows(), flg_plus);
});
Def<traits::mju_cholFactorBand>(
pymodule, [](Eigen::Ref<EigenVectorX> mat, int ntotal, int nband,
int ndense, mjtNum diagadd, mjtNum diagmul) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
return InterceptMjErrors(::mju_cholFactorBand)(
mat.data(), ntotal, nband, ndense, diagadd, diagmul);
});
Def<traits::mju_cholSolveBand>(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorX> vec, int ntotal, int nband,
int ndense) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have (ntotal-ndense)*nband + "
"ndense*ntotal elements");
}
if (res.size() != ntotal) {
throw py::type_error("size of res should equal ntotal");
}
if (vec.size() != ntotal) {
throw py::type_error("size of vec should equal ntotal");
}
return InterceptMjErrors(::mju_cholSolveBand)(
res.data(), mat.data(), vec.data(), ntotal, nband, ndense);
});
Def<traits::mju_band2Dense>(
pymodule,
[](Eigen::Ref<EigenArrayXX> res, Eigen::Ref<const EigenVectorX> mat,
int ntotal, int nband, int ndense, mjtByte flg_sym) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (res.rows() != ntotal) {
throw py::type_error("res should have ntotal rows");
}
if (res.cols() != ntotal) {
throw py::type_error("res should have ntotal columns");
}
return InterceptMjErrors(::mju_band2Dense)(
res.data(), mat.data(), ntotal, nband, ndense, flg_sym);
});
Def<traits::mju_dense2Band>(pymodule, [](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenArrayXX> mat,
int ntotal, int nband, int ndense) {
int nRes = (ntotal - ndense) * nband + ndense * ntotal;
if (res.size() != nRes) {
throw py::type_error(
"res must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (mat.rows() != ntotal) {
throw py::type_error("mat should have ntotal rows");
}
if (mat.cols() != ntotal) {
throw py::type_error("mat should have ntotal columns");
}
return InterceptMjErrors(::mju_dense2Band)(res.data(), mat.data(), ntotal,
nband, ndense);
});
Def<traits::mju_bandMulMatVec>(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenArrayXX> vec, int ntotal, int nband, int ndense,
int nVec, mjtByte flg_sym) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (res.rows() != ntotal) {
throw py::type_error("res should have ntotal rows");
}
if (res.cols() != nVec) {
throw py::type_error("res should have nVec columns");
}
if (vec.rows() != ntotal) {
throw py::type_error("vec should have ntotal rows");
}
if (vec.cols() != nVec) {
throw py::type_error("vec should have nVec columns");
}
return InterceptMjErrors(::mju_bandMulMatVec)(res.data(), mat.data(),
vec.data(), ntotal, nband,
ndense, nVec, flg_sym);
});
Def<traits::mju_bandDiag>(pymodule);
Def<traits::mju_eig3>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_boxQP, "n")(
pymodule,
src/engine/engine_util_solve.c
+282
View File
@@ -300,6 +300,288 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
return rank;
}
//---------------------------- banded Cholesky -----------------------------------------------------
// band-dense Cholesky decomposition
// returns minimum value in the factorized diagonal, or 0 if rank-deficient
// mat has (ntotal-ndense) x nband + ndense x ntotal elements
// the first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive
// the second ndense x ntotal store the band part as entire dense rows
// add diagadd+diagmul*mat_ii to diagonal before factorization
mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense,
mjtNum diagadd, mjtNum diagmul) {
int nsparse = ntotal - ndense;
mjtNum mindiag = -1;
// sparse part, including sparse-sparse and sparse-dense
for (int j=0; j<nsparse; j++) {
// number of non-zeros left of (j,j)
int width_jj = mjMIN(j, nband-1);
// number of non-zeros below (j,j), sparse part
int height = mjMIN(nsparse-j-1, nband-1);
// address of (j,j)
int adr_jj = (j+1)*nband-1;
// compute L(j,j), before sqrt
mjtNum left_ij = width_jj>0 ? mju_dot(mat+adr_jj-width_jj, mat+adr_jj-width_jj, width_jj) : 0;
mjtNum Ljj = diagadd + diagmul*mat[adr_jj] + mat[adr_jj] - left_ij;
// update mindiag
if (Ljj<mindiag || mindiag<0) {
mindiag = Ljj;
}
// stop if rank-deficient
if (Ljj<mjMINVAL) {
return 0;
}
// compute Ljj, scale = 1/Ljj
Ljj = mju_sqrt(Ljj);
mjtNum scale = 1/Ljj;
// compute L(i,j) for i>j, sparse part
for (int i=j+1; i<=j+height; i++) {
// number of non-zeros left of (i,j)
int width_ij = mjMIN(j, nband-1-i+j);
// address of (i,j)
int adr_ij = (i+1)*nband-1-i+j;
// in-place computation of L(i,j)
left_ij = width_ij>0 ? mju_dot(mat+adr_jj-width_ij, mat+adr_ij-width_ij, width_ij) : 0;
mat[adr_ij] = scale * (mat[adr_ij] - left_ij);
}
// compute L(i,j) for i>j, dense part
for (int i=nsparse; i<ntotal; i++) {
// address of (i,j)
int adr_ij = nsparse*nband + (i-nsparse)*ntotal + j;
// in-place computation of L(i,j)
// number of non-zeros left of (i,j) now equals width_jj
left_ij = width_jj>0 ? mju_dot(mat+adr_jj-width_jj, mat+adr_ij-width_jj, width_jj) : 0;
mat[adr_ij] = scale * (mat[adr_ij] - left_ij);
}
// save L(j,j)
mat[adr_jj] = Ljj;
}
// dense part
for (int j=nsparse; j<ntotal; j++) {
// address of (j,j)
int adr_jj = nsparse*nband + (j-nsparse)*ntotal + j;
// compute Ljj
mjtNum Ljj = diagadd + diagmul*mat[adr_jj] + mat[adr_jj] -
mju_dot(mat+adr_jj-j, mat+adr_jj-j, j);
// update mindiag
if (Ljj<mindiag || mindiag<0) {
mindiag = Ljj;
}
// stop if rank-deficient
if (Ljj<mjMINVAL) {
return 0;
}
// compute Ljj, scale = 1/Ljj
Ljj = mju_sqrt(Ljj);
mjtNum scale = 1/Ljj;
// compute L(i,j) for i>j
for (int i=j+1; i<ntotal; i++) {
// address of off-diagonal element
int adr_ij = adr_jj + ntotal*(i-j);
// in-place computation of L(i,j)
mat[adr_ij] = scale * (mat[adr_ij] - mju_dot(mat+adr_jj-j, mat+adr_ij-j, j));
}
// save L(j,j)
mat[adr_jj] = Ljj;
}
return mindiag;
}
// solve with band-Cholesky decomposition
void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense) {
int width, height, nsparse = ntotal - ndense;
// copy into result if different
if (res!=vec) {
mju_copy(res, vec, ntotal);
}
//------- forward substitution: solve L*res = vec
// sparse part
for (int i=0; i<nsparse; i++) {
// number of non-zeros left of (i,i)
width = mjMIN(i, nband-1);
if (width) {
res[i] -= mju_dot(mat+(i+1)*nband-1-width, res+i-width, width);
}
// diagonal
res[i] /= mat[(i+1)*nband-1];
}
// dense part
for (int i=nsparse; i<ntotal; i++) {
res[i] -= mju_dot(mat+nsparse*nband+(i-nsparse)*ntotal, res, i);
// diagonal
res[i] /= mat[nsparse*nband+(i-nsparse)*ntotal+i];
}
//------- backward substitution: solve L'*res = res
// dense part
for (int i=ntotal-1; i>=nsparse; i--) {
for (int j=i+1; j<ntotal; j++) {
res[i] -= mat[nsparse*nband+(j-nsparse)*ntotal+i] * res[j];
}
// diagonal
res[i] /= mat[nsparse*nband+(i-nsparse)*ntotal+i];
}
// sparse part
for (int i=nsparse-1; i>=0; i--) {
// number of non-zeros below (i,i), sparse part
height = mjMIN(nsparse-1-i, nband-1);
// sparse rows
for (int j=i+1; j<=i+height; j++)
res[i] -= mat[(j+1)*nband-1-(j-i)] * res[j];
// dense rows
for (int j=nsparse; j<ntotal; j++)
res[i] -= mat[nsparse*nband+(j-nsparse)*ntotal+i] * res[j];
// diagonal
res[i] /= mat[(i+1)*nband-1];
}
}
// address of diagonal element i in band-dense matrix representation
int mju_bandDiag(int i, int ntotal, int nband, int ndense) {
int nsparse = ntotal-ndense;
// sparse part
if (i<nsparse) {
return i*nband + nband-1;
}
// dense part
else {
return nsparse*nband + (i-nsparse)*ntotal + i;
}
}
// convert band matrix to dense matrix
void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
mjtByte flg_sym) {
int nsparse = ntotal-ndense;
// clear all
mju_zero(res, ntotal*ntotal);
// sparse part
for(int i=0; i<nsparse; i++) {
// number of non-zeros left of (i,i)
int width = mjMIN(i, nband-1);
// copy data
mju_copy(res + i*ntotal + i-width, mat + (i+1)*nband - (width+1), width+1);
}
// dense part
for(int i=nsparse; i<ntotal; i++) {
mju_copy(res + i*ntotal, mat + nsparse*nband + (i-nsparse)*ntotal, i+1);
}
// make symmetric
if (flg_sym) {
for(int i=0; i<ntotal; i++) {
for (int j=i+1; j<ntotal; j++) {
res[i*ntotal + j] = res[j*ntotal + i];
}
}
}
}
// convert dense matrix to band matrix
void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense) {
int nsparse = ntotal-ndense;
// sparse part
for(int i=0; i<nsparse; i++) {
// number of non-zeros left of (i,i)
int width = mjMIN(i, nband-1);
// copy data
mju_copy(res + (i+1)*nband - (width+1), mat + i*ntotal + i-width, width+1);
}
// dense part
for(int i=nsparse; i<ntotal; i++) {
mju_copy(res + nsparse*nband + (i-nsparse)*ntotal, mat + i*ntotal, i+1);
}
}
// multiply band-diagonal matrix with vector
void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym) {
int nsparse = ntotal-ndense;
// handle multiple vectors
for(int j=0; j<nvec; j++ ) {
// precompute pointer to corresponding vector in vec and res
const mjtNum* vec_j = vec + ntotal*j;
mjtNum* res_j = res + ntotal*j;
// sparse part
for(int i=0; i<nsparse; i++) {
int width = mjMIN(i+1, nband);
int adr = i*nband + nband - width;
int offset = mjMAX(0, i-nband+1);
res_j[i] = mju_dot(mat+adr, vec_j+offset, width); // lower triangle
if (flg_sym) {
// strict upper triangle
mju_addToScl(res_j+offset, mat+adr, vec_j[i], width-1);
}
}
// dense part
for(int i=nsparse; i<ntotal; i++) {
int adr = nsparse*nband + (i-nsparse)*ntotal;
res_j[i] = mju_dot(mat+adr, vec_j, i+1);
if (flg_sym) {
// strict upper triangle
mju_addToScl(res_j, mat+adr, vec_j[i], i);
}
}
}
}
//------------------------------ LU factorization --------------------------------------------------
src/engine/engine_util_solve.h
+28 -2
View File
@@ -34,8 +34,7 @@ MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
// sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank'
// mat must have uncompressed layout; rownnz is modified to end at diagonal
int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
int* rownnz, int* rowadr, int* colind,
int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, int* rownnz, int* rowadr, int* colind,
mjData* d);
// sparse reverse-order Cholesky solve
@@ -48,6 +47,33 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind,
mjData* d);
// band-dense Cholesky decomposition
// returns minimum value in the factorized diagonal, or 0 if rank-deficient
// mat has (ntotal-ndense) x nband + ndense x ntotal elements
// the first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive
// the second ndense x ntotal store the band part as entire dense rows
// add diagadd+diagmul*mat_ii to diagonal before factorization
MJAPI mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense,
mjtNum diagadd, mjtNum diagmul);
// solve (mat*mat')*res = vec with band-Cholesky decomposition
MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense);
// convert banded matrix to dense matrix, fill upper triangle if flg_sym>0
MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
mjtByte flg_sym);
// convert dense matrix to banded matrix
MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
// multiply band-diagonal matrix with vector, include upper triangle if flg_sym>0
MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
// address of diagonal element i in band-dense matrix representation
MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
// sparse reverse-order LU factorization, no fill-in (assuming tree topology)
// LU = L + U; original = (U+I) * L; scratch is size n
void mju_factorLUSparse(mjtNum *LU, int n, int* scratch,
test/engine/engine_util_solve_test.cc
+292
View File
@@ -29,10 +29,16 @@ namespace mujoco {
namespace {
using ::testing::DoubleEq;
using ::testing::Pointwise;
using ::testing::DoubleNear;
using ::std::string;
using ::std::setw;
using QCQP2Test = MujocoTest;
std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(array, array + n);
}
TEST_F(QCQP2Test, DegenerateAMatrix) {
// A 2x2 matrix with determinant zero.
const mjtNum Ain[9] { 6, -15, 2, -5 };
@@ -347,5 +353,291 @@ TEST_F(BoxQPTest, BoundedQPvariations) {
mju_free(upper);
}
// ------------------------- band matrices -------------------------------------
using BandMatrixTest = MujocoTest;
// utility: random "arrowhead", banded-then-dense SPD matrix
// optional random vector and diagonal regularizer
void randomBanded(mjtNum* H, int nTotal, int nBand, int nDense, int seed,
mjtNum* vec = nullptr, mjtNum reg = 0) {
// make distribution using seed
std::mt19937_64 rng;
rng.seed(seed);
std::normal_distribution<double> dist(0, 1);
// allocate square root
mjtNum* sqrtH = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal*nTotal);
// sample
for (int i=0; i < nTotal; i++) {
if (vec) vec[i] = dist(rng);
for (int j=0; j < nTotal; j++) {
sqrtH[nTotal*i+j] = dist(rng);
}
}
// make SPD matrix H
mju_mulMatTMat(H, sqrtH, sqrtH, nTotal, nTotal, nTotal);
// set zeros
int nSparse = nTotal-nDense;
for (int i=0; i < nSparse; i++) {
int nzeros = mjMAX(0, i + 1 - nBand);
for (int j=0; j < nzeros; j++) {
H[nTotal*i + j] = 0;
H[nTotal*j + i] = 0;
}
}
// add regularizer to diagonal
for (int i=0; i < nTotal; i++) {
H[nTotal*i + i] += reg;
}
mju_free(sqrtH);
}
// test banded-vector diagonal values
TEST_F(BandMatrixTest, Diagonal) {
int seed = 1;
int nTotal = 8;
for (int nBand : {1, 3}) {
for (int nDense : {0, 2}) {
// allocate
int nB = (nTotal-nDense)*nBand + nDense*nTotal;
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB);
int nH = nTotal*nTotal;
mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
// make random banded SPD matrix, dense representation
randomBanded(H, nTotal, nBand, nDense, seed++);
// convert to banded representation
mju_dense2Band(B, H, nTotal, nBand, nDense);
// expect diagonals to be equal
for (int i=0; i < nTotal; i++) {
EXPECT_EQ(H[i*nTotal + i], B[mju_bandDiag(i, nTotal, nBand, nDense)]);
}
mju_free(H);
mju_free(B);
}
}
}
// test conversion of banded <-> dense
TEST_F(BandMatrixTest, Conversion) {
int seed = 1;
int nTotal = 8;
for (int nBand : {0, 1, 3}) {
for (int nDense : {0, 2}) {
// allocate
int nB = (nTotal-nDense)*nBand + nDense*nTotal;
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB);
int nH = nTotal*nTotal;
mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
// make random banded SPD matrix, dense
randomBanded(H, nTotal, nBand, nDense, seed++);
// convert to banded
mju_dense2Band(B, H, nTotal, nBand, nDense);
// convert back to dense
mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/1);
// expect exact equality
EXPECT_EQ(AsVector(H, nH), AsVector(H1, nH));
mju_free(H1);
mju_free(H);
mju_free(B);
}
}
}
// test banded-vector multiplication
TEST_F(BandMatrixTest, Multiplication) {
int seed = 1;
int nTotal = 8;
for (int nBand : {0, 1, 3}) {
for (int nDense : {0, 2}) {
// allocate
int nB = (nTotal-nDense)*nBand + nDense*nTotal;
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB);
int nH = nTotal*nTotal;
mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
// make random banded SPD matrix, dense
randomBanded(H, nTotal, nBand, nDense, seed++, vec);
// multiply dense
mju_mulMatVec(res, H, vec, nTotal, nTotal);
// convert to banded, multiply
mju_dense2Band(B, H, nTotal, nBand, nDense);
mju_bandMulMatVec(res1, B, vec, nTotal, nBand, nDense,
/*nVec=*/1, /*flg_sym=*/1);
// expect numerical equality
mjtNum eps = 1e-12;
EXPECT_THAT(AsVector(res, nTotal),
Pointwise(DoubleNear(eps), AsVector(res1, nTotal)));
mju_free(res1);
mju_free(res);
mju_free(vec);
mju_free(H);
mju_free(B);
}
}
}
// test banded factorization and vector product with factor
TEST_F(BandMatrixTest, Factorization) {
int seed = 1;
int nTotal = 8;
for (int nBand : {1, 3}) {
for (int nDense : {0, 2}) {
for (mjtNum diagadd : {0.0, 1.0}) {
for (int diagmul : {0.0, 1.3}) {
// allocate
int nB = (nTotal-nDense)*nBand + nDense*nTotal;
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB);
int nH = nTotal*nTotal;
mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
// make random banded matrix, dense representation
// add regularizer to ensure PD
randomBanded(H, nTotal, nBand, nDense, seed++, vec, /*reg=*/nTotal);
// convert to banded
mju_dense2Band(B, H, nTotal, nBand, nDense);
// apply diagadd and diagmul
for (int i=0; i < nTotal; i++) {
H[nTotal*i + i] += diagadd + diagmul*H[nTotal*i + i];
}
// in-place dense factorization
int rank = mju_cholFactor(H, nTotal, /*mindiag=*/0);
// expect factorization to have succeeded
EXPECT_EQ(rank, nTotal);
// banded factorization
mjtNum minDiag = mju_cholFactorBand(B, nTotal, nBand, nDense,
diagadd, diagmul);
// expect factorization to have succeeded
EXPECT_GT(minDiag, 0);
// convert back to dense, lower triangle only
mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/0);
// zero upper triangle of H (unused)
for (int i=0; i < nTotal-1; i++) {
mju_zero(H + nTotal*i + i + 1, nTotal - i - 1);
}
// expect numerical equality
mjtNum eps = 1e-12;
EXPECT_THAT(AsVector(H, nH),
Pointwise(DoubleNear(eps), AsVector(H1, nH)));
// multiply dense
mju_mulMatVec(res, H, vec, nTotal, nTotal);
// multiply sparse, only lower triangle
mju_bandMulMatVec(res1, B, vec, nTotal, nBand, nDense,
/*nVec=*/1, /*flg_sym=*/0);
// expect numerical equality
EXPECT_THAT(AsVector(res, nTotal),
Pointwise(DoubleNear(eps), AsVector(res1, nTotal)));
mju_free(res1);
mju_free(res);
mju_free(vec);
mju_free(H1);
mju_free(H);
mju_free(B);
}
}
}
}
}
// test banded solve
TEST_F(BandMatrixTest, Solve) {
int seed = 1;
int nTotal = 8;
for (int nBand : {1, 3}) {
for (int nDense : {0, 2}) {
// allocate
int nB = (nTotal-nDense)*nBand + nDense*nTotal;
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB);
int nH = nTotal*nTotal;
mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH);
mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal);
// make random banded matrix, dense representation
// add regularizer to ensure PD
randomBanded(H, nTotal, nBand, nDense, seed++, vec, /*reg=*/nTotal);
// convert to banded
mju_dense2Band(B, H, nTotal, nBand, nDense);
// in-place dense factorization
int rank = mju_cholFactor(H, nTotal, /*mindiag=*/0);
// expect factorization to have succeeded
EXPECT_EQ(rank, nTotal);
// banded factorization
mjtNum minDiag = mju_cholFactorBand(B, nTotal, nBand, nDense,
/*diagadd=*/0, /*diagmul=*/0);
// expect factorization to have succeeded
EXPECT_GT(minDiag, 0);
// convert back to dense, lower triangle only
mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/0);
// solve with dense
mju_cholSolve(res, H, vec, nTotal);
// solve with banded
mju_cholSolveBand(res1, B, vec, nTotal, nBand, nDense);
// expect numerical equality
mjtNum eps = 1e-12;
EXPECT_THAT(AsVector(res, nTotal),
Pointwise(DoubleNear(eps), AsVector(res1, nTotal)));
mju_free(res1);
mju_free(res);
mju_free(vec);
mju_free(H1);
mju_free(H);
mju_free(B);
}
}
}
} // namespace
} // namespace mujoco
unity/Runtime/Bindings/MjBindings.cs
+18
View File
@@ -3733,6 +3733,24 @@ public static unsafe extern void mju_cholSolve(double* res, double* mat, double*
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mju_cholUpdate(double* mat, double* x, int n, int flg_plus);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern double mju_cholFactorBand(double* mat, int ntotal, int nband, int ndense, double diagadd, double diagmul);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_cholSolveBand(double* res, double* mat, double* vec, int ntotal, int nband, int ndense);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_band2Dense(double* res, double* mat, int ntotal, int nband, int ndense, byte flg_sym);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_dense2Band(double* res, double* mat, int ntotal, int nband, int ndense);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_bandMulMatVec(double* res, double* mat, double* vec, int ntotal, int nband, int ndense, int nvec, byte flg_sym);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mju_bandDiag(int i, int ntotal, int nband, int ndense);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mju_eig3(double* eigval, double* eigvec, double* quat, double* mat);