Refactor convex collision code

[kbayes]

This refactor aligns the convex collision code to be in parity with C Mujoco code:

• Merge primitive collision functions into convex collision detection (ccd)
• Move collision table to collision_driver.py
• Merge rest of hfield logic into collision_convex.py

[adenzler-nvidia]

@kbayes could you provide some context as to why this PR is necessary/your decision making? I fear this might make it even harder to have a separate collision library.

[kbayes]

@adenzler-nvidia yes, please take a closer look in case that I'm misunderstanding something. I was hoping that this PR would make things easier actually. How we're currently doing things is by having two endpoints for calling narrowphase: primitive and general. I think this simplifies things by having one endpoint: convex_narrowphase. And have its logic be to use a primitive test otherwise fallback to the general test. This removes a lot of redundant code, and the core functions aren't altered so I didn't think it would impact Newton.

[adenzler-nvidia]

understood - yes looking closer at this it doesn't change any of the core functions.

My separate feedback is that I'm unsure whether it's the right thing to launch a kernel per pair type for all the possible pair types - leading to n x naconmax threads being launched over the course of this, where before it was just naconmax threads for the entirety of the primitive collisions. There is a downside to this as well due to the branching involved, so the solution is probably somewhere in the middle?

[kbayes]

Yes the difference is that the primitive kernel runs a batch of naconmax and branches on each geom type pair while the convex kernel runs a batch of n x naconmax where n is statically the number of unique geom type pairs. I think we should do one or the other but not both.

[erikfrey]

@kbayes I'm a bit confused by the organization scheme here... it looks like you're putting all the narrowphase functions into collision_convex.py - but then should we call it narrowphase.py?

Stepping back, one of the nice properties of the various implementations of MuJoCo is that we (kind of, sort of) have parity in function and file names. I can look in engine_collision_driver.c and see more or less the same functionality as here in collision_driver.py - it would be nice to maintain that property.

Is there a way to maintain that property and still achieve what you're going for here? Feel free to grab me in person.

[kbayes]

@adenzler-nvidia @erikfrey did some benchmarking but found no difference in timing (that includes the geom diverse mixed.xml model).

In terms of where I'm going with this as this is a first step. This PR better aligns with the table logic in engine_collision_driver.c. The next steps will be to delete the legacy GJK logic, and move the kernel setup out of collision_convex and into collision_driver.py, and keep collision_convex.py has a shared util for the primitive and gjk code.

[erikfrey]

OK I think I see what's going on here, I think this is directionally good but still has a few architectural issues.

Probably the best starting point for discussion is: I think it might still be better to have distinct narrowphase helpers for functionally distinct parts of the code, e.g. primitive vs. convex vs sdf... not sure where hfield lives in that hierarchy, if it should be grouped with convex or not.

My main argument for this is that, apparently, kernel launches are somewhat sensitive to the number of parameters, and I've been warned by folks on the Warp team that some of our kernel launches are starting to break older CUDA versions that only supported smaller stack sizes.

So in general, prefer fewer inputs/outputs where possible/practical.

Besides that, any big change like this, please do run the benchmarks before and after and let us know if you're seeing changes eitehr to step time or JIT time. You can compare head to yours by doing:

to test your branch:

asv run ccd_refactor^!

against main:

asv run

Thanks for your patience - this is a big refactor, please grab me in chat if you think it'll save more thrashing.

[kbayes]

From my perspective the user shouldn't care if a convex pair was handled by a primitive function or ccd (as long as it was handled in a performant and accurate manner). The current way duplicates logic to segregate geom pairs into two categories which doesn't necessary align with anything (i.e. box-box pairs can use ccd or the box-box primitive code). And it makes more sense to unify these into one collision function table just like in mujoco c. We could separate out the primitives into their own kernel launch but that doesn't change the size of input / outputs nor give us any improvements in performance (actually, there might be a slight edge to unify them), so why not just merge them? SDFs are fundamentally different and should be handled separately.