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

Add mjData.L, sparse Cholesky factor of Newton solver Hessian.

Browse Source 
PiperOrigin-RevId: 683669555
Change-Id: I18d7530b5d68d1d48cb95d4faa0150c434109638
This commit is contained in:
Yuval Tassa
2024-10-08 10:12:10 -07:00
committed by Copybara-Service
parent 39af655c1e
commit b66da3d369
9 changed files with 99 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+10 -2
View File
@@ -164,12 +164,13 @@ struct mjData_ {
mjTimerStat timer[mjNTIMER]; // timer statistics
// variable sizes
int ncon; // number of detected contacts
int ne; // number of equality constraints
int nf; // number of friction constraints
int nl; // number of limit constraints
int nefc; // number of constraints
int nnzJ; // number of non-zeros in constraint Jacobian
int ncon; // number of detected contacts
int nnzL; // number of non-zeros in Newton Cholesky factor
int nisland; // number of detected constraint islands
// global properties
@@ -385,12 +386,19 @@ struct mjData_ {
int* island_efcadr; // start address in island_efcind (nisland x 1)
int* island_efcind; // island constraint indices (nefc x 1)
// computed by mj_projectConstraint (dual solver)
// computed by mj_projectConstraint (PGS solver)
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
int* efc_AR_rowadr; // row start address in colind array (nefc x 1)
int* efc_AR_colind; // column indices in sparse AR (nefc x nefc)
mjtNum* efc_AR; // J*inv(M)*J' + R (nefc x nefc)
// computed by mj_fwdConstraint (Newton solver)
int* L_rownnz; // number of non-zeros in Hessian factor L rows (nv x 1)
int* L_rowadr; // row start address in colind array (nv x 1)
int* L_colind; // column indices in sparse AR (nnzL x 1)
mjtNum* L; // chol(M + J'*diag(efc_D)*J) (nnzL x 1)
mjtNum* Lcone; // L with cone contributions (nnzL x 1)
//-------------------- arena-allocated: POSITION, VELOCITY dependent
// computed by mj_fwdVelocity/mj_referenceConstraint
include/mujoco/mjdata.h
+10 -2
View File
@@ -192,12 +192,13 @@ struct mjData_ {
mjTimerStat timer[mjNTIMER]; // timer statistics
// variable sizes
int ncon; // number of detected contacts
int ne; // number of equality constraints
int nf; // number of friction constraints
int nl; // number of limit constraints
int nefc; // number of constraints
int nnzJ; // number of non-zeros in constraint Jacobian
int ncon; // number of detected contacts
int nnzL; // number of non-zeros in Newton Cholesky factor
int nisland; // number of detected constraint islands
// global properties
@@ -413,12 +414,19 @@ struct mjData_ {
int* island_efcadr; // start address in island_efcind (nisland x 1)
int* island_efcind; // island constraint indices (nefc x 1)
// computed by mj_projectConstraint (dual solver)
// computed by mj_projectConstraint (PGS solver)
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
int* efc_AR_rowadr; // row start address in colind array (nefc x 1)
int* efc_AR_colind; // column indices in sparse AR (nefc x nefc)
mjtNum* efc_AR; // J*inv(M)*J' + R (nefc x nefc)
// computed by mj_fwdConstraint (Newton solver)
int* L_rownnz; // number of non-zeros in Hessian factor L rows (nv x 1)
int* L_rowadr; // row start address in colind array (nv x 1)
int* L_colind; // column indices in sparse AR (nnzL x 1)
mjtNum* L; // chol(M + J'*diag(efc_D)*J) (nnzL x 1)
mjtNum* Lcone; // L with cone contributions (nnzL x 1)
//-------------------- arena-allocated: POSITION, VELOCITY dependent
// computed by mj_fwdVelocity/mj_referenceConstraint
include/mujoco/mjxmacro.h
+13 -3
View File
@@ -680,7 +680,7 @@
X( mjContact, contact, MJ_D(ncon), 1 )
// array fields of mjData that are used in the primal problem
#define MJDATA_ARENA_POINTERS_PRIMAL \
#define MJDATA_ARENA_POINTERS_SOLVER \
X( int, efc_type, MJ_D(nefc), 1 ) \
X( int, efc_id, MJ_D(nefc), 1 ) \
X( int, efc_J_rownnz, MJ_D(nefc), 1 ) \
@@ -707,6 +707,14 @@
X( mjtNum, efc_force, MJ_D(nefc), 1 ) \
X( int, efc_state, MJ_D(nefc), 1 )
// array fields of mjData that are used in the dual problem
#define MJDATA_ARENA_POINTERS_NEWTON \
X( int, L_rownnz, MJ_M(nv), 1 ) \
X( int, L_rowadr, MJ_M(nv), 1 ) \
X( int, L_colind, MJ_D(nnzL), 1 ) \
X( mjtNum, L, MJ_D(nnzL), 1 ) \
X( mjtNum, Lcone, MJ_D(nnzL), 1 )
// array fields of mjData that are used in the dual problem
#define MJDATA_ARENA_POINTERS_DUAL \
X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \
@@ -729,7 +737,8 @@
// array fields of mjData that live in d->arena
#define MJDATA_ARENA_POINTERS \
MJDATA_ARENA_POINTERS_CONTACT \
MJDATA_ARENA_POINTERS_PRIMAL \
MJDATA_ARENA_POINTERS_SOLVER \
MJDATA_ARENA_POINTERS_NEWTON \
MJDATA_ARENA_POINTERS_DUAL \
MJDATA_ARENA_POINTERS_ISLAND
@@ -747,12 +756,13 @@
X( int, maxuse_con ) \
X( int, maxuse_efc ) \
X( int, solver_nisland ) \
X( int, ncon ) \
X( int, ne ) \
X( int, nf ) \
X( int, nl ) \
X( int, nefc ) \
X( int, nnzJ ) \
X( int, ncon ) \
X( int, nnzL ) \
X( int, nisland ) \
X( mjtNum, time ) \
X( uintptr_t, threadpool )
introspect/structs.py
+42 -2
View File
@@ -4296,6 +4296,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='timer statistics',
),
StructFieldDecl(
name='ncon',
type=ValueType(name='int'),
doc='number of detected contacts',
),
StructFieldDecl(
name='ne',
type=ValueType(name='int'),
@@ -4322,9 +4327,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='number of non-zeros in constraint Jacobian',
),
StructFieldDecl(
name='ncon',
name='nnzL',
type=ValueType(name='int'),
doc='number of detected contacts',
doc='number of non-zeros in Newton Cholesky factor',
),
StructFieldDecl(
name='nisland',
@@ -5282,6 +5287,41 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc="J*inv(M)*J' + R (nefc x nefc)", # pylint: disable=line-too-long
),
StructFieldDecl(
name='L_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of non-zeros in Hessian factor L rows (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='L_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='row start address in colind array (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='L_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='column indices in sparse AR (nnzL x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='L',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc="chol(M + J'*diag(efc_D)*J) (nnzL x 1)", # pylint: disable=line-too-long
),
StructFieldDecl(
name='Lcone',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='L with cone contributions (nnzL x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='efc_vel',
type=PointerType(
python/mujoco/functions.cc
+2 -2
View File
@@ -1441,7 +1441,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->nefc = 0;
data->contact = static_cast<raw::MjContact*>(data->arena);
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
MJDATA_ARENA_POINTERS_DUAL
#undef X
};
@@ -1469,7 +1469,7 @@ PYBIND11_MODULE(_functions, pymodule) {
throw FatalError("insufficient arena memory available"); \
}
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
if (mj_isDual(d.model().get())) {
MJDATA_ARENA_POINTERS_DUAL
}
python/mujoco/structs.cc
+11 -5
View File
@@ -41,6 +41,7 @@
#include <mujoco/mujoco.h>
#include "errors.h"
#include "function_traits.h"
#include "indexer_xmacro.h"
#include "indexers.h"
#include "private.h"
#include "raw.h"
@@ -743,11 +744,12 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
X(solver);
X(timer);
X(warning);
X(ncon);
X(ne);
X(nf);
X(nnzJ);
X(nnzL);
X(nefc);
X(ncon);
X(nisland);
X(time);
X(energy);
@@ -772,7 +774,8 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
}
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
// MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists
if (mj_isDual(this->model_->get())) {
MJDATA_ARENA_POINTERS_DUAL
}
@@ -819,11 +822,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
X(solver);
X(timer);
X(warning);
X(ncon);
X(ne);
X(nf);
X(nnzJ);
X(nnzL);
X(nefc);
X(ncon);
X(nisland);
X(time);
X(energy);
@@ -850,7 +854,8 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
}
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
// MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists
if (is_dual) {
MJDATA_ARENA_POINTERS_DUAL
}
@@ -2009,7 +2014,8 @@ This is useful for example when the MJB is not available as a file on disk.)"));
return InitPyArray(X_ARRAY_SHAPE(dim0, dim1), d.get()->var, d.owner()); \
});
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
MJDATA_ARENA_POINTERS_NEWTON
MJDATA_ARENA_POINTERS_DUAL
MJDATA_ARENA_POINTERS_ISLAND
src/engine/engine_core_constraint.c
+2 -3
View File
@@ -71,11 +71,10 @@ static int arenaAllocEfc(const mjModel* m, mjData* d) {
return 0; \
}
MJDATA_ARENA_POINTERS_PRIMAL
MJDATA_ARENA_POINTERS_SOLVER
if (mj_isDual(m)) {
MJDATA_ARENA_POINTERS_DUAL
}
#undef X
#undef MJ_M
@@ -1918,7 +1917,7 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
// driver: call all functions above
void mj_makeConstraint(const mjModel* m, mjData* d) {
// clear sizes
d->ne = d->nf = d->nl = d->nefc = d->nnzJ = 0;
d->ne = d->nf = d->nl = d->nefc = d->nnzJ = d->nnzL = 0;
// disabled or Jacobian not allocated: return
if (mjDISABLED(mjDSBL_CONSTRAINT)) {
src/engine/engine_io.c
+2 -1
View File
@@ -1768,12 +1768,13 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
mju_zero(d->solver_fwdinv, 2);
// clear variable sizes
d->ncon = 0;
d->ne = 0;
d->nf = 0;
d->nl = 0;
d->nefc = 0;
d->nnzJ = 0;
d->ncon = 0;
d->nnzL = 0;
d->nisland = 0;
// clear global properties
unity/Runtime/Bindings/MjBindings.cs
+7 -1
View File
@@ -4841,12 +4841,13 @@ public unsafe struct mjData_ {
public mjTimerStat_ timer12;
public mjTimerStat_ timer13;
public mjTimerStat_ timer14;
public int ncon;
public int ne;
public int nf;
public int nl;
public int nefc;
public int nnzJ;
public int ncon;
public int nnzL;
public int nisland;
public double time;
public fixed double energy[2];
@@ -4984,6 +4985,11 @@ public unsafe struct mjData_ {
public int* efc_AR_rowadr;
public int* efc_AR_colind;
public double* efc_AR;
public int* L_rownnz;
public int* L_rowadr;
public int* L_colind;
public double* L;
public double* Lcone;
public double* efc_vel;
public double* efc_aref;
public double* efc_b;