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

Sparse ten_J and ten_J_colind

Browse Source 
PiperOrigin-RevId: 869712136
Change-Id: Id2979684aa0c39552cb8453852483cbac55f0ea3
This commit is contained in:
Taylor Howell
2026-02-13 06:55:42 -08:00
committed by Copybara-Service
parent 277f469239
commit 5903d4826f
21 changed files with 118 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+3 -2
View File
@@ -291,8 +291,8 @@ struct mjData_ {
int* ten_wrapnum; // number of wrap points in path (ntendon x 1)
int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1)
int* ten_J_rowadr; // row start address in colind array (ntendon x 1)
int* ten_J_colind; // column indices in sparse Jacobian (ntendon x nv)
mjtNum* ten_J; // tendon Jacobian (ntendon x nv)
int* ten_J_colind; // column indices in sparse Jacobian (nJten x 1)
mjtNum* ten_J; // tendon Jacobian (nJten x 1)
mjtNum* ten_length; // tendon lengths (ntendon x 1)
int* wrap_obj; // geom id; -1: site; -2: pulley (nwrap x 2)
mjtNum* wrap_xpos; // Cartesian 3D points in all paths (nwrap x 6)
@@ -1096,6 +1096,7 @@ struct mjModel_ {
// sizes set after mjModel construction
mjtSize nnames_map; // number of slots in the names hash map
mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix
mjtSize nJten; // number of non-zeros in sparse ten_J matrix
mjtSize ngravcomp; // number of bodies with nonzero gravcomp
mjtSize nemax; // number of potential equality-constraint rows
mjtSize njmax; // number of available rows in constraint Jacobian (legacy)
include/mujoco/mjdata.h
+2 -2
View File
@@ -325,8 +325,8 @@ struct mjData_ {
int* ten_wrapnum; // number of wrap points in path (ntendon x 1)
int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1)
int* ten_J_rowadr; // row start address in colind array (ntendon x 1)
int* ten_J_colind; // column indices in sparse Jacobian (ntendon x nv)
mjtNum* ten_J; // tendon Jacobian (ntendon x nv)
int* ten_J_colind; // column indices in sparse Jacobian (nJten x 1)
mjtNum* ten_J; // tendon Jacobian (nJten x 1)
mjtNum* ten_length; // tendon lengths (ntendon x 1)
int* wrap_obj; // geom id; -1: site; -2: pulley (nwrap x 2)
mjtNum* wrap_xpos; // Cartesian 3D points in all paths (nwrap x 6)
include/mujoco/mjmodel.h
+1
View File
@@ -754,6 +754,7 @@ struct mjModel_ {
// sizes set after mjModel construction
mjtSize nnames_map; // number of slots in the names hash map
mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix
mjtSize nJten; // number of non-zeros in sparse ten_J matrix
mjtSize ngravcomp; // number of bodies with nonzero gravcomp
mjtSize nemax; // number of potential equality-constraint rows
mjtSize njmax; // number of available rows in constraint Jacobian (legacy)
include/mujoco/mjxmacro.h
+3 -7
View File
@@ -240,6 +240,7 @@
X( npaths ) \
X( nnames_map ) \
X( nJmom ) \
X( nJten ) \
X( ngravcomp ) \
X( nemax ) \
X( njmax ) \
@@ -798,11 +799,6 @@
//-------------------------------- mjData ----------------------------------------------------------
// define symbols needed in MJDATA_POINTERS (corresponding to number of columns)
#define MJDATA_POINTERS_PREAMBLE( m ) \
int nv = m->nv;
// pointer fields of mjData
// XNV means that the field is not required to construct mjvScene
// (by default we define XNV to be the same as X)
@@ -855,8 +851,8 @@
X ( int, ten_wrapnum, ntendon, 1 ) \
X ( int, ten_J_rownnz, ntendon, 1 ) \
X ( int, ten_J_rowadr, ntendon, 1 ) \
X ( int, ten_J_colind, ntendon, MJ_M(nv) ) \
X ( mjtNum, ten_J, ntendon, MJ_M(nv) ) \
X ( int, ten_J_colind, nJten, 1 ) \
X ( mjtNum, ten_J, nJten, 1 ) \
X ( mjtNum, ten_length, ntendon, 1 ) \
X ( int, wrap_obj, nwrap, 2 ) \
X ( mjtNum, wrap_xpos, nwrap, 6 ) \
mjx/mujoco/mjx/_src/io.py
+42 -3
View File
@@ -768,8 +768,8 @@ def _make_data_c(
'flexedge_length': (nflexedge, float_),
'ten_J_rownnz': (m.ntendon, np.int32),
'ten_J_rowadr': (m.ntendon, np.int32),
'ten_J_colind': (m.ntendon, m.nv, np.int32),
'ten_J': (m.ntendon, m.nv, float_),
'ten_J_colind': (m.nJten, np.int32),
'ten_J': (m.nJten, float_),
'ten_wrapadr': (m.ntendon, np.int32),
'ten_wrapnum': (m.ntendon, np.int32),
'wrap_obj': (m.nwrap, 2, np.int32),
@@ -1099,7 +1099,7 @@ def _put_data_jax(
# MJX does not support islanding, so only transfer the first solver_niter
impl_fields['solver_niter'] = impl_fields['solver_niter'][0]
# convert sparse representation of actuator_moment to dense matrix
# convert sparse actuator_moment to dense matrix
moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense(
moment,
@@ -1110,6 +1110,22 @@ def _put_data_jax(
)
impl_fields['actuator_moment'] = moment
# convert ten_J to dense matrix
if mujoco.mj_isSparse(m):
ten_J = np.zeros((m.ntendon, m.nv))
mujoco.mju_sparse2dense(
ten_J,
d.ten_J,
d.ten_J_rownnz,
d.ten_J_rowadr,
d.ten_J_colind,
)
elif m.ntendon:
ten_J = d.ten_J.reshape((m.ntendon, m.nv))
else:
ten_J = np.zeros((m.ntendon, m.nv))
impl_fields['ten_J'] = ten_J
contact, contact_map = _put_contact(d.contact, dim, efc_address)
# pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
@@ -1568,6 +1584,29 @@ def _get_data_into(
result_i.actuator_moment[:] = actuator_moment
continue
# MuJoCo ten_J is sparse, MJX uses a dense representation.
if field.name == 'ten_J':
ten_j_rownnz = np.zeros(m.ntendon, dtype=np.int32)
ten_j_rowadr = np.zeros(m.ntendon, dtype=np.int32)
ten_j_colind = np.zeros(m.nJten, dtype=np.int32)
ten_j = np.zeros(m.nJten)
if m.ntendon:
if d_i.impl == types.Impl.JAX:
mujoco.mju_dense2sparse(
ten_j,
d_i._impl.ten_J,
ten_j_rownnz,
ten_j_rowadr,
ten_j_colind,
)
else:
ten_j = d_i._impl.ten_J
result_i.ten_J_rownnz[:] = ten_j_rownnz
result_i.ten_J_rowadr[:] = ten_j_rowadr
result_i.ten_J_colind[:] = ten_j_colind
result_i.ten_J[:] = ten_j
continue
if hasattr(d_i._impl, field.name):
value = getattr(d_i._impl, field.name)
else:
mjx/mujoco/mjx/_src/smooth_test.py
+17 -2
View File
@@ -121,7 +121,9 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_forward(m, d)
# tendon
dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx._impl, 'ten_J')
# with dense jacobian mode, d.ten_J is already dense (ntendon*nv,), just reshape
ten_J = d.ten_J.reshape((m.ntendon, m.nv))
_assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
_assert_attr_eq(d, dx, 'ten_length')
# transmission
dx = jax.jit(mjx.transmission)(mx, dx)
@@ -395,7 +397,20 @@ class TendonTest(parameterized.TestCase):
dx = jax.jit(mjx.forward)(mx, dx)
_assert_eq(d.ten_length, dx.ten_length, 'ten_length')
_assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
# convert ten_J for comparison based on jacobian mode
if mujoco.mj_isSparse(m):
ten_J = np.zeros((m.ntendon, m.nv))
mujoco.mju_sparse2dense(
ten_J,
d.ten_J,
d.ten_J_rownnz,
d.ten_J_rowadr,
d.ten_J_colind,
)
else:
# dense mode: just reshape
ten_J = d.ten_J.reshape((m.ntendon, m.nv))
_assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
_assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
_assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
_assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
mjx/mujoco/mjx/_src/types.py
+2
View File
@@ -722,6 +722,7 @@ class Model(PyTreeNode):
nC: number of non-zeros in C matrix
nD: number of non-zeros in D matrix
nJmom: number of non-zeros in Jacobian momentum matrix
nJten: number of non-zeros in sparse tendon Jacobian
ngravcomp: number of bodies with nonzero gravcomp
nuserdata: number of elements in userdata
nsensordata: number of elements in sensor data vector
@@ -772,6 +773,7 @@ class Model(PyTreeNode):
nC: int # pylint:disable=invalid-name
nD: int # pylint:disable=invalid-name
nJmom: int # pylint:disable=invalid-name
nJten: int # pylint:disable=invalid-name
ngravcomp: int
nuserdata: int
nsensordata: int
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+12 -1
View File
@@ -1163,7 +1163,18 @@ def get_data_into(
# tendon
result.ten_length[:] = d.ten_length.numpy()[world_id]
result.ten_J[:] = d.ten_J.numpy()[world_id]
# TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
if BLEEDING_EDGE_MUJOCO:
ten_J = d.ten_J.numpy()[world_id]
mujoco.mju_dense2sparse(
result.ten_J,
ten_J,
result.ten_J_rownnz,
result.ten_J_rowadr,
result.ten_J_colind,
)
else:
result.ten_J[:] = d.ten_J.numpy()[world_id]
result.ten_wrapadr[:] = d.ten_wrapadr.numpy()[world_id]
result.ten_wrapnum[:] = d.ten_wrapnum.numpy()[world_id]
result.wrap_obj[:] = d.wrap_obj.numpy()[world_id]
mjx/mujoco/mjx/warp/forward_test.py
+2 -1
View File
@@ -149,7 +149,8 @@ class ForwardTest(parameterized.TestCase):
tu.assert_attr_eq(dx, d, 'cam_xpos')
tu.assert_eq(dx.cam_xmat, d.cam_xmat.reshape((-1, 3, 3)), 'cam_xmat')
tu.assert_attr_eq(dx, d, 'ten_length')
tu.assert_attr_eq(dx._impl, d, 'ten_J')
ten_J = d.ten_J.reshape((m.ntendon, m.nv))
tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')
tu.assert_attr_eq(dx._impl, d, 'wrap_xpos')
python/mujoco/introspect/structs.py
+7 -2
View File
@@ -1242,6 +1242,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=ValueType(name='mjtSize'),
doc='number of non-zeros in sparse actuator_moment matrix',
),
StructFieldDecl(
name='nJten',
type=ValueType(name='mjtSize'),
doc='number of non-zeros in sparse ten_J matrix',
),
StructFieldDecl(
name='ngravcomp',
type=ValueType(name='mjtSize'),
@@ -5829,7 +5834,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='int'),
),
doc='column indices in sparse Jacobian',
array_extent=('ntendon', 'nv'),
array_extent=('nJten',),
),
StructFieldDecl(
name='ten_J',
@@ -5837,7 +5842,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'),
),
doc='tendon Jacobian',
array_extent=('ntendon', 'nv'),
array_extent=('nJten',),
),
StructFieldDecl(
name='ten_length',
python/mujoco/structs_wrappers.cc
-2
View File
@@ -747,7 +747,6 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
// Write buffer and arena contents
{
MJDATA_POINTERS_PREAMBLE((this->model_->get()))
#define X(type, name, nr, nc) \
WriteBytes(output, ptr_->name, \
@@ -827,7 +826,6 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
// Read buffer and arena contents
{
MJDATA_POINTERS_PREAMBLE((&m))
#define X(type, name, nr, nc) \
ReadBytes(input, d->name, sizeof(type) * (m.nr) * (nc));
simulate/simulate.cc
-3
View File
@@ -669,9 +669,6 @@ void UpdateWatch(mj::Simulate* sim, const mjModel* m, const mjData* d) {
sim->ui0.sect[SECT_WATCH].item[2].multi.nelem = 1;
mju::strcpy_arr(sim->ui0.sect[SECT_WATCH].item[2].multi.name[0], "invalid field");
// prepare symbols needed by xmacro
MJDATA_POINTERS_PREAMBLE(m);
// find specified field in mjData arrays, update value
#define X(TYPE, NAME, NR, NC) \
if (!mju::strcmp_arr(#NAME, sim->field) && \
src/engine/engine_io.c
+3 -11
View File
@@ -224,7 +224,7 @@ void mj_makeModel(mjModel** dest,
// CHECK SIZE PARAMETERS
{
// dummy variables for MJMODEL_SIZES set after mjModel construction
int nnames_map=0, nJmom=0, ngravcomp=0, nemax=0, njmax=0, nconmax=0;
int nnames_map = 0, nJmom = 0, nJten = 0, ngravcomp = 0, nemax = 0, njmax = 0, nconmax=0;
int nuserdata=0, nsensordata=0, npluginstate=0, nhistory=0, narena=0, nbuffer=0;
// sizes must be non-negative and fit in int, except for the byte arrays texdata and textdata
@@ -240,10 +240,10 @@ void mj_makeModel(mjModel** dest,
return; \
}
MJMODEL_SIZES
#undef X
#undef X
// suppress unused variable warnings
(void)nnames_map; (void)nJmom; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax;
(void)nnames_map; (void)nJmom; (void)nJten; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax;
(void)nuserdata; (void)nsensordata; (void)npluginstate; (void)nhistory; (void)narena;
(void)nbuffer;
}
@@ -964,9 +964,6 @@ void mj_makeDofDofMaps(int nv, int nM, int nC, int nD,
static void mj_setPtrData(const mjModel* m, mjData* d) {
char* ptr = (char*)d->buffer;
// prepare symbols needed by xmacro
MJDATA_POINTERS_PREAMBLE(m);
// assign pointers with padding
#define X(type, name, nr, nc) \
d->name = (type*)(ptr + SKIP((intptr_t)ptr)); \
@@ -1047,9 +1044,6 @@ void mj_makeRawData(mjData** dest, const mjModel* m) {
mjERROR("could not allocate mjData");
}
// prepare symbols needed by xmacro
MJDATA_POINTERS_PREAMBLE(m);
// compute buffer size
d->nbuffer = 0;
d->buffer = d->arena = NULL;
@@ -1160,7 +1154,6 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
// copy buffer
{
MJDATA_POINTERS_PREAMBLE(m)
if (flg_all) {
#define X(type, name, nr, nc) \
memcpy((char*)dest->name, (const char*)src->name, sizeof(type)*(m->nr)*nc);
@@ -1346,7 +1339,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
#ifdef ADDRESS_SANITIZER
{
#define X(type, name, nr, nc) memset(d->name, (int)debug_value, sizeof(type)*(m->nr)*(nc));
MJDATA_POINTERS_PREAMBLE(m)
MJDATA_POINTERS
#undef X
}
src/experimental/platform/helpers.cc
-1
View File
@@ -100,7 +100,6 @@ void SaveToWebp(int width, int height, const std::byte* data,
const void* GetValue(const mjModel* model, const mjData* data,
const char* field, int index) {
MJDATA_POINTERS_PREAMBLE(model);
#define X(TYPE, NAME, NR, NC) \
if (!std::strcmp(#NAME, field) && !std::strcmp(#TYPE, "mjtNum")) { \
if (index >= 0 && index < model->NR * NC) { \
src/user/user_model.cc
+12
View File
@@ -3226,7 +3226,16 @@ int mjCModel::CountNJmom(const mjModel* m) {
return count;
}
// compute non-zeros in ten_J matrix
int mjCModel::CountNJten(const mjModel* m) {
int nv = m->nv;
int ntendon = m->ntendon;
// conservative upper bound: each tendon can have at most nv non-zeros
// TODO(taylorhowell): compute tighter bound
int count = ntendon * nv;
return count;
}
// copy objects outside kinematic tree
void mjCModel::CopyObjects(mjModel* m) {
@@ -5037,6 +5046,9 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
// compute non-zeros in actuator_moment
m->nJmom = nJmom = CountNJmom(m);
// compute non-zeros in ten_J
m->nJten = nJten = CountNJten(m);
// scale mass
if (compiler.settotalmass > 0) {
mj_setTotalmass(m, compiler.settotalmass);
src/user/user_model.h
+2
View File
@@ -130,6 +130,7 @@ class mjCModel_ : public mjsElement {
mjtSize nC; // number of non-zeros in reduced sparse dof-dof matrix
mjtSize nD; // number of non-zeros in sparse dof-dof matrix
mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix
mjtSize nJten; // number of non-zeros in sparse ten_J matrix
// statistics, as computed by mj_setConst
double meaninertia_auto; // mean diagonal inertia, as computed by mj_setConst
@@ -368,6 +369,7 @@ class mjCModel : public mjCModel_, private mjSpec {
void FinalizeSimple(mjModel* m); // finalize simple bodies/dofs including tendon information
void CopyPlugins(mjModel*); // copy plugin data
int CountNJmom(const mjModel* m); // compute number of non-zeros in actuator_moment matrix
int CountNJten(const mjModel* m); // compute number of non-zeros in ten_J matrix
// remove plugins that are not referenced by any object
void RemovePlugins();
test/engine/engine_core_smooth_test.cc
-1
View File
@@ -73,7 +73,6 @@ TEST_F(CoreSmoothTest, MjDataWorldBodyValuesAreInitialized) {
mj_rnePostConstraint(model, data);
{
MJDATA_POINTERS_PREAMBLE(model)
#define X(type, name, d0, d1) \
if constexpr (std::string_view(#d0) == "nbody") { \
absl::Span<type> values(data->name, model->d0 * d1); \
test/engine/engine_forward_test.cc
-1
View File
@@ -725,7 +725,6 @@ TEST_F(ForwardTest, NormalizeQuats) {
mj_step(model, data_n);
// expect everything to match
MJDATA_POINTERS_PREAMBLE(model)
#define X(type, name, nr, nc) \
for (int i = 0; i < model->nr; i++) \
for (int j = 0; j < nc; j++) \
test/engine/engine_thread_test.cc
-2
View File
@@ -71,7 +71,6 @@ TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
// compare the mjData's.
{
MJDATA_POINTERS_PREAMBLE((model))
#define X(type, name, nr, nc) \
EXPECT_EQ(std::memcmp(data->name, data_threaded->name, \
sizeof(type)*(model->nr)*(nc)), \
@@ -127,7 +126,6 @@ TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
// compare the mjData's.
{
MJDATA_POINTERS_PREAMBLE((model))
#define X(type, name, nr, nc) \
EXPECT_EQ(std::memcmp(data->name, data_threaded->name, \
sizeof(type)*(model->nr)*(nc)), \
unity/Runtime/Bindings/MjBindings.cs
+1
View File
@@ -5371,6 +5371,7 @@ public unsafe struct mjModel_ {
public Int64 npaths;
public Int64 nnames_map;
public Int64 nJmom;
public Int64 nJten;
public Int64 ngravcomp;
public Int64 nemax;
public Int64 njmax;
wasm/codegen/generated/bindings.cc
+9 -2
View File
@@ -4078,6 +4078,12 @@ struct MjModel {
void set_nJmom(int value) {
ptr_->nJmom = static_cast<mjtSize>(value);
}
int nJten() const {
return static_cast<int>(ptr_->nJten);
}
void set_nJten(int value) {
ptr_->nJten = static_cast<mjtSize>(value);
}
int ngravcomp() const {
return static_cast<int>(ptr_->ngravcomp);
}
@@ -6650,10 +6656,10 @@ struct MjData {
return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_J_rowadr));
}
emscripten::val ten_J_colind() const {
return emscripten::val(emscripten::typed_memory_view(model->ntendon * model->nv, ptr_->ten_J_colind));
return emscripten::val(emscripten::typed_memory_view(model->nJten, ptr_->ten_J_colind));
}
emscripten::val ten_J() const {
return emscripten::val(emscripten::typed_memory_view(model->ntendon * model->nv, ptr_->ten_J));
return emscripten::val(emscripten::typed_memory_view(model->nJten, ptr_->ten_J));
}
emscripten::val ten_length() const {
return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_length));
@@ -11897,6 +11903,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
.property("nJfe", &MjModel::nJfe, &MjModel::set_nJfe, reference())
.property("nJfv", &MjModel::nJfv, &MjModel::set_nJfv, reference())
.property("nJmom", &MjModel::nJmom, &MjModel::set_nJmom, reference())
.property("nJten", &MjModel::nJten, &MjModel::set_nJten, reference())
.property("nM", &MjModel::nM, &MjModel::set_nM, reference())
.property("na", &MjModel::na, &MjModel::set_na, reference())
.property("name_actuatoradr", &MjModel::name_actuatoradr)