Adds raycast implementation and tests

mujoco_warp/__init__.py

@@ -35,6 +35,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/io.py

@@ -137,6 +137,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.nwrap = mjm.nwrap
   m.nsensor = mjm.nsensor
   m.nsensordata = mjm.nsensordata
+  m.nmeshvert = mjm.nmeshvert
+  m.nmeshface = mjm.nmeshface
   m.nlsp = mjm.opt.ls_iterations  # TODO(team): how to set nlsp?
   m.npair = mjm.npair
   m.nexclude = mjm.nexclude
@@ -281,10 +283,15 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       tile_beg = tile_corners[i]
       tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
       tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
-    qLD_tile = np.concatenate([tiles[sz] for sz in sorted(tiles.keys())])
-    tile_off = [0] + [len(tiles[sz]) for sz in sorted(tiles.keys())]
-    qLD_tileadr = np.cumsum(tile_off)[:-1]
-    qLD_tilesize = np.array(sorted(tiles.keys()))
+    if tiles:
+      qLD_tile = np.concatenate([tiles[sz] for sz in sorted(tiles.keys())])
+      tile_off = [0] + [len(tiles[sz]) for sz in sorted(tiles.keys())]
+      qLD_tileadr = np.cumsum(tile_off)[:-1]
+      qLD_tilesize = np.array(sorted(tiles.keys()))
+    else:
+      qLD_tile = np.array([], dtype=int)
+      qLD_tileadr = np.array([], dtype=int)
+      qLD_tilesize = np.array([], dtype=int)
 
   # tiles for actuator_moment - needs nu + nv tile size and offset
   actuator_moment_offset_nv = np.empty(shape=(0,), dtype=int)
@@ -297,7 +304,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     # how many actuators for each tree
     tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
     tree_id = mjm.dof_treeid[tile_corners]
-    num_trees = int(np.max(tree_id))
+    num_trees = int(np.max(tree_id)) if len(tree_id) > 0 else 0
     tree = mjm.body_treeid[mjm.jnt_bodyid[mjm.actuator_trnid[:, 0]]]
     counts, ids = np.histogram(tree, bins=np.arange(0, num_trees + 2))
     acts_per_tree = dict(zip([int(i) for i in ids], [int(i) for i in counts]))
@@ -416,12 +423,17 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.geom_friction = wp.array(mjm.geom_friction, dtype=wp.vec3, ndim=1)
   m.geom_margin = wp.array(mjm.geom_margin, dtype=wp.float32, ndim=1)
   m.geom_gap = wp.array(mjm.geom_gap, dtype=wp.float32, ndim=1)
+  m.geom_rgba = wp.array(mjm.geom_rgba, dtype=wp.vec4, ndim=1)
   m.geom_aabb = wp.array(mjm.geom_aabb, dtype=wp.vec3, ndim=3)
   m.geom_rbound = wp.array(mjm.geom_rbound, dtype=wp.float32, ndim=1)
   m.geom_dataid = wp.array(mjm.geom_dataid, dtype=wp.int32, ndim=1)
+  m.geom_group = wp.array(mjm.geom_group, dtype=wp.int32, ndim=1)
+  m.geom_matid = wp.array(mjm.geom_matid, dtype=wp.int32, ndim=1)
   m.mesh_vertadr = wp.array(mjm.mesh_vertadr, dtype=wp.int32, ndim=1)
   m.mesh_vertnum = wp.array(mjm.mesh_vertnum, dtype=wp.int32, ndim=1)
   m.mesh_vert = wp.array(mjm.mesh_vert, dtype=wp.vec3, ndim=1)
+  m.mesh_faceadr = wp.array(mjm.mesh_faceadr, dtype=wp.int32, ndim=1)
+  m.mesh_face = wp.array(mjm.mesh_face, dtype=wp.vec3i, ndim=1)
   m.eq_type = wp.array(mjm.eq_type, dtype=wp.int32, ndim=1)
   m.eq_obj1id = wp.array(mjm.eq_obj1id, dtype=wp.int32, ndim=1)
   m.eq_obj2id = wp.array(mjm.eq_obj2id, dtype=wp.int32, ndim=1)
@@ -508,6 +520,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.pair_friction = wp.array(mjm.pair_friction, dtype=types.vec5, ndim=1)
   m.condim_max = np.max(mjm.geom_condim)  # TODO(team): get max after filtering
 
+  m.mat_rgba = wp.array(mjm.mat_rgba, dtype=wp.vec4, ndim=1)
+
   # tendon
   m.tendon_adr = wp.array(mjm.tendon_adr, dtype=wp.int32, ndim=1)
   m.tendon_num = wp.array(mjm.tendon_num, dtype=wp.int32, ndim=1)