diff --git a/doc/includes/references.h b/doc/includes/references.h index bbc75adb..04db929f 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -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) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index d9b539b4..57337d69 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -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) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 7118ec54..45e99516 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -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) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index efac7fd1..c4db5515 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -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 ) \ diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 9e3c1d5e..68d1a69a 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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), diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 0ade3ad2..08e2a3f2 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -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 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index a902247e..c7e00d2a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -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: diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index bf27c58d..808aed09 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -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( diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 303375aa..bc7ba7f3 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -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, diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index efc521be..3b206438 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -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 diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 600acf31..e810601a 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -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); } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 5ec3203d..b62babf7 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -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); } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index fdd8ef94..535a99f9 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -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; diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 64a520fa..78071b01 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -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); diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 34387de4..60a49135 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -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); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index d59ea93d..0705d069 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -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; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 5e00d5f2..ed60e6b5 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -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)