Model batching v2

contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py

@@ -136,8 +136,8 @@ def test_all_issues(
 @kernel
 def test_no_issues(
     # Model:
-    qpos0: wp.array(dtype=wp.float32, ndim=1),
-    geom_pos: wp.array(dtype=wp.vec3, ndim=1),
+    qpos0: wp.array(dtype=wp.float32, ndim=2),
+    geom_pos: wp.array(dtype=wp.vec3, ndim=2),
     # Data in:
     qpos_in: wp.array(dtype=wp.float32, ndim=2),
     qvel_in: wp.array(dtype=wp.float32, ndim=2),

mujoco_warp/_src/collision_box.py

@@ -210,7 +210,7 @@ def box_box_kernel(
     worldid = d.collision_worldid[bp_idx]
 
     geoms, margin, gap, condim, friction, solref, solreffriction, solimp = (
-      contact_params(m, d, tid)
+      contact_params(m, d, tid, worldid)
     )
 
     # transformations
@@ -220,8 +220,8 @@ def box_box_kernel(
     trans_atob = b_mat_inv @ (a_pos - b_pos)
     rot_atob = b_mat_inv @ a_mat
 
-    a_size = m.geom_size[ga]
-    b_size = m.geom_size[gb]
+    a_size = m.geom_size[worldid, ga]
+    b_size = m.geom_size[worldid, gb]
     a = box(rot_atob, trans_atob, a_size)
     b = box(wp.identity(3, wp.float32), wp.vec3(0.0), b_size)
 
@@ -312,7 +312,10 @@ def box_box_kernel(
       for i in range(4):
         pos[i] = pos[idx]
 
-    margin = wp.max(m.geom_margin[ga], m.geom_margin[gb])
+    margin = wp.max(
+      m.geom_margin[worldid, ga],
+      m.geom_margin[worldid, gb],
+    )
     for i in range(4):
       pos_glob = b_mat @ pos[i] + b_pos
       n_glob = b_mat @ sep_axis

mujoco_warp/_src/collision_convex.py

@@ -754,7 +754,7 @@ def _gjk_epa_sparse(m: Model, d: Data):
 
     worldid = d.collision_worldid[tid]
     geoms, margin, gap, condim, friction, solref, solreffriction, solimp = (
-      contact_params(m, d, tid)
+      contact_params(m, d, tid, worldid)
     )
 
     g1 = geoms[0]
@@ -763,10 +763,13 @@ def _gjk_epa_sparse(m: Model, d: Data):
     if m.geom_type[g1] != geomtype1 or m.geom_type[g2] != geomtype2:
       return
 
-    geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
-    geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
+    geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid)
+    geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid)
 
-    margin = wp.max(m.geom_margin[g1], m.geom_margin[g2])
+    margin = wp.max(
+      m.geom_margin[worldid, g1],
+      m.geom_margin[worldid, g2],
+    )
 
     simplex, normal = _gjk(m, geom1, geom2)
 

mujoco_warp/_src/collision_driver.py

@@ -32,12 +32,12 @@
 
 @wp.func
 def _sphere_filter(m: Model, d: Data, geom1: int, geom2: int, worldid: int) -> bool:
-  margin1 = m.geom_margin[geom1]
-  margin2 = m.geom_margin[geom2]
+  margin1 = m.geom_margin[worldid, geom1]
+  margin2 = m.geom_margin[worldid, geom2]
   pos1 = d.geom_xpos[worldid, geom1]
   pos2 = d.geom_xpos[worldid, geom2]
-  size1 = m.geom_rbound[geom1]
-  size2 = m.geom_rbound[geom2]
+  size1 = m.geom_rbound[worldid, geom1]
+  size2 = m.geom_rbound[worldid, geom2]
 
   bound = size1 + size2 + wp.max(margin1, margin2)
   dif = pos2 - pos1
@@ -105,13 +105,13 @@ def _sap_project(m: Model, d: Data, direction: wp.vec3):
   worldid, geomid = wp.tid()
 
   xpos = d.geom_xpos[worldid, geomid]
-  rbound = m.geom_rbound[geomid]
+  rbound = m.geom_rbound[worldid, geomid]
 
   if rbound == 0.0:
     # geom is a plane
     rbound = MJ_MAXVAL
 
-  radius = rbound + m.geom_margin[geomid]
+  radius = rbound + m.geom_margin[worldid, geomid]
   center = wp.dot(direction, xpos)
 
   d.sap_projection_lower[worldid, geomid] = center - radius

mujoco_warp/_src/collision_primitive.py

@@ -44,14 +44,15 @@ def _geom(
   m: Model,
   geom_xpos: wp.array(dtype=wp.vec3),
   geom_xmat: wp.array(dtype=wp.mat33),
+  worldid: int,
 ) -> Geom:
   geom = Geom()
   geom.pos = geom_xpos[gid]
   rot = geom_xmat[gid]
   geom.rot = rot
-  geom.size = m.geom_size[gid]
+  geom.size = m.geom_size[worldid, gid]
   geom.normal = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2])  # plane
-  dataid = m.geom_dataid[gid]
+  dataid = m.geom_dataid[worldid, gid]
   if dataid >= 0:
     geom.vertadr = m.mesh_vertadr[dataid]
     geom.vertnum = m.mesh_vertnum[dataid]
@@ -737,44 +738,53 @@ def plane_cylinder(
 
 
 @wp.func
-def contact_params(m: Model, d: Data, cid: int):
+def contact_params(m: Model, d: Data, cid: int, worldid: int):
   geoms = d.collision_pair[cid]
   pairid = d.collision_pairid[cid]
 
   if pairid > -1:
-    margin = m.pair_margin[pairid]
-    gap = m.pair_gap[pairid]
+    margin = m.pair_margin[worldid, pairid]
+    gap = m.pair_gap[worldid, pairid]
     condim = m.pair_dim[pairid]
-    friction = m.pair_friction[pairid]
-    solref = m.pair_solref[pairid]
-    solreffriction = m.pair_solreffriction[pairid]
-    solimp = m.pair_solimp[pairid]
+    friction = m.pair_friction[worldid, pairid]
+    solref = m.pair_solref[worldid, pairid]
+    solreffriction = m.pair_solreffriction[worldid, pairid]
+    solimp = m.pair_solimp[worldid, pairid]
   else:
     g1 = geoms[0]
     g2 = geoms[1]
 
-    p1 = m.geom_priority[g1]
-    p2 = m.geom_priority[g2]
+    p1 = m.geom_priority[worldid, g1]
+    p2 = m.geom_priority[worldid, g2]
 
-    solmix1 = m.geom_solmix[g1]
-    solmix2 = m.geom_solmix[g2]
+    solmix1 = m.geom_solmix[worldid, g1]
+    solmix2 = m.geom_solmix[worldid, g2]
 
     mix = solmix1 / (solmix1 + solmix2)
     mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
     mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
     mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
     mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0))
 
-    margin = wp.max(m.geom_margin[g1], m.geom_margin[g2])
-    gap = wp.max(m.geom_gap[g1], m.geom_gap[g2])
+    margin = wp.max(
+      m.geom_margin[worldid, g1],
+      m.geom_margin[worldid, g2],
+    )
+    gap = wp.max(
+      m.geom_gap[worldid, g1],
+      m.geom_gap[worldid, g2],
+    )
 
     condim1 = m.geom_condim[g1]
     condim2 = m.geom_condim[g2]
     condim = wp.where(
       p1 == p2, wp.max(condim1, condim2), wp.where(p1 > p2, condim1, condim2)
     )
 
-    geom_friction = wp.max(m.geom_friction[g1], m.geom_friction[g2])
+    geom_friction = wp.max(
+      m.geom_friction[worldid, g1],
+      m.geom_friction[worldid, g2],
+    )
     friction = vec5(
       geom_friction[0],
       geom_friction[0],
@@ -783,14 +793,19 @@ def contact_params(m: Model, d: Data, cid: int):
       geom_friction[2],
     )
 
-    if m.geom_solref[g1].x > 0.0 and m.geom_solref[g2].x > 0.0:
-      solref = mix * m.geom_solref[g1] + (1.0 - mix) * m.geom_solref[g2]
+    if m.geom_solref[worldid, g1].x > 0.0 and m.geom_solref[worldid, g2].x > 0.0:
+      solref = (
+        mix * m.geom_solref[worldid, g1] + (1.0 - mix) * m.geom_solref[worldid, g2]
+      )
     else:
-      solref = wp.min(m.geom_solref[g1], m.geom_solref[g2])
+      solref = wp.min(
+        m.geom_solref[worldid, g1],
+        m.geom_solref[worldid, g2],
+      )
 
     solreffriction = wp.vec2(0.0, 0.0)
 
-    solimp = mix * m.geom_solimp[g1] + (1.0 - mix) * m.geom_solimp[g2]
+    solimp = mix * m.geom_solimp[worldid, g1] + (1.0 - mix) * m.geom_solimp[worldid, g2]
 
   return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
 
@@ -1258,16 +1273,16 @@ def _primitive_narrowphase(
   if tid >= d.ncollision[0]:
     return
 
+  worldid = d.collision_worldid[tid]
+
   geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
-    m, d, tid
+    m, d, tid, worldid
   )
   g1 = geoms[0]
   g2 = geoms[1]
 
-  worldid = d.collision_worldid[tid]
-
-  geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
-  geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
+  geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid)
+  geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid)
 
   type1 = m.geom_type[g1]
   type2 = m.geom_type[g2]