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

warp/render/render_usd.py

@@ -28,9 +28,12 @@ def _usd_set_xform(xform, pos: tuple, rot: tuple, scale: tuple, time):
 
     xform_ops = xform.GetOrderedXformOps()
 
-    xform_ops[0].Set(Gf.Vec3d(float(pos[0]), float(pos[1]), float(pos[2])), time)
-    xform_ops[1].Set(Gf.Quatf(float(rot[3]), float(rot[0]), float(rot[1]), float(rot[2])), time)
-    xform_ops[2].Set(Gf.Vec3d(float(scale[0]), float(scale[1]), float(scale[2])), time)
+    if pos is not None:
+        xform_ops[0].Set(Gf.Vec3d(float(pos[0]), float(pos[1]), float(pos[2])), time)
+    if rot is not None:
+        xform_ops[1].Set(Gf.Quatf(float(rot[3]), float(rot[0]), float(rot[1]), float(rot[2])), time)
+    if scale is not None:
+        xform_ops[2].Set(Gf.Vec3d(float(scale[0]), float(scale[1]), float(scale[2])), time)
 
 
 # transforms a cylinder such that it connects the two points pos0, pos1
@@ -104,24 +107,20 @@ class UsdRenderer:
         elif up_axis == "Z":
             UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
 
-        # add default lights
-        light_0 = UsdLux.DistantLight.Define(stage, "/light_0")
-        light_0.GetPrim().CreateAttribute("intensity", Sdf.ValueTypeNames.Float, custom=False).Set(2500.0)
-        light_0.GetPrim().CreateAttribute("color", Sdf.ValueTypeNames.Color3f, custom=False).Set(
-            Gf.Vec3f(0.98, 0.85, 0.7)
-        )
-
-        UsdGeom.Xform(light_0.GetPrim()).AddRotateYOp().Set(value=(70.0))
-        UsdGeom.Xform(light_0.GetPrim()).AddRotateXOp().Set(value=(-45.0))
-
-        light_1 = UsdLux.DistantLight.Define(stage, "/light_1")
-        light_1.GetPrim().CreateAttribute("intensity", Sdf.ValueTypeNames.Float, custom=False).Set(2500.0)
-        light_1.GetPrim().CreateAttribute("color", Sdf.ValueTypeNames.Color3f, custom=False).Set(
-            Gf.Vec3f(0.62, 0.82, 0.98)
-        )
-
-        UsdGeom.Xform(light_1.GetPrim()).AddRotateYOp().Set(value=(-70.0))
-        UsdGeom.Xform(light_1.GetPrim()).AddRotateXOp().Set(value=(-45.0))
+        dome_light = UsdLux.DomeLight.Define(stage, "/dome_light")
+        dome_light.AddRotateXYZOp().Set((-90.0, -30.0, 0.0))
+        dome_light.GetEnableColorTemperatureAttr().Set(True)
+        dome_light.GetColorTemperatureAttr().Set(6150.0)
+        dome_light.GetIntensityAttr().Set(1.0)
+        dome_light.GetExposureAttr().Set(9.0)
+        dome_light.GetPrim().CreateAttribute("visibleInPrimaryRay", Sdf.ValueTypeNames.Bool).Set(False)
+
+        distant_light = UsdLux.DistantLight.Define(stage, "/distant_light")
+        distant_light.AddRotateXYZOp().Set((-35.0, 45.0, 0.0))
+        distant_light.GetEnableColorTemperatureAttr().Set(True)
+        distant_light.GetColorTemperatureAttr().Set(7250.0)
+        distant_light.GetIntensityAttr().Set(1.0)
+        distant_light.GetExposureAttr().Set(10.0)
 
     def begin_frame(self, time):
         self.time = round(time * self.fps)
@@ -155,7 +154,11 @@ class UsdRenderer:
         rot: tuple,
         scale: tuple = (1.0, 1.0, 1.0),
         color: tuple = (1.0, 1.0, 1.0),
+        custom_index: int = -1,
+        visible: bool = True,
     ):
