warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/sim/import_usd.py

@@ -10,95 +10,80 @@ import numpy as np
 import re
 
 import warp as wp
-from . import ModelBuilder
 
 
 def parse_usd(
-    filename,
-    builder: ModelBuilder,
+    source,
+    builder,
     default_density=1.0e3,
     only_load_enabled_rigid_bodies=False,
     only_load_enabled_joints=True,
     default_ke=1e5,
     default_kd=250.0,
     default_kf=500.0,
-    default_mu=0.0,
+    default_mu=0.6,
     default_restitution=0.0,
     default_thickness=0.0,
     joint_limit_ke=100.0,
     joint_limit_kd=10.0,
+    invert_rotations=False,
     verbose=False,
     ignore_paths=[],
-    export_usda=False,
 ):
+    """
+    Parses a Universal Scene Description (USD) stage containing UsdPhysics schema definitions for rigid-body articulations and adds the bodies, shapes and joints to the given ModelBuilder.
+
+    The USD description has to be either a path (file name or URL), or an existing USD stage instance that implements the `UsdStage <https://openusd.org/dev/api/class_usd_stage.html>`_ interface.
+
+    Args:
+        source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance.
+        builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
+        default_density (float): The default density to use for bodies without a density attribute.
+        only_load_enabled_rigid_bodies (bool): If True, only rigid bodies which do not have `physics:rigidBodyEnabled` set to False are loaded.
+        only_load_enabled_joints (bool): If True, only joints which do not have `physics:jointEnabled` set to False are loaded.
+        default_ke (float): The default contact stiffness to use, only considered by SemiImplicitIntegrator.
+        default_kd (float): The default contact damping to use, only considered by SemiImplicitIntegrator.
+        default_kf (float): The default friction stiffness to use, only considered by SemiImplicitIntegrator.
+        default_mu (float): The default friction coefficient to use if a shape has not friction coefficient defined.
+        default_restitution (float): The default coefficient of restitution to use if a shape has not coefficient of restitution defined.
+        default_thickness (float): The thickness to add to the shape geometry.
+        joint_limit_ke (float): The default stiffness to use for joint limits, only considered by SemiImplicitIntegrator.
+        joint_limit_kd (float): The default damping to use for joint limits, only considered by SemiImplicitIntegrator.
+        invert_rotations (bool): If True, inverts any rotations defined in the shape transforms.
+        verbose (bool): If True, print additional information about the parsed USD file.
+        ignore_paths (List[str]): A list of regular expressions matching prim paths to ignore.
+
+    Returns:
+        dict: Dictionary with the following entries:
+
+        .. list-table::
+            :widths: 25 75
+
+            * - "fps"
+              - USD stage frames per second
+            * - "duration"
+              - Difference between end time code and start time code of the USD stage
+            * - "up_axis"
+              - Upper-case string of the stage's up axis ("X", "Y", or "Z")
+            * - "path_shape_map"
+              - Mapping from prim path (str) of the UsdGeom to the respective shape index in :class:`ModelBuilder`
+            * - "path_body_map"
+              - Mapping from prim path (str) of a rigid body prim (e.g. that implements the PhysicsRigidBodyAPI) to the respective body index in :class:`ModelBuilder`
+            * - "path_shape_scale"
+              - Mapping from prim path (str) of the UsdGeom to its respective 3D world scale
+            * - "mass_unit"
+              - The stage's Kilograms Per Unit (KGPU) definition (1.0 by default)
+            * - "linear_unit"
+              - The stage's Meters Per Unit (MPU) definition (1.0 by default)
+
+
+    Note:
+        This importer is experimental and only supports a subset of the USD Physics schema. Please report any issues you encounter.
+    """
     try:
         from pxr import Usd, UsdGeom, UsdPhysics
     except ImportError:
