warp-lang 1.7.2__py3-none-win_amd64.whl → 1.8.0__py3-none-win_amd64.whl

warp/render/render_usd.py

@@ -80,7 +80,7 @@ class UsdRenderer:
     """A USD renderer"""
 
     def __init__(self, stage, up_axis="Y", fps=60, scaling=1.0):
-        """Construct a UsdRenderer object
+        """Construct a UsdRenderer object.
 
         Args:
             model: A simulation model
@@ -173,12 +173,13 @@ class UsdRenderer:
     def add_shape_instance(
         self,
         name: str,
-        shape,
+        shape: int,
         body,
         pos: tuple,
         rot: tuple,
         scale: tuple = (1.0, 1.0, 1.0),
-        color: tuple = (1.0, 1.0, 1.0),
+        color1=None,
+        color2=None,
         custom_index: int = -1,
         visible: bool = True,
     ):
@@ -201,12 +202,15 @@ class UsdRenderer:
         rot: tuple,
         width: float,
         length: float,
-        color: tuple = None,
-        parent_body: str = None,
+        color: tuple[float, float, float] = (1.0, 1.0, 1.0),
+        color2=None,
+        parent_body: str | None = None,
         is_template: bool = False,
+        u_scaling: float = 1.0,
+        v_scaling: float = 1.0,
+        visible: bool = True,
     ):
-        """
-        Render a plane with the given dimensions.
+        """Render a plane with the given dimensions.
 
         Args:
             name: Name of the plane
@@ -254,6 +258,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(plane, pos, rot, (1.0, 1.0, 1.0), 0.0)
 
+        plane.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_ground(self, size: float = 100.0, plane=None):
@@ -302,9 +307,10 @@ class UsdRenderer:
         pos: tuple,
         rot: tuple,
         radius: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Debug helper to add a sphere for visualization
 
@@ -343,6 +349,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(sphere, pos, rot, (1.0, 1.0, 1.0), self.time)
 