+        if not visible:
+            return
         sdf_path = self._resolve_path(name, body)
         instance = self._shape_constructors[shape.name].Define(self.stage, sdf_path)
         instance.GetPrim().GetReferences().AddInternalReference(shape)
@@ -228,7 +231,7 @@ class UsdRenderer:
 
         return prim_path
 
-    def render_ground(self, size: float = 100.0):
+    def render_ground(self, size: float = 100.0, plane=None):
         from pxr import UsdGeom
 
         mesh = UsdGeom.Mesh.Define(self.stage, self.root.GetPath().AppendChild("ground"))
@@ -243,6 +246,23 @@ class UsdRenderer:
         elif self.up_axis == "Z":
             points = ((-size, -size, 0.0), (size, -size, 0.0), (size, size, 0.0), (-size, size, 0.0))
             normals = ((0.0, 0.0, 1.0), (0.0, 0.0, 1.0), (0.0, 0.0, 1.0), (0.0, 0.0, 1.0))
+        if plane is not None:
+            normal = np.array(plane[:3])
+            normal /= np.linalg.norm(normal)
+            pos = plane[3] * normal
+            axis_up = [0.0, 0.0, 0.0]
+            axis_up["XYZ".index(self.up_axis)] = 1.0
+            if np.allclose(normal, axis_up):
+                # no rotation necessary
+                q = (0.0, 0.0, 0.0, 1.0)
+            else:
+                c = np.cross(normal, axis_up)
+                angle = np.arcsin(np.linalg.norm(c))
+                axis = np.abs(c) / np.linalg.norm(c)
+                q = wp.quat_from_axis_angle(axis, angle)
+            tf = wp.transform(pos, q)
+            points = [wp.transform_point(tf, p) for p in points]
+            normals = [wp.transform_vector(tf, n) for n in normals]
         counts = (4,)
         indices = [0, 1, 2, 3]
 
@@ -252,7 +272,7 @@ class UsdRenderer:
         mesh.GetFaceVertexIndicesAttr().Set(indices)
 
     def render_sphere(
-        self, name: str, pos: tuple, rot: tuple, radius: float, parent_body: str = None, is_template: bool = False
+        self, name: str, pos: tuple, rot: tuple, radius: float, parent_body: str = None, is_template: bool = False, color: tuple = None
     ):
         """Debug helper to add a sphere for visualization
 
@@ -260,9 +280,10 @@ class UsdRenderer:
             pos: The position of the sphere
             radius: The radius of the sphere
             name: A name for the USD prim on the stage
+            color: The color of the sphere
         """
 
-        from pxr import UsdGeom, Sdf
+        from pxr import Gf, UsdGeom, Sdf
 
         if is_template:
             prim_path = self._resolve_path(name, parent_body, is_template)
@@ -282,6 +303,9 @@ class UsdRenderer:
 
         sphere.GetRadiusAttr().Set(radius, self.time)
 
+        if color is not None:
+            sphere.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
         self._shape_constructors[name] = UsdGeom.Sphere
 
         if not is_template:
@@ -298,6 +322,7 @@ class UsdRenderer:
         half_height: float,
         parent_body: str = None,
         is_template: bool = False,
+        color: tuple = None,
     ):
         """
         Debug helper to add a capsule for visualization
@@ -307,9 +332,10 @@ class UsdRenderer:
             radius: The radius of the capsule
             half_height: The half height of the capsule
             name: A name for the USD prim on the stage
+            color: The color of the capsule
         """
 
-        from pxr import UsdGeom, Sdf
+        from pxr import Gf, UsdGeom, Sdf
 
         if is_template:
             prim_path = self._resolve_path(name, parent_body, is_template)
@@ -331,10 +357,13 @@ class UsdRenderer:
         capsule.GetHeightAttr().Set(float(half_height * 2.0))
         capsule.GetAxisAttr().Set("Y")
 
+        if color is not None:
+            capsule.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
         self._shape_constructors[name] = UsdGeom.Capsule
 
         if not is_template:
