Merge remote-tracking branch 'upstream/main' into tendon_friction

Author: thowell

conftest.py

@@ -24,10 +24,13 @@ def pytest_addoption(parser):
     default=False,
     help="run tests with cuda error checking",
   )
+  parser.addoption("--lineinfo", action="store_true", default=False, help="add lineinfo to warp kernel")
 
 
 def pytest_configure(config):
   if config.getoption("--cpu"):
     wp.set_device("cpu")
   if config.getoption("--verify_cuda"):
     wp.config.verify_cuda = True
+  if config.getoption("--lineinfo"):
+    wp.config.lineinfo = True

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/__init__.py

@@ -34,6 +34,7 @@
 from ._src.io import put_data as put_data
 from ._src.io import put_model as put_model
 from ._src.passive import passive as passive
+from ._src.ray import ray as ray
 from ._src.sensor import sensor_acc as sensor_acc
 from ._src.sensor import sensor_pos as sensor_pos
 from ._src.sensor import sensor_vel as sensor_vel

mujoco_warp/_src/broadphase_test.py

@@ -85,7 +85,7 @@ def test_sap_broadphase(self):
       (6, 7),
     ]
 
-    _, _, m, d = test_util.fixture(xml=_SAP_MODEL)
+    _, _, m, d = test_util.fixture(xml=_SAP_MODEL, qpos0=True)
 
     mjwarp.sap_broadphase(m, d)
 

mujoco_warp/_src/collision_box.py

@@ -192,20 +192,20 @@ def _box_box(
   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_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

@@ -709,23 +709,23 @@ def gjk_epa_sparse(
     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_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]
@@ -781,28 +782,32 @@ def gjk_epa_sparse(
       return
 
     geom1 = _geom(
+      geom_type,
       geom_dataid,
-      geom_size,
+      geom_size[worldid],
       mesh_vertadr,
       mesh_vertnum,
+      mesh_vert,
       geom_xpos_in,
       geom_xmat_in,
       worldid,
       g1,
     )
 
     geom2 = _geom(
+      geom_type,
       geom_dataid,
-      geom_size,
+      geom_size[worldid],
       mesh_vertadr,
       mesh_vertnum,
+      mesh_vert,
       geom_xpos_in,
       geom_xmat_in,
       worldid,
       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

@@ -17,7 +17,6 @@
 
 import warp as wp
 
-from .collision_box import box_box_narrowphase
 from .collision_convex import gjk_narrowphase
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MAXVAL
@@ -32,8 +31,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 +41,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 +123,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 +137,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 +181,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 +342,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:
@@ -446,4 +445,3 @@ def collision(m: Model, d: Data):
   # TODO(team) switch between collision functions and GJK/EPA here
   gjk_narrowphase(m, d)
   primitive_narrowphase(m, d)
-  box_box_narrowphase(m, d)