add utilities for tendon wrapping

mujoco_warp/_src/io.py

@@ -38,6 +38,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     if unsupported.any():
       raise NotImplementedError(f"{field_str} {field[unsupported]} not supported.")
 
+  if np.isin(mjm.wrap_type, [types.WrapType.SPHERE, types.WrapType.CYLINDER]).any():
+    raise NotImplementedError("Tendon geom wrapping not supported.")
+
   if mjm.nplugin > 0:
     raise NotImplementedError("Plugins are unsupported.")
 

mujoco_warp/_src/math.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Tuple
+from typing import Any, Tuple
 
 import warp as wp
 
@@ -213,7 +213,7 @@ def make_frame(a: wp.vec3):
 
 
 @wp.func
-def normalize_with_norm(x: wp.vec3):
+def normalize_with_norm(x: Any):
   norm = wp.length(x)
   if norm == 0.0:
     return x, 0.0

mujoco_warp/_src/types.py

@@ -346,11 +346,15 @@ class WrapType(enum.IntEnum):
   Members:
     JOINT: constant moment arm
     SITE: pass through site
+    SPHERE: wrap around sphere
+    CYLINDER: wrap around (infinite) cylinder
   """
 
   JOINT = mujoco.mjtWrap.mjWRAP_JOINT
   SITE = mujoco.mjtWrap.mjWRAP_SITE
-  # unsupported: PULLEY, SPHERE, CYLINDER
+  SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE
+  CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER
+  # unsupported: PULLEY
 
 
 class vec5f(wp.types.vector(length=5, dtype=float)):