Port constraint solver implicit diff (2/3) onto current google-deepmind/main - #10
Open
johnnynunez wants to merge 134 commits into
Open
Port constraint solver implicit diff (2/3) onto current google-deepmind/main#10johnnynunez wants to merge 134 commits into
johnnynunez wants to merge 134 commits into
Conversation
* Fix multi flex indexing * Fix kernel order * Add multiflex test * Move flex test, fix types ordering * Fix kernel args ordering * Add newline EOF * Cleanup multiflex xml
* Optimize qderiv_actuator_passive_actuation * Add a UT for sparse and coupled actuator * Autoformatter * Address PR review comment * Address PR review comment (2) * fix merge --------- Co-authored-by: Taylor Howell <taylorhowell@google.com>
benchmark/humanoid -> benchmarks/humanoid in testspeed, viewer, and render docstrings.
* Add 1d flex rendering, and add flex to bvh * WIP refactor of flex rendering * Fix and cleanup post merge * Minor fixes and ruff * More cleanup of io for new flex * Remove unused values * Add multi flex test xml * Fix multiflex bug * Fix multi flex indexing * Fix kernel order * Add multiflex test * Move flex test, fix types ordering * Fix kernel args ordering * Add newline EOF * Cleanup multiflex xml * Update rope flex test * Fix model arg * Small fixes * Ruff * Kernel args ordering * Render kernel args ordering * Cleanup tid naming in render/bvh kernels * Restructure accumulate vertex normals * Fix formatting * Cleanup * Fix bvh tests
* 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.
* two-pass GJK+EPA to reduce naccdmax memory Split ccd() in collision_gjk.py into a GJK phase and an EPA phase, with a thin wrapper preserving the original signature. In the non-hfield kernel, the EPA scratch slot is now allocated only when GJK has confirmed penetration, so naccdmax (= nworld * nccdmax) tracks the count of actually-penetrating pairs rather than all broadphase pairs. The hfield kernel keeps the original single-pass behavior (multiple prisms per thread make a two-pass split materially more complex; the wins there are smaller). Lower nccdmax in the aloha benchmark configs accordingly: pot 12 -> 1, cloth 12 -> 1, clutter 320 -> 64. Benchmarks on a single device: | benchmark | step before | step after | mem before | mem after | |---------------|--------------|--------------|-------------|------------| | aloha_pot | 1281 ns | 1291 ns | 2032 MiB | 592 MiB | | aloha_cloth | 17.27 ms | 17.26 ms | 412 MiB | 412 MiB | | aloha_clutter | 329 us | 307 us | 2820 MiB | 772 MiB | Cloth memory is unchanged because nworld=32 makes the EPA scratch trivially small either way. * Address PR comments: remove is_discrete parameter passing and restore deleted comments in GJK phase.
* Fix missing skybox error catch * Add block_dim for render kernel * Small cleanup of camid naming * Add static branch elimination to render function * Refactor cast ray and cast ray first hit * Minor formatting
…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>
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
20 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.
Same treatment as the determinism chain: your
mark/autodifferentiation2is 131 commits behind upstream main, so here's a semantic port onto the rebased 1/3 (#7). Your design is preserved exactly —solver_implicit_adjoint(H v = adj_qacc → adj_qacc_smooth = M v), the three GPU tile-Cholesky adjoint paths, and theData.solver_h/solver_hfactor/solver_Jarefretained state.Adaptations forced by upstream drift:
create_blocked_cholesky_funcwas replaced upstream by the fusedcreate_blocked_cholesky_factorize_solve_func—_adjoint_cholesky_full_blockednow uses the fused variantupdate_constraint_gauss_cost(the pre-Refactor CLI tools and generate benchmark documentation google-deepmind/mujoco_warp#1301 solver context withctx.costno longer exists; gauss cost lives in the linesearch tiles now)support.next_act(current main'sderivative.pyimports it; your branch had moved it into forward.py)Data+ docstring reordered (currenttypes_testenforces warp-only fields after MuJoCo's field order)test_solver_retained_statereferenced an undefinedSPARSE_CONSTRAINT_JACOBIANglobal →m.is_sparseVerification: 1078 passed / 23 skipped including your 4
GradSolverAdjointTestcases, plus an independent cross-check — your GPU adjoint mapped to force space matches central finite differences through the full GPU solver to 8e-5 max relative error on a persistent-contact scene.I previously opened #9 with my own host-side IFT implementation before properly reading your 2/3 — yours supersedes it (GPU end-to-end, tape-integrated); I'll close #9 and can contribute its FD-validation harness as extra test coverage here if you want it. 3/3 (smooth contact AD) port coming next.