warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/dlpack.py

@@ -5,6 +5,9 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+# Python specification for DLpack:
+# https://dmlc.github.io/dlpack/latest/python_spec.html
+
 import warp
 import ctypes
 
@@ -14,77 +17,104 @@ from warp.thirdparty.dlpack import (
     DLDeviceType,
     DLDataType,
     DLDataTypeCode,
-    DLTensor,
     _c_str_dltensor,
 )
 
-ctypes.pythonapi.PyMem_RawMalloc.restype = ctypes.c_void_p
-ctypes.pythonapi.PyMem_RawFree.argtypes = [ctypes.c_void_p]
+_c_str_used_dltensor = b"used_dltensor"
 
-ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
-ctypes.pythonapi.PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+PyMem_RawMalloc = ctypes.pythonapi.PyMem_RawMalloc
+PyMem_RawMalloc.argtypes = [ctypes.c_size_t]
+PyMem_RawMalloc.restype = ctypes.c_void_p
 
-ctypes.pythonapi.PyCapsule_IsValid.restype = ctypes.c_int
-ctypes.pythonapi.PyCapsule_IsValid.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyMem_RawFree = ctypes.pythonapi.PyMem_RawFree
+PyMem_RawFree.argtypes = [ctypes.c_void_p]
+PyMem_RawFree.restype = None
 
-ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
-ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+Py_IncRef = ctypes.pythonapi.Py_IncRef
+Py_IncRef.argtypes = [ctypes.py_object]
+Py_IncRef.restype = None
 
+Py_DecRef = ctypes.pythonapi.Py_DecRef
+Py_DecRef.argtypes = [ctypes.py_object]
+Py_DecRef.restype = None
 
-class _Holder:
-    def __init__(self, wp_array) -> None:
-        self.wp_array = wp_array
+PyCapsule_Destructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
 
-    def _as_manager_ctx(self) -> ctypes.c_void_p:
-        py_obj = ctypes.py_object(self)
-        py_obj_ptr = ctypes.pointer(py_obj)
-        ctypes.pythonapi.Py_IncRef(py_obj)
-        ctypes.pythonapi.Py_IncRef(ctypes.py_object(py_obj_ptr))
-        return ctypes.cast(py_obj_ptr, ctypes.c_void_p)
+PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, PyCapsule_Destructor]
+PyCapsule_New.restype = ctypes.py_object
 
+PyCapsule_IsValid = ctypes.pythonapi.PyCapsule_IsValid
+PyCapsule_IsValid.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyCapsule_IsValid.restype = ctypes.c_int
 
-@ctypes.CFUNCTYPE(None, ctypes.c_void_p)
-def _warp_array_deleter(handle: ctypes.c_void_p) -> None:
-    """A function to deallocate the memory of a Warp array."""
+PyCapsule_GetPointer = ctypes.pythonapi.PyCapsule_GetPointer
+PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyCapsule_GetPointer.restype = ctypes.c_void_p
+
+PyCapsule_SetName = ctypes.pythonapi.PyCapsule_SetName
+PyCapsule_SetName.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyCapsule_SetName.restype = ctypes.c_int
+
+
+class _DLPackTensorHolder:
+    """Class responsible for deleting DLManagedTensor memory after ownership is transferred from a capsule."""
 
-    dl_managed_tensor = DLManagedTensor.from_address(handle)
-    py_obj_ptr = ctypes.cast(dl_managed_tensor.manager_ctx, ctypes.POINTER(ctypes.py_object))
-    py_obj = py_obj_ptr.contents
-    ctypes.pythonapi.Py_DecRef(py_obj)
-    ctypes.pythonapi.Py_DecRef(ctypes.py_object(py_obj_ptr))
-    ctypes.pythonapi.PyMem_RawFree(handle)
+    def __init__(self, mem_ptr):
+        self.mem_ptr = mem_ptr
+    
+    def __del__(self):
+        managed_tensor = DLManagedTensor.from_address(self.mem_ptr)
+        if managed_tensor.deleter:
+            managed_tensor.deleter(self.mem_ptr)
 
 
 @ctypes.CFUNCTYPE(None, ctypes.c_void_p)
