Model Batching V4

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=float),
-    geom_pos: wp.array(dtype=wp.vec3),
+    qpos0: wp.array2d(dtype=float),
+    geom_pos: wp.array2d(dtype=wp.vec3),
     # Data in:
     qpos_in: wp.array2d(dtype=float),
     qvel_in: wp.array2d(dtype=float),

mujoco_warp/_src/collision_box.py

@@ -191,21 +191,21 @@ def _box_box(
   # Model:
   geom_type: wp.array(dtype=int),
   geom_condim: wp.array(dtype=int),
-  geom_priority: wp.array(dtype=int),
-  geom_solmix: wp.array(dtype=float),
-  geom_solref: wp.array(dtype=wp.vec2),
-  geom_solimp: wp.array(dtype=vec5),
-  geom_size: wp.array(dtype=wp.vec3),
-  geom_friction: wp.array(dtype=wp.vec3),
-  geom_margin: wp.array(dtype=float),
-  geom_gap: wp.array(dtype=float),
+  geom_priority: wp.array2d(dtype=int),
+  geom_solmix: wp.array2d(dtype=float),
+  geom_solref: wp.array2d(dtype=wp.vec2),
+  geom_solimp: wp.array2d(dtype=vec5),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_friction: wp.array2d(dtype=wp.vec3),
+  geom_margin: wp.array2d(dtype=float),
+  geom_gap: wp.array2d(dtype=float),
   pair_dim: wp.array(dtype=int),
-  pair_solref: wp.array(dtype=wp.vec2),
-  pair_solreffriction: wp.array(dtype=wp.vec2),
-  pair_solimp: wp.array(dtype=vec5),
-  pair_margin: wp.array(dtype=float),
-  pair_gap: wp.array(dtype=float),
-  pair_friction: wp.array(dtype=vec5),
+  pair_solref: wp.array2d(dtype=wp.vec2),
+  pair_solreffriction: wp.array2d(dtype=wp.vec2),
+  pair_solimp: wp.array2d(dtype=vec5),
+  pair_margin: wp.array2d(dtype=float),
+  pair_gap: wp.array2d(dtype=float),
+  pair_friction: wp.array2d(dtype=vec5),
   # Data in:
   nconmax_in: int,
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -262,6 +262,7 @@ def _box_box(
       collision_pair_in,
       collision_pairid_in,
       tid,
+      worldid,
     )
 
     # transformations
@@ -271,8 +272,8 @@ def _box_box(
     trans_atob = b_mat_inv @ (a_pos - b_pos)
     rot_atob = b_mat_inv @ a_mat
 
-    a_size = geom_size[ga]
-    b_size = geom_size[gb]
+    a_size = geom_size[worldid, ga]
+    b_size = 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)
 
@@ -363,7 +364,7 @@ def _box_box(
       for i in range(4):
         pos[i] = pos[idx]
 
-    margin = wp.max(geom_margin[ga], geom_margin[gb])
+    margin = wp.max(geom_margin[worldid, ga], 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

@@ -708,24 +708,24 @@ def gjk_epa_sparse(
     geom_type: wp.array(dtype=int),
     geom_condim: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
-    geom_priority: wp.array(dtype=int),
-    geom_solmix: wp.array(dtype=float),
-    geom_solref: wp.array(dtype=wp.vec2),
-    geom_solimp: wp.array(dtype=vec5),
-    geom_size: wp.array(dtype=wp.vec3),
-    geom_friction: wp.array(dtype=wp.vec3),
-    geom_margin: wp.array(dtype=float),
-    geom_gap: wp.array(dtype=float),
+    geom_priority: wp.array2d(dtype=int),
+    geom_solmix: wp.array2d(dtype=float),
+    geom_solref: wp.array2d(dtype=wp.vec2),
+    geom_solimp: wp.array2d(dtype=vec5),
+    geom_size: wp.array2d(dtype=wp.vec3),
+    geom_friction: wp.array2d(dtype=wp.vec3),
+    geom_margin: wp.array2d(dtype=float),
+    geom_gap: wp.array2d(dtype=float),
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
     mesh_vert: wp.array(dtype=wp.vec3),
     pair_dim: wp.array(dtype=int),
-    pair_solref: wp.array(dtype=wp.vec2),
-    pair_solreffriction: wp.array(dtype=wp.vec2),
-    pair_solimp: wp.array(dtype=vec5),
-    pair_margin: wp.array(dtype=float),
-    pair_gap: wp.array(dtype=float),
-    pair_friction: wp.array(dtype=vec5),
+    pair_solref: wp.array2d(dtype=wp.vec2),
+    pair_solreffriction: wp.array2d(dtype=wp.vec2),
+    pair_solimp: wp.array2d(dtype=vec5),
+    pair_margin: wp.array2d(dtype=float),
+    pair_gap: wp.array2d(dtype=float),
+    pair_friction: wp.array2d(dtype=vec5),
     # Data in:
     nconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -772,6 +772,7 @@ def gjk_epa_sparse(
       collision_pair_in,
       collision_pairid_in,
       tid,
+      worldid,
     )
 
     g1 = geoms[0]
@@ -782,7 +783,7 @@ def gjk_epa_sparse(
 
     geom1 = _geom(
       geom_dataid,
-      geom_size,
+      geom_size[worldid],
       mesh_vertadr,
       mesh_vertnum,
       geom_xpos_in,
@@ -793,7 +794,7 @@ def gjk_epa_sparse(
 
     geom2 = _geom(
       geom_dataid,
-      geom_size,
+      geom_size[worldid],
       mesh_vertadr,
       mesh_vertnum,
       geom_xpos_in,
@@ -802,7 +803,7 @@ def gjk_epa_sparse(
       g2,
     )
 
-    margin = wp.max(geom_margin[g1], geom_margin[g2])
+    margin = wp.max(geom_margin[worldid, g1], geom_margin[worldid, g2])
 
     simplex, normal = _gjk(mesh_vert, geom1, geom2)
 

mujoco_warp/_src/collision_driver.py

@@ -32,8 +32,8 @@
 @wp.func
 def _sphere_filter(
   # Model:
-  geom_rbound: wp.array(dtype=float),
-  geom_margin: wp.array(dtype=float),
+  geom_rbound: wp.array2d(dtype=float),
+  geom_margin: wp.array2d(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
@@ -42,12 +42,12 @@ def _sphere_filter(
   geom2: int,
   worldid: int,
 ) -> bool:
-  margin1 = geom_margin[geom1]
-  margin2 = geom_margin[geom2]
+  margin1 = geom_margin[worldid, geom1]
+  margin2 = geom_margin[worldid, geom2]
   pos1 = geom_xpos_in[worldid, geom1]
   pos2 = geom_xpos_in[worldid, geom2]
-  size1 = geom_rbound[geom1]
-  size2 = geom_rbound[geom2]
+  size1 = geom_rbound[worldid, geom1]
+  size2 = geom_rbound[worldid, geom2]
 
   bound = size1 + size2 + wp.max(margin1, margin2)
   dif = pos2 - pos1
@@ -124,8 +124,8 @@ def _upper_tri_index(n: int, i: int, j: int) -> int:
 @wp.kernel
 def _sap_project(
   # Model:
-  geom_rbound: wp.array(dtype=float),
-  geom_margin: wp.array(dtype=float),
+  geom_rbound: wp.array2d(dtype=float),
+  geom_margin: wp.array2d(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
@@ -138,13 +138,13 @@ def _sap_project(
   worldid, geomid = wp.tid()
 
   xpos = geom_xpos_in[worldid, geomid]
-  rbound = geom_rbound[geomid]
+  rbound = geom_rbound[worldid, geomid]
 
   if rbound == 0.0:
     # geom is a plane
     rbound = MJ_MAXVAL
 
-  radius = rbound + geom_margin[geomid]
+  radius = rbound + geom_margin[worldid, geomid]
   center = wp.dot(direction_in, xpos)
 
   sap_projection_lower_out[worldid, geomid] = center - radius
@@ -182,8 +182,8 @@ def _sap_broadphase(
   # Model:
   ngeom: int,
   geom_type: wp.array(dtype=int),
-  geom_rbound: wp.array(dtype=float),
-  geom_margin: wp.array(dtype=float),
+  geom_rbound: wp.array2d(dtype=float),
+  geom_margin: wp.array2d(dtype=float),
   nxn_pairid: wp.array(dtype=int),
   # Data in:
   nworld_in: int,
@@ -343,8 +343,8 @@ def sap_broadphase(m: Model, d: Data):
 def _nxn_broadphase(
   # Model:
   geom_type: wp.array(dtype=int),
-  geom_rbound: wp.array(dtype=float),
-  geom_margin: wp.array(dtype=float),
+  geom_rbound: wp.array2d(dtype=float),
+  geom_margin: wp.array2d(dtype=float),
   nxn_geom_pair: wp.array(dtype=wp.vec2i),
   nxn_pairid: wp.array(dtype=int),
   # Data in: