Skip to content

Proposal: remove quadratic native island-discovery scratch for large tree counts #3388

Description

@teerthsharma

The feature and motivation

Native mj_island currently materializes an ntree x ntree byte adjacency matrix and an
uncompressed ntree x ntree integer column-index array before flood fill. The same construction was
still present on main at dc7581a (checked 2026-07-10):
engine_island.c.

This makes temporary stack memory quadratic even when the constraint incidence is sparse. I would
like to ask whether maintainers would support a native-only PR that constructs connected components
directly from constraint/tree incidence, without materializing the tree adjacency matrix.

I tested a file-local union-by-canonical-minimum-root prototype. It preserves ascending,
deterministic island IDs and leaves the existing mj_floodFill API and tests in place. For a
generated singleton family (one free body and one body-to-world connect per tree, giving 3*ntree
scalar rows), the measured peak mj_island stack follows the source allocation formulas exactly:

ntree current dense path direct-incidence prototype ratio
64 22,816 B 1,056 B 21.6x
256 336,928 B 4,128 B 81.6x
1,024 5,279,776 B 16,416 B 321.6x
4,096 84,033,568 B 65,568 B 1,281.6x

For this family, the current peak is 5*ntree^2 + 36*ntree + 32 bytes and the prototype peak is
16*ntree + 32 bytes. A focused 1,024-tree test also fails on the current path with a 2 MiB arena
and passes unchanged with the prototype.

The complete synthetic matrix supports investigating the direct stage, but not a general
end-to-end speed claim. These are paired baseline/candidate geometric speedups over four topology
families and four tree counts, with 99% bootstrap intervals; values above one favor the prototype:

measured scope speedup 99% CI pairs
direct mj_island 1.6847x [1.5494, 1.8333] 480
mj_fwdPosition 1.1782x [1.1409, 1.2174] 480
full mj_step 1.0355x [1.0254, 1.0460] 480

The prespecified full-step gate required a 99% lower bound of at least 1.05x; it fails, and even the
upper bound is below 1.05x. All four pooled full-step topology-family lower bounds exceed the 0.98
regression threshold, but 19 individual workload/mode cells have lower bounds below 0.98. This is
therefore stage-level evidence under a pooled-family policy, not a per-cell no-regression result.

A separate Callgrind run at 4,096 trees collected 58,071,749 instructions across three matched
mj_island invocations on the current path and 9,092,275 on the prototype (6.39x). The three dense
adjacency clears account for approximately 50.33 million instructions, or 86.7% of the current-path
total. This profile includes setup-time island calls as well as the explicit probe call; it is a
separate mechanistic check, not another timing sample. Both Callgrind processes emitted a nonfatal
brk segment overflow warning.

Correctness checks completed so far:

  • exact agreement between dense reference, sparse-incidence BFS, and DSU on 10,000 deterministic
    generated hypergraphs;
  • exact equality of all island topology/count/map arrays on 16 generated models plus shipped
    humanoid and flex fixtures;
  • zero recorded relative error in qpos, qvel, and qacc for those tested horizons;
  • production-path tests for static-first incidence, repeated scalar rows, and per-row flex equality
    rescanning, with mutation tests;
  • 15/15 focused CTest cases and 300/300 shuffled repetitions;
  • complete output checksums matched across repeated arena rewinds for all four generated topology
    families at 1,024 and 4,096 trees.
  • a fresh final-commit rerun passed the 18-model runtime parity matrix with exact island arrays and
    zero recorded trajectory error.

This relates to the large-scene scaling direction in #1671, but the proposal is limited to native
island discovery.

Questions before preparing a PR:

  1. Is a native-only change acceptable, or should this be coordinated with MuJoCo Warp first?
  2. Is removing quadratic scratch worth the measured small-model stage regression, or would a
    dense/direct threshold be preferable?
  3. Would maintainers prefer retaining the current O(nefc) first-tree scratch, instead of the
    prototype's second constraint-incidence scan?
  4. Which real models and platforms should be required in the benchmark matrix for a PR?

Alternatives

  • Keep the dense path below a tree-count threshold and use direct incidence above it. This may
    preserve current small-model performance but adds a second production path.
  • Build sparse CSR adjacency and keep flood fill. This is conventional, but requires edge counting,
    allocation, and duplicate handling before traversal.
  • Use direct-incidence DSU while retaining efc_tree scratch. This avoids the second constraint scan
    and may reduce the 64-tree regression at the cost of O(nefc) temporary memory.
  • Keep the existing implementation. It is competitive at small ntree, but retains quadratic
    scratch and scans.

Additional context

  • Tested MuJoCo source: f04b0fcd3d5a8f1af798c1fb18f7abac8e89c270; current main was checked
    separately and retains the dense construction.
  • Environment: Docker on WSL2, Ubuntu 24.04, GCC 13.3, Release -O3 -DNDEBUG, LTO, AVX, one pinned
    CPU.
  • Timing protocol: five warmup process pairs and 30 measured process pairs, randomized within each
    pair. Each process loaded identical MJCF and ran the same mode-specific preparation and timed
    region. Direct-stage runs restored the same post-constraint arena boundary before each
    mj_island; position and step runs timed their corresponding native entry points. Full output
    checksums matched within every reported pair.
  • Full-step baseline-first and candidate-first estimates were 1.0363x and 1.0347x. An independent
    workload-adjusted order-effect audit found 1.0076x with 99% CI [0.9947, 1.0207], so no order bias
    was detected; this is not proof of absence.
  • The timing matrix is synthetic. The measured full-step effect misses the prespecified claim gate.
  • The full MuJoCo suite, sanitizers, non-Linux platforms, and real user workloads have not yet been
    validated.
  • MuJoCo Warp is a separate implementation. It was not modified or benchmarked here, and this work
    is not a Warp port. Its current island code also uses dense tree adjacency, which is why I am
    asking whether scope should be coordinated rather than assuming native-only ownership.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions