warp-lang 1.2.2__py3-none-manylinux2014_aarch64.whl → 1.3.0__py3-none-manylinux2014_aarch64.whl

warp/sim/articulation.py

@@ -164,12 +164,10 @@ def invert_3d_rotational_dofs(
     return angles, velocities
 
 
-@wp.kernel
-def eval_articulation_fk(
-    articulation_start: wp.array(dtype=int),
-    articulation_mask: wp.array(
-        dtype=int
-    ),  # used to enable / disable FK for an articulation, if None then treat all as enabled
+@wp.func
+def eval_single_articulation_fk(
+    joint_start: int,
+    joint_end: int,
     joint_q: wp.array(dtype=float),
     joint_qd: wp.array(dtype=float),
     joint_q_start: wp.array(dtype=int),
@@ -187,16 +185,6 @@ def eval_articulation_fk(
     body_q: wp.array(dtype=wp.transform),
     body_qd: wp.array(dtype=wp.spatial_vector),
 ):
-    tid = wp.tid()
-
-    # early out if disabling FK for this articulation
-    if articulation_mask:
-        if articulation_mask[tid] == 0:
-            return
-
-    joint_start = articulation_start[tid]
-    joint_end = articulation_start[tid + 1]
-
     for i in range(joint_start, joint_end):
         parent = joint_parent[i]
         child = joint_child[i]
@@ -374,6 +362,118 @@ def eval_articulation_fk(
         body_qd[child] = v_wc
 
 
+# implementation where mask is an integer array
+@wp.kernel
+def eval_articulation_fk(
+    articulation_start: wp.array(dtype=int),
+    articulation_mask: wp.array(
+        dtype=int
+    ),  # used to enable / disable FK for an articulation, if None then treat all as enabled
+    joint_q: wp.array(dtype=float),
+    joint_qd: wp.array(dtype=float),
+    joint_q_start: wp.array(dtype=int),
+    joint_qd_start: wp.array(dtype=int),
+    joint_type: wp.array(dtype=int),
+    joint_parent: wp.array(dtype=int),
+    joint_child: wp.array(dtype=int),
+    joint_X_p: wp.array(dtype=wp.transform),
+    joint_X_c: wp.array(dtype=wp.transform),
+    joint_axis: wp.array(dtype=wp.vec3),
+    joint_axis_start: wp.array(dtype=int),
+    joint_axis_dim: wp.array(dtype=int, ndim=2),
+    body_com: wp.array(dtype=wp.vec3),
+    # outputs
+    body_q: wp.array(dtype=wp.transform),
+    body_qd: wp.array(dtype=wp.spatial_vector),
+):
+    tid = wp.tid()
+
+    # early out if disabling FK for this articulation
+    if articulation_mask:
+        if articulation_mask[tid] == 0:
+            return
+
+    joint_start = articulation_start[tid]
+    joint_end = articulation_start[tid + 1]
+
+    eval_single_articulation_fk(
+        joint_start,
+        joint_end,
+        joint_q,
+        joint_qd,
+        joint_q_start,
+        joint_qd_start,
+        joint_type,
+        joint_parent,
+        joint_child,
+        joint_X_p,
+        joint_X_c,
+        joint_axis,
+        joint_axis_start,
+        joint_axis_dim,
+        body_com,
+        # outputs
+        body_q,
+        body_qd,
+    )
+
+
+# overload where mask is a bool array
+@wp.kernel
+def eval_articulation_fk(
+    articulation_start: wp.array(dtype=int),
+    articulation_mask: wp.array(
+        dtype=bool
+    ),  # used to enable / disable FK for an articulation, if None then treat all as enabled
+    joint_q: wp.array(dtype=float),
+    joint_qd: wp.array(dtype=float),
+    joint_q_start: wp.array(dtype=int),
+    joint_qd_start: wp.array(dtype=int),
+    joint_type: wp.array(dtype=int),
+    joint_parent: wp.array(dtype=int),
+    joint_child: wp.array(dtype=int),
+    joint_X_p: wp.array(dtype=wp.transform),
+    joint_X_c: wp.array(dtype=wp.transform),
+    joint_axis: wp.array(dtype=wp.vec3),
+    joint_axis_start: wp.array(dtype=int),
+    joint_axis_dim: wp.array(dtype=int, ndim=2),
+    body_com: wp.array(dtype=wp.vec3),
+    # outputs
+    body_q: wp.array(dtype=wp.transform),
+    body_qd: wp.array(dtype=wp.spatial_vector),
+):
+    tid = wp.tid()
+
+    # early out if disabling FK for this articulation
+    if articulation_mask:
+        if not articulation_mask[tid]:
+            return
+
+    joint_start = articulation_start[tid]
+    joint_end = articulation_start[tid + 1]
+
+    eval_single_articulation_fk(
+        joint_start,
+        joint_end,
+        joint_q,
+        joint_qd,
+        joint_q_start,
+        joint_qd_start,
+        joint_type,
+        joint_parent,
+        joint_child,
+        joint_X_p,
+        joint_X_c,
+        joint_axis,
+        joint_axis_start,
+        joint_axis_dim,
+        body_com,
+        # outputs
+        body_q,
+        body_qd,
+    )
+
+
 # updates state body information based on joint coordinates
 def eval_fk(model, joint_q, joint_qd, mask, state):
     """
@@ -383,7 +483,7 @@ def eval_fk(model, joint_q, joint_qd, mask, state):
         model (Model): The model to evaluate.
         joint_q (array): Generalized joint position coordinates, shape [joint_coord_count], float
         joint_qd (array): Generalized joint velocity coordinates, shape [joint_dof_count], float
-        mask (array): The mask to use to enable / disable FK for an articulation. If None then treat all as enabled, shape [articulation_count], int
+        mask (array): The mask to use to enable / disable FK for an articulation. If None then treat all as enabled, shape [articulation_count], int/bool
         state (State): The state to update.
     """
     wp.launch(

warp/sim/collide.py

@@ -859,10 +859,9 @@ def broadphase_collision_pairs(
                 contact_shape0[index + num_contacts_a + i] = actual_shape_b
                 contact_shape1[index + num_contacts_a + i] = actual_shape_a
                 contact_point_id[index + num_contacts_a + i] = i
-            contact_point_limit[pair_index_ab] = 2
-            if mesh_contact_max > 0:
+            if mesh_contact_max > 0 and contact_point_limit and pair_index_ba < contact_point_limit.shape[0]:
                 num_contacts_b = wp.min(mesh_contact_max, num_contacts_b)
-            contact_point_limit[pair_index_ba] = num_contacts_b
+                contact_point_limit[pair_index_ba] = num_contacts_b
             return
         else:
             num_contacts = 2
@@ -877,13 +876,11 @@ def broadphase_collision_pairs(
                 contact_shape0[index + i] = shape_a
                 contact_shape1[index + i] = shape_b
                 contact_point_id[index + i] = i
-            contact_point_limit[pair_index_ab] = 12
             # allocate contact points from box B against A
             for i in range(12):
                 contact_shape0[index + 12 + i] = shape_b
                 contact_shape1[index + 12 + i] = shape_a
                 contact_point_id[index + 12 + i] = i
-            contact_point_limit[pair_index_ba] = 12
             return
         elif actual_type_b == wp.sim.GEO_MESH:
             num_contacts_a = 8
@@ -908,10 +905,9 @@ def broadphase_collision_pairs(
                 contact_shape1[index + num_contacts_a + i] = actual_shape_a
                 contact_point_id[index + num_contacts_a + i] = i
 
-            contact_point_limit[pair_index_ab] = num_contacts_a
-            if mesh_contact_max > 0:
+            if mesh_contact_max > 0 and contact_point_limit and pair_index_ba < contact_point_limit.shape[0]:
                 num_contacts_b = wp.min(mesh_contact_max, num_contacts_b)
-            contact_point_limit[pair_index_ba] = num_contacts_b
+                contact_point_limit[pair_index_ba] = num_contacts_b
             return
         elif actual_type_b == wp.sim.GEO_PLANE:
             if geo.scale[actual_shape_b][0] == 0.0 and geo.scale[actual_shape_b][1] == 0.0:
@@ -947,11 +943,13 @@ def broadphase_collision_pairs(
                 contact_shape1[index + num_contacts_a + i] = actual_shape_a
                 contact_point_id[index + num_contacts_a + i] = i
 
-            if mesh_contact_max > 0:
+            if mesh_contact_max > 0 and contact_point_limit:
                 num_contacts_a = wp.min(mesh_contact_max, num_contacts_a)
                 num_contacts_b = wp.min(mesh_contact_max, num_contacts_b)
-            contact_point_limit[pair_index_ab] = num_contacts_a
-            contact_point_limit[pair_index_ba] = num_contacts_b
+                if pair_index_ab < contact_point_limit.shape[0]:
+                    contact_point_limit[pair_index_ab] = num_contacts_a
+                if pair_index_ba < contact_point_limit.shape[0]:
+                    contact_point_limit[pair_index_ba] = num_contacts_b
         return
     elif actual_type_a == wp.sim.GEO_PLANE:
         return  # no plane-plane contacts
@@ -969,8 +967,11 @@ def broadphase_collision_pairs(
             contact_shape0[cp_index] = actual_shape_a
             contact_shape1[cp_index] = actual_shape_b
             contact_point_id[cp_index] = i
-        contact_point_limit[pair_index_ab] = num_contacts
-        contact_point_limit[pair_index_ba] = 0
+        if contact_point_limit:
+            if pair_index_ab < contact_point_limit.shape[0]:
+                contact_point_limit[pair_index_ab] = num_contacts
+            if pair_index_ba < contact_point_limit.shape[0]:
+                contact_point_limit[pair_index_ba] = 0
 
 
 @wp.kernel
@@ -1005,12 +1006,14 @@ def handle_contact_pairs(
     if shape_a == shape_b:
         return
 
+    if contact_point_limit:
+        pair_index = shape_a * num_shapes + shape_b
+        contact_limit = contact_point_limit[pair_index]
+        if contact_pairwise_counter[pair_index] >= contact_limit:
+            # reached limit of contact points per contact pair
+            return
+
     point_id = contact_point_id[tid]
-    pair_index = shape_a * num_shapes + shape_b
-    contact_limit = contact_point_limit[pair_index]
-    if contact_pairwise_counter[pair_index] >= contact_limit:
-        # reached limit of contact points per contact pair
-        return
 
     rigid_a = shape_body[shape_a]
     X_wb_a = wp.transform_identity()
@@ -1404,15 +1407,16 @@ def handle_contact_pairs(
 
     d = distance - thickness
     if d < rigid_contact_margin:
-        pair_contact_id = limited_counter_increment(
-            contact_pairwise_counter, pair_index, contact_tids, tid, contact_limit
-        )
-        if pair_contact_id == -1:
-            # wp.printf("Reached contact point limit %d >= %d for shape pair %d and %d (pair_index: %d)\n",
-            #           contact_pairwise_counter[pair_index], contact_limit, shape_a, shape_b, pair_index)
-            # reached contact point limit
-            return
-        index = limited_counter_increment(contact_count, 0, contact_tids, tid, -1)
+        if contact_pairwise_counter:
+            pair_contact_id = limited_counter_increment(
+                contact_pairwise_counter, pair_index, contact_tids, tid, contact_limit
+            )
+            if pair_contact_id == -1:
+                # wp.printf("Reached contact point limit %d >= %d for shape pair %d and %d (pair_index: %d)\n",
+                #           contact_pairwise_counter[pair_index], contact_limit, shape_a, shape_b, pair_index)
+                # reached contact point limit
+                return
+        index = counter_increment(contact_count, 0, contact_tids, tid)
         contact_shape0[index] = shape_a
         contact_shape1[index] = shape_b
         # transform from world into body frame (so the contact point includes the shape transform)
@@ -1550,14 +1554,16 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
                 model.rigid_contact_normal = wp.clone(model.rigid_contact_normal)
                 model.rigid_contact_thickness = wp.clone(model.rigid_contact_thickness)
                 model.rigid_contact_count = wp.zeros_like(model.rigid_contact_count)
-                model.rigid_contact_pairwise_counter = wp.zeros_like(model.rigid_contact_pairwise_counter)
                 model.rigid_contact_tids = wp.zeros_like(model.rigid_contact_tids)
                 model.rigid_contact_shape0 = wp.empty_like(model.rigid_contact_shape0)
                 model.rigid_contact_shape1 = wp.empty_like(model.rigid_contact_shape1)
+                if model.rigid_contact_pairwise_counter is not None:
+                    model.rigid_contact_pairwise_counter = wp.zeros_like(model.rigid_contact_pairwise_counter)
             else:
                 model.rigid_contact_count.zero_()
-                model.rigid_contact_pairwise_counter.zero_()
                 model.rigid_contact_tids.zero_()
+                if model.rigid_contact_pairwise_counter is not None:
+                    model.rigid_contact_pairwise_counter.zero_()
             model.rigid_contact_shape0.fill_(-1)
             model.rigid_contact_shape1.fill_(-1)
 

warp/sim/model.py

@@ -561,7 +561,7 @@ class Model:
         joint_X_p (array): Joint transform in parent frame, shape [joint_count, 7], float
         joint_X_c (array): Joint mass frame in child frame, shape [joint_count, 7], float
         joint_axis (array): Joint axis in child frame, shape [joint_axis_count, 3], float
-        joint_armature (array): Armature for each joint axis (only used by :class:`FeatherstoneIntegrator`), shape [joint_count], float
+        joint_armature (array): Armature for each joint axis (only used by :class:`FeatherstoneIntegrator`), shape [joint_dof_count], float
         joint_target_ke (array): Joint stiffness, shape [joint_axis_count], float
         joint_target_kd (array): Joint damping, shape [joint_axis_count], float
         joint_axis_start (array): Start index of the first axis per joint, shape [joint_count], int
@@ -1014,11 +1014,23 @@ class Model:
             target.rigid_contact_broad_shape0 = wp.zeros(self.rigid_contact_max, dtype=wp.int32)
             target.rigid_contact_broad_shape1 = wp.zeros(self.rigid_contact_max, dtype=wp.int32)
 
-            max_pair_count = self.shape_count * self.shape_count
-            # maximum number of contact points per contact pair
-            target.rigid_contact_point_limit = wp.zeros(max_pair_count, dtype=wp.int32)
-            # currently found contacts per contact pair
-            target.rigid_contact_pairwise_counter = wp.zeros(max_pair_count, dtype=wp.int32)
+            if self.rigid_mesh_contact_max > 0:
+                # add additional buffers to track how many contact points are generated per contact pair
+                # (significantly increases memory usage, only enable if mesh contacts need to be pruned)
+                if self.shape_count >= 46340:
+                    # clip the number of potential contacts to avoid signed 32-bit integer overflow
+                    # i.e. when the number of shapes exceeds sqrt(2**31 - 1)
+                    max_pair_count = 2**31 - 1
+                else:
+                    max_pair_count = self.shape_count * self.shape_count
+                # maximum number of contact points per contact pair
+                target.rigid_contact_point_limit = wp.zeros(max_pair_count, dtype=wp.int32)
+                # currently found contacts per contact pair
+                target.rigid_contact_pairwise_counter = wp.zeros(max_pair_count, dtype=wp.int32)
+            else:
+                target.rigid_contact_point_limit = None
+                target.rigid_contact_pairwise_counter = None
+
             # ID of thread that found the current contact point
             target.rigid_contact_tids = wp.zeros(self.rigid_contact_max, dtype=wp.int32)
 
@@ -1193,7 +1205,6 @@ class ModelBuilder:
         self.body_shapes = {}  # mapping from body to shapes
 
         # rigid joints
-        self.joint = {}
         self.joint_parent = []  # index of the parent body                      (constant)
         self.joint_parents = {}  # mapping from joint to parent bodies
         self.joint_child = []  # index of the child body                       (constant)
@@ -1233,7 +1244,7 @@ class ModelBuilder:
         self.joint_axis_total_count = 0
 
         self.up_vector = wp.vec3(up_vector)
-        self.up_axis = wp.vec3(np.argmax(np.abs(up_vector)))
+        self.up_axis = int(np.argmax(np.abs(up_vector)))
         self.gravity = gravity
         # indicates whether a ground plane has been created
         self._ground_created = False
@@ -2262,8 +2273,30 @@ class ModelBuilder:
             enabled=enabled,
         )
 
-    def plot_articulation(self, plot_shapes=True):
-        """Plots the model's articulation."""
+    def plot_articulation(
+        self,
+        show_body_names=True,
+        show_joint_names=True,
+        show_joint_types=True,
+        plot_shapes=True,
+        show_shape_types=True,
+        show_legend=True,
+    ):
+        """
+        Visualizes the model's articulation graph using matplotlib and networkx.
+        Uses the spring layout algorithm from networkx to arrange the nodes.
+        Bodies are shown as orange squares, shapes are shown as blue circles.
+
+        Args:
+            show_body_names (bool): Whether to show the body names or indices
+            show_joint_names (bool): Whether to show the joint names or indices
+            show_joint_types (bool): Whether to show the joint types
+            plot_shapes (bool): Whether to render the shapes connected to the rigid bodies
+            show_shape_types (bool): Whether to show the shape geometry types
+            show_legend (bool): Whether to show a legend
+        """
+        import matplotlib.pyplot as plt
+        import networkx as nx
 
         def joint_type_str(type):
             if type == JOINT_FREE:
@@ -2286,18 +2319,88 @@ class ModelBuilder:
                 return "distance"
             return "unknown"
 
-        vertices = ["world"] + self.body_name
+        def shape_type_str(type):
+            if type == GEO_SPHERE:
+                return "sphere"
+            if type == GEO_BOX:
+                return "box"
+            if type == GEO_CAPSULE:
+                return "capsule"
+            if type == GEO_CYLINDER:
+                return "cylinder"
+            if type == GEO_CONE:
+                return "cone"
+            if type == GEO_MESH:
+                return "mesh"
+            if type == GEO_SDF:
+                return "sdf"
+            if type == GEO_PLANE:
+                return "plane"
+            if type == GEO_NONE:
+                return "none"
+            return "unknown"
+
+        if show_body_names:
+            vertices = ["world"] + self.body_name
+        else:
+            vertices = ["-1"] + [str(i) for i in range(self.body_count)]
         if plot_shapes:
-            vertices += [f"shape_{i}" for i in range(self.shape_count)]
+            for i in range(self.shape_count):
+                shape_label = f"shape_{i}"
+                if show_shape_types:
+                    shape_label += f"\n({shape_type_str(self.shape_geo_type[i])})"
+                vertices.append(shape_label)
         edges = []
         edge_labels = []
         for i in range(self.joint_count):
-            edges.append((self.joint_child[i] + 1, self.joint_parent[i] + 1))
-            edge_labels.append(f"{self.joint_name[i]}\n({joint_type_str(self.joint_type[i])})")
+            edge = (self.joint_child[i] + 1, self.joint_parent[i] + 1)
+            edges.append(edge)
+            if show_joint_names:
+                joint_label = self.joint_name[i]
+            else:
+                joint_label = str(i)
+            if show_joint_types:
+                joint_label += f"\n({joint_type_str(self.joint_type[i])})"
+            edge_labels.append(joint_label)
+
         if plot_shapes:
             for i in range(self.shape_count):
                 edges.append((len(self.body_name) + i + 1, self.shape_body[i] + 1))
-        wp.plot_graph(vertices, edges, edge_labels=edge_labels)
+
+        # plot graph
+        G = nx.Graph()
+        for i in range(len(vertices)):
+            G.add_node(i, label=vertices[i])
+        for i in range(len(edges)):
+            label = edge_labels[i] if i < len(edge_labels) else ""
+            G.add_edge(edges[i][0], edges[i][1], label=label)
+        pos = nx.spring_layout(G)
+        nx.draw_networkx_edges(G, pos, node_size=0, edgelist=edges[: self.joint_count])
+        # render body vertices
+        draw_args = {"node_size": 100}
+        bodies = nx.subgraph(G, list(range(self.body_count + 1)))
+        nx.draw_networkx_nodes(bodies, pos, node_color="orange", node_shape="s", **draw_args)
+        if plot_shapes:
+            # render shape vertices
+            shapes = nx.subgraph(G, list(range(self.body_count + 1, len(vertices))))
+            nx.draw_networkx_nodes(shapes, pos, node_color="skyblue", **draw_args)
+            nx.draw_networkx_edges(
+                G, pos, node_size=0, edgelist=edges[self.joint_count :], edge_color="gray", style="dashed"
+            )
+        edge_labels = nx.get_edge_attributes(G, "label")
+        nx.draw_networkx_edge_labels(
+            G, pos, edge_labels=edge_labels, font_size=6, bbox={"alpha": 0.6, "color": "w", "lw": 0}
+        )
+        # add node labels
+        nx.draw_networkx_labels(G, pos, dict(enumerate(vertices)), font_size=6)
+        if show_legend:
+            plt.plot([], [], "s", color="orange", label="body")
+            plt.plot([], [], "k-", label="joint")
+            if plot_shapes:
+                plt.plot([], [], "o", color="skyblue", label="shape")
+                plt.plot([], [], "k--", label="shape-body connection")
+            plt.legend(loc="upper left", fontsize=6)
+        plt.show()
 
     def collapse_fixed_joints(self, verbose=wp.config.verbose):
         """Removes fixed joints from the model and merges the bodies they connect. This is useful for simplifying the model for faster and more stable simulation."""
@@ -2342,7 +2445,6 @@ class ModelBuilder:
                 "type": self.joint_type[i],
                 "q": self.joint_q[q_start : q_start + q_dim],
                 "qd": self.joint_qd[qd_start : qd_start + qd_dim],
-                "act": self.joint_act[qd_start : qd_start + qd_dim],
                 "armature": self.joint_armature[qd_start : qd_start + qd_dim],
                 "q_start": q_start,
                 "qd_start": qd_start,
@@ -2371,6 +2473,7 @@ class ModelBuilder:
                         "limit_kd": self.joint_limit_kd[j],
                         "limit_lower": self.joint_limit_lower[j],
                         "limit_upper": self.joint_limit_upper[j],
+                        "act": self.joint_act[j],
                     }
                 )
 
@@ -2525,7 +2628,6 @@ class ModelBuilder:
             self.joint_qd_start.append(len(self.joint_qd))
             self.joint_q.extend(joint["q"])
             self.joint_qd.extend(joint["qd"])
-            self.joint_act.extend(joint["act"])
             self.joint_armature.extend(joint["armature"])
             self.joint_enabled.append(joint["enabled"])
             self.joint_linear_compliance.append(joint["linear_compliance"])
@@ -2543,6 +2645,7 @@ class ModelBuilder:
                 self.joint_limit_upper.append(axis["limit_upper"])
                 self.joint_limit_ke.append(axis["limit_ke"])
                 self.joint_limit_kd.append(axis["limit_kd"])
+                self.joint_act.append(axis["act"])
 
     # muscles
     def add_muscle(
@@ -4466,7 +4569,9 @@ class ModelBuilder:
 
             # enable ground plane
             m.ground_plane = wp.array(self._ground_params["plane"], dtype=wp.float32, requires_grad=requires_grad)
-            m.gravity = np.array(self.up_vector) * self.gravity
+            m.gravity = np.array(self.up_vector, dtype=wp.float32) * self.gravity
+            m.up_axis = self.up_axis
+            m.up_vector = np.array(self.up_vector, dtype=wp.float32)
 
             m.enable_tri_collisions = False
 

warp/sim/render.py

@@ -67,13 +67,15 @@ def CreateSimRenderer(renderer):
             path,
             scaling=1.0,
             fps=60,
-            up_axis="Y",
+            up_axis=None,
             show_rigid_contact_points=False,
             contact_points_radius=1e-3,
             show_joints=False,
             **render_kwargs,
         ):
             # create USD stage
+            if up_axis is None:
+                up_axis = "XYZ"[model.up_axis]
             super().__init__(path, scaling=scaling, fps=fps, up_axis=up_axis, **render_kwargs)
             self.scaling = scaling
             self.cam_axis = "XYZ".index(up_axis.upper())