-        raise ImportError("Failed to import pxr. Please install USD.")
-
-    if filename.startswith("http://") or filename.startswith("https://"):
-        # download file
-        import requests, os, datetime
-
-        response = requests.get(filename, allow_redirects=True)
-        if response.status_code != 200:
-            raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
-        file = response.content
-        dot = os.path.extsep
-        base = os.path.basename(filename)
-        url_folder = os.path.dirname(filename)
-        base_name = dot.join(base.split(dot)[:-1])
-        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-        folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
-        os.makedirs(folder_name, exist_ok=True)
-        target_filename = os.path.join(folder_name, base)
-        with open(target_filename, "wb") as f:
-            f.write(file)
-
-        stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
-        stage_str = stage.GetRootLayer().ExportToString()
-        print(f"Downloaded USD file to {target_filename}.")
-        if export_usda:
-            usda_filename = os.path.join(folder_name, base_name + ".usda")
-            with open(usda_filename, "w") as f:
-                f.write(stage_str)
-                print(f"Exported USDA file to {usda_filename}.")
-
-        # parse referenced USD files like `references = @./franka_collisions.usd@`
-        downloaded = set()
-        for match in re.finditer(r"references.=.@(.*?)@", stage_str):
-            refname = match.group(1)
-            if refname.startswith("./"):
-                refname = refname[2:]
-            if refname in downloaded:
-                continue
-            try:
-                response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
-                if response.status_code != 200:
-                    print(f"Failed to download reference {refname}. Status code: {response.status_code}")
-                    continue
-                file = response.content
-                refdir = os.path.dirname(refname)
-                if refdir:
-                    os.makedirs(os.path.join(folder_name, refdir), exist_ok=True)
-                ref_filename = os.path.join(folder_name, refname)
-                with open(ref_filename, "wb") as f:
-                    f.write(file)
-                downloaded.add(refname)
-                print(f"Downloaded USD reference {refname} to {ref_filename}.")
-                if export_usda:
-                    ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
-                    ref_stage_str = ref_stage.GetRootLayer().ExportToString()
-                    base = os.path.basename(ref_filename)
-                    base_name = dot.join(base.split(dot)[:-1])
-                    usda_filename = os.path.join(folder_name, base_name + ".usda")
-                    with open(usda_filename, "w") as f:
-                        f.write(ref_stage_str)
-                        print(f"Exported USDA file to {usda_filename}.")
-            except:
-                print(f"Failed to download {refname}.")
-
-        filename = target_filename
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
 
     def get_attribute(prim, name):
         if "*" in name:
@@ -127,7 +112,10 @@ def parse_usd(
         if not attr or not attr.HasAuthoredValue():
             return default
         val = attr.Get()
-        quat = wp.quat(*val.imaginary, val.real)
+        if invert_rotations:
+            quat = wp.quat(*val.imaginary, -val.real)
+        else:
+            quat = wp.quat(*val.imaginary, val.real)
         l = wp.length(quat)
         if np.isfinite(l) and l > 0.0:
             return quat
@@ -153,20 +141,31 @@ def parse_usd(
         axis["XYZ".index(s.upper())] = 1.0
         return axis
 
-    stage = Usd.Stage.Open(filename, Usd.Stage.LoadAll)
-    if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
-        mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
+    if isinstance(source, str):
+        stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
     else:
-        mass_unit = 1.0
-    if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
-        linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
-    else:
-        linear_unit = 1.0
+        stage = source
+
+    mass_unit = 1.0
+    try:
+        if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
+            mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
+    except:
+        pass
+    linear_unit = 1.0
+    try:
+        if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
+            linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
+    except:
+        pass
 
     def parse_xform(prim):
         xform = UsdGeom.Xform(prim)
         mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
-        rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
+        if invert_rotations:
+            rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].T.flatten()))
+        else:
+            rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
         pos = mat[3, :3] * linear_unit
         scale = np.ones(3, dtype=np.float32)
         for op in xform.GetOrderedXformOps():