-            _usd_set_xform(capsule, pos, rot, (1.0, 1.0, 1.0), 0.0)
+            _usd_set_xform(capsule, pos, rot, (1.0, 1.0, 1.0), self.time)
 
         return prim_path
 
@@ -347,6 +376,7 @@ class UsdRenderer:
         half_height: float,
         parent_body: str = None,
         is_template: bool = False,
+        color: tuple = None,
     ):
         """
         Debug helper to add a cylinder for visualization
@@ -356,9 +386,10 @@ class UsdRenderer:
             radius: The radius of the cylinder
             half_height: The half height of the cylinder
             name: A name for the USD prim on the stage
+            color: The color of the cylinder
         """
 
-        from pxr import UsdGeom, Sdf
+        from pxr import Gf, UsdGeom, Sdf
 
         if is_template:
             prim_path = self._resolve_path(name, parent_body, is_template)
@@ -380,10 +411,13 @@ class UsdRenderer:
         cylinder.GetHeightAttr().Set(float(half_height * 2.0))
         cylinder.GetAxisAttr().Set("Y")
 
+        if color is not None:
+            cylinder.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
         self._shape_constructors[name] = UsdGeom.Cylinder
 
         if not is_template:
-            _usd_set_xform(cylinder, pos, rot, (1.0, 1.0, 1.0), 0.0)
+            _usd_set_xform(cylinder, pos, rot, (1.0, 1.0, 1.0), self.time)
 
         return prim_path
 
@@ -396,6 +430,7 @@ class UsdRenderer:
         half_height: float,
         parent_body: str = None,
         is_template: bool = False,
+        color: tuple = None,
     ):
         """
         Debug helper to add a cone for visualization
@@ -405,9 +440,10 @@ class UsdRenderer:
             radius: The radius of the cone
             half_height: The half height of the cone
             name: A name for the USD prim on the stage
+            color: The color of the cone
         """
 
-        from pxr import UsdGeom, Sdf
+        from pxr import Gf, UsdGeom, Sdf
 
         if is_template:
             prim_path = self._resolve_path(name, parent_body, is_template)
@@ -429,25 +465,29 @@ class UsdRenderer:
         cone.GetHeightAttr().Set(float(half_height * 2.0))
         cone.GetAxisAttr().Set("Y")
 
+        if color is not None:
+            cone.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
         self._shape_constructors[name] = UsdGeom.Cone
 
         if not is_template:
-            _usd_set_xform(cone, pos, rot, (1.0, 1.0, 1.0), 0.0)
+            _usd_set_xform(cone, pos, rot, (1.0, 1.0, 1.0), self.time)
 
         return prim_path
 
     def render_box(
-        self, name: str, pos: tuple, rot: tuple, extents: tuple, parent_body: str = None, is_template: bool = False
+        self, name: str, pos: tuple, rot: tuple, extents: tuple, parent_body: str = None, is_template: bool = False, color: tuple = None
     ):
         """Debug helper to add a box for visualization
 
         Args:
-            pos: The position of the sphere
-            extents: The radius of the sphere
+            pos: The position of the box
+            extents: The radius of the box
             name: A name for the USD prim on the stage
+            color: The color of the box
         """
 
-        from pxr import UsdGeom, Sdf, Gf, Vt
+        from pxr import UsdGeom, Sdf, Gf
 
         if is_template:
             prim_path = self._resolve_path(name, parent_body, is_template)
@@ -465,11 +505,14 @@ class UsdRenderer:
             cube = UsdGeom.Cube.Define(self.stage, cube_path)
             _usd_add_xform(cube)
 
+        if color is not None:
+            cube.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
         self._shape_constructors[name] = UsdGeom.Cube
         self._shape_custom_scale[name] = extents
 
         if not is_template:
-            _usd_set_xform(cube, pos, rot, extents, 0.0)
+            _usd_set_xform(cube, pos, rot, extents, self.time)
 
         return prim_path
 
@@ -662,11 +705,10 @@ class UsdRenderer:
 
                 instancer = UsdGeom.Points.Define(self.stage, instancer_path)
 
