From bb0b5ca5214487c4e529b4a66b2e150a3097eec1 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 18 Feb 2026 12:35:08 +0000 Subject: [PATCH 01/48] lint --- .github/workflows/lint.yml | 20 ++++++++++++++++++++ .pre-commit-config.yaml | 6 ++++++ 2 files changed, 26 insertions(+) create mode 100644 .github/workflows/lint.yml create mode 100644 .pre-commit-config.yaml diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..8a2aa9a2 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,20 @@ +# Run pre-commit hooks on files changed in the PR only. +name: lint + +on: + pull_request: + +jobs: + pre-commit: + name: pre-commit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - uses: pre-commit/action@v3.0.1 + with: + extra_args: --from-ref ${{ github.event.pull_request.base.sha }} --to-ref ${{ github.event.pull_request.head.sha }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..a674a832 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,6 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer From 59bc4fb88edc221d78501d710bce679feb78e877 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Wed, 18 Feb 2026 15:19:00 -0500 Subject: [PATCH 02/48] Fix typo and improve wording in render doc --- doc/mjwarp/index.rst | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 4fe9dfe3..e85011c0 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -467,9 +467,9 @@ subset of fields. Batch Rendering =============== -MJWarp provides a high-throughput ray-tracing batch renderer built on +MJWarp provides a high-throughput ray tracing batch renderer built on `Warp's accelerated BVHs `__ for -rendering worlds with multiple cameras in parallel on device. +rendering worlds with multiple cameras in parallel. Key features: @@ -490,8 +490,8 @@ Key features: Basic Usage ---------- -Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures -and rendering and output buffers. +Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures, +rendering specific fields, and output buffers. .. code-block:: python @@ -526,9 +526,9 @@ followed by :func:`mjw.render ` to write to output buffers. mjw.render(m, d, rc) The output buffers contain stacked pixels for all cameras with shape `(nworld, npixel)` and RGB data is -packed into one `unit32` variable. `RenderContext.rgb_adr` and `RenderContext.depth_adr` provide per-camera indexing. +packed into one `uint32` variable. `RenderContext.rgb_adr` and `RenderContext.depth_adr` provide per-camera indexing. For convenience, :func:`mjw.get_rgb ` and :func:`mjw.get_depth ` -provide per-camera batched post-processing. +return processed and reshaped RGB and depth data for a given camera batched for all worlds. .. code-block:: python @@ -555,12 +555,12 @@ For benchmark results across a variety of scenes, see the Notes ----- -- **Meshes**: Rendering computation scales with mesh complexity. A primitive is expected to have better - performance (i.e., higher throughput) compared to a similar sized :ref:`mesh` or +- **Meshes**: Rendering computation scales with mesh complexity, specifically the number of vertices and faces. A primitive is expected to have better + performance (i.e., higher throughput) compared to a similar-sized :ref:`mesh` or :ref:`heightfield `. -- **Flex**: Currently limited to 2D and 3D :ref:`flex` objects. Performance is expected to improved as +- **Flex**: Currently limited to 2D and 3D :ref:`flex` objects. Performance is expected to improve as this feature is further developed. -- **Scaling**: Rendering scales linearly with resolution (total number of pixels) and number of cameras. +- **Scaling**: Rendering scales linearly with resolution (total pixel count) and camera count (parallel structure). .. _mjwFAQ: From 3c5986fe502d00d5561d6d8f47e5502a12fc2eb7 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Thu, 19 Feb 2026 07:20:31 -0500 Subject: [PATCH 03/48] Minor edit to MJWarp batch rendering documentation --- doc/mjwarp/index.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index e85011c0..baafac32 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -467,7 +467,7 @@ subset of fields. Batch Rendering =============== -MJWarp provides a high-throughput ray tracing batch renderer built on +MJWarp provides a batch renderer for high-throughput ray tracing built on `Warp's accelerated BVHs `__ for rendering worlds with multiple cameras in parallel. @@ -560,7 +560,7 @@ Notes :ref:`heightfield `. - **Flex**: Currently limited to 2D and 3D :ref:`flex` objects. Performance is expected to improve as this feature is further developed. -- **Scaling**: Rendering scales linearly with resolution (total pixel count) and camera count (parallel structure). +- **Scaling**: Rendering scales linearly with resolution (total pixel count) and camera count. .. _mjwFAQ: From 10524c28dde74b915ac8cfab5bc165bb845d8f6c Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Thu, 19 Feb 2026 07:05:56 -0800 Subject: [PATCH 04/48] Track and restore viscous pause state every frame. PiperOrigin-RevId: 872382150 Change-Id: I96bddcfb2a0a204c6505e9795dfff97f09834855 --- src/experimental/platform/step_control.cc | 49 +++++++++++++++-------- src/experimental/platform/step_control.h | 5 --- 2 files changed, 32 insertions(+), 22 deletions(-) diff --git a/src/experimental/platform/step_control.cc b/src/experimental/platform/step_control.cc index 4c8b58cf..38c97fe2 100644 --- a/src/experimental/platform/step_control.cc +++ b/src/experimental/platform/step_control.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -30,6 +31,33 @@ static mjtNum Timer() { return Milliseconds(Clock::now() - start).count(); } +// Updates key viscous pause parameters restores them when done. +struct ViscousPauseState { + ViscousPauseState(mjModel* model) : model(model) { + if (model) { + mju_copy3(gravity, model->opt.gravity); + viscosity = model->opt.viscosity; + disableflags = model->opt.disableflags; + mju_zero3(model->opt.gravity); + model->opt.viscosity = 10; + model->opt.disableflags |= mjDSBL_SPRING; + } + } + + ~ViscousPauseState() { + if (model) { + mju_copy3(model->opt.gravity, gravity); + model->opt.viscosity = viscosity; + model->opt.disableflags = disableflags; + } + } + mjModel* model; + mjtNum gravity[3]; + mjtNum viscosity; + int disableflags; +}; + + StepControl::StepControl() { mjcb_time = Timer; } float StepControl::GetSpeedMeasured() const { return speed_measured_; } @@ -56,23 +84,6 @@ void StepControl::SetNoiseParameters(float ctrl_noise_scale, } void StepControl::SetPauseState(PauseState state, mjModel* m) { - if (pause_state_ == PauseState::kViscousPaused && - state != PauseState::kViscousPaused && m) { - mju_copy3(m->opt.gravity, saved_gravity_); - m->opt.viscosity = saved_viscosity_; - m->opt.disableflags = saved_disableflags_; - } - - if (state == PauseState::kViscousPaused && - pause_state_ != PauseState::kViscousPaused && m) { - mju_copy3(saved_gravity_, m->opt.gravity); - saved_viscosity_ = m->opt.viscosity; - saved_disableflags_ = m->opt.disableflags; - mju_zero3(m->opt.gravity); - m->opt.viscosity = 10; - m->opt.disableflags |= mjDSBL_SPRING; - } - pause_state_ = state; } @@ -81,6 +92,10 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d) { return Status::kOk; } + std::optional viscous_pause_state; + if (m && pause_state_ == PauseState::kViscousPaused) { + viscous_pause_state.emplace(m); + } if (pause_state_ == PauseState::kNormalPaused) { // When we eventually unpause, we need to make sure we sync to immediately diff --git a/src/experimental/platform/step_control.h b/src/experimental/platform/step_control.h index 916a23c7..7cebbd1b 100644 --- a/src/experimental/platform/step_control.h +++ b/src/experimental/platform/step_control.h @@ -111,11 +111,6 @@ class StepControl { PauseState pause_state_ = PauseState::kUnpaused; - // Viscous pause state variables - mjtNum saved_gravity_[3] = {0}; - mjtNum saved_viscosity_ = 0; - int saved_disableflags_ = 0; - // Perform only a single step on the next call to Advance() if the simulation // is paused. bool single_step_ = false; From 9efe41c0c167f4299899ca443030aedf1823129a Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 19 Feb 2026 09:53:11 -0800 Subject: [PATCH 05/48] Remove dense code path for tendon Jacobian PiperOrigin-RevId: 872444531 Change-Id: I6180101abc49469a72aee8bec5726e1e94f142ec --- mjx/mujoco/mjx/_src/io.py | 4 +- mjx/mujoco/mjx/_src/smooth_test.py | 31 ++--- mjx/mujoco/mjx/warp/forward_test.py | 9 +- src/engine/engine_core_constraint.c | 38 +++--- src/engine/engine_core_smooth.c | 179 +++++++------------------ src/engine/engine_derivative.c | 8 +- src/engine/engine_forward.c | 10 +- src/engine/engine_passive.c | 22 ++- src/engine/engine_print.c | 16 +-- src/engine/engine_setconst.c | 11 +- test/engine/engine_core_smooth_test.cc | 14 +- 11 files changed, 124 insertions(+), 218 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index e683633b..8feef14d 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1111,7 +1111,7 @@ def _put_data_jax( impl_fields['actuator_moment'] = moment # convert ten_J to dense matrix - if mujoco.mj_isSparse(m): + if m.ntendon: ten_J = np.zeros((m.ntendon, m.nv)) mujoco.mju_sparse2dense( ten_J, @@ -1120,8 +1120,6 @@ def _put_data_jax( 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 diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 667b3d63..ce8c2588 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -121,8 +121,14 @@ class SmoothTest(absltest.TestCase): mujoco.mj_forward(m, d) # tendon dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d)) - # with dense jacobian mode, d.ten_J is already dense (ntendon*nv,), just reshape - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) + 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, + ) _assert_eq(ten_J, dx._impl.ten_J, 'ten_J') _assert_attr_eq(d, dx, 'ten_length') # transmission @@ -397,19 +403,14 @@ class TendonTest(parameterized.TestCase): dx = jax.jit(mjx.forward)(mx, dx) _assert_eq(d.ten_length, dx.ten_length, 'ten_length') - # 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)) + 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, + ) _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') diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 020e924a..30eb21cf 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -149,7 +149,14 @@ 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') - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) + 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, + ) 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') diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index ab6abd39..d550766c 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -567,7 +567,6 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // copy Jacobian: sparse or dense if (issparse) { - // add first or second chain if (j == 0) { NV = d->ten_J_rownnz[id[j]]; mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); @@ -578,7 +577,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { mju_copy(jac[j], d->ten_J+d->ten_J_rowadr[id[j]], NV2); } } else { - mju_copy(jac[j], d->ten_J+id[j]*nv, nv); + mju_sparse2dense(jac[j], d->ten_J, 1, nv, d->ten_J_rownnz+id[j], d->ten_J_rowadr+id[j], d->ten_J_colind); } } } @@ -737,11 +736,17 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { if (m->tendon_frictionloss[i] > 0) { int efcadr = d->nefc; // add constraint - mj_addConstraint(m, d, d->ten_J + (issparse ? d->ten_J_rowadr[i] : i*nv), - 0, 0, m->tendon_frictionloss[i], - 1, mjCNSTR_FRICTION_TENDON, i, - issparse ? d->ten_J_rownnz[i] : 0, - issparse ? d->ten_J_colind+d->ten_J_rowadr[i] : NULL); + if (issparse) { + mj_addConstraint(m, d, d->ten_J + d->ten_J_rowadr[i], + 0, 0, m->tendon_frictionloss[i], + 1, mjCNSTR_FRICTION_TENDON, i, + d->ten_J_rownnz[i], + d->ten_J_colind+d->ten_J_rowadr[i]); + } else { + mju_sparse2dense(jac, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); + mj_addConstraint(m, d, jac, 0, 0, m->tendon_frictionloss[i], + 1, mjCNSTR_FRICTION_TENDON, i, 0, NULL); + } // set tendon_efcadr if (d->tendon_efcadr[i] == -1) { d->tendon_efcadr[i] = efcadr; @@ -877,19 +882,20 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { // detect tendon limit if (dist < margin) { - // prepare Jacobian: sparse or dense + // prepare Jacobian + int efcadr = d->nefc; if (issparse) { mju_scl(jac, d->ten_J+d->ten_J_rowadr[i], -side, d->ten_J_rownnz[i]); + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_TENDON, i, + d->ten_J_rownnz[i], + d->ten_J_colind+d->ten_J_rowadr[i]); } else { - mju_scl(jac, d->ten_J+i*nv, -side, nv); + mju_sparse2dense(jac, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); + mju_scl(jac, jac, -side, nv); + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_TENDON, i, 0, NULL); } - - // add constraint - int efcadr = d->nefc; - mj_addConstraint(m, d, jac, &dist, &margin, 0, - 1, mjCNSTR_LIMIT_TENDON, i, - issparse ? d->ten_J_rownnz[i] : 0, - issparse ? d->ten_J_colind+d->ten_J_rowadr[i] : NULL); // set tendon_efcadr if (d->tendon_efcadr[i] == -1) { d->tendon_efcadr[i] = efcadr; diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ba5c7254..1d3fed48 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -904,7 +904,7 @@ void mj_flex(const mjModel* m, mjData* d) { // compute tendon lengths and moments void mj_tendon(const mjModel* m, mjData* d) { - int issparse = mj_isSparse(m), nv = m->nv, nten = m->ntendon; + int nv = m->nv, nten = m->ntendon; int *rownnz = d->ten_J_rownnz, *rowadr = d->ten_J_rowadr, *colind = d->ten_J_colind; mjtNum *L = d->ten_length, *J = d->ten_J; @@ -913,28 +913,22 @@ void mj_tendon(const mjModel* m, mjData* d) { } // allocate stack arrays - int *chain = NULL, *buf_ind = NULL; - mjtNum *jac1, *jac2, *jacdif, *tmp, *sparse_buf = NULL; + int *chain, *buf_ind; + mjtNum *jac1, *jac2, *jacdif, *tmp, *sparse_buf; mj_markStack(d); jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); tmp = mjSTACKALLOC(d, nv, mjtNum); - if (issparse) { - chain = mjSTACKALLOC(d, nv, int); - buf_ind = mjSTACKALLOC(d, nv, int); - sparse_buf = mjSTACKALLOC(d, nv, mjtNum); - } + chain = mjSTACKALLOC(d, nv, int); + buf_ind = mjSTACKALLOC(d, nv, int); + sparse_buf = mjSTACKALLOC(d, nv, mjtNum); // clear results mju_zero(L, nten); - // clear Jacobian: sparse or dense - if (issparse) { - mju_zeroInt(rownnz, nten); - } else { - mju_zero(J, nten*nv); - } + // clear Jacobian + mju_zeroInt(rownnz, nten); // sleep filtering int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; @@ -954,9 +948,7 @@ void mj_tendon(const mjModel* m, mjData* d) { int tendon_num = m->tendon_num[i]; // sparse Jacobian row init - if (issparse) { - rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); - } + rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); // process fixed tendon if (m->wrap_type[adr] == mjWRAP_JOINT) { @@ -969,17 +961,10 @@ void mj_tendon(const mjModel* m, mjData* d) { L[i] += m->wrap_prm[adr+j] * d->qpos[m->jnt_qposadr[k]]; // add to moment - if (issparse) { - rownnz[i] = mju_combineSparse(J+rowadr[i], &m->wrap_prm[adr+j], 1, 1, - rownnz[i], 1, - colind+rowadr[i], &m->jnt_dofadr[k], - sparse_buf, buf_ind); - } - - // add to moment: dense - else { - J[i*nv + m->jnt_dofadr[k]] = m->wrap_prm[adr+j]; - } + rownnz[i] = mju_combineSparse(J+rowadr[i], &m->wrap_prm[adr+j], 1, 1, + rownnz[i], 1, + colind+rowadr[i], &m->jnt_dofadr[k], + sparse_buf, buf_ind); } continue; @@ -1060,40 +1045,23 @@ void mj_tendon(const mjModel* m, mjData* d) { mji_sub3(dif, wpnt+3*k+3, wpnt+3*k); mju_normalize3(dif); - // sparse - if (issparse) { - // get endpoint Jacobians, subtract - int NV = mj_jacDifPair(m, d, chain, - wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, - jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); + // get endpoint Jacobians, subtract + int NV = mj_jacDifPair(m, d, chain, + wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); - // no dofs: skip - if (!NV) { - continue; - } - - // apply chain rule to compute tendon Jacobian - mju_mulMatTVec(tmp, jacdif, dif, 3, NV); - - // add to existing - rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor, - rownnz[i], NV, colind+rowadr[i], - chain, sparse_buf, buf_ind); + // no dofs: skip + if (!NV) { + continue; } - // dense - else { - // get endpoint Jacobians, subtract - mj_jac(m, d, jac1, 0, wpnt+3*k, wbody[k]); - mj_jac(m, d, jac2, 0, wpnt+3*k+3, wbody[k+1]); - mju_sub(jacdif, jac2, jac1, 3*nv); + // apply chain rule to compute tendon Jacobian + mju_mulMatTVec(tmp, jacdif, dif, 3, NV); - // apply chain rule to compute tendon Jacobian - mju_mulMatTVec(tmp, jacdif, dif, 3, nv); - - // add to existing - mju_addToScl(J + i*nv, tmp, 1/divisor, nv); - } + // add to existing + rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor, + rownnz[i], NV, colind+rowadr[i], + chain, sparse_buf, buf_ind); } } @@ -1449,28 +1417,14 @@ void mj_transmission(const mjModel* m, mjData* d) { case mjTRN_TENDON: // tendon length[i] = d->ten_length[id]*gear[0]; - // moment: sparse or dense - if (issparse) { - // sparsity + // moment + { int ten_J_rownnz = d->ten_J_rownnz[id]; int ten_J_rowadr = d->ten_J_rowadr[id]; rownnz[i] = ten_J_rownnz; mju_copyInt(colind + adr, d->ten_J_colind + ten_J_rowadr, ten_J_rownnz); mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz); - } else { - mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv); - - // sparsity (compress) - nnz = 0; - for (int j = 0; j < nv; j++) { - if (moment[adr+j]) { - moment[adr+nnz] = moment[adr+j]; - colind[adr+nnz] = j; - nnz++; - } - } - rownnz[i] = nnz; } break; @@ -1743,7 +1697,7 @@ void mj_transmission(const mjModel* m, mjData* d) { // add tendon armature to M void mj_tendonArmature(const mjModel* m, mjData* d) { - int nv = m->nv, ntendon = m->ntendon, issparse = mj_isSparse(m); + int nv = m->nv, ntendon = m->ntendon; const int* M_rownnz = m->M_rownnz; const int* M_rowadr = m->M_rowadr; const int* M_colind = m->M_colind; @@ -1762,47 +1716,25 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { continue; } - // dense - if (!issparse) { - // M += armature * ten_J' * ten_J - mjtNum* ten_J = d->ten_J + nv*k; - for (int i=0; i < nv; i++) { - mjtNum ten_J_i = ten_J[i]; - if (!ten_J_i) { - continue; - } + // get sparse info for tendon k + int J_rowadr = d->ten_J_rowadr[k]; + int J_rownnz = d->ten_J_rownnz[k]; + const int* J_colind = d->ten_J_colind + J_rowadr; + mjtNum* ten_J = d->ten_J + J_rowadr; - // M[i,:] += armature * ten_J[i] * ten_J - int start = M_rowadr[i]; - int end = start + M_rownnz[i]; - for (int adr = start; adr < end; adr++) { - d->M[adr] += armature * ten_J_i * ten_J[M_colind[adr]]; - } + // M += armature * ten_J' * ten_J + for (int j=0; j < J_rownnz; j++) { + mjtNum ten_J_i = ten_J[j]; + if (!ten_J_i) { + continue; } - } - // sparse - else { - // get sparse info for tendon k - int J_rowadr = d->ten_J_rowadr[k]; - int J_rownnz = d->ten_J_rownnz[k]; - const int* J_colind = d->ten_J_colind + J_rowadr; - mjtNum* ten_J = d->ten_J + J_rowadr; - - // M += armature * ten_J' * ten_J - for (int j=0; j < J_rownnz; j++) { - mjtNum ten_J_i = ten_J[j]; - if (!ten_J_i) { - continue; - } - - // M[i,:] += armature * ten_J[i] * ten_J - int i = J_colind[j]; - int M_adr = M_rowadr[i]; - mju_addToSclSparseInc(d->M + M_adr, ten_J, - M_rownnz[i], M_colind + M_adr, - J_rownnz, J_colind, armature * ten_J_i); - } + // M[i,:] += armature * ten_J[i] * ten_J + int i = J_colind[j]; + int M_adr = M_rowadr[i]; + mju_addToSclSparseInc(d->M + M_adr, ten_J, + M_rownnz[i], M_colind + M_adr, + J_rownnz, J_colind, armature * ten_J_i); } } } @@ -2686,7 +2618,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { // add bias force due to tendon armature void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; - int ntendon = m->ntendon, nv = m->nv, issparse = mj_isSparse(m); + int ntendon = m->ntendon, nv = m->nv; mjtNum* ten_Jdot = NULL; mj_markStack(d); @@ -2716,20 +2648,13 @@ void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { mjtNum coef = armature * mju_dot(ten_Jdot, d->qvel, nv); if (coef) { - // dense - if (!issparse) { - mju_addToScl(qfrc, d->ten_J + nv*i, coef, nv); - } - // sparse - else { - int nnz = d->ten_J_rownnz[i]; - int adr = d->ten_J_rowadr[i]; - const int* colind = d->ten_J_colind + adr; - const mjtNum* ten_J = d->ten_J + adr; - for (int j=0; j < nnz; j++) { - qfrc[colind[j]] += coef * ten_J[j]; - } + int nnz = d->ten_J_rownnz[i]; + int adr = d->ten_J_rowadr[i]; + const int* colind = d->ten_J_colind + adr; + const mjtNum* ten_J = d->ten_J + adr; + for (int j=0; j < nnz; j++) { + qfrc[colind[j]] += coef * ten_J[j]; } } } diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 6c37710f..67b0d02e 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -1775,12 +1775,8 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { continue; } - // add sparse or dense - if (mj_isSparse(m)) { - addJTBJSparse(m, d, d->ten_J, &B, 1, i, d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind); - } else { - addJTBJ(m, d, d->ten_J+i*nv, &B, 1); - } + // add sparse + addJTBJSparse(m, d, d->ten_J, &B, 1, i, d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind); } } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 5b336890..6491cdca 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -229,13 +229,9 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) { mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL); - // tendon velocity: dense or sparse - if (mj_isSparse(m)) { - mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, - d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL); - } else { - mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv); - } + // tendon velocity: always sparse + mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, + d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL); // actuator velocity: always sparse if (!mjDISABLED(mjDSBL_ACTUATION)) { diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index d52e478c..5e26c2c3 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -119,7 +119,6 @@ static void mj_springdamper(const mjModel* m, mjData* d) { int nv = m->nv, ntendon = m->ntendon; int has_spring = !mjDISABLED(mjDSBL_SPRING); int has_damping = !mjDISABLED(mjDSBL_DAMPER); - int issparse = mj_isSparse(m); int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; int nbody = sleep_filter ? d->nbody_awake : m->nbody; @@ -472,20 +471,15 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // compute damper linear force along tendon mjtNum frc_damper = -damping * d->ten_velocity[i]; - // transform to joint torque, add to qfrc_{spring, damper}: dense or sparse - if (issparse) { - if (frc_spring || frc_damper) { - int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; - for (int j=d->ten_J_rowadr[i]; j < end; j++) { - int k = d->ten_J_colind[j]; - mjtNum J = d->ten_J[j]; - d->qfrc_spring[k] += J * frc_spring; - d->qfrc_damper[k] += J * frc_damper; - } + // transform to joint torque, add to qfrc_{spring, damper} + if (frc_spring || frc_damper) { + int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; + for (int j=d->ten_J_rowadr[i]; j < end; j++) { + int k = d->ten_J_colind[j]; + mjtNum J = d->ten_J[j]; + d->qfrc_spring[k] += J * frc_spring; + d->qfrc_damper[k] += J * frc_damper; } - } else { - if (frc_spring) mju_addToScl(d->qfrc_spring, d->ten_J+i*nv, frc_spring, nv); - if (frc_damper) mju_addToScl(d->qfrc_damper, d->ten_J+i*nv, frc_damper, nv); } } } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 5cee8f49..0a927715 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1441,16 +1441,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format); printArray2d("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format); - if (!mj_isSparse(m)) { - printArray2d("TEN_MOMENT", m->ntendon, m->nv, d->ten_J, fp, float_format); - } else { - mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL, - d->ten_J_rownnz, NULL, d->ten_J_colind, fp); - printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp); - printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp); - printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz, - d->ten_J_rowadr, d->ten_J_colind, fp, float_format); - } + mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL, + d->ten_J_rownnz, NULL, d->ten_J_colind, fp); + printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp); + printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp); + printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz, + d->ten_J_rowadr, d->ten_J_colind, fp, float_format); for (int i=0; i < m->ntendon; i++) { fprintf(fp, "TENDON %d: %d wrap points\n", i, d->ten_wrapnum[i]); for (int j=0; j < d->ten_wrapnum[i]; j++) { diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 7b5fe6bb..435291a4 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -756,16 +756,7 @@ static void set0(mjModel* m, mjData* d) { // compute tendon_invweight0 for (int i=0; i < m->ntendon; i++) { - // make dense vector into tmp - if (mj_isSparse(m)) { - mju_zero(tmp, nv); - int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; - for (int j=d->ten_J_rowadr[i]; j < end; j++) { - tmp[d->ten_J_colind[j]] = d->ten_J[j]; - } - } else { - mju_copy(tmp, d->ten_J+i*nv, nv); - } + mju_sparse2dense(tmp, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); // solve into tmp+nv mj_solveM(m, d, tmp+nv, tmp, 1); diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 53f36b38..7b292bfe 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -253,15 +253,11 @@ TEST_F(CoreSmoothTest, TendonArmature) { // add tendon inertias to M2 using outer product for (int j=0; j < m->ntendon; j++) { // get tendon Jacobian - if (mj_isSparse(m)) { - int rowadr = d->ten_J_rowadr[j]; - int* rownnz = d->ten_J_rownnz + j; - int zero = 0; - mju_sparse2dense(ten_J.data(), d->ten_J + rowadr, 1, nv, - rownnz, &zero, d->ten_J_colind + rowadr); - } else { - mju_copy(ten_J.data(), d->ten_J + j*nv, nv); - } + int rowadr = d->ten_J_rowadr[j]; + int* rownnz = d->ten_J_rownnz + j; + int zero = 0; + mju_sparse2dense(ten_J.data(), d->ten_J + rowadr, 1, nv, + rownnz, &zero, d->ten_J_colind + rowadr); // get tendon inertia only, using outer product mju_mulMatMat(ten_M.data(), ten_J.data(), ten_J.data(), nv, 1, nv); From 84431fbf7e6381519b7331e3a12550f7abc35a2e Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 19 Feb 2026 11:07:41 -0800 Subject: [PATCH 06/48] Improved upper bound for `nJten` PiperOrigin-RevId: 872479828 Change-Id: I750d56320d6ace6145475d378cb9d0c30919438d --- src/user/user_model.cc | 42 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/src/user/user_model.cc b/src/user/user_model.cc index c2f19df6..2bddb3a8 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -3231,9 +3231,45 @@ 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; + std::vector dof_bitmap(nv, false); + int count = 0; + for (int i = 0; i < ntendon; i++) { + int adr = m->tendon_adr[i]; + int num = m->tendon_num[i]; + + if (m->wrap_type[adr] == mjWRAP_JOINT) { + count += num; + continue; + } + + std::fill(dof_bitmap.begin(), dof_bitmap.end(), false); + for (int j = 0; j < num; j++) { + int type = m->wrap_type[adr + j]; + int bodyid = -1; + if (type == mjWRAP_SITE) { + bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; + } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { + bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; + } + if (bodyid > 0) { + int bid = bodyid; + while (bid > 0) { + int bdofadr = m->body_dofadr[bid]; + int bdofnum = m->body_dofnum[bid]; + for (int k = 0; k < bdofnum; k++) { + dof_bitmap[bdofadr + k] = true; + } + bid = m->body_parentid[bid]; + } + } + } + + // only count unique dofs + for (int j = 0; j < nv; j++) { + count += dof_bitmap[j]; + } + } + return count; } From 62a32386d6d37429f3583bc7e7433fc45031b387 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Feb 2026 13:13:41 -0800 Subject: [PATCH 07/48] Add pmap support for mjx-warp render. PiperOrigin-RevId: 872533248 Change-Id: I828303866c31d33d2367dd9fba4eb42028c67540 --- mjx/mujoco/mjx/_src/io.py | 8 +++- mjx/mujoco/mjx/_src/render_util.py | 4 +- mjx/mujoco/mjx/_src/render_util_test.py | 8 ++-- mjx/mujoco/mjx/warp/bvh.py | 2 +- mjx/mujoco/mjx/warp/collision_driver.py | 1 + mjx/mujoco/mjx/warp/forward.py | 1 + mjx/mujoco/mjx/warp/io.py | 31 ++++++++++--- mjx/mujoco/mjx/warp/render.py | 4 +- mjx/mujoco/mjx/warp/smooth.py | 1 + mjx/mujoco/mjx/warp/types.py | 6 ++- mjx/mujoco/mjx/warp/visualize_render.py | 61 ++++++++++++++++++++++++- 11 files changed, 108 insertions(+), 19 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 8feef14d..4199936a 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1943,6 +1943,7 @@ def set_state( def create_render_context( mjm: mujoco.MjModel, nworld: int, + devices: Optional[Sequence[str]] = None, **kwargs, ): """Creates a render context. @@ -1953,6 +1954,9 @@ def create_render_context( because Warp creates arrays of size nworld that are not exposed to JAX. Thus we cannot use JAX transforms like vmap with the render context. + devices: optional list of device names (e.g. ['cuda:0', 'cuda:1']). + If provided, rendering workloads are sharded across these devices. + By default, devices is None and the default device from wp.get_device(None) is used. **kwargs: forwarded to the render context constructor. Returns: @@ -1960,4 +1964,6 @@ def create_render_context( """ _check_warp_installed() from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error - return mjxw_io.create_render_context(mjm, nworld=nworld, **kwargs) + return mjxw_io.create_render_context( + mjm, nworld=nworld, devices=devices, **kwargs + ) diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index 58e98ef4..cfe6bf0b 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -44,7 +44,7 @@ def get_rgb( else: raise RuntimeError('Warp not installed.') - warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key] + warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)] rgb_adr = int(warp_rc.rgb_adr.numpy()[cam_id]) width = int(warp_rc.cam_res.numpy()[cam_id][0]) height = int(warp_rc.cam_res.numpy()[cam_id][1]) @@ -84,7 +84,7 @@ def get_depth( import mujoco.mjx.warp.render as mjxw_render # pylint: disable=g-import-not-at-top # pytype: disable=import-error else: raise RuntimeError('Warp not installed.') - warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key] + warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)] depth_adr = int(warp_rc.depth_adr.numpy()[cam_id]) width = int(warp_rc.cam_res.numpy()[cam_id][0]) height = int(warp_rc.cam_res.numpy()[cam_id][1]) diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py index e62658ac..ed82f5c0 100644 --- a/mjx/mujoco/mjx/_src/render_util_test.py +++ b/mjx/mujoco/mjx/_src/render_util_test.py @@ -56,7 +56,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))(rc, 0, rgb_data) @@ -70,7 +70,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): rgb = jax.jit( jax.vmap(render_util.get_rgb, in_axes=(None, None, 0)), @@ -87,7 +87,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))( rc, 0, depth_data, 5.0 @@ -103,7 +103,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): depth = jax.jit( jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None)), diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 573c24e4..a86bad63 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -88,7 +88,7 @@ def _refit_bvh_shim( _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos _d.nworld = nworld - render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + render_context = _MJX_RENDER_CONTEXT_BUFFERS[(rc_id, wp.get_device().ordinal)] dummy.zero_() mjwarp.refit_bvh(_m, _d, render_context) diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 35a85720..aa88f118 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _collision_shim( # Model diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 349e5ed2..db52a8cd 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _forward_shim( # Model diff --git a/mjx/mujoco/mjx/warp/io.py b/mjx/mujoco/mjx/warp/io.py index eb9af42c..db0f2df9 100644 --- a/mjx/mujoco/mjx/warp/io.py +++ b/mjx/mujoco/mjx/warp/io.py @@ -19,26 +19,43 @@ import threading import mujoco from mujoco.mjx.warp.types import RenderContext import mujoco.mjx.third_party.mujoco_warp as mjw +import warp as wp _MJX_RENDER_CONTEXT_COUNTER = 0 _MJX_RENDER_CONTEXT_LOCK = threading.Lock() _MJX_RENDER_CONTEXT_BUFFERS = {} +def _create_context(mjm, nworld, device, **kwargs): + with wp.ScopedDevice(device): + ctx = mjw.create_render_context(mjm=mjm, nworld=nworld, **kwargs) + ctx.rgb_data_shape = ctx.rgb_data.shape + ctx.depth_data_shape = ctx.depth_data.shape + ctx.rgb_data = None + ctx.depth_data = None + return ctx + + def create_render_context( mjm: mujoco.MjModel, nworld: int, + devices: list[str | None] | None = None, **kwargs, ): - rc = mjw.create_render_context(mjm=mjm, nworld=nworld, **kwargs) - rc.rgb_data_shape = rc.rgb_data.shape - rc.depth_data_shape = rc.depth_data.shape - rc.rgb_data = None - rc.depth_data = None - global _MJX_RENDER_CONTEXT_COUNTER + + if not devices: + devices = [None] + + contexts = [_create_context(mjm, nworld, d, **kwargs) for d in devices] + with _MJX_RENDER_CONTEXT_LOCK: _MJX_RENDER_CONTEXT_COUNTER += 1 key = _MJX_RENDER_CONTEXT_COUNTER - _MJX_RENDER_CONTEXT_BUFFERS[key] = rc + for d, ctx in zip(devices, contexts): + ordinal = wp.get_device(d).ordinal + _MJX_RENDER_CONTEXT_BUFFERS[(key, ordinal)] = ctx + if (key, None) not in _MJX_RENDER_CONTEXT_BUFFERS: + # save the first context as the default context + _MJX_RENDER_CONTEXT_BUFFERS[(key, None)] = contexts[0] return RenderContext(key, _owner=True) diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 091f72d0..73ced206 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -110,14 +110,14 @@ def _render_shim( _d.light_xdir = light_xdir _d.light_xpos = light_xpos _d.nworld = nworld - render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + render_context = _MJX_RENDER_CONTEXT_BUFFERS[(rc_id, wp.get_device().ordinal)] render_context.rgb_data = rgb render_context.depth_data = depth mjwarp.render(_m, _d, render_context) def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): - render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[ctx.key] + render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[(ctx.key, None)] output_dims = { 'rgb': render_ctx.rgb_data_shape, 'depth': render_ctx.depth_data_shape, diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 217cf873..6e1c0ea8 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _kinematics_shim( # Model diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index d39d4861..03d403a4 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -126,7 +126,11 @@ class RenderContext: if lock is None or buffers is None: return with lock: - buffers.pop(self.key, None) + keys_to_remove = [ + k for k in buffers.keys() if isinstance(k, tuple) and k[0] == self.key + ] + for k in keys_to_remove: + buffers.pop(k, None) class StatisticWarp(PyTreeNode): diff --git a/mjx/mujoco/mjx/warp/visualize_render.py b/mjx/mujoco/mjx/warp/visualize_render.py index 42ed02f1..6a3b5afa 100644 --- a/mjx/mujoco/mjx/warp/visualize_render.py +++ b/mjx/mujoco/mjx/warp/visualize_render.py @@ -57,6 +57,9 @@ _WP_KERNEL_CACHE_DIR = flags.DEFINE_string( '/tmp/wp_kernel_cache_dir_visualize_render', 'warp kernel cache directory', ) +_PMAP = flags.DEFINE_boolean( + 'pmap', False, 'also render with pmap across GPUs and compare' +) _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1} jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS) @@ -107,6 +110,7 @@ def _main(_: Sequence[str]): print(f' camera_id : {_CAMERA_ID.value}') print(f' use_textures: {_USE_TEXTURES.value}') print(f' use_shadows : {_USE_SHADOWS.value}') + print(f' pmap : {_PMAP.value}') print(f' output_dir : {_OUTPUT_DIR.value}\n') mx = mjx.put_model(m, impl='warp') @@ -143,7 +147,6 @@ def _main(_: Sequence[str]): enabled_geom_groups=[0, 1, 2], ) - print('rendering...') dx_batch = jax_jit(jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None)))( mx, dx_batch, rc ) @@ -187,6 +190,62 @@ def _main(_: Sequence[str]): ) _save_tiled(depth_rgb, depth_tiled_path) + if _PMAP.value: + ndevices = jax.local_device_count() + nworld = _NWORLD.value + nworld_per_device = nworld // ndevices + assert nworld >= ndevices and nworld % ndevices == 0, ( + f'--pmap requires nworld ({nworld}) divisible by device count' + f' ({ndevices})' + ) + print(f'\nrendering (pmap across {ndevices} devices)...') + + device_strs = [f'cuda:{i}' for i in range(ndevices)] + + pmap_rc = io.create_render_context( + mjm=m, + nworld=nworld_per_device, + devices=device_strs, + cam_res=(_WIDTH.value, _HEIGHT.value), + use_textures=_USE_TEXTURES.value, + use_shadows=_USE_SHADOWS.value, + render_rgb=True, + render_depth=True, + enabled_geom_groups=[0, 1, 2], + ) + + devices = jax.local_devices()[:ndevices] + mesh = jax.sharding.Mesh(np.array(devices), axis_names=('i',)) + P = jax.sharding.PartitionSpec + sharded = jax.sharding.NamedSharding(mesh, P('i')) + + def safe_shard(x, sharding): + # Go through CPU to avoid P2P DMA issues on certain machines. + x_cpu = jax.device_put(x, jax.devices('cpu')[0]) + if x_cpu.ndim > 0 and x_cpu.shape[0] == nworld: + reshaped = x_cpu.reshape(ndevices, nworld_per_device, *x_cpu.shape[1:]) + else: + reshaped = jp.stack([x_cpu] * ndevices) + return jax.device_put(reshaped, sharding) + + dx_pmap = jax.tree.map(lambda x: safe_shard(x, sharded), dx_batch) + mx_pmap = jax.tree.map(lambda x: safe_shard(x, sharded), mx) + + def inner(mx, dx): + dx = bvh.refit_bvh(mx, dx, pmap_rc) + out = render.render(mx, dx, pmap_rc) + return render_util.get_rgb(pmap_rc, _CAMERA_ID.value, out[0]) + + inner = jax.vmap(inner, in_axes=(None, 0)) + out = jax.pmap(inner)(mx_pmap, dx_pmap) + + pmap_rgb = jax.device_put(out, jax.devices('cpu')[0]).reshape(-1, *out.shape[2:]) + + pmap_tiled_path = os.path.join( + _OUTPUT_DIR.value, f'pmap_tiled_{_CAMERA_ID.value}.png' + ) + _save_tiled(pmap_rgb, pmap_tiled_path) + print('\ndone.') From 86bca81d18cc4612615b42e079694cbe5632a3e4 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Feb 2026 13:17:23 -0800 Subject: [PATCH 08/48] Update mjx docs and changelog. PiperOrigin-RevId: 872535000 Change-Id: I51b05a22ac553c34b436de63f43c7efca6607724 --- doc/changelog.rst | 9 +++++++++ doc/mjx.rst | 29 ++++++++++++++++++++++++++--- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 79a1929e..eeb6fab0 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,6 +2,15 @@ Changelog ========= +Upcoming version (not yet released) +----------------------------------- + +MJX +^^^ + +- Add batch rendering support for MJX-Warp. See the :ref:`MJX-Warp batch rendering` section for details. + + Version 3.5.0 (February 12, 2026) --------------------------------- diff --git a/doc/mjx.rst b/doc/mjx.rst index 8825d869..a68d57b1 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -207,6 +207,8 @@ excessive graph captures in the JAX-Warp FFI layer. - 0.65M +.. _MjxWarpBatchRendering: + MJX-Warp Batch Rendering ~~~~~~~~~~~~~~~~~~~~~~~~ @@ -218,9 +220,9 @@ Note that the number of parallel worlds (``nworld``) is fixed when creating the .. code-block:: python - from mujoco.mjx import io + from mujoco.mjx import create_render_context - rc = io.create_render_context( + rc = create_render_context( mjm=m, nworld=nworld, cam_res=(width, height), @@ -247,12 +249,33 @@ volume hierarchy (BVH) and executing the raycaster: pixels, _ = mjx.render(mx, d, rc) # 3. Extract the RGB tensor for the first camera (index 0) - rgb = get_rgb(rc, pixels, 0) + rgb = get_rgb(rc, 0, pixels) # CAVEAT: Always return or use the updated `d` in your computation graph. # Otherwise, JAX's dead-code elimination will optimize away the refit_bvh call! return rgb, d +Multi-GPU rendering with ``pmap`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To render across multiple GPUs, create a render context **per device** by passing ``devices`` to +:func:`create_render_context `. + +.. code-block:: python + + ndevices = jax.local_device_count() + nworld_per_device = nworld // ndevices + + # Create one render context for all devices + rc = create_render_context( + mjm=m, + nworld=nworld_per_device, + devices=[f'cuda:{i}' for i in range(ndevices)], + cam_res=(width, height), + ) + +Then use ``jax.pmap`` to parallelize the rendering across devices. See the complete example in +`visualize_render.py `__. .. _MjxJAX: From fe50efd41a3c3dd50352b1626141e2ab0522994e Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 19 Feb 2026 13:51:54 -0800 Subject: [PATCH 09/48] Sparse tendon Jacobian update for MuJoCo Warp PiperOrigin-RevId: 872550625 Change-Id: Ifc5afacb048bec36ef6455072b188396db958002 --- mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) 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 9ee93d3d..05bd1d2b 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -952,13 +952,9 @@ def put_data( d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float) - if mujoco.mj_isSparse(mjm): - ten_J = np.zeros((mjm.ntendon, mjm.nv)) - mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1)) - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) - else: - ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv)) - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) + ten_J = np.zeros((mjm.ntendon, mjm.nv)) + mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1)) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) From 8ba68cee3e28f266cfbe69340f529c77fdde766d Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Thu, 19 Feb 2026 15:31:32 -0800 Subject: [PATCH 10/48] Limit the number of flex contacts per collision pair. This change introduces a filter for flex-related contacts, ensuring that no more than mjMAXCONPAIR contacts are kept for each geom-flex, flex-flex, flex internal, and flex self-collision pair. The contacts are sorted using farthest-point sampling: - Starts with the deepest penetrating contact - Iteratively selects the contact farthest from already-selected contacts - Produces a spatially distributed set of contacts PiperOrigin-RevId: 872593388 Change-Id: I50b233a0dc66da6a297852c258ba514f131c8f23 --- src/engine/engine_collision_driver.c | 90 ++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index ac12d0e4..ceaac7ef 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -259,6 +259,78 @@ static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* mjSORT(contactSort, mjContact, contactcompare); +// filter flex contacts based on distance +static void filterFlexContacts(mjData* d, int ncon_before) { + int n = d->ncon - ncon_before; + if (n <= mjMAXCONPAIR) { + return; + } + + mjContact* contacts = d->contact + ncon_before; + + mj_markStack(d); + mjtByte* selected = mjSTACKALLOC(d, n, mjtByte); + mjtNum* min_dist = mjSTACKALLOC(d, n, mjtNum); + memset(selected, 0, n); + + for (int i = 0; i < n; i++) { + min_dist[i] = mjMAXVAL; + } + + // start with the deepest penetrating contact + int nselected = 0; + int best = 0; + mjtNum bestdist = -contacts[0].dist; + for (int i = 1; i < n; i++) { + if (-contacts[i].dist > bestdist) { + bestdist = -contacts[i].dist; + best = i; + } + } + + while (nselected < mjMAXCONPAIR && best >= 0) { + selected[best] = 1; + mjtNum* bestpos = contacts[best].pos; + + int nextbest = -1; + mjtNum nextbestdist = -1; + for (int i = 0; i < n; i++) { + if (selected[i]) continue; + + mjtNum dx = contacts[i].pos[0] - bestpos[0]; + mjtNum dy = contacts[i].pos[1] - bestpos[1]; + mjtNum dz = contacts[i].pos[2] - bestpos[2]; + mjtNum d2 = dx*dx + dy*dy + dz*dz; + if (d2 < min_dist[i]) { + min_dist[i] = d2; + } + if (min_dist[i] > nextbestdist) { + nextbestdist = min_dist[i]; + nextbest = i; + } + } + + if (nselected < mjMAXCONPAIR - 1) { + mjContact temp = contacts[nselected]; + contacts[nselected] = contacts[best]; + contacts[best] = temp; + + if (nextbest == nselected) { + nextbest = best; + } + } + + nselected++; + best = nextbest; + } + + mj_freeStack(d); + + d->ncon = ncon_before + nselected; + resetArena(d); +} + + // main collision function void mj_collision(const mjModel* m, mjData* d) { @@ -362,6 +434,12 @@ void mj_collision(const mjModel* m, mjData* d) { mj_collideTree(m, d, bf1, bf2, merged, startadr, pairadr); int ncon_after = d->ncon; + // filter flex contacts (limit per geom-flex or flex-flex pair) + if (bf1 >= nbody || bf2 >= nbody) { + filterFlexContacts(d, ncon_before); + ncon_after = d->ncon; + } + // sort contacts int n = ncon_after - ncon_before; if (n > 1) { @@ -400,15 +478,19 @@ void mj_collision(const mjModel* m, mjData* d) { // plane special processing if (m->geom_type[g] == mjGEOM_PLANE) { + int ncon_before = d->ncon; mj_collidePlaneFlex(m, d, g, f); + filterFlexContacts(d, ncon_before); continue; } // collide geom with flex elements + int ncon_before = d->ncon; int elemnum = m->flex_elemnum[f]; for (int e=0; e < elemnum; e++) { mj_collideGeomElem(m, d, g, f, e); } + filterFlexContacts(d, ncon_before); } } @@ -418,11 +500,13 @@ void mj_collision(const mjModel* m, mjData* d) { int f2 = bf2 - nbody; // collide elements of two flexes + int ncon_before = d->ncon; for (int e1=0; e1 < m->flex_elemnum[f1]; e1++) { for (int e2=0; e2 < m->flex_elemnum[f2]; e2++) { mj_collideElems(m, d, f1, e1, f2, e2); } } + filterFlexContacts(d, ncon_before); } } } @@ -439,11 +523,15 @@ void mj_collision(const mjModel* m, mjData* d) { if (!m->flex_rigid[f] && (m->flex_contype[f] & m->flex_conaffinity[f])) { // internal collisions if (m->flex_internal[f]) { + int ncon_before = d->ncon; mj_collideFlexInternal(m, d, f); + filterFlexContacts(d, ncon_before); } // active element collisions if (m->flex_selfcollide[f] != mjFLEXSELF_NONE) { + int ncon_before = d->ncon; + // element-element: midphase if (!mjDISABLED(mjDSBL_MIDPHASE) && m->flex_selfcollide[f] != mjFLEXSELF_NARROW && @@ -470,6 +558,8 @@ void mj_collision(const mjModel* m, mjData* d) { } } } + + filterFlexContacts(d, ncon_before); } } } From ca0552c83b1ffb6d92f43388e9d4439cc3deb6e9 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Thu, 19 Feb 2026 16:33:52 -0800 Subject: [PATCH 11/48] Cleanup Studio toolbar PiperOrigin-RevId: 872618794 Change-Id: I2e8f99dc47937416c8516ed1df7821369211b12f --- src/experimental/platform/imgui_widgets.h | 3 + src/experimental/studio/app.cc | 365 ++++++++-------------- 2 files changed, 140 insertions(+), 228 deletions(-) diff --git a/src/experimental/platform/imgui_widgets.h b/src/experimental/platform/imgui_widgets.h index e8988369..29f5005f 100644 --- a/src/experimental/platform/imgui_widgets.h +++ b/src/experimental/platform/imgui_widgets.h @@ -28,11 +28,14 @@ namespace mujoco::platform { // FontAwesome icon codes. +static constexpr const char ICON_FA_ADJUST[] = "\xEF\x81\x82"; static constexpr const char ICON_FA_ARROWS[] = "\xEF\x81\x87"; static constexpr const char ICON_FA_CAMERA[] = "\xEF\x80\xBD"; static constexpr const char ICON_FA_CARET_LEFT[] = "\xEF\x83\x99"; static constexpr const char ICON_FA_CARET_RIGHT[] = "\xEF\x83\x9A"; static constexpr const char ICON_FA_CHECK_SQUARE_O[] = "\xEF\x81\x9D"; +static constexpr const char ICON_FA_CIRCLE[] = "\xEF\x84\x91"; +static constexpr const char ICON_FA_CIRCLE_O[] = "\xEF\x84\x8C"; static constexpr const char ICON_FA_COMMENT[] = "\xEF\x83\xA5"; static constexpr const char ICON_FA_COPY[] = "\xEF\x83\x85"; static constexpr const char ICON_FA_DIAMOND[] = "\xEF\x88\x99"; diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 4c250f36..adef9a42 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -24,7 +24,6 @@ #include #include #include -#include #include #include #include @@ -78,14 +77,13 @@ static constexpr const char* ICON_LABEL = platform::ICON_FA_COMMENT; static constexpr const char* ICON_RESET_MODEL = platform::ICON_FA_UNDO; static constexpr const char* ICON_FRAME = platform::ICON_FA_ARROWS; static constexpr const char* ICON_CAMERA = platform::ICON_FA_CAMERA; -static constexpr const char* ICON_DARKMODE = platform::ICON_FA_MOON; -static constexpr const char* ICON_LIGHTMODE = platform::ICON_FA_SUN; -static constexpr const char* ICON_CLASSICMODE = platform::ICON_FA_DIAMOND; +static constexpr const char* ICON_DARKMODE = platform::ICON_FA_CIRCLE; +static constexpr const char* ICON_LIGHTMODE = platform::ICON_FA_CIRCLE_O; +static constexpr const char* ICON_CLASSICMODE = platform::ICON_FA_ADJUST; static constexpr const char* ICON_PREV_FRAME = platform::ICON_FA_CARET_LEFT; static constexpr const char* ICON_NEXT_FRAME = platform::ICON_FA_CARET_RIGHT; static constexpr const char* ICON_CURR_FRAME = platform::ICON_FA_FAST_FORWARD; static constexpr const char* ICON_SPEED = platform::ICON_FA_TACHOMETER; -static constexpr const char* ICON_DELETE = platform::ICON_FA_TRASH_CAN; // UI labels for mjtLabel. static constexpr const char* kLabelNames[] = { @@ -107,8 +105,7 @@ static constexpr std::array kPercentRealTime = { }; // clang-format on -App::App(Config config) - : rng_(std::random_device()()), ini_path_(std::move(config.ini_path)) { +App::App(Config config) : ini_path_(std::move(config.ini_path)) { platform::Window::Config window_config; window_config.renderer_backend = platform::Renderer::GetBackend(); window_config.offscreen_mode = config.offscreen_mode; @@ -135,13 +132,6 @@ void App::ClearModel() { step_error_ = ""; } -void App::Recompile() { - mj_recompile(model_holder_->spec(), model_holder_->vfs(), - model_holder_->model(), model_holder_->data()); - const int state_size = mj_stateSize(model(), mjSTATE_INTEGRATION); - history_.Init(state_size); -} - void App::RequestModelLoad(std::string model_file) { pending_load_ = std::move(model_file); } @@ -149,7 +139,6 @@ void App::RequestModelLoad(std::string model_file) { void App::RequestModelReload() { if (model_kind_ == kModelFromFile) { pending_load_ = model_path_; - preserve_camera_on_load_ = true; } } @@ -200,16 +189,6 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { const int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); history_.Init(state_size); - if (!preserve_camera_on_load_) { - const int model_cam = model->vis.global.cameraid; - if (model_cam >= 0 && model_cam < model->ncam) { - ui_.camera_idx = platform::SetCamera(model, &camera_, model_cam); - } else { - mjv_defaultFreeCamera(model, &camera_); - } - } - preserve_camera_on_load_ = false; - // Initialize the speed based on the model's default real-time setting. float min_error = FLT_MAX; const float desired = mju_log(100 * model->vis.global.realtime); @@ -386,11 +365,6 @@ void App::ProcessPendingLoads() { } } - if (spec_op_) { - spec_op_(); - spec_op_ = nullptr; - } - // Check plugins to see if we need to load a new model. platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { if (plugin->get_model_to_load) { @@ -408,19 +382,6 @@ void App::ProcessPendingLoads() { }); } -void App::SpecDeleteSelectedElement() { - spec_op_ = [this]() { - mjs_delete(spec(), tmp_.element); - if (tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select == tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - } - tmp_.element = nullptr; - tmp_.element_id = -1; - Recompile(); - }; -} - void App::HandleWindowEvents() { const std::string drop_file = window_->GetDropFile(); if (!drop_file.empty()) { @@ -513,28 +474,15 @@ void App::HandleMouseEvents() { perturb_.flexselect = picked.flex; perturb_.skinselect = picked.skin; - // Select the corresponding element in the spec. - tmp_.element = nullptr; - tmp_.element_id = -1; - if (has_spec()) { - mjsElement* element = mjs_firstElement(spec(), mjOBJ_BODY); - while (element) { - if (mjs_getId(element) == picked.body) { - tmp_.element = element; - tmp_.element_id = picked.body; - break; - } - element = mjs_nextElement(spec(), element); - } - } - // Compute the local position of the selected object in the world. mjtNum tmp[3]; mju_sub3(tmp, picked.point, data()->xpos + 3 * picked.body); mju_mulMatTVec(perturb_.localpos, data()->xmat + 9 * picked.body, tmp, 3, 3); } else { - mjv_defaultPerturb(&perturb_); + perturb_.select = 0; + perturb_.flexselect = -1; + perturb_.skinselect = -1; } } @@ -623,8 +571,6 @@ void App::HandleKeyboardEvents() { } } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) { ResetPhysics(); - } else if (ImGui_IsChordJustPressed(ImGuiKey_Delete)) { - SpecDeleteSelectedElement(); } else if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) { SelectParentPerturb(model(), perturb_); } else if (ImGui_IsChordJustPressed(ImGuiKey_F1)) { @@ -715,61 +661,6 @@ void App::HandleKeyboardEvents() { ToggleFlag(vis_options_.geomgroup[4]); } else if (ImGui_IsChordJustPressed(ImGuiKey_5)) { ToggleFlag(vis_options_.geomgroup[5]); - } else if (ImGui_IsChordJustPressed(ImGuiKey_Enter | ImGuiMode_CtrlShift)) { - if (has_spec()) { - spec_op_ = [this]() { - mjsBody* world = mjs_findBody(spec(), "world"); - if (!world) return; - mjsBody* body = mjs_addBody(world, nullptr); - if (!body) return; - mjsJoint* joint = mjs_addJoint(body, nullptr); - if (!joint) return; - mjsGeom* geom = mjs_addGeom(body, nullptr); - if (!geom) return; - - // Set body position slightly in front of the camera. - mjtNum pos[3]; - mjtNum dir[3]; - mjtNum up[3]; - mjv_cameraFrame(pos, dir, up, nullptr, data(), &camera_); - - static int counter = 0; - std::string name = "projectile" + std::to_string(counter++); - mjs_setName(body->element, name.c_str()); - - body->mass = 10.0; - body->pos[0] = pos[0] + dir[0] * 0.2; - body->pos[1] = pos[1] + dir[1] * 0.2; - body->pos[2] = pos[2] + dir[2] * 0.2; - geom->type = mjGEOM_BOX; - geom->size[0] = 0.13365; - geom->size[1] = 0.13365; - geom->size[2] = 0.13365; - geom->rgba[0] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[1] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[2] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[3] = 1.0; - - joint->type = mjJNT_FREE; - - Recompile(); - - // Give the newly added body a velocity in the direction of the camera. - int bodyid = mj_name2id(model(), mjOBJ_BODY, name.c_str()); - if (bodyid >= 0) { - int jntid = model()->body_jntadr[bodyid]; - if (jntid >= 0 && model()->jnt_type[jntid] == mjJNT_FREE) { - int qveladr = model()->jnt_dofadr[jntid]; - if (qveladr >= 0) { - mjtNum speed = 10.0; // Magnitude of the initial velocity. - data()->qvel[qveladr + 0] = dir[0] * speed + up[0]; - data()->qvel[qveladr + 1] = dir[1] * speed + up[1]; - data()->qvel[qveladr + 2] = dir[2] * speed + up[2]; - } - } - } - }; - } } else if (has_model()) { if (ImGui_IsChordJustPressed(ImGuiKey_Escape)) { ui_.camera_idx = @@ -891,10 +782,16 @@ void App::BuildGui() { MainMenuGui(); - if (ImGui::Begin("ToolBar")) { - ToolBarGui(); + { + ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(0, 0)); + if (ImGui::Begin("ToolBar")) { + ImGui::PopStyleVar(); + ToolBarGui(); + } else { + ImGui::PopStyleVar(); + } + ImGui::End(); } - ImGui::End(); { platform::ScopedStyle style; @@ -1191,8 +1088,7 @@ void App::SpecExplorerGui() { label = "(" + prefix + " " + std::to_string(id) + ")"; } - const bool selected = (tmp_.element == element); - if (ImGui::Selectable(label.c_str(), selected)) { + if (ImGui::Selectable(label.c_str(), false)) { tmp_.element = element; tmp_.element_id = id; } @@ -1202,24 +1098,42 @@ void App::SpecExplorerGui() { }; if (ImGui::TreeNodeEx("Bodies", flags)) { - display_group(mjOBJ_BODY, "Body"); + // We don't use `display_group` here because we do additional selection + // logic tied to the `perturb_` field. + mjsElement* element = mjs_firstElement(spec(), mjOBJ_BODY); + while (element) { + const int id = mjs_getId(element); + + const mjString* name = mjs_getName(element); + std::string label = *name; + if (label.empty()) { + label = "(Body " + std::to_string(id) + ")"; + } + + if (ImGui::Selectable(label.c_str(), (id == perturb_.select), + ImGuiSelectableFlags_AllowDoubleClick)) { + tmp_.element = element; + tmp_.element_id = id; + } + if (ImGui::IsItemHovered() && + ImGui::IsMouseDoubleClicked(ImGuiMouseButton_Left)) { + perturb_.select = id; + } + + element = mjs_nextElement(spec(), element); + } ImGui::TreePop(); } + if (ImGui::TreeNodeEx("Joints", flags)) { display_group(mjOBJ_JOINT, "Joint"); ImGui::TreePop(); } + if (ImGui::TreeNodeEx("Sites", flags)) { display_group(mjOBJ_SITE, "Site"); ImGui::TreePop(); } - - // If we selected a body, then select the same body for the perturb object. - if (tmp_.element && tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select != tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - perturb_.select = tmp_.element_id; - } } void App::PropertiesGui() { @@ -1228,31 +1142,22 @@ void App::PropertiesGui() { return; } - if (ImGui::BeginTable("##PropertiesHeader", 2)) { - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 20); - ImGui::TableNextColumn(); - ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); - ImGui::TableNextColumn(); - if (tmp_.element->elemtype == mjOBJ_BODY) { - if (ImGui::SmallButton(ICON_DELETE)) { - SpecDeleteSelectedElement(); - } - } - ImGui::EndTable(); - } - ImGui::Separator(); - switch (tmp_.element->elemtype) { case mjOBJ_BODY: + ImGui::Text("Body"); + ImGui::Separator(); platform::BodyPropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; case mjOBJ_JOINT: + ImGui::Text("Joint"); + ImGui::Separator(); platform::JointPropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; case mjOBJ_SITE: + ImGui::Text("Site"); + ImGui::Separator(); platform::SitePropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; @@ -1375,6 +1280,18 @@ void App::HelpGui() { ImGui::Columns(); } +struct SpeedStatus { + bool misaligned; + float measured; +}; + +static SpeedStatus IsSpeedMisaligned( + const platform::StepControl& step_control) { + const float desired = step_control.GetSpeed(); + const float measured = step_control.GetSpeedMeasured(); + return {std::abs(measured - desired) > 0.1f * desired, measured}; +} + void App::ToolBarGui() { if (ImGui::BeginTable("##ToolBarTable", 2)) { platform::ScopedStyle style; @@ -1384,22 +1301,30 @@ void App::ToolBarGui() { const int combo_flags = ImGuiComboFlags_NoArrowButton; const float scale = ImGui::GetWindowDpiScale(); - const float right_width = 520.f * scale; const ImVec2 button_size(48.f * scale, 32.f * scale); const ImVec2 play_button_size(80.f * scale, 32.f * scale); + const float label_width = GetExpectedLabelWidth(); + const float copy_btn_width = ImGui::CalcTextSize(ICON_COPY_CAMERA).x + + ImGui::GetStyle().FramePadding.x * 2; + const float theme_width = ImGui::CalcTextSize(ICON_LIGHTMODE).x + + ImGui::GetStyle().FramePadding.x * 2; + const float sp = ImGui::GetStyle().ItemSpacing.x; + const float right_width = label_width + sp + label_width + sp + + label_width + sp + copy_btn_width + sp + + theme_width; + const float separator_width = .2f * button_size.x; + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, right_width); ImGui::TableNextColumn(); - ImGui::Text("%s", " "); // Combined (Unload, Reload) widget { style.Var(ImGuiStyleVar_FrameRounding, 2.0f); // Unload button. - ImGui::SameLine(); { const ImColor a = red; const ImColor h(a.Value.x, a.Value.y, a.Value.z, a.Value.w * 0.6f); @@ -1421,25 +1346,19 @@ void App::ToolBarGui() { ImGui::SetItemTooltip("%s", "Reload"); } - ImGui::SameLine(0, 0); - ImGui::Text(" "); - // Reset button. - ImGui::SameLine(); + ImGui::SameLine(0, separator_width); if (ImGui::Button(ICON_RESET_MODEL, button_size)) { ResetPhysics(); } ImGui::SetItemTooltip("%s", "Reset"); - ImGui::SameLine(0, 0); - ImGui::Text(" "); - // Combined (Normal Pause, Viscous Pause, Play) widget { style.Var(ImGuiStyleVar_FrameRounding, 2.0f); // Normal pause button. - ImGui::SameLine(); + ImGui::SameLine(0, separator_width); ImColor paused_color = yellow; bool paused = step_control_.GetPauseState() == PauseState::kNormalPaused; if (platform::ImGui_ColorButton(ICON_PAUSE, paused, paused_color, @@ -1480,18 +1399,23 @@ void App::ToolBarGui() { } } - ImGui::SameLine(); - ImGui::Text("%s", " |"); - // Speed selection. ImGui::SameLine(); - ImGui::Text("%s", ICON_SPEED); - ImGui::SetItemTooltip("%s", "Playback Speed"); - - ImGui::SameLine(); - ImGui::SetNextItemWidth(50.0f * scale); - if (ImGui::BeginCombo("##Speed", kPercentRealTime[tmp_.speed_index], - combo_flags)) { + float pad_y = (button_size.y - ImGui::GetFontSize()) * 0.5f; + ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, + ImVec2(ImGui::GetStyle().FramePadding.x + 5.f, pad_y)); + const auto [misaligned, measured] = IsSpeedMisaligned(step_control_); + char speed_preview[64]; + if (misaligned) { + snprintf(speed_preview, sizeof(speed_preview), "%s%s (%-4.1f%%)", + ICON_SPEED, kPercentRealTime[tmp_.speed_index], measured); + } else { + snprintf(speed_preview, sizeof(speed_preview), "%s%s", ICON_SPEED, + kPercentRealTime[tmp_.speed_index]); + } + ImGui::SetNextItemWidth(ImGui::CalcTextSize(speed_preview).x + + ImGui::GetStyle().FramePadding.x * 2); + if (ImGui::BeginCombo("##Speed", speed_preview, combo_flags)) { for (int n = 0; n < kPercentRealTime.size(); n++) { if (ImGui::Selectable(kPercentRealTime[n], (tmp_.speed_index == n))) { SetSpeedIndex(n); @@ -1499,17 +1423,30 @@ void App::ToolBarGui() { } ImGui::EndCombo(); } - ImGui::SetItemTooltip("%s", "Playback Speed"); + ImGui::PopStyleVar(); + if (misaligned) { + ImGui::SetItemTooltip("%s", "Desired Speed (Measured Speed)"); + } else { + ImGui::SetItemTooltip("%s", "Desired Speed"); + } + + ImGui::TableNextColumn(); + ImGui::SetCursorPosY(ImGui::GetCursorPosY() + + (button_size.y - ImGui::GetFrameHeight()) * 0.5f); // Camera selection. - std::vector cameras = GetCameraNames(); - ImGui::TableNextColumn(); - ImGui::Text("%s", ICON_CAMERA); - ImGui::SetItemTooltip("%s", "Camera"); - ImGui::SameLine(); + if (ImGui::Button(ICON_COPY_CAMERA)) { + std::string camera_string = platform::CameraToString(data(), &camera_); + platform::MaybeSaveToClipboard(camera_string); + } + ImGui::SetItemTooltip("%s", "Copy Camera"); + ImGui::SameLine(0, 0); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); int camera_idx = ui_.camera_idx - platform::kTumbleCameraIdx; - if (ImGui::BeginCombo("##Camera", cameras[camera_idx], combo_flags)) { + std::vector cameras = GetCameraNames(); + std::string camera_preview = + std::string(ICON_CAMERA) + " " + cameras[camera_idx]; + if (ImGui::BeginCombo("##Camera", camera_preview.c_str(), combo_flags)) { for (int n = 0; n < cameras.size(); n++) { if (ImGui::Selectable(cameras[n], (camera_idx == n))) { ui_.camera_idx = platform::SetCamera(model(), &camera_, @@ -1519,25 +1456,13 @@ void App::ToolBarGui() { ImGui::EndCombo(); } ImGui::SetItemTooltip("%s", "Camera"); - ImGui::SameLine(); - if (ImGui::Button(ICON_COPY_CAMERA)) { - std::string camera_string = platform::CameraToString(data(), &camera_); - platform::MaybeSaveToClipboard(camera_string); - } - ImGui::SetItemTooltip("%s", "Copy Camera"); - - ImGui::SameLine(); - ImGui::Text("%s", " |"); // Label selection. - ImGui::SameLine(); - ImGui::Text("%s", ICON_LABEL); - ImGui::SetItemTooltip("%s", "Label"); - ImGui::SameLine(); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); - if (ImGui::BeginCombo("##Label", kLabelNames[vis_options_.label], - combo_flags)) { + std::string label_preview = + std::string(ICON_LABEL) + " " + kLabelNames[vis_options_.label]; + if (ImGui::BeginCombo("##Label", label_preview.c_str(), combo_flags)) { for (int n = 0; n < IM_ARRAYSIZE(kLabelNames); n++) { if (ImGui::Selectable(kLabelNames[n], (vis_options_.label == n))) { vis_options_.label = n; @@ -1547,18 +1472,12 @@ void App::ToolBarGui() { } ImGui::SetItemTooltip("%s", "Label"); - ImGui::SameLine(); - ImGui::Text("%s", " |"); - // Frame selection. - ImGui::SameLine(); - ImGui::Text("%s", ICON_FRAME); - ImGui::SetItemTooltip("%s", "Frame"); - ImGui::SameLine(); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); - if (ImGui::BeginCombo("##Frame", kFrameNames[vis_options_.frame], - combo_flags)) { + std::string frame_preview = + std::string(ICON_FRAME) + " " + kFrameNames[vis_options_.frame]; + if (ImGui::BeginCombo("##Frame", frame_preview.c_str(), combo_flags)) { for (int n = 0; n < IM_ARRAYSIZE(kFrameNames); n++) { if (ImGui::Selectable(kFrameNames[n], (vis_options_.frame == n))) { vis_options_.frame = n; @@ -1568,31 +1487,31 @@ void App::ToolBarGui() { } ImGui::SetItemTooltip("%s", "Frame"); + // Theme selection. ImGui::SameLine(); - ImGui::Text("%s", " |"); - - // Style selection. - ImGui::SameLine(); - switch (ui_.theme) { - case platform::GuiTheme::kLight: - if (ImGui::Button(ICON_LIGHTMODE)) { - SetupTheme(platform::GuiTheme::kDark); + const char* theme_icons[] = {ICON_LIGHTMODE, ICON_DARKMODE, + ICON_CLASSICMODE}; + const char* theme_tooltips[] = {"Light Mode", "Dark Mode", "Classic Mode"}; + const platform::GuiTheme theme_values[] = { + platform::GuiTheme::kLight, + platform::GuiTheme::kDark, + platform::GuiTheme::kClassic, + }; + int theme_idx = static_cast(ui_.theme); + ImGui::SetNextItemWidth(ImGui::CalcTextSize(theme_icons[0]).x + + ImGui::GetStyle().FramePadding.x * 2); + if (ImGui::BeginCombo("##Theme", theme_icons[theme_idx], combo_flags)) { + for (int n = 0; n < IM_ARRAYSIZE(theme_icons); n++) { + if (ImGui::Selectable(theme_icons[n], (theme_idx == n))) { + SetupTheme(theme_values[n]); } - ImGui::SetItemTooltip("%s", "Switch to Dark Mode"); - break; - case platform::GuiTheme::kDark: - if (ImGui::Button(ICON_DARKMODE)) { - SetupTheme(platform::GuiTheme::kClassic); + if (ImGui::IsItemHovered()) { + ImGui::SetTooltip("%s", theme_tooltips[n]); } - ImGui::SetItemTooltip("%s", "Switch to Classic Mode"); - break; - case platform::GuiTheme::kClassic: - if (ImGui::Button(ICON_CLASSICMODE)) { - SetupTheme(platform::GuiTheme::kLight); - } - ImGui::SetItemTooltip("%s", "Switch to Light Mode"); - break; + } + ImGui::EndCombo(); } + ImGui::SetItemTooltip("%s", "Theme"); ImGui::EndTable(); } @@ -1613,17 +1532,7 @@ void App::StatusBarGui() { } else if (step_control_.GetPauseState() == PauseState::kNormalPaused) { ImGui::Text("Paused"); } else { - const float desired_realtime = step_control_.GetSpeed(); - const float measured_realtime = step_control_.GetSpeedMeasured(); - const float realtime_offset = - mju_abs(measured_realtime - desired_realtime); - const bool misaligned = realtime_offset > 0.1 * desired_realtime; - if (misaligned) { - ImGui::Text("Running: %g%% (%-4.1f%%)", desired_realtime, - measured_realtime); - } else { - ImGui::Text("Running: %g%%", desired_realtime); - } + ImGui::Text("Running"); } if (!step_error_.empty()) { From 8f56d5eefa00e5e41b2fa72d01bcf0e7ff3530d2 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Thu, 19 Feb 2026 19:15:44 -0800 Subject: [PATCH 12/48] Minor style fixes PiperOrigin-RevId: 872669434 Change-Id: I794173b9fad3a286cbf57d246b2f8672412a0e20 --- simulate/simulate.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/simulate/simulate.cc b/simulate/simulate.cc index be7318c5..022b9842 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -1584,6 +1584,7 @@ void UiEvent(mjuiState* state) { // rendering section else if (it && it->sectionid==SECT_RENDERING) { + // only update the camera when the camera itself changed if (it->pdata == &sim->camera) { if (sim->camera==0) { @@ -1603,6 +1604,7 @@ void UiEvent(mjuiState* state) { sim->cam.fixedcamid = sim->camera - 2; } } + // copy camera spec to clipboard (as MJCF element) if (it->itemid == 3) { CopyCamera(sim); From 82694b520b03e759dbe6885546f7341d0d370c95 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 01:27:20 -0800 Subject: [PATCH 13/48] Re-apply changes that were accidentally deleted. PiperOrigin-RevId: 872782055 Change-Id: If2f403835cd2d90c88abe9ed82380dfca5b3d9c8 --- src/experimental/studio/app.cc | 117 +++++++++++++++++++++++---------- 1 file changed, 82 insertions(+), 35 deletions(-) diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index adef9a42..30a2166e 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -84,6 +85,7 @@ static constexpr const char* ICON_PREV_FRAME = platform::ICON_FA_CARET_LEFT; static constexpr const char* ICON_NEXT_FRAME = platform::ICON_FA_CARET_RIGHT; static constexpr const char* ICON_CURR_FRAME = platform::ICON_FA_FAST_FORWARD; static constexpr const char* ICON_SPEED = platform::ICON_FA_TACHOMETER; +static constexpr const char* ICON_DELETE = platform::ICON_FA_TRASH_CAN; // UI labels for mjtLabel. static constexpr const char* kLabelNames[] = { @@ -105,7 +107,8 @@ static constexpr std::array kPercentRealTime = { }; // clang-format on -App::App(Config config) : ini_path_(std::move(config.ini_path)) { +App::App(Config config) + : rng_(std::random_device()()), ini_path_(std::move(config.ini_path)) { platform::Window::Config window_config; window_config.renderer_backend = platform::Renderer::GetBackend(); window_config.offscreen_mode = config.offscreen_mode; @@ -132,6 +135,13 @@ void App::ClearModel() { step_error_ = ""; } +void App::Recompile() { + mj_recompile(model_holder_->spec(), model_holder_->vfs(), + model_holder_->model(), model_holder_->data()); + const int state_size = mj_stateSize(model(), mjSTATE_INTEGRATION); + history_.Init(state_size); +} + void App::RequestModelLoad(std::string model_file) { pending_load_ = std::move(model_file); } @@ -139,6 +149,7 @@ void App::RequestModelLoad(std::string model_file) { void App::RequestModelReload() { if (model_kind_ == kModelFromFile) { pending_load_ = model_path_; + preserve_camera_on_load_ = true; } } @@ -189,6 +200,16 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { const int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); history_.Init(state_size); + if (!preserve_camera_on_load_) { + const int model_cam = model->vis.global.cameraid; + if (model_cam >= 0 && model_cam < model->ncam) { + ui_.camera_idx = platform::SetCamera(model, &camera_, model_cam); + } else { + mjv_defaultFreeCamera(model, &camera_); + } + } + preserve_camera_on_load_ = false; + // Initialize the speed based on the model's default real-time setting. float min_error = FLT_MAX; const float desired = mju_log(100 * model->vis.global.realtime); @@ -200,6 +221,10 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { SetSpeedIndex(i); } } +if (spec_op_) { + spec_op_(); + spec_op_ = nullptr; + } platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { if (plugin->post_model_loaded) { @@ -382,6 +407,19 @@ void App::ProcessPendingLoads() { }); } +void App::SpecDeleteSelectedElement() { + spec_op_ = [this]() { + mjs_delete(spec(), tmp_.element); + if (tmp_.element->elemtype == mjOBJ_BODY && + perturb_.select == tmp_.element_id) { + mjv_defaultPerturb(&perturb_); + } + tmp_.element = nullptr; + tmp_.element_id = -1; + Recompile(); + }; +} + void App::HandleWindowEvents() { const std::string drop_file = window_->GetDropFile(); if (!drop_file.empty()) { @@ -474,15 +512,28 @@ void App::HandleMouseEvents() { perturb_.flexselect = picked.flex; perturb_.skinselect = picked.skin; + // Select the corresponding element in the spec. + tmp_.element = nullptr; + tmp_.element_id = -1; + if (has_spec()) { + mjsElement* element = mjs_firstElement(spec(), mjOBJ_BODY); + while (element) { + if (mjs_getId(element) == picked.body) { + tmp_.element = element; + tmp_.element_id = picked.body; + break; + } + element = mjs_nextElement(spec(), element); + } + } + // Compute the local position of the selected object in the world. mjtNum tmp[3]; mju_sub3(tmp, picked.point, data()->xpos + 3 * picked.body); mju_mulMatTVec(perturb_.localpos, data()->xmat + 9 * picked.body, tmp, 3, 3); } else { - perturb_.select = 0; - perturb_.flexselect = -1; - perturb_.skinselect = -1; + mjv_defaultPerturb(&perturb_); } } @@ -571,6 +622,8 @@ void App::HandleKeyboardEvents() { } } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) { ResetPhysics(); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Delete)) { + SpecDeleteSelectedElement(); } else if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) { SelectParentPerturb(model(), perturb_); } else if (ImGui_IsChordJustPressed(ImGuiKey_F1)) { @@ -1088,7 +1141,8 @@ void App::SpecExplorerGui() { label = "(" + prefix + " " + std::to_string(id) + ")"; } - if (ImGui::Selectable(label.c_str(), false)) { + const bool selected = (tmp_.element == element); + if (ImGui::Selectable(label.c_str(), selected)) { tmp_.element = element; tmp_.element_id = id; } @@ -1098,30 +1152,7 @@ void App::SpecExplorerGui() { }; if (ImGui::TreeNodeEx("Bodies", flags)) { - // We don't use `display_group` here because we do additional selection - // logic tied to the `perturb_` field. - mjsElement* element = mjs_firstElement(spec(), mjOBJ_BODY); - while (element) { - const int id = mjs_getId(element); - - const mjString* name = mjs_getName(element); - std::string label = *name; - if (label.empty()) { - label = "(Body " + std::to_string(id) + ")"; - } - - if (ImGui::Selectable(label.c_str(), (id == perturb_.select), - ImGuiSelectableFlags_AllowDoubleClick)) { - tmp_.element = element; - tmp_.element_id = id; - } - if (ImGui::IsItemHovered() && - ImGui::IsMouseDoubleClicked(ImGuiMouseButton_Left)) { - perturb_.select = id; - } - - element = mjs_nextElement(spec(), element); - } + display_group(mjOBJ_BODY, "Body"); ImGui::TreePop(); } @@ -1134,6 +1165,13 @@ void App::SpecExplorerGui() { display_group(mjOBJ_SITE, "Site"); ImGui::TreePop(); } + + // If we selected a body, then select the same body for the perturb object. + if (tmp_.element && tmp_.element->elemtype == mjOBJ_BODY && + perturb_.select != tmp_.element_id) { + mjv_defaultPerturb(&perturb_); + perturb_.select = tmp_.element_id; + } } void App::PropertiesGui() { @@ -1142,22 +1180,31 @@ void App::PropertiesGui() { return; } + if (ImGui::BeginTable("##PropertiesHeader", 2)) { + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 20); + ImGui::TableNextColumn(); + ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); + ImGui::TableNextColumn(); + if (tmp_.element->elemtype == mjOBJ_BODY) { + if (ImGui::SmallButton(ICON_DELETE)) { + SpecDeleteSelectedElement(); + } + } + ImGui::EndTable(); + } + ImGui::Separator(); + switch (tmp_.element->elemtype) { case mjOBJ_BODY: - ImGui::Text("Body"); - ImGui::Separator(); platform::BodyPropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; case mjOBJ_JOINT: - ImGui::Text("Joint"); - ImGui::Separator(); platform::JointPropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; case mjOBJ_SITE: - ImGui::Text("Site"); - ImGui::Separator(); platform::SitePropertiesGui(model(), data(), tmp_.element, tmp_.element_id); break; From f0bf4032544e82a6203f0791247bc7c1c94e21aa Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 02:08:16 -0800 Subject: [PATCH 14/48] Add a delete button to the highlighted element. PiperOrigin-RevId: 872795277 Change-Id: I63afa9c87b425ebed5d20126e9909ad95532d559 --- src/experimental/studio/app.cc | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 30a2166e..2454821f 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -1130,7 +1131,8 @@ void App::SpecExplorerGui() { const ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; - auto display_group = [this](mjtObj type, const std::string& prefix) { + auto display_group = [this](mjtObj type, const std::string& prefix, + std::function delete_callback = {}) { mjsElement* element = mjs_firstElement(spec(), type); while (element) { const int id = mjs_getId(element); @@ -1142,17 +1144,28 @@ void App::SpecExplorerGui() { } const bool selected = (tmp_.element == element); - if (ImGui::Selectable(label.c_str(), selected)) { + if (ImGui::Selectable(label.c_str(), selected, + ImGuiSelectableFlags_AllowOverlap)) { tmp_.element = element; tmp_.element_id = id; } + if (selected && delete_callback) { + // Right-align the delete button. + const float button_width = ImGui::CalcTextSize(ICON_DELETE).x + + ImGui::GetStyle().FramePadding.x * 2.0f; + ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); + if (ImGui::SmallButton(ICON_DELETE)) { + delete_callback(); + } + } + element = mjs_nextElement(spec(), element); } }; if (ImGui::TreeNodeEx("Bodies", flags)) { - display_group(mjOBJ_BODY, "Body"); + display_group(mjOBJ_BODY, "Body", [this] { SpecDeleteSelectedElement(); }); ImGui::TreePop(); } From 28ad603f6be70061cbfc295e2f334ff91b5564f7 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 02:42:13 -0800 Subject: [PATCH 15/48] Use templates and macros to reduce boilerplate. PiperOrigin-RevId: 872805912 Change-Id: I6fc1abe283bc809a95d660abc072906b0810aaeb --- src/experimental/platform/plugin.cc | 107 ++++++++++---------------- src/experimental/platform/plugin.h | 31 +++----- src/experimental/platform/renderer.cc | 2 +- src/experimental/studio/app.cc | 12 +-- 4 files changed, 59 insertions(+), 93 deletions(-) diff --git a/src/experimental/platform/plugin.cc b/src/experimental/platform/plugin.cc index 13c262c1..fb790f63 100644 --- a/src/experimental/platform/plugin.cc +++ b/src/experimental/platform/plugin.cc @@ -12,89 +12,62 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "experimental/platform/plugin.h" + #include #include #include -#include "experimental/platform/plugin.h" #include "engine/engine_global_table.h" +using GuiPlugin = mujoco::platform::GuiPlugin; +using ModelPlugin = mujoco::platform::ModelPlugin; + namespace mujoco::platform { -void RegisterGuiPlugin(const GuiPlugin* plugin) { - if (plugin->name == nullptr || plugin->name[0] == '\0') { +template +void RegisterPlugin(T plugin) { + if (plugin.name == nullptr || plugin.name[0] == '\0') { mju_error("Plugin name must not be empty or null."); } - GlobalTable::GetSingleton().AppendIfUnique(*plugin); + GlobalTable::GetSingleton().AppendIfUnique(plugin); } -void ForEachGuiPlugin(const std::function& fn) { - auto& table = GlobalTable::GetSingleton(); +template +void ForEachPlugin(const std::function& fn) { + auto& table = mujoco::GlobalTable::GetSingleton(); for (int i = 0; i < table.count(); ++i) { - const GuiPlugin* plugin = table.GetAtSlot(i); - fn(const_cast(plugin)); - } -} - -void RegisterModelPlugin(const ModelPlugin* plugin) { - if (plugin->name == nullptr || plugin->name[0] == '\0') { - mju_error("Plugin name must not be empty or null."); - } - GlobalTable::GetSingleton().AppendIfUnique(*plugin); -} - -void ForEachModelPlugin(const std::function& fn) { - auto& table = GlobalTable::GetSingleton(); - for (int i = 0; i < table.count(); ++i) { - const ModelPlugin* plugin = table.GetAtSlot(i); - fn(const_cast(plugin)); + const T* plugin = table.GetAtSlot(i); + fn(const_cast(plugin)); } } } // namespace mujoco::platform -using mujoco::GlobalTable; -using GuiPlugin = mujoco::platform::GuiPlugin; -using ModelPlugin = mujoco::platform::ModelPlugin; +#define MUJOCO_SPECIALIZE_PLUGIN(PLUGIN, NAME) \ + template <> \ + const char* mujoco::GlobalTable::HumanReadableTypeName() { \ + return NAME; \ + } \ + template <> \ + std::string_view mujoco::GlobalTable::ObjectKey(const PLUGIN& p) { \ + return std::string_view(p.name); \ + } \ + template <> \ + bool mujoco::GlobalTable::ObjectEqual(const PLUGIN& p1, \ + const PLUGIN& p2) { \ + return CaseInsensitiveEqual(p1.name, p2.name); \ + } \ + template <> \ + bool mujoco::GlobalTable::CopyObject(PLUGIN& dst, const PLUGIN& src, \ + ErrorMessage& err) { \ + dst = src; \ + return true; \ + } \ + namespace mujoco::platform { \ + template void RegisterPlugin(PLUGIN plugin); \ + template void ForEachPlugin(const std::function& fn); \ + } -template <> -const char* GlobalTable::HumanReadableTypeName() { - return "gui plugin"; -} - -template <> -std::string_view GlobalTable::ObjectKey(const GuiPlugin& plugin) { - return std::string_view(plugin.name); -} - -template <> -bool GlobalTable::ObjectEqual(const GuiPlugin& p1, const GuiPlugin& p2) { - return CaseInsensitiveEqual(p1.name, p2.name); -} - -template <> -bool GlobalTable::CopyObject(GuiPlugin& dst, const GuiPlugin& src, ErrorMessage& err) { - dst = src; - return true; -} - -template <> -const char* GlobalTable::HumanReadableTypeName() { - return "model plugin"; -} - -template <> -std::string_view GlobalTable::ObjectKey(const ModelPlugin& plugin) { - return std::string_view(plugin.name); -} - -template <> -bool GlobalTable::ObjectEqual(const ModelPlugin& p1, const ModelPlugin& p2) { - return CaseInsensitiveEqual(p1.name, p2.name); -} - -template <> -bool GlobalTable::CopyObject(ModelPlugin& dst, const ModelPlugin& src, ErrorMessage& err) { - dst = src; - return true; -} +MUJOCO_SPECIALIZE_PLUGIN(GuiPlugin, "gui plugin"); +MUJOCO_SPECIALIZE_PLUGIN(ModelPlugin, "model plugin"); diff --git a/src/experimental/platform/plugin.h b/src/experimental/platform/plugin.h index 824b9793..f10fd641 100644 --- a/src/experimental/platform/plugin.h +++ b/src/experimental/platform/plugin.h @@ -20,10 +20,16 @@ namespace mujoco::platform { -// Important: Do not inherit from these plugin structs. They are copied by value -// and therefore any derived classes will be sliced. We assume plugins are -// effectively globals and so any pointers will be valid for the lifetime -// of the process. +// Registers plugins with the global registry. The plugins must have a +// case-insensitive unique name for the plugin type. Note that plugins are +// copied by value, so do not use inheritance. +template +void RegisterPlugin(T plugin); + +// Executes the given function for each registered plugin of type T. +template +void ForEachPlugin(const std::function& fn); + // Plugin for processing custom UI windows. The plugin will be listed in the // "Plugins" main menu and, when selected, an ImGui window will be opened with @@ -47,8 +53,9 @@ struct GuiPlugin final { void* data = nullptr; }; +// Plugin for loading and updating models. struct ModelPlugin final { - using GetModelToLoadFn = const char* (*)(ModelPlugin * self, int* size, + using GetModelToLoadFn = const char* (*)(ModelPlugin* self, int* size, char* content_type, int content_type_size, char* model_name, @@ -75,20 +82,6 @@ struct ModelPlugin final { void* data = nullptr; }; -// Registers a plugin with a global registry. The plugin must have a -// case-insensitive unique name. -void RegisterGuiPlugin(const GuiPlugin* plugin); - -// Executes the given function for each registered plugin. -void ForEachGuiPlugin(const std::function& fn); - -// Registers a plugin with a global registry. The plugin must have a -// case-insensitive unique name. -void RegisterModelPlugin(const ModelPlugin* plugin); - -// Executes the given function for each registered plugin. -void ForEachModelPlugin(const std::function& fn); - } // namespace mujoco::platform #endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_ diff --git a/src/experimental/platform/renderer.cc b/src/experimental/platform/renderer.cc index b1b39e52..6fe16730 100644 --- a/src/experimental/platform/renderer.cc +++ b/src/experimental/platform/renderer.cc @@ -219,6 +219,6 @@ mjPLUGIN_LIB_INIT { plugin.update = [](mujoco::platform::GuiPlugin* self) { mjr_updateGui(nullptr); }; - mujoco::platform::RegisterGuiPlugin(&plugin); + mujoco::platform::RegisterPlugin(plugin); } #endif diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 2454821f..6da86d2f 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -227,7 +227,7 @@ if (spec_op_) { spec_op_ = nullptr; } - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->post_model_loaded) { plugin->post_model_loaded(plugin, model_path_.c_str()); } @@ -271,7 +271,7 @@ void App::UpdatePhysics() { } bool stepped = false; - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->do_update) { if (plugin->do_update(plugin, model(), data())) { stepped = true; @@ -392,7 +392,7 @@ void App::ProcessPendingLoads() { } // Check plugins to see if we need to load a new model. - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->get_model_to_load) { char model_name[1000] = ""; char content_type[1000] = ""; @@ -781,7 +781,7 @@ void App::LoadSettings() { platform::KeyValues plugin_names = platform::ReadIniSection(settings, "[Studio][Plugins]"); - platform::ForEachGuiPlugin([&](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { auto it = plugin_names.find(plugin->name); if (it != plugin_names.end()) { plugin->active = std::stoi(it->second) != 0; @@ -797,7 +797,7 @@ void App::SaveSettings() { platform::AppendIniSection(settings, "[Studio][UX]", ui_.ToDict()); platform::KeyValues plugin_names; - platform::ForEachGuiPlugin([&](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { plugin_names[plugin->name] = std::to_string((int)plugin->active); }); platform::AppendIniSection(settings, "[Studio][Plugins]", plugin_names); @@ -977,7 +977,7 @@ void App::BuildGui() { ImGui::End(); } - platform::ForEachGuiPlugin([](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([](auto* plugin) { if (!plugin->update) { return; } From 0083cf7fa81df8d45d8051560900094f1d6f9133 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 03:18:55 -0800 Subject: [PATCH 16/48] Refactor the projectile launcher into a plugin. Add a simple GUI for it that allows users to enable/disable the key binding as well as control some parameters. PiperOrigin-RevId: 872816645 Change-Id: I8abb2803e17b4f6195d447220c436d4701b24db8 --- src/experimental/platform/CMakeLists.txt | 1 + .../platform/object_launcher_plugin.cc | 214 ++++++++++++++++++ src/experimental/platform/plugin.cc | 4 + src/experimental/platform/plugin.h | 39 ++++ src/experimental/studio/app.cc | 126 ++++++----- src/experimental/studio/app.h | 3 - 6 files changed, 329 insertions(+), 58 deletions(-) create mode 100644 src/experimental/platform/object_launcher_plugin.cc diff --git a/src/experimental/platform/CMakeLists.txt b/src/experimental/platform/CMakeLists.txt index 6429db47..2cdc5f4e 100644 --- a/src/experimental/platform/CMakeLists.txt +++ b/src/experimental/platform/CMakeLists.txt @@ -50,6 +50,7 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} interaction.h model_holder.cc model_holder.h + object_launcher_plugin.cc picture_gui.h picture_gui.cc plugin.cc diff --git a/src/experimental/platform/object_launcher_plugin.cc b/src/experimental/platform/object_launcher_plugin.cc new file mode 100644 index 00000000..0ba23894 --- /dev/null +++ b/src/experimental/platform/object_launcher_plugin.cc @@ -0,0 +1,214 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include +#include +#include "experimental/platform/imgui_widgets.h" +#include "experimental/platform/plugin.h" + +namespace mujoco::studio { + +class ObjectLauncher { + public: + ObjectLauncher() : rng_(std::random_device{}()) {} + + void UpdateGui() { + using platform::ImGui_Input; + + ImGui::Checkbox("Enable Key Binding (Ctrl+Shift+Enter)", &enabled_); + ImGui_Input("Size", &size_, {0.01f, 1.0f, 0.01, 0.1}); + ImGui_Input("Speed", &speed_, {0.01f, 100.0f, 0.1, 1.0}); + ImGui_Input("Mass", &mass_, {0.01f, 100.0f, 0.01, 0.1}); + ImGui_Input("Life", &lifetime_, {0.0f, 60.0f, 0.1, 1.0}); + + int shape = type_ == mjGEOM_BOX ? 0 : 1; + const char* names[] = {"Box", "Sphere"}; + ImGui::Combo("Shape", &shape, names, 2); + type_ = shape == 0 ? mjGEOM_BOX : mjGEOM_SPHERE; + + if (ImGui::Button("Launch", ImVec2(-1.0f, 0.0f))) { + active_ = true; + } + + if (ImGui::Button("Clear")) { + for (auto& object : objects_) { + object.expiration = -1; + } + } + } + + void HandleKeyboardEvent() { if (enabled_) active_ = true; } + + bool UpdateSpecPreCompile(mjSpec* spec, const mjModel* model, + const mjData* data, const mjvCamera* camera) { + // Remove expired objects. + auto it = std::remove_if( + objects_.begin(), objects_.end(), [&](const ObjectInfo& o) { + const bool expired = + o.body_id >= 0 && o.expiration != 0 && o.expiration < data->time; + if (expired) { + mjsBody* body = mjs_findBody(spec, o.name.c_str()); + if (body) { + mjs_delete(spec, body->element); + } + } + return expired; + }); + if (it != objects_.end()) { + objects_.erase(it, objects_.end()); + return true; + } + + if (!active_) return false; + active_ = false; + + mjsBody* world = mjs_findBody(spec, "world"); + if (!world) return false; + mjsBody* body = mjs_addBody(world, nullptr); + if (!body) return false; + mjsJoint* joint = mjs_addJoint(body, nullptr); + if (!joint) return false; + mjsGeom* geom = mjs_addGeom(body, nullptr); + if (!geom) return false; + + ObjectInfo& object = objects_.emplace_back(); + object.name = "projectile" + std::to_string(counter_++);; + object.expiration = data->time + lifetime_; + + mjtNum pos[3]; + mjtNum dir[3]; + mjtNum up[3]; + mjv_cameraFrame(pos, dir, up, nullptr, data, camera); + mjs_setName(body->element, object.name.c_str()); + + joint->type = mjJNT_FREE; + body->mass = mass_; + geom->type = type_; + geom->size[0] = size_; + geom->size[1] = size_; + geom->size[2] = size_; + // Slightly in front of the camera. + body->pos[0] = pos[0] + (dir[0] * 0.1); + body->pos[1] = pos[1] + (dir[1] * 0.1); + body->pos[2] = pos[2] + (dir[2] * 0.1); + // Randomize the color. + geom->rgba[0] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[1] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[2] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[3] = 1.0; + // Launch it slightly upwards to get a nice arc. + launch_vel_[0] = (dir[0] * speed_) + up[0]; + launch_vel_[1] = (dir[1] * speed_) + up[1]; + launch_vel_[2] = (dir[2] * speed_) + up[2]; + return true; + } + + void UpdateSpecPostCompile(const mjSpec* spec, const mjModel* model, + mjData* data) { + if (objects_.empty()) { + return; + } + + ObjectInfo& object = objects_.back(); + if (object.launched) { + return; + } + object.launched = true; + + const int body_id = mj_name2id(model, mjOBJ_BODY, object.name.c_str()); + if (body_id < 0) { + return; + } + int joint_id = model->body_jntadr[body_id]; + if (joint_id < 0 || model->jnt_type[joint_id] != mjJNT_FREE) { + return; + } + int qvel_addr = model->jnt_dofadr[joint_id]; + if (qvel_addr < 0) { + return; + } + object.body_id = body_id; + data->qvel[qvel_addr + 0] = launch_vel_[0]; + data->qvel[qvel_addr + 1] = launch_vel_[1]; + data->qvel[qvel_addr + 2] = launch_vel_[2]; + } + + private: + struct ObjectInfo { + std::string name; + int body_id = -1; + mjtNum expiration = 0; + bool launched = false; + }; + + std::mt19937 rng_; + int counter_ = 0; + bool enabled_ = false; + bool active_ = false; + mjtNum size_ = 0.13365; + mjtNum speed_ = 10.0; + mjtNum mass_ = 10.0; + mjtNum lifetime_ = 5.0; + mjtGeom type_ = mjGEOM_BOX; + mjtNum launch_vel_[3] = {0, 0, 0}; + std::vector objects_; +}; + +} // namespace mujoco::studio + +mjPLUGIN_LIB_INIT { + using mujoco::studio::ObjectLauncher; + + static ObjectLauncher plugin; + + mujoco::platform::GuiPlugin gui; + gui.data = &plugin; + gui.name = "ObjectLauncher"; + gui.update = [](mujoco::platform::GuiPlugin* self) { + auto* plugin = static_cast(self->data); + plugin->UpdateGui(); + }; + mujoco::platform::RegisterPlugin(gui); + + mujoco::platform::KeyHandlerPlugin key_handler; + key_handler.data = &plugin; + key_handler.name = "ObjectLauncher"; + key_handler.key_chord = ImGuiKey_Enter | ImGuiMod_Ctrl | ImGuiMod_Shift; + key_handler.on_key_pressed = [](mujoco::platform::KeyHandlerPlugin* self) { + auto* plugin = static_cast(self->data); + plugin->HandleKeyboardEvent(); + }; + mujoco::platform::RegisterPlugin(key_handler); + + mujoco::platform::SpecEditorPlugin spec_editor; + spec_editor.data = &plugin; + spec_editor.name = "ObjectLauncher"; + spec_editor.pre_compile = [](mujoco::platform::SpecEditorPlugin* self, + mjSpec* spec, const mjModel* model, + const mjData* data, const mjvCamera* camera) { + auto* plugin = static_cast(self->data); + return plugin->UpdateSpecPreCompile(spec, model, data, camera); + }; + spec_editor.post_compile = [](mujoco::platform::SpecEditorPlugin* self, + const mjSpec* spec, const mjModel* model, + mjData* data) { + auto* plugin = static_cast(self->data); + return plugin->UpdateSpecPostCompile(spec, model, data); + }; + mujoco::platform::RegisterPlugin(spec_editor); +} diff --git a/src/experimental/platform/plugin.cc b/src/experimental/platform/plugin.cc index fb790f63..3ddb6e1b 100644 --- a/src/experimental/platform/plugin.cc +++ b/src/experimental/platform/plugin.cc @@ -22,6 +22,8 @@ using GuiPlugin = mujoco::platform::GuiPlugin; using ModelPlugin = mujoco::platform::ModelPlugin; +using KeyHandlerPlugin = mujoco::platform::KeyHandlerPlugin; +using SpecEditorPlugin = mujoco::platform::SpecEditorPlugin; namespace mujoco::platform { @@ -71,3 +73,5 @@ void ForEachPlugin(const std::function& fn) { MUJOCO_SPECIALIZE_PLUGIN(GuiPlugin, "gui plugin"); MUJOCO_SPECIALIZE_PLUGIN(ModelPlugin, "model plugin"); +MUJOCO_SPECIALIZE_PLUGIN(KeyHandlerPlugin, "key handler plugin"); +MUJOCO_SPECIALIZE_PLUGIN(SpecEditorPlugin, "spec editor plugin"); diff --git a/src/experimental/platform/plugin.h b/src/experimental/platform/plugin.h index f10fd641..5e47fae9 100644 --- a/src/experimental/platform/plugin.h +++ b/src/experimental/platform/plugin.h @@ -82,6 +82,45 @@ struct ModelPlugin final { void* data = nullptr; }; +// Plugin for handling custom keyboard events. +struct KeyHandlerPlugin final { + using OnKeyPressedFn = void (*)(KeyHandlerPlugin* self); + + // The name of the plugin; must be unique. + const char* name = ""; + + // The ImGui key codes for the key combination that triggers the plugin. + int key_chord = 0; + + // The function to be called when the above key combination is pressed. + OnKeyPressedFn on_key_pressed = nullptr; + + // Optional data pointer. + void* data = nullptr; +}; + +// Plugin for editing the mjSpec. +struct SpecEditorPlugin final { + using PreCompileFn = bool (*)(SpecEditorPlugin* self, mjSpec* spec, + const mjModel* model, const mjData* data, + const mjvCamera* camera); + using PostCompileFn = void (*)(SpecEditorPlugin* self, const mjSpec* spec, + const mjModel* model, mjData* data); + + // The name of the plugin; must be unique. + const char* name = ""; + + // Callback that edits the spec. If it returns true, then the spec will be + // recompiled and `post_compile` will be called with the result. + PreCompileFn pre_compile = nullptr; + + // Callback that is called after the spec has been recompiled. + PostCompileFn post_compile = nullptr; + + // Optional data pointer. + void* data = nullptr; +}; + } // namespace mujoco::platform #endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_ diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 6da86d2f..352e8c8a 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -25,7 +25,6 @@ #include #include #include -#include #include #include #include @@ -108,8 +107,7 @@ static constexpr std::array kPercentRealTime = { }; // clang-format on -App::App(Config config) - : rng_(std::random_device()()), ini_path_(std::move(config.ini_path)) { +App::App(Config config) : ini_path_(std::move(config.ini_path)) { platform::Window::Config window_config; window_config.renderer_backend = platform::Renderer::GetBackend(); window_config.offscreen_mode = config.offscreen_mode; @@ -222,10 +220,6 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { SetSpeedIndex(i); } } -if (spec_op_) { - spec_op_(); - spec_op_ = nullptr; - } platform::ForEachPlugin([&](auto* plugin) { if (plugin->post_model_loaded) { @@ -391,6 +385,23 @@ void App::ProcessPendingLoads() { } } + if (spec_op_) { + spec_op_(); + spec_op_ = nullptr; + } + + // Allow plugins to edit the spec as well. + platform::ForEachPlugin([&](auto* plugin) { + if (plugin->pre_compile) { + if (plugin->pre_compile(plugin, spec(), model(), data(), &camera_)) { + Recompile(); + if (plugin->post_compile) { + plugin->post_compile(plugin, spec(), model(), data()); + } + }; + } + }); + // Check plugins to see if we need to load a new model. platform::ForEachPlugin([&](auto* plugin) { if (plugin->get_model_to_load) { @@ -715,60 +726,65 @@ void App::HandleKeyboardEvents() { ToggleFlag(vis_options_.geomgroup[4]); } else if (ImGui_IsChordJustPressed(ImGuiKey_5)) { ToggleFlag(vis_options_.geomgroup[5]); - } else if (has_model()) { - if (ImGui_IsChordJustPressed(ImGuiKey_Escape)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, platform::kTumbleCameraIdx); - } else if (ImGui_IsChordJustPressed(ImGuiKey_LeftBracket)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, ui_.camera_idx - 1); - } else if (ImGui_IsChordJustPressed(ImGuiKey_RightBracket)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, ui_.camera_idx + 1); + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_Escape)) { + ui_.camera_idx = + platform::SetCamera(model(), &camera_, platform::kTumbleCameraIdx); + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_LeftBracket)) { + ui_.camera_idx = platform::SetCamera(model(), &camera_, ui_.camera_idx - 1); + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_RightBracket)) { + ui_.camera_idx = platform::SetCamera(model(), &camera_, ui_.camera_idx + 1); + // WASD camera controls for free camera. + } else if (is_freecam_wasd && + (ImGui::IsKeyDown(ImGuiKey_W) || ImGui::IsKeyDown(ImGuiKey_S) || + ImGui::IsKeyDown(ImGuiKey_A) || ImGui::IsKeyDown(ImGuiKey_D) || + ImGui::IsKeyDown(ImGuiKey_Q) || ImGui::IsKeyDown(ImGuiKey_E))) { + bool moved = false; + + // Move (dolly) forward/backward using W and S keys. + if (ImGui::IsKeyDown(ImGuiKey_W)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, tmp_.cam_speed); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_S)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, -tmp_.cam_speed); + moved = true; } - // WASD camera controls for free camera. - if (is_freecam_wasd) { - bool moved = false; + // Strafe (truck) left/right using A and D keys. + if (ImGui::IsKeyDown(ImGuiKey_A)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, -tmp_.cam_speed, 0); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_D)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, tmp_.cam_speed, 0); + moved = true; + } - // Move (dolly) forward/backward using W and S keys. - if (ImGui::IsKeyDown(ImGuiKey_W)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, tmp_.cam_speed); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_S)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, -tmp_.cam_speed); - moved = true; + // Move (pedestal) up/down using Q and E keys. + if (ImGui::IsKeyDown(ImGuiKey_Q)) { + MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, tmp_.cam_speed); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_E)) { + MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, -tmp_.cam_speed); + moved = true; + } + + if (moved) { + tmp_.cam_speed += 0.001f; + + const float max_speed = ImGui::GetIO().KeyShift ? 0.1 : 0.01f; + if (tmp_.cam_speed > max_speed) { + tmp_.cam_speed = max_speed; } - - // Strafe (truck) left/right using A and D keys. - if (ImGui::IsKeyDown(ImGuiKey_A)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, -tmp_.cam_speed, 0); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_D)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, tmp_.cam_speed, 0); - moved = true; - } - - // Move (pedestal) up/down using Q and E keys. - if (ImGui::IsKeyDown(ImGuiKey_Q)) { - MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, tmp_.cam_speed); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_E)) { - MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, -tmp_.cam_speed); - moved = true; - } - - if (moved) { - tmp_.cam_speed += 0.001f; - - const float max_speed = ImGui::GetIO().KeyShift ? 0.1 : 0.01f; - if (tmp_.cam_speed > max_speed) { - tmp_.cam_speed = max_speed; + } else { + tmp_.cam_speed = 0.001f; + } + } else { + platform::ForEachPlugin([&](auto* plugin) { + if (plugin->key_chord && plugin->on_key_pressed) { + if (ImGui_IsChordJustPressed(plugin->key_chord)) { + plugin->on_key_pressed(plugin); } - } else { - tmp_.cam_speed = 0.001f; } - } + }); } } diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index 170aadd6..d2cce3a5 100644 --- a/src/experimental/studio/app.h +++ b/src/experimental/studio/app.h @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include @@ -224,8 +223,6 @@ class App { bool has_model() const { return model_holder_ && model_holder_->model(); } bool has_data() const { return model_holder_ && model_holder_->data(); } - std::mt19937 rng_; - std::string ini_path_; std::string model_name_; // Used if model_kind_ is kModelFromBuffer. std::string model_path_; From f71e215c9a105ce0433612fbc326bb4fad47ed0a Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 04:23:16 -0800 Subject: [PATCH 17/48] Include missing header. PiperOrigin-RevId: 872835150 Change-Id: I670bddb907b594788bda143c5175bbcb4b427903 --- src/experimental/platform/object_launcher_plugin.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/experimental/platform/object_launcher_plugin.cc b/src/experimental/platform/object_launcher_plugin.cc index 0ba23894..ad065bae 100644 --- a/src/experimental/platform/object_launcher_plugin.cc +++ b/src/experimental/platform/object_launcher_plugin.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include From 8c9ba9bcf5534125862896ab61d4cb98d6aa4bbd Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 04:59:04 -0800 Subject: [PATCH 18/48] Disable plugin on windows. PiperOrigin-RevId: 872845569 Change-Id: I4d80c80a47358a2882b2c63c2d655fe374afe65d --- src/experimental/platform/CMakeLists.txt | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/experimental/platform/CMakeLists.txt b/src/experimental/platform/CMakeLists.txt index 2cdc5f4e..a180b148 100644 --- a/src/experimental/platform/CMakeLists.txt +++ b/src/experimental/platform/CMakeLists.txt @@ -50,7 +50,6 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} interaction.h model_holder.cc model_holder.h - object_launcher_plugin.cc picture_gui.h picture_gui.cc plugin.cc @@ -68,6 +67,13 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} window.h ) +if(NOT WINDOWS) + target_sources(${MUJOCO_PLATFORM_TARGET_NAME} + PUBLIC + object_launcher_plugin.cc + ) +endif() + if(APPLE) set_source_files_properties(window_osx.mm PROPERTIES COMPILE_FLAGS "-x objective-c++") From f9a39413cf19683487713fa26b71f3ff3db1dbc8 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 20 Feb 2026 05:30:12 -0800 Subject: [PATCH 19/48] Fix: Include ancestor DOFs in flex interpolation factorization. The reduced dense factorization for flex interpolation now considers all DOFs in the kinematic chain of the body containing the flex, using mj_bodyChain, instead of only the DOFs directly associated with that body. This is necessary for correctly handling pinned flexes when their parent body is part of a larger kinematic structure. PiperOrigin-RevId: 872854468 Change-Id: Idbe9fb459084dde9e8eb1076c70dbb685c1b0bdb --- src/engine/engine_forward.c | 54 +++++++++++++++++++++---- test/engine/engine_forward_test.cc | 64 ++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 7 deletions(-) diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 6491cdca..c38e6508 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -27,6 +27,7 @@ #include "engine/engine_core_constraint.h" #include "engine/engine_core_smooth.h" #include "engine/engine_derivative.h" +#include "engine/engine_core_util.h" #include "engine/engine_inverse.h" #include "engine/engine_island.h" #include "engine/engine_macro.h" @@ -1189,14 +1190,38 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // flex: reduced dense factorization if (has_flex_interp && !sleep_filter) { + // temporary allocations for body chain + int* chain_dofs = mjSTACKALLOC(d, nv, int); + int* seen_dof = mjSTACKALLOC(d, nv, int); + mju_fillInt(seen_dof, 0, nv); + // identify flex DOFs + // For pinned nodes (body_dofnum==0): use bodyChain to include parent DOFs + // For regular flex nodes: use body_dofadr for one-way coupling for (int f=0; f < m->nflex; f++) { if (m->flex_interp[f]) { int nodenum = m->flex_nodenum[f]; int nodeadr = m->flex_nodeadr[f]; for (int n=0; n < nodenum; n++) { int b = m->flex_nodebodyid[nodeadr + n]; - nflexdofs += m->body_dofnum[b]; + int chain_nnz; + if (m->body_dofnum[b] == 0) { + // Pinned node: use bodyChain to get parent DOFs + chain_nnz = mj_bodyChain(m, b, chain_dofs); + } else { + // Regular flex node: use body's own DOFs only + chain_nnz = m->body_dofnum[b]; + for (int j = 0; j < chain_nnz; j++) { + chain_dofs[j] = m->body_dofadr[b] + j; + } + } + for (int i=0; i < chain_nnz; i++) { + int dof = chain_dofs[i]; + if (!seen_dof[dof]) { + seen_dof[dof] = 1; + nflexdofs++; + } + } } } } @@ -1207,19 +1232,34 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { int* global2local = mjSTACKALLOC(d, nv, int); mju_fillInt(global2local, -1, nv); + // collect unique DOFs in order int cnt = 0; + mju_fillInt(seen_dof, 0, nv); for (int f=0; f < m->nflex; f++) { if (m->flex_interp[f]) { int nodenum = m->flex_nodenum[f]; int nodeadr = m->flex_nodeadr[f]; for (int n=0; n < nodenum; n++) { int b = m->flex_nodebodyid[nodeadr + n]; - int dofnum = m->body_dofnum[b]; - int dofadr = m->body_dofadr[b]; - for (int j=0; j < dofnum; j++) { - flex_dof_indices[cnt] = dofadr + j; - global2local[dofadr + j] = cnt; - cnt++; + int chain_nnz; + if (m->body_dofnum[b] == 0) { + // Pinned node: use bodyChain to get parent DOFs + chain_nnz = mj_bodyChain(m, b, chain_dofs); + } else { + // Regular flex node: use body's own DOFs only + chain_nnz = m->body_dofnum[b]; + for (int j = 0; j < chain_nnz; j++) { + chain_dofs[j] = m->body_dofadr[b] + j; + } + } + for (int i=0; i < chain_nnz; i++) { + int dof = chain_dofs[i]; + if (!seen_dof[dof]) { + seen_dof[dof] = 1; + flex_dof_indices[cnt] = dof; + global2local[dof] = cnt; + cnt++; + } } } } diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index 58e0fe73..79df5022 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -17,6 +17,7 @@ #include "src/engine/engine_forward.h" #include "src/engine/engine_derivative.h" +#include #include #include #include @@ -1816,5 +1817,68 @@ TEST_F(ForwardTest, FlexParentCoupling) { mj_deleteModel(model); } + +TEST_F(ForwardTest, TrilinearPinnedParentWithFreejoint) { + static constexpr char xml[] = R"( + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(m, NotNull()) << error.data(); + mjData* d = mj_makeData(m); + + int parent_id = mj_name2id(m, mjOBJ_BODY, "parent"); + ASSERT_GT(parent_id, 0); + + EXPECT_EQ(m->nflexnode, 8); + EXPECT_EQ(m->body_dofnum[parent_id], 0) << "parent body should have 0 DOFs"; + + int freejoint_body = m->body_parentid[parent_id]; + EXPECT_EQ(m->body_dofnum[freejoint_body], 6) << "freejoint body has 6 DOFs"; + + mj_resetData(m, d); + mj_forward(m, d); + + for (int i = 0; i < 500; i++) { + mj_step(m, d); + + ASSERT_FALSE(mju_isBad(d->qpos[0])) + << "Simulation became unstable at step " << i; + ASSERT_FALSE(mju_isBad(d->qvel[0])) + << "Velocity became unstable at step " << i; + + for (int j = 0; j < m->nq; j++) { + ASSERT_LT(mju_abs(d->qpos[j]), 100.0) + << "Position exploded at step " << i << ", qpos[" << j + << "]=" << d->qpos[j]; + } + for (int j = 0; j < m->nv; j++) { + ASSERT_LT(mju_abs(d->qvel[j]), 1000.0) + << "Velocity exploded at step " << i << ", qvel[" << j + << "]=" << d->qvel[j]; + } + } + + mj_deleteData(d); + mj_deleteModel(m); +} + } // namespace } // namespace mujoco From 7d4b615edb748216cad5f8eb149c5623ac7a6a61 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 20 Feb 2026 06:09:15 -0800 Subject: [PATCH 20/48] Use correct WIN32 define, not WINDOWS. PiperOrigin-RevId: 872866601 Change-Id: I35bb5cc0fdef80cf816a63bb7fc1619ef424c0d0 --- src/experimental/platform/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/experimental/platform/CMakeLists.txt b/src/experimental/platform/CMakeLists.txt index a180b148..d4a58d23 100644 --- a/src/experimental/platform/CMakeLists.txt +++ b/src/experimental/platform/CMakeLists.txt @@ -67,7 +67,7 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} window.h ) -if(NOT WINDOWS) +if(NOT WIN32) target_sources(${MUJOCO_PLATFORM_TARGET_NAME} PUBLIC object_launcher_plugin.cc @@ -85,7 +85,7 @@ elseif(UNIX AND NOT APPLE) target_sources(${MUJOCO_PLATFORM_TARGET_NAME} PUBLIC file_dialog_zenity.cc ) -elseif(WINDOWS) +elseif(WIN32) target_sources(${MUJOCO_PLATFORM_TARGET_NAME} PUBLIC file_dialog_win.cc ) From e5a236774d38e9ae79100673d3a48537bed14449 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 20 Feb 2026 09:20:38 -0800 Subject: [PATCH 21/48] Add has_side_effect for pmap bvh refit in mjx-warp. PiperOrigin-RevId: 872934001 Change-Id: If48cab439718b9e783c4b1da4fb1b1a08521a602 --- .../mjx/third_party/warp/_src/jax_experimental/ffi.py | 6 +++++- mjx/mujoco/mjx/warp/bvh.py | 3 ++- mjx/mujoco/mjx/warp/collision_driver.py | 1 + mjx/mujoco/mjx/warp/ffi.py | 2 ++ mjx/mujoco/mjx/warp/forward.py | 2 ++ mjx/mujoco/mjx/warp/render.py | 1 + mjx/mujoco/mjx/warp/smooth.py | 2 ++ 7 files changed, 15 insertions(+), 2 deletions(-) diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py index 43f63397..9e2e58af 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py +++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py @@ -434,6 +434,7 @@ class FfiCallable: stage_out_argnames, graph_cache_max, module_preload_mode, + has_side_effect=False, ): self.func = func self.name = generate_unique_name(func) @@ -442,6 +443,7 @@ class FfiCallable: self.graph_mode = graph_mode self.output_dims = output_dims self.module_preload_mode = module_preload_mode + self.has_side_effect = has_side_effect self.first_array_arg = None self.call_id = 0 self.call_descriptors = {} @@ -613,7 +615,7 @@ class FfiCallable: out_types, vmap_method=vmap_method, input_output_aliases=self.input_output_aliases, - # has_side_effect=True, # force this function to execute even if outputs aren't used + has_side_effect=self.has_side_effect, ) # preload on the specified devices @@ -1379,6 +1381,7 @@ def jax_callable( stage_out_argnames=None, graph_cache_max: int | None = None, module_preload_mode: ModulePreloadMode = ModulePreloadMode.CURRENT_DEVICE, + has_side_effect: bool = False, ): """Create a JAX callback from an annotated Python function. @@ -1449,6 +1452,7 @@ def jax_callable( stage_out_argnames, graph_cache_max, module_preload_mode, + has_side_effect, ) _FFI_CALLABLE_REGISTRY[key] = callable else: diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index a86bad63..529145ee 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -104,6 +104,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=True, ) out = jf( d.qpos.shape[0], @@ -122,7 +123,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): d.geom_xpos, ctx.key, ) - d = d.tree_replace({'time': d.time + out[0]}) + d = d.tree_replace({}) return d diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index aa88f118..2ae01123 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -289,6 +289,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data): ]), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index 6d0b6d7e..a485cbbb 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/mjx/mujoco/mjx/warp/ffi.py @@ -103,6 +103,7 @@ def jax_callable_variadic_tuple( in_out_argnames: Optional[Sequence[str]] = None, stage_in_argnames: Optional[Sequence[str]] = None, stage_out_argnames: Optional[Sequence[str]] = None, + has_side_effect: bool = False, ): """Wraps a JAX callable to support variadic tuples and dataclasses.""" @@ -134,6 +135,7 @@ def jax_callable_variadic_tuple( in_out_argnames=in_out_argnames, stage_in_argnames=stage_in_argnames, stage_out_argnames=stage_out_argnames, + has_side_effect=has_side_effect, ) flat_args, in_tree = jax.tree.flatten(args) diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index db52a8cd..e48219a8 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -1279,6 +1279,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], @@ -3070,6 +3071,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 73ced206..3b2d3aee 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -145,6 +145,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): ]), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 6e1c0ea8..0e1ec813 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -214,6 +214,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], @@ -408,6 +409,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): ]), stage_out_argnames=set(['ten_length']), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], From 0f36b0fc2e6a878629452d98ad4abce263647e99 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 20 Feb 2026 10:28:15 -0800 Subject: [PATCH 22/48] Import NVIDIA/warp from GitHub. PiperOrigin-RevId: 872969263 Change-Id: Ic890cb53c32593c158a1f3926906b492c7f85348 --- mjx/cuda_requirements.txt | 10 ++++---- .../warp/_src/jax_experimental/ffi.py | 24 ++++++++++++++----- mjx/pyproject.toml | 2 +- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt index deca223c..0f9ae0d5 100644 --- a/mjx/cuda_requirements.txt +++ b/mjx/cuda_requirements.txt @@ -16,8 +16,8 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \ jax-cuda12-pjrt==0.4.30; python_version == '3.9' \ --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \ --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb -warp-lang==1.11.0 \ - --hash=sha256:3a4f1c9a6e721d7de7d6dad6b242c54afaf20c6e14a767c0da03e5e963fcc13c \ - --hash=sha256:524dce20de6162ba25333552168ebf430973050e00d9f8116b8df41a60d25d6e \ - --hash=sha256:1ae6cfc226107f96e4d495b41a3dab32488e8ee8f074b0e1bcaf22e7fb8c904d \ - --hash=sha256:80d8493cbe243a3510134f3af289646d7bd7484217a30ecf565d676466ef8a5e +warp-lang==1.11.1 \ + --hash=sha256:1ad11f1fa775269e991a3d55039152c8a504baf86701c849b485cb8e66c49d15 \ + --hash=sha256:8b098f41e71d421d80ee7562e38aa8380ff6b0d3b4c6ee866cfbdef733ac5bdc \ + --hash=sha256:5d0904b0eefcc81f39ba65375427a3de99006088aa43e24a9011263f07d0cd07 \ + --hash=sha256:15dc10aa51fb0fdbe1ca16d52e5fadca35a47ffd9d0c636826506f96bb2e7c41 diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py index 9e2e58af..f5c925dd 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py +++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py @@ -62,11 +62,22 @@ def check_jax_version(): class GraphMode(IntEnum): - NONE = 0 # don't capture a graph - JAX = 1 # let JAX capture a graph - WARP = 2 # let Warp capture a graph - WARP_STAGED = 3 # use Warp graph with staging buffers, copy inside of the graph - WARP_STAGED_EX = 4 # use Warp graph with staging buffers, copy outside of the graph + """CUDA graph capture modes for :func:`warp.jax_experimental.jax_callable`. + + These modes control whether JAX or Warp captures a CUDA graph, and whether + staging buffers are used when capturing with Warp. + """ + + NONE = 0 + """Disable graph capture. Use when operations are not CUDA-graph compatible (for example, host synchronization).""" + JAX = 1 + """Let JAX capture the graph so the callable can be used as a subgraph within a larger JAX capture.""" + WARP = 2 + """Let Warp capture the graph and replay it for matching buffer addresses.""" + WARP_STAGED = 3 + """Capture a Warp graph using staging buffers and insert memcpy nodes inside the graph.""" + WARP_STAGED_EX = 4 + """Capture a Warp graph using staging buffers and perform memcpy outside the graph.""" class ModulePreloadMode(IntEnum): @@ -682,12 +693,13 @@ class FfiCallable: assert num_outputs == self.num_outputs cuda_stream = get_stream_from_callframe(call_frame.contents) + device_ordinal = get_device_ordinal_from_callframe(call_frame.contents) if self.graph_mode == GraphMode.WARP: # check if we already captured an identical call ip = [inputs[i].contents.data for i in self.array_input_indices] op = [outputs[i].contents.data for i in self.array_output_indices] - capture_key = hash((call_id, *ip, *op)) + capture_key = hash((device_ordinal, call_id, *ip, *op)) capture = self.captures.get(capture_key) # launch existing graph diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index 0348ce44..bbfbf7e9 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ [project.optional-dependencies] warp = [ - "warp-lang==1.11.0", + "warp-lang==1.11.1", ] [project.scripts] From c428f675e73e7eb3d21662b16f1b305450731ff1 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Fri, 20 Feb 2026 12:44:39 -0800 Subject: [PATCH 23/48] MuJoCo Warp documentation: Update installation instructions PiperOrigin-RevId: 873032865 Change-Id: I7f75de9ed0d9239b47bc80b321f213aa4014fda6 --- doc/mjwarp/index.rst | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 06e6ad7b..ddef203f 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -87,24 +87,25 @@ via Warp - if this feature is important to you, please chime in on this issue Installation ============ -The beta version of MuJoCo Warp is installed from GitHub. Please note that the beta version of MuJoCo Warp does not -support all versions of MuJoCo, Warp, CUDA, NVIDIA drivers, etc. +**From PyPI:** + +.. code-block:: shell + + pip install mujoco-warp + +**From source:** .. code-block:: shell git clone https://github.com/google-deepmind/mujoco_warp.git cd mujoco_warp - python3 -m venv env - source env/bin/activate - pip install --upgrade pip - pip install uv - uv pip install -e .[dev,cuda] + uv sync --all-extras -Test the Installation +To make sure everything is working: .. code-block:: shell - pytest + uv run pytest -n 8 .. _MJW_Usage: From 9d6d9089e54ed815dae48da711e02b7fce54d93a Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 21 Feb 2026 09:14:24 -0800 Subject: [PATCH 24/48] Improve solver parameter documentation, in particular for frictional constraints. PiperOrigin-RevId: 873364458 Change-Id: Ie4da21abc2ee54d9eaaefc81b3d2624a520ad96e --- doc/XMLreference.rst | 8 +++--- doc/computation/index.rst | 23 +++++++++++------ doc/modeling.rst | 54 +++++++++++++++++++++++++++------------ 3 files changed, 58 insertions(+), 27 deletions(-) diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 5d08dee3..2e12173c 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -2263,7 +2263,8 @@ rotations as unit quaternions. .. _body-joint-solimpfriction: :at:`solreffriction`, :at:`solimpfriction` - Constraint solver parameters for simulating dry friction. See :ref:`CSolver`. + Constraint solver parameters for simulating dry friction. + See also :ref:`Friction`. .. _body-joint-stiffness: @@ -4074,7 +4075,7 @@ friction can only be created with this element. Note that as with other :at:`solreffriction` attributes, the constraint violation is identically 0. Therefore, when using positive semantics :at:`solreffriction[1]` is ignored, while for negative semantics :at:`solreffriction[0]` is - ignored. See :ref:`CSolver` for more details. + ignored. See :ref:`Friction` for more details. .. _contact-pair-margin: @@ -4936,7 +4937,8 @@ length X, as in the clip on the right of `this example model .. _tendon-spatial-solimpfriction: :at:`solreffriction`, :at:`solimpfriction` - Constraint solver parameters for simulating dry friction in the tendon. See :ref:`CSolver`. + Constraint solver parameters for simulating dry friction in the tendon. + See also :ref:`Friction`. .. _tendon-spatial-margin: diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 2534e498..422e4860 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -1335,6 +1335,8 @@ It is a vector with dimensionality :math:`\nq` satisfying :math:`0` section of the +Modeling chapter. -To summarize, the user specifies the vectors of impedance coefficients :math:`0 0` -and stiffness coefficients :math:`k > 0`. The quantities :math:`R, \ar` are then computed by MuJoCo as shown above, and -the selected optimization algorithm is applied to solve problem :eq:`eq:dual`. As explained in the :ref:`solver -parameters ` section of the Modeling chapter, MuJoCo offers additional automation for setting :math:`d, b, k` -so as to achieve critical damping, or model a soft contact layer by varying :math:`d` with distance. +To summarize, the constraint behavior is determined by three per-constraint quantities: impedance :math:`0 0` and stiffness :math:`k \geq 0`. These are computed from the :at:`solimp` and :at:`solref` attributes as +described in the :ref:`solver parameters ` section of the Modeling chapter, which also offers additional +automation (e.g., achieving critical damping, or varying :math:`d` with distance to model a soft contact layer). The +quantities :math:`R, \ar` are then computed from :eq:`eq:impedance_R` and :eq:`eq:aref`, and the selected optimization +algorithm is applied to solve problem :eq:`eq:dual`. .. _soCones: diff --git a/doc/modeling.rst b/doc/modeling.rst index fee781e9..fc5979fd 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -274,7 +274,8 @@ experiment interactively with parameter settings or implement continuation metho Here we focus on a single scalar constraint. Using slightly different notation from the Computation chapter, let :math:`\ac` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in friction dimensions), :math:`k` and :math:`b` the stiffness and damping of the virtual spring used to define the -reference acceleration :math:`\ar = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`\au` the +reference acceleration :math:`\ar = -b v - k r` (see :eq:`eq:aref`). +Let :math:`d` be the constraint impedance, and :math:`\au` the acceleration in the absence of constraint force. Our earlier analysis revealed that the dynamics in constraint space are approximately @@ -344,8 +345,7 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec constraint becomes active; for contacts this margin is :ref:`margin`-:ref:`gap`. Limit and contact constraints are active when :math:`r < 0` (penetration). - For friction loss or friction dimensions of elliptic cones, the violation :math:`r` is identically zero, so - only :math:`d(0)` affects these constraints, all other :at:`solimp` values are ignored. + For frictional constraints, see :ref:`Friction`. .. _solimp0: @@ -385,25 +385,27 @@ There are two formats for this attribute, determined by the sign of the numbers. specification is considered to be in the :math:`(\text{timeconst}, \text{dampratio})` format. If negative it is in the "direct" :math:`(-\text{stiffness}, -\text{damping})` format. -Frictional constraints whose residual is identically 0 have first-order dynamics and the mass-spring-damper analysis -below does not apply. In this case the time constant is the rate of exponential decay of the constraint velocity, -and the damping ratio is ignored. Equivalently, in the direct format, the :math:`\text{stiffness}` is ignored. +For frictional constraints, the mass-spring-damper analysis below does not directly apply; +see :ref:`Friction`. **solref :** real(2), "0.02 1" We first describe the default, positive-value format where the two numbers are :math:`(\text{timeconst}, \text{dampratio})`. + .. _soRefScaling: + The idea here is to re-parameterize the model in terms of the time constant and damping ratio of a mass-spring-damper - system. By "time constant" we mean the inverse of the natural frequency times the damping ratio. In this case we use - a mass-spring-damper model to compute :math:`k, b` after suitable scaling. Note that the effective stiffness - :math:`d(r) \cdot k` and damping :math:`d(r) \cdot b` are scaled by the impedance :math:`d(r)` which is a function of - the distance :math:`r`. Thus we cannot always achieve the specified mass-spring-damper properties, unless we - completely undo the scaling by :math:`d`. But the latter is undesirable because it would ruin the interpolating - property, in particular the limit :math:`d=0` would no longer disable the constraint. Instead we scale the stiffness - and damping so that the damping ratio remains constant, while the time constant increases when :math:`d(r)` gets - smaller. The scaling formulas are + system. By "time constant" we mean the inverse of the natural frequency times the damping ratio. Now recall that the + products :math:`d \cdot k` and :math:`d \cdot b` in :eq:`eq:constraint` are the effective stiffness and damping in + constraint space. Because the impedance :math:`d(r)` varies with the + position residual :math:`r`, we cannot achieve constant mass-spring-damper properties; completely undoing the scaling + by :math:`d` is undesirable because the limit :math:`d = 0` would no longer disable the constraint. Instead, we + absorb one factor of :math:`d(r)` into :math:`k` (but not into :math:`b`), so that the damping ratio remains constant + while the time constant scales with :math:`d(r)`. The formulas are .. math:: + :label: eq:solref_standard + \begin{aligned} b &= 2 / (d_\text{width}\cdot \text{timeconst}) \\ k &= d(r) / (d_\text{width}^2 \cdot \text{timeconst}^2 \cdot \text{dampratio}^2) \\ @@ -414,7 +416,7 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: can go unstable. This is enforced internally, unless the :ref:`refsafe` attribute of :ref:`flag ` is set to false. The :math:`\text{dampratio}` parameter would normally be set to 1, corresponding to critical damping. Smaller values result in under-damped or bouncy constraints, while larger values result in - over-damped constraints. Combining the above formula with :eq:`eq:constraint`, we can derive the following result. + over-damped constraints. Combining :eq:`eq:solref_standard` with :eq:`eq:constraint`, we can derive the following If the reference acceleration is given using the positive number format and the impedance is constant :math:`d = d_0 = d_\text{width}`, then the penetration depth at rest is @@ -427,12 +429,14 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: interact. The scaling formulas are .. math:: + :label: eq:solref_direct + \begin{aligned} b &= \text{damping} / d_\text{width} \\ k &= \text{stiffness} \cdot d(r) / d_\text{width}^2 \\ \end{aligned} - Similarly to the above derivation, if the reference acceleration is given using the negative number format and the + Similarly to the derivation following :eq:`eq:solref_standard`, if the reference acceleration is given using the impedance is constant, then the penetration depth at rest is .. math:: @@ -449,6 +453,24 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: A :math:`\text{dampratio}` of 1 in the positive-value format is equivalent to :math:`\text{damping} = 2 \sqrt{ \text{stiffness} }` in the direct format. +.. _CSolverFriction: + +Friction +^^^^^^^^ + +Friction loss constraints (in joints and tendons) and friction dimensions of elliptic contact cones have zero position +violation: :math:`r \equiv 0`. This simplifies the constraint model (see also :ref:`soParameters`): + +- The **impedance** is always :math:`d_0` (:at:`solimp[0]`), since :math:`d(r)` is evaluated at :math:`r=0`. + The sigmoid shape parameters (:math:`\text{width}`, :math:`\text{midpoint}`, :math:`\text{power}`) have no effect. +- The dynamics are **first-order** (exponential decay of constraint velocity, no spring): the stiffness :math:`k` is + always 0. +- In the standard :at:`solref` format, the time constant controls exponential velocity decay. The damping ratio is + ignored (it only appears in the :math:`k` formula). +- In the direct :at:`solref` format, the damping (second value) is used but the stiffness (first value) is ignored. +- :math:`d_\text{width}` (:at:`solimp[1]`) still affects the damping :math:`b` as a scaling denominator + (:eq:`eq:solref_standard`, :eq:`eq:solref_direct`), even though it does not affect the impedance. + .. _CContact: Contact parameters From 6b9e0b9e15bf1e6a5d75e3e084b3181aa94466fc Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 23 Feb 2026 02:23:07 -0800 Subject: [PATCH 25/48] MJX API documentation: add `Any` to exclude-members PiperOrigin-RevId: 873954713 Change-Id: Iafbdb571c4e5b0491d7d91159e84874c58c01a71 --- doc/mjx_api.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/mjx_api.rst b/doc/mjx_api.rst index e84a8d90..291674e3 100644 --- a/doc/mjx_api.rst +++ b/doc/mjx_api.rst @@ -7,6 +7,7 @@ MJX API :special-members: False :private-members: False :exclude-members: + Any, __init__, __format__, __new__, From ba0a5fe4db43038f0838c06eaeda1eec915639f5 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 23 Feb 2026 02:23:31 -0800 Subject: [PATCH 26/48] Fix section underline length in MJWarp docs PiperOrigin-RevId: 873954878 Change-Id: I40b668b24d04721d8fd4bb61bb5ac7b902ea31f1 --- doc/mjwarp/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index ddef203f..ed5af221 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -489,7 +489,7 @@ Key features: `Warp's BVHs `__. Basic Usage ----------- +----------- Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures, rendering specific fields, and output buffers. From 2bf9539234e0db5e92c7fbb7acb40f63c9e19a12 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 23 Feb 2026 05:19:13 -0800 Subject: [PATCH 27/48] MJX put_data with impl='warp' PiperOrigin-RevId: 874017611 Change-Id: I23a25727d41cd888d90f2ea15b2b534eba5583be --- mjx/mujoco/mjx/_src/io.py | 56 ++++++++++++++++++------ mjx/mujoco/mjx/_src/io_test.py | 78 +++++++++++++++++++++++++--------- mjx/mujoco/mjx/viewer.py | 12 ++---- 3 files changed, 104 insertions(+), 42 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 4199936a..670c93e1 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -903,17 +903,6 @@ def _make_data_warp( data = jax.device_put(data, device=device) - with wp.ScopedDevice('cuda:0'): # pylint: disable=undefined-variable - # Warm-up the warp kernel cache. - # TODO(robotics-simulation): remove this warmup compilation once warp - # stops unloading modules during XLA graph capture for tile kernels. - # pylint: disable=undefined-variable - dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, njmax=njmax) - mw = mjwp.put_model(m) - _ = mjwp.step(mw, dw) - # pylint: enable=undefined-variable - del dw, mw - return data @@ -1363,6 +1352,44 @@ def _put_data_cpp( return _strip_weak_type(data) +def _put_data_warp( + m: mujoco.MjModel, + d: mujoco.MjData, + device: Optional[jax.Device] = None, + naconmax: Optional[int] = None, + njmax: Optional[int] = None, +) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" + + with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable + dw = mjwp.put_data(m, d, nworld=1, naconmax=naconmax, njmax=njmax) # pylint: disable=undefined-variable + + fields = _put_data_public_fields(d) + for k in fields: + if not hasattr(dw, k): + continue + field = _wp_to_np_type(getattr(dw, k)) + if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access + field = field.reshape(field.shape[1:]) + fields[k] = field + + impl_fields = {} + for k in mjxw.types.DataWarp.__annotations__.keys(): + field = _get_nested_attr(dw, k, split='__') + field = _wp_to_np_type(field) + if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access + field = field.reshape(field.shape[1:]) + impl_fields[k] = field + + data = types.Data( + **fields, + _impl=mjxw.types.DataWarp(**impl_fields), + ) + + data = jax.device_put(data, device=device) + return data + + def put_data( m: mujoco.MjModel, d: mujoco.MjData, @@ -1393,7 +1420,6 @@ def put_data( an mjx.Data placed on device DeprecationWarning: if nconmax is used """ - del njmax if nconmax is not None: warnings.warn( 'nconmax will be deprecated in mujoco-mjx>=3.5. Use naconmax instead.', @@ -1410,8 +1436,10 @@ def put_data( return _put_data_cpp( m, d, device, dummy_arg_for_batching=dummy_arg_for_batching ) - - # TODO(robotics-team): implement put_data_warp + elif impl == types.Impl.WARP: + _check_warp_installed() + naconmax = nconmax if naconmax is None else naconmax + return _put_data_warp(m, d, device, naconmax, njmax) raise NotImplementedError( f'put_data for implementation "{impl}" not implemented yet.' diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 8382a6cf..1cd0afe9 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -494,9 +494,15 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d._impl.contact__dist.shape[0], 9) self.assertEqual(d._impl.efc__pos.shape[0], 23) - @parameterized.parameters('jax', 'c', 'cpp') + @parameterized.parameters('jax', 'c', 'cpp', 'warp') def test_put_data(self, impl: str): """Test that put_data puts the correct data for dense and sparse.""" + if impl == 'warp': + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp is not installed.') + if not mjx_io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device.') + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) @@ -517,6 +523,15 @@ class DataIOTest(parameterized.TestCase): ) ) + # xmat, ximat, geom_xmat are all shape transformed + np.testing.assert_allclose(dx.xmat.reshape((-1, 9)), d.xmat) + np.testing.assert_allclose(dx.ximat.reshape((-1, 9)), d.ximat) + np.testing.assert_allclose(dx.geom_xmat.reshape((-1, 9)), d.geom_xmat) + np.testing.assert_allclose(dx.site_xmat.reshape((-1, 9)), d.site_xmat) + + # tendon length is correct + np.testing.assert_allclose(dx.ten_length, d.ten_length) + if impl == 'jax': # check that qM is transformed properly qm = np.zeros((m.nv, m.nv), dtype=np.float64) @@ -530,6 +545,21 @@ class DataIOTest(parameterized.TestCase): self.assertTrue(hasattr(dx._impl, 'pointer_lo')) self.assertTrue(hasattr(dx._impl, 'pointer_hi')) return # cpp does not populate other fields in _impl + elif impl == 'warp': + qm = np.zeros((m.nv, m.nv), dtype=np.float64) + mujoco.mj_fullM(m, qm, d.qM) + np.testing.assert_allclose(dx._impl.qM, qm) + # TODO(taylorhowell): test efc__J + np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3]) + + # tendon impl data is correct + np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,))) + np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,))) + np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2))) + np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6))) + + if impl == 'warp': + return # 4 contacts, 2 for each capsule against the plane self.assertEqual(dx._impl.contact.dist.shape, (4,)) @@ -542,23 +572,6 @@ class DataIOTest(parameterized.TestCase): ) np.testing.assert_allclose(dx._impl.contact.frame[1:], 0) - # xmat, ximat, geom_xmat are all shape transformed - self.assertEqual(dx.xmat.shape, (3, 3, 3)) - self.assertEqual(dx.ximat.shape, (3, 3, 3)) - self.assertEqual(dx.geom_xmat.shape, (3, 3, 3)) - self.assertEqual(dx.site_xmat.shape, (1, 3, 3)) - np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat) - np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat) - np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat) - np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat) - - # tendon data is correct - np.testing.assert_allclose(dx.ten_length, d.ten_length) - np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,))) - np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,))) - np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2))) - np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6))) - # efc_ are also shape transformed and padded self.assertEqual(dx._impl.efc_J.shape, (45, 8)) # nefc, nv d_efc_j = d.efc_J.reshape((-1, 8)) @@ -583,7 +596,9 @@ class DataIOTest(parameterized.TestCase): d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) dx_sparse = mjx.put_data(m, d, impl=impl) - np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8) + np.testing.assert_allclose( + dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8 + ) # check sparse mass matrices are correct np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8) @@ -604,6 +619,31 @@ class DataIOTest(parameterized.TestCase): elif impl == 'c': np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8) + def test_put_data_warp_ndim(self): + """Tests that put_data produces expected dimensions for Warp fields.""" + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp is not installed.') + if not mjx_io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device.') + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d, impl='warp') + + def check_ndim(path, x): + k = _get_name_from_path(path) + if k not in mjxw_types._NDIM['Data']: + return + is_batched = mjxw_types._BATCH_DIM['Data'][k] + expected_ndim = mjxw_types._NDIM['Data'][k] - is_batched + if not hasattr(x, 'ndim'): + return + msg = f'Field {k} has ndim {x.ndim} but expected {expected_ndim}' + self.assertEqual(x.ndim, expected_ndim, msg) + + _ = jax.tree.map_with_path(check_ndim, dx) + @parameterized.parameters( ('jax', False), ('jax', True), ('c', False), ('c', True) ) diff --git a/mjx/mujoco/mjx/viewer.py b/mjx/mujoco/mjx/viewer.py index 075e243e..ddd88696 100644 --- a/mjx/mujoco/mjx/viewer.py +++ b/mjx/mujoco/mjx/viewer.py @@ -88,15 +88,9 @@ def _main(argv: Sequence[str]) -> None: m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value) d = mujoco.MjData(m) mx = mjx.put_model(m, impl=_IMPL.value) - if _IMPL.value == 'warp': - # TODO(btaba): use put_data. - dx = mjx.make_data( - m, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value - ) - else: - dx = mjx.put_data( - m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value - ) + dx = mjx.put_data( + m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value + ) print(f'Default backend: {jax.default_backend()}') step_fn = mjx.step From 37e993f67e3d01e49b02f94f5637f34533f1f0b2 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 23 Feb 2026 07:10:49 -0800 Subject: [PATCH 28/48] Improve upper bound for `nJmom` by computing maximum number of dofs for tendon transmissions PiperOrigin-RevId: 874057077 Change-Id: I83736719cc2d91a899cd5997ea4a91d21cf06c9e --- src/user/user_model.cc | 80 ++++++++++++++++++++++-------------------- src/user/user_model.h | 2 ++ 2 files changed, 44 insertions(+), 38 deletions(-) diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 2bddb3a8..1c3e7e24 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -3171,7 +3171,46 @@ void mjCModel::CopyPlugins(mjModel* m) { -// compute non-zeros in actuator_moment matrix +// compute number of dofs for a given tendon +int mjCModel::CountTendonDofs(const mjModel* m, int id) { + std::vector dof_used(m->nv, false); + int nv = m->nv; + int adr = m->tendon_adr[id]; + int num = m->tendon_num[id]; + + if (m->wrap_type[adr] == mjWRAP_JOINT) { + return num; + } + + std::fill(dof_used.begin(), dof_used.end(), false); + for (int j = 0; j < num; j++) { + int type = m->wrap_type[adr + j]; + int bodyid = -1; + if (type == mjWRAP_SITE) { + bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; + } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { + bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; + } + if (bodyid > 0) { + int bid = bodyid; + while (bid > 0) { + int bdofadr = m->body_dofadr[bid]; + int bdofnum = m->body_dofnum[bid]; + for (int k = 0; k < bdofnum; k++) { + dof_used[bdofadr + k] = true; + } + bid = m->body_parentid[bid]; + } + } + } + + int count = 0; + for (int j = 0; j < nv; j++) { + count += dof_used[j]; + } + return count; +} + int mjCModel::CountNJmom(const mjModel* m) { int nu = m->nu; int nv = m->nv; @@ -3206,7 +3245,7 @@ int mjCModel::CountNJmom(const mjModel* m) { break; case mjTRN_TENDON: - count += nv; + count += CountTendonDofs(m, id); break; case mjTRN_SITE: @@ -3228,46 +3267,11 @@ int mjCModel::CountNJmom(const mjModel* m) { // compute non-zeros in ten_J matrix int mjCModel::CountNJten(const mjModel* m) { - int nv = m->nv; int ntendon = m->ntendon; - std::vector dof_bitmap(nv, false); int count = 0; for (int i = 0; i < ntendon; i++) { - int adr = m->tendon_adr[i]; - int num = m->tendon_num[i]; - - if (m->wrap_type[adr] == mjWRAP_JOINT) { - count += num; - continue; - } - - std::fill(dof_bitmap.begin(), dof_bitmap.end(), false); - for (int j = 0; j < num; j++) { - int type = m->wrap_type[adr + j]; - int bodyid = -1; - if (type == mjWRAP_SITE) { - bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; - } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { - bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; - } - if (bodyid > 0) { - int bid = bodyid; - while (bid > 0) { - int bdofadr = m->body_dofadr[bid]; - int bdofnum = m->body_dofnum[bid]; - for (int k = 0; k < bdofnum; k++) { - dof_bitmap[bdofadr + k] = true; - } - bid = m->body_parentid[bid]; - } - } - } - - // only count unique dofs - for (int j = 0; j < nv; j++) { - count += dof_bitmap[j]; - } + count += CountTendonDofs(m, i); } return count; diff --git a/src/user/user_model.h b/src/user/user_model.h index 93e4d3f4..2a3ac355 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -368,6 +368,8 @@ class mjCModel : public mjCModel_, private mjSpec { void CopyTree(mjModel*); // copy objects inside kinematic tree void FinalizeSimple(mjModel* m); // finalize simple bodies/dofs including tendon information void CopyPlugins(mjModel*); // copy plugin data + int CountTendonDofs(const mjModel* m, // compute number of dofs for a given tendon + int id); 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 From a043df6bdf5c11cf375e78f8210f642bdd61f97a Mon Sep 17 00:00:00 2001 From: Michael Moss Date: Mon, 23 Feb 2026 07:28:45 -0800 Subject: [PATCH 29/48] Update supported Python versions and remove obsolete 3.9 references. PiperOrigin-RevId: 874063785 Change-Id: Ia4b2956b632aae5f0211fa766b6d413edcba8fe2 --- README.md | 2 +- mjx/pyproject.toml | 4 ++-- mjx/requirements.txt | 28 +--------------------------- python/build_requirements.txt | 12 ------------ python/make_sdist_requirements.txt | 2 -- python/pyproject.toml | 4 ++-- 6 files changed, 6 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index 9eabc498..98bee9e2 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,7 @@ Users who wish to build MuJoCo from source should consult the [build from source] section of the documentation. However, note that the commit at the tip of the `main` branch may be unstable. -### Python (>= 3.9) +### Python (>= 3.10) The native Python bindings, which come pre-packaged with a copy of MuJoCo, can be installed from [PyPI] via: diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index bbfbf7e9..4ee26dd5 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -17,13 +17,13 @@ classifiers = [ "Intended Audience :: Science/Research", "Natural Language :: English", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ "absl-py", "etils[epath]", diff --git a/mjx/requirements.txt b/mjx/requirements.txt index f0ef43d6..4e5b56f0 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -2,15 +2,11 @@ absl-py==2.1.0 \ --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 etils[epath]==1.10.0; python_version >= '3.10' \ --hash=sha256:0777fe60a234b4c65ca53470fc64f2dd2d0c6bca7fcc623fdaa8d7fa5a317098 -etils[epath]==1.5.2; python_version == '3.9' \ - --hash=sha256:6dc882d355e1e98a5d1a148d6323679dc47c9a5792939b9de72615aa4737eb0b jax==0.5.3; python_version >= '3.10' and (sys_platform != 'darwin' or platform_machine != 'x86_64') \ --hash=sha256:1483dc237b4f47e41755d69429e8c3c138736716147cd43bb2b99b259d4e3c41 \ --hash=sha256:f17fcb0fd61dc289394af6ce4de2dada2312f2689bb0d73642c6f026a95fbb2c jax==0.4.38; python_version >= '3.10' and sys_platform == 'darwin' and platform_machine == 'x86_64' \ --hash=sha256:78987306f7041ea8500d99df1a17c33ed92620c2268c4c3677fb24e06712be64 -jax==0.4.30; python_version == '3.9' \ - --hash=sha256:289b30ae03b52f7f4baf6ef082a9f4e3e29c1080e22d13512c5ecf02d5f1a55b jaxlib==0.5.3; python_version >= '3.10' and (sys_platform != 'darwin' or platform_machine != 'x86_64') \ --hash=sha256:48ff5c89fb8a0fe04d475e9ddc074b4879a91d7ab68a51cec5cd1e87f81e6c47 \ --hash=sha256:972400db4af6e85270d81db5e6e620d31395f0472e510c50dfcd4cb3f72b7220 \ @@ -49,12 +45,6 @@ jaxlib==0.4.38; python_version >= '3.10' and sys_platform == 'darwin' and platfo --hash=sha256:248cca3771ebf24b070f49701364ceada33e6139445b06c782cca5ac5ad92bf4 \ --hash=sha256:2ce77ba8cda9259a4bca97afc1c722e4291a6c463a63f8d372c6edc85117d625 \ --hash=sha256:4103db0b3a38a5dc132741237453c24d8547290a22079ba1b577d6c88c95300a -jaxlib==0.4.30; python_version == '3.9' \ - --hash=sha256:54987e97a22db70f3829b437b9329e4799d653634bacc8b398554d3b90c76b2a \ - --hash=sha256:f74a6b0e09df4b5e2ee399ebb9f0e01190e26e84ccb0a758fadb516415c07f18 \ - --hash=sha256:11602d5556e8baa2f16314c36518e9be4dfae0c2c256a361403fb29dc9dc79a4 \ - --hash=sha256:3d31e01191ce8052bd611aaf16ff967d8d0ec0b63f1ea4b199020cecb248d667 \ - --hash=sha256:ea3a00005faafbe3c18b178d3b534208b3b4027b2be6230227e7b87ce399fc29 pip==25.2 \ --hash=sha256:578283f006390f85bb6282dffb876454593d637f5d1be494b5202ce4877e71f2 \ --hash=sha256:6d67a2b4e7f14d8b31b8b52648866fa717f45a1eb70e83002f4331d07e953717 @@ -91,12 +81,6 @@ scipy==1.14.1; python_version >= '3.10' \ --hash=sha256:8bddf15838ba768bb5f5083c1ea012d64c9a444e16192762bd858f1e126196d0 \ --hash=sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3 \ --hash=sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389 -scipy==1.13.1; python_version == '3.9' \ - --hash=sha256:392e4ec766654852c25ebad4f64e4e584cf19820b980bc04960bca0b0cd6eaa2 \ - --hash=sha256:637e98dcf185ba7f8e663e122ebf908c4702420477ae52a04f9908707456ba4d \ - --hash=sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004 \ - --hash=sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24 \ - --hash=sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5 setuptools==78.1.1 \ --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d @@ -116,8 +100,6 @@ zipp==3.21.0 \ --hash=sha256:ac1bbe05fd2991f160ebce24ffbac5f6d11d83dc90891255885223d42b3cd931 # Transitive dependencies of jax and jaxlib -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b ml-dtypes==0.5.0 \ --hash=sha256:cb5cc7b25acabd384f75bbd78892d0c724943f3e2e1986254665a1aa10982e07 \ --hash=sha256:54415257f00eb44fbcc807454efac3356f75644f1cbfc2d4e5522a72ae1dacab \ @@ -167,15 +149,7 @@ numpy==2.1.3; python_version >= '3.10' \ --hash=sha256:6a4825252fcc430a182ac4dee5a505053d262c807f8a924603d411f6718b88fd \ --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ - --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff \ -numpy==2.0.2; python_version == '3.9' \ - --hash=sha256:a3f4ab0caa7f053f6797fcd4e1e25caee367db3112ef2b6ef82d749530768c73 \ - --hash=sha256:f26b258c385842546006213344c50655ff1555a9338e2e5e02a0756dc3e803dd \ - --hash=sha256:1e795a8be3ddbac43274f18588329c72939870a16cae810c2b73461c40718ab1 \ - --hash=sha256:97032a27bd9d8988b9a97a8c4d2c9f2c15a81f61e2f21404d7e8ef00cb5be729 \ - --hash=sha256:2b2955fa6f11907cf7a70dab0d0755159bca87755e831e47932367fc8f2f2d0b \ - --hash=sha256:423e89b23490805d2a5a96fe40ec507407b8ee786d66f7328be214f9679df6dd \ - --hash=sha256:9059e10581ce4093f735ed23f3b9d283b9d517ff46009ddd485f1747eb22653c + --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff opt-einsum==3.4.0 \ --hash=sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 7e2ec063..49b3750f 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -6,8 +6,6 @@ build==1.2.2.post1 \ --hash=sha256:1d61c0887fa860c01971625baae8bdd338e517b836a2f70dd1f7aa3a6b2fc5b5 etils[epath]==1.10.0; python_version >= '3.10' \ --hash=sha256:0777fe60a234b4c65ca53470fc64f2dd2d0c6bca7fcc623fdaa8d7fa5a317098 -etils[epath]==1.5.2; python_version == '3.9' \ - --hash=sha256:6dc882d355e1e98a5d1a148d6323679dc47c9a5792939b9de72615aa4737eb0b glfw==2.9.0 \ --hash=sha256:9aa3ae51601601c53838315bd2a03efb1e6bebecd072b2f64ddbd0b2556d511a \ --hash=sha256:8e4fbff88e4e953bb969b6813195d5de4641f886530cc8083897e56b00bf2c8e \ @@ -46,14 +44,6 @@ numpy==2.1.3; python_version >= '3.10' \ --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff -numpy==2.0.2; python_version == '3.9' \ - --hash=sha256:a3f4ab0caa7f053f6797fcd4e1e25caee367db3112ef2b6ef82d749530768c73 \ - --hash=sha256:f26b258c385842546006213344c50655ff1555a9338e2e5e02a0756dc3e803dd \ - --hash=sha256:1e795a8be3ddbac43274f18588329c72939870a16cae810c2b73461c40718ab1 \ - --hash=sha256:97032a27bd9d8988b9a97a8c4d2c9f2c15a81f61e2f21404d7e8ef00cb5be729 \ - --hash=sha256:2b2955fa6f11907cf7a70dab0d0755159bca87755e831e47932367fc8f2f2d0b \ - --hash=sha256:423e89b23490805d2a5a96fe40ec507407b8ee786d66f7328be214f9679df6dd \ - --hash=sha256:9059e10581ce4093f735ed23f3b9d283b9d517ff46009ddd485f1747eb22653c pip==24.3.1 \ --hash=sha256:3790624780082365f47549d032f3770eeb2b1e8bd1f7b2e02dace1afa361b4ed PyOpenGL==3.1.7 \ @@ -73,8 +63,6 @@ pyelftools==0.31; platform_system == 'Linux' \ # Transitive dependencies of build colorama==0.4.6; platform_system == 'Windows' \ --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6 -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b packaging==24.2 \ --hash=sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759 pyproject_hooks==1.2.0 \ diff --git a/python/make_sdist_requirements.txt b/python/make_sdist_requirements.txt index f961524c..da33bbb9 100644 --- a/python/make_sdist_requirements.txt +++ b/python/make_sdist_requirements.txt @@ -11,8 +11,6 @@ setuptools==78.1.1 \ # Transitive dependencies of build colorama==0.4.6; platform_system == 'Windows' \ --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6 -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b packaging==24.2 \ --hash=sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759 pyproject_hooks==1.2.0 \ diff --git a/python/pyproject.toml b/python/pyproject.toml index f23b9bf5..f3f10bb7 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -9,7 +9,7 @@ authors = [ {name = "Google DeepMind", email = "mujoco@deepmind.com"}, ] description = "MuJoCo Physics Simulator" -requires-python = ">=3.9" +requires-python = ">=3.10" license = "Apache-2.0" classifiers = [ "Development Status :: 5 - Production/Stable", @@ -17,10 +17,10 @@ classifiers = [ "Intended Audience :: Science/Research", "Natural Language :: English", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", ] dependencies = [ From 6ec808e2ce3af289ab3ddea6f6628eb11243e245 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 23 Feb 2026 07:38:02 -0800 Subject: [PATCH 30/48] Fix attach for spatial tendon. Fixes #3119. PiperOrigin-RevId: 874067521 Change-Id: If8d415278dde2bccf269cc90e8c176e3d7edefa6 --- doc/changelog.rst | 5 ++ src/user/user_objects.cc | 8 ++- test/user/user_api_test.cc | 127 +++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 2 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index eeb6fab0..21552163 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,11 @@ MJX - Add batch rendering support for MJX-Warp. See the :ref:`MJX-Warp batch rendering` section for details. +Bug fixes +^^^^^^^^^ + +- Fixed a bug where :ref:`mjs_attach` silently dropped spatial tendons with wrapping geometries that had no + ``sidesite`` attribute (:issue:`3119`, reported by :github:user:`tomstewart89`). Version 3.5.0 (February 12, 2026) --------------------------------- diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 9dfa3346..4b5ee2a4 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -6267,12 +6267,16 @@ void mjCTendon::ResolveReferences(const mjCModel* m) { try { // look for wrapped element with namespace path[i]->name = prefix + pname + suffix; - path[i]->sidesite = prefix + psidesite + suffix; + if (!psidesite.empty()) { + path[i]->sidesite = prefix + psidesite + suffix; + } path[i]->ResolveReferences(m); } catch(mjCError) { // remove namespace from wrap names path[i]->name = pname; - path[i]->sidesite = psidesite; + if (!psidesite.empty()) { + path[i]->sidesite = psidesite; + } path[i]->ResolveReferences(m); nfailure++; } diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 73e36a48..e9f519e5 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -1079,6 +1079,133 @@ TEST_F(MujocoTest, AttachSame) { mj_deleteModel(m_expected); } +TEST_F(MujocoTest, AttachSpatialTendonWithoutSidesite) { + static constexpr char xml_parent[] = R"( + + + + + + + )"; + + static constexpr char xml_child[] = R"( + + + + + + + + + + + + + + + + + + + + + + )"; + + std::array er; + mjSpec* parent = mj_parseXMLString(xml_parent, 0, er.data(), er.size()); + ASSERT_THAT(parent, NotNull()) << er.data(); + mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size()); + ASSERT_THAT(child, NotNull()) << er.data(); + + mjsBody* parent_body = mjs_findBody(parent, "parent_body"); + ASSERT_THAT(parent_body, NotNull()); + mjsSite* attach_site = mjs_addSite(parent_body, 0); + mjs_setName(attach_site->element, "attach_site"); + + mjs_attach(attach_site->element, + mjs_findBody(child, "child_body")->element, "", "_child"); + + EXPECT_THAT(mjs_findElement(parent, mjOBJ_TENDON, + "tendon_with_sidesite_child"), NotNull()); + EXPECT_THAT(mjs_findElement(parent, mjOBJ_TENDON, + "tendon_without_sidesite_child"), NotNull()); + + mjModel* model = mj_compile(parent, nullptr); + ASSERT_THAT(model, NotNull()) << mjs_getError(parent); + EXPECT_EQ(model->ntendon, 2); + + mj_deleteModel(model); + mj_deleteSpec(parent); + mj_deleteSpec(child); +} + +TEST_F(MujocoTest, AttachSpatialTendonGitHubIssue3119) { + static constexpr char parent_xml[] = R"( + + + + + + + )"; + + static constexpr char child_xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + )"; + + std::array er; + mjSpec* parent_spec = + mj_parseXMLString(parent_xml, 0, er.data(), er.size()); + ASSERT_THAT(parent_spec, NotNull()) << er.data(); + mjSpec* child_spec = + mj_parseXMLString(child_xml, 0, er.data(), er.size()); + ASSERT_THAT(child_spec, NotNull()) << er.data(); + + mjsBody* parent_body = mjs_findBody(parent_spec, "parent_body"); + ASSERT_THAT(parent_body, NotNull()); + mjsSite* attach_site = mjs_addSite(parent_body, 0); + mjs_setName(attach_site->element, "attach_site"); + + mjs_attach(attach_site->element, + mjs_findBody(child_spec, "child_body")->element, + "", "_child"); + + EXPECT_THAT(mjs_findElement(parent_spec, mjOBJ_TENDON, + "tendon_with_sidesite_child"), NotNull()); + EXPECT_THAT(mjs_findElement(parent_spec, mjOBJ_TENDON, + "tendon_without_sidesite_child"), NotNull()); + + mjModel* model = mj_compile(parent_spec, nullptr); + ASSERT_THAT(model, NotNull()) << mjs_getError(parent_spec); + EXPECT_EQ(model->ntendon, 2); + + mj_deleteModel(model); + mj_deleteSpec(parent_spec); + mj_deleteSpec(child_spec); +} + TEST_F(MujocoTest, AttachDifferent) { std::array er; mjtNum tol = 0; From 82e92cbcaae55b381a34de58be84b5a3e8c18093 Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Mon, 23 Feb 2026 12:14:37 -0800 Subject: [PATCH 31/48] Remove PluginTest and make MujocoTest load plugins. MujocoTest now loads plugins from MUJOCO_PLUGIN_DIR if set. PluginTest is removed; all tests use MujocoTest directly. testspeed binary loads plugins from MUJOCO_PLUGIN_DIR. This is in preparation for moving common asset format parsing (obj, msh, stl, etc.) where we will always want to load those plugins. PiperOrigin-RevId: 874194428 Change-Id: Id90805a9ba5de4627911b56d8b9c4ab4e1b29310 --- sample/testspeed.cc | 6 +++++ test/engine/engine_plugin_test.cc | 6 ++--- test/engine/engine_ray_test.cc | 2 +- test/fixture.h | 32 +++++++++++++---------- test/plugin/actuator/pid_test.cc | 2 +- test/plugin/elasticity/elasticity_test.cc | 2 +- test/user/user_api_test.cc | 28 ++++++++++---------- test/xml/xml_native_writer_test.cc | 2 +- 8 files changed, 45 insertions(+), 35 deletions(-) diff --git a/sample/testspeed.cc b/sample/testspeed.cc index f14e6572..544edfc6 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -187,6 +187,12 @@ int main(int argc, char** argv) { nthread = mjMAX(1, mjMIN(maxthread, nthread)); npoolthread = mjMAX(1, mjMIN(maxthread, npoolthread)); + // load plugins from MUJOCO_PLUGIN_DIR if set + const char* plugin_dir = std::getenv("MUJOCO_PLUGIN_DIR"); + if (plugin_dir) { + mj_loadAllPluginLibraries(plugin_dir, nullptr); + } + // get filename, determine file type std::string filename(argv[1]); bool binary = (filename.find(".mjb") != std::string::npos); // NOLINT diff --git a/test/engine/engine_plugin_test.cc b/test/engine/engine_plugin_test.cc index 40ab13d4..180ab9c8 100644 --- a/test/engine/engine_plugin_test.cc +++ b/test/engine/engine_plugin_test.cc @@ -376,10 +376,10 @@ int RegisterNoAttributePlugin() { return mjp_registerPlugin(&plugin); } -class EnginePluginTest : public PluginTest { +class EnginePluginTest : public MujocoTest { public: // register all plugins - EnginePluginTest() : PluginTest() { + EnginePluginTest() : MujocoTest() { RegisterSensorPlugin(); for (int i = 1; i <= kNumFakePlugins; ++i) { @@ -466,7 +466,7 @@ TEST_F(MujocoTest, EmptyPluginDisallowed) { mj_deleteModel(m); } -TEST_F(PluginTest, FirstPartyPlugins) { +TEST_F(MujocoTest, FirstPartyPlugins) { EXPECT_THAT(mjp_pluginCount(), kNumTruePlugins); } diff --git a/test/engine/engine_ray_test.cc b/test/engine/engine_ray_test.cc index c55583f2..31433a33 100644 --- a/test/engine/engine_ray_test.cc +++ b/test/engine/engine_ray_test.cc @@ -82,7 +82,7 @@ using ::testing::DoubleNear; using ::testing::ElementsAre; using ::testing::NotNull; using ::testing::Pointwise; -using RayTest = PluginTest; +using RayTest = MujocoTest; TEST_F(RayTest, NoExclusions) { char error[1024]; diff --git a/test/fixture.h b/test/fixture.h index f9d0cb1f..64759b08 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -16,9 +16,12 @@ #define MUJOCO_TEST_FIXTURE_H_ #include +#include // IWYU pragma: keep +#include // IWYU pragma: keep #include #include #include +#include // IWYU pragma: keep #include #include #include @@ -54,6 +57,21 @@ class MujocoErrorTestGuard { // trigger a test failure. class MujocoTest : public ::testing::Test { public: + MujocoTest() { + static std::once_flag flag; + std::call_once(flag, []() { + const char* plugin_dir = std::getenv("MUJOCO_PLUGIN_DIR"); + if (plugin_dir) { + mj_loadAllPluginLibraries( + plugin_dir, +[](const char* filename, int first, int count) { + std::printf("Plugins registered by library '%s':\n", filename); + for (int i = first; i < first + count; ++i) { + std::printf(" %s\n", mjp_getPluginAtSlot(i)->name); + } + }); + } + }); + } ~MujocoTest() { mj_freeLastXML(); } private: @@ -183,20 +201,6 @@ class MockFilesystem { std::string dir_; // current directory }; -// Installs all plugins -class PluginTest : public MujocoTest { - public: - // load plugin library - PluginTest() : MujocoTest() { - mj_loadAllPluginLibraries( - std::string(std::getenv("MUJOCO_PLUGIN_DIR")).c_str(), +[](const char* filename, int first, int count) { - std::printf("Plugins registered by library '%s':\n", filename); - for (int i = first; i < first + count; ++i) { - std::printf(" %s\n", mjp_getPluginAtSlot(i)->name); - } - }); - } -}; } // namespace mujoco #endif // MUJOCO_TEST_FIXTURE_H_ diff --git a/test/plugin/actuator/pid_test.cc b/test/plugin/actuator/pid_test.cc index 8dfac5cf..e5fcb32c 100644 --- a/test/plugin/actuator/pid_test.cc +++ b/test/plugin/actuator/pid_test.cc @@ -29,7 +29,7 @@ namespace mujoco { namespace { -using PidTest = PluginTest; +using PidTest = MujocoTest; using ::testing::DoubleNear; using ::testing::HasSubstr; using ::testing::IsNull; diff --git a/test/plugin/elasticity/elasticity_test.cc b/test/plugin/elasticity/elasticity_test.cc index c7a06439..d2817787 100644 --- a/test/plugin/elasticity/elasticity_test.cc +++ b/test/plugin/elasticity/elasticity_test.cc @@ -26,7 +26,7 @@ namespace mujoco { namespace { -using ElasticityTest = PluginTest; +using ElasticityTest = MujocoTest; // -------------------------------- cable ----------------------------------- diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index e9f519e5..82cb2a08 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -171,7 +171,7 @@ TEST_F(MujocoTest, TreeTraversal) { mj_deleteSpec(spec); } -TEST_F(PluginTest, ActivatePlugin) { +TEST_F(MujocoTest, ActivatePlugin) { mjSpec* spec = mj_makeSpec(); mjs_activatePlugin(spec, "mujoco.elasticity.cable"); @@ -196,7 +196,7 @@ TEST_F(PluginTest, ActivatePlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, DeletePlugin) { +TEST_F(MujocoTest, DeletePlugin) { mjSpec* spec = mj_makeSpec(); mjs_activatePlugin(spec, "mujoco.pid"); @@ -267,7 +267,7 @@ static constexpr char xml_plugin_2[] = R"( )"; -TEST_F(PluginTest, AttachPlugin) { +TEST_F(MujocoTest, AttachPlugin) { std::array err; mjSpec* parent = mj_parseXMLString(xml_plugin_1, 0, err.data(), err.size()); ASSERT_THAT(parent, NotNull()) << err.data(); @@ -316,7 +316,7 @@ TEST_F(PluginTest, AttachPlugin) { mj_deleteSpec(spec_3); } -TEST_F(PluginTest, DetachPlugin) { +TEST_F(MujocoTest, DetachPlugin) { std::array err; mjSpec* parent = mj_parseXMLString(xml_plugin_1, 0, err.data(), err.size()); ASSERT_THAT(parent, NotNull()) << err.data(); @@ -342,7 +342,7 @@ TEST_F(PluginTest, DetachPlugin) { mj_deleteSpec(child); } -TEST_F(PluginTest, AttachExplicitPlugin) { +TEST_F(MujocoTest, AttachExplicitPlugin) { static constexpr char xml_parent[] = R"( @@ -396,7 +396,7 @@ TEST_F(PluginTest, AttachExplicitPlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, ReplicatePlugin) { +TEST_F(MujocoTest, ReplicatePlugin) { static constexpr char xml[] = R"( @@ -429,7 +429,7 @@ TEST_F(PluginTest, ReplicatePlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, ReplicateExplicitPlugin) { +TEST_F(MujocoTest, ReplicateExplicitPlugin) { static constexpr char xml[] = R"( @@ -484,7 +484,7 @@ TEST_F(MujocoTest, RecompileFails) { mj_deleteSpec(spec); } -TEST_F(PluginTest, ModifyShellInertiaFails) { +TEST_F(MujocoTest, ModifyShellInertiaFails) { static constexpr char xml[] = R"( @@ -515,7 +515,7 @@ TEST_F(PluginTest, ModifyShellInertiaFails) { } // ------------------- test recompilation multiple files ----------------------- -TEST_F(PluginTest, RecompileCompare) { +TEST_F(MujocoTest, RecompileCompare) { mjtNum tol = 0; std::string field = ""; @@ -606,7 +606,7 @@ TEST_F(PluginTest, RecompileCompare) { } } -TEST_F(PluginTest, RecompileEdit) { +TEST_F(MujocoTest, RecompileEdit) { static constexpr char xml[] = R"( @@ -640,7 +640,7 @@ TEST_F(PluginTest, RecompileEdit) { // ------------------- test cache with modified assets ------------------------- -TEST_F(PluginTest, RecompileCompareObjCache) { +TEST_F(MujocoTest, RecompileCompareObjCache) { static constexpr char xml[] = R"( @@ -716,7 +716,7 @@ static constexpr uint8_t tex2[] = { 0x82 }; -TEST_F(PluginTest, RecompileComparePngCache) { +TEST_F(MujocoTest, RecompileComparePngCache) { static constexpr char xml[] = R"( @@ -753,7 +753,7 @@ TEST_F(PluginTest, RecompileComparePngCache) { mj_deleteVFS(vfs.get()); } -TEST_F(PluginTest, DisableCache) { +TEST_F(MujocoTest, DisableCache) { static constexpr char xml[] = R"( @@ -791,7 +791,7 @@ TEST_F(PluginTest, DisableCache) { // -------------------------------- test textures ------------------------------ -TEST_F(PluginTest, TextureFromBuffer) { +TEST_F(MujocoTest, TextureFromBuffer) { mjSpec* spec = mj_makeSpec(); mjsTexture* t1 = mjs_addTexture(spec); diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index 71c250e6..20fd7a14 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -44,7 +44,7 @@ using ::testing::HasSubstr; using ::testing::Not; using ::testing::NotNull; -using XMLWriterTest = PluginTest; +using XMLWriterTest = MujocoTest; static const char* const kNonRgbTextureXMLPath = "xml/testdata/hfield_png_nonrgb.xml"; From c791f950cbb480576f78126c78f3cc353dd2f0c6 Mon Sep 17 00:00:00 2001 From: MatiasManevi Date: Mon, 23 Feb 2026 18:41:47 -0300 Subject: [PATCH 32/48] Enable COOP/COEP headers for wasm Vite dev servers --- wasm/demo_app/vite.demo.config.ts | 4 ++++ wasm/tests/sandbox/vite.sandbox.config.ts | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/wasm/demo_app/vite.demo.config.ts b/wasm/demo_app/vite.demo.config.ts index b84e280f..a5a0217a 100644 --- a/wasm/demo_app/vite.demo.config.ts +++ b/wasm/demo_app/vite.demo.config.ts @@ -23,5 +23,9 @@ export default defineConfig({ }, server: { open: true, + headers: { + "Cross-Origin-Opener-Policy": "same-origin", + "Cross-Origin-Embedder-Policy": "require-corp", + }, }, }) diff --git a/wasm/tests/sandbox/vite.sandbox.config.ts b/wasm/tests/sandbox/vite.sandbox.config.ts index 123f412b..01312ba1 100644 --- a/wasm/tests/sandbox/vite.sandbox.config.ts +++ b/wasm/tests/sandbox/vite.sandbox.config.ts @@ -23,5 +23,9 @@ export default defineConfig({ }, server: { open: true, + headers: { + "Cross-Origin-Opener-Policy": "same-origin", + "Cross-Origin-Embedder-Policy": "require-corp", + }, }, }) From 5fc580a8ba5194cf502810facf73a4a4481dcdd9 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 24 Feb 2026 06:03:24 -0800 Subject: [PATCH 33/48] Refactor SDF coefficient computation into mjCOctree. The logic for computing Signed Distance Function (SDF) coefficients using TriangleMeshDistance is moved from `mjCPlugin::Make` in `user_mesh.cc` to a new method `mjCOctree::ComputeSdfCoeffs` in `user_objects.cc`. This encapsulates the SDF computation within the octree class. PiperOrigin-RevId: 874572414 Change-Id: Ia2365c9c4527a252ff1d1bb43fe2c2593761e471 --- src/user/user_mesh.cc | 48 +---------------------------------- src/user/user_objects.cc | 54 ++++++++++++++++++++++++++++++++++++++++ src/user/user_objects.h | 4 +++ 3 files changed, 59 insertions(+), 47 deletions(-) diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 5d060117..70d78bed 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -34,7 +34,6 @@ #include #include "user/user_api.h" -#include #ifdef MUJOCO_TINYOBJLOADER_IMPL #define TINYOBJLOADER_IMPLEMENTATION @@ -780,52 +779,7 @@ void mjCMesh::TryCompile(const mjVFS* vfs) { // compute sdf coefficients if (!plugin.active) { - tmd::TriangleMeshDistance sdf(vert_.data(), nvert(), face_.data(), nface()); - - std::vector coeffs(octree_.NumVerts()); - std::vector processed(octree_.NumVerts(), false); - std::deque queue; - - if (octree_.NumNodes() > 0) { - queue.push_back(0); // start traversal from the root node - } - - while (!queue.empty()) { - int node_idx = queue.front(); - queue.pop_front(); - - for (int j = 0; j < 8; ++j) { - int vert_id = octree_.VertId(node_idx, j); - if (processed[vert_id]) { - continue; - } - if (octree_.Hang(vert_id).empty()) { - coeffs[vert_id] = sdf.signed_distance(octree_.Vert(vert_id)).distance; - } else { - double sum_coeff = 0; - for (int dep_id : octree_.Hang(vert_id)) { - sum_coeff += coeffs[dep_id]; - if (!processed[dep_id]) { - throw mjCError(this, "sdf coefficient computation failed"); - } - } - coeffs[vert_id] = sum_coeff / octree_.Hang(vert_id).size(); - } - processed[vert_id] = true; - } - - for (int child_idx : octree_.Children(node_idx)) { - if (child_idx != -1) { - queue.push_back(child_idx); - } - } - } - - for (int i = 0; i < octree_.NumNodes(); ++i) { - for (int j = 0; j < 8; j++) { - octree_.AddCoeff(i, j, coeffs[octree_.VertId(i, j)]); - } - } + octree_.ComputeSdfCoeffs(vert_.data(), nvert(), face_.data(), nface()); } } diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 4b5ee2a4..06045081 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -47,6 +47,7 @@ #include "user/user_model.h" #include "user/user_resource.h" #include "user/user_util.h" +#include namespace { namespace mju = ::mujoco::util; @@ -621,6 +622,59 @@ void mjCOctree::CreateOctree(const double aamm[6]) { } +// compute SDF coefficients at octree vertices using triangle mesh distance +void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, + const int* face, int nface) { + tmd::TriangleMeshDistance sdf(vert, static_cast(nvert), + face, static_cast(nface)); + + std::vector coeffs(NumVerts()); + std::vector processed(NumVerts(), false); + std::deque queue; + + if (NumNodes() > 0) { + queue.push_back(0); + } + + while (!queue.empty()) { + int node_idx = queue.front(); + queue.pop_front(); + + // compute SDF coefficients at the 8 vertices of the octree node + for (int j = 0; j < 8; ++j) { + int vert_id = VertId(node_idx, j); + if (processed[vert_id]) { + continue; + } + if (Hang(vert_id).empty()) { + coeffs[vert_id] = sdf.signed_distance(Vert(vert_id)).distance; + } else { + double sum_coeff = 0; + for (int dep_id : Hang(vert_id)) { + sum_coeff += coeffs[dep_id]; + } + coeffs[vert_id] = sum_coeff / Hang(vert_id).size(); + } + processed[vert_id] = true; + } + + // add children to the queue + for (int child_idx : Children(node_idx)) { + if (child_idx != -1) { + queue.push_back(child_idx); + } + } + } + + // copy coefficients to the octree nodes + for (int i = 0; i < NumNodes(); ++i) { + for (int j = 0; j < 8; j++) { + AddCoeff(i, j, coeffs[VertId(i, j)]); + } + } +} + + static double dot2(const double* a, const double* b) { return a[0] * b[0] + a[1] * b[1]; } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 081791d7..85a1d5e1 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -301,6 +301,10 @@ class mjCOctree : public mjCOctree_ { face_.clear(); } void AddCoeff(int n, int v, double coeff) { node_[n].coeff[v] = coeff; } + double Coeff(int n, int v) const { return node_[n].coeff[v]; } + + // compute SDF coefficients at octree vertices using triangle mesh distance + void ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface); private: void Make(std::vector& elements); From 1d6ff2cecec1290f0610818482476113ee3787ae Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Tue, 24 Feb 2026 07:57:01 -0800 Subject: [PATCH 34/48] Display Bodies as a tree in Spec Explorer. Display all other elements under "Elements" or "Assets" group. Show properties for mjSpec, mjModel, and mjData for selected elements. PiperOrigin-RevId: 874614195 Change-Id: I128c2bec844983ac5d35a3a0b1e58b8a5b0109cb --- src/experimental/platform/CMakeLists.txt | 2 + src/experimental/platform/gui.cc | 109 +-- src/experimental/platform/gui.h | 8 - src/experimental/platform/gui_spec.cc | 847 +++++++++++++++++++++ src/experimental/platform/gui_spec.h | 37 + src/experimental/platform/imgui_widgets.cc | 226 ++++++ src/experimental/platform/imgui_widgets.h | 80 ++ src/experimental/studio/app.cc | 194 +++-- src/experimental/studio/app.h | 13 +- 9 files changed, 1290 insertions(+), 226 deletions(-) create mode 100644 src/experimental/platform/gui_spec.cc create mode 100644 src/experimental/platform/gui_spec.h diff --git a/src/experimental/platform/CMakeLists.txt b/src/experimental/platform/CMakeLists.txt index d4a58d23..8ba7aa6f 100644 --- a/src/experimental/platform/CMakeLists.txt +++ b/src/experimental/platform/CMakeLists.txt @@ -42,6 +42,8 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} file_dialog.h gui.cc gui.h + gui_spec.cc + gui_spec.h helpers.cc helpers.h imgui_widgets.cc diff --git a/src/experimental/platform/gui.cc b/src/experimental/platform/gui.cc index 5cfc9f2b..2d29eeca 100644 --- a/src/experimental/platform/gui.cc +++ b/src/experimental/platform/gui.cc @@ -470,7 +470,7 @@ void StateGui(const mjModel* model, mjData* data, std::vector& state, ImVec2(0, ImGui::GetTextLineHeightWithSpacing() * 20))) { ImGui::TableSetupColumn("Index"); ImGui::TableSetupColumn("Name"); - ImGui::TableSetupColumn("Value"); + ImGui::TableSetupColumn("Value", ImGuiTableColumnFlags_WidthStretch); ImGui::TableSetupScrollFreeze(0, 1); ImGui::TableHeadersRow(); @@ -1107,111 +1107,4 @@ void StatsGui(const mjModel* model, const mjData* data, bool paused, ImGui::Columns(); } -void BodyPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsBody* body = mjs_asBody(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(body->element); - if (name.empty()) { - name = "(Body " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - - ImGui::Text("Name"); - ImGui::Text("xpos[0]"); - ImGui::Text("xpos[1]"); - ImGui::Text("xpos[2]"); - ImGui::Text("xquat[0]"); - ImGui::Text("xquat[1]"); - ImGui::Text("xquat[2]"); - ImGui::Text("xquat[3]"); - ImGui::Text("mass"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); - ImGui::Text("%f", data->xpos[3*id+0]); - ImGui::Text("%f", data->xpos[3*id+1]); - ImGui::Text("%f", data->xpos[3*id+2]); - ImGui::Text("%f", data->xquat[4*id+0]); - ImGui::Text("%f", data->xquat[4*id+1]); - ImGui::Text("%f", data->xquat[4*id+2]); - ImGui::Text("%f", data->xquat[4*id+3]); - ImGui::Text("%f", model->body_mass[id]); -} - -void JointPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsJoint* joint = mjs_asJoint(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(joint->element); - if (name.empty()) { - name = "(Joint " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - ImGui::Text("Name"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); -} - -void SitePropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsSite* site = mjs_asSite(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(site->element); - if (name.empty()) { - name = "(Joint " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - ImGui::Text("Name"); - ImGui::Text("site_xpos[0]"); - ImGui::Text("site_xpos[1]"); - ImGui::Text("site_xpos[2]"); - ImGui::Text("site_xmat[0]"); - ImGui::Text("site_xmat[1]"); - ImGui::Text("site_xmat[2]"); - ImGui::Text("site_xmat[3]"); - ImGui::Text("site_xmat[4]"); - ImGui::Text("site_xmat[5]"); - ImGui::Text("site_xmat[6]"); - ImGui::Text("site_xmat[7]"); - ImGui::Text("site_xmat[8]"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); - ImGui::Text("%f", data->site_xpos[3*id+0]); - ImGui::Text("%f", data->site_xpos[3*id+1]); - ImGui::Text("%f", data->site_xpos[3*id+2]); - ImGui::Text("%f", data->site_xmat[4*id+0]); - ImGui::Text("%f", data->site_xmat[4*id+1]); - ImGui::Text("%f", data->site_xmat[4*id+2]); - ImGui::Text("%f", data->site_xmat[4*id+3]); - ImGui::Text("%f", data->site_xmat[4*id+4]); - ImGui::Text("%f", data->site_xmat[4*id+5]); - ImGui::Text("%f", data->site_xmat[4*id+6]); - ImGui::Text("%f", data->site_xmat[4*id+7]); - ImGui::Text("%f", data->site_xmat[4*id+8]); -} - } // namespace mujoco::platform diff --git a/src/experimental/platform/gui.h b/src/experimental/platform/gui.h index 2341023e..7dbdae09 100644 --- a/src/experimental/platform/gui.h +++ b/src/experimental/platform/gui.h @@ -116,14 +116,6 @@ void CountsGui(const mjModel* model, mjData* data); // FPS needs to be tracked by the caller and passed here to be displayed. void StatsGui(const mjModel* model, const mjData* data, bool paused, float fps); -// UX for displaying properties of various mjSpec elements. -void BodyPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); -void JointPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); -void SitePropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); - } // namespace mujoco::platform #endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_H_ diff --git a/src/experimental/platform/gui_spec.cc b/src/experimental/platform/gui_spec.cc new file mode 100644 index 00000000..eeb74581 --- /dev/null +++ b/src/experimental/platform/gui_spec.cc @@ -0,0 +1,847 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/platform/gui_spec.h" + +#include +#include + +#include +#include +#include +#include "experimental/platform/imgui_widgets.h" + +// Define the mujoco X macros to add fields to the ImGui_DataTable. +// We limit the fields to the ones with a matching element by comparing the +// array size field (e.g. nbody) with the MATCH constexpr value. +#define X(TYPE, NAME, NELEM, SIZE) \ + if constexpr (#NELEM == MATCH) table(#NAME, ptr->NAME, SIZE); + +// Simple wrapper around MJMODEL_POINTERS that prepares values we need for the +// X macro above. +#define MJMODEL_POINTERS_X(M) \ + constexpr std::string_view MATCH(#M); \ + const auto* ptr = model; \ + MJMODEL_POINTERS + +// Simple wrapper around MJDATA_POINTERS that prepares values we need for the +// X macro above. +#define MJDATA_POINTERS_X(M) \ + constexpr std::string_view MATCH(#M); \ + const auto* ptr = data; \ + MJDATA_POINTERS + +namespace mujoco::platform { + +// Returns the index of the element in the spec. This is different from +// mjs_getId which returns the runtime ID of an element. +static int GetElementIndexInSpec(mjsElement* element) { + int n = 0; + mjSpec* spec = mjs_getSpec(element); + mjsElement* iter = mjs_firstElement(spec, element->elemtype); + while (iter) { + if (iter == element) { + return n; + } + iter = mjs_nextElement(spec, iter); + ++n; + } + return -1; +} + +// Returns a name for the element; either the element has a name, or we +// construct a unique name from the element's id (using mjs_getId) or index +// (using GetElementIndexInSpec). +static std::string ElementName(mjsElement* element) { + const mjString* name = mjs_getName(element); + std::string label = *name; + if (label.empty()) { + int id = mjs_getId(element); + if (id == -1) { + id = GetElementIndexInSpec(element); + } + const char* type_name = mju_type2Str(element->elemtype); + label = "(" + std::string(type_name) + " " + std::to_string(id) + ")"; + } + return label; +} + +static void QuatOrOrientation(ImGui_DataTable& table, const double quat[4], + const mjsOrientation& orientation, + const char* quat_name, const char* alt_name) { + auto alt = + [&](const char* label) { return std::string(alt_name) + "." + label; }; + + switch (orientation.type) { + case mjORIENTATION_QUAT: + table(quat_name, quat, 4); + break; + case mjORIENTATION_AXISANGLE: + table(alt("axisangle").c_str(), orientation.axisangle, 4); + break; + case mjORIENTATION_XYAXES: + table(alt("xyaxes").c_str(), orientation.xyaxes, 6); + break; + case mjORIENTATION_ZAXIS: + table(alt("zaxis").c_str(), orientation.zaxis, 3); + break; + case mjORIENTATION_EULER: + table(alt("euler").c_str(), orientation.euler, 3); + break; + } +} + +static void AddDeleteButton(mjsElement* element, + const SpecElementCallbackFn& on_delete) { + if (on_delete) { + // Right-align the delete button. + const float button_width = ImGui::CalcTextSize(ICON_FA_TRASH_CAN).x + + ImGui::GetStyle().FramePadding.x * 2.0f; + ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); + if (ImGui::SmallButton(ICON_FA_TRASH_CAN)) { + on_delete(element); + } + } +} + +static void SelectableElement(mjsElement* element, + mjsElement** selected_element, + const SpecElementCallbackFn& on_delete) { + constexpr ImGuiSelectableFlags flags = ImGuiSelectableFlags_AllowOverlap; + + const std::string name = ElementName(element); + const bool selected = (element == *selected_element); + if (ImGui::Selectable(name.c_str(), selected, flags)) { + *selected_element = element; + } + if (selected) { + AddDeleteButton(element, on_delete); + } +} + +static void BodyChildrenGui(const char* heading, mjtObj type, + mjsElement** element, mjsBody* body, + const SpecElementCallbackFn& on_delete) { + mjsElement* iter = mjs_firstChild(body, type, 0); + if (!iter) { + return; + } + + constexpr ImGuiTreeNodeFlags tree_flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_DrawLinesFull; + if (ImGui::TreeNodeEx(heading, tree_flags)) { + while (iter) { + SelectableElement(iter, element, on_delete); + iter = mjs_nextChild(body, iter, 0); + } + ImGui::TreePop(); + } +} + +static void ElementListGui(const char* heading, mjtObj type, + mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete) { + mjsElement* iter = mjs_firstElement(spec, type); + if (!iter) { + return; + } + + constexpr ImGuiTreeNodeFlags tree_flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + if (ImGui::TreeNodeEx(heading, tree_flags)) { + while (iter) { + SelectableElement(iter, element, on_delete); + iter = mjs_nextElement(spec, iter); + } + ImGui::TreePop(); + } +} + +static void BodyTreeGuiRecursive(mjsElement** element, mjsBody* body, + const SpecElementCallbackFn& on_delete) { + const std::string label = ElementName(body->element); + + ImGui::PushID(body); + + ImGuiTreeNodeFlags flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed | + ImGuiTreeNodeFlags_DrawLinesFull | ImGuiTreeNodeFlags_AllowOverlap; + + if (*element == body->element) { + flags |= ImGuiTreeNodeFlags_Selected; + } + + const bool tree_open = ImGui::TreeNodeEx(label.c_str(), flags); + if (ImGui::IsItemClicked()) { + *element = body->element; + } + if (*element == body->element) { + AddDeleteButton(body->element, on_delete); + } + + if (tree_open) { + mjsElement* iter = mjs_firstChild(body, mjOBJ_BODY, 0); + while (iter) { + BodyTreeGuiRecursive(element, mjs_asBody(iter), on_delete); + iter = mjs_nextChild(body, iter, 0); + } + + BodyChildrenGui("Frames", mjOBJ_FRAME, element, body, on_delete); + BodyChildrenGui("Sites", mjOBJ_SITE, element, body, on_delete); + BodyChildrenGui("Joints", mjOBJ_JOINT, element, body, on_delete); + BodyChildrenGui("Geoms", mjOBJ_GEOM, element, body, on_delete); + BodyChildrenGui("Lights", mjOBJ_LIGHT, element, body, on_delete); + BodyChildrenGui("Cameras", mjOBJ_CAMERA, element, body, on_delete); + + ImGui::TreePop(); + } + + ImGui::PopID(); +} + +void SpecExplorerGui(mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete) { + const ImGuiTreeNodeFlags flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + + if (ImGui::TreeNodeEx("Body Tree", flags)) { + mjsElement* root = mjs_firstElement(spec, mjOBJ_BODY); + if (root) { + mjsBody* body = mjs_asBody(root); + if (body) { + BodyTreeGuiRecursive(element, body, on_delete); + } + } + ImGui::TreePop(); + } + + auto list = [&](const char* heading, mjtObj type) { + ElementListGui(heading, type, element, spec, on_delete); + }; + + ImGui::PushID(spec); + + // Non-tree elements. + if (ImGui::TreeNodeEx("Elements", flags)) { + list("Actuators", mjOBJ_ACTUATOR); + list("Sensors", mjOBJ_SENSOR); + list("Flexes", mjOBJ_FLEX); + list("Tendons", mjOBJ_TENDON); + list("Pair", mjOBJ_PAIR); + list("Exclude", mjOBJ_EXCLUDE); + list("Equality", mjOBJ_EQUALITY); + list("Numeric", mjOBJ_NUMERIC); + list("Text", mjOBJ_TEXT); + list("Tuple", mjOBJ_TUPLE); + list("Key", mjOBJ_KEY); + list("Default", mjOBJ_DEFAULT); + ImGui::TreePop(); + } + + // Assets. + if (ImGui::TreeNodeEx("Assets", flags)) { + list("Meshes", mjOBJ_MESH); + list("Height Fields", mjOBJ_HFIELD); + list("Skins", mjOBJ_SKIN); + list("Textures", mjOBJ_TEXTURE); + list("Materials", mjOBJ_MATERIAL); + ImGui::TreePop(); + } + + ImGui::PopID(); +} + +void ElementSpecGui(const mjSpec* spec, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + + switch (element->elemtype) { + case mjOBJ_BODY: { + const mjsBody* body = mjs_asBody(element); + table("childclass", body->childclass, 1); // childclass name + table("pos", body->pos, 3); // frame position + QuatOrOrientation(table, body->quat, body->alt, "quat", "alt"); // frame orientation + table("ipos", body->ipos, 3); // inertial frame position + QuatOrOrientation(table, body->iquat, body->ialt, "iquat", "ialt"); // inertial frame orientation + table("mass", body->mass, 1); // mass + table("inertia", body->inertia, 3); // diagonal inertia (in i-frame) + table("fullinertia", body->fullinertia, 6); // non-axis-aligned inertia matrix + table("mocap", body->mocap, 1); // is this a mocap body + table("gravcomp", body->gravcomp, 1); // gravity compensation + table("explicitinertial", body->explicitinertial, 1); // whether to save the body with explicit inertial clause + table("sleep", body->sleep, 1); // sleep policy + table("info", body->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_JOINT: { + mjsJoint* joint = mjs_asJoint(element); + table("pos", joint->pos, 3); // anchor position + table("axis", joint->axis, 3); // joint axis + table("ref", joint->ref, 1); // value at reference configuration: qpos0 + table("align", joint->align, 1); // align free joint with body com (mjtAlignFree) + table("stiffness", joint->stiffness, 1); // stiffness coefficient + table("springref", joint->springref, 1); // spring reference value: qpos_spring + table("springdamper", joint->springdamper, 2); // timeconst, dampratio + table("limited", joint->limited, 1); // does joint have limits (mjtLimited) + table("range", joint->range, 2); // joint limits + table("margin", joint->margin, 1); // margin value for joint limit detection + table("solref_limit", joint->solref_limit, mjNREF); // solver reference: joint limits + table("solimp_limit", joint->solimp_limit, mjNIMP); // solver impedance: joint limits + table("actfrclimited", joint->actfrclimited, 1); // are actuator forces on joint limited (mjtLimited) + table("actfrcrange", joint->actfrcrange, 2); // actuator force limits + table("armature", joint->armature, 1); // armature inertia (mass for slider) + table("damping", joint->damping, 1); // damping coefficient + table("frictionloss", joint->frictionloss, 1); // friction loss + table("solref_friction", joint->solref_friction, mjNREF); // solver reference: dof friction + table("solimp_friction", joint->solimp_friction, mjNIMP); // solver impedance: dof friction + table("group", joint->group, 1); // group + table("actgravcomp", joint->actgravcomp, 1); // is gravcomp force applied via actuators + table("info", joint->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_ACTUATOR: { + mjsActuator* actuator = mjs_asActuator(element); + table("gaintype", actuator->gaintype, 1); // gain type + table("gainprm", actuator->gainprm, mjNGAIN); // gain parameters + table("biastype", actuator->biastype, 1); // bias type + table("biasprm", actuator->biasprm, mjNGAIN); // bias parameters + table("dyntype", actuator->dyntype, 1); // dynamics type + table("dynprm", actuator->dynprm, mjNDYN); // dynamics parameters + table("actdim", actuator->actdim, 1); // number of activation variables + table("actearly", actuator->actearly, 1); // apply next activations to qfrc + table("trntype", actuator->trntype, 1); // transmission type + table("gear", actuator->gear, 6); // length and transmitted force scaling + table("target", actuator->target, 1); // name of transmission target + table("refsite", actuator->refsite, 1); // reference site, for site transmission + table("slidersite", actuator->slidersite, 1); // site defining cylinder, for slider-crank + table("cranklength", actuator->cranklength, 1); // crank length, for slider-crank + table("lengthrange", actuator->lengthrange, 2); // transmission length range + table("inheritrange", actuator->inheritrange, 1); // automatic range setting for position and intvelocity + table("ctrllimited", actuator->ctrllimited, 1); // are control limits defined (mjtLimited) + table("ctrlrange", actuator->ctrlrange, 2); // control range + table("forcelimited", actuator->forcelimited, 1); // are force limits defined (mjtLimited) + table("forcerange", actuator->forcerange, 2); // force range + table("actlimited", actuator->actlimited, 1); // are activation limits defined (mjtLimited) + table("actrange", actuator->actrange, 2); // activation range + table("group", actuator->group, 1); // group + table("nsample", actuator->nsample, 1); // number of samples in history buffer + table("interp", actuator->interp, 1); // interpolation order (0=ZOH, 1=linear, 2=cubic) + table("delay", actuator->delay, 1); // delay time in seconds; 0: no delay + table("info", actuator->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SENSOR: { + mjsSensor* sensor = mjs_asSensor(element); + table("type", sensor->type, 1); // type of sensor + table("objtype", sensor->objtype, 1); // type of sensorized object + table("objname", sensor->objname, 1); // name of sensorized object + table("reftype", sensor->reftype, 1); // type of referenced object + table("refname", sensor->refname, 1); // name of referenced object + table("intprm", sensor->intprm, mjNSENS); // integer parameters + table("datatype", sensor->datatype, 1); // data type for sensor measurement + table("needstage", sensor->needstage, 1); // compute stage needed to simulate sensor + table("dim", sensor->dim, 1); // number of scalar outputs + table("cutoff", sensor->cutoff, 1); // cutoff for real and positive datatypes + table("noise", sensor->noise, 1); // noise stdev + table("nsample", sensor->nsample, 1); // number of samples in history buffer + table("interp", sensor->interp, 1); // interpolation order (0=ZOH, 1=linear, 2=cubic) + table("delay", sensor->delay, 1); // delay time in seconds + table("interval", sensor->interval, 2); // [period, time_prev] in seconds + table("info", sensor->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SITE: { + mjsSite* site = mjs_asSite(element); + table("pos", site->pos, 3); // position + QuatOrOrientation(table, site->quat, site->alt, "quat", "alt"); // orientation + table("fromto", site->fromto, 6); // alternative for capsule, cylinder, box, ellipsoid + table("size", site->size, 3); // geom size + table("type", site->type, 1); // geom type + table("material", site->material, 1); // name of material + table("group", site->group, 1); // group + table("rgba", site->rgba, 4); // rgba when material is omitted + table("info", site->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_FRAME: { + mjsFrame* frame = mjs_asFrame(element); + table("childclass", frame->childclass, 1); // childclass name + table("pos", frame->pos, 3); // position + QuatOrOrientation(table, frame->quat, frame->alt, "quat", "alt"); // orientation + table("info", frame->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_GEOM: { + mjsGeom* geom = mjs_asGeom(element); + table("type", geom->type, 1); // geom type + table("pos", geom->pos, 3); // position + QuatOrOrientation(table, geom->quat, geom->alt, "quat", "alt"); // orientation + table("fromto", geom->fromto, 6); // alternative for capsule, cylinder, box, ellipsoid + table("size", geom->size, 3); // type-specific size + table("contype", geom->contype, 1); // contact type + table("conaffinity", geom->conaffinity, 1); // contact affinity + table("condim", geom->condim, 1); // contact dimensionality + table("priority", geom->priority, 1); // contact priority + table("friction", geom->friction, 3); // one-sided friction coefficients: slide, roll, spin + table("solmix", geom->solmix, 1); // solver mixing for contact pairs + table("solref", geom->solref, mjNREF); // solver reference + table("solimp", geom->solimp, mjNIMP); // solver impedance + table("margin", geom->margin, 1); // margin for contact detection + table("gap", geom->gap, 1); // include in solver if dist < margin-gap + table("mass", geom->mass, 1); // used to compute density + table("density", geom->density, 1); // used to compute mass and inertia from volume or surface + table("typeinertia", geom->typeinertia, 1); // selects between surface and volume inertia + table("fluid_ellipsoid", geom->fluid_ellipsoid, 1); // whether ellipsoid-fluid model is active + table("fluid_coefs", geom->fluid_coefs, 5); // ellipsoid-fluid interaction coefs + table("material", geom->material, 1); // name of material + table("rgba", geom->rgba, 4); // rgba when material is omitted + table("group", geom->group, 1); // group + table("hfieldname", geom->hfieldname, 1); // heightfield attached to geom + table("meshname", geom->meshname, 1); // mesh attached to geom + table("fitscale", geom->fitscale, 1); // scale mesh uniformly + table("info", geom->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_LIGHT: { + mjsLight* light = mjs_asLight(element); + table("pos", light->pos, 3); // position + table("dir", light->dir, 3); // direction + table("mode", light->mode, 1); // tracking mode + table("targetbody", light->targetbody, 1); // target body for targeting + table("active", light->active, 1); // is light active + table("type", light->type, 1); // type of light + table("texture", light->texture, 1); // texture name for image lights + table("castshadow", light->castshadow, 1); // does light cast shadows + table("bulbradius", light->bulbradius, 1); // bulb radius, for soft shadows + table("intensity", light->intensity, 1); // intensity, in candelas + table("range", light->range, 1); // range of effectiveness + table("attenuation", light->attenuation, 3); // OpenGL attenuation (quadratic model) + table("cutoff", light->cutoff, 1); // OpenGL cutoff + table("exponent", light->exponent, 1); // OpenGL exponent + table("ambient", light->ambient, 3); // ambient color + table("diffuse", light->diffuse, 3); // diffuse color + table("specular", light->specular, 3); // specular color + table("info", light->info, 1); // message appended to compiler errorsx + break; + } + case mjOBJ_CAMERA: { + mjsCamera* camera = mjs_asCamera(element); + table("pos", camera->pos, 3); // position + QuatOrOrientation(table, camera->quat, camera->alt, "quat", "alt"); // orientation + table("mode", camera->mode, 1); // tracking mode + table("targetbody", camera->targetbody, 1); // target body for tracking/targeting + table("proj", camera->proj, 1); // camera projection type + table("resolution", camera->resolution, 2); // resolution (pixel) + table("output", camera->output, 1); // bit flags for output type + table("fovy", camera->fovy, 1); // y-field of view + table("ipd", camera->ipd, 1); // inter-pupillary distance + table("intrinsic", camera->intrinsic, 4); // camera intrinsics (length) + table("sensor_size", camera->sensor_size, 2); // sensor size (length) + table("focal_length", camera->focal_length, 2); // focal length (length) + table("focal_pixel", camera->focal_pixel, 2); // focal length (pixel) + table("principal_length", camera->principal_length, 2); // principal point (length) + table("principal_pixel", camera->principal_pixel, 2); // principal point (pixel) + table("info", camera->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_MESH: { + mjsMesh* mesh = mjs_asMesh(element); + table("content_type", mesh->content_type, 1); // content type of file + table("file", mesh->file, 1); // mesh file + table("refpos", mesh->refpos, 3); // reference position + table("refquat", mesh->refquat, 4); // reference orientation + table("scale", mesh->scale, 3); // rescale mesh + table("inertia", mesh->inertia, 1); // inertia type (convex, legacy, exact, shell) + table("smoothnormal", mesh->smoothnormal, 1); // do not exclude large-angle faces from normals + table("needsdf", mesh->needsdf, 1); // compute sdf from mesh + table("maxhullvert", mesh->maxhullvert, 1); // maximum vertex count for the convex hull + table("material", mesh->material, 1); // name of material + table("info", mesh->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_HFIELD: { + mjsHField* hfield = mjs_asHField(element); + table("content_type", hfield->content_type, 1); // content type of file + table("file", hfield->file, 1); // file: (nrow, ncol, [elevation data]) + table("size", hfield->size, 4); // hfield size (ignore referencing geom size) + table("nrow", hfield->nrow, 1); // number of rows + table("ncol", hfield->ncol, 1); // number of columns + table("info", hfield->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SKIN: { + mjsSkin* skin = mjs_asSkin(element); + table("file", skin->file, 1); // skin file + table("material", skin->material, 1); // name of material used for rendering + table("rgba", skin->rgba, 4); // rgba when material is omitted + table("inflate", skin->inflate, 1); // inflate in normal direction + table("group", skin->group, 1); // group for visualization + table("info", skin->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_FLEX: { + mjsFlex* flex = mjs_asFlex(element); + table("contype", flex->contype, 1); // contact type + table("conaffinity", flex->conaffinity, 1); // contact affinity + table("condim", flex->condim, 1); // contact dimensionality + table("priority", flex->priority, 1); // contact priority + table("friction", flex->friction, 3); // one-sided friction coefficients: slide, roll, spin + table("solmix", flex->solmix, 1); // solver mixing for contact pairs + table("solref", flex->solref, mjNREF); // solver reference + table("solimp", flex->solimp, mjNIMP); // solver impedance + table("margin", flex->margin, 1); // margin for contact detection + table("gap", flex->gap, 1); // include in solver if distdim, 1); // element dimensionality + table("radius", flex->radius, 1); // radius around primitive element + table("size", flex->size, 3); // vertex bounding box half sizes in qpos0 + table("internal", flex->internal, 1); // enable internal collisions + table("flatskin", flex->flatskin, 1); // render flex skin with flat shading + table("selfcollide", flex->selfcollide, 1); // mode for flex self collision + table("vertcollide", flex->vertcollide, 1); // mode for vertex collision + table("passive", flex->passive, 1); // mode for passive collisions + table("activelayers", flex->activelayers, 1); // number of active element layers in 3D + table("group", flex->group, 1); // group for visualization + table("edgestiffness", flex->edgestiffness, 1); // edge stiffness + table("edgedamping", flex->edgedamping, 1); // edge damping + table("rgba", flex->rgba, 4); // rgba when material is omitted + table("material", flex->material, 1); // name of material used for rendering + table("young", flex->young, 1); // Young's modulus + table("poisson", flex->poisson, 1); // Poisson's ratio + table("damping", flex->damping, 1); // Rayleigh's damping + table("thickness", flex->thickness, 1); // thickness (2D only) + table("elastic2d", flex->elastic2d, 1); // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both + table("info", flex->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_TENDON: { + mjsTendon* tendon = mjs_asTendon(element); + table("stiffness", tendon->stiffness, 1); // stiffness coefficient + table("springlength", tendon->springlength, 2); // spring resting length; {-1, -1}: use qpos_spring + table("damping", tendon->damping, 1); // damping coefficient + table("frictionloss", tendon->frictionloss, 1); // friction loss + table("solref_friction", tendon->solref_friction, mjNREF); // solver reference: tendon friction + table("solimp_friction", tendon->solimp_friction, mjNIMP); // solver impedance: tendon friction + table("armature", tendon->armature, 1); // inertia associated with tendon velocity + table("limited", tendon->limited, 1); // does tendon have limits (mjtLimited) + table("actfrclimited", tendon->actfrclimited, 1); // does tendon have actuator force limits + table("range", tendon->range, 2); // length limits + table("actfrcrange", tendon->actfrcrange, 2); // actuator force limits + table("margin", tendon->margin, 1); // margin value for tendon limit detection + table("solref_limit", tendon->solref_limit, mjNREF); // solver reference: tendon limits + table("solimp_limit", tendon->solimp_limit, mjNIMP); // solver impedance: tendon limits + table("material", tendon->material, 1); // name of material for rendering + table("width", tendon->width, 1); // width for rendering + table("rgba", tendon->rgba, 4); // rgba when material is omitted + table("group", tendon->group, 1); // group + table("info", tendon->info, 1); // message appended to errors + break; + } + case mjOBJ_TEXTURE: { + mjsTexture* texture = mjs_asTexture(element); + table("type", texture->type, 1); // texture type + table("colorspace", texture->colorspace, 1); // colorspace + table("builtin", texture->builtin, 1); // builtin type (mjtBuiltin) + table("mark", texture->mark, 1); // mark type (mjtMark) + table("rgb1", texture->rgb1, 3); // first color for builtin + table("rgb2", texture->rgb2, 3); // second color for builtin + table("markrgb", texture->markrgb, 3); // mark color + table("random", texture->random, 1); // probability of random dots + table("height", texture->height, 1); // height in pixels (square for cube and skybox) + table("width", texture->width, 1); // width in pixels + table("nchannel", texture->nchannel, 1); // number of channels + table("content_type", texture->content_type, 1); // content type of file + table("file", texture->file, 1); // png file to load; use for all sides of cube + table("gridsize", texture->gridsize, 2); // size of grid for composite file; (1,1)-repeat + // TODO: table("gridlayout", texture->gridlayout, 12); // row-major: L,R,F,B,U,D for faces; . for unused + table("cubefiles", texture->cubefiles, 1); // different file for each side of the cube + table("hflip", texture->hflip, 1); // horizontal flip + table("vflip", texture->vflip, 1); // vertical flip + table("info", texture->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_MATERIAL: { + mjsMaterial* material = mjs_asMaterial(element); + table("textures", material->textures, 1); // names of textures (empty: none) + table("texuniform", material->texuniform, 1); // make texture cube uniform + table("texrepeat", material->texrepeat, 2); // texture repetition for 2D mapping + table("emission", material->emission, 1); // emission + table("specular", material->specular, 1); // specular + table("shininess", material->shininess, 1); // shininess + table("reflectance", material->reflectance, 1); // reflectance + table("metallic", material->metallic, 1); // metallic + table("roughness", material->roughness, 1); // roughness + table("rgba", material->rgba, 4); // rgba + table("info", material->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_PAIR: { + mjsPair* pair = mjs_asPair(element); + table("geomname1", pair->geomname1, 1); // name of geom 1 + table("geomname2", pair->geomname2, 1); // name of geom 2 + table("condim", pair->condim, 1); // contact dimensionality + table("solref", pair->solref, mjNREF); // solver reference, normal direction + table("solreffriction", pair->solreffriction, mjNREF); // solver reference, frictional directions + table("solimp", pair->solimp, mjNIMP); // solver impedance + table("margin", pair->margin, 1); // margin for contact detection + table("gap", pair->gap, 1); // include in solver if distfriction, 5); // full contact friction + table("info", pair->info, 1); // message appended to errors + break; + } + case mjOBJ_EQUALITY: { + mjsEquality* equality = mjs_asEquality(element); + table("type", equality->type, 1); // constraint type + table("data", equality->data, mjNEQDATA); // type-dependent data + table("active", equality->active, 1); // is equality initially active + table("name1", equality->name1, 1); // name of object 1 + table("name2", equality->name2, 1); // name of object 2 + table("objtype", equality->objtype, 1); // type of both objects + table("solref", equality->solref, mjNREF); // solver reference + table("solimp", equality->solimp, mjNIMP); // solver impedance + table("info", equality->info, 1); // message appended to errors + break; + } + case mjOBJ_EXCLUDE: { + mjsExclude* exclude = mjs_asExclude(element); + table("bodyname1", exclude->bodyname1, 1); // name of geom 1 + table("bodyname2", exclude->bodyname2, 1); // name of geom 2 + table("info", exclude->info, 1); // message appended to errors + break; + } + case mjOBJ_NUMERIC: { + mjsNumeric* numeric = mjs_asNumeric(element); + table("data", numeric->data, 1); // initialization data + table("size", numeric->size, 1); // array size, can be bigger than data size + table("info", numeric->info, 1); // message appended to errors + break; + } + case mjOBJ_TEXT: { + mjsText* text = mjs_asText(element); + table("data", text->data, 1); // text string + table("info", text->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_TUPLE: { + mjsTuple* tuple = mjs_asTuple(element); + table("objtype", tuple->objtype, 1); // object types + table("objname", tuple->objname, 1); // object names + table("objprm", tuple->objprm, 1); // object parameters + table("info", tuple->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_KEY: { + mjsKey* key = mjs_asKey(element); + table("time", key->time, 1); // time + table("qpos", key->qpos, 1); // qpos + table("qvel", key->qvel, 1); // qvel + table("act", key->act, 1); // act + table("mpos", key->mpos, 1); // mocap pos + table("mquat", key->mquat, 1); // mocap quat + table("ctrl", key->ctrl, 1); // ctrl + table("info", key->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_PLUGIN: { + mjsPlugin* plugin = mjs_asPlugin(element); + table("name", plugin->name, 1); // instance name + table("plugin_name", plugin->plugin_name, 1); // plugin name + table("active", plugin->active, 1); // is the plugin active + table("info", plugin->info, 1); // message appended to compiler errors + break; + } + default: + // ignore other types + break; + } +} + +void ElementModelGui(const mjModel* model, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + table.SetArrayIndex(mjs_getId(element)); + + MJMODEL_POINTERS_PREAMBLE(model); + switch (element->elemtype) { + case mjOBJ_BODY: { + table.SetPrefix("body_"); + MJMODEL_POINTERS_X(nbody) + break; + } + case mjOBJ_JOINT: { + table.SetPrefix("jnt_"); + MJMODEL_POINTERS_X(njnt) + break; + } + case mjOBJ_ACTUATOR: { + table.SetPrefix("actuator_"); + MJMODEL_POINTERS_X(nu) + break; + } + case mjOBJ_SENSOR: { + table.SetPrefix("sensor_"); + MJMODEL_POINTERS_X(nsensor) + break; + } + case mjOBJ_SITE: { + table.SetPrefix("site_"); + MJMODEL_POINTERS_X(nsite) + break; + } + case mjOBJ_GEOM: { + table.SetPrefix("geom_"); + MJMODEL_POINTERS_X(ngeom) + break; + } + case mjOBJ_LIGHT: { + table.SetPrefix("light_"); + MJMODEL_POINTERS_X(nlight) + break; + } + case mjOBJ_CAMERA: { + table.SetPrefix("cam_"); + MJMODEL_POINTERS_X(ncam) + break; + } + case mjOBJ_MESH: { + table.SetPrefix("mesh_"); + MJMODEL_POINTERS_X(nmesh) + break; + } + case mjOBJ_HFIELD: { + table.SetPrefix("hfield_"); + MJMODEL_POINTERS_X(nhfield) + break; + } + case mjOBJ_SKIN: { + table.SetPrefix("skin_"); + MJMODEL_POINTERS_X(nskin) + break; + } + case mjOBJ_FLEX: { + table.SetPrefix("flex_"); + MJMODEL_POINTERS_X(nflex) + break; + } + case mjOBJ_TENDON: { + table.SetPrefix("tendon_"); + MJMODEL_POINTERS_X(ntendon) + break; + } + case mjOBJ_TEXTURE: { + table.SetPrefix("tex_"); + MJMODEL_POINTERS_X(ntex) + break; + } + case mjOBJ_MATERIAL: { + table.SetPrefix("mat_"); + MJMODEL_POINTERS_X(nmat) + break; + } + default: + // ignore other types + break; + } +} + +void ElementDataGui(const mjData* data, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + table.SetArrayIndex(mjs_getId(element)); + + switch (element->elemtype) { + case mjOBJ_BODY: { + MJDATA_POINTERS_X(nbody) + break; + } + case mjOBJ_JOINT: { + MJDATA_POINTERS_X(njnt) + break; + } + case mjOBJ_SITE: { + table.SetPrefix("site_"); + MJDATA_POINTERS_X(nsite) + break; + } + case mjOBJ_GEOM: { + table.SetPrefix("geom_"); + MJDATA_POINTERS_X(ngeom) + break; + } + case mjOBJ_CAMERA: { + table.SetPrefix("cam_"); + MJDATA_POINTERS_X(ncam) + break; + } + case mjOBJ_LIGHT: { + table.SetPrefix("light_"); + MJDATA_POINTERS_X(nlight) + break; + } + case mjOBJ_SENSOR: { + table.SetPrefix("sensor_"); + MJDATA_POINTERS_X(nsensor) + break; + } + case mjOBJ_MESH: { + table.SetPrefix("mesh_"); + MJDATA_POINTERS_X(nmesh) + break; + } + case mjOBJ_HFIELD: { + table.SetPrefix("hfield_"); + MJDATA_POINTERS_X(nhfield) + break; + } + case mjOBJ_SKIN: { + table.SetPrefix("skin_"); + MJDATA_POINTERS_X(nskin) + break; + } + case mjOBJ_FLEX: { + table.SetPrefix("flex_"); + MJDATA_POINTERS_X(nflex) + break; + } + case mjOBJ_TENDON: { + table.SetPrefix("ten_"); + MJDATA_POINTERS_X(ntendon) + break; + } + case mjOBJ_TEXTURE: { + table.SetPrefix("tex_"); + MJDATA_POINTERS_X(ntex) + break; + } + case mjOBJ_MATERIAL: { + table.SetPrefix("mat_"); + MJDATA_POINTERS_X(nmat) + break; + } + default: + break; + } +} +} // namespace mujoco::platform diff --git a/src/experimental/platform/gui_spec.h b/src/experimental/platform/gui_spec.h new file mode 100644 index 00000000..e1094583 --- /dev/null +++ b/src/experimental/platform/gui_spec.h @@ -0,0 +1,37 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ +#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ + +#include + +#include + +namespace mujoco::platform { + +using SpecElementCallbackFn = std::function; + +// UX for displaying the spec as a tree. +void SpecExplorerGui(mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete); + +// UX for displaying the properties of an mjSpec element. +void ElementSpecGui(const mjSpec* spec, mjsElement* element); +void ElementDataGui(const mjData* data, mjsElement* element); +void ElementModelGui(const mjModel* model, mjsElement* element); + +} // namespace mujoco::platform + +#endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ diff --git a/src/experimental/platform/imgui_widgets.cc b/src/experimental/platform/imgui_widgets.cc index 10a7b374..def513e9 100644 --- a/src/experimental/platform/imgui_widgets.cc +++ b/src/experimental/platform/imgui_widgets.cc @@ -14,9 +14,13 @@ #include "experimental/platform/imgui_widgets.h" +#include +#include #include #include +#include #include +#include #include #include @@ -57,6 +61,228 @@ KeyValues ReadIniSection(const std::string& contents, return key_values; } +ImGui_DataTable::ImGui_DataTable(float w1, float w2) { + ImGui::BeginTable("##PropertiesTable", 2); + const float width = ImGui::GetContentRegionAvail().x; + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, width * w1); + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, width * w2); +} +ImGui_DataTable::~ImGui_DataTable() { ImGui::EndTable(); } + +void ImGui_DataTable::SetArrayIndex(int index) { index_ = index; } + +void ImGui_DataTable::SetPrefix(const char* prefix) { + prefix_ = strlen(prefix); +} + +void ImGui_DataTable::operator()(const char* label, const uintptr_t* ptr, + int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("(%s)", &ptr[index_ + i] ? "[ptr]" : "null"); + } +} + +void ImGui_DataTable::operator()(const char* label, const char* ptr, int n) { + if (n == 1) { + MakeLabel(label); + ImGui::Text("%s", &ptr[index_]); + } else { + mju_error("char cannot be converted to a vector"); + } +} + +void ImGui_DataTable::operator()(const char* label, const mjtByte* ptr, int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("%s", ptr[index_ + i] ? "true" : "false"); + } +} + +void ImGui_DataTable::operator()(const char* label, const mjtByte& val, int n) { + MakeLabel(label, 0, 1); + ImGui::Text("%s", val ? "true" : "false"); +} + +void ImGui_DataTable::operator()(const char* label, const mjtSize* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const int* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const float* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const double* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const mjtSize& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const int& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const float& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const double& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const std::string* ptr, + int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("%s", ptr[i].c_str()); + } +} + +void ImGui_DataTable::operator()(const char* label, + const std::vector* ptr, int n) { + if (n == 1) { + for (int i = 0; i < ptr->size(); ++i) { + MakeLabel(label, i, ptr->size()); + ImGui::Text("%s", ptr->at(i).c_str()); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +void ImGui_DataTable::operator()(const char* label, const std::vector* ptr, + int n) { + if (n == 1) { + const int size = ptr->size(); + if (size == 0) { + (*this)(label, "[empty]", 1); + } else { + std::string tmp = "[" + std::to_string(size) + " values]"; + (*this)(label, tmp.c_str(), 1); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +void ImGui_DataTable::operator()(const char* label, + const std::vector* ptr, int n) { + if (n == 1) { + const int size = ptr->size(); + if (size == 0) { + (*this)(label, "[empty]", 1); + } else { + std::string tmp = "[" + std::to_string(size) + " values]"; + (*this)(label, tmp.c_str(), 1); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +template +void ImGui_DataTable::Scalar(const char* label, const T& value, int n) { + if (n == 1) { + Numeric(label, &value, n); + } else { + mju_error("scalar cannot be converted to a vector"); + } +} + +template +void ImGui_DataTable::Numeric(const char* label, const T* ptr, int n) { + const T* addr = ptr + index_ * n; + + using U = std::conditional_t, float, int>; + + // special treatment for NaNs. + if constexpr (std::is_same_v) { + if (*addr != *addr) { + MakeLabel(label); + ImGui::Text("nan"); + return; + } + } + + constexpr const char* fmt1 = + std::is_floating_point_v ? "%f" : "%d"; + constexpr const char* fmt2 = + std::is_floating_point_v ? "%f %f" : "%d %d"; + constexpr const char* fmt3 = + std::is_floating_point_v ? "%f %f %f" : "%d %d %d"; + constexpr const char* fmt4 = + std::is_floating_point_v ? "%f %f %f %f" : "%d %d %d %d"; + + auto text1 = [&](int offset) { + ImGui::Text(fmt1, (U)(addr[offset])); + }; + auto text2 = [&](int offset) { + ImGui::Text(fmt2, (U)(addr[offset + 0]), (U)(addr[offset + 1])); + }; + auto text3 = [&](int offset) { + ImGui::Text(fmt3, (U)(addr[offset + 0]), (U)(addr[offset + 1]), + (U)(addr[offset + 2])); + }; + auto text4 = [&](int offset) { + ImGui::Text(fmt4, (U)(addr[offset + 0]), (U)(addr[offset + 1]), + (U)(addr[offset + 2]), (U)(addr[offset + 3])); + }; + + if (n == 1) { + MakeLabel(label); + text1(0); + } else if (n == 2) { + MakeLabel(label); + text2(0); + } else if (n == 3) { + MakeLabel(label); + text3(0); + } else if (n == 4) { + MakeLabel(label); + text4(0); + } else if (n == 6) { + MakeLabel(label); + text3(0); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(3); + } else if (n == 9) { + MakeLabel(label); + text3(0); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(3); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(6); + } else { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + text1(i); + } + } +} + +void ImGui_DataTable::MakeLabel(const char* label, int index, int total) { + if (total == 1) { + ImGui::TableNextColumn(); + ImGui::Text("%s", &label[prefix_]); + ImGui::TableNextColumn(); + } else { + const std::string tmp = + std::string(&label[prefix_]) + "[" + std::to_string(index) + "]"; + ImGui::TableNextColumn(); + ImGui::Text("%s", tmp.c_str()); + ImGui::TableNextColumn(); + } +} + bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max) { float f = *value; const bool res = ImGui::SliderFloat(name, &f, min, max); diff --git a/src/experimental/platform/imgui_widgets.h b/src/experimental/platform/imgui_widgets.h index 29f5005f..7fd64ca8 100644 --- a/src/experimental/platform/imgui_widgets.h +++ b/src/experimental/platform/imgui_widgets.h @@ -15,11 +15,13 @@ #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_IMGUI_WIDGETS_H_ #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_IMGUI_WIDGETS_H_ +#include #include #include #include #include #include +#include #include #include @@ -112,6 +114,12 @@ struct ScopedStyle { return *this; } + ScopedStyle& Color(ImGuiCol col, ImGuiCol col2) { + ImGui::PushStyleColor(col, CurrentColor(col2)); + ++num_colors; + return *this; + } + ScopedStyle& Var(ImGuiStyleVar var, float value) { ImGui::PushStyleVar(var, value); ++num_vars; @@ -124,6 +132,10 @@ struct ScopedStyle { return *this; } + ImVec4 CurrentColor(ImGuiCol col) { + return ImGui::GetStyle().Colors[col]; + } + void Reset() { ImGui::PopStyleVar(num_vars); ImGui::PopStyleColor(num_colors); @@ -135,6 +147,74 @@ struct ScopedStyle { int num_vars = 0; }; +// Helper for displaying rows of key/value pairs in an ImGui table. +// +// Designed specifically to be used to display mjSpec, mjModel, and mjData +// values. +// +// To add a value to the table, call the operator() function with the label, +// the value, and (optionally) the dimensionality of the value (e.g. for vectors +// and matrices). To support generic code, even scalar values should be passed +// to operator() with n = 1. +class ImGui_DataTable { + public: + // Starts the table (i.e. ImGui::BeginTable()) with two columns of the + // specified widths. + ImGui_DataTable(float w1 = 0.25f, float w2 = 0.75f); + + // Ends the table (e.g. ImGui::EndTable(). + ~ImGui_DataTable(); + + ImGui_DataTable(const ImGui_DataTable& other) = delete; + ImGui_DataTable& operator=(const ImGui_DataTable& other) = delete; + + // Sets the offset into an array of values (e.g. for pointers in mjModel and + // mjData). This is only used for the display functions that take a pointer. + void SetArrayIndex(int index); + + // Sets the prefix that will be removed from all labels. Note: that we simply + // remove the first N characters of the label without actually comparing + // against this prefix. + void SetPrefix(const char* prefix); + + // Displays a labelled value in the table. + void operator()(const char* label, const uintptr_t* ptr, int n); + void operator()(const char* label, const char* ptr, int n); + void operator()(const char* label, const mjtByte* ptr, int n); + void operator()(const char* label, const mjtSize* ptr, int n); + void operator()(const char* label, const int* ptr, int n); + void operator()(const char* label, const float* ptr, int n); + void operator()(const char* label, const double* ptr, int n); + + // Displays a single scalar value in the table. Assumes n == 1. This should + // only be used for mjSpec objects and, therefore, will ignore the array index + // if set. + void operator()(const char* label, const mjtByte& val, int n); + void operator()(const char* label, const mjtSize& val, int n); + void operator()(const char* label, const int& val, int n); + void operator()(const char* label, const float& val, int n); + void operator()(const char* label, const double& val, int n); + + // Overloads for C++ container types. Assumes its only used for mjSpec objects + // and, therefore, will ignore the array index if set. + void operator()(const char* label, const std::string* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + + private: + template + void Numeric(const char* label, const T* ptr, int n); + + template + void Scalar(const char* label, const T& value, int n); + + void MakeLabel(const char* label, int index = 0, int total = 1); + + int prefix_ = 0; + int index_ = 0; +}; + // ImGui Slider that supports both float and double types. bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max); diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 352e8c8a..bade8680 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -36,6 +36,7 @@ #include #include "experimental/platform/file_dialog.h" #include "experimental/platform/gui.h" +#include "experimental/platform/gui_spec.h" #include "experimental/platform/helpers.h" #include "experimental/platform/imgui_widgets.h" #include "experimental/platform/interaction.h" @@ -132,6 +133,7 @@ void App::ClearModel() { tmp_ = UiTempState(); load_error_ = ""; step_error_ = ""; + edit_error_ = ""; } void App::Recompile() { @@ -257,6 +259,7 @@ void App::ResetPhysics() { mj_resetData(model(), data()); mj_forward(model(), data()); step_error_ = ""; + edit_error_ = ""; } void App::UpdatePhysics() { @@ -419,15 +422,42 @@ void App::ProcessPendingLoads() { }); } -void App::SpecDeleteSelectedElement() { - spec_op_ = [this]() { - mjs_delete(spec(), tmp_.element); - if (tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select == tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - } - tmp_.element = nullptr; +void App::SpecSelectElement(mjsElement* element) { + tmp_.element = element; + if (tmp_.element == nullptr) { tmp_.element_id = -1; + } else { + tmp_.element_id = mjs_getId(tmp_.element); + + // If we selected a body, then select the same body for perturb. + if (tmp_.element->elemtype == mjOBJ_BODY && + perturb_.select != tmp_.element_id) { + mjv_defaultPerturb(&perturb_); + perturb_.select = tmp_.element_id; + } + } +} + +void App::SpecDeleteElement(mjsElement* element) { + if (element == nullptr) { + return; + } + // Only bodies can be deleted for now... + if (element->elemtype != mjOBJ_BODY) { + edit_error_ = "WARNING: Only bodies can be deleted (for now...)"; + return; + } + spec_op_ = [this, element]() { + mjs_delete(spec(), element); + if (tmp_.element == element) { + tmp_.element = nullptr; + tmp_.element_id = -1; + } + if (element->elemtype == mjOBJ_BODY) { + if (perturb_.select == mjs_getId(element)) { + mjv_defaultPerturb(&perturb_); + } + } Recompile(); }; } @@ -635,7 +665,7 @@ void App::HandleKeyboardEvents() { } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) { ResetPhysics(); } else if (ImGui_IsChordJustPressed(ImGuiKey_Delete)) { - SpecDeleteSelectedElement(); + SpecDeleteElement(tmp_.element); } else if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) { SelectParentPerturb(model(), perturb_); } else if (ImGui_IsChordJustPressed(ImGuiKey_F1)) { @@ -896,7 +926,7 @@ void App::BuildGui() { if (explorer_is_open && tmp_.element != nullptr) { if (ImGui::Begin("Properties")) { - PropertiesGui(); + SpecPropertiesGui(); } ImGui::End(); } @@ -1129,117 +1159,62 @@ void App::DataInspectorGui() { ImGui::EndChild(); } -void DisplayElementTree(mjsElement* element) { - const mjString* name = mjs_getName(element); - if (name->empty()) { - ImGui::Text("(unnamed)"); - } else { - ImGui::Text("%s", name->c_str()); - } -} - void App::SpecExplorerGui() { if (!has_spec()) { ImGui::Text("No mjSpec loaded."); return; } - const ImGuiTreeNodeFlags flags = - ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + auto on_delete = [this](mjsElement* element) { SpecDeleteElement(element); }; - auto display_group = [this](mjtObj type, const std::string& prefix, - std::function delete_callback = {}) { - mjsElement* element = mjs_firstElement(spec(), type); - while (element) { - const int id = mjs_getId(element); - - const mjString* name = mjs_getName(element); - std::string label = *name; - if (label.empty()) { - label = "(" + prefix + " " + std::to_string(id) + ")"; - } - - const bool selected = (tmp_.element == element); - if (ImGui::Selectable(label.c_str(), selected, - ImGuiSelectableFlags_AllowOverlap)) { - tmp_.element = element; - tmp_.element_id = id; - } - - if (selected && delete_callback) { - // Right-align the delete button. - const float button_width = ImGui::CalcTextSize(ICON_DELETE).x + - ImGui::GetStyle().FramePadding.x * 2.0f; - ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); - if (ImGui::SmallButton(ICON_DELETE)) { - delete_callback(); - } - } - - element = mjs_nextElement(spec(), element); - } - }; - - if (ImGui::TreeNodeEx("Bodies", flags)) { - display_group(mjOBJ_BODY, "Body", [this] { SpecDeleteSelectedElement(); }); - ImGui::TreePop(); - } - - if (ImGui::TreeNodeEx("Joints", flags)) { - display_group(mjOBJ_JOINT, "Joint"); - ImGui::TreePop(); - } - - if (ImGui::TreeNodeEx("Sites", flags)) { - display_group(mjOBJ_SITE, "Site"); - ImGui::TreePop(); - } - - // If we selected a body, then select the same body for the perturb object. - if (tmp_.element && tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select != tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - perturb_.select = tmp_.element_id; + mjsElement* element = tmp_.element; + platform::SpecExplorerGui(&element, spec(), on_delete); + if (element != tmp_.element) { + SpecSelectElement(element); } } -void App::PropertiesGui() { - if (tmp_.element == nullptr) { - ImGui::Text("No element selected."); - return; - } +void App::SpecPropertiesGui() { + platform::ScopedStyle style; - if (ImGui::BeginTable("##PropertiesHeader", 2)) { - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 20); - ImGui::TableNextColumn(); - ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); - ImGui::TableNextColumn(); - if (tmp_.element->elemtype == mjOBJ_BODY) { - if (ImGui::SmallButton(ICON_DELETE)) { - SpecDeleteSelectedElement(); - } - } - ImGui::EndTable(); + ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); + ImGui::SameLine(); + ImGui::Text("(%d)", tmp_.element_id); + + ImGui::SameLine(120); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); } + if (ImGui::SmallButton("S")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kSpec; + } + ImGui::SetItemTooltip("Spec"); + style.Reset(); + ImGui::SameLine(); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); + } + if (ImGui::SmallButton("M")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kModel; + } + ImGui::SetItemTooltip("Model"); + style.Reset(); + ImGui::SameLine(); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kData) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); + } + if (ImGui::SmallButton("D")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kData; + } + ImGui::SetItemTooltip("Data"); + style.Reset(); ImGui::Separator(); - - switch (tmp_.element->elemtype) { - case mjOBJ_BODY: - platform::BodyPropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - case mjOBJ_JOINT: - platform::JointPropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - case mjOBJ_SITE: - platform::SitePropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - default: - // ignore other types - break; + if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { + platform::ElementSpecGui(spec(), tmp_.element); + } else if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { + platform::ElementModelGui(model(), tmp_.element); + } else { + platform::ElementDataGui(data(), tmp_.element); } } @@ -1617,6 +1592,9 @@ void App::StatusBarGui() { } else if (!load_error_.empty()) { ImGui::SameLine(); ImGui::Text(" | Load Error: %s", load_error_.c_str()); + } else if (!edit_error_.empty()) { + ImGui::SameLine(); + ImGui::Text(" | Edit Error: %s", edit_error_.c_str()); } ImGui::TableNextColumn(); diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index d2cce3a5..66332e70 100644 --- a/src/experimental/studio/app.h +++ b/src/experimental/studio/app.h @@ -91,6 +91,12 @@ class App { kModelFromBuffer, }; + enum class SpecPropertiesMode { + kSpec, + kModel, + kData, + }; + // UI state that is persisted across application runs struct UiState { char watch_field[1000] = "qpos"; @@ -136,6 +142,7 @@ class App { std::vector speed_names; // Spec Properties. + SpecPropertiesMode spec_prop_mode = SpecPropertiesMode::kSpec; mjsElement* element = nullptr; int element_id = -1; @@ -209,9 +216,10 @@ class App { void ModelOptionsGui(); void DataInspectorGui(); void SpecExplorerGui(); - void PropertiesGui(); + void SpecPropertiesGui(); - void SpecDeleteSelectedElement(); + void SpecSelectElement(mjsElement* element); + void SpecDeleteElement(mjsElement* element); float GetExpectedLabelWidth(); std::vector GetCameraNames(); @@ -228,6 +236,7 @@ class App { std::string model_path_; std::string load_error_; std::string step_error_; + std::string edit_error_; std::optional pending_load_; bool preserve_camera_on_load_ = false; ModelKind model_kind_ = kEmptyModel; From 94acdb33bd12f35a8dc6bc79d435fea34e81a6cc Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Tue, 24 Feb 2026 08:16:39 -0800 Subject: [PATCH 35/48] Fix Impl.CPP memory leaks due to mjModel/mjData not clearing from pytree. PiperOrigin-RevId: 874622402 Change-Id: I38a00633dea77d940e4b0422f93e9456ad84daf9 --- mjx/mujoco/mjx/_src/io.py | 89 +++++++++++++++++++++++++--------- mjx/mujoco/mjx/_src/io_test.py | 27 +++++------ mjx/mujoco/mjx/_src/types.py | 2 - 3 files changed, 80 insertions(+), 38 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 670c93e1..d4e0d1fa 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -494,9 +494,12 @@ def _put_model_warp( return _strip_weak_type(model) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_model_cpp( m: mujoco.MjModel, device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" @@ -517,9 +520,11 @@ def _put_model_cpp( c_pointers_impl = types.ModelCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _model=m, ) + if keepalive_refs is not None: + keepalive_refs[addr] = m + model = types.Model( **{k: copy.copy(v) for k, v in fields.items()}, _impl=c_pointers_impl ) @@ -532,6 +537,7 @@ def put_model( device: Optional[jax.Device] = None, impl: Optional[Union[str, types.Impl]] = None, graph_mode: Optional[mjxw.types.GraphMode] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model. @@ -541,6 +547,9 @@ def put_model( impl: implementation to use graph_mode: CUDA graph capture mode (for Warp only). Use GraphMode enum from warp._src.jax_experimental.ffi. GraphMode.WARP is the default mode. + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl to keep the model alive. Returns: an mjx.Model placed on device @@ -560,7 +569,7 @@ def put_model( graph_mode = graph_mode or getattr(mjxw.types.GraphMode, 'WARP') return _put_model_warp(m, graph_mode, device) elif impl == types.Impl.CPP: - return _put_model_cpp(m, device) + return _put_model_cpp(m, device, keepalive_refs=keepalive_refs) else: raise ValueError(f'Unsupported implementation: {impl}') @@ -906,9 +915,12 @@ def _make_data_warp( return data +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _make_data_cpp( m: Union[types.Model, mujoco.MjModel], device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data for the CPP implementation.""" if isinstance(m, mujoco.MjModel): @@ -918,7 +930,13 @@ def _make_data_cpp( m_impl = m._impl # pylint: disable=protected-access if not isinstance(m_impl, types.ModelCPP): raise ValueError(f'Expected ModelCPP impl, got {type(m_impl)}') - mj_model = m_impl._model # pylint: disable=protected-access + model_addr = int(m_impl.pointer_lo) | (int(m_impl.pointer_hi) << 32) + if keepalive_refs is None or model_addr not in keepalive_refs: + raise ValueError( + 'keepalive_refs must be provided and contain the model when calling' + ' _make_data_cpp with a types.Model.' + ) + mj_model = keepalive_refs[model_addr] # Create the raw MuJoCo data mj_data = mujoco.MjData(mj_model) @@ -933,9 +951,11 @@ def _make_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=[mj_data], ) + if keepalive_refs is not None: + keepalive_refs[addr] = [mj_data] + data = types.Data( _impl=c_pointers_impl, **fields, @@ -952,6 +972,7 @@ def make_data( nconmax: Optional[int] = None, naconmax: Optional[int] = None, njmax: Optional[int] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data. @@ -969,6 +990,9 @@ def make_data( `naconmax` argument to set the upper bound for the number of contacts across all worlds, rather than the `nconmax` argument from MuJoCo Warp. njmax: maximum number of constraints to allocate for warp across all worlds + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl when passing a types.Model. Returns: an initialized mjx.Data placed on device @@ -997,7 +1021,7 @@ def make_data( elif impl == types.Impl.C: return _make_data_c(m, device) elif impl == types.Impl.CPP: - return _make_data_cpp(m, device) + return _make_data_cpp(m, device, keepalive_refs=keepalive_refs) elif impl == types.Impl.WARP: _check_warp_installed() naconmax = nconmax if naconmax is None else naconmax @@ -1300,11 +1324,14 @@ def _put_data_c( return _strip_weak_type(data) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_data_cpp( m: mujoco.MjModel, d: mujoco.MjData, device: Optional[jax.Device] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" @@ -1317,6 +1344,8 @@ def _put_data_cpp( mujoco.mj_copyData(new_d, m, d) data_list.append(new_d) addr = new_d._address + if keepalive_refs is not None: + keepalive_refs[addr] = new_d # To ensure that we retain the full pointer even if jax.config.enable_x64 is # set to True, we store the pointer as two 32-bit values. In the FFI call, # we combine the two values into a single pointer value. @@ -1341,7 +1370,6 @@ def _put_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=data_list, ) data = types.Data( @@ -1399,6 +1427,7 @@ def put_data( naconmax: Optional[int] = None, njmax: Optional[int] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data. @@ -1415,6 +1444,8 @@ def put_data( njmax: maximum number of constraints to allocate for warp dummy_arg_for_batching: dummy argument to use for batching in cpp implementation + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Returns: an mjx.Data placed on device @@ -1434,7 +1465,11 @@ def put_data( return _put_data_c(m, d, device) elif impl == types.Impl.CPP: return _put_data_cpp( - m, d, device, dummy_arg_for_batching=dummy_arg_for_batching + m, + d, + device, + dummy_arg_for_batching=dummy_arg_for_batching, + keepalive_refs=keepalive_refs, ) elif impl == types.Impl.WARP: _check_warp_installed() @@ -1692,15 +1727,18 @@ def _get_data_into( mujoco.mj_factorM(m, result_i) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _get_data_into_cpp( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """Gets mjx.Data from CPP impl into an existing mujoco.MjData or list. For the CPP implementation, the mjx.Data wraps underlying mujoco.MjData - objects that are stored in DataCPP._data. This function simply copies the + objects that are stored in keepalive_refs. This function simply copies the data from those underlying MjData objects to the result using mj_copyData. """ @@ -1712,13 +1750,8 @@ def _get_data_into_cpp( if not isinstance(d_impl, types.DataCPP): raise ValueError(f'Expected DataCPP impl, got {type(d_impl)}') - mj_data_list = d_impl._data # pylint: disable=protected-access - - if batch_size > len(mj_data_list): - raise ValueError( - f'Batch size {batch_size} exceeds number of underlying MjData objects ' - f'({len(mj_data_list)}). Cannot copy data.' - ) + if keepalive_refs is None: + raise ValueError('keepalive_refs must be provided for CPP implementation.') # Verify that the underlying MjData state matches the mjx.Data state # Ideally we'd use mj_getState and get_state here but that requires an @@ -1728,7 +1761,19 @@ def _get_data_into_cpp( d_i: types.Data = ( jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d ) - src_data = mj_data_list[i] + result_i = result[i] if batched else result + + if batched: + addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32) + else: + addr_i = int(d_impl.pointer_lo) | (int(d_impl.pointer_hi) << 32) + + if addr_i not in keepalive_refs: + raise ValueError( + f'Address {addr_i} not found in keepalive_refs. ' + 'Ensure keepalive_refs from the original compile() is passed.' + ) + src_data = keepalive_refs[addr_i] needs_syncing = False for field in fields_to_check: @@ -1745,9 +1790,6 @@ def _get_data_into_cpp( src_data.mocap_quat[:] = d_i.mocap_quat mujoco.mj_kinematics(m, src_data) - for i in range(batch_size): - result_i = result[i] if batched else result - src_data = mj_data_list[i] mujoco.mj_copyData(result_i, m, src_data) @@ -1755,6 +1797,7 @@ def get_data_into( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" is_batched = isinstance(result, list) @@ -1770,7 +1813,7 @@ def get_data_into( return _get_data_into(result, m, d) if d.impl == types.Impl.CPP: - return _get_data_into_cpp(result, m, d) + return _get_data_into_cpp(result, m, d, keepalive_refs=keepalive_refs) if d.impl == types.Impl.WARP: return _get_data_into_warp(result, m, d) @@ -1781,7 +1824,9 @@ def get_data_into( def get_data( - m: mujoco.MjModel, d: types.Data + m: mujoco.MjModel, + d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> Union[mujoco.MjData, List[mujoco.MjData]]: """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" batched = len(d.qpos.shape) > 1 @@ -1792,7 +1837,7 @@ def get_data( else: result = mujoco.MjData(m) - get_data_into(result, m, d) + get_data_into(result, m, d, keepalive_refs=keepalive_refs) return result diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 1cd0afe9..9910e9b5 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -768,9 +768,10 @@ class DataIOTest(parameterized.TestCase): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d, impl=impl) + keepalive = {} if impl == 'cpp' else None + dx = mjx.put_data(m, d, impl=impl, keepalive_refs=keepalive) d_2 = mujoco.MjData(m) - mjx.get_data_into(d_2, m, dx) + mjx.get_data_into(d_2, m, dx, keepalive_refs=keepalive) # check a few fields np.testing.assert_allclose(d_2.qpos, d.qpos) @@ -931,12 +932,11 @@ class DataIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) # Test cases for `_resolve_impl_and_device` where the device is @@ -1265,12 +1265,11 @@ class StateIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) def test_get_set_state(self): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index fcdc41d6..9f25a845 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -538,7 +538,6 @@ class ModelCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _model: mujoco.MjModel class DataCPP(PyTreeNode): @@ -548,7 +547,6 @@ class DataCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _data: list[Any] = dataclasses.field(default_factory=list, repr=False) class ModelC(PyTreeNode): From 09c7633a76b0de1636e2e993433beacd65bf7484 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 24 Feb 2026 10:44:41 -0800 Subject: [PATCH 36/48] Add caching for flex stiffness computation. This change introduces methods to cache and load the computed stiffness matrix and nodal masses for flex elements. The cache key is generated based on material properties (Young's modulus, Poisson's ratio), interpolation order, and the bounding box of the flex. This avoids redundant expensive computations when compiling models with identical flex definitions. PiperOrigin-RevId: 874689771 Change-Id: I6d1e299baf294a984b507e2a2602b2c068610536 --- src/user/user_mesh.cc | 93 ++++++++++++++++++++++++++++++++++++- src/user/user_objects.h | 8 ++++ test/user/user_flex_test.cc | 40 ++++++++++++++++ 3 files changed, 140 insertions(+), 1 deletion(-) diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 70d78bed..5df01943 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -4053,6 +4053,88 @@ void mjCFlex::ResolveReferences(const mjCModel* m) { } +std::string mjCFlex::ComputeStiffnessCacheKey() const { + std::size_t hash = 0; + auto combine = [&hash](std::size_t v) { + hash ^= v + 0x9e3779b9 + (hash << 6) + (hash >> 2); + }; + + combine(std::hash{}(young)); + combine(std::hash{}(poisson)); + combine(std::hash{}(order_)); + + // compute bounding box from vertex positions + if (!vert_.empty()) { + double minx = vert_[0], maxx = vert_[0]; + double miny = vert_[1], maxy = vert_[1]; + double minz = vert_[2], maxz = vert_[2]; + for (std::size_t i = 3; i < vert_.size(); i += 3) { + minx = std::min(minx, vert_[i]); + maxx = std::max(maxx, vert_[i]); + miny = std::min(miny, vert_[i + 1]); + maxy = std::max(maxy, vert_[i + 1]); + minz = std::min(minz, vert_[i + 2]); + maxz = std::max(maxz, vert_[i + 2]); + } + combine(std::hash{}(maxx - minx)); + combine(std::hash{}(maxy - miny)); + combine(std::hash{}(maxz - minz)); + } + + for (std::size_t i = 0; i < vert_.size(); i += std::max(1, (int)vert_.size()/100)) { + combine(std::hash{}(vert_[i])); + } + + for (std::size_t i = 0; i < shell.size(); i += std::max(1, (int)shell.size()/50)) { + combine(std::hash{}(shell[i])); + } + + return "flex_stiffness:" + std::to_string(hash); +} + + +bool mjCFlex::LoadCachedStiffness() { + mjCCache* cache = reinterpret_cast(mj_getCache()->impl_); + if (!cache) return false; + + std::string key = ComputeStiffnessCacheKey(); + + auto load_fn = [this](const void* data) { + const auto* cached = static_cast*>(data); + stiffness = *cached; + return true; + }; + + mjResource dummy_resource{}; + dummy_resource.name = const_cast(key.c_str()); + dummy_resource.timestamp[0] = '\0'; + + return cache->PopulateData(key, &dummy_resource, load_fn); +} + + +void mjCFlex::CacheStiffness() { + mjCCache* cache = reinterpret_cast(mj_getCache()->impl_); + if (!cache || stiffness.empty()) return; + + std::string key = ComputeStiffnessCacheKey(); + + auto* cached = new std::vector(stiffness); + + std::size_t size = sizeof(*cached) + sizeof(double) * stiffness.size(); + + std::shared_ptr cached_data(cached, [](const void* data) { + delete static_cast*>(data); + }); + + mjResource dummy_resource{}; + dummy_resource.name = const_cast(key.c_str()); + dummy_resource.timestamp[0] = '\0'; + + cache->Insert("", key, &dummy_resource, cached_data, size); +} + + // compiler void mjCFlex::Compile(const mjVFS* vfs) { CopyFromSpec(); @@ -4295,7 +4377,6 @@ void mjCFlex::Compile(const mjVFS* vfs) { if (min_size > nelem) { throw mjCError(this, "Trilinear dofs are require at least %d elements", "", min_size); } - ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_); } // geometrically nonlinear elasticity @@ -4345,6 +4426,16 @@ void mjCFlex::Compile(const mjVFS* vfs) { // create shell fragments and element-vertex collision pairs CreateShellPair(); + // compute linear stiffness for interpolated elements (cached) + bool stiffness_cached = false; + if (young > 0 && interpolated) { + stiffness_cached = LoadCachedStiffness(); + } + + if (!stiffness_cached && young > 0 && interpolated) { + ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_); + } + // create bounding volume hierarchy CreateBVH(); diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 85a1d5e1..2e5ad168 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -992,6 +992,9 @@ class mjCFlex_ : public mjCBase { std::vector spec_elem_; std::vector spec_texcoord_; std::vector spec_elemtexcoord_; + + // caching + std::vector cached_stiffness_; // cached stiffness matrix }; class mjCFlex: public mjCFlex_, private mjsFlex { @@ -1040,6 +1043,11 @@ class mjCFlex: public mjCFlex_, private mjsFlex { std::vector node0_; // node Cartesian positions int order_ = 0; // interpolation order + + // stiffness caching + std::string ComputeStiffnessCacheKey() const; + bool LoadCachedStiffness(); + void CacheStiffness(); }; diff --git a/test/user/user_flex_test.cc b/test/user/user_flex_test.cc index 8db36823..841bcdac 100644 --- a/test/user/user_flex_test.cc +++ b/test/user/user_flex_test.cc @@ -456,6 +456,46 @@ TEST_F(UserFlexTest, StiffnessMatrix) { mj_deleteModel(m); } +TEST_F(UserFlexTest, StiffnessCacheDiffersByGeometry) { + std::array error; + + // Create two flexes with same material but different bounding boxes + static constexpr char xml_small[] = R"( + + + + + + + + + )"; + + static constexpr char xml_large[] = R"( + + + + + + + + + )"; + + mjModel* m_small = LoadModelFromString(xml_small, error.data(), error.size()); + ASSERT_THAT(m_small, NotNull()) << error.data(); + + mjModel* m_large = LoadModelFromString(xml_large, error.data(), error.size()); + ASSERT_THAT(m_large, NotNull()) << error.data(); + + // Same number of nodes but different stiffness due to different geometry + EXPECT_EQ(m_small->nflexnode, m_large->nflexnode); + EXPECT_NE(m_small->flex_stiffness[0], m_large->flex_stiffness[0]); + + mj_deleteModel(m_small); + mj_deleteModel(m_large); +} + TEST_F(UserFlexTest, LoadTexture) { const std::string xml_path = GetTestDataFilePath("user/testdata/textured_torus_flex.xml"); From a83fa7299df69bc63673925c56b1ec066d3a1877 Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Tue, 24 Feb 2026 13:08:14 -0800 Subject: [PATCH 37/48] Add support for loading XML models with VFS in WASM bindings. PiperOrigin-RevId: 874756898 Change-Id: If2c30aef4c690cf0852587261ca32a3b91448621 --- wasm/codegen/generated/bindings.cc | 16 ++++++++-- wasm/codegen/generators/structs.py | 13 +++++--- wasm/codegen/templates/bindings.cc | 11 ++++++- wasm/tests/bindings_test.ts | 49 ++++++++++++++++++++++++++++-- 4 files changed, 79 insertions(+), 10 deletions(-) diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index c49629ee..d8b69429 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -8329,7 +8329,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr mj_loadXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper_1(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { @@ -8338,6 +8338,15 @@ std::unique_ptr mj_loadXML_wrapper(std::string filename) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_loadXML_wrapper_2(std::string filename, const MjVFS& vfs) { + char error[1000]; + mjModel *model = mj_loadXML(filename.c_str(), vfs.get(), error, sizeof(error)); + if (!model) { + mju_error("Loading error: %s\n", error); + } + return std::unique_ptr(new MjModel(model)); +} + void mj_saveModel_wrapper(const MjModel& m, const StringOrNull& filename, const val& buffer) { UNPACK_NULLABLE_STRING(filename); UNPACK_NULLABLE_VALUE(uint8_t, buffer); @@ -11581,8 +11590,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("useexisting", &MjLROpt::useexisting, &MjLROpt::set_useexisting, reference()) .property("uselimit", &MjLROpt::uselimit, &MjLROpt::set_uselimit, reference()); emscripten::class_("MjModel") - .class_function("mj_loadXML", &mj_loadXML_wrapper, take_ownership()) - .class_function("mj_loadBinary", &mj_loadModel_wrapper, take_ownership()) + .class_function("mj_loadXML", emscripten::select_overload(std::string)>(&mj_loadXML_wrapper_1)) + .class_function("mj_loadXML", emscripten::select_overload(std::string, const MjVFS&)>(&mj_loadXML_wrapper_2)) + .class_function("mj_loadModel", &mj_loadModel_wrapper) .constructor() // Binds the functions on MjModel that return accessors. #define X_ACCESSOR(NAME, Name, OBJTYPE, field_name, nfield) \ diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py index f6282693..0517dba9 100644 --- a/wasm/codegen/generators/structs.py +++ b/wasm/codegen/generators/structs.py @@ -549,14 +549,19 @@ def _build_struct_bindings( MJDATA_ACCESSORS #undef X_ACCESSOR""".lstrip()) elif w == "MjModel": - f1 = common.wrapped_function_name( - introspect_functions.FUNCTIONS["mj_loadXML"] + builder.line( + '.class_function("mj_loadXML",' + " emscripten::select_overload(std::string)>(&mj_loadXML_wrapper_1))" + ) + builder.line( + '.class_function("mj_loadXML",' + " emscripten::select_overload(std::string," + " const MjVFS&)>(&mj_loadXML_wrapper_2))" ) - builder.line(f'.class_function("mj_loadXML", &{f1}, take_ownership())') f2 = common.wrapped_function_name( introspect_functions.FUNCTIONS["mj_loadModel"] ) - builder.line(f'.class_function("mj_loadBinary", &{f2}, take_ownership())') + builder.line(f'.class_function("mj_loadModel", &{f2})') builder.line(".constructor()") builder.line(""" // Binds the functions on MjModel that return accessors. diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index 8e29a080..8f2a1ecf 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -714,7 +714,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr mj_loadXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper_1(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { @@ -723,6 +723,15 @@ std::unique_ptr mj_loadXML_wrapper(std::string filename) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_loadXML_wrapper_2(std::string filename, const MjVFS& vfs) { + char error[1000]; + mjModel *model = mj_loadXML(filename.c_str(), vfs.get(), error, sizeof(error)); + if (!model) { + mju_error("Loading error: %s\n", error); + } + return std::unique_ptr(new MjModel(model)); +} + void mj_saveModel_wrapper(const MjModel& m, const StringOrNull& filename, const val& buffer) { UNPACK_NULLABLE_STRING(filename); UNPACK_NULLABLE_VALUE(uint8_t, buffer); diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index 8fe9afe0..0ed482cc 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -2512,7 +2512,7 @@ describe('MuJoCo WASM Bindings', () => { vfs = new mujoco.MjVFS(); vfs.addBuffer(objFilename, new TextEncoder().encode(cube1)); - binaryModel = mujoco.MjModel.mj_loadBinary(mjbFilename, vfs); + binaryModel = mujoco.MjModel.mj_loadModel(mjbFilename, vfs); assertExists(binaryModel); expect(mujoco.mj_sizeModel(binaryModel)) @@ -2542,7 +2542,7 @@ describe('MuJoCo WASM Bindings', () => { const bufSize = mujoco.mj_sizeModel(model!); vfs = new mujoco.MjVFS(); - binaryModel = mujoco.MjModel.mj_loadBinary(mjbFilename, vfs); + binaryModel = mujoco.MjModel.mj_loadModel(mjbFilename, vfs); assertExists(binaryModel); expect(mujoco.mj_sizeModel(binaryModel)).toEqual(bufSize); @@ -2558,4 +2558,49 @@ describe('MuJoCo WASM Bindings', () => { } }); + it('should load XML with assets from VFS', () => { + const xml = ` + + + + + + + + `; + + const cube1 = ` + v -1 -1 1 + v 1 -1 1 + v -1 1 1 + v 1 1 1 + v -1 1 -1 + v 1 1 -1 + v -1 -1 -1 + v 1 -1 -1`; + + const xmlFilename = '/tmp/with_vfs.xml'; + writeXMLFile(xmlFilename, xml); + + let model: MjModel|null = null; + let vfs: MjVFS|null = null; + try { + vfs = new mujoco.MjVFS(); + vfs.addBuffer('cube.obj', new TextEncoder().encode(cube1)); + assertExists(vfs); + + model = mujoco.MjModel.mj_loadXML(xmlFilename, vfs); + assertExists(model); + expect(model.nmesh).toBe(1); + + const meshId = + mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH.value, 'cube'); + expect(meshId).toBeGreaterThanOrEqual(0); + } finally { + model?.delete(); + vfs?.delete(); + unlinkXMLFile(xmlFilename); + } + }); + }); From 40e000784940f550d5a53c932837a952dceb29ad Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 25 Feb 2026 00:53:38 -0800 Subject: [PATCH 38/48] Restructure numeric constants table in API docs Break the monolithic numeric constants table into 6 sub-grouped tables: Version, Engine constants, Array sizes, Visualization, Rendering, UI constants. Fix stale values: mjMAXUITEXT (500->300), mjMAXUIMULTI (20->35), mjMAXUIEDIT (5->7), mjMAXUIRECT (15->25). Add missing constants: mjMAXLIGHT, mjMAXMATERIAL. Fix placeholder description for mjMAXFLEXNODES. PiperOrigin-RevId: 875010581 Change-Id: Id47846de4d98b62cb96087b72ecd28358d5db9c7 --- doc/APIreference/APIglobals.rst | 158 ++++++++++++++++++------ doc/APIreference/functions.rst | 4 +- doc/APIreference/functions_override.rst | 4 +- doc/XMLreference.rst | 4 +- doc/changelog.rst | 4 +- doc/computation/index.rst | 14 +-- doc/mjwarp/index.rst | 2 +- doc/modeling.rst | 2 +- doc/programming/index.rst | 2 +- doc/programming/simulation.rst | 4 +- 10 files changed, 140 insertions(+), 58 deletions(-) diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 2292549e..656c6a07 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -368,10 +368,41 @@ Numeric constants ^^^^^^^^^^^^^^^^^ Many integer constants were already documented in the primitive types above. In addition, the header files define -several other constants documented here. Unless indicated otherwise, each entry in the table below is defined in -`mjmodel.h `_. Note that some extended key -codes are defined in `mjui.h `_ which are not -shown in the table below. Their names are in the format ``mjKEY_XXX``. They correspond to GLFW key codes. +several other constants documented here. Note that some extended key codes are defined in +`mjui.h `_ which are not shown below. Their +names are in the format ``mjKEY_XXX``. They correspond to GLFW key codes. + + +.. _glNumericVersion: + +Version +~~~~~~~ + +Defined in `mujoco.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description + * - ``mjVERSION_HEADER`` + - 3005001 + - The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P" + using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001. + The API function :ref:`mj_version` returns a number with the same meaning + but for the compiled library. See + `VERSIONING.md `__ for details. + + +.. _glNumericEngine: + +Engine constants +~~~~~~~~~~~~~~~~ + +Defined in `mjmodel.h `_ unless +indicated otherwise. .. list-table:: :widths: 2 1 8 @@ -384,6 +415,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - 1E-15 - The minimal value allowed in any denominator, and in general any mathematical operation where 0 is not allowed. In almost all cases, MuJoCo silently clamps smaller values to mjMINVAL. + Defined in `mjtnum.h `_. * - ``mjPI`` - :math:`\pi` - The value of :math:`\pi`. This is used in various trigonometric functions, and also for conversion from degrees @@ -416,17 +448,38 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr is raised and ray casting may not be possible. For a balanced hierarchy, this implies 1E15 bounding volumes. * - ``mjMAXFLEXNODES`` - 27 - - Some number by Alessio that needs documentation. I guess it's related to trilinear flexes? + - The maximum number of nodes in a trilinear flex element. * - ``mjMINAWAKE`` - 10 - The minimum number of timesteps that must pass after a tree is awoken, before it is allowed to go back to sleep. + * - ``mjMAXTHREAD`` + - 128 + - Maximum number of OS threads that can be used in a thread pool. + Defined in `mjthread.h `_. + + +.. _glNumericSizes: + +Array sizes +~~~~~~~~~~~ + +Defined in `mjmodel.h `_. These constants +correspond to array sizes which we have not fully settled. There may be reasons to increase them in the future, so as to +accommodate extra parameters needed for more elaborate computations. This is why we maintain them as symbolic constants +that can be easily changed, as opposed to the array size for representing quaternions for example -- which has no reason +to change. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNEQDATA`` - 11 - The maximal number of real-valued parameters used to define each equality constraint. Determines the size of - ``mjModel.eq_data``. This and the next five constants correspond to array sizes which we have not fully settled. - There may be reasons to increase them in the future, so as to accommodate extra parameters needed for more - elaborate computations. This is why we maintain them as symbolic constants that can be easily changed, as opposed - to the array size for representing quaternions for example -- which has no reason to change. + ``mjModel.eq_data``. * - ``mjNDYN`` - 10 - The maximal number of real-valued parameters used to define the activation dynamics of each actuator. @@ -464,75 +517,104 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - The number of islands for which solver statistics can be stored in ``mjData.solver``. This array is used to store diagnostic information about each iteration of the constraint solver. The actual number of islands for which the solver was run is given by ``mjData.nsolver_island``. + + +.. _glNumericVisualization: + +Visualization +~~~~~~~~~~~~~ + +Defined in `mjvisualize.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNGROUP`` - 6 - The number of geom, site, joint, tendon and actuator groups whose rendering can be enabled and disabled via :ref:`mjvOption`. - Defined in `mjvisualize.h `_. + * - ``mjMAXLIGHT`` + - 100 + - The maximum number of lights in a scene. * - ``mjMAXOVERLAY`` - 500 - The maximal number of characters in overlay text for rendering. - Defined in `mjvisualize.h `_. * - ``mjMAXLINE`` - 100 - The maximal number of lines per 2D figure (:ref:`mjvFigure`). - Defined in `mjvisualize.h `_. * - ``mjMAXLINEPNT`` - 1001 - The maximal number of points in each line in a 2D figure. Note that the buffer ``mjvFigure.linepnt`` has length ``2*mjMAXLINEPNT`` because each point has X and Y coordinates. - Defined in `mjvisualize.h `_. * - ``mjMAXPLANEGRID`` - 200 - The maximal number of grid lines in each dimension for rendering planes. - Defined in `mjvisualize.h `_. + + +.. _glNumericRendering: + +Rendering +~~~~~~~~~ + +Defined in `mjrender.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNAUX`` - 10 - Number of auxiliary buffers that can be allocated in mjrContext. - Defined in `mjrender.h `_. * - ``mjMAXTEXTURE`` - 1000 - Maximum number of textures allowed. - Defined in `mjrender.h `_. - * - ``mjMAXTHREAD`` - - 128 - - Maximum number OS threads that can be used in a thread pool. - Defined in `mjthread.h `_. + * - ``mjMAXMATERIAL`` + - 1000 + - Maximum number of materials with textures. + + +.. _glNumericUI: + +UI constants +~~~~~~~~~~~~ + +Defined in `mjui.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjMAXUISECT`` - 10 - Maximum number of UI sections. - Defined in `mjui.h `_. * - ``mjMAXUIITEM`` - 200 - Maximum number of items per UI section. - Defined in `mjui.h `_. * - ``mjMAXUITEXT`` - - 500 + - 300 - Maximum number of characters in UI fields 'edittext' and 'other'. - Defined in `mjui.h `_. * - ``mjMAXUINAME`` - 40 - Maximum number of characters in any UI name. - Defined in `mjui.h `_. * - ``mjMAXUIMULTI`` - - 20 + - 35 - Maximum number of radio and select items in UI group. - Defined in `mjui.h `_. * - ``mjMAXUIEDIT`` - - 5 + - 7 - Maximum number of elements in UI edit list. - Defined in `mjui.h `_. * - ``mjMAXUIRECT`` - - 15 + - 25 - Maximum number of UI rectangles. - Defined in `mjui.h `_. - * - ``mjVERSION_HEADER`` - - 3005001 - - The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P" - using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001. - Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning - but for the compiled library. See - `VERSIONING.md `__ for details. .. _Macros: diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 5cbb95c5..bec885ae 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -1252,7 +1252,7 @@ Intersect ray ``pnt+x*vec, x >= 0`` with geoms. - If ``normal`` is not NULL, write the surface normal at the intersection point. The normal always points **out of the geometry**, regardless of the ray's direction (i.e., including rays hitting the surface from the inside). - Exclude geoms in body with id ``bodyexclude``, use -1 to include all bodies. -- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass +- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass NULL to skip geom group exclusion. - If ``flg_static`` is 0, static geoms will be excluded. @@ -2911,7 +2911,7 @@ each corresponding to one item. The last (unused) item has its type set to -1, t after the end of the last used section. There is also another version of this function (:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by -:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. .. _mjui_addToSection: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 1e0d5bbf..88e90085 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -403,7 +403,7 @@ Intersect ray ``pnt+x*vec, x >= 0`` with geoms. - If ``normal`` is not NULL, write the surface normal at the intersection point. The normal always points **out of the geometry**, regardless of the ray's direction (i.e., including rays hitting the surface from the inside). - Exclude geoms in body with id ``bodyexclude``, use -1 to include all bodies. -- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass +- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass NULL to skip geom group exclusion. - If ``flg_static`` is 0, static geoms will be excluded. @@ -444,7 +444,7 @@ each corresponding to one item. The last (unused) item has its type set to -1, t after the end of the last used section. There is also another version of this function (:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by -:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. .. _mjui_update: diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 2e12173c..5b5392a9 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -1249,7 +1249,7 @@ The full list of processing steps applied by the compiler to each mesh is as fol transformations in ``mjModel.mesh_{pos, quat, scale}``. #. Construct the convex hull if specified; #. Find the centroid of all triangle faces, and construct the union-of-pyramids representation. Triangles whose area is - too small (below the :ref:`mjMINVAL ` value of 1E-14) result in compile error; + too small (below the :ref:`mjMINVAL ` value of 1E-14) result in compile error; #. Compute the center of mass and inertia matrix of the union-of-pyramids. Use eigenvalue decomposition to find the principal axes of inertia. Center and align the mesh, saving the translational and rotational offsets for subsequent geom-related computations. @@ -1518,7 +1518,7 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe and other geoms (except for planes and other height fields which are not supported) are computed by first selecting the sub-grid of prisms that could collide with the geom based on its bounding box, and then using the general convex collider. The number of possible contacts between a height field and a geom is limited to 50 - (:ref:`mjMAXCONPAIR `); any contacts beyond that are discarded. To avoid penetration due to discarded + (:ref:`mjMAXCONPAIR `); any contacts beyond that are discarded. To avoid penetration due to discarded contacts, the spatial features of the height field should be large compared to the geoms it collides with. .. _asset-hfield-name: diff --git a/doc/changelog.rst b/doc/changelog.rst index 21552163..74178e54 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -1463,8 +1463,8 @@ General Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. If the solver does not use islands, only row 0 is filled. - - The new constant :ref:`mjNISLAND` was set to 20. - - :ref:`mjNSOLVER` was reduced from 1000 to 200. + - The new constant :ref:`mjNISLAND` was set to 20. + - :ref:`mjNSOLVER` was reduced from 1000 to 200. - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer vectors of length ``mjNISLAND``. diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 422e4860..537d9884 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -1621,17 +1621,17 @@ will generate up to 1 contact or with ``multiccd`` up to 4 contacts. | **1** * - HField - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldSDF | :ref:`sdf_initpoints ` * - Sphere @@ -1762,7 +1762,7 @@ sleeping mechanism is provided in the :ref:`Simulation chapter` but her Sleeping can occur in one of two ways: - **Automatic:** A tree whose maximum velocity in absolute value is less than the - :ref:`tolerance ` for :ref:`mjMINAWAKE ` time steps is marked as "ready to sleep". + :ref:`tolerance ` for :ref:`mjMINAWAKE ` time steps is marked as "ready to sleep". If all trees in an island are ready to sleep, they are put to sleep during state advancement. - **Initialized asleep:** By setting the :ref:`body/sleep` attribute of a tree root to "init", it is marked as "initialized-asleep" and put to sleep during :ref:`mjData` initialization. diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index ed5af221..591061cc 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -724,7 +724,7 @@ Warnings are provided when memory requirements exceed existing allocations durin setting `m.opt.contact_sensor_maxmatch`. Alternatively, refactor the contact sensor matching criteria, for example if the 2 geoms of interest are known, specify ``geom1`` and ``geom2``. - ``height field collision overflow``: The number of potential contacts generated by a height field exceeds - :ref:`mjMAXCONPAIR ` and some contacts are ignored. To resolve this warning, reduce the height field + :ref:`mjMAXCONPAIR ` and some contacts are ignored. To resolve this warning, reduce the height field resolution or reduce the size of the geom interacting with the height field. Compilation diff --git a/doc/modeling.rst b/doc/modeling.rst index fc5979fd..127bc1a0 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -301,7 +301,7 @@ We begin by explaining the constraint impedance :math:`d`. at rest. Impedance is set using the :at:`solimp` attribute. Recall that :math:`d` must lie between 0 and 1; internally MuJoCo clamps it to the range [:ref:`mjMINIMP mjMAXIMP -`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced +`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced acceleration :math:`\au` and reference acceleration :math:`\ar`. The user can set :math:`d` to a constant, or take advantage of its interpolating property and make it position-dependent, i.e., a function of the constraint violation :math:`r`. Position-dependent impedance can be used to model soft contact layers around objects, or define diff --git a/doc/programming/index.rst b/doc/programming/index.rst index 7d91d013..5d246f78 100644 --- a/doc/programming/index.rst +++ b/doc/programming/index.rst @@ -200,7 +200,7 @@ The situation is more subtle if existing code was developed with a certain versi compiled and linked with a different version. If the definitions of the API functions used in that code have changed, either the compiler or the linker will generate errors. But even if the function definitions have not changed, it may still be a good idea to assert that the software version is the same. To this end, the main header (mujoco.h) defines -the symbol :ref:`mjVERSION_HEADER ` and the library provides the function +the symbol :ref:`mjVERSION_HEADER ` and the library provides the function :ref:`mj_version`. Thus the header and library versions can be compared with: .. code-block:: C diff --git a/doc/programming/simulation.rst b/doc/programming/simulation.rst index da6c2b16..c172b785 100644 --- a/doc/programming/simulation.rst +++ b/doc/programming/simulation.rst @@ -1098,7 +1098,7 @@ implementation details. The high level sleep state of :ref:`trees` is described by ``mjData.tree_asleep`` (though see caveat below). A negative value means a tree is awake, non-negative means asleep. Maximally awake trees are given the value - |-| (1 |-| -+ |-| :ref:`mjMINAWAKE`), and for every timestep where their velocity falls below the sleep :ref:`tolerance ++ |-| :ref:`mjMINAWAKE`), and for every timestep where their velocity falls below the sleep :ref:`tolerance `, this integer is incremented, up to -1, which means "ready to sleep". If all trees in an island are ready to sleep, they are put to sleep during state advancement and their associated values in ``tree_asleep`` are set to a (non-negative) index cycle: the "sleeping island". If any tree in the island is woken, all are woken. @@ -1213,7 +1213,7 @@ Notes **Provisional choices** Some implementation choices are provisional and subject to change. - A concrete example is the decision to hard-code the value of :ref:`mjMINAWAKE` instead of exposing it to + A concrete example is the decision to hard-code the value of :ref:`mjMINAWAKE` instead of exposing it to the user as a runtime option. This was done for two reasons. First, in our experiments, we've found that changing this value is equivalent to changing the :ref:`sleep_tolerance`, which is the more useful knob. Second, one could argue for a time-to-sleep semantic that is in units of time rather than an integer number of From 27dc92b0567375d4d7466228e11cba49806cabe7 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 25 Feb 2026 02:50:53 -0800 Subject: [PATCH 39/48] Fix typo in modeling.rst. PiperOrigin-RevId: 875053694 Change-Id: Idb6d09d2713f2cc2a09d2b0912a227fe9ac9396a --- doc/modeling.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/modeling.rst b/doc/modeling.rst index 127bc1a0..cf9b056c 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -1712,7 +1712,7 @@ dedicated section :ref:`therein`. 1. :ref:`Timestep`: Try to increase the simulation timestep. As explained at the end of the :ref:`Numerical Integration` section, the timestep is the single most important parameter in any model. The default value is chosen for stability rather than efficiency, and can often be increased. At some point, - increasing it further will cause diveregence, so the optimal timestep is the largest timestep at which divergence + increasing it further will cause divergence, so the optimal timestep is the largest timestep at which divergence never happens or is very rare. The actual value is model-dependent. 2. :ref:`Integrator`: Choose your integrator according to the recommendations at the end of the :ref:`Numerical Integration` section. The default recommended choice is the ``implicitfast`` From ae7efead09eca112efd5d913d40ae8b83ead952a Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Wed, 25 Feb 2026 03:10:42 -0800 Subject: [PATCH 40/48] Show MuJoCo version in the help menu. PiperOrigin-RevId: 875061430 Change-Id: I4d4ce32eebd8ca4af4622e9b4d1cd50d180d0b96 --- src/experimental/studio/app.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index bade8680..6d33f5af 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -1777,6 +1777,9 @@ void App::MainMenuGui() { if (ImGui::MenuItem("ImPlot Demo")) { tmp_.implot_demo = !tmp_.implot_demo; } + ImGui::Separator(); + std::string version = "Version " + std::string(mj_versionString()); + ImGui::MenuItem(version.c_str()); ImGui::EndMenu(); } ImGui::EndMainMenuBar(); From d0fc1c2c8327eeebeb6d96b9be6d974cd3e38f99 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 25 Feb 2026 03:16:02 -0800 Subject: [PATCH 41/48] MuJoCo Warp documentation: SDF plugins PiperOrigin-RevId: 875063400 Change-Id: Iefa014e3c046a965722c633df0b19771ab481375 --- doc/mjwarp/index.rst | 79 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 78 insertions(+), 1 deletion(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 591061cc..4bf4cffc 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -489,7 +489,7 @@ Key features: `Warp's BVHs `__. Basic Usage ------------ +---------- Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures, rendering specific fields, and output buffers. @@ -825,3 +825,80 @@ Additional MJWarp-only options are available: A new :ref:`graph capture ` may be necessary after modifying an :class:`mjw.Option ` field in order for the updated setting to take effect. + +SDF plugins +----------- + +SDF collisions support plugins. The following example for +`plugin/sdf/bowl.xml `__ illustrates how to implement the +SDF plugin implementation in `bowl.cc `__: + +.. code-block:: python + + import mujoco_warp as mjw + import warp as wp + + # distance function + @wp.func + def bowl(p: wp.vec3, attr: wp.vec3) -> float: + """Signed distance function for a bowl shape. + + attr[0] = height + attr[1] = radius + attr[2] = thickness + """ + height = attr[0] + radius = attr[1] + thick = attr[2] + width = wp.sqrt(radius * radius - height * height) + + # q = (norm_xy(p), p.z) + q0 = wp.sqrt(p[0] * p[0] + p[1] * p[1]) + q1 = p[2] + + # qdiff = q - (width, height) + qdiff0 = q0 - width + qdiff1 = q1 - height + + if height * q0 < width * q1: + dist = wp.sqrt(qdiff0 * qdiff0 + qdiff1 * qdiff1) + else: + q_norm = wp.sqrt(q0 * q0 + q1 * q1) + dist = wp.abs(q_norm - radius) + + return dist - thick + + + # gradient of distance function + @wp.func + def bowl_sdf_grad(p: wp.vec3, attr: wp.vec3) -> wp.vec3: + """Gradient of bowl SDF via finite differences.""" + eps = float(1e-6) + f0 = bowl(p, attr) + + px = wp.vec3(p[0] + eps, p[1], p[2]) + py = wp.vec3(p[0], p[1] + eps, p[2]) + pz = wp.vec3(p[0], p[1], p[2] + eps) + + grad = wp.vec3( + (bowl(px, attr) - f0) / eps, + (bowl(py, attr) - f0) / eps, + (bowl(pz, attr) - f0) / eps, + ) + return grad + + + # register the bowl SDF plugin + @wp.func + def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float: + return bowl(p, attr) + + + @wp.func + def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: + return bowl_sdf_grad(p, attr) + + + # override the module-level hooks + mjw._src.collision_sdf.user_sdf = user_sdf + mjw._src.collision_sdf.user_sdf_grad = user_sdf_grad From cd8bfb90c3eff2b40904a6ed856735a2f761d8c6 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Wed, 25 Feb 2026 03:39:37 -0800 Subject: [PATCH 42/48] Improve polygon creation logic in compiler. PiperOrigin-RevId: 875070301 Change-Id: I06e851f3450035fff68b99997406d6ea471ce3b8 --- src/user/user_mesh.cc | 134 +++++++++++++---------- test/engine/engine_collision_gjk_test.cc | 2 +- test/fixture.cc | 16 +-- 3 files changed, 82 insertions(+), 70 deletions(-) diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 5df01943..c1e1a027 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -2691,17 +2691,17 @@ class MeshPolygon { public: // constructors (need starting face) MeshPolygon(const double v1[3], const double v2[3], const double v3[3], - int v1i, int v2i, int v3i); + int v1i, int v2i, int v3i, double theta, double phi); MeshPolygon() = delete; MeshPolygon(const MeshPolygon&) = delete; MeshPolygon& operator=(const MeshPolygon&) = delete; + MeshPolygon(MeshPolygon&&) = default; + MeshPolygon& operator=(MeshPolygon&&) = default; - void InsertFace(int v1, int v2, int v3); // insert a face into the polygon - std::vector> Paths() const; // return trace of the polygons - const double* Normal() const { return normal_; } // return the normal of the polygon - - // return the ith component of the normal of the polygon - double Normal(int i) const { return normal_[i]; } + void InsertFace(int v1, int v2, int v3); // insert a face into the polygon + std::vector> Paths() const; // return trace of the polygons + const double* Normal() const { return normal_; } // return the normal of the polygon + double Normal(int i) const { return normal_[i]; } // return the i-th component of the normal private: std::vector> edges_; @@ -2714,40 +2714,46 @@ class MeshPolygon { void CombineIslands(int& island1, int& island2); }; +bool MeshPolygonKey(std::pair& angles, const double v1[3], const double v2[3], + const double v3[3], double angle_tol) { + double diff12[3] = {v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2]}; + double diff13[3] = {v3[0] - v1[0], v3[1] - v1[1], v3[2] - v1[2]}; + double normal[3], norm; - -MeshPolygon::MeshPolygon(const double v1[3], const double v2[3], const double v3[3], - int v1i, int v2i, int v3i) { - mjuu_makenormal(normal_, v1, v2, v3); - edges_ = {{v1i, v2i}, {v2i, v3i}, {v3i, v1i}}; - nisland_ = 1; - islands_ = {0, 0, 0}; -} - - - -// comparison operator for std::set -bool PolygonCmp(const MeshPolygon& p1, const MeshPolygon& p2) { - const double* n1 = p1.Normal(); - const double* n2 = p2.Normal(); - double dot3 = n1[0] * n2[0] + n1[1] * n2[1] + n1[2] * n2[2]; - - // TODO(kylebayes): The tolerance should be a parameter set the user, as it should be optimized - // from mesh to mesh. - if (dot3 > 0.99999872) { + mjuu_crossvec(normal, diff12, diff13); + if ((norm = std::sqrt(mjuu_dot3(normal, normal))) < mjMINVAL) { return false; } - if (std::abs(n1[0] - n2[0]) > mjMINVAL) { - return n1[0] > n2[0]; + // atan2 is sensitive to sign of 0.0, adding 0.0 to enforcing only positive 0.0 + normal[0] = (normal[0] / norm) + 0.0; + normal[1] = (normal[1] / norm) + 0.0; + normal[2] = (normal[2] / norm) + 0.0; + double rtheta = 0.0, rphi = 0.0; + + // clamp normal to be in valid range for acos + if (std::abs(normal[2]) > 1.0 - 1e-7) { + if (normal[2] < 0) rphi = std::round(mjPI / angle_tol); + angles = std::make_pair(rtheta, rphi); + return true; } - if (std::abs(n1[1] - n2[1]) > mjMINVAL) { - return n1[1] > n2[1]; - } - if (std::abs(n1[2] - n2[2]) > mjMINVAL) { - return n1[2] > n2[2]; - } - return false; + // rounded azimuthal and polar angles + rtheta = std::round(std::atan2(normal[1], normal[0]) / angle_tol); + rphi = std::round(std::acos(normal[2]) / angle_tol); + angles = std::make_pair(rtheta, rphi); + return true; +} + + +MeshPolygon::MeshPolygon(const double v1[3], const double v2[3], const double v3[3], + int v1i, int v2i, int v3i, double theta, double phi) { + normal_[0] = std::cos(theta) * std::sin(phi); + normal_[1] = std::sin(theta) * std::sin(phi); + normal_[2] = std::cos(phi); + + edges_ = {{v1i, v2i}, {v2i, v3i}, {v3i, v1i}}; + nisland_ = 1; + islands_ = {0, 0, 0}; } @@ -2846,8 +2852,8 @@ void MeshPolygon::InsertFace(int v1, int v2, int v3) { // return the transverse vertices of the polygon, multiple paths possible if not connected -std::vector > MeshPolygon::Paths() const { - std::vector > paths; +std::vector> MeshPolygon::Paths() const { + std::vector> paths; // shortcut if polygon is just a triangular face if (edges_.size() == 3) { return {{edges_[0].first, edges_[1].first, edges_[2].first}}; @@ -2899,9 +2905,20 @@ std::vector > MeshPolygon::Paths() const { +// hash function for std::pair +struct PairHash { + template + std::size_t operator() (const std::pair& pair) const { + return std::hash()(pair.first) ^ std::hash()(pair.second); + } +}; + + + // merge coplanar mesh triangular faces into polygonal sides to represent the geometry of the mesh void mjCMesh::MakePolygons() { - std::set polygons(PolygonCmp); + constexpr double kAngleTol = 0.01; + std::unordered_map, MeshPolygon, PairHash> mesh_polygons; polygons_.clear(); polygon_normals_.clear(); polygon_map_.clear(); @@ -2923,22 +2940,30 @@ void mjCMesh::MakePolygons() { // process each face for (int i = 0; i < nfaces; i++) { - double* v1 = &vert_[3*faces[3*i + 0]]; - double* v2 = &vert_[3*faces[3*i + 1]]; - double* v3 = &vert_[3*faces[3*i + 2]]; + int vi1 = faces[3*i + 0]; + int vi2 = faces[3*i + 1]; + int vi3 = faces[3*i + 2]; + double* v1 = &vert_[3*vi1]; + double* v2 = &vert_[3*vi2]; + double* v3 = &vert_[3*vi3]; - MeshPolygon face(v1, v2, v3, faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); - auto it = polygons.find(face); - if (it == polygons.end()) { - polygons.emplace(v1, v2, v3, faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); + std::pair key; + if (!MeshPolygonKey(key, v1, v2, v3, kAngleTol)) { + continue; + } + auto it = mesh_polygons.find(key); + if (it == mesh_polygons.end()) { + double theta = kAngleTol * key.first; + double phi = kAngleTol * key.second; + mesh_polygons.emplace(key, MeshPolygon(v1, v2, v3, vi1, vi2, vi3, theta, phi)); } else { - MeshPolygon& p = const_cast(*it); - p.InsertFace(faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); + it->second.InsertFace(vi1, vi2, vi3); } } - for (const auto& polygon : polygons) { - std::vector > paths = polygon.Paths(); + for (const auto& pair : mesh_polygons) { + const MeshPolygon& polygon = pair.second; + std::vector> paths = polygon.Paths(); // separate the polygons if they were grouped together for (const auto& path : paths) { @@ -3388,15 +3413,6 @@ void mjCSkin::LoadSKN(mjResource* resource) { //-------------------------- nonlinear elasticity -------------------------------------------------- -// hash function for std::pair -struct PairHash -{ - template - std::size_t operator() (const std::pair& pair) const { - return std::hash()(pair.first) ^ std::hash()(pair.second); - } -}; - // simplex connectivity constexpr int eledge[3][6][2] = {{{ 0, 1}, {-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}}, diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc index d06a5b02..cfcba486 100644 --- a/test/engine/engine_collision_gjk_test.cc +++ b/test/engine/engine_collision_gjk_test.cc @@ -1373,7 +1373,7 @@ TEST_F(MjGjkTest, BoxMesh) { std::vector dir, pos; mjtNum dist; int ncons = Penetration(status, dist, dir, pos, model, data, g2, g1, 0, 1000); - + EXPECT_EQ(model->nmeshpoly, 7); EXPECT_EQ(ncons, 4); mj_deleteData(data); mj_deleteModel(model); diff --git a/test/fixture.cc b/test/fixture.cc index 4b18832d..f621330e 100644 --- a/test/fixture.cc +++ b/test/fixture.cc @@ -238,16 +238,12 @@ mjtNum CompareModel(const mjModel* m1, const mjModel* m2, MJMODEL_POINTERS_PREAMBLE(m1); // compare ints, exclude nbuffer because it hides the actual difference -// TODO(kylebayes): re-enable poly comparisons. -#define X(name) \ - if constexpr (std::string_view(#name) != "nbuffer" && \ - std::string_view(#name) != "nmeshpolymap" && \ - std::string_view(#name) != "nmeshpolyvert" && \ - std::string_view(#name) != "nmeshpoly") { \ - if (m1->name != m2->name) { \ - maxdif = std::abs((long)m1->name - (long)m2->name); \ - field = #name; \ - } \ +#define X(name) \ + if constexpr (std::string_view(#name) != "nbuffer") { \ + if (m1->name != m2->name) { \ + maxdif = std::abs((long)m1->name - (long)m2->name); \ + field = #name; \ + } \ } MJMODEL_SIZES #undef X From 84205f7dd066b1f6a8c49f32ff38d191bc0b7cb4 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Wed, 25 Feb 2026 03:52:48 -0800 Subject: [PATCH 43/48] Isolate libccd logic in mjc_penetration. PiperOrigin-RevId: 875073780 Change-Id: I71bd39695aa667110fffff47f34cd98faa562cc6 --- src/engine/engine_collision_convex.c | 203 +++++++++++++-------------- 1 file changed, 97 insertions(+), 106 deletions(-) diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c index c298b6d9..1339128d 100644 --- a/src/engine/engine_collision_convex.c +++ b/src/engine/engine_collision_convex.c @@ -45,44 +45,76 @@ static void ccd_free(void* data, void* buffer) { mj_freeStack((mjData*)data); } -// call libccd or nativeccd to recover penetration info -static int mjc_penetration(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, - const ccd_t* ccd, ccd_real_t* depth, ccd_vec3_t* dir, ccd_vec3_t* pos) { - // fallback to MPR - if (mjDISABLED(mjDSBL_NATIVECCD)) { - return ccdMPRPenetration(obj1, obj2, ccd, depth, dir, pos); +// ccd prism first dir +static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) { + ccdVec3Set(vec, 0, 0, 1); +} + +// wrapper around libccd; returns number of collisions found +static int libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum* dist, + mjtNum dir[3], mjtNum pos[3]) { + ccd_t ccd; + CCD_INIT(&ccd); + ccd.mpr_tolerance = m->opt.ccd_tolerance; + ccd.epa_tolerance = m->opt.ccd_tolerance; // use MPR tolerance for EPA + ccd.max_iterations = m->opt.ccd_iterations; + ccd.support1 = mjccd_support; + ccd.support2 = mjccd_support; + ccd.center1 = mjccd_center; + ccd.center2 = mjccd_center; + if (obj1->geom_type == mjGEOM_HFIELD) { + ccd.first_dir = prism_firstdir; + } + if (obj2->geom_type == mjGEOM_HFIELD) { + ccd.first_dir = prism_firstdir; } + ccd_real_t ccd_depth; + ccd_vec3_t ccd_dir, ccd_pos; + int ret = ccdMPRPenetration(obj1, obj2, &ccd, &ccd_depth, &ccd_dir, &ccd_pos); + *dist = -ccd_depth; + mji_copy3(dir, ccd_dir.v); + mji_copy3(pos, ccd_pos.v); + if (ret == 0 && dir[0] == 0 && dir[1] == 0 && dir[2] == 0) { + return 0; + } + return ret == 0; +} + + +// find penetration info between two geoms; returns number of collisions found +static int mjc_penetration(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum* dist, + mjtNum dir[3], mjtNum pos[3]) { + if (mjDISABLED(mjDSBL_NATIVECCD)) { + return libccd_wrapper(m, obj1, obj2, dist, dir, pos); + } + + // nativeccd mjCCDConfig config; mjCCDStatus status; + mjtNum d; // distance returned by mjc_ccd // set config - config.max_iterations = ccd->max_iterations; - config.tolerance = ccd->mpr_tolerance; + config.max_iterations = m->opt.ccd_iterations; + config.tolerance = m->opt.ccd_tolerance; config.max_contacts = 1; config.dist_cutoff = 0; // no geom distances needed config.context = (void*)obj1->data; config.alloc = ccd_allocate; config.free = ccd_free; - mjtNum dist = mjc_ccd(&config, &status, obj1, obj2); - if (dist < 0) { - if (depth) *depth = -dist; - if (dir) { - mju_sub3(dir->v, status.x1, status.x2); - mju_normalize3(dir->v); - } - if (pos) { - pos->v[0] = 0.5 * (status.x1[0] + status.x2[0]); - pos->v[1] = 0.5 * (status.x1[1] + status.x2[1]); - pos->v[2] = 0.5 * (status.x1[2] + status.x2[2]); - } - return 0; + if ((d = mjc_ccd(&config, &status, obj1, obj2)) < 0) { + *dist = d; + + mju_sub3(dir, status.x1, status.x2); + mju_normalize3(dir); + + pos[0] = 0.5 * (status.x1[0] + status.x2[0]); + pos[1] = 0.5 * (status.x1[1] + status.x2[1]); + pos[2] = 0.5 * (status.x1[2] + status.x2[2]); + return 1; } - if (depth) *depth = 0; - if (dir) mju_zero3(dir->v); - if (pos) mju_zero3(dir->v); - return 1; + return 0; } @@ -98,39 +130,34 @@ void mjc_center(mjtNum res[3], const mjCCDObj *obj) { int e = obj->elem; int v = obj->vert; + if (obj->geom_type == mjGEOM_HFIELD) { + mju_zero3(res); + for (int i=0; i < 6; i++) { + mji_addTo3(res, obj->prism[i]); + } + mju_scl3(res, res, 1.0/6.0); + return; + } + // return geom position if (g >= 0) { mji_copy3(res, obj->data->geom_xpos + 3*g); + return; } // return flex element position - else if (e >= 0) { + if (e >= 0) { mji_copy3(res, obj->data->flexelem_aabb + 6*(obj->model->flex_elemadr[f]+e)); + return; } // return flex vertex position - else { + if (f >= 0) { mji_copy3(res, obj->data->flexvert_xpos + 3*(obj->model->flex_vertadr[f]+v)); + return; } } - -// prism center function -static void mjc_prism_center(mjtNum res[3], const mjCCDObj* obj) { - // compute mean - mju_zero3(res); - for (int i=0; i < 6; i++) { - mji_addTo3(res, obj->prism[i]); - } - mju_scl3(res, res, 1.0/6.0); -} - - -// ccd prism center function -static void mjccd_prism_center(const void *obj, ccd_vec3_t *center) { - mjc_prism_center(center->v, (const mjCCDObj*) obj); -} - // ------------------------------------ Support functions ----------------------------------------- // transform a vector from global to local frame @@ -681,6 +708,10 @@ void mjccd_support(const void *_obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) { } break; + case mjGEOM_HFIELD: + mjc_prism_support(res, obj, dir); + return; + default: mjERROR("ccd support function is undefined for geom type %d", m->geom_type[g]); } @@ -697,12 +728,6 @@ void mjccd_support(const void *_obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) { mji_addTo3(res, d->geom_xpos+3*g); } - -// libccd prism support function -static void mjccd_prism_support(const void *obj, const ccd_vec3_t *dir, ccd_vec3_t *vec) { - mjc_prism_support(vec->v, (mjCCDObj*) obj, dir->v); -} - // ------------------------------------------------------------------------------------------------ // initialize a CCD object @@ -747,7 +772,7 @@ void mjc_initCCDObj(mjCCDObj* obj, const mjModel* m, const mjData* d, int g, mjt obj->support = mjc_boxSupport; break; case mjGEOM_HFIELD: - obj->center = mjc_prism_center; + obj->center = mjc_center; obj->support = mjc_prism_support; int hid = m->geom_dataid[g]; @@ -1140,11 +1165,6 @@ int mjc_PlaneConvex(const mjModel* m, const mjData* d, //---------------------------- heightfield collisions --------------------------------------------- -// ccd prism first dir -static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) { - ccdVec3Set(vec, 0, 0, 1); -} - // add vertex to prism static inline void addVert(mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) { @@ -1289,16 +1309,6 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d, rmin = mjMAX(0, rmin); rmax = mjMIN(nrow-1, rmax); - // CCD collision testing - - ccd_t ccd; - mjc_initCCD(&ccd, m); - ccd.first_dir = prism_firstdir; - ccd.center1 = mjccd_prism_center; - ccd.center2 = mjccd_center; - ccd.support1 = mjccd_prism_support; - ccd.support2 = mjccd_support; - // geom margin needed for actual collision test obj2.margin = margin; @@ -1325,14 +1335,12 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d, } // run penetration function, save contact - ccd_vec3_t dirccd, vecccd; - ccd_real_t depth; - if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0 - && !ccdVec3Eq(&dirccd, ccd_vec3_origin)) { + mjtNum dist; + if (mjc_penetration(m, &obj1, &obj2, &dist, dir, pos)) { // fill in contact data, transform to global coordinates - con[ncon].dist = -depth; - mji_mulMatVec3(con[ncon].frame, mat1, dirccd.v); - mji_mulMatVec3(con[ncon].pos, mat1, vecccd.v); + con[ncon].dist = dist; + mji_mulMatVec3(con[ncon].frame, mat1, dir); + mji_mulMatVec3(con[ncon].pos, mat1, pos); mji_addTo3(con[ncon].pos, pos1); mju_zero3(con[ncon].frame+3); @@ -1660,7 +1668,7 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, mjtNum xmin, xmax, ymin, ymax, zmin, zmax; int dr[2], cnt, rmin, rmax, cmin, cmax; mjCCDObj obj1; - obj1.center = mjc_prism_center; + obj1.center = mjc_center; obj1.support = mjc_prism_support; // get hfield info @@ -1682,13 +1690,9 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, mjtNum* ecenter = d->flexelem_aabb + 6*(m->flex_elemadr[f]+e); // ccd-related - ccd_vec3_t dirccd, vecccd; - ccd_real_t depth; mjCCDObj obj2; mjc_initCCDObj(&obj2, m, d, -1, margin); mjc_setCCDObjFlex(&obj2, f, e, -1); - ccd_t ccd; - //------------------------------------- AABB computation, box-box test // save elem vertices, transform to hfield frame @@ -1743,18 +1747,6 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, //------------------------------------- collision testing - // init ccd structure - CCD_INIT(&ccd); - ccd.first_dir = prism_firstdir; - ccd.center1 = mjccd_prism_center; - ccd.center2 = mjccd_center; - ccd.support1 = mjccd_prism_support; - ccd.support2 = mjccd_support; - - // set ccd parameters - ccd.max_iterations = m->opt.ccd_iterations; - ccd.mpr_tolerance = m->opt.ccd_tolerance; - // compute real-valued grid step, and triangulation direction dx = (2.0*hsize[0]) / (ncol-1); dy = (2.0*hsize[1]) / (nrow-1); @@ -1782,23 +1774,22 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, } // run ccd, save contact - if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0) { - if (!ccdVec3Eq(&dirccd, ccd_vec3_origin)) { - // fill in contact data, transform to global coordinates - con[cnt].dist = -depth; - mji_mulMatVec3(con[cnt].frame, hmat, dirccd.v); - mji_mulMatVec3(con[cnt].pos, hmat, vecccd.v); - mji_addTo3(con[cnt].pos, hpos); - mju_zero3(con[cnt].frame+3); + mjtNum dist, dir[3], pos[3]; + if (mjc_penetration(m, &obj1, &obj2, &dist, dir, pos)) { + // fill in contact data, transform to global coordinates + con[cnt].dist = dist; + mji_mulMatVec3(con[cnt].frame, hmat, dir); + mji_mulMatVec3(con[cnt].pos, hmat, pos); + mji_addTo3(con[cnt].pos, hpos); + mju_zero3(con[cnt].frame+3); - // count, stop if max number reached - cnt++; - if (cnt >= mjMAXCONPAIR) { - r = rmax+1; - c = cmax+1; - k = 3; - break; - } + // count, stop if max number reached + cnt++; + if (cnt >= mjMAXCONPAIR) { + r = rmax+1; + c = cmax+1; + k = 3; + break; } } } From 253494739abf746548cab500f93774d8f071048e Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Wed, 25 Feb 2026 04:08:31 -0800 Subject: [PATCH 44/48] Read the clear color from the model if available PiperOrigin-RevId: 875079722 Change-Id: Id5a9c359d19f46d1c0bb5fbf5043db0ee4b1b9b6 --- .../filament/filament/filament_context.cc | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/experimental/filament/filament/filament_context.cc b/src/experimental/filament/filament/filament_context.cc index e02defb9..844d3cf1 100644 --- a/src/experimental/filament/filament/filament_context.cc +++ b/src/experimental/filament/filament/filament_context.cc @@ -40,6 +40,7 @@ #include "experimental/filament/filament/gui_view.h" #include "experimental/filament/filament/imgui_editor.h" #include "experimental/filament/filament/object_manager.h" +#include "experimental/filament/filament/model_util.h" #include "experimental/filament/filament/scene_view.h" #include "experimental/filament/filament/texture_util.h" #include "experimental/filament/render_context_filament.h" @@ -91,30 +92,29 @@ FilamentContext::FilamentContext(const mjrFilamentConfig* config, filament::Renderer::ClearOptions opts; opts.clear = true; opts.discard = true; - opts.clearColor = {0, 0, 0, 1}; + opts.clearColor = ReadElement(model_, "filament.clearColor", + filament::math::float4(0, 0, 0, 1)); renderer_->setClearOptions(opts); // Copy parameters from model to context. - if (model_) { - context_->shadowClip = model_->stat.extent * model_->vis.map.shadowclip; - context_->shadowScale = model_->vis.map.shadowscale; - context_->offWidth = model_->vis.global.offwidth; - context_->offHeight = model_->vis.global.offheight; - context_->offSamples = model_->vis.quality.offsamples; - context_->fogStart = - (float)(model_->stat.extent * model_->vis.map.fogstart); - context_->fogEnd = (float)(model_->stat.extent * model_->vis.map.fogend); - context_->fogRGBA[0] = model_->vis.rgba.fog[0]; - context_->fogRGBA[1] = model_->vis.rgba.fog[1]; - context_->fogRGBA[2] = model_->vis.rgba.fog[2]; - context_->fogRGBA[3] = model_->vis.rgba.fog[3]; - context_->lineWidth = model_->vis.global.linewidth; - context_->shadowSize = model_->vis.quality.shadowsize; - context_->readPixelFormat = 0x1907; // 0x1907 = GL_RGB; - context_->ntexture = model_->ntex; - for (int i = 0; i < model_->ntex; ++i) { - context_->textureType[i] = model_->tex_type[i]; - } + context_->shadowClip = model_->stat.extent * model_->vis.map.shadowclip; + context_->shadowScale = model_->vis.map.shadowscale; + context_->offWidth = model_->vis.global.offwidth; + context_->offHeight = model_->vis.global.offheight; + context_->offSamples = model_->vis.quality.offsamples; + context_->fogStart = + (float)(model_->stat.extent * model_->vis.map.fogstart); + context_->fogEnd = (float)(model_->stat.extent * model_->vis.map.fogend); + context_->fogRGBA[0] = model_->vis.rgba.fog[0]; + context_->fogRGBA[1] = model_->vis.rgba.fog[1]; + context_->fogRGBA[2] = model_->vis.rgba.fog[2]; + context_->fogRGBA[3] = model_->vis.rgba.fog[3]; + context_->lineWidth = model_->vis.global.linewidth; + context_->shadowSize = model_->vis.quality.shadowsize; + context_->readPixelFormat = 0x1907; // 0x1907 = GL_RGB; + context_->ntexture = model_->ntex; + for (int i = 0; i < model_->ntex; ++i) { + context_->textureType[i] = model_->tex_type[i]; } scene_view_ = std::make_unique(engine_, object_manager_.get()); From 52ed96bc3ac8464f3d5d51003141ec6c6cd3cd09 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 25 Feb 2026 04:31:37 -0800 Subject: [PATCH 45/48] Move tendon sparsity fields from mjData to mjModel PiperOrigin-RevId: 875087590 Change-Id: I1a5489d2d2011795ee38b09d547e68acd72e3cc7 --- doc/includes/references.h | 8 +- include/mujoco/mjdata.h | 3 - include/mujoco/mjmodel.h | 5 +- include/mujoco/mjxmacro.h | 8 +- mjx/mujoco/mjx/_src/io.py | 12 +-- mjx/mujoco/mjx/_src/smooth_test.py | 12 +-- .../mjx/third_party/mujoco_warp/_src/io.py | 15 ++-- mjx/mujoco/mjx/warp/forward_test.py | 6 +- python/mujoco/indexer_xmacro.h | 9 +- python/mujoco/introspect/structs.py | 58 ++++++------ src/engine/engine_core_constraint.c | 46 +++++----- src/engine/engine_core_smooth.c | 46 +++++----- src/engine/engine_derivative.c | 2 +- src/engine/engine_forward.c | 2 +- src/engine/engine_io.c | 26 +++--- src/engine/engine_io.h | 8 +- src/engine/engine_passive.c | 6 +- src/engine/engine_print.c | 12 +-- src/engine/engine_setconst.c | 90 ++++++++++++++++++- src/user/user_model.cc | 36 +++++++- src/user/user_model.h | 2 +- test/engine/engine_core_smooth_test.cc | 12 +-- test/user/user_model_test.cc | 42 +++++++++ unity/Runtime/Bindings/MjBindings.cs | 8 +- wasm/codegen/generated/bindings.cc | 36 ++++---- 25 files changed, 332 insertions(+), 178 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 6a30e704..9bc868ab 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -289,9 +289,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_tendon int* ten_wrapadr; // start address of tendon's path (ntendon x 1) 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 (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) @@ -1070,6 +1067,7 @@ struct mjModel_ { mjtSize nexclude; // number of excluded geom pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nwrap; // number of wrap objects in all tendon paths mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric custom fields @@ -1096,7 +1094,6 @@ 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) @@ -1490,6 +1487,9 @@ struct mjModel_ { int* tendon_group; // group for visibility (ntendon x 1) int* tendon_treenum; // number of trees along tendon's path (ntendon x 1) int* tendon_treeid; // first two trees along tendon's path (ntendon x 2) + 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 (nJten x 1) mjtByte* tendon_limited; // does tendon have length limits (ntendon x 1) mjtByte* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1) mjtNum* tendon_width; // width for rendering (ntendon x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 314dfc0e..b7a1e635 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -323,9 +323,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_tendon int* ten_wrapadr; // start address of tendon's path (ntendon x 1) 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 (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) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 20151f48..97de5c2f 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -728,6 +728,7 @@ struct mjModel_ { mjtSize nexclude; // number of excluded geom pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nwrap; // number of wrap objects in all tendon paths mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric custom fields @@ -754,7 +755,6 @@ 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) @@ -1148,6 +1148,9 @@ struct mjModel_ { int* tendon_group; // group for visibility (ntendon x 1) int* tendon_treenum; // number of trees along tendon's path (ntendon x 1) int* tendon_treeid; // first two trees along tendon's path (ntendon x 2) + 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 (nJten x 1) mjtByte* tendon_limited; // does tendon have length limits (ntendon x 1) mjtByte* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1) mjtNum* tendon_width; // width for rendering (ntendon x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index cba09e7c..3894c898 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -215,6 +215,7 @@ X( nexclude ) \ X( neq ) \ X( ntendon ) \ + X( nJten ) \ X( nwrap ) \ X( nsensor ) \ X( nnumeric ) \ @@ -239,7 +240,6 @@ X( npaths ) \ X( nnames_map ) \ X( nJmom ) \ - X( nJten ) \ X( ngravcomp ) \ X( nemax ) \ X( njmax ) \ @@ -629,6 +629,9 @@ X ( int, tendon_group, ntendon, 1 ) \ X ( int, tendon_treenum, ntendon, 1 ) \ X ( int, tendon_treeid, ntendon, 2 ) \ + X ( int, ten_J_rownnz, ntendon, 1 ) \ + X ( int, ten_J_rowadr, ntendon, 1 ) \ + X ( int, ten_J_colind, nJten, 1 ) \ X ( mjtByte, tendon_limited, ntendon, 1 ) \ X ( mjtByte, tendon_actfrclimited, ntendon, 1 ) \ X ( mjtNum, tendon_width, ntendon, 1 ) \ @@ -848,9 +851,6 @@ X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ X ( int, ten_wrapadr, ntendon, 1 ) \ 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, nJten, 1 ) \ X ( mjtNum, ten_J, nJten, 1 ) \ X ( mjtNum, ten_length, ntendon, 1 ) \ X ( int, wrap_obj, nwrap, 2 ) \ diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index d4e0d1fa..0ba4d41c 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1129,9 +1129,9 @@ def _put_data_jax( mujoco.mju_sparse2dense( ten_J, d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, ) else: ten_J = np.zeros((m.ntendon, m.nv)) @@ -1251,6 +1251,9 @@ def _put_data_c( for f in types.DataC.fields() if hasattr(d, f.name) } + for f in types.DataC.fields(): + if not hasattr(d, f.name) and hasattr(m, f.name): + impl_fields[f.name] = getattr(m, f.name) # TODO(stunya): support islanding via C impl. impl_fields['solver_niter'] = impl_fields['solver_niter'][0] @@ -1662,9 +1665,6 @@ def _get_data_into( ) 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 diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index ce8c2588..7f0d8c09 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -125,9 +125,9 @@ class SmoothTest(absltest.TestCase): mujoco.mju_sparse2dense( ten_J, d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, ) _assert_eq(ten_J, dx._impl.ten_J, 'ten_J') _assert_attr_eq(d, dx, 'ten_length') @@ -407,9 +407,9 @@ class TendonTest(parameterized.TestCase): mujoco.mju_sparse2dense( ten_J, d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, ) _assert_eq(ten_J, dx._impl.ten_J, 'ten_J') _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum') 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 05bd1d2b..e6533dd6 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -952,9 +952,12 @@ def put_data( d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float) - ten_J = np.zeros((mjm.ntendon, mjm.nv)) - mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1)) - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) + if mjm.ntendon: + ten_J = np.zeros((mjm.ntendon, mjm.nv)) + mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjm.ten_J_rownnz, mjm.ten_J_rowadr, mjm.ten_J_colind.reshape(-1)) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) + else: + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), 0.0), dtype=float) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) @@ -1165,9 +1168,9 @@ def get_data_into( mujoco.mju_dense2sparse( result.ten_J, ten_J, - result.ten_J_rownnz, - result.ten_J_rowadr, - result.ten_J_colind, + mjm.ten_J_rownnz, + mjm.ten_J_rowadr, + mjm.ten_J_colind, ) else: result.ten_J[:] = d.ten_J.numpy()[world_id] diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 30eb21cf..3eb936a7 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -153,9 +153,9 @@ class ForwardTest(parameterized.TestCase): mujoco.mju_sparse2dense( ten_J, d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, ) tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J') tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr') diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h index 9f26000d..524f9376 100644 --- a/python/mujoco/indexer_xmacro.h +++ b/python/mujoco/indexer_xmacro.h @@ -259,7 +259,10 @@ X( mjtNum, tendon, _length0, ntendon, 1 ) \ X( mjtNum, tendon, _invweight0, ntendon, 1 ) \ X( mjtNum, tendon, _user, ntendon, MJ_M(nuser_tendon) ) \ - X( float, tendon, _rgba, ntendon, 4 ) + X( float, tendon, _rgba, ntendon, 4 ) \ + X( int, ten_, J_rownnz, ntendon, 1 ) \ + X( int, ten_, J_rowadr, ntendon, 1 ) \ + X( int, ten_, J_colind, ntendon, MJ_M(nv) ) #define MJMODEL_TEXTURE \ X( int, tex_, type, ntex, 1 ) \ @@ -402,11 +405,7 @@ #define MJDATA_TENDON \ X( int, ten_, wrapadr , ntendon, 1 ) \ 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_, length , ntendon, 1 ) \ - X( mjtNum, ten_, J , ntendon, MJ_M(nv) ) \ X( mjtNum, ten_, velocity, ntendon, 1 ) #define MJDATA_VIEW_GROUPS \ diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index c8f6774f..433b8167 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -1122,6 +1122,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='mjtSize'), doc='number of tendons', ), + StructFieldDecl( + name='nJten', + type=ValueType(name='mjtSize'), + doc='number of non-zeros in sparse ten_J matrix', + ), StructFieldDecl( name='nwrap', type=ValueType(name='mjtSize'), @@ -1242,11 +1247,6 @@ 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'), @@ -3927,6 +3927,30 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc="first two trees along tendon's path", array_extent=('ntendon', 2), ), + StructFieldDecl( + name='ten_J_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in Jacobian row', + array_extent=('ntendon',), + ), + StructFieldDecl( + name='ten_J_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in colind array', + array_extent=('ntendon',), + ), + StructFieldDecl( + name='ten_J_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in sparse Jacobian', + array_extent=('nJten',), + ), StructFieldDecl( name='tendon_limited', type=PointerType( @@ -5812,30 +5836,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='number of wrap points in path', array_extent=('ntendon',), ), - StructFieldDecl( - name='ten_J_rownnz', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of non-zeros in Jacobian row', - array_extent=('ntendon',), - ), - StructFieldDecl( - name='ten_J_rowadr', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='row start address in colind array', - array_extent=('ntendon',), - ), - StructFieldDecl( - name='ten_J_colind', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='column indices in sparse Jacobian', - array_extent=('nJten',), - ), StructFieldDecl( name='ten_J', type=PointerType( diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index d550766c..982463d7 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -568,16 +568,16 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // copy Jacobian: sparse or dense if (issparse) { if (j == 0) { - NV = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); - mju_copy(jac[j], d->ten_J+d->ten_J_rowadr[id[j]], NV); + NV = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV); + mju_copy(jac[j], d->ten_J+m->ten_J_rowadr[id[j]], NV); } else { - NV2 = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); - mju_copy(jac[j], d->ten_J+d->ten_J_rowadr[id[j]], NV2); + NV2 = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain2, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV2); + mju_copy(jac[j], d->ten_J+m->ten_J_rowadr[id[j]], NV2); } } else { - mju_sparse2dense(jac[j], d->ten_J, 1, nv, d->ten_J_rownnz+id[j], d->ten_J_rowadr+id[j], d->ten_J_colind); + mju_sparse2dense(jac[j], d->ten_J, 1, nv, m->ten_J_rownnz+id[j], m->ten_J_rowadr+id[j], m->ten_J_colind); } } } @@ -737,13 +737,13 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { int efcadr = d->nefc; // add constraint if (issparse) { - mj_addConstraint(m, d, d->ten_J + d->ten_J_rowadr[i], + mj_addConstraint(m, d, d->ten_J + m->ten_J_rowadr[i], 0, 0, m->tendon_frictionloss[i], 1, mjCNSTR_FRICTION_TENDON, i, - d->ten_J_rownnz[i], - d->ten_J_colind+d->ten_J_rowadr[i]); + m->ten_J_rownnz[i], + m->ten_J_colind+m->ten_J_rowadr[i]); } else { - mju_sparse2dense(jac, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); + mju_sparse2dense(jac, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); mj_addConstraint(m, d, jac, 0, 0, m->tendon_frictionloss[i], 1, mjCNSTR_FRICTION_TENDON, i, 0, NULL); } @@ -885,13 +885,13 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { // prepare Jacobian int efcadr = d->nefc; if (issparse) { - mju_scl(jac, d->ten_J+d->ten_J_rowadr[i], -side, d->ten_J_rownnz[i]); + mju_scl(jac, d->ten_J+m->ten_J_rowadr[i], -side, m->ten_J_rownnz[i]); mj_addConstraint(m, d, jac, &dist, &margin, 0, 1, mjCNSTR_LIMIT_TENDON, i, - d->ten_J_rownnz[i], - d->ten_J_colind+d->ten_J_rowadr[i]); + m->ten_J_rownnz[i], + m->ten_J_colind+m->ten_J_rowadr[i]); } else { - mju_sparse2dense(jac, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); + mju_sparse2dense(jac, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); mju_scl(jac, jac, -side, nv); mj_addConstraint(m, d, jac, &dist, &margin, 0, 1, mjCNSTR_LIMIT_TENDON, i, 0, NULL); @@ -1738,11 +1738,11 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { } } else { if (!j) { - NV = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); + NV = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV); } else { - NV2 = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); + NV2 = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain2, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV2); } } } @@ -1840,8 +1840,8 @@ static int mj_nf(const mjModel* m, const mjData* d, int *nnz) { for (int i=0; i < ntendon; i++) { if (m->tendon_frictionloss[i] > 0) { - nf += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); - if (nnz) *nnz += d->ten_J_rownnz[i]; + nf += mj_addConstraintCount(m, 1, m->ten_J_rownnz[i]); + if (nnz) *nnz += m->ten_J_rownnz[i]; } } @@ -1908,8 +1908,8 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { for (int i=0; i < ntendon; i++) { int count = tendonLimit(m, d->ten_length, i); for (int j = 0; j < count; j++) { - nl += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); - if (nnz) *nnz += d->ten_J_rownnz[i]; + nl += mj_addConstraintCount(m, 1, m->ten_J_rownnz[i]); + if (nnz) *nnz += m->ten_J_rownnz[i]; } } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 1d3fed48..47e0fdee 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -905,7 +905,7 @@ void mj_flex(const mjModel* m, mjData* d) { // compute tendon lengths and moments void mj_tendon(const mjModel* m, mjData* d) { int nv = m->nv, nten = m->ntendon; - int *rownnz = d->ten_J_rownnz, *rowadr = d->ten_J_rowadr, *colind = d->ten_J_colind; + const int *rownnz = m->ten_J_rownnz, *rowadr = m->ten_J_rowadr, *colind = m->ten_J_colind; mjtNum *L = d->ten_length, *J = d->ten_J; if (!nten) { @@ -913,22 +913,20 @@ void mj_tendon(const mjModel* m, mjData* d) { } // allocate stack arrays - int *chain, *buf_ind; - mjtNum *jac1, *jac2, *jacdif, *tmp, *sparse_buf; + int *chain; + mjtNum *jac1, *jac2, *jacdif, *tmp; mj_markStack(d); jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); tmp = mjSTACKALLOC(d, nv, mjtNum); chain = mjSTACKALLOC(d, nv, int); - buf_ind = mjSTACKALLOC(d, nv, int); - sparse_buf = mjSTACKALLOC(d, nv, mjtNum); // clear results mju_zero(L, nten); // clear Jacobian - mju_zeroInt(rownnz, nten); + mju_zero(J, m->nJten); // sleep filtering int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; @@ -947,9 +945,6 @@ void mj_tendon(const mjModel* m, mjData* d) { d->ten_wrapnum[i] = 0; int tendon_num = m->tendon_num[i]; - // sparse Jacobian row init - rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); - // process fixed tendon if (m->wrap_type[adr] == mjWRAP_JOINT) { // process all defined joints @@ -960,11 +955,10 @@ void mj_tendon(const mjModel* m, mjData* d) { // add to length L[i] += m->wrap_prm[adr+j] * d->qpos[m->jnt_qposadr[k]]; - // add to moment - rownnz[i] = mju_combineSparse(J+rowadr[i], &m->wrap_prm[adr+j], 1, 1, - rownnz[i], 1, - colind+rowadr[i], &m->jnt_dofadr[k], - sparse_buf, buf_ind); + mjtNum coef = 1; + int dofadr = m->jnt_dofadr[k]; + mju_combineSparseInc(J + rowadr[i], &coef, m->nv, 1, m->wrap_prm[adr+j], + rownnz[i], 1, colind + rowadr[i], &dofadr); } continue; @@ -1059,9 +1053,9 @@ void mj_tendon(const mjModel* m, mjData* d) { mju_mulMatTVec(tmp, jacdif, dif, 3, NV); // add to existing - rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor, - rownnz[i], NV, colind+rowadr[i], - chain, sparse_buf, buf_ind); + mju_combineSparseInc(J+rowadr[i], tmp, nv, 1, 1/divisor, + rownnz[i], NV, colind+rowadr[i], + chain); } } @@ -1419,10 +1413,10 @@ void mj_transmission(const mjModel* m, mjData* d) { // moment { - int ten_J_rownnz = d->ten_J_rownnz[id]; - int ten_J_rowadr = d->ten_J_rowadr[id]; + int ten_J_rownnz = m->ten_J_rownnz[id]; + int ten_J_rowadr = m->ten_J_rowadr[id]; rownnz[i] = ten_J_rownnz; - mju_copyInt(colind + adr, d->ten_J_colind + ten_J_rowadr, ten_J_rownnz); + mju_copyInt(colind + adr, m->ten_J_colind + ten_J_rowadr, ten_J_rownnz); mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz); } @@ -1717,9 +1711,9 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { } // get sparse info for tendon k - int J_rowadr = d->ten_J_rowadr[k]; - int J_rownnz = d->ten_J_rownnz[k]; - const int* J_colind = d->ten_J_colind + J_rowadr; + int J_rowadr = m->ten_J_rowadr[k]; + int J_rownnz = m->ten_J_rownnz[k]; + const int* J_colind = m->ten_J_colind + J_rowadr; mjtNum* ten_J = d->ten_J + J_rowadr; // M += armature * ten_J' * ten_J @@ -2649,9 +2643,9 @@ void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { if (coef) { // sparse - int nnz = d->ten_J_rownnz[i]; - int adr = d->ten_J_rowadr[i]; - const int* colind = d->ten_J_colind + adr; + int nnz = m->ten_J_rownnz[i]; + int adr = m->ten_J_rowadr[i]; + const int* colind = m->ten_J_colind + adr; const mjtNum* ten_J = d->ten_J + adr; for (int j=0; j < nnz; j++) { qfrc[colind[j]] += coef * ten_J[j]; diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 67b0d02e..09dd4903 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -1776,7 +1776,7 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { } // add sparse - addJTBJSparse(m, d, d->ten_J, &B, 1, i, d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind); + addJTBJSparse(m, d, d->ten_J, &B, 1, i, m->ten_J_rownnz, m->ten_J_rowadr, m->ten_J_colind); } } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index c38e6508..f2a7c3c4 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -232,7 +232,7 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) { // tendon velocity: always sparse mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, - d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL); + m->ten_J_rownnz, m->ten_J_rowadr, m->ten_J_colind, NULL); // actuator velocity: always sparse if (!mjDISABLED(mjDSBL_ACTUATION)) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 512c7661..f0cf467f 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -211,10 +211,10 @@ void mj_makeModel(mjModel** dest, mjtSize nmeshpoly, mjtSize nmeshpolyvert, mjtSize nmeshpolymap, mjtSize nskin, mjtSize nskinvert, mjtSize nskintexvert, mjtSize nskinface, mjtSize nskinbone, mjtSize nskinbonevert, mjtSize nhfield, mjtSize nhfielddata, mjtSize ntex, mjtSize ntexdata, - mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nwrap, - mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, mjtSize ntextdata, - mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, mjtSize nplugin, - mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, + mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nJten, + mjtSize nwrap, mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, + mjtSize ntextdata, mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, + mjtSize nplugin, mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, mjtSize nuser_site, mjtSize nuser_cam, mjtSize nuser_tendon, mjtSize nuser_actuator, mjtSize nuser_sensor, mjtSize nnames, mjtSize npaths) { intptr_t offset = 0; @@ -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, nJten = 0, ngravcomp = 0, nemax = 0, njmax = 0, nconmax=0; + int nnames_map = 0, nJmom = 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 @@ -243,7 +243,7 @@ void mj_makeModel(mjModel** dest, #undef X // suppress unused variable warnings - (void)nnames_map; (void)nJmom; (void)nJten; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax; + (void)nnames_map; (void)nJmom; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax; (void)nuserdata; (void)nsensordata; (void)npluginstate; (void)nhistory; (void)narena; (void)nbuffer; } @@ -323,6 +323,7 @@ void mj_makeModel(mjModel** dest, m->nexclude = nexclude; m->neq = neq; m->ntendon = ntendon; + m->nJten = nJten; m->nwrap = nwrap; m->nsensor = nsensor; m->nnumeric = nnumeric; @@ -410,11 +411,11 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) { src->nskin, src->nskinvert, src->nskintexvert, src->nskinface, src->nskinbone, src->nskinbonevert, src->nhfield, src->nhfielddata, src->ntex, src->ntexdata, src->nmat, src->npair, src->nexclude, - src->neq, src->ntendon, src->nwrap, src->nsensor, src->nnumeric, - src->nnumericdata, src->ntext, src->ntextdata, src->ntuple, - src->ntupledata, src->nkey, src->nmocap, src->nplugin, src->npluginattr, - src->nuser_body, src->nuser_jnt, src->nuser_geom, src->nuser_site, - src->nuser_cam, src->nuser_tendon, src->nuser_actuator, + src->neq, src->ntendon, src->nJten, src->nwrap, src->nsensor, + src->nnumeric, src->nnumericdata, src->ntext, src->ntextdata, + src->ntuple, src->ntupledata, src->nkey, src->nmocap, src->nplugin, + src->npluginattr, src->nuser_body, src->nuser_jnt, src->nuser_geom, + src->nuser_site, src->nuser_cam, src->nuser_tendon, src->nuser_actuator, src->nuser_sensor, src->nnames, src->npaths); } if (!dest) { @@ -597,7 +598,8 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) { sizes[49], sizes[50], sizes[51], sizes[52], sizes[53], sizes[54], sizes[55], sizes[56], sizes[57], sizes[58], sizes[59], sizes[60], sizes[61], sizes[62], sizes[63], sizes[64], sizes[65], sizes[66], sizes[67], sizes[68], sizes[69], - sizes[70], sizes[71], sizes[72], sizes[73], sizes[74], sizes[75], sizes[76]); + sizes[70], sizes[71], sizes[72], sizes[73], sizes[74], sizes[75], sizes[76], + sizes[77]); // mj_makeModel may fail if the input buffer has invalid sizes if (!m) { diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 918ad5a3..0af1058f 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -58,10 +58,10 @@ void mj_makeModel(mjModel** dest, mjtSize nmeshpoly, mjtSize nmeshpolyvert, mjtSize nmeshpolymap, mjtSize nskin, mjtSize nskinvert, mjtSize nskintexvert, mjtSize nskinface, mjtSize nskinbone, mjtSize nskinbonevert, mjtSize nhfield, mjtSize nhfielddata, mjtSize ntex, mjtSize ntexdata, - mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nwrap, - mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, mjtSize ntextdata, - mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, mjtSize nplugin, - mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, + mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nJten, + mjtSize nwrap, mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, + mjtSize ntextdata, mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, + mjtSize nplugin, mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, mjtSize nuser_site, mjtSize nuser_cam, mjtSize nuser_tendon, mjtSize nuser_actuator, mjtSize nuser_sensor, mjtSize nnames, mjtSize npaths); diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 5e26c2c3..60446c13 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -473,9 +473,9 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // transform to joint torque, add to qfrc_{spring, damper} if (frc_spring || frc_damper) { - int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; - for (int j=d->ten_J_rowadr[i]; j < end; j++) { - int k = d->ten_J_colind[j]; + int end = m->ten_J_rowadr[i] + m->ten_J_rownnz[i]; + for (int j=m->ten_J_rowadr[i]; j < end; j++) { + int k = m->ten_J_colind[j]; mjtNum J = d->ten_J[j]; d->qfrc_spring[k] += J * frc_spring; d->qfrc_damper[k] += J * frc_damper; diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 0a927715..a8f407d9 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1441,12 +1441,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format); printArray2d("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format); - mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL, - d->ten_J_rownnz, NULL, d->ten_J_colind, fp); - printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp); - printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp); - printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz, - d->ten_J_rowadr, d->ten_J_colind, fp, float_format); + mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, m->ten_J_rowadr, NULL, + m->ten_J_rownnz, NULL, m->ten_J_colind, fp); + printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, m->ten_J_rownnz, fp); + printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, m->ten_J_rowadr, fp); + printSparse("TEN_J", d->ten_J, m->ntendon, m->ten_J_rownnz, + m->ten_J_rowadr, m->ten_J_colind, fp, float_format); for (int i=0; i < m->ntendon; i++) { fprintf(fp, "TENDON %d: %d wrap points\n", i, d->ten_wrapnum[i]); for (int j=0; j < d->ten_wrapnum[i]; j++) { diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 435291a4..6dd01a94 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -282,6 +282,93 @@ static void setFixed(mjModel* m, mjData* d) { mj_freeStack(d); } +// compute tendon Jacobian sparsity +static void makeTendonSparse(mjModel* m) { + int ntendon = m->ntendon; + int* rownnz = m->ten_J_rownnz; + int* rowadr = m->ten_J_rowadr; + int* colind = m->ten_J_colind; + + if (!ntendon) { + return; + } + + // clear + mju_zeroInt(rownnz, ntendon); + mju_zeroInt(rowadr, ntendon); + + // compute rownnz, rowadr, and colind for each tendon + for (int i = 0; i < ntendon; i++) { + rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); + int adr = m->tendon_adr[i]; + int num = m->tendon_num[i]; + + // joint tendon: each wrap object is a joint, colind is its dofadr + if (m->wrap_type[adr] == mjWRAP_JOINT) { + for (int j = 0; j < num; j++) { + colind[rowadr[i] + j] = m->jnt_dofadr[m->wrap_objid[adr + j]]; + } + rownnz[i] = num; + } else { + // spatial tendon: collect used dofs from wrap object bodies + int nnz = 0; + for (int j = 0; j < num; j++) { + int type = m->wrap_type[adr + j]; + + // get body id from site or geom wrap object + int bodyid = -1; + if (type == mjWRAP_SITE) { + bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; + } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { + bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; + } + + // walk up the body tree, collecting used dofs + if (bodyid > 0) { + int bid = bodyid; + while (bid > 0) { + int bdofadr = m->body_dofadr[bid]; + int bdofnum = m->body_dofnum[bid]; + for (int k = 0; k < bdofnum; k++) { + int dof = bdofadr + k; + + // check if dof already in colind + int found = 0; + for (int l = 0; l < nnz; l++) { + if (colind[rowadr[i] + l] == dof) { + found = 1; + break; + } + } + + // append new dof + if (!found) { + colind[rowadr[i] + nnz] = dof; + nnz++; + } + } + bid = m->body_parentid[bid]; + } + } + } + rownnz[i] = nnz; + } + + // sort colind for this tendon + int nnz = rownnz[i]; + for (int j = 0; j < nnz - 1; j++) { + for (int k = j + 1; k < nnz; k++) { + // swap out-of-order entries + if (colind[rowadr[i] + k] < colind[rowadr[i] + j]) { + int tmp = colind[rowadr[i] + j]; + colind[rowadr[i] + j] = colind[rowadr[i] + k]; + colind[rowadr[i] + k] = tmp; + } + } + } + } +} + // compute flex sparsity: flexedge_J_{rowadr,rownnz,colind} and flexvert_J_{rowadr,rownnz} static void makeFlexSparse(mjModel* m, mjData* d) { int nv = m->nv; @@ -552,6 +639,7 @@ static void mj_alignFlex(mjModel* m, mjData* d) { // set quantities that depend on qpos0 static void set0(mjModel* m, mjData* d) { + makeTendonSparse(m); makeFlexSparse(m, d); mj_alignFlex(m, d); int nv = m->nv; @@ -756,7 +844,7 @@ static void set0(mjModel* m, mjData* d) { // compute tendon_invweight0 for (int i=0; i < m->ntendon; i++) { - mju_sparse2dense(tmp, d->ten_J, 1, nv, d->ten_J_rownnz+i, d->ten_J_rowadr+i, d->ten_J_colind); + mju_sparse2dense(tmp, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); // solve into tmp+nv mj_solveM(m, d, tmp+nv, tmp, 1); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 1c3e7e24..6c642352 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -907,6 +907,7 @@ void mjCModel::ComputeSparseSizes() { // no dofs, quick return if (nv == 0) { nM = nD = nB = nC = 0; + nJten = 0; return; } @@ -1084,6 +1085,37 @@ void mjCModel::ComputeSparseSizes() { } } nC = nOD + nv; + + nJten = 0; + if (nv > 0) { + std::vector dof_bitmap(nv, false); + for (const auto* tendon : tendons_) { + if (!tendon->path.empty() && + tendon->path[0]->Type() == mjWRAP_JOINT) { + nJten += tendon->path.size(); + continue; + } + + std::fill(dof_bitmap.begin(), dof_bitmap.end(), false); + for (const auto* wrap : tendon->path) { + int bodyid = GetBodyIdFromWrap(wrap); + if (bodyid > 0) { + mjCBody* b = bodies_[bodyid]; + while (b && b->id > 0) { + for (const auto* jnt : b->joints) { + for (int k = 0; k < jnt->nv(); k++) { + dof_bitmap[jnt->dofadr_ + k] = true; + } + } + b = b->GetParent(); + } + } + } + for (int j = 0; j < nv; j++) { + nJten += dof_bitmap[j]; + } + } + } } @@ -5073,7 +5105,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly, nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert, nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude, - neq, ntendon, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata, + neq, ntendon, nJten, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata, ntuple, ntupledata, nkey, nmocap, nplugin, npluginattr, nuser_body, nuser_jnt, nuser_geom, nuser_site, nuser_cam, nuser_tendon, nuser_actuator, nuser_sensor, nnames, npaths); @@ -5105,8 +5137,6 @@ 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) { diff --git a/src/user/user_model.h b/src/user/user_model.h index 2a3ac355..025f8097 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -73,6 +73,7 @@ class mjCModel_ : public mjsElement { mjtSize nexclude; // number of excluded body pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric fields mjtSize ntext; // number of text fields @@ -130,7 +131,6 @@ 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 diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 7b292bfe..0449cb4f 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -154,12 +154,8 @@ TEST_F(CoreSmoothTest, FixedTendonSortedIndices) { mjData* data = mj_makeData(model); mj_fwdPosition(model, data); - int rowadr = data->ten_J_rowadr[0]; - int* colind = data->ten_J_colind + rowadr; - mjtNum* J = data->ten_J + rowadr; - + mjtNum* J = data->ten_J; EXPECT_THAT(vector(J, J + 3), ElementsAre(1, 2, 3)); - EXPECT_THAT(vector(colind, colind + 3), ElementsAre(0, 1, 2)); mj_deleteData(data); mj_deleteModel(model); @@ -253,11 +249,11 @@ TEST_F(CoreSmoothTest, TendonArmature) { // add tendon inertias to M2 using outer product for (int j=0; j < m->ntendon; j++) { // get tendon Jacobian - int rowadr = d->ten_J_rowadr[j]; - int* rownnz = d->ten_J_rownnz + j; + int rowadr = m->ten_J_rowadr[j]; + int* rownnz = m->ten_J_rownnz + j; int zero = 0; mju_sparse2dense(ten_J.data(), d->ten_J + rowadr, 1, nv, - rownnz, &zero, d->ten_J_colind + rowadr); + rownnz, &zero, m->ten_J_colind + rowadr); // get tendon inertia only, using outer product mju_mulMatMat(ten_M.data(), ten_J.data(), ten_J.data(), nv, 1, nv); diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 8fd84caa..bfd68c0e 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -199,6 +199,48 @@ TEST_F(UserModelTest, ActuatorSparsity) { mj_deleteModel(m); } +TEST_F(UserModelTest, FixedTendonSparsity) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + )"; + mjModel* m = LoadModelFromString(xml); + ASSERT_THAT(m, NotNull()); + + EXPECT_EQ(m->nJten, 3); + EXPECT_EQ(m->ten_J_rownnz[0], 3); + EXPECT_EQ(m->ten_J_rowadr[0], 0); + EXPECT_EQ(m->wrap_type[m->tendon_adr[0]], mjWRAP_JOINT); + + int rowadr = m->ten_J_rowadr[0]; + int* colind = m->ten_J_colind + rowadr; + EXPECT_THAT(std::vector(colind, colind + 3), ElementsAre(0, 1, 2)); + + mj_deleteModel(m); +} + TEST_F(UserModelTest, NestedZeroMassBodiesOK) { static constexpr char xml[] = R"( diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 949dc889..0bd6e798 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4994,9 +4994,6 @@ public unsafe struct mjData_ { public double* bvh_aabb_dyn; public int* ten_wrapadr; public int* ten_wrapnum; - public int* ten_J_rownnz; - public int* ten_J_rowadr; - public int* ten_J_colind; public double* ten_J; public double* ten_length; public int* wrap_obj; @@ -5347,6 +5344,7 @@ public unsafe struct mjModel_ { public Int64 nexclude; public Int64 neq; public Int64 ntendon; + public Int64 nJten; public Int64 nwrap; public Int64 nsensor; public Int64 nnumeric; @@ -5371,7 +5369,6 @@ 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; @@ -5712,6 +5709,9 @@ public unsafe struct mjModel_ { public int* tendon_group; public int* tendon_treenum; public int* tendon_treeid; + public int* ten_J_rownnz; + public int* ten_J_rowadr; + public int* ten_J_colind; public byte* tendon_limited; public byte* tendon_actfrclimited; public double* tendon_width; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index d8b69429..3aa71a8c 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -3934,6 +3934,12 @@ struct MjModel { void set_ntendon(int value) { ptr_->ntendon = static_cast(value); } + int nJten() const { + return static_cast(ptr_->nJten); + } + void set_nJten(int value) { + ptr_->nJten = static_cast(value); + } int nwrap() const { return static_cast(ptr_->nwrap); } @@ -4078,12 +4084,6 @@ struct MjModel { void set_nJmom(int value) { ptr_->nJmom = static_cast(value); } - int nJten() const { - return static_cast(ptr_->nJten); - } - void set_nJten(int value) { - ptr_->nJten = static_cast(value); - } int ngravcomp() const { return static_cast(ptr_->ngravcomp); } @@ -5125,6 +5125,15 @@ struct MjModel { emscripten::val tendon_treeid() const { return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon * 2, ptr_->tendon_treeid)); } + emscripten::val ten_J_rownnz() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->ten_J_rownnz)); + } + emscripten::val ten_J_rowadr() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->ten_J_rowadr)); + } + emscripten::val ten_J_colind() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nJten, ptr_->ten_J_colind)); + } emscripten::val tendon_limited() const { return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->tendon_limited)); } @@ -6649,15 +6658,6 @@ struct MjData { emscripten::val ten_wrapnum() const { return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_wrapnum)); } - emscripten::val ten_J_rownnz() const { - return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_J_rownnz)); - } - emscripten::val ten_J_rowadr() const { - 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->nJten, ptr_->ten_J_colind)); - } emscripten::val ten_J() const { return emscripten::val(emscripten::typed_memory_view(model->nJten, ptr_->ten_J)); } @@ -11550,9 +11550,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("subtree_com", &MjData::subtree_com) .property("subtree_linvel", &MjData::subtree_linvel) .property("ten_J", &MjData::ten_J) - .property("ten_J_colind", &MjData::ten_J_colind) - .property("ten_J_rowadr", &MjData::ten_J_rowadr) - .property("ten_J_rownnz", &MjData::ten_J_rownnz) .property("ten_length", &MjData::ten_length) .property("ten_velocity", &MjData::ten_velocity) .property("ten_wrapadr", &MjData::ten_wrapadr) @@ -12100,6 +12097,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("skin_vertadr", &MjModel::skin_vertadr) .property("skin_vertnum", &MjModel::skin_vertnum) .property("stat", &MjModel::stat, reference()) + .property("ten_J_colind", &MjModel::ten_J_colind) + .property("ten_J_rowadr", &MjModel::ten_J_rowadr) + .property("ten_J_rownnz", &MjModel::ten_J_rownnz) .property("tendon_actfrclimited", &MjModel::tendon_actfrclimited) .property("tendon_actfrcrange", &MjModel::tendon_actfrcrange) .property("tendon_adr", &MjModel::tendon_adr) From 22e3217fc0f394c5355f4cf851130ff04dcca790 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 25 Feb 2026 04:46:49 -0800 Subject: [PATCH 46/48] Replace TriangleMeshDistance with custom BVH-based SDF computation. The `mjCOctree::ComputeSdfCoeffs` function has been replaced by `mjCOctree::ComputeSdf`. The new implementation removes the dependency on the `triangle_mesh_distance` library. Instead, it uses a provided `mjCBoundingVolumeHierarchy` to perform closest point queries on the mesh. The signed distance is computed by finding the closest triangle face via the BVH and determining the sign based on the dot product of the vector from the closest point on the triangle to the octree vertex and the triangle's normal. Optional Laplacian smoothing has been added to the SDF coefficients. This helps smooth out potential discontinuities at octree level boundaries. PiperOrigin-RevId: 875093374 Change-Id: I6fa53243a2dda107ccf255bc6b876ca5a4554d4f --- cmake/MujocoDependencies.cmake | 30 ---- src/user/user_mesh.cc | 2 +- src/user/user_objects.cc | 263 +++++++++++++++++++++++++++++++-- src/user/user_objects.h | 14 +- test/user/user_objects_test.cc | 145 ++++++++++++++++++ 5 files changed, 411 insertions(+), 43 deletions(-) diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake index 699bbb15..d2404bc1 100644 --- a/cmake/MujocoDependencies.cmake +++ b/cmake/MujocoDependencies.cmake @@ -58,11 +58,6 @@ set(MUJOCO_DEP_VERSION_benchmark CACHE STRING "Version of `benchmark` to be fetched." ) -set(MUJOCO_DEP_VERSION_TriangleMeshDistance - 2cb643de1436e1ba8e2be49b07ec5491ac604457 - CACHE STRING "Version of `TriangleMeshDistance` to be fetched." -) - mark_as_advanced(MUJOCO_DEP_VERSION_lodepng) mark_as_advanced(MUJOCO_DEP_VERSION_MarchingCubeCpp) mark_as_advanced(MUJOCO_DEP_VERSION_tinyxml2) @@ -73,7 +68,6 @@ mark_as_advanced(MUJOCO_DEP_VERSION_Eigen3) mark_as_advanced(MUJOCO_DEP_VERSION_abseil) mark_as_advanced(MUJOCO_DEP_VERSION_gtest) mark_as_advanced(MUJOCO_DEP_VERSION_benchmark) -mark_as_advanced(MUJOCO_DEP_VERSION_TriangleMeshDistance) include(FetchContent) include(FindOrFetch) @@ -202,30 +196,6 @@ if(CMAKE_POLICY_VERSION_MINIMUM_LOCALLY_DEFINED) unset(CMAKE_POLICY_VERSION_MINIMUM_LOCALLY_DEFINED) endif() -if(NOT TARGET trianglemeshdistance) - FetchContent_Declare( - trianglemeshdistance - GIT_REPOSITORY https://github.com/InteractiveComputerGraphics/TriangleMeshDistance.git - GIT_TAG ${MUJOCO_DEP_VERSION_TriangleMeshDistance} - ) - - FetchContent_GetProperties(trianglemeshdistance) - if(NOT trianglemeshdistance_POPULATED) - FetchContent_Populate(trianglemeshdistance) - # Patch the source code to silence a warning/error related to a loop variable creating a copy. - # Since this is a header only library this fix is less intrusive than disabling the warning for - # any target including the header. - set(TMD_HEADER ${trianglemeshdistance_SOURCE_DIR}/TriangleMeshDistance/include/tmd/TriangleMeshDistance.h) - file(READ ${TMD_HEADER} TMD_CONTENT) - string(REPLACE - "for (const auto edge_count : edges_count) {" - "for (const auto& edge_count : edges_count) {" - TMD_CONTENT "${TMD_CONTENT}") - file(WRITE ${TMD_HEADER} "${TMD_CONTENT}") - include_directories(${trianglemeshdistance_SOURCE_DIR}) - endif() -endif() - set(ENABLE_DOUBLE_PRECISION ON) set(CCD_HIDE_ALL_SYMBOLS ON) diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index c1e1a027..812eb0df 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -779,7 +779,7 @@ void mjCMesh::TryCompile(const mjVFS* vfs) { // compute sdf coefficients if (!plugin.active) { - octree_.ComputeSdfCoeffs(vert_.data(), nvert(), face_.data(), nface()); + octree_.ComputeSdfCoeffs(vert_.data(), nvert(), face_.data(), nface(), tree_); } } diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 06045081..a5c77e9c 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -25,10 +25,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -47,7 +49,6 @@ #include "user/user_model.h" #include "user/user_resource.h" #include "user/user_util.h" -#include namespace { namespace mju = ::mujoco::util; @@ -375,6 +376,7 @@ void mjCBoundingVolumeHierarchy::RemoveInactiveVolumes(int nmax) { bvleaf_.erase(bvleaf_.begin() + nmax, bvleaf_.end()); } + const mjCBoundingVolume* mjCBoundingVolumeHierarchy::AddBoundingVolume(int id, int contype, int conaffinity, const double* pos, const double* quat, @@ -553,6 +555,7 @@ void mjCOctree::CopyLevel(int* level) const { } } + void mjCOctree::CopyChild(int* child) const { for (int i = 0; i < node_.size(); ++i) { for (int j = 0; j < 8; ++j) { @@ -561,6 +564,7 @@ void mjCOctree::CopyChild(int* child) const { } } + void mjCOctree::CopyAabb(mjtNum* aabb) const { for (int i = 0; i < node_.size(); ++i) { aabb[i * 6 + 0] = (node_[i].aamm[0] + node_[i].aamm[3]) / 2; @@ -572,6 +576,7 @@ void mjCOctree::CopyAabb(mjtNum* aabb) const { } } + void mjCOctree::CopyCoeff(mjtNum* coeff) const { for (int i = 0; i < node_.size(); ++i) { for (int j = 0; j < 8; ++j) { @@ -580,6 +585,7 @@ void mjCOctree::CopyCoeff(mjtNum* coeff) const { } } + void mjCOctree::SetFace(const std::vector& vert, const std::vector& face) { for (int i = 0; i < face.size(); i += 3) { std::array v0 = {vert[3*face[i+0]], vert[3*face[i+0]+1], vert[3*face[i+0]+2]}; @@ -622,14 +628,207 @@ void mjCOctree::CreateOctree(const double aamm[6]) { } -// compute SDF coefficients at octree vertices using triangle mesh distance -void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, - const int* face, int nface) { - tmd::TriangleMeshDistance sdf(vert, static_cast(nvert), - face, static_cast(nface)); +namespace { - std::vector coeffs(NumVerts()); - std::vector processed(NumVerts(), false); +double pointBoxDistSq(const double* p, const mjtNum* aabb) { + double dist_sq = 0; + for (int i = 0; i < 3; ++i) { + double lo = aabb[i] - aabb[i + 3]; + double hi = aabb[i] + aabb[i + 3]; + if (p[i] < lo) { + dist_sq += (lo - p[i]) * (lo - p[i]); + } else if (p[i] > hi) { + dist_sq += (p[i] - hi) * (p[i] - hi); + } + } + return dist_sq; +} + + +// compute squared distance between point p and triangle (v0, v1, v2), +// and return barycentric coordinates (u,v) of the closest point +double pointTriDistSqWithUV(const double* p, const double* v0, const double* v1, + const double* v2, double& out_u, double& out_v) { + double ab[3] = {v1[0] - v0[0], v1[1] - v0[1], v1[2] - v0[2]}; + double ac[3] = {v2[0] - v0[0], v2[1] - v0[1], v2[2] - v0[2]}; + double ap[3] = {p[0] - v0[0], p[1] - v0[1], p[2] - v0[2]}; + + // the closest point on the triangle is determined by partitioning space into Voronoi regions + double d1 = ab[0]*ap[0] + ab[1]*ap[1] + ab[2]*ap[2]; + double d2 = ac[0]*ap[0] + ac[1]*ap[1] + ac[2]*ap[2]; + + // region A (vertex v0) + if (d1 <= 0 && d2 <= 0) { + out_u = 0; out_v = 0; + return ap[0]*ap[0] + ap[1]*ap[1] + ap[2]*ap[2]; + } + + double bp[3] = {p[0] - v1[0], p[1] - v1[1], p[2] - v1[2]}; + double d3 = ab[0]*bp[0] + ab[1]*bp[1] + ab[2]*bp[2]; + double d4 = ac[0]*bp[0] + ac[1]*bp[1] + ac[2]*bp[2]; + + // region B (vertex v1) + if (d3 >= 0 && d4 <= d3) { + out_u = 1; out_v = 0; + return bp[0]*bp[0] + bp[1]*bp[1] + bp[2]*bp[2]; + } + + // region AB (edge v0-v1) + double vc = d1*d4 - d3*d2; + if (vc <= 0 && d1 >= 0 && d3 <= 0) { + double u = d1 / (d1 - d3); + out_u = u; out_v = 0; + double closest[3] = {v0[0] + u*ab[0], v0[1] + u*ab[1], v0[2] + u*ab[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + double cp[3] = {p[0] - v2[0], p[1] - v2[1], p[2] - v2[2]}; + double d5 = ab[0]*cp[0] + ab[1]*cp[1] + ab[2]*cp[2]; + double d6 = ac[0]*cp[0] + ac[1]*cp[1] + ac[2]*cp[2]; + + // region C (vertex v2) + if (d6 >= 0 && d5 <= d6) { + out_u = 0; out_v = 1; + return cp[0]*cp[0] + cp[1]*cp[1] + cp[2]*cp[2]; + } + + // region AC (edge v0-v2) + double vb = d5*d2 - d1*d6; + if (vb <= 0 && d2 >= 0 && d6 <= 0) { + double v = d2 / (d2 - d6); + out_u = 0; out_v = v; + double closest[3] = {v0[0] + v*ac[0], v0[1] + v*ac[1], v0[2] + v*ac[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + // region BC (edge v1-v2) + double va = d3*d6 - d5*d4; + if (va <= 0 && (d4 - d3) >= 0 && (d5 - d6) >= 0) { + double w = (d4 - d3) / ((d4 - d3) + (d5 - d6)); + out_u = 1 - w; out_v = w; + double bc[3] = {v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2]}; + double closest[3] = {v1[0] + w*bc[0], v1[1] + w*bc[1], v1[2] + w*bc[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + // region ABC (inside triangle) + double denom = 1.0 / (va + vb + vc); + double u = vb * denom; + double v = vc * denom; + out_u = u; out_v = v; + double closest[3] = {v0[0] + u*ab[0] + v*ac[0], + v0[1] + u*ab[1] + v*ac[1], + v0[2] + u*ab[2] + v*ac[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); +} + + +// query BVH for closest face to point p, return distance, face index and barycentric coordinates +void queryClosestBVHWithFace(const mjtNum* bvh, const int* child, const int* nodeid, + const double* vert, const int* face, int node_idx, + const double* p, double& best_dist_sq, + int& best_face, double& best_u, double& best_v) { + const mjtNum* aabb = &bvh[node_idx * 6]; + if (pointBoxDistSq(p, aabb) >= best_dist_sq) return; + + int left = child[node_idx * 2]; + int right = child[node_idx * 2 + 1]; + + if (left == -1 && right == -1) { + int fi = nodeid[node_idx]; + if (fi >= 0) { + const double* v0 = vert + face[fi * 3 + 0] * 3; + const double* v1 = vert + face[fi * 3 + 1] * 3; + const double* v2 = vert + face[fi * 3 + 2] * 3; + double u, v; + double dist_sq = pointTriDistSqWithUV(p, v0, v1, v2, u, v); + if (dist_sq < best_dist_sq) { + best_dist_sq = dist_sq; + best_face = fi; + best_u = u; + best_v = v; + } + } + return; + } + + if (left >= 0) { + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, left, p, + best_dist_sq, best_face, best_u, best_v); + } + if (right >= 0) { + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, right, p, + best_dist_sq, best_face, best_u, best_v); + } +} + + +double querySignedDistance(const mjtNum* bvh, const int* child, const int* nodeid, + int nbvh, const double* point, + const double* vert, const int* face) { + if (nbvh == 0) { + return 0; + } + + double best_dist_sq = 1e20; + int best_face = -1; + double best_u = 0, best_v = 0; + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, 0, point, + best_dist_sq, best_face, best_u, best_v); + double dist = std::sqrt(best_dist_sq); + + double sign = 1.0; + if (best_face >= 0) { + const double* v0 = vert + face[best_face * 3 + 0] * 3; + const double* v1 = vert + face[best_face * 3 + 1] * 3; + const double* v2 = vert + face[best_face * 3 + 2] * 3; + + double e1[3] = {v1[0]-v0[0], v1[1]-v0[1], v1[2]-v0[2]}; + double e2[3] = {v2[0]-v0[0], v2[1]-v0[1], v2[2]-v0[2]}; + double normal[3] = { + e1[1]*e2[2] - e1[2]*e2[1], + e1[2]*e2[0] - e1[0]*e2[2], + e1[0]*e2[1] - e1[1]*e2[0] + }; + + double closest[3] = { + v0[0] + best_u*(v1[0]-v0[0]) + best_v*(v2[0]-v0[0]), + v0[1] + best_u*(v1[1]-v0[1]) + best_v*(v2[1]-v0[1]), + v0[2] + best_u*(v1[2]-v0[2]) + best_v*(v2[2]-v0[2]) + }; + + double u[3] = {point[0]-closest[0], point[1]-closest[1], point[2]-closest[2]}; + double dot = u[0]*normal[0] + u[1]*normal[1] + u[2]*normal[2]; + double normal_len = mjuu_normvec(normal, 3); + double eps = 1e-12 * normal_len * dist; + sign = (dot > eps) ? 1.0 : -1.0; + } + + return sign * dist; +} + +} // namespace + + +double mjCBoundingVolumeHierarchy::QuerySignedDistance( + const double* point, const double* vert, const int* face) const { + return querySignedDistance(bvh_.data(), child_.data(), nodeid_.data(), + nbvh_, point, vert, face); +} + + +void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface, + const mjCBoundingVolumeHierarchy& tree) { + std::vector coeffs(nvert_, 0.0); + std::vector processed(nvert_, false); std::deque queue; if (NumNodes() > 0) { @@ -647,8 +846,16 @@ void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, continue; } if (Hang(vert_id).empty()) { - coeffs[vert_id] = sdf.signed_distance(Vert(vert_id)).distance; + // transform from octree frame (body inertial) back to mesh frame + double p_mesh[3]; + mjuu_rotVecQuat(p_mesh, Vert(vert_id), iquat_); + p_mesh[0] += ipos_[0]; + p_mesh[1] += ipos_[1]; + p_mesh[2] += ipos_[2]; + + coeffs[vert_id] = tree.QuerySignedDistance(p_mesh, vert, face); } else { + // hanging node: interpolate from parents double sum_coeff = 0; for (int dep_id : Hang(vert_id)) { sum_coeff += coeffs[dep_id]; @@ -666,6 +873,42 @@ void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, } } + // optional Laplacian smoothing (smooths octree level transitions) + if (smoothing_iterations_ > 0) { + // build vertex neighbor graph from octree connectivity + std::vector> neighbors(nvert_); + for (int i = 0; i < NumNodes(); ++i) { + static const int edges[12][2] = { + {0, 1}, {2, 3}, {4, 5}, {6, 7}, + {0, 2}, {1, 3}, {4, 6}, {5, 7}, + {0, 4}, {1, 5}, {2, 6}, {3, 7} + }; + for (const auto& edge : edges) { + int v0 = VertId(i, edge[0]); + int v1 = VertId(i, edge[1]); + neighbors[v0].insert(v1); + neighbors[v1].insert(v0); + } + } + + // apply Laplacian smoothing + const double alpha = 0.2; + std::vector sdf_new(nvert_); + for (int iter = 0; iter < smoothing_iterations_; ++iter) { + for (int i = 0; i < nvert_; ++i) { + if (neighbors[i].empty()) { + sdf_new[i] = coeffs[i]; + } else { + double avg = 0; + for (int j : neighbors[i]) avg += coeffs[j]; + avg /= neighbors[i].size(); + sdf_new[i] = (1 - alpha) * coeffs[i] + alpha * avg; + } + } + std::swap(coeffs, sdf_new); + } + } + // copy coefficients to the octree nodes for (int i = 0; i < NumNodes(); ++i) { for (int j = 0; j < 8; j++) { @@ -680,7 +923,7 @@ static double dot2(const double* a, const double* b) { } -// From M. Schwarz and H.-P. Seidel, "Fast Parallel Surface and Solid Voxelization on GPUs". +// from M. Schwarz and H.-P. Seidel, "Fast Parallel Surface and Solid Voxelization on GPUs". static bool boxTriangle(const Triangle& v, const double aamm[6]) { // bounding box tests for (int i = 0; i < 3; i++) { diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 2e5ad168..8399a146 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -200,6 +200,10 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ { + sizeof(int) * nodeid_.size() + sizeof(int) * level_.size(); } + // query signed distance from point to mesh surface + double QuerySignedDistance(const double* point, const double* vert, + const int* face) const; + private: // internal class used during BVH construction, for partial sorting of bounding volumes struct BVElement { @@ -275,6 +279,7 @@ struct mjCOctree_ { std::vector> hang_; // hanging nodes status (nvert x 1) double ipos_[3] = {0, 0, 0}; double iquat_[4] = {1, 0, 0, 0}; + int smoothing_iterations_ = 0; // Laplacian smoothing iterations (0 = disabled) }; class mjCOctree : public mjCOctree_ { @@ -303,8 +308,13 @@ class mjCOctree : public mjCOctree_ { void AddCoeff(int n, int v, double coeff) { node_[n].coeff[v] = coeff; } double Coeff(int n, int v) const { return node_[n].coeff[v]; } - // compute SDF coefficients at octree vertices using triangle mesh distance - void ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface); + // Set number of Laplacian smoothing iterations (0 = disabled, default) + void SetSmoothingIterations(int iterations) { smoothing_iterations_ = iterations; } + int SmoothingIterations() const { return smoothing_iterations_; } + + // compute SDF coefficients via BVH queries, optionally with Laplacian smoothing + void ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface, + const mjCBoundingVolumeHierarchy& tree); private: void Make(std::vector& elements); diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc index 41528d23..a86fc6b6 100644 --- a/test/user/user_objects_test.cc +++ b/test/user/user_objects_test.cc @@ -2682,5 +2682,150 @@ TEST_F(UserObjectsTest, ZeroMass) { mj_deleteModel(model); } + +// ------------- test Octree SDF computation ----------------------------------- + +using OctreeSDFTest = MujocoTest; + +TEST_F(OctreeSDFTest, SphereSDF) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + EXPECT_GT(model->nmesh, 0); + EXPECT_EQ(model->geom_type[0], mjGEOM_SDF); + + int geom_id = 0; + int mesh_id = model->geom_dataid[geom_id]; + mjSDF sdf; + const mjpPlugin* null_plugin = nullptr; + sdf.plugin = &null_plugin; + sdf.id = &mesh_id; + sdf.type = mjSDFTYPE_SINGLE; + sdf.geomtype = (mjtGeom*)(model->geom_type + geom_id); + + // Analytic SDF for unit sphere: distance = |p| - 1 + auto analyticSdf = [](const mjtNum* p) -> double { + return mju_sqrt(p[0]*p[0] + p[1]*p[1] + p[2]*p[2]) - 1.0; + }; + + int sign_errors = 0; + int total_points = 0; + double sum_sq_error = 0.0; + + // Test grid of points + for (double x = -2.0; x <= 2.0; x += 0.5) { + for (double y = -2.0; y <= 2.0; y += 0.5) { + for (double z = -2.0; z <= 2.0; z += 0.5) { + mjtNum p[3] = {x, y, z}; + double sdf_dist = mjc_distance(model, data, &sdf, p); + double gt_dist = analyticSdf(p); + + if ((sdf_dist < 0) != (gt_dist < 0)) { + sign_errors++; + } + + double error = sdf_dist - gt_dist; + sum_sq_error += error * error; + total_points++; + } + } + } + + double rmse = mju_sqrt(sum_sq_error / total_points); + + EXPECT_LT(sign_errors, total_points / 200) + << "No more than 0.5% of points should have sign errors"; + EXPECT_LT(rmse, 0.11) << "RMSE should be less than 0.11"; + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(OctreeSDFTest, TorusSDF) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + EXPECT_GT(model->nmesh, 0); + EXPECT_EQ(model->geom_type[0], mjGEOM_SDF); + + int geom_id = 0; + int mesh_id = model->geom_dataid[geom_id]; + mjSDF sdf; + const mjpPlugin* null_plugin = nullptr; + sdf.plugin = &null_plugin; + sdf.id = &mesh_id; + sdf.type = mjSDFTYPE_SINGLE; + sdf.geomtype = (mjtGeom*)(model->geom_type + geom_id); + + // Analytic SDF for torus: distance = |p_proj| - r, where p_proj is + // projection of p onto circle of radius R, and r is minor radius. + // R=1, r=0.3 + auto analyticSdf = [](const mjtNum* p) -> double { + double xy = mju_sqrt(p[0]*p[0] + p[1]*p[1]); + double vec[2] = {xy - 1.0, p[2]}; + return mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]) - 0.3; + }; + + int sign_errors = 0; + int total_points = 0; + double sum_sq_error = 0.0; + + // Test grid of points + for (double x = -2.0; x <= 2.0; x += 0.5) { + for (double y = -2.0; y <= 2.0; y += 0.5) { + for (double z = -2.0; z <= 2.0; z += 0.5) { + mjtNum p[3] = {x, y, z}; + double sdf_dist = mjc_distance(model, data, &sdf, p); + double gt_dist = analyticSdf(p); + + if ((sdf_dist < 0) != (gt_dist < 0)) { + sign_errors++; + } + + double error = sdf_dist - gt_dist; + sum_sq_error += error * error; + total_points++; + } + } + } + + double rmse = mju_sqrt(sum_sq_error / total_points); + + EXPECT_LT(sign_errors, total_points / 20) + << "No more than 5% of points should have sign errors"; + EXPECT_LT(rmse, 0.52) << "RMSE should be close to 0.516"; + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco From 757758dd2e610755bacd6305cd8c047e53b46790 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Wed, 25 Feb 2026 05:43:22 -0800 Subject: [PATCH 47/48] Rename all filament mjr_ functions to mjrf_. Rename stubs.cc to mjr_compat.cc. This library acts like a "compatibility layer" that allows the filament renderer to be implemented behind the mjr API. PiperOrigin-RevId: 875114054 Change-Id: I5c47581cacbcadff4c78d2ad15728cc1c3a394ad --- src/experimental/filament/CMakeLists.txt | 2 +- .../filament/{stubs.cc => mjr_compat.cc} | 41 ++++++++++++++++++ .../filament/render_context_filament.cc | 42 ++++++++++--------- .../filament/render_context_filament.h | 39 +++++++++-------- src/experimental/platform/renderer.cc | 14 +++---- 5 files changed, 93 insertions(+), 45 deletions(-) rename src/experimental/filament/{stubs.cc => mjr_compat.cc} (65%) diff --git a/src/experimental/filament/CMakeLists.txt b/src/experimental/filament/CMakeLists.txt index 99854dfc..17ebec93 100644 --- a/src/experimental/filament/CMakeLists.txt +++ b/src/experimental/filament/CMakeLists.txt @@ -21,7 +21,7 @@ target_compile_definitions(${MUJOCO_FILAMENT_TARGET_NAME} PRIVATE MJ_STATIC) target_sources(${MUJOCO_FILAMENT_TARGET_NAME} PUBLIC - stubs.cc + mjr_compat.cc render_context_filament.h render_context_filament.cc filament/buffer_util.cc diff --git a/src/experimental/filament/stubs.cc b/src/experimental/filament/mjr_compat.cc similarity index 65% rename from src/experimental/filament/stubs.cc rename to src/experimental/filament/mjr_compat.cc index 9d2d099c..9d06dea6 100644 --- a/src/experimental/filament/stubs.cc +++ b/src/experimental/filament/mjr_compat.cc @@ -13,9 +13,50 @@ // limitations under the License. #include +#include "experimental/filament/render_context_filament.h" + +// This library implements the entirety of mujoco's mjr API. You can link this +// library with your application (instead of the "classic" mujoco renderer) to +// use the same APIs but with Filament rendering instead. +// +// However, you should consider using filament's mjrf API directly as it will +// provide you with access to more features and optimizations. extern "C" { +// mjr functions that are supported by the filament renderer. + +void mjr_defaultContext(mjrContext* con) { + mjrf_defaultContext(con); +} +void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale) { + mjrf_makeContext(m, con, fontscale); +} +void mjr_freeContext(mjrContext* con) { + mjrf_freeContext(con); +} +void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { + mjrf_render(viewport, scn, con); +} +void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { + mjrf_uploadMesh(m, con, meshid); +} +void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { + mjrf_uploadTexture(m, con, texid); +} +void mjr_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { + mjrf_uploadHField(m, con, hfieldid); +} +void mjr_setBuffer(int framebuffer, mjrContext* con) { + mjrf_setBuffer(framebuffer, con); +} +void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, + const mjrContext* con) { + mjrf_readPixels(rgb, depth, viewport, con); +} + +// mjr functions that are NOT supported by the filament renderer. + void mjr_setAux(int index, const mjrContext* con) { mju_error("mjr_setAux not implemented."); } diff --git a/src/experimental/filament/render_context_filament.cc b/src/experimental/filament/render_context_filament.cc index 4695599f..3955a816 100644 --- a/src/experimental/filament/render_context_filament.cc +++ b/src/experimental/filament/render_context_filament.cc @@ -32,17 +32,17 @@ static mujoco::FilamentContext* g_filament_context = nullptr; static void CheckFilamentContext() { if (g_filament_context == nullptr) { - mju_error("Missing context; did you call mjr_makeFilamentContext?"); + mju_error("Missing context; did you call mjrf_makeFilamentContext?"); } } extern "C" { -void mjr_defaultFilamentConfig(mjrFilamentConfig* config) { +void mjrf_defaultFilamentConfig(mjrFilamentConfig* config) { memset(config, 0, sizeof(mjrFilamentConfig)); } -void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, +void mjrf_makeFilamentContext(const mjModel* m, mjrContext* con, const mjrFilamentConfig* config) { // TODO: Support multiple contexts and multiple threads. For now, we'll just // assume a single, global context. @@ -52,16 +52,18 @@ void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, g_filament_context = new mujoco::FilamentContext(config, m, con); } -void mjr_defaultContext(mjrContext* con) { memset(con, 0, sizeof(mjrContext)); } - -void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale) { - mjr_freeContext(con); - mjrFilamentConfig cfg; - mjr_defaultFilamentConfig(&cfg); - mjr_makeFilamentContext(m, con, &cfg); +void mjrf_defaultContext(mjrContext* con) { + memset(con, 0, sizeof(mjrContext)); } -void mjr_freeContext(mjrContext* con) { +void mjrf_makeContext(const mjModel* m, mjrContext* con, int fontscale) { + mjr_freeContext(con); + mjrFilamentConfig cfg; + mjrf_defaultFilamentConfig(&cfg); + mjrf_makeFilamentContext(m, con, &cfg); +} + +void mjrf_freeContext(mjrContext* con) { // mjr_freeContext may be called multiple times. if (g_filament_context) { delete g_filament_context; @@ -70,50 +72,50 @@ void mjr_freeContext(mjrContext* con) { mjr_defaultContext(con); } -void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { +void mjrf_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { CheckFilamentContext(); g_filament_context->Render(viewport, scn, con); } -void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { +void mjrf_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { CheckFilamentContext(); g_filament_context->UploadMesh(m, meshid); } -void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { +void mjrf_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { CheckFilamentContext(); g_filament_context->UploadTexture(m, texid); } -void mjr_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { +void mjrf_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { CheckFilamentContext(); g_filament_context->UploadHeightField(m, hfieldid); } -void mjr_setBuffer(int framebuffer, mjrContext* con) { +void mjrf_setBuffer(int framebuffer, mjrContext* con) { CheckFilamentContext(); g_filament_context->SetFrameBuffer(framebuffer); } -void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, +void mjrf_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, const mjrContext* con) { CheckFilamentContext(); g_filament_context->ReadPixels(viewport, rgb, depth); } -uintptr_t mjr_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, +uintptr_t mjrf_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, int width, int height, int bpp, const mjrContext* con) { CheckFilamentContext(); return g_filament_context->UploadGuiImage(tex_id, pixels, width, height, bpp); } -double mjr_getFrameRate(const mjrContext* con) { +double mjrf_getFrameRate(const mjrContext* con) { CheckFilamentContext(); return g_filament_context->GetFrameRate(); } -void mjr_updateGui(const mjrContext* con) { +void mjrf_updateGui(const mjrContext* con) { if (g_filament_context != nullptr) { g_filament_context->UpdateGui(); } diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h index c0423e4a..aef8f182 100644 --- a/src/experimental/filament/render_context_filament.h +++ b/src/experimental/filament/render_context_filament.h @@ -26,6 +26,9 @@ extern "C" { #endif +// IMPORTANT: This API should still be considered experimental and is likely +// change frequently. + typedef enum mjtGraphicsApi_ { // backend graphics API to use mjGFX_DEFAULT = 0, // default based on platform mjGFX_OPENGL, // OpenGL (desktop) @@ -43,35 +46,37 @@ struct mjrFilamentConfig { bool enable_gui; }; -void mjr_defaultFilamentConfig(mjrFilamentConfig* config); +void mjrf_defaultFilamentConfig(mjrFilamentConfig* config); -void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, - const mjrFilamentConfig* config); +void mjrf_makeFilamentContext(const mjModel* m, mjrContext* con, + const mjrFilamentConfig* config); -MJAPI void mjr_defaultContext(mjrContext* con); +void mjrf_defaultContext(mjrContext* con); -MJAPI void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale); +void mjrf_makeContext(const mjModel* m, mjrContext* con, int fontscale); -MJAPI void mjr_freeContext(mjrContext* con); +void mjrf_freeContext(mjrContext* con); -MJAPI void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con); +void mjrf_render(mjrRect viewport, mjvScene* scn, const mjrContext* con); -MJAPI void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid); +void mjrf_uploadMesh(const mjModel* m, const mjrContext* con, int meshid); -MJAPI void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid); +void mjrf_uploadTexture(const mjModel* m, const mjrContext* con, int texid); -MJAPI void mjr_setBuffer(int framebuffer, mjrContext* con); +void mjrf_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid); -MJAPI void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, - const mjrContext* con); +void mjrf_setBuffer(int framebuffer, mjrContext* con); -double mjr_getFrameRate(const mjrContext* con); +void mjrf_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, + const mjrContext* con); -uintptr_t mjr_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, - int width, int height, int bpp, - const mjrContext* con); +double mjrf_getFrameRate(const mjrContext* con); -void mjr_updateGui(const mjrContext* con); +uintptr_t mjrf_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, + int width, int height, int bpp, + const mjrContext* con); + +void mjrf_updateGui(const mjrContext* con); #if defined(__cplusplus) } // extern "C" diff --git a/src/experimental/platform/renderer.cc b/src/experimental/platform/renderer.cc index 6fe16730..2188d86f 100644 --- a/src/experimental/platform/renderer.cc +++ b/src/experimental/platform/renderer.cc @@ -56,7 +56,7 @@ void Renderer::Init(const mjModel* model) { mjr_makeContext(model, &render_context_, mjFONTSCALE_150); #else mjrFilamentConfig render_config; - mjr_defaultFilamentConfig(&render_config); + mjrf_defaultFilamentConfig(&render_config); render_config.native_window = native_window_; render_config.enable_gui = true; #if defined(MUJOCO_RENDERER_FILAMENT_OPENGL) @@ -66,7 +66,7 @@ void Renderer::Init(const mjModel* model) { #elif defined(MUJOCO_RENDERER_FILAMENT_VULKAN) render_config.graphics_api = mjGFX_VULKAN; #endif - mjr_makeFilamentContext(model, &render_context_, &render_config); + mjrf_makeFilamentContext(model, &render_context_, &render_config); #endif mjv_defaultScene(&scene_); @@ -171,9 +171,9 @@ int Renderer::UploadImage(int texture_id, const std::byte* pixels, int width, #if defined(MUJOCO_RENDERER_CLASSIC_OPENGL) return 0; #else - return mjr_uploadGuiImage(texture_id, - reinterpret_cast(pixels), - width, height, bpp, &render_context_); + return mjrf_uploadGuiImage(texture_id, + reinterpret_cast(pixels), + width, height, bpp, &render_context_); #endif } @@ -191,7 +191,7 @@ void Renderer::UpdateFps() { frames_ = 0; } #else - fps_ = mjr_getFrameRate(&render_context_); + fps_ = mjrf_getFrameRate(&render_context_); #endif } @@ -217,7 +217,7 @@ mjPLUGIN_LIB_INIT { mujoco::platform::GuiPlugin plugin; plugin.name = "Filament"; plugin.update = [](mujoco::platform::GuiPlugin* self) { - mjr_updateGui(nullptr); + mjrf_updateGui(nullptr); }; mujoco::platform::RegisterPlugin(plugin); } From 940cab25082dfa4b10a3fceade9bd35340185c22 Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Wed, 25 Feb 2026 06:50:08 -0800 Subject: [PATCH 48/48] Rename MuJoCo WASM output artifacts to `mujoco.*`. This change updates the Emscripten build to output files named `mujoco.js`, `mujoco.wasm`, and `mujoco.d.ts` instead of `mujoco_wasm.*`. The internal CMake target name remains `mujoco_wasm` to avoid conflicts. The package name in `package.json` is also updated to `mujoco`. New enum tests are added. PiperOrigin-RevId: 875140314 Change-Id: Ie03e0d1db2d03dd8582d5bfcb260b3d33549c2c9 --- wasm/CMakeLists.txt | 10 ++++++++-- wasm/codegen/tests/enums_test_generator.py | 4 ++-- wasm/demo_app/app.ts | 2 +- wasm/package-lock.json | 4 ++-- wasm/package.json | 2 +- wasm/tests/bindings_test.ts | 4 ++-- wasm/tests/enums_test.ts | 16 ++++++++++++++-- wasm/tests/sandbox/main.ts | 4 ++-- 8 files changed, 32 insertions(+), 14 deletions(-) diff --git a/wasm/CMakeLists.txt b/wasm/CMakeLists.txt index 2f2c77a8..0adcf7a0 100644 --- a/wasm/CMakeLists.txt +++ b/wasm/CMakeLists.txt @@ -46,13 +46,19 @@ set(EMCC_LINKER_FLAGS "-s DISABLE_EXCEPTION_CATCHING=0" "-gsource-map" "-g" - "--emit-tsd mujoco_wasm.d.ts" + "--emit-tsd mujoco.d.ts" ) string (REPLACE ";" " " EMCC_LINKER_FLAGS_STR "${EMCC_LINKER_FLAGS}") add_executable(mujoco_wasm ${MUJOCO_WASM_FILES}) -set_target_properties(mujoco_wasm PROPERTIES LINK_FLAGS "${EMCC_LINKER_FLAGS_STR}") +# Keep the internal target name distinct to avoid colliding with the native +# `mujoco` library target, but emit artifacts named `mujoco.*` by setting the +# output name. Also apply the emscripten linker flags to the wasm target. +set_target_properties(mujoco_wasm PROPERTIES + LINK_FLAGS "${EMCC_LINKER_FLAGS_STR}" + OUTPUT_NAME "mujoco" +) target_link_libraries(mujoco_wasm ccd lodepng mujoco tinyxml2 qhullstatic_r) diff --git a/wasm/codegen/tests/enums_test_generator.py b/wasm/codegen/tests/enums_test_generator.py index f380aabc..862b15e3 100644 --- a/wasm/codegen/tests/enums_test_generator.py +++ b/wasm/codegen/tests/enums_test_generator.py @@ -26,8 +26,8 @@ def generate_typescript_enum_tests(): output = textwrap.dedent("""\ import 'jasmine'; - import { MainModule } from "../dist/mujoco_wasm" - import loadMujoco from "../dist/mujoco_wasm.js" + import { MainModule } from "../dist/mujoco" + import loadMujoco from "../dist/mujoco.js" let mujoco: MainModule; diff --git a/wasm/demo_app/app.ts b/wasm/demo_app/app.ts index 3e04f42c..9e5a1a9f 100644 --- a/wasm/demo_app/app.ts +++ b/wasm/demo_app/app.ts @@ -14,7 +14,7 @@ import * as THREE from "three" import { OrbitControls } from "three/examples/jsm/controls/OrbitControls.js" -import loadMujoco from "../dist/mujoco_wasm.js" +import loadMujoco from "../dist/mujoco.js" declare function loadMujoco(): Promise; diff --git a/wasm/package-lock.json b/wasm/package-lock.json index 859f393d..a21344d4 100644 --- a/wasm/package-lock.json +++ b/wasm/package-lock.json @@ -1,11 +1,11 @@ { - "name": "mujoco_wasm", + "name": "mujoco", "version": "1.0.0-alpha.1", "lockfileVersion": 3, "requires": true, "packages": { "": { - "name": "mujoco_wasm", + "name": "mujoco", "version": "1.0.0-alpha.1", "license": "Apache-2.0", "devDependencies": { diff --git a/wasm/package.json b/wasm/package.json index 6384da1b..fd2587b0 100644 --- a/wasm/package.json +++ b/wasm/package.json @@ -1,5 +1,5 @@ { - "name": "mujoco_wasm", + "name": "mujoco", "version": "1.0.0-alpha.1", "description": "MuJoCo JavaScript Bindings", "directories": { diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index 0ed482cc..4cff1ec9 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -17,9 +17,9 @@ import 'jasmine'; import {MainModule, MjContact, MjContactVec, MjData, MjLROpt, MjModel, MjOption, MjsGeom, MjSolverStat, MjSpec, MjStatistic, MjTimerStat, MjvCamera, MjvFigure, MjvGeom, MjvGLCamera, MjvLight, MjvOption, MjvPerturb, MjvScene, -MjWarningStat, MjVFS, Uint8Buffer} from '../dist/mujoco_wasm.js'; +MjWarningStat, MjVFS, Uint8Buffer} from '../dist/mujoco.js'; -import loadMujoco from '../dist/mujoco_wasm.js' +import loadMujoco from '../dist/mujoco.js' function assertExists(value: T | null | undefined, message?: string): asserts value is T { diff --git a/wasm/tests/enums_test.ts b/wasm/tests/enums_test.ts index 7c966e92..d80fb48e 100644 --- a/wasm/tests/enums_test.ts +++ b/wasm/tests/enums_test.ts @@ -14,8 +14,8 @@ import 'jasmine'; -import { MainModule } from "../dist/mujoco_wasm" -import loadMujoco from "../dist/mujoco_wasm.js" +import { MainModule } from "../dist/mujoco" +import loadMujoco from "../dist/mujoco.js" let mujoco: MainModule; @@ -40,6 +40,10 @@ describe('Enums', () => { expect(mujoco.mjtGeom).toBeDefined(); }); + it('mjtProjection should exist', () => { + expect(mujoco.mjtProjection).toBeDefined(); + }); + it('mjtCamLight should exist', () => { expect(mujoco.mjtCamLight).toBeDefined(); }); @@ -120,6 +124,14 @@ describe('Enums', () => { expect(mujoco.mjtConDataField).toBeDefined(); }); + it('mjtRayDataField should exist', () => { + expect(mujoco.mjtRayDataField).toBeDefined(); + }); + + it('mjtCamOutBit should exist', () => { + expect(mujoco.mjtCamOutBit).toBeDefined(); + }); + it('mjtSameFrame should exist', () => { expect(mujoco.mjtSameFrame).toBeDefined(); }); diff --git a/wasm/tests/sandbox/main.ts b/wasm/tests/sandbox/main.ts index 79702e5f..60fc5e95 100644 --- a/wasm/tests/sandbox/main.ts +++ b/wasm/tests/sandbox/main.ts @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -import { MainModule, MjData, MjModel } from "../../dist/mujoco_wasm" -import loadMujoco from "../../dist/mujoco_wasm.js" +import { MainModule, MjData, MjModel } from "../../dist/mujoco" +import loadMujoco from "../../dist/mujoco.js" declare function loadMujoco(): Promise;