@@ -247,12 +246,18 @@ def parse_usd(
         else:
             joint_data["linear_axes"].append(axis)
 
-    upaxis = str2axis(UsdGeom.GetStageUpAxis(stage))
+    axis_str = "Y"
+    try:
+        axis_str = UsdGeom.GetStageUpAxis(stage)
+    except:
+        pass
+    upaxis = str2axis(axis_str)
 
     shape_types = {"Cube", "Sphere", "Mesh", "Capsule", "Plane", "Cylinder", "Cone"}
 
     path_body_map = {}
     path_shape_map = {}
+    path_shape_scale = {}
     # maps prim path name to its world transform
     path_world_poses = {}
     # transform from body frame to where the actual joint child frame is
@@ -347,21 +352,27 @@ def parse_usd(
             materials[path] = material
 
         elif type_name == "PhysicsScene":
-            scene = UsdPhysics.Scene(prim)
-            g_vec = scene.GetGravityDirectionAttr()
-            g_mag = scene.GetGravityMagnitudeAttr()
-            if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
-                builder.gravity = g_mag.Get() * linear_unit
-            if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
-                builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)  # TODO flip sign?
-            else:
-                builder.up_vector = upaxis
+            try:
+                scene = UsdPhysics.Scene(prim)
+                g_vec = scene.GetGravityDirectionAttr()
+                g_mag = scene.GetGravityMagnitudeAttr()
+                if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
+                    builder.gravity = -np.abs(g_mag.Get() * linear_unit)
+                if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
+                    builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)
+                    if np.any(builder.up_vector < 0.0):
+                        builder.up_vector = -builder.up_vector
+                else:
+                    builder.up_vector = upaxis
+            except:
+                pass
 
     def parse_prim(prim, incoming_xform, incoming_scale, parent_body: int = -1):
         nonlocal builder
         nonlocal joint_data
         nonlocal path_body_map
         nonlocal path_shape_map
+        nonlocal path_shape_scale
         nonlocal path_world_poses
         nonlocal prim_joint_xforms
         nonlocal path_collision_filters