-def _warp_pycapsule_deleter(handle: ctypes.c_void_p) -> None:
-    """A function to deallocate a pycapsule that wraps a Warp array."""
+def _dlpack_tensor_deleter(managed_ptr) -> None:
+    """A function to deallocate a DLManagedTensor."""
 
-    pycapsule: ctypes.py_object = ctypes.cast(handle, ctypes.py_object)
-    if ctypes.pythonapi.PyCapsule_IsValid(pycapsule, _c_str_dltensor):
-        dl_managed_tensor = ctypes.pythonapi.PyCapsule_GetPointer(pycapsule, _c_str_dltensor)
-        _warp_array_deleter(dl_managed_tensor)
-        ctypes.pythonapi.PyCapsule_SetDestructor(pycapsule, None)
+    managed_tensor = DLManagedTensor.from_address(managed_ptr)
 
+    # unreference the source array
+    manager = ctypes.cast(managed_tensor.manager_ctx, ctypes.py_object)
+    ctypes.pythonapi.Py_DecRef(manager)
 
-def device_to_dlpack(wp_device) -> DLDevice:
-    d = warp.get_device(wp_device)
-
-    if d.is_cpu:
-        device_type = DLDeviceType.kDLCPU
-        device_id = 0
-    elif d.is_cuda:
-        device_type = DLDeviceType.kDLCUDA
-        device_id = d.ordinal
-    else:
-        raise RuntimeError("Unhandled device type converting to dlpack")
+    # free the DLManagedTensor memory, including shape and strides
+    PyMem_RawFree(ctypes.c_void_p(managed_ptr))
+
+
+@PyCapsule_Destructor
+def _dlpack_capsule_deleter(ptr) -> None:
+    """Destructor for a capsule holding a DLManagedTensor."""
+
+    capsule = ctypes.cast(ptr, ctypes.py_object)
 
+    if ctypes.pythonapi.PyCapsule_IsValid(capsule, _c_str_dltensor):
+        managed_ptr = ctypes.pythonapi.PyCapsule_GetPointer(capsule, _c_str_dltensor)
+        managed_tensor = DLManagedTensor.from_address(managed_ptr)
+        if managed_tensor.deleter:
+            managed_tensor.deleter(managed_ptr)
+
+
+def _device_to_dlpack(wp_device: warp.context.Device) -> DLDevice:
     dl_device = DLDevice()
-    dl_device.device_type = device_type
-    dl_device.device_id = device_id
+
+    if wp_device.is_cpu:
+        dl_device.device_type = DLDeviceType.kDLCPU
+        dl_device.device_id = 0
+    elif wp_device.is_cuda:
+        dl_device.device_type = DLDeviceType.kDLCUDA
+        dl_device.device_id = wp_device.ordinal
+    else:
+        raise RuntimeError(f"Invalid device type converting to dlpack: {wp_device}")
 
     return dl_device
 
 
+def device_to_dlpack(wp_device) -> DLDevice:
+    return _device_to_dlpack(warp.get_device(wp_device))
+
+
 def dtype_to_dlpack(wp_dtype) -> DLDataType:
     if wp_dtype == warp.int8:
         return (DLDataTypeCode.kDLInt, 8, 1)
@@ -149,13 +179,15 @@ def dtype_from_dlpack(dl_dtype):
 
 
 def device_from_dlpack(dl_device):
+    assert warp.context.runtime is not None, "Warp not initialized, call wp.init() before use"
+
     if dl_device.device_type.value == DLDeviceType.kDLCPU or dl_device.device_type.value == DLDeviceType.kDLCUDAHost:
