refactor collision sensor

[thowell]

refactor collision sensor to follow @erikfrey's suggestion to utilize the collision driver

• augments nxn_pairid/nxn_pairid_filtered/collision_pairid with two new elements [1]: unique collision id corresponding to row in new field Data.sensor_collision, [2]: unique geom collision id since collision sensors can be specified with body semantics and potentially have multiple geoms that will need to be evaluted and compared
• add field sensor_collision to Data. this field will be computed by the collision driver and _sensor_pos will read from this field and utilize the information to compute values for sensordata
• augments nxn_pairid/nxn_pairid_filtered/collision_pairid with a new element (making the array type wp.vec2i): a unique collision id corresponding to an element in the new field Data.collision.
• add field collision to Data. this field is computed by the collision driver and _sensor_pos will read from this field and utilize the information to compute collision sensor values for sensordata. each element in this field is a vec7. each element is distance (1) and fromto (6).
• adds type ContactType for a new bitflag field Contact.type that is utilized by constraints and sensors.

notes:

• each collision sensor currently iterates over all contacts in order to compute values for sensordata. a TODO has been added for improving performance here by iterating over all contacts once for all of the collision sensors
• Data.contact and Data.nacon can now include and account for contacts that don't contribute to the constraints

[erikfrey]

Not a full review yet, but one of my questions could involve some code change if you agree, so let's start with this.

[erikfrey]

this is starting to get a little tricky to reason about... what if we introduce an enum that's like:

class NxNType(enum.IntEnum):
  DEFAULT = 0
  SKIP = 1
  PAIR = 2
  SENSOR = 3

Then separately we have an nxn_objid that corresponds to the collision id or sensor id?

If want to keep it as a single array that's fine, but this way it can be a vec2i (type, objid) instead of a vec3i - WDYT?

[thowell]

not sure if the proposed enum.IntEnum will work in all cases we need. for example:

• contact pair and distance sensor have the same geoms
• exclude contact and distance sensor have the same geoms

we could enumerate all possibilities, for example adding SKIPSENSOR (this would skip contact but write collision sensor), etc or maybe utilize an enum.IntFlag instead?

---

i think we probably need three numbers in some form in order to write contacts and write collision sensor data during the narrowphase

• [0]: contact pair id, -1 if no pair id, -2 if contact is excluded/skipped
• [1]: sensor id, -1 if no sensor

[2]: collision sensors have both geom and body semantics. we need some way of organizing geom data by body when we write collision sensor data. the current implementation writes to sensor_collision[worldid, [1], [2]] using the second and third numbers. an alternative approach might be to have a counter and to utilize atomic adds to increment it. another approach might be to have the third number be the body id and then search over sensor_collision for body id matches in the sensor function.

---

do we think there is an approach that only requires two numbers (ie, type and objid)?

[erikfrey]

Sure, we can use IntFlag. Not quite sure why you need to store the bodyid in nxn_pairid - can't you get it fromgeom_bodyid?

[kevinzakka]

Can confirm this compiles much much faster and in principle resolves #747.

[thowell]

@erikfrey the pr and pr description are updated.

[erikfrey]

OK this is lovely and moves the resulting code is easier to understand!

I know you had some concerns that sweeping across all contacts to find just the sensor contacts might be expensive. Did you run any benchmarks?

Assuming you're satisfied with the perf, I think this code looks nice. Just had one nit.

[thowell]

benchmarking with humanoid/humanoid.xml

mjwarp-testspeed benchmark/humanoid/humanoid.xml --nconmax=24 --njmax=64 --event_trace

Summary for 8192 parallel rollouts

Total JIT time: 0.33 s
Total simulation time: 3.04 s
Total steps per second: 2,692,632
Total realtime factor: 13,463.16 x
Total time per step: 371.38 ns
Total converged worlds: 8192 / 8192

    sensor_pos: 0.17

adding a distance sensor

  <sensor>
    <distance geom1="hand_left" geom2="hand_right"/>
  </sensor>

note: needed to increase nconmax=24 (benchmark) -> nconmax=25 to prevent overflow

mjwarp-testspeed benchmark/humanoid/humanoid.xml --nconmax=25 --njmax=64 --event_trace

Summary for 8192 parallel rollouts

Total JIT time: 0.34 s
Total simulation time: 8.46 s
Total steps per second: 968,807
Total realtime factor: 4,844.04 x
Total time per step: 1032.20 ns
Total converged worlds: 8192 / 8192

    sensor_pos: 660.98

---

to improve performance we can add a kernel launch over contacts that writes to a temporary array. will follow up with new timing results when the updated implementation is ready.

[thowell]

the updated implementation aggregates contact information for collision sensors before writing to sensordata

• adds field Contact.geomcollisionid to identify unique geom pair collisions -- this information is important because colliders can generate more than one contact per geom pair (eg, plane<>capsule: 2)
• adds a new kernel _sensor_collision that is launched over contacts by sensor_pos. this organizes and writes distance and fromto contact information by world, collision sensor, and geom collision id into a temporary array sensor_collision. this array is utilized to write into sensordata for distance, normal, and fromto sensors.

---

humanoid + distance sensor

mjwarp-testspeed benchmark/humanoid/humanoid.xml --nconmax=25 --njmax=64 --event_trace

Summary for 8192 parallel rollouts

Total JIT time: 0.34 s
Total simulation time: 3.06 s
Total steps per second: 2,675,142
Total realtime factor: 13,375.71 x
Total time per step: 373.81 ns
Total converged worlds: 8192 / 8192

    sensor_pos: 2.41

[erikfrey]

Looks great, one last nit and let's get it in.