Port smooth contact autodifferentiation (3/3) onto current google-deepmind/main - #11
Open
johnnynunez wants to merge 131 commits into
Open
Conversation
* Add cubql constructor for mesh/hfield * Ruff * Gate to specific nightly * Make check_version inline
Bumps and [brace-expansion](https://github.com/juliangruber/brace-expansion). These dependencies needed to be updated together. Updates `brace-expansion` from 2.0.2 to 2.0.3 - [Release notes](https://github.com/juliangruber/brace-expansion/releases) - [Commits](juliangruber/brace-expansion@v2.0.2...v2.0.3) Updates `brace-expansion` from 5.0.3 to 5.0.5 - [Release notes](https://github.com/juliangruber/brace-expansion/releases) - [Commits](juliangruber/brace-expansion@v2.0.2...v2.0.3) Updates `brace-expansion` from 1.1.12 to 1.1.13 - [Release notes](https://github.com/juliangruber/brace-expansion/releases) - [Commits](juliangruber/brace-expansion@v2.0.2...v2.0.3) --- updated-dependencies: - dependency-name: brace-expansion dependency-version: 2.0.3 dependency-type: indirect - dependency-name: brace-expansion dependency-version: 5.0.5 dependency-type: indirect - dependency-name: brace-expansion dependency-version: 1.1.13 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
* Fix flex bugs for multi flex and 1d * Update kernel arg order * Fix kernel arg order
Bumps [pygments](https://github.com/pygments/pygments) from 2.19.2 to 2.20.0. - [Release notes](https://github.com/pygments/pygments/releases) - [Changelog](https://github.com/pygments/pygments/blob/master/CHANGES) - [Commits](pygments/pygments@2.19.2...2.20.0) --- updated-dependencies: - dependency-name: pygments dependency-version: 2.20.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
* Fused solve_LD_sparse kernels * Fix bug in new fused kernel * Fix formatting * Fix new kernel on CPU * Change to use wp.block_dim
* Expose Callback in imports. * update import order --------- Co-authored-by: Taylor Howell <taylorhowell@google.com>
* per-world meshes * geom mass fixes * create a copy of spec in per_world_mesh * updates for renderer * update per_world_mesh
…ind#1273) Warp 1.12 introduced bracket syntax for array type annotations (wp.array[X] instead of wp.array(dtype=X)). The parenthesized form causes false positives from static type checkers that interpret it as a function call. This migrates all annotations to bracket syntax. Transformations applied: - wp.array(dtype=X) -> wp.array[X] - wp.array(dtype=X, ndim=N) -> wp.arrayNd[X] - wp.arrayNd(dtype=X) -> wp.arrayNd[X] Runtime constructor calls (wp.array(dtype=X, shape=...)) are unchanged. The kernel analyzer now enforces bracket syntax: a new ParenthesizedArraySyntax issue flags any kernel/func parameter still using the old form. Expected types from types.py field annotations (both custom array() and wp.array(dtype=X)) are normalized to bracket syntax. Code in io.py and tests that used isinstance(f.type, wp.array) to identify array field specs is updated to also handle _ArrayAnnotation, which is what bracket syntax produces at class definition time.
* write only lower triangle in _JTDAJ_sparse The H matrix assembly in _JTDAJ_sparse was writing both triangles (two atomic adds per off-diagonal pair), but the Cholesky factorization only reads the lower triangle. Write a single entry to the lower triangle instead, halving the number of atomic adds for off-diagonal elements. * use min/max for single atomic add instead of branch
* Add mjviser web viewer backend (--viewer=viser). * Address reviewer comments.
Bumps [pytest](https://github.com/pytest-dev/pytest) from 8.4.2 to 9.0.3. - [Release notes](https://github.com/pytest-dev/pytest/releases) - [Changelog](https://github.com/pytest-dev/pytest/blob/main/CHANGELOG.rst) - [Commits](pytest-dev/pytest@8.4.2...9.0.3) --- updated-dependencies: - dependency-name: pytest dependency-version: 9.0.3 dependency-type: direct:production ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Warp's module="unique" kernel dedup now correctly handles all the cases that cache_kernel was working around. The decorator and its _KERNEL_CACHE dict are no longer needed. Warp computes a content hash from the kernel source, resolved wp.static() expressions, and transitively referenced functions. Identical factory arguments produce identical hashes, so the same kernel object is returned without recompilation. The previous attempt to remove this (google-deepmind#1282, reverted in google-deepmind#1287) failed because the branch was missing the wp.array annotation migration (google-deepmind#1273), which caused different kernel source text and thus different CUDA codegen. On current main all 727 GPU tests pass without the decorator, both sequentially and with pytest-xdist -n 2.
…pmind#1297) * skip Cholesky refactorization when no constraints changed In the incremental Newton solver (pyramidal cone), ~88% of solver iterations have zero constraint state changes. The blocked Cholesky kernel previously refactorized H every iteration even when H was unchanged. Since the blocked path already caches L in hfactor, we can skip the factorization and go straight to the solve. Add update_gradient_cholesky_blocked_skip_unchanged, a variant of the blocked Cholesky kernel that checks changed_efc_count per world and only refactorizes when > 0. Use it from _update_gradient_incremental for the blocked path (nv > 32). Benchmarks (RTX PRO 6000, 8192 worlds): - three_humanoids (nv=81): solver kernel 1.9x faster, ~2% step speedup - G1 on plane (nv=35): solver kernel 1.3x faster, ~2% step speedup - humanoid (nv=27): unchanged (uses tiled path, not blocked) * fix docstring line length * Address review feedback on skip-refactor-k0 - Remove @cache_kernel decorator (now redundant after google-deepmind#1296 landed upstream; decorator no longer exists). - Rename 'cached L in hfactor' to 'cached factorization in hfactor' in the _cholesky_factorize_solve docstring per thowell's suggestion.
Bumps [lxml](https://github.com/lxml/lxml) from 6.0.2 to 6.1.0. - [Release notes](https://github.com/lxml/lxml/releases) - [Changelog](https://github.com/lxml/lxml/blob/master/CHANGES.txt) - [Commits](lxml/lxml@lxml-6.0.2...lxml-6.1.0) --- updated-dependencies: - dependency-name: lxml dependency-version: 6.1.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…eepmind#1398) * Parallelize CG beta calculation using atomic accumulations * Address reviewer comments: write unified zeroing kernels for beta and beta_den * Fix kernel parameter comments to satisfy MuJoCo Warp JIT compiler requirements
* Plumb lighting params to warp renderer and utilize to reach parity with opengl renderer make background color controllable in create_render_context api make light params batchable make material params batchable pr comments PR comments PR comments add render sanity tests remove lighting field prop test make max shininess a constant * linter fix * use max shininess
* Fix line search cost precision Evaluate line-search candidates relative to alpha zero so small improvements are not lost when the absolute objective is large. Store the accepted line-search improvement directly for convergence instead of recomputing it from absolute costs after constraint update. * Remove unused solver cost tracking Solver convergence now uses the line-search improvement directly, so the accumulated objective cost is no longer consumed by runtime code. Drop the internal cost workspace and the cost atomics from constraint update. * Remove unused solver Gauss tracking Drop the stored Gauss cost from solver contexts now that line search tracks shifted objective improvements directly. This removes the per-iteration Gauss update kernels and keeps the line-search Gauss polynomial delta-only. * Fix inactive shifted line-search costs Preserve the zero-cost baseline for one-sided rows that start inactive and become active during line search. Keep the shifted quadratic form fast with an explicit offset, and use cost-only helpers where alpha-zero subtraction does not need derivatives. * Adapt context types to main * Adapt island line search to cleanup * Adapt line-search test to main * Fix island line-search cost rebase * Format solver module * Collapse shifted helpers and clarify offset Drops _eval_pt_direct_shifted (no-offset) into the offset variant with offset=0.0, and renames the offset variant to drop the redundant _offset suffix. Same for the 3-alpha pair. Adds an offset variable name and a comment at the four limit/contact branches that motivates why the offset is needed when alpha=0 was inactive. * Apply ruff-format to solver * Address line-search review comments Hoist the per-contact reads in the parallel line-search kernel so the impratio index, friction, and quad values are loaded once and shared by both elliptic-cost evaluations instead of being re-read for each call. This matches the iterative kernel. Rename the test helper update_mujoco_constraints to update_constraints and rename test_linesearch_accepts_sub_ulp_improvement to test_linesearch_accepts_sub_float32_improvement to avoid the unexplained "ulp" abbreviation.
…ogle-deepmind#1402) * Disable MathDx GEMM for blocked Cholesky tile_matmul The blocked Cholesky performs small (16x16) upper-convention (U^T U) rank-k updates. Warp's register-blocked scalar GEMM is faster than cuBLASDx for that left-transpose pattern at this tile size, and skipping cuBLASDx also avoids its LTO compile cost. Set enable_mathdx_gemm=False via per-kernel module_options on the blocked-Cholesky kernels. Also switch the existing dense-JTDAJ disable to the same per-kernel mechanism and remove the now-unused scoped_mathdx_gemm_disabled helper, so the option lives on the kernel definition rather than mutating the global warp config around each launch. * Tune JTDAJ_dense block_dim to 128 for the scalar GEMM The dense JTDAJ tile_matmul runs on Warp's scalar GEMM (MathDx disabled). Its optimal block_dim for the scalar path is 128, whereas 96 was tuned for the cuBLASDx path. Only affects dense models (nv <= 32).
* Optimize h.zero_ call * Update memory allocation * Remove dangerous return --------- Co-authored-by: Taylor Howell <taylorhowell@google.com>
The island line search was the last solver path still comparing absolute constraint costs. When absolute costs are large, the accepted improvement can fall below float32 resolution, so previous_cost - current_cost rounds to zero and the line search refuses an otherwise-valid step. Evaluate line-search candidates as deltas from alpha=0, matching the monolithic iterative and parallel kernels: drop the constant gauss cost, make the alpha=0 point the zero-cost reference, evaluate each constraint contribution as cost(alpha) - cost(0) via the shifted helpers, and compare candidate costs against zero. Gradient and hessian are offset-free and are left unchanged, so the Newton step and bracketing logic are untouched. Add a direct island line-search unit test that drives the kernel on a single island whose large gauss cost hides a sub-resolution improvement.
_update_gradient_init_h_sparse wrote the mass matrix into the lower triangle of h via M_elemid[i, j], which is only populated where the column is an ancestor of the row (the lower triangle). The constraint term from _JTDAJ_sparse and the Cholesky factorization both use the upper triangle, so the upper-triangle Hessian was assembled without the mass matrix, producing an incorrect Newton Hessian. Transpose the lookup to M_elemid[j, i] so M lands in the upper triangle, and skip writing the lower triangle entirely since it is never read. The solver still converges with the wrong Hessian (Newton degrades toward gradient descent), so existing tests pass; the symptom is a large performance regression on the sparse path.
…model degrees of freedom (google-deepmind#1404)
The dense-Jacobian branch of _update_gradient_JTCJ_island skips an iteration when J[ic1, idof_i] is zero, but in the dim1 != dim2 case the swap-pair contribution hcone * J[ic2, idof_i] * J[ic1, idof_j] still needs to write — the swap-pair-only block under 'if dim1id != dim2id' only ran when J[ic1, i] AND J[ic2, j] were both nonzero. Cells where exactly one of those four J entries was zero silently dropped the corresponding cone contribution. Per cell H[a, b] for an off-diagonal cone pair (dim1, dim2) the contribution should be hcone * (J[ic1, a] * J[ic2, b] + J[ic2, a] * J[ic1, b]) which the monolithic _update_gradient_JTCJ_dense computes correctly (solver.py:2700-2704). This change keeps the original loop (which already handles the first term) and adds a second loop with the same shape but ic1/ic2 swapped, gated on dim1id != dim2id, for the second term. Both loops keep the aggressive J==0 early-skip from the original. The existing scalar _cholesky_factorize_solve_island absorbs the resulting indefinite per-island H via its mid-factorization clamp (s <= 1e-6 -> s = 1e-6), so no existing test fails. The bug becomes visible when the per-island solve uses any Cholesky variant without diagonal clamping (e.g. wp.tile_cholesky on the per-island H sub-block), producing NaN qacc on configurations such as constraints.xml keyframe 2 with elliptic cone + Newton + dense Jacobian.
* Apply changes from cl/925389681: Implement tiled kernels for CG solver helpers * Fix pre-submit failures: remove unused import, re-order solve_cg_finalize parameters and adjust launch arguments * Apply ruff formatting fixes to solver.py * Specialize CG block dimensions and resolve review comments * Fix solver context initialization for island solver * fixed formatting * add _ * Remove dead kernels, redundant zeroing; dynamic block_dim * fixed formatting * remove block_dim from else branch
…d#1413) The sparse elliptic-cone Hessian assembly (JTCJ) previously ran output-stationary: one thread per (contact, dense dof-pair), scanning the contact's column indices to locate each pair. When nefc >> nv the overwhelming majority of those dof-pairs do not appear in the contact's support, so the kernel spent almost all of its time scanning and skipping. Restructure it to be input-stationary: launch one thread per (contact, support-pair), decode the pair index directly into the two participating dofs, register-sum the cone block, and accumulate into H with a single atomic add. This touches only the pairs that actually contribute and exposes far more parallelism. The pair dimension is bounded by jtcj_max_pairs, derived once in io.py from the deepest geom-body dof chain. The math is unchanged. Co-authored-by: Alain Denzler <adenzler@nvidia.com>
The unrolled _process_joint_vel introduced for warp AD compatibility dropped the explicit zeroing of cdof_dot rows 0-2 for free joints, leaving stale or uninitialized values (caught by smooth_test.py::test_com_vel which prefills with inf). Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Semantic port of mar-yan24's mark/autodifferentiation2 (131 commits behind) onto the rebased autodiff branch. Preserves the design intact: - solver_implicit_adjoint: tape-backward hook solving H v = adj_qacc with the retained Newton Hessian, writing adj_qacc_smooth = M v - GPU tile-Cholesky adjoint paths (small-nv single tile, blocked solve-only with stored factor, blocked full factorize+solve) - Data.solver_h / solver_hfactor / solver_Jaref retained state, allocated in make_data/put_data, aliased by the solver context Adaptations to current main: - create_blocked_cholesky_func no longer exists upstream (replaced by fused create_blocked_cholesky_factorize_solve_func in 083e5ef); _adjoint_cholesky_full_blocked uses the fused func - dropped his update_constraint_gauss_cost kernel (belonged to the pre-google-deepmind#1301 solver context which carried ctx.cost; current main computes gauss cost inside the linesearch tiles) - dropped his support.py next_act removal (current main's derivative.py imports it) - solver retained fields placed at the end of Data (types_test enforces warp-only fields after MuJoCo fields, with matching docstring order) - test_solver_retained_state used an undefined SPARSE_CONSTRAINT_JACOBIAN global; replaced with m.is_sparse Verification: full suite 1078 passed / 23 skipped, including his 4 GradSolverAdjointTest cases. Cross-validated independently: the GPU adjoint mapped to force space (M^-1 M H^-1 g) matches central finite differences through the full GPU solver to 8e-5 max relative error on a persistent-contact scene. Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Semantic port of mar-yan24's mark/autodifferentiation3 (123 commits behind) onto the 2/3 port. Preserves the design intact: - collision_smooth.py: smooth analytic distance functions overwrite discrete contact geometry (dist/pos/frame) during taped forwards; differentiable constraint assembly (smooth_contact_to_efc) recomputes efc.J/efc.pos on an AD-visible path. Plane-sphere, sphere-sphere, sphere-capsule, capsule-capsule, plane-capsule supported; other geom pairs pass through with zero gradient. - adjoint.py: efc-level adjoint kernels (_efc_J_grad_kernel, _efc_pos_grad_kernel) connecting the phase-2 KKT adjoint to contact geometry; smooth/surrogate friction adjoint variants. - forward.py: tape-aware Phase 3 hooks after collision and make_constraint; _record_solver_adjoint/_record_fwd_accel_adjoint/ _record_euler_damp_adjoint tape callbacks; _isolate_intermediates_for_ad per-substep array isolation; freejoint zerograd fix (is_free flag instead of continue) was already present upstream. - grad.py: COLLISION_GRAD_FIELDS registry + dotted-path _resolve_field. Adaptations to current main: - _record_euler_damp_adjoint captures damp_deriv (polynomial damping, computed from the substep's qvel at record time) instead of the old constant dof_damping, matching the forward euler solve; sparse kernel call updated to the M_rownnz/M_rowadr signature - _qfrc_smooth factory kernel re-enabled for backward (pure elementwise sum; main had enable_backward=False which severed the qfrc_smooth -> qfrc_actuator/ctrl gradient chain) - jac_dof calls updated for the body_isdofancestor parameter - _solve_LD_sparse: kept main's fused fast path for non-grad solves, adopted the nograd-copy + manual-adjoint design for grad solves - warmstart copy uses d.qacc (solver solution) not integrator-local qacc, restoring upstream/sleep semantics; sleep.wake retained in forward() - step-level xpos-loss tests refresh kinematics inside the tape (step() does not recompute xpos post-integration; FD agrees it would otherwise be identically zero) and the nonzero-gradient guard is scaled by the FD norm (single-step dL/dctrl ~ dt^2 ~ 2e-7 < the old 1e-6 threshold) - d.qM -> d.M renames; dropped stale pre-google-deepmind#1301 hunks (gauss_cost kernel, old next_act relocation, old collision_flex/passive/types drift) Verification: full suite 1093 passed / 23 skipped, including all 29 grad tests (solver adjoint, euler damp stress, smooth contact, friction surrogate, freejoint kinematics). Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
* filter duplicate contacts * update aloha cloth benchmark * update rollout_cloth.webp * merge contact
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
19 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Completes the autodiff chain port:
mark/autodifferentiation3(123 commits behind) onto the 2/3 port (#10), which sits on the 1/3 rebase (#7). Your design is preserved exactly:collision_smooth.py(smooth analytic distances overwriting discrete contact geometry + differentiablesmooth_contact_to_efc), the efc-level adjoint kernels connecting the phase-2 KKT adjoint to contact geometry, the friction surrogate adjoints,COLLISION_GRAD_FIELDS,_isolate_intermediates_for_ad, and the tape callback architecture (_record_solver_adjoint/_record_fwd_accel_adjoint/_record_euler_damp_adjoint).Stack: #7 (1/3 rebase) → #10 (2/3 port) → this.
Adaptations forced by upstream drift (all documented in the commit)
_record_euler_damp_adjoint: main moved to polynomial damping (_compute_damping_derivfrom the substep's qvel) and changed_euler_damp_qfrc_sparseto theM_rownnz/M_rowadr + damp_derivsignature. The adjoint now capturesdamp_derivat record time so backward applies exactly the forward'sM + dt·Dtransform (reduces to your constant-damping behavior when dampingpoly is unset)._qfrc_smooth: main rewrote it as a factory withenable_backward=False, severing theqfrc_smooth → qfrc_actuator → ctrlchain — re-enabled (pure elementwise sum, auto-adjoint safe). This was the root cause of 13 test failures.jac_dof: newbody_isdofancestorparameter threaded through_smooth_contact_to_efc_kernel._solve_LD_sparse: kept main's fused fast path for non-grad solves; your nograd-copy + manual-adjoint design used when gradients are active.d.qacc(solver solution, not integrator-local qacc),sleep.wakeretained inforward().step()doesn't recompute xpos post-integration — FD confirms the gradient would otherwise be identically zero), and the nonzero-grad guard is scaled by the FD norm (single-step dL/dctrl ≈ dt² ≈ 2e-7, below the old absolute 1e-6 threshold).Verification
1093 passed / 23 skipped — including all 29 grad tests: solver adjoint (contact step, dense jacobian, identity), euler damp stress (high damping sparse+dense), smooth contact AD, friction surrogate boundedness, freejoint kinematics zerograd. E2E: gradient flows
loss → qpos → … → ctrlthrough a full contact step on the Newton solver with smooth recompute active.With this, the full chain (1/3 smooth dynamics → 2/3 solver implicit diff → 3/3 smooth contact AD) is reviewable against today's upstream, same as the determinism stack (#4/#6/#8).