-        return "cpu"
+        return warp.context.runtime.cpu_device
     elif (
         dl_device.device_type.value == DLDeviceType.kDLCUDA
         or dl_device.device_type.value == DLDeviceType.kDLCUDAManaged
     ):
-        return f"cuda:{dl_device.device_id}"
+        return warp.context.runtime.cuda_devices[dl_device.device_id]
     else:
         raise RuntimeError(f"Unknown device type from dlpack: {dl_device.device_type.value}")
 
@@ -167,39 +199,29 @@ def shape_to_dlpack(shape):
 
 def strides_to_dlpack(strides, dtype):
     # convert from byte count to element count
-    s = []
-    for i in range(len(strides)):
-        s.append(int(int(strides[i]) / int(warp.types.type_size_in_bytes(dtype))))
-
-    a = (ctypes.c_int64 * len(strides))(*s)
+    ndim = len(strides)
+    a = (ctypes.c_int64 * ndim)()
+    dtype_size = warp.types.type_size_in_bytes(dtype)
+    for i in range(ndim):
+        a[i] = strides[i] // dtype_size
     return a
 
 
 def to_dlpack(wp_array: warp.array):
     """Convert a Warp array to another type of dlpack compatible array.
 
-    Parameters
-    ----------
-    np_array : np.ndarray
-        The source numpy array that will be converted.
+    Args:
+        wp_array: The source Warp array that will be converted.
 
-    Returns
-    -------
-    pycapsule : PyCapsule
-        A pycapsule containing a DLManagedTensor that can be converted
-        to other array formats without copying the underlying memory.
+    Returns:
+        A capsule containing a DLManagedTensor that can be converted
+        to another array type without copying the underlying memory.
     """
 
     # DLPack does not support structured arrays
     if isinstance(wp_array.dtype, warp.codegen.Struct):
         raise RuntimeError("Cannot convert structured Warp arrays to DLPack.")
 
-    holder = _Holder(wp_array)
-
-    # allocate DLManagedTensor
-    size = ctypes.c_size_t(ctypes.sizeof(DLManagedTensor))
-    dl_managed_tensor = DLManagedTensor.from_address(ctypes.pythonapi.PyMem_RawMalloc(size))
-
     # handle vector types
     if hasattr(wp_array.dtype, "_wp_scalar_type_"):
         # vector type, flatten the dimensions into one tuple
@@ -215,33 +237,60 @@ def to_dlpack(wp_array: warp.array):
         target_shape = wp_array.shape
         target_strides = wp_array.strides
 
-    # store the shape and stride arrays with the holder to prevent them from getting deallocated
-    holder._shape = shape_to_dlpack(target_shape)
-    holder._strides = strides_to_dlpack(target_strides, target_dtype)
-
     if wp_array.pinned:
-        dl_device = DLDeviceType.kDLCUDAHost
+        dl_device = DLDevice()
+        dl_device.device_type = DLDeviceType.kDLCUDAHost
+        dl_device.device_id = 0
+    else:
+        dl_device = _device_to_dlpack(wp_array.device)
+
+    # allocate DLManagedTensor, shape, and strides together
+    managed_tensor_size = ctypes.sizeof(DLManagedTensor)
+    padding = managed_tensor_size & 7
+    shape_size = target_ndim * 8
+    mem_size = managed_tensor_size + padding + 2 * shape_size
+    mem_ptr = PyMem_RawMalloc(mem_size)
+    assert mem_ptr, "Failed to allocate memory for DLManagedTensor"
+
+    # set managed tensor attributes
+    managed_tensor = DLManagedTensor.from_address(mem_ptr)
+    managed_tensor.dl_tensor.data = wp_array.ptr
+    managed_tensor.dl_tensor.device = dl_device
+    managed_tensor.dl_tensor.ndim = target_ndim
+    managed_tensor.dl_tensor.dtype = dtype_to_dlpack(target_dtype)
+    managed_tensor.dl_tensor.byte_offset = 0
+
+    # shape
+    shape_offset = managed_tensor_size + padding
+    shape_ptr = ctypes.cast(mem_ptr + shape_offset, ctypes.POINTER(ctypes.c_int64))
+    for i in range(target_ndim):
+        shape_ptr[i] = target_shape[i]
+    managed_tensor.dl_tensor.shape = shape_ptr
+
+    # strides, if not contiguous
+    if wp_array.is_contiguous:
+        managed_tensor.dl_tensor.strides = None
     else:
