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

Move flexedge sparse Jacobian indices to mjModel.

Browse Source 
PiperOrigin-RevId: 855718305
Change-Id: Ic136f4efa6735b934288f65c332d71ee8b0f6da5
This commit is contained in:
Alessio Quaglino
2026-01-13 07:32:31 -08:00
committed by Copybara-Service
parent 09a6fb7d4d
commit a49576e469
17 changed files with 70 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+4 -4
View File
@@ -272,9 +272,6 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_flex
mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3)
mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6)
int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1)
int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1)
int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv)
mjtNum* flexedge_J; // flex edge Jacobian (nflexedge x nv)
mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1)
mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6)
@@ -1336,12 +1333,15 @@ struct mjModel_ {
mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1)
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtByte* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
int* flex_bvhnum; // number of bounding volumes (nflex x 1)
int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1)
int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1)
int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv)
float* flex_rgba; // rgba when material is omitted (nflex x 4)
float* flex_texcoord; // vertex texture coordinates (nflextexcoord x 2)
include/mujoco/mjdata.h
-3
View File
@@ -306,9 +306,6 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_flex
mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3)
mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6)
int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1)
int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1)
int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv)
mjtNum* flexedge_J; // flex edge Jacobian (nflexedge x nv)
mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1)
mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6)
include/mujoco/mjmodel.h
+4 -1
View File
@@ -1001,12 +1001,15 @@ struct mjModel_ {
mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1)
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtByte* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
int* flex_bvhnum; // number of bounding volumes (nflex x 1)
int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1)
int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1)
int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv)
float* flex_rgba; // rgba when material is omitted (nflex x 4)
float* flex_texcoord; // vertex texture coordinates (nflextexcoord x 2)
include/mujoco/mjxmacro.h
+3 -3
View File
@@ -397,6 +397,9 @@
X ( mjtByte, flex_flatskin, nflex, 1 ) \
X ( int, flex_bvhadr, nflex, 1 ) \
X ( int, flex_bvhnum, nflex, 1 ) \
X ( int, flexedge_J_rownnz, nflexedge, 1 ) \
X ( int, flexedge_J_rowadr, nflexedge, 1 ) \
X ( int, flexedge_J_colind, nflexedge, MJ_M(nv) ) \
X ( float, flex_rgba, nflex, 4 ) \
X ( float, flex_texcoord, nflextexcoord, 2 )
@@ -730,9 +733,6 @@
X ( mjtNum, cinert, nbody, 10 ) \
X ( mjtNum, flexvert_xpos, nflexvert, 3 ) \
X ( mjtNum, flexelem_aabb, nflexelem, 6 ) \
X ( int, flexedge_J_rownnz, nflexedge, 1 ) \
X ( int, flexedge_J_rowadr, nflexedge, 1 ) \
X ( int, flexedge_J_colind, nflexedge, MJ_M(nv) ) \
X ( mjtNum, flexedge_J, nflexedge, MJ_M(nv) ) \
X ( mjtNum, flexedge_length, nflexedge, 1 ) \
X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \
mjx/mujoco/mjx/_src/io.py
-3
View File
@@ -754,9 +754,6 @@ def _make_data_c(
'light_xdir': (m.nlight, 3, float_),
'flexvert_xpos': (nflexvert, 3, float_),
'flexelem_aabb': (nflexelem, 6, float_),
'flexedge_J_rownnz': (nflexedge, np.int32),
'flexedge_J_rowadr': (nflexedge, np.int32),
'flexedge_J_colind': (nflexedge, m.nv, np.int32),
'flexedge_J': (nflexedge, m.nv, float_),
'flexedge_length': (nflexedge, float_),
'ten_J_rownnz': (m.ntendon, np.int32),
mjx/mujoco/mjx/_src/types.py
+3 -3
View File
@@ -626,6 +626,9 @@ class ModelC(PyTreeNode):
flex_centered: jax.Array
flex_bvhadr: jax.Array
flex_bvhnum: jax.Array
flexedge_J_rownnz: jax.Array
flexedge_J_rowadr: jax.Array
flexedge_J_colind: jax.Array
mesh_polynum: jax.Array
mesh_polyadr: jax.Array
mesh_polynormal: jax.Array
@@ -1047,9 +1050,6 @@ class DataC(PyTreeNode):
cinert: jax.Array
flexvert_xpos: jax.Array
flexelem_aabb: jax.Array
flexedge_J_rownnz: jax.Array # pylint:disable=invalid-name
flexedge_J_rowadr: jax.Array # pylint:disable=invalid-name
flexedge_J_colind: jax.Array # pylint:disable=invalid-name
flexedge_J: jax.Array # pylint:disable=invalid-name
flexedge_length: jax.Array
bvh_aabb_dyn: jax.Array
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+1 -1
View File
@@ -858,7 +858,7 @@ def put_data(
d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
flexedge_J = np.zeros((mjm.nflexedge, mjm.nv))
mujoco.mju_sparse2dense(
flexedge_J, mjd.flexedge_J.reshape(-1), mjd.flexedge_J_rownnz, mjd.flexedge_J_rowadr, mjd.flexedge_J_colind.reshape(-1)
flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1)
)
d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float)
else:
python/mujoco/introspect/structs.py
+24 -24
View File
@@ -2971,6 +2971,30 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='number of bounding volumes',
array_extent=('nflex',),
),
StructFieldDecl(
name='flexedge_J_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of non-zeros in Jacobian row',
array_extent=('nflexedge',),
),
StructFieldDecl(
name='flexedge_J_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='row start address in colind array',
array_extent=('nflexedge',),
),
StructFieldDecl(
name='flexedge_J_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='column indices in sparse Jacobian',
array_extent=('nflexedge', 'nv'),
),
StructFieldDecl(
name='flex_rgba',
type=PointerType(
@@ -5564,30 +5588,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='flex element bounding boxes (center, size)',
array_extent=('nflexelem', 6),
),
StructFieldDecl(
name='flexedge_J_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of non-zeros in Jacobian row',
array_extent=('nflexedge',),
),
StructFieldDecl(
name='flexedge_J_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='row start address in colind array',
array_extent=('nflexedge',),
),
StructFieldDecl(
name='flexedge_J_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='column indices in sparse Jacobian',
array_extent=('nflexedge', 'nv'),
),
StructFieldDecl(
name='flexedge_J',
type=PointerType(
src/engine/engine_core_constraint.c
+3 -3
View File
@@ -627,10 +627,10 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
// add constraint: sparse or dense
if (issparse) {
mj_addConstraint(m, d, d->flexedge_J+d->flexedge_J_rowadr[e], cpos, 0, 0,
mj_addConstraint(m, d, d->flexedge_J+m->flexedge_J_rowadr[e], cpos, 0, 0,
1, mjCNSTR_EQUALITY, i,
d->flexedge_J_rownnz[e],
d->flexedge_J_colind+d->flexedge_J_rowadr[e]);
m->flexedge_J_rownnz[e],
m->flexedge_J_colind+m->flexedge_J_rowadr[e]);
} else {
mj_addConstraint(m, d, d->flexedge_J+e*nv, cpos, 0, 0,
1, mjCNSTR_EQUALITY, i,
src/engine/engine_core_smooth.c
+2 -2
View File
@@ -536,7 +536,7 @@ void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) {
// compute flex-related quantities
void mj_flex(const mjModel* m, mjData* d) {
int nv = m->nv, issparse = mj_isSparse(m);
int* rowadr = d->flexedge_J_rowadr, *rownnz = d->flexedge_J_rownnz;
int* rowadr = m->flexedge_J_rowadr, *rownnz = m->flexedge_J_rownnz;
mjtNum* J = d->flexedge_J;
// skip if no flexes
@@ -720,7 +720,7 @@ void mj_flex(const mjModel* m, mjData* d) {
// copy sparsity info
rownnz[ebase+e] = NV;
mju_copyInt(d->flexedge_J_colind + rowadr[ebase+e], chain, NV);
mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV);
}
// dense edge Jacobian
src/engine/engine_derivative.c
+1 -1
View File
@@ -1498,7 +1498,7 @@ void mjd_passive_vel(const mjModel* m, mjData* d) {
// add sparse or dense
if (mj_isSparse(m)) {
addJTBJSparse(m, d, d->flexedge_J, &B, 1, e,
d->flexedge_J_rownnz, d->flexedge_J_rowadr, d->flexedge_J_colind);
m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind);
} else {
addJTBJ(m, d, d->flexedge_J+e*nv, &B, 1);
}
src/engine/engine_forward.c
+1 -1
View File
@@ -221,7 +221,7 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) {
// flexedge velocity: dense or sparse
if (mj_isSparse(m)) {
mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge,
d->flexedge_J_rownnz, d->flexedge_J_rowadr, d->flexedge_J_colind, NULL);
m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL);
} else {
mju_mulMatVec(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->nv);
}
src/engine/engine_passive.c
+3 -3
View File
@@ -414,9 +414,9 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
// transform to joint torque, add to qfrc_{spring, damper}: dense or sparse
if (issparse) {
int end = d->flexedge_J_rowadr[e] + d->flexedge_J_rownnz[e];
for (int j=d->flexedge_J_rowadr[e]; j < end; j++) {
int colind = d->flexedge_J_colind[j];
int end = m->flexedge_J_rowadr[e] + m->flexedge_J_rownnz[e];
for (int j=m->flexedge_J_rowadr[e]; j < end; j++) {
int colind = m->flexedge_J_colind[j];
mjtNum J = d->flexedge_J[j];
d->qfrc_spring[colind] += J * frc_spring;
d->qfrc_damper[colind] += J * frc_damper;
src/engine/engine_print.c
+3 -5
View File
@@ -1284,12 +1284,10 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
printArray2d("FLEXEDGE_J", m->nflexedge, m->nv, d->flexedge_J, fp, float_format);
} else {
mj_printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv,
d->flexedge_J_rowadr, NULL, d->flexedge_J_rownnz, NULL, d->flexedge_J_colind,
m->flexedge_J_rowadr, NULL, m->flexedge_J_rownnz, NULL, m->flexedge_J_colind,
fp);
printArray2dInt("FLEXEDGE_J_ROWNNZ", m->nflexedge, 1, d->flexedge_J_rownnz, fp);
printArray2dInt("FLEXEDGE_J_ROWADR", m->nflexedge, 1, d->flexedge_J_rowadr, fp);
printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, d->flexedge_J_rownnz,
d->flexedge_J_rowadr, d->flexedge_J_colind, fp, float_format);
printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, m->flexedge_J_rownnz,
m->flexedge_J_rowadr, m->flexedge_J_colind, fp, float_format);
}
printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format);
src/engine/engine_setconst.c
+3 -3
View File
@@ -470,9 +470,9 @@ static void set0(mjModel* m, mjData* d) {
// make dense vector into tmp
if (mj_isSparse(m)) {
mju_zero(tmp, nv);
int end = d->flexedge_J_rowadr[i] + d->flexedge_J_rownnz[i];
for (int j=d->flexedge_J_rowadr[i]; j < end; j++) {
tmp[d->flexedge_J_colind[j]] = d->flexedge_J[j];
int end = m->flexedge_J_rowadr[i] + m->flexedge_J_rownnz[i];
for (int j=m->flexedge_J_rowadr[i]; j < end; j++) {
tmp[m->flexedge_J_colind[j]] = d->flexedge_J[j];
}
} else {
mju_copy(tmp, d->flexedge_J+i*nv, nv);
unity/Runtime/Bindings/MjBindings.cs
+3 -3
View File
@@ -4981,9 +4981,6 @@ public unsafe struct mjData_ {
public double* cinert;
public double* flexvert_xpos;
public double* flexelem_aabb;
public int* flexedge_J_rownnz;
public int* flexedge_J_rowadr;
public int* flexedge_J_colind;
public double* flexedge_J;
public double* flexedge_length;
public double* bvh_aabb_dyn;
@@ -5586,6 +5583,9 @@ public unsafe struct mjModel_ {
public byte* flex_flatskin;
public int* flex_bvhadr;
public int* flex_bvhnum;
public int* flexedge_J_rownnz;
public int* flexedge_J_rowadr;
public int* flexedge_J_colind;
public float* flex_rgba;
public float* flex_texcoord;
public int* mesh_vertadr;
wasm/codegen/generated/bindings.cc
+12 -12
View File
@@ -4441,6 +4441,15 @@ struct MjModel {
emscripten::val flex_bvhnum() const {
return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_bvhnum));
}
emscripten::val flexedge_J_rownnz() const {
return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge, ptr_->flexedge_J_rownnz));
}
emscripten::val flexedge_J_rowadr() const {
return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge, ptr_->flexedge_J_rowadr));
}
emscripten::val flexedge_J_colind() const {
return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge * ptr_->nv, ptr_->flexedge_J_colind));
}
emscripten::val flex_rgba() const {
return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 4, ptr_->flex_rgba));
}
@@ -6152,15 +6161,6 @@ struct MjData {
emscripten::val flexelem_aabb() const {
return emscripten::val(emscripten::typed_memory_view(model->nflexelem * 6, ptr_->flexelem_aabb));
}
emscripten::val flexedge_J_rownnz() const {
return emscripten::val(emscripten::typed_memory_view(model->nflexedge, ptr_->flexedge_J_rownnz));
}
emscripten::val flexedge_J_rowadr() const {
return emscripten::val(emscripten::typed_memory_view(model->nflexedge, ptr_->flexedge_J_rowadr));
}
emscripten::val flexedge_J_colind() const {
return emscripten::val(emscripten::typed_memory_view(model->nflexedge * model->nv, ptr_->flexedge_J_colind));
}
emscripten::val flexedge_J() const {
return emscripten::val(emscripten::typed_memory_view(model->nflexedge * model->nv, ptr_->flexedge_J));
}
@@ -10826,9 +10826,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
.property("energy", &MjData::energy)
.property("eq_active", &MjData::eq_active)
.property("flexedge_J", &MjData::flexedge_J)
.property("flexedge_J_colind", &MjData::flexedge_J_colind)
.property("flexedge_J_rowadr", &MjData::flexedge_J_rowadr)
.property("flexedge_J_rownnz", &MjData::flexedge_J_rownnz)
.property("flexedge_length", &MjData::flexedge_length)
.property("flexedge_velocity", &MjData::flexedge_velocity)
.property("flexelem_aabb", &MjData::flexelem_aabb)
@@ -11153,6 +11150,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
.property("flex_vertadr", &MjModel::flex_vertadr)
.property("flex_vertbodyid", &MjModel::flex_vertbodyid)
.property("flex_vertnum", &MjModel::flex_vertnum)
.property("flexedge_J_colind", &MjModel::flexedge_J_colind)
.property("flexedge_J_rowadr", &MjModel::flexedge_J_rowadr)
.property("flexedge_J_rownnz", &MjModel::flexedge_J_rownnz)
.property("flexedge_invweight0", &MjModel::flexedge_invweight0)
.property("flexedge_length0", &MjModel::flexedge_length0)
.property("flexedge_rigid", &MjModel::flexedge_rigid)