height field: improve pruning and modify subgrid iterations

[thowell]

update pruning to utilizing support functions for computing bounding box (mujoco reference)

• this should resolve the benchmark/apollo_apptronik/scene_hfield.xml printing warnings

modify subgrid iterations (mujoco reference)

apollo hfield

mjwarp-testspeed benchmark/apptronik_apollo/scene_hfield.xml --nconmax=32 --njmax=128 --nworld=8192

Loading model from: benchmark/apptronik_apollo/scene_hfield.xml...
  nbody: 37 nv: 25 ngeom: 19 nu: 19 is_sparse: False
  broadphase: NXN broadphase_filter: PLANE|SPHERE|OBB
  solver: NEWTON cone: PYRAMIDAL iterations: 100 iterative linesearch iterations: 50
  integrator: EULER graph_conditional: True
Data
  nworld: 8192 naconmax: 262144 njmax: 128

Rolling out 1000 steps at dt = 0.005...

Summary for 8192 parallel rollouts

Total JIT time: 38.30 s
Total simulation time: 6.75 s
Total steps per second: 1,214,517
Total realtime factor: 6,072.59 x
Total time per step: 823.37 ns
Total converged worlds: 8192 / 8192

apollo flat

mjwarp-testspeed benchmark/apptronik_apollo/scene_flat.xml --nconmax=32 --njmax=128 --nworld=8192

Loading model from: benchmark/apptronik_apollo/scene_flat.xml...
  nbody: 37 nv: 25 ngeom: 19 nu: 19 is_sparse: False
  broadphase: NXN broadphase_filter: PLANE|SPHERE|OBB
  solver: NEWTON cone: PYRAMIDAL iterations: 100 iterative linesearch iterations: 50
  integrator: EULER graph_conditional: True
Data
  nworld: 8192 naconmax: 262144 njmax: 128

Rolling out 1000 steps at dt = 0.005...

Summary for 8192 parallel rollouts

Total JIT time: 4.58 s
Total simulation time: 4.10 s
Total steps per second: 1,999,975
Total realtime factor: 9,999.88 x
Total time per step: 500.01 ns
Total converged worlds: 8192 / 8192

flat / hfield: 1999975 / 1214517 = 1.6467245826941903

---

mjx flat / hfield reference

mjx does not support apollo hfield

    raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
NotImplementedError: (mjtGeom.mjGEOM_HFIELD, mjtGeom.mjGEOM_CAPSULE) margin/gap not implemented.

so we test barkour

hfield

mjx-testspeed --mjcf=../../../../../mujoco_menagerie/google_barkour_vb/scene_hfield_mjx.xml --batch_size=8192

Summary for 8192 parallel rollouts

 Total JIT time: 38.19 s
 Total simulation time: 1.93 s
 Total steps per second: 4236413
 Total realtime factor: 8472.83 x
 Total time per step: 0.24 µs

flat

mjx-testspeed --mjcf=../../../../../mujoco_menagerie/google_barkour_vb/scene_mjx.xml --batch_size=8192

Summary for 8192 parallel rollouts

 Total JIT time: 32.52 s
 Total simulation time: 1.12 s
 Total steps per second: 7305322
 Total realtime factor: 14610.64 x
 Total time per step: 0.14 µs

flat / hfield: 7305322 / 4236413 = 1.7244121382877449