-                instancer.CreatePrimvar("displayColor", Sdf.ValueTypeNames.Float3Array, "vertex", 1)
                 if radius_is_scalar:
-                    instancer.GetWidthsAttr().Set([radius] * len(points))
+                    instancer.GetWidthsAttr().Set([radius * 2.0] * len(points))
                 else:
-                    instancer.GetWidthsAttr().Set(radius)
+                    instancer.GetWidthsAttr().Set(radius * 2.0)
 
         if colors is None:
             instancer.GetPositionsAttr().Set(points, self.time)
@@ -694,6 +736,6 @@ class UsdRenderer:
         try:
             self.stage.Save()
             return True
-        except:
-            print("Failed to save USD stage")
+        except Exception as e:
+            print("Failed to save USD stage:", e)
             return False

warp/sim/__init__.py

@@ -6,7 +6,7 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 
-from .model import State, Model, ModelBuilder, Mesh, SDF
+from .model import State, Control, Model, ModelBuilder, Mesh, SDF
 
 from .model import GEO_SPHERE
 from .model import GEO_BOX
@@ -17,13 +17,13 @@ from .model import GEO_MESH
 from .model import GEO_SDF
 from .model import GEO_PLANE
 from .model import GEO_NONE
-from .model import ModelShapeGeometry
 
-from .model import JOINT_MODE_LIMIT
+from .model import JOINT_MODE_FORCE
 from .model import JOINT_MODE_TARGET_POSITION
 from .model import JOINT_MODE_TARGET_VELOCITY
 
 from .model import JointAxis
+from .model import ModelShapeGeometry
 from .model import ModelShapeMaterials
 
 from .model import JOINT_PRISMATIC
@@ -36,8 +36,10 @@ from .model import JOINT_UNIVERSAL
 from .model import JOINT_DISTANCE
 from .model import JOINT_D6
 
+from .integrator import Integrator, integrate_particles, integrate_bodies
 from .integrator_euler import SemiImplicitIntegrator
-from .integrator_euler import VariationalImplicitIntegrator
+
+from .integrator_featherstone import FeatherstoneIntegrator
 
 from .integrator_xpbd import XPBDIntegrator
 
@@ -48,3 +50,10 @@ from .import_mjcf import parse_mjcf
 from .import_urdf import parse_urdf
 from .import_snu import parse_snu
 from .import_usd import parse_usd, resolve_usd_from_url
+
+from .utils import velocity_at_point
+from .utils import quat_from_euler
+from .utils import quat_to_euler
+from .utils import load_mesh
+
+from .inertia import transform_inertia

warp/sim/articulation.py

@@ -105,7 +105,7 @@ def compute_3d_rotational_dofs(
     local_1 = wp.quat_rotate(q_off, wp.vec3(0.0, 1.0, 0.0))
     local_2 = wp.quat_rotate(q_off, wp.vec3(0.0, 0.0, 1.0))
 
-    # reconstruct rotation axes, todo: can probably use fact that rz'*ry'*rx' == rx*ry*rz to avoid some work here
+    # reconstruct rotation axes
     axis_0 = local_0
     q_0 = wp.quat_from_axis_angle(axis_0, q0)
 
@@ -331,10 +331,9 @@ def eval_articulation_fk(
             iq = q_start + lin_axis_count
             iqd = qd_start + lin_axis_count
             if ang_axis_count == 1:
-                axis = joint_axis[axis_start + lin_axis_count + 0]
-                qi = wp.quat_from_axis_angle(axis, joint_q[q_start + lin_axis_count + 0])
-                rot = qi * rot
-                vel_w = joint_qd[qd_start + lin_axis_count + 0] * axis
+                axis = joint_axis[ia]
+                rot = wp.quat_from_axis_angle(axis, joint_q[iq])
+                vel_w = joint_qd[iqd] * axis
             if ang_axis_count == 2:
                 rot, vel_w = compute_2d_rotational_dofs(
                     joint_axis[ia + 0],