# Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. # NVIDIA CORPORATION and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. import warp as wp import warp.sim import warp.render from collections import defaultdict import numpy as np from warp.render.utils import solidify_mesh, tab10_color_map # TODO allow NaNs in Warp kernels NAN = wp.constant(-1.0e8) @wp.kernel def compute_contact_points( body_q: wp.array(dtype=wp.transform), shape_body: wp.array(dtype=int), contact_shape0: wp.array(dtype=int), contact_shape1: wp.array(dtype=int), contact_point0: wp.array(dtype=wp.vec3), contact_point1: wp.array(dtype=wp.vec3), # outputs contact_pos0: wp.array(dtype=wp.vec3), contact_pos1: wp.array(dtype=wp.vec3), ): tid = wp.tid() shape_a = contact_shape0[tid] shape_b = contact_shape1[tid] if shape_a == shape_b: contact_pos0[tid] = wp.vec3(NAN, NAN, NAN) contact_pos1[tid] = wp.vec3(NAN, NAN, NAN) return body_a = shape_body[shape_a] body_b = shape_body[shape_b] X_wb_a = wp.transform_identity() X_wb_b = wp.transform_identity() if body_a >= 0: X_wb_a = body_q[body_a] if body_b >= 0: X_wb_b = body_q[body_b] contact_pos0[tid] = wp.transform_point(X_wb_a, contact_point0[tid]) contact_pos1[tid] = wp.transform_point(X_wb_b, contact_point1[tid]) def CreateSimRenderer(renderer): class SimRenderer(renderer): use_unique_colors = True def __init__( self, model: warp.sim.Model, path, scaling=1.0, fps=60, up_axis="Y", show_rigid_contact_points=False, contact_points_radius=1e-3, show_joints=False, **render_kwargs, ): # create USD stage super().__init__(path, scaling=scaling, fps=fps, up_axis=up_axis, **render_kwargs) self.scaling = scaling self.cam_axis = "XYZ".index(up_axis.upper()) self.show_rigid_contact_points = show_rigid_contact_points self.show_joints = show_joints self.contact_points_radius = contact_points_radius self.populate(model) def populate(self, model: warp.sim.Model): self.skip_rendering = False self.model = model self.num_envs = model.num_envs self.body_names = [] if self.show_rigid_contact_points and model.rigid_contact_max: self.contact_points0 = wp.array( np.zeros((model.rigid_contact_max, 3)), dtype=wp.vec3, device=model.device ) self.contact_points1 = wp.array( np.zeros((model.rigid_contact_max, 3)), dtype=wp.vec3, device=model.device ) self.contact_points0_colors = [(1.0, 0.5, 0.0)] * model.rigid_contact_max self.contact_points1_colors = [(0.0, 0.5, 1.0)] * model.rigid_contact_max self.body_env = [] # mapping from body index to its environment index env_id = 0 self.bodies_per_env = model.body_count // self.num_envs # create rigid body nodes for b in range(model.body_count): body_name = f"body_{b}_{self.model.body_name[b].replace(' ', '_')}" self.body_names.append(body_name) self.register_body(body_name) if b > 0 and b % self.bodies_per_env == 0: env_id += 1 self.body_env.append(env_id) # create rigid shape children if self.model.shape_count: # mapping from hash of geometry to shape ID self.geo_shape = {} self.instance_count = 0 self.body_name = {} # mapping from body name to its body ID self.body_shapes = defaultdict(list) # mapping from body index to its shape IDs shape_body = model.shape_body.numpy() shape_geo_src = model.shape_geo_src shape_geo_type = model.shape_geo.type.numpy() shape_geo_scale = model.shape_geo.scale.numpy() shape_geo_thickness = model.shape_geo.thickness.numpy() shape_geo_is_solid = model.shape_geo.is_solid.numpy() shape_transform = model.shape_transform.numpy() p = np.zeros(3, dtype=np.float32) q = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) scale = np.ones(3) color = np.ones(3) # loop over shapes excluding the ground plane for s in range(model.shape_count - 1): geo_type = shape_geo_type[s] geo_scale = [float(v) for v in shape_geo_scale[s]] geo_thickness = float(shape_geo_thickness[s]) geo_is_solid = bool(shape_geo_is_solid[s]) geo_src = shape_geo_src[s] if self.use_unique_colors: color = self._get_new_color() name = f"shape_{s}" # shape transform in body frame body = int(shape_body[s]) if body >= 0 and body < len(self.body_names): body = self.body_names[body] else: body = None # shape transform in body frame X_bs = wp.transform_expand(shape_transform[s]) # check whether we can instance an already created shape with the same geometry geo_hash = hash((int(geo_type), geo_src, *geo_scale, geo_thickness, geo_is_solid)) if geo_hash in self.geo_shape: shape = self.geo_shape[geo_hash] else: if geo_type == warp.sim.GEO_PLANE: if s == model.shape_count - 1 and not model.ground: continue # hide ground plane # plane mesh width = geo_scale[0] if geo_scale[0] > 0.0 else 100.0 length = geo_scale[1] if geo_scale[1] > 0.0 else 100.0 shape = self.render_plane( name, p, q, width, length, color, parent_body=body, is_template=True ) elif geo_type == warp.sim.GEO_SPHERE: shape = self.render_sphere(name, p, q, geo_scale[0], parent_body=body, is_template=True) elif geo_type == warp.sim.GEO_CAPSULE: shape = self.render_capsule( name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True ) elif geo_type == warp.sim.GEO_CYLINDER: shape = self.render_cylinder( name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True ) elif geo_type == warp.sim.GEO_CONE: shape = self.render_cone( name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True ) elif geo_type == warp.sim.GEO_BOX: shape = self.render_box(name, p, q, geo_scale, parent_body=body, is_template=True) elif geo_type == warp.sim.GEO_MESH: if not geo_is_solid: faces, vertices = solidify_mesh(geo_src.indices, geo_src.vertices, geo_thickness) else: faces, vertices = geo_src.indices, geo_src.vertices shape = self.render_mesh( name, vertices, faces, pos=p, rot=q, scale=geo_scale, colors=[color], parent_body=body, is_template=True, ) elif geo_type == warp.sim.GEO_SDF: continue self.geo_shape[geo_hash] = shape self.add_shape_instance(name, shape, body, X_bs.p, X_bs.q, scale) self.instance_count += 1 if self.show_joints and model.joint_count: joint_type = model.joint_type.numpy() joint_axis = model.joint_axis.numpy() joint_axis_start = model.joint_axis_start.numpy() joint_axis_dim = model.joint_axis_dim.numpy() joint_parent = model.joint_parent.numpy() joint_child = model.joint_child.numpy() joint_tf = model.joint_X_p.numpy() shape_collision_radius = model.shape_collision_radius.numpy() y_axis = wp.vec3(0., 1., 0.) color = (1., 0., 1.) shape = self.render_arrow( "joint_arrow", None, None, base_radius=0.01, base_height=0.4, cap_radius=0.02, cap_height=0.1, parent_body=None, is_template=True, color=color ) for i, t in enumerate(joint_type): if t not in { warp.sim.JOINT_REVOLUTE, # warp.sim.JOINT_PRISMATIC, warp.sim.JOINT_UNIVERSAL, warp.sim.JOINT_COMPOUND, warp.sim.JOINT_D6, }: continue tf = joint_tf[i] body = int(joint_parent[i]) # if body == -1: # continue num_linear_axes = int(joint_axis_dim[i][0]) num_angular_axes = int(joint_axis_dim[i][1]) # find a good scale for the arrow based on the average radius # of the shapes attached to the joint child body scale = np.ones(3) child = int(joint_child[i]) if child >= 0: radii = [] for s in model.body_shapes[child]: radii.append(shape_collision_radius[s]) if len(radii) > 0: scale *= np.mean(radii) * 2.0 for a in range(num_linear_axes, num_linear_axes + num_angular_axes): index = joint_axis_start[i] + a axis = joint_axis[index] if np.linalg.norm(axis) < 1e-6: continue p = wp.vec3(tf[:3]) q = wp.quat(tf[3:]) # compute rotation between axis and y axis = axis / np.linalg.norm(axis) q = q * wp.quat_between_vectors(wp.vec3(axis), y_axis) name = f"joint_{i}_{a}" self.add_shape_instance(name, shape, body, p, q, scale, color1=color, color2=color) self.instance_count += 1 if model.ground: self.render_ground() if hasattr(self, "complete_setup"): self.complete_setup() def _get_new_color(self): return tab10_color_map(self.instance_count) def render(self, state: warp.sim.State): if self.skip_rendering: return if self.model.particle_count: particle_q = state.particle_q.numpy() # render particles self.render_points("particles", particle_q, radius=self.model.particle_radius.numpy()) # render tris if self.model.tri_count: self.render_mesh("surface", particle_q, self.model.tri_indices.numpy().flatten()) # render springs if self.model.spring_count: self.render_line_list("springs", particle_q, self.model.spring_indices.numpy().flatten(), [], 0.05) # render muscles if self.model.muscle_count: body_q = state.body_q.numpy() muscle_start = self.model.muscle_start.numpy() muscle_links = self.model.muscle_bodies.numpy() muscle_points = self.model.muscle_points.numpy() muscle_activation = self.model.muscle_activation.numpy() # for s in self.skeletons: # # for mesh, link in s.mesh_map.items(): # # if link != -1: # # X_sc = wp.transform_expand(self.state.body_X_sc[link].tolist()) # # #self.renderer.add_mesh(mesh, "../assets/snu/OBJ/" + mesh + ".usd", X_sc, 1.0, self.render_time) # # self.renderer.add_mesh(mesh, "../assets/snu/OBJ/" + mesh + ".usd", X_sc, 1.0, self.render_time) for m in range(self.model.muscle_count): start = int(muscle_start[m]) end = int(muscle_start[m + 1]) points = [] for w in range(start, end): link = muscle_links[w] point = muscle_points[w] X_sc = wp.transform_expand(body_q[link][0]) points.append(wp.transform_point(X_sc, point).tolist()) self.render_line_strip( name=f"muscle_{m}", vertices=points, radius=0.0075, color=(muscle_activation[m], 0.2, 0.5) ) # update bodies if self.model.body_count: self.update_body_transforms(state.body_q) if self.show_rigid_contact_points and self.model.rigid_contact_max: wp.launch( kernel=compute_contact_points, dim=self.model.rigid_contact_max, inputs=[ state.body_q, self.model.shape_body, self.model.rigid_contact_shape0, self.model.rigid_contact_shape1, self.model.rigid_contact_point0, self.model.rigid_contact_point1, ], outputs=[ self.contact_points0, self.contact_points1, ], device=self.model.device, ) self.render_points( "contact_points0", self.contact_points0.numpy(), radius=self.contact_points_radius * self.scaling, colors=self.contact_points0_colors, ) self.render_points( "contact_points1", self.contact_points1.numpy(), radius=self.contact_points_radius * self.scaling, colors=self.contact_points1_colors, ) return SimRenderer SimRendererUsd = CreateSimRenderer(wp.render.UsdRenderer) SimRendererOpenGL = CreateSimRenderer(wp.render.OpenGLRenderer) SimRenderer = SimRendererUsd