@@ -374,13 +385,14 @@ def parse_usd(
                 return
 
         type_name = str(prim.GetTypeName())
-        if (
-            type_name.endswith("Joint")
-            or type_name.endswith("Light")
-            or type_name.endswith("Scene")
-            or type_name.endswith("Material")
-        ):
+        if type_name.endswith("Joint") or type_name.endswith("Light") or type_name.endswith("Material"):
             return
+        if verbose:
+            print(f"parse_prim {prim.GetPath()} ({type_name})")
+        if type_name == "PhysicsScene":
+            # in case the PhysicsScene has bodies as children...
+            for child in prim.GetChildren():
+                parse_prim(child, incoming_xform, incoming_scale, parent_body)
 
         schemas = set(prim.GetAppliedSchemas())
         children_refs = prim.GetChildren()
@@ -696,6 +708,7 @@ def parse_usd(
 
             path_body_map[path] = body_id
             path_shape_map[path] = shape_id
+            path_shape_scale[path] = scale
 
             if prim.HasRelationship("physics:filteredPairs"):
                 other_paths = prim.GetRelationship("physics:filteredPairs").GetTargets()
@@ -761,5 +774,93 @@ def parse_usd(
     return {
         "fps": stage.GetFramesPerSecond(),
         "duration": stage.GetEndTimeCode() - stage.GetStartTimeCode(),
-        "up_axis": UsdGeom.GetStageUpAxis(stage).lower(),
+        "up_axis": UsdGeom.GetStageUpAxis(stage).upper(),
+        "path_shape_map": path_shape_map,
+        "path_body_map": path_body_map,
+        "path_shape_scale": path_shape_scale,
+        "mass_unit": mass_unit,
+        "linear_unit": linear_unit,
     }
+
+
+def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
+    """
+    Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
+
+    Args:
+        url (str): URL to the USD file.
+        target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
+        export_usda (bool): If True, converts each downloaded USD file to USDA and saves the additional USDA file in the target folder with the same base name as the original USD file.
+
+    Returns:
+        str: File path to the downloaded USD file.
+    """
+    import requests
+    import datetime
+    import os
+
+    try:
+        from pxr import Usd
+    except ImportError:
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
+
+    response = requests.get(url, allow_redirects=True)
+    if response.status_code != 200:
+        raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
+    file = response.content
+    dot = os.path.extsep
+    base = os.path.basename(url)
+    url_folder = os.path.dirname(url)
+    base_name = dot.join(base.split(dot)[:-1])
+    if target_folder_name is None:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        target_folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
+    os.makedirs(target_folder_name, exist_ok=True)
+    target_filename = os.path.join(target_folder_name, base)
+    with open(target_filename, "wb") as f:
+        f.write(file)
+
+    stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
+    stage_str = stage.GetRootLayer().ExportToString()
+    print(f"Downloaded USD file to {target_filename}.")
+    if export_usda:
+        usda_filename = os.path.join(target_folder_name, base_name + ".usda")
+        with open(usda_filename, "w") as f:
+            f.write(stage_str)
+            print(f"Exported USDA file to {usda_filename}.")
+
+    # parse referenced USD files like `references = @./franka_collisions.usd@`
+    downloaded = set()
+    for match in re.finditer(r"references.=.@(.*?)@", stage_str):
+        refname = match.group(1)
+        if refname.startswith("./"):
+            refname = refname[2:]
+        if refname in downloaded:
+            continue
+        try:
+            response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
+            if response.status_code != 200:
+                print(f"Failed to download reference {refname}. Status code: {response.status_code}")
+                continue
+            file = response.content
+            refdir = os.path.dirname(refname)
+            if refdir:
+                os.makedirs(os.path.join(target_folder_name, refdir), exist_ok=True)
+            ref_filename = os.path.join(target_folder_name, refname)
+            if not os.path.exists(ref_filename):
+                with open(ref_filename, "wb") as f:
+                    f.write(file)
+            downloaded.add(refname)
+            print(f"Downloaded USD reference {refname} to {ref_filename}.")
+            if export_usda:
+                ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
+                ref_stage_str = ref_stage.GetRootLayer().ExportToString()
+                base = os.path.basename(ref_filename)
+                base_name = dot.join(base.split(dot)[:-1])
+                usda_filename = os.path.join(target_folder_name, base_name + ".usda")
+                with open(usda_filename, "w") as f:
+                    f.write(ref_stage_str)
+                    print(f"Exported USDA file to {usda_filename}.")
+        except Exception:
+            print(f"Failed to download {refname}.")
+    return target_filename

warp/sim/inertia.py

@@ -149,9 +149,9 @@ def compute_sphere_inertia(density: float, r: float) -> tuple:
     m = density * v
     Ia = 2.0 / 5.0 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ia, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ia, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_capsule_inertia(density: float, r: float, h: float) -> tuple:
@@ -177,9 +177,9 @@ def compute_capsule_inertia(density: float, r: float, h: float) -> tuple:
     Ia = mc * (0.25 * r * r + (1.0 / 12.0) * h * h) + ms * (0.4 * r * r + 0.375 * r * h + 0.25 * h * h)
     Ib = (mc * 0.5 + ms * 0.4) * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_cylinder_inertia(density: float, r: float, h: float) -> tuple:
@@ -200,9 +200,9 @@ def compute_cylinder_inertia(density: float, r: float, h: float) -> tuple:
     Ia = 1 / 12 * m * (3 * r * r + h * h)
     Ib = 1 / 2 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_cone_inertia(density: float, r: float, h: float) -> tuple:
@@ -223,9 +223,9 @@ def compute_cone_inertia(density: float, r: float, h: float) -> tuple:
     Ia = 1 / 20 * m * (3 * r * r + 2 * h * h)
     Ib = 3 / 10 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_box_inertia(density: float, w: float, h: float, d: float) -> tuple:
@@ -249,17 +249,16 @@ def compute_box_inertia(density: float, w: float, h: float, d: float) -> tuple:
     Ib = 1.0 / 12.0 * m * (w * w + d * d)
     Ic = 1.0 / 12.0 * m * (w * w + h * h)
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ic]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ic]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_mesh_inertia(
     density: float, vertices: list, indices: list, is_solid: bool = True, thickness: Union[List[float], float] = 0.001
 ) -> tuple:
     """Computes mass, center of mass, 3x3 inertia matrix, and volume for a mesh."""
-    com = np.mean(vertices, 0)
-    com_warp = wp.vec3(com[0], com[1], com[2])
+    com = wp.vec3(np.mean(vertices, 0))
 
     num_tris = len(indices) // 3
 
@@ -279,7 +278,7 @@ def compute_mesh_inertia(
             kernel=compute_solid_mesh_inertia,
             dim=num_tris,
             inputs=[
-                com_warp,
+                com,
                 weight,
                 wp.array(indices, dtype=int),
                 wp.array(vertices, dtype=wp.vec3),
@@ -294,7 +293,7 @@ def compute_mesh_inertia(
             kernel=compute_hollow_mesh_inertia,
             dim=num_tris,
             inputs=[
-                com_warp,
+                com,
                 weight,
                 wp.array(indices, dtype=int),
                 wp.array(vertices, dtype=wp.vec3),
@@ -304,14 +303,14 @@ def compute_mesh_inertia(
         )
 
     # Extract mass and inertia and save to class attributes.
-    mass = mass_warp.numpy()[0] * density
-    I = I_warp.numpy()[0] * density
+    mass = float(mass_warp.numpy()[0] * density)
+    I = wp.mat33(*(I_warp.numpy()[0] * density))
     volume = vol_warp.numpy()[0]
     return mass, com, I, volume
 
 
 def transform_inertia(m, I, p, q):
-    R = np.array(wp.quat_to_matrix(q)).reshape(3, 3)
+    R = wp.quat_to_matrix(q)
 
     # Steiner's theorem
-    return R @ I @ R.T + m * (np.dot(p, p) * np.eye(3) - np.outer(p, p))
+    return R @ I @ wp.transpose(R) + m * (wp.dot(p, p) * wp.mat33(np.eye(3)) - wp.outer(p, p))

warp/sim/integrator_euler.py

@@ -12,11 +12,10 @@ models + state forward in time.
 
 import warp as wp
 
-from .model import PARTICLE_FLAG_ACTIVE, ModelShapeMaterials, ModelShapeGeometry
-
+from .collide import triangle_closest_point_barycentric
+from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry, ModelShapeMaterials
 from .optimizer import Optimizer
 from .particles import eval_particle_forces
-from .collide import triangle_closest_point_barycentric
 from .utils import quat_decompose, quat_twist
 
 
@@ -29,6 +28,7 @@ def integrate_particles(
     particle_flags: wp.array(dtype=wp.uint32),
     gravity: wp.vec3,
     dt: float,
+    v_max: float,
     x_new: wp.array(dtype=wp.vec3),
     v_new: wp.array(dtype=wp.vec3),
 ):
@@ -44,6 +44,10 @@ def integrate_particles(
 
     # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
     v1 = v0 + (f0 * inv_mass + gravity * wp.step(0.0 - inv_mass)) * dt
+    # enforce velocity limit to prevent instability
+    v1_mag = wp.length(v1)
+    if v1_mag > v_max:
+        v1 *= v_max / v1_mag
     x1 = x0 + v1 * dt
 
     x_new[tid] = x1
@@ -837,10 +841,10 @@ def eval_particle_ground_contacts(
 def eval_particle_contacts(
     particle_x: wp.array(dtype=wp.vec3),
     particle_v: wp.array(dtype=wp.vec3),
-    particle_radius: wp.array(dtype=float),
-    particle_flags: wp.array(dtype=wp.uint32),
     body_q: wp.array(dtype=wp.transform),
     body_qd: wp.array(dtype=wp.spatial_vector),
+    particle_radius: wp.array(dtype=float),
+    particle_flags: wp.array(dtype=wp.uint32),
     body_com: wp.array(dtype=wp.vec3),
     shape_body: wp.array(dtype=int),
     shape_materials: ModelShapeMaterials,
@@ -1664,10 +1668,10 @@ def compute_forces(model, state, particle_f, body_f, requires_grad):
             inputs=[
                 state.particle_q,
                 state.particle_qd,
-                model.particle_radius,
-                model.particle_flags,
                 state.body_q,
                 state.body_qd,
+                model.particle_radius,
+                model.particle_flags,
                 model.body_com,
                 model.shape_body,
                 model.shape_materials,
@@ -1825,6 +1829,7 @@ class SemiImplicitIntegrator:
                         model.particle_flags,
                         model.gravity,
                         dt,
+                        model.particle_max_velocity,
                     ],
                     outputs=[state_out.particle_q, state_out.particle_qd],
                     device=model.device,
@@ -1913,6 +1918,7 @@ def init_state(model, state_in, state_out, dt):
             model.particle_flags,
             model.gravity,
             dt,
+            model.particle_max_velocity,
         ],
         outputs=[state_out.particle_q, state_out.particle_qd],
         device=model.device,
@@ -1956,7 +1962,7 @@ class VariationalImplicitIntegrator:
                     compute_forces(model, state_out, self.particle_f, None)
                     compute_residual(model, state_in, state_out, self.particle_f, dfdx, dt)
 
-                # initialize oututput state using the input velocity to create 'predicted state'
+                # initialize output state using the input velocity to create 'predicted state'
                 init_state(model, state_in, state_out, dt)
 
                 # our optimization variable