island discovery: edge discovery#1101
Conversation
Model fields: - ntree: number of kinematic trees - tree_dofadr: start address of tree's DOFs (ntree,) - tree_dofnum: number of DOFs in tree (ntree,) - tree_bodynum: number of bodies in tree (ntree,) - body_treeid: id of body's tree; -1 for static (nbody,) - dof_treeid: id of dof's tree (nv,) All fields are imported directly from mjModel via the dataclass field matching pattern in put_model(). Test: test_tree_structure_fields verifies all fields match MuJoCo.
- _find_tree_edges kernel scans Jacobian for tree-tree connections - On-device deduplication using radix sort and prefix sum compaction - Helper kernels: _compute_edge_keys, _init_indices, _mark_unique_edges, _compact_edges - find_tree_edges host function orchestrates the full pipeline - 5 tests: single edge, self-edge, chain, deduplication, no constraints
thowell
force-pushed
the
edge_discovery
branch
from
2f7ebb8 to
18c6394
Compare
yuvaltassa
left a comment
yuvaltassa
left a comment
There was a problem hiding this comment.
I'm not sure this is the right approach on GPU. In the C code and this PR we compute the incidence matrix which is then used to compute the adjacency matrix K. But the sparse adjacency matrix is probably overkill and if K, which is just ntree x ntree is allocated densely then you can compute it directly in parallel quite efficiently IMO.
Imagine that all your threads are just scanning the Jacobian rows in parallel. Every time I see a nonzero in tree i and tree j I set K[i,j] = 1, done.
Let's discuss offline?
mujoco_warp/_src/island.py
Outdated
| # skip continuation rows (same constraint type and id as previous) | ||
| # this avoids duplicate edges from multi-row constraints (e.g., 3D contacts) | ||
| if efcid > 0: | ||
| if efc_type_in[worldid, efcid] == efc_type_in[worldid, efcid - 1]: |
There was a problem hiding this comment.
are we sure that type + id uniquely identifies the column pattern? I think I made that assumption at some point and discovered it was wrong, but I can't quite remember the counterexample. I might be wrong and this could be correct.
thowell
force-pushed
the
edge_discovery
branch
2 times, most recently
from
fb62446 to
6b5ce5e
Compare
|
previous implementation was incorrect/incomplete and has been refactored with different parallelism |
thowell
force-pushed
the
edge_discovery
branch
from
6b5ce5e to
4826c91
Compare
thowell
force-pushed
the
edge_discovery
branch
from
4826c91 to
aff7f4a
Compare
2 participants
implements the host function
tree_edgesand a collection of kernels to compute tree-tree edge structure in the constraint Jacobiantypes.Constraint.J("efc_J")depends on #1100