-        dl_device = device_to_dlpack(wp_array.device)
-
-    # set Tensor attributes
-    dl_managed_tensor.dl_tensor.data = wp_array.ptr
-    dl_managed_tensor.dl_tensor.device = dl_device
-    dl_managed_tensor.dl_tensor.ndim = target_ndim
-    dl_managed_tensor.dl_tensor.dtype = dtype_to_dlpack(target_dtype)
-    dl_managed_tensor.dl_tensor.shape = holder._shape
-    dl_managed_tensor.dl_tensor.strides = holder._strides
-    dl_managed_tensor.dl_tensor.byte_offset = 0
-    dl_managed_tensor.manager_ctx = holder._as_manager_ctx()
-    dl_managed_tensor.deleter = _warp_array_deleter
-
-    pycapsule = ctypes.pythonapi.PyCapsule_New(
-        ctypes.byref(dl_managed_tensor),
+        stride_offset = shape_offset + shape_size
+        stride_ptr = ctypes.cast(mem_ptr + stride_offset, ctypes.POINTER(ctypes.c_int64))
+        dtype_size = warp.types.type_size_in_bytes(target_dtype)
+        for i in range(target_ndim):
+            stride_ptr[i] = target_strides[i] // dtype_size
+        managed_tensor.dl_tensor.strides = stride_ptr
+
+    # DLManagedTensor holds a reference to the source array
+    managed_tensor.manager_ctx = id(wp_array)
+    Py_IncRef(wp_array)
+
+    managed_tensor.deleter = _dlpack_tensor_deleter
+
+    capsule = PyCapsule_New(
+        ctypes.byref(managed_tensor),
         _c_str_dltensor,
-        _warp_pycapsule_deleter,
+        _dlpack_capsule_deleter,
     )
 
-    return pycapsule
+    return capsule
 
 
 def dtype_is_compatible(dl_dtype, wp_dtype):
@@ -272,35 +321,28 @@ def dtype_is_compatible(dl_dtype, wp_dtype):
         raise RuntimeError(f"Unsupported dlpack dtype {(str(dl_dtype.type_code), dl_dtype.bits)}")
 
 
-def from_dlpack(pycapsule, dtype=None) -> warp.array:
-    """Convert a dlpack tensor into a numpy array without copying.
+def _from_dlpack(capsule, dtype=None) -> warp.array:
+    """Convert a DLPack capsule into a Warp array without copying.
 
-    Parameters
-    ----------
-    pycapsule : PyCapsule
-        A pycapsule wrapping a dlpack tensor that will be converted.
+    Args:
+        capsule: A DLPack capsule wrapping an external array or tensor.
+        dtype: An optional Warp data type to interpret the source data.
 
-    Returns
-    -------
-    np_array : np.ndarray
-        A new numpy array that uses the same underlying memory as the input
-        pycapsule.
+    Returns:
+        A new Warp array that uses the same underlying memory as the input capsule.
     """
 
-    assert ctypes.pythonapi.PyCapsule_IsValid(pycapsule, _c_str_dltensor)
-    dl_managed_tensor = ctypes.pythonapi.PyCapsule_GetPointer(pycapsule, _c_str_dltensor)
-    dl_managed_tensor_ptr = ctypes.cast(dl_managed_tensor, ctypes.POINTER(DLManagedTensor))
-    dl_managed_tensor = dl_managed_tensor_ptr.contents
+    assert PyCapsule_IsValid(capsule, _c_str_dltensor), "Invalid capsule"
+    mem_ptr = PyCapsule_GetPointer(capsule, _c_str_dltensor)
+    managed_tensor = DLManagedTensor.from_address(mem_ptr)
 
-    dlt = dl_managed_tensor.dl_tensor
-    assert isinstance(dlt, DLTensor)
+    dlt = managed_tensor.dl_tensor
 
     device = device_from_dlpack(dlt.device)
-
     pinned = dlt.device.device_type.value == DLDeviceType.kDLCUDAHost
-
     shape = tuple(dlt.shape[dim] for dim in range(dlt.ndim))
 
+    # strides, if not contiguous
     itemsize = dlt.dtype.bits // 8
     if dlt.strides:
         strides = tuple(dlt.strides[dim] * itemsize for dim in range(dlt.ndim))
@@ -348,10 +390,53 @@ def from_dlpack(pycapsule, dtype=None) -> warp.array:
         raise RuntimeError(f"Incompatible data types: {dlt.dtype} and {dtype}")
 
     a = warp.types.array(
-        ptr=dlt.data, dtype=dtype, shape=shape, strides=strides, copy=False, owner=False, device=device, pinned=pinned
+        ptr=dlt.data, dtype=dtype, shape=shape, strides=strides, copy=False, device=device, pinned=pinned
     )
 
-    # keep a reference to the capsule so it doesn't get deleted
-    a._pycapsule = pycapsule
+    # take ownership of the DLManagedTensor
+    a._dlpack_tensor_holder = _DLPackTensorHolder(mem_ptr)
+
+    # rename the capsule so that it no longer owns the DLManagedTensor
+    PyCapsule_SetName(capsule, _c_str_used_dltensor)
 
     return a
+
+
+def from_dlpack(source, dtype=None) -> warp.array:
+    """Convert a source array or DLPack capsule into a Warp array without copying.
+
+    Args:
+        source: A DLPack-compatible array or PyCapsule
+        dtype: An optional Warp data type to interpret the source data.
+
+    Returns:
+        A new Warp array that uses the same underlying memory as the input
+        pycapsule.
+    """
+
+    # See https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.array.__dlpack__.html
+
+    if hasattr(source, "__dlpack__"):
+        device_type, device_id = source.__dlpack_device__()
+        # Check if the source lives on a CUDA device
+        if device_type in (DLDeviceType.kDLCUDA, DLDeviceType.kDLCUDAManaged):
+            # Assume that the caller will use the array on its device's current stream.
+            # Note that we pass 1 for the null stream, per DLPack spec.
+            cuda_stream = warp.get_cuda_device(device_id).stream.cuda_stream or 1
+        elif device_type == DLDeviceType.kDLCPU:
+            # No stream sync for CPU arrays.
+            cuda_stream = None
+        elif device_type == DLDeviceType.kDLCUDAHost:
+            # For pinned memory, we sync with the current CUDA device's stream.
+            # Note that we pass 1 for the null stream, per DLPack spec.
+            cuda_stream = warp.get_cuda_device().stream.cuda_stream or 1
+        else:
+            raise TypeError("Unsupported source device")
+
+        capsule = source.__dlpack__(stream=cuda_stream)
+
+    else:
+        # legacy behaviour, assume source is a capsule
+        capsule = source
+
+    return _from_dlpack(capsule, dtype=dtype)

warp/examples/assets/cartpole.urdf

@@ -0,0 +1,110 @@
+<?xml version="1.0"?>
+<robot name="cartpole">
+
+  <link name="slider">
+    <visual>
+      <geometry>
+        <box size="0.03 8 0.03"/>
+      </geometry>
+      <material name="slider_mat">
+        <color rgba="0.9 0.6 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.03 8 0.03"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <joint name="slider_to_cart" type="prismatic">
+    <axis xyz="0 1 0"/>
+    <origin xyz="0 0 0"/>
+    <parent link="slider"/>
+    <child link="cart"/>
+    <limit effort="1000.0" lower="-4" upper="4" velocity="100"/>
+  </joint>
+  
+  <link name="cart">
+    <visual>
+      <geometry>
+        <box size="0.2 0.25 0.2"/>
+      </geometry>
+      <material name="cart_mat">
+        <color rgba="0.3 0.5 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+          <box size="0.2 0.25 0.2"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <mass value="1"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <joint name="cart_to_pole1" type="continuous">
+    <axis xyz="1 0 0"/>
+    <origin xyz="0.12 0 0"/>
+    <parent link="cart"/>
+    <child link="pole1"/>
+    <limit effort="1000.0" velocity="8"/>
+  </joint>
+  
+  <link name="pole1">
+    <visual>
+      <geometry>
+        <box size="0.04 0.06 1.0"/>	
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+      <material name="pole_mat">
+        <color rgba="0.1 0.1 0.3 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.04 0.06 1.0"/>	
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+    </collision>
+    <inertial>
+      <origin xyz="0 0 0.47"/>
+      <mass value="0.25"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <link name="pole2">
+    <visual>
+      <geometry>
+        <box size="0.04 0.06 1.0"/> 
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+      <material name="pole_mat">
+        <color rgba="0.1 0.1 0.3 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.04 0.06 1.0"/> 
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+    </collision>
+    <inertial>
+      <origin xyz="0 0 0.47"/>
+      <mass value="0.25"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <joint name="pole1_to_pole2" type="continuous">
+    <axis xyz="1 0 0"/>
+    <origin xyz="0.0 0 1.0"/>
+    <parent link="pole1"/>
+    <child link="pole2"/>
+    <limit effort="1000.0" velocity="8"/>
+  </joint>
+
+</robot>

warp/examples/assets/cube.usda

@@ -0,0 +1,42 @@
+#usda 1.0
+(
+    doc = "Blender v2.92.0"
+    metersPerUnit = 1
+    upAxis = "Z"
+    defaultPrim = "Cube"
+)
+
+def Xform "Cube"
+{
+    matrix4d xformOp:transform = ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) )
+    uniform token[] xformOpOrder = ["xformOp:transform"]
+
+    def Mesh "Cube_001"
+    {
+        int[] faceVertexCounts = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
+        int[] faceVertexIndices = [1, 2, 0, 3, 6, 2, 7, 4, 6, 5, 0, 4, 6, 0, 2, 3, 5, 7, 1, 3, 2, 3, 7, 6, 7, 5, 4, 5, 1, 0, 6, 4, 0, 3, 1, 5]
+        normal3f[] normals = [(-1, 0, 0), (-1, 0, 0), (-1, 0, 0), (0, 1, -0), (0, 1, -0), (0, 1, -0), (1, 0, -0), (1, 0, -0), (1, 0, -0), (0, -1, 0), (0, -1, 0), (0, -1, 0), (0, 0, -1), (0, 0, -1), (0, 0, -1), (0, 0, 1), (0, 0, 1), (0, 0, 1), (-1, 0, 0), (-1, 0, 0), (-1, 0, 0), (0, 1, -0), (0, 1, -0), (0, 1, -0), (1, 0, 0), (1, 0, 0), (1, 0, 0), (0, -1, 0), (0, -1, 0), (0, -1, 0), (0, 0, -1), (0, 0, -1), (0, 0, -1), (0, -0, 1), (0, -0, 1), (0, -0, 1)] (
+            interpolation = "faceVarying"
+        )
+        point3f[] points = [(-1, -1, -1), (-1, -1, 1), (-1, 1, -1), (-1, 1, 1), (1, -1, -1), (1, -1, 1), (1, 1, -1), (1, 1, 1)]
+        uniform token subdivisionScheme = "none"
+    }
+}
+
+def Xform "Camera"
+{
+    matrix4d xformOp:transform = ( (0.6859206557273865, 0.7276763319969177, 0, 0), (-0.32401347160339355, 0.305420845746994, 0.8953956365585327, 0), (0.6515582203865051, -0.6141703724861145, 0.44527140259742737, 0), (7.358891487121582, -6.925790786743164, 4.958309173583984, 1) )
+    uniform token[] xformOpOrder = ["xformOp:transform"]
+
+    def Camera "Camera"
+    {
+        float2 clippingRange = (0.1, 100)
+        float focalLength = 50
+        float horizontalAperture = 36
+        float horizontalApertureOffset = 0
+        token projection = "perspective"
+        float verticalAperture = 20.25
+        float verticalApertureOffset = 0
+    }
+}
+

warp/examples/assets/nv_ant.xml

@@ -0,0 +1,92 @@
+<mujoco model="ant">
+  <custom>
+    <numeric data="0.0 0.0 0.55 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0" name="init_qpos"/>
+  </custom>
+
+  <default>
+    <joint armature="0.01" damping="0.1" limited="true"/>
+    <geom condim="3" density="5.0" friction="1.5 0.1 0.1" margin="0.01" rgba="0.97 0.38 0.06 1"/>
+  </default>
+
+  <compiler inertiafromgeom="true" angle="degree"/>
+
+  <option timestep="0.016" iterations="50" tolerance="1e-10" solver="Newton" jacobian="dense" cone="pyramidal"/>
+
+  <size nconmax="50" njmax="200" nstack="10000"/>
+  <visual>
+      <map force="0.1" zfar="30"/>
+      <rgba haze="0.15 0.25 0.35 1"/>
+      <quality shadowsize="2048"/>
+      <global offwidth="800" offheight="800"/>
+  </visual>
+
+  <asset>
+      <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="512"/> 
+      <texture name="texplane" type="2d" builtin="checker" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2" width="512" height="512" mark="cross" markrgb=".8 .8 .8"/>
+      <texture name="texgeom" type="cube" builtin="flat" mark="cross" width="127" height="1278" 
+          rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1" random="0.01"/>  
+
+      <material name="matplane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true"/>
+      <material name="matgeom" texture="texgeom" texuniform="true" rgba="0.8 0.6 .4 1"/>
+  </asset>
+
+  <worldbody>
+    <geom name="floor" pos="0 0 0" size="0 0 .25" type="plane" material="matplane" condim="3"/>
+
+    <light directional="false" diffuse=".2 .2 .2" specular="0 0 0" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
+    <light mode="targetbodycom" target="torso" directional="false" diffuse=".8 .8 .8" specular="0.3 0.3 0.3" pos="0 0 4.0" dir="0 0 -1"/>
+
+    <body name="torso" pos="0 0 0.75">
+      <geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere"/>
+      <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+      <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule"/>
+      <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule"/>
+      <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule" rgba=".999 .2 .02 1"/>
+
+      <joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free"/>
+      <body name="front_left_leg" pos="0.2 0.2 0">
+        <joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        <body pos="0.2 0.2 0" name="front_left_foot">
+          <joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 100" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        </body>
+      </body>
+      <body name="front_right_leg" pos="-0.2 0.2 0">
+        <joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule"/>
+        <body pos="-0.2 0.2 0" name="front_right_foot">
+          <joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-100 -30" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule"/>
+        </body>
+      </body>
+      <body name="left_back_leg" pos="-0.2 -0.2 0">
+        <joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule"/>
+        <body pos="-0.2 -0.2 0" name="left_back_foot">
+          <joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-100 -30" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule"/>
+        </body>
+      </body>
+      <body name="right_back_leg" pos="0.2 -0.2 0">
+        <joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        <body pos="0.2 -0.2 0" name="right_back_foot">
+          <joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 100" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="15"/>
+  </actuator>
+</mujoco>