+        sphere.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_capsule(
@@ -352,12 +359,13 @@ class UsdRenderer:
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        up_axis: int = 1,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
-        """
-        Debug helper to add a capsule for visualization
+        """Debug helper to add a capsule for visualization
 
         Args:
             pos: The position of the capsule
@@ -387,7 +395,7 @@ class UsdRenderer:
 
         capsule.GetRadiusAttr().Set(float(radius))
         capsule.GetHeightAttr().Set(float(half_height * 2.0))
-        capsule.GetAxisAttr().Set("Y")
+        capsule.GetAxisAttr().Set(UP_AXIS_TOKEN[up_axis])
 
         if color is not None:
             capsule.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
@@ -397,6 +405,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(capsule, pos, rot, (1.0, 1.0, 1.0), self.time)
 
+        capsule.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_cylinder(
@@ -406,12 +415,13 @@ class UsdRenderer:
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        up_axis: int = 1,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
-        """
-        Debug helper to add a cylinder for visualization
+        """Debug helper to add a cylinder for visualization
 
         Args:
             pos: The position of the cylinder
@@ -441,7 +451,7 @@ class UsdRenderer:
 
         cylinder.GetRadiusAttr().Set(float(radius))
         cylinder.GetHeightAttr().Set(float(half_height * 2.0))
-        cylinder.GetAxisAttr().Set("Y")
+        cylinder.GetAxisAttr().Set(UP_AXIS_TOKEN[up_axis])
 
         if color is not None:
             cylinder.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
@@ -451,6 +461,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(cylinder, pos, rot, (1.0, 1.0, 1.0), self.time)
 
+        cylinder.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_cone(
@@ -460,12 +471,13 @@ class UsdRenderer:
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        up_axis: int = 1,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
-        """
-        Debug helper to add a cone for visualization
+        """Debug helper to add a cone for visualization
 
         Args:
             pos: The position of the cone
@@ -495,7 +507,7 @@ class UsdRenderer:
 
         cone.GetRadiusAttr().Set(float(radius))
         cone.GetHeightAttr().Set(float(half_height * 2.0))
-        cone.GetAxisAttr().Set("Y")
+        cone.GetAxisAttr().Set(UP_AXIS_TOKEN[up_axis])
 
         if color is not None:
             cone.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
@@ -505,6 +517,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(cone, pos, rot, (1.0, 1.0, 1.0), self.time)
 
+        cone.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_box(
@@ -513,9 +526,10 @@ class UsdRenderer:
         pos: tuple,
         rot: tuple,
         extents: tuple,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Debug helper to add a box for visualization
 
@@ -553,9 +567,10 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(cube, pos, rot, extents, self.time)
 
+        cube.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
-    def render_ref(self, name: str, path: str, pos: tuple, rot: tuple, scale: tuple, color: tuple = None):
+    def render_ref(self, name: str, path: str, pos: tuple, rot: tuple, scale: tuple, color: tuple | None = None):
         from pxr import Gf, Usd, UsdGeom
 
         ref_path = "/root/" + name
@@ -586,8 +601,10 @@ class UsdRenderer:
         rot=(0.0, 0.0, 0.0, 1.0),
         scale=(1.0, 1.0, 1.0),
         update_topology=False,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
+        smooth_shading: bool = True,
+        visible: bool = True,
     ):
         from pxr import Sdf, UsdGeom
 
@@ -635,6 +652,7 @@ class UsdRenderer:
         if not is_template:
             _usd_set_xform(mesh, pos, rot, scale, self.time)
 
+        mesh.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
         return prim_path
 
     def render_arrow(
@@ -644,12 +662,12 @@ class UsdRenderer:
         rot: tuple,
         base_radius: float,
         base_height: float,
-        cap_radius: float = None,
-        cap_height: float = None,
-        parent_body: str = None,
+        cap_radius: float | None = None,
+        cap_height: float | None = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: tuple[float, float, float] = None,
+        color: tuple[float, float, float] | None = None,
         visible: bool = True,
     ):
         from pxr import Gf, Sdf, UsdGeom
@@ -709,7 +727,7 @@ class UsdRenderer:
         name: str,
         vertices,
         indices,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
         radius: float = 0.01,
         visible: bool = True,
     ):
@@ -762,9 +780,18 @@ class UsdRenderer:
         instancer.GetScalesAttr().Set(line_scales, self.time)
         instancer.GetProtoIndicesAttr().Set([0] * num_lines, self.time)
 
-    #      instancer.GetPrimvar("displayColor").Set(line_colors, time)
+        # instancer.GetPrimvar("displayColor").Set(line_colors, time)
+
+        instancer.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
 
-    def render_line_strip(self, name: str, vertices, color: tuple, radius: float = 0.01):
+    def render_line_strip(
+        self,
+        name: str,
+        vertices,
+        color: tuple[float, float, float] | None = None,
+        radius: float = 0.01,
+        visible: bool = True,
+    ):
         from pxr import Gf, UsdGeom
 
         num_lines = int(len(vertices) - 1)
@@ -807,44 +834,80 @@ class UsdRenderer:
         instancer_capsule = UsdGeom.Capsule.Get(self.stage, instancer.GetPath().AppendChild("capsule"))
         instancer_capsule.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
 
+        instancer.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
+
     def render_points(self, name: str, points, radius, colors=None, as_spheres: bool = True, visible: bool = True):
-        from pxr import Gf, UsdGeom, Vt
+        from pxr import Sdf, UsdGeom, Vt
 
-        instancer_path = self.root.GetPath().AppendChild(name)
-        instancer = UsdGeom.PointInstancer.Get(self.stage, instancer_path)
-        radius_is_scalar = np.isscalar(radius)
-        if not instancer:
-            if colors is None or len(colors) == 3:
+        instancer_path = self._resolve_path(name)
+
+        if np.isscalar(radius):
+            radius_interp = "constant"
+        else:
+            radius_interp = "vertex"
+
+        if colors is None:
+            color_interp = "constant"
+        elif len(colors) == 3 and all(np.isscalar(x) for x in colors):
+            color_interp = "constant"
+            is_single_color = True
+        else:
+            color_interp = "vertex"
+            is_single_color = False
+
+        if as_spheres:
+            instancer = UsdGeom.PointInstancer.Get(self.stage, instancer_path)
+            if not instancer:
                 instancer = UsdGeom.PointInstancer.Define(self.stage, instancer_path)
-                instancer_sphere = UsdGeom.Sphere.Define(self.stage, instancer.GetPath().AppendChild("sphere"))
-                if radius_is_scalar:
-                    instancer_sphere.GetRadiusAttr().Set(radius)
-                else:
-                    instancer_sphere.GetRadiusAttr().Set(1.0)
-                    instancer.GetScalesAttr().Set(np.tile(radius, (3, 1)).T)
+                sphere = UsdGeom.Sphere.Define(self.stage, instancer.GetPath().AppendChild("sphere"))
+                primvar_api = UsdGeom.PrimvarsAPI(instancer)
+
+                instancer.CreatePrototypesRel().SetTargets((sphere.GetPath(),))
+                instancer.CreateProtoIndicesAttr().Set(Vt.IntArray((0,) * len(points)))
 
                 if colors is not None:
-                    instancer_sphere.GetDisplayColorAttr().Set([Gf.Vec3f(colors)], self.time)
+                    primvar_api.CreatePrimvar("displayColor", Sdf.ValueTypeNames.Color3fArray, color_interp, 1)
 
-                instancer.CreatePrototypesRel().SetTargets([instancer_sphere.GetPath()])
-                instancer.CreateProtoIndicesAttr().Set(Vt.IntArray((0,) * len(points)))
+            instancer.GetPositionsAttr().Set(points, self.time)
 
-                # set identity rotations
-                quats = [Gf.Quath(1.0, 0.0, 0.0, 0.0)] * len(points)
-                instancer.GetOrientationsAttr().Set(quats, self.time)
+            if radius_interp == "constant":
+                radius = np.tile(radius, (3, len(points))).T
             else:
-                instancer = UsdGeom.Points.Define(self.stage, instancer_path)
+                radius = np.tile(radius, (3, 1)).T
 
-                if radius_is_scalar:
-                    instancer.GetWidthsAttr().Set([radius * 2.0] * len(points))
-                else:
-                    instancer.GetWidthsAttr().Set(radius * 2.0)
+            instancer.GetScalesAttr().Set(radius, self.time)
 
-        if colors is None or len(colors) == 3:
-            instancer.GetPositionsAttr().Set(points, self.time)
+            if colors is not None:
+                if is_single_color:
+                    colors = (colors,)
+
+                primvar_api = UsdGeom.PrimvarsAPI(instancer)
+                primvar_api.GetPrimvar("displayColor").Set(colors, self.time)
         else:
+            instancer = UsdGeom.Points.Get(self.stage, instancer_path)
+            if not instancer:
+                instancer = UsdGeom.Points.Define(self.stage, instancer_path)
+
+                UsdGeom.Primvar(instancer.GetWidthsAttr()).SetInterpolation(radius_interp)
+                UsdGeom.Primvar(instancer.GetDisplayColorAttr()).SetInterpolation(color_interp)
+
             instancer.GetPointsAttr().Set(points, self.time)
-            instancer.GetDisplayColorAttr().Set(colors, self.time)
+
+            if np.isscalar(radius):
+                widths = (radius * 2.0,)
+            else:
+                widths = np.array(radius) * 2.0
+
+            instancer.GetWidthsAttr().Set(widths, self.time)
+
+            if colors is not None:
+                if len(colors) == 3:
+                    colors = (colors,)
+
+                instancer.GetDisplayColorAttr().Set(colors, self.time)
+
+        instancer.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
+        return instancer.GetPath()
 
     def update_body_transforms(self, body_q):
         from pxr import Sdf, UsdGeom

warp/sim/__init__.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from warp.utils import warn
+
 from .articulation import eval_fk, eval_ik
 from .collide import collide
 from .import_mjcf import parse_mjcf
@@ -63,3 +65,10 @@ from .utils import (
     quat_to_euler,
     velocity_at_point,
 )
+
+warn(
+    "The `warp.sim` module is deprecated and will be removed in v1.10. "
+    "Please transition to using the forthcoming Newton library instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)