{ "cells": [ { "cell_type": "markdown", "id": "cell-title", "metadata": {}, "source": [ "# Simplicits with Franka Robot\n", "\n", "This notebook demonstrates a coupled robot–soft-body simulation. A Franka Emika Panda\n", "robot arm (driven by `SolverFeatherstone`) manipulates two deformable Simplicits cubes\n", "(driven by `SimplicitsSolver`) using Jacobian-based inverse kinematics. The robot follows\n", "a pre-scripted sequence of end-effector key poses while the soft bodies deform on contact." ] }, { "cell_type": "markdown", "id": "cell-imports-md", "metadata": {}, "source": [ "## Imports\n", "\n", "We import Kaolin, Newton, Warp, and **k3d** for interactive 3D visualization.\n", "`SolverFeatherstone` drives articulated robot dynamics; `eval_fk` computes\n", "forward kinematics for Jacobian computation. `newton.utils.download_asset`\n", "fetches the Franka URDF." ] }, { "cell_type": "code", "execution_count": null, "id": "4c324686-be9a-491f-b6b2-1ac65de2d18f", "metadata": {}, "outputs": [], "source": [ "!pip install -q newton[examples]==1.2.0" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-imports", "metadata": {}, "outputs": [], "source": [ "from types import SimpleNamespace\n", "import os\n", "import sys\n", "import threading\n", "from pathlib import Path\n", "import numpy as np\n", "import torch\n", "import warp as wp\n", "from IPython.display import display\n", "from ipywidgets import Button, VBox\n", "\n", "import newton\n", "import newton.utils\n", "from newton import Model, ModelBuilder, State, eval_fk\n", "from newton.solvers import SolverFeatherstone\n", "from newton.math import transform_twist\n", "\n", "import kaolin\n", "from kaolin.physics.simplicits import SimplicitsObject\n", "from kaolin.experimental.newton.builder import SimplicitsModelBuilder\n", "from kaolin.experimental.newton.solver import SimplicitsSolver\n", "\n", "sys.path.append(str(Path(\"..\")))\n", "from tutorial_common import COMMON_DATA_DIR\n", "\n", "wp.clear_kernel_cache()\n", "\n", "try:\n", " import k3d\n", "except ModuleNotFoundError:\n", " !pip install -q k3d\n", " import k3d\n", " raise RuntimeError(\"k3d has been installed. Do a hard refresh of the page (Ctrl + Shift + R) to avoid issue with k3d\")" ] }, { "cell_type": "markdown", "id": "cell-config-md", "metadata": {}, "source": [ "## Configuration\n", "\n", "Key constants:\n", "- **`SOFT_YOUNGS_MODULUS` / `POISSON_RATIO` / `DENSITY`** — elastic material parameters.\n", "- **`FRAME_DT`** — simulation timestep (1/60 s per frame, 1 substep).\n", "- **`COLLISION_PARTICLE_RADIUS`** — radius for Simplicits particle collision detection.\n", "- **`NULL_SPACE_GAIN`** — gain for null-space posture control.\n", "- **`GRIPPER_CLAMP_CLOSE/OPEN`** — gripper activation values in the key pose table." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-config", "metadata": {}, "outputs": [], "source": [ "# Simulation parameters\n", "FRAME_DT = 1.0 / 60.0\n", "FLOOR_PLANE = 0.0\n", "\n", "# Simplicits parameters\n", "SOFT_YOUNGS_MODULUS = 1e5\n", "POISSON_RATIO = 0.45\n", "DENSITY = 500\n", "APPROX_VOLUME = 0.5\n", "NUM_SAMPLES = 10000\n", "\n", "# Contact parameters\n", "SOFT_CONTACT_RADIUS = 0.01\n", "SOFT_CONTACT_MARGIN = 0.05\n", "SOFT_CONTACT_KE = 10000\n", "SOFT_CONTACT_KD = 2e-3\n", "ROBOT_FRICTION = 1.0\n", "SELF_CONTACT_FRICTION = 0.25\n", "SOFT_CONTACT_MAX = 1000000\n", "\n", "# Collision parameters\n", "COLLISION_PARTICLE_RADIUS = 0.1\n", "DETECTION_RATIO = 1.5\n", "IMPENETRABLE_BARRIER_RATIO = 0.25\n", "COLLISION_PENALTY = 1000.0\n", "MAX_CONTACT_PAIRS = 10000\n", "COLLISION_FRICTION = 0.5\n", "\n", "# Robot control parameters\n", "NULL_SPACE_GAIN = 1.0\n", "GRIPPER_CLAMP_CLOSE = 0.06\n", "GRIPPER_CLAMP_OPEN = 0.8" ] }, { "cell_type": "markdown", "id": "cell-kernels-md", "metadata": {}, "source": [ "## Warp Kernel and Jacobian\n", "\n", "- **`compute_ee_delta`** — GPU kernel that computes the 6-DOF spatial error\n", " (position + rotation) between the current end-effector pose and the target.\n", "- **`compute_body_jacobian`** — pure-Python function using `wp.Tape` autodiff\n", " to compute the velocity Jacobian of the end effector w.r.t. joint velocities." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-kernels", "metadata": {}, "outputs": [], "source": [ "@wp.kernel\n", "def compute_ee_delta(\n", " body_q: wp.array(dtype=wp.transform),\n", " offset: wp.transform,\n", " body_id: int,\n", " bodies_per_env: int,\n", " target: wp.transform,\n", " ee_delta: wp.array(dtype=wp.spatial_vector),\n", "):\n", " \"\"\"Compute the pose difference between end effector and target.\"\"\"\n", " env_id = wp.tid()\n", " tf = body_q[bodies_per_env * env_id + body_id] * offset\n", " pos = wp.transform_get_translation(tf)\n", " pos_des = wp.transform_get_translation(target)\n", " pos_diff = pos_des - pos\n", " rot = wp.transform_get_rotation(tf)\n", " rot_des = wp.transform_get_rotation(target)\n", " ang_diff = rot_des * wp.quat_inverse(rot)\n", " ee_delta[env_id] = wp.spatial_vector(\n", " pos_diff[0], pos_diff[1], pos_diff[2], ang_diff[0], ang_diff[1], ang_diff[2]\n", " )\n", "\n", "\n", "def compute_body_jacobian(\n", " model: Model,\n", " joint_q: wp.array,\n", " joint_qd: wp.array,\n", " body_id,\n", " offset=None,\n", " velocity: bool = True,\n", " include_rotation: bool = False,\n", "):\n", " \"\"\"Compute the Jacobian of a body with respect to joint coordinates via autodiff.\"\"\"\n", " if isinstance(body_id, str):\n", " body_id = model.body_name.get(body_id)\n", " if offset is None:\n", " offset = wp.transform_identity()\n", "\n", " joint_q.requires_grad = True\n", " joint_qd.requires_grad = True\n", "\n", " if velocity:\n", " @wp.kernel\n", " def compute_body_out(body_qd: wp.array(dtype=wp.spatial_vector), body_out: wp.array(dtype=float)):\n", " mv = transform_twist(offset, body_qd[body_id])\n", " if wp.static(include_rotation):\n", " for i in range(6):\n", " body_out[i] = mv[i]\n", " else:\n", " for i in range(3):\n", " body_out[i] = mv[3 + i]\n", " in_dim = model.joint_dof_count\n", " out_dim = 6 if include_rotation else 3\n", " else:\n", " @wp.kernel\n", " def compute_body_out(body_q: wp.array(dtype=wp.transform), body_out: wp.array(dtype=float)):\n", " tf = body_q[body_id] * offset\n", " if wp.static(include_rotation):\n", " for i in range(7):\n", " body_out[i] = tf[i]\n", " else:\n", " for i in range(3):\n", " body_out[i] = tf[i]\n", " in_dim = model.joint_coord_count\n", " out_dim = 7 if include_rotation else 3\n", "\n", " out_state = model.state(requires_grad=True)\n", " body_out = wp.empty(out_dim, dtype=float, requires_grad=True)\n", " tape = wp.Tape()\n", " with tape:\n", " eval_fk(model, joint_q, joint_qd, out_state)\n", " wp.launch(compute_body_out, 1, inputs=[\n", " out_state.body_qd if velocity else out_state.body_q], outputs=[body_out])\n", "\n", " def onehot(i):\n", " x = np.zeros(out_dim, dtype=np.float32)\n", " x[i] = 1.0\n", " return wp.array(x)\n", "\n", " J = np.empty((out_dim, in_dim), dtype=wp.float32)\n", " for i in range(out_dim):\n", " tape.backward(grads={body_out: onehot(i)})\n", " J[i] = joint_qd.grad.numpy() if velocity else joint_q.grad.numpy()\n", " tape.zero()\n", " return J.astype(np.float32)" ] }, { "cell_type": "markdown", "id": "cell-object-md", "metadata": {}, "source": [ "## Simplicits Cube Object\n", "\n", "`get_simplicits_cube()` loads a unit cube, scales by 0.25, samples interior points,\n", "and creates a single handle `SimplicitsObject`. Returns `(sim_obj, mesh)` so that `mesh`\n", "vertices and faces can be reused for k3d visualization." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-object", "metadata": {}, "outputs": [], "source": [ "def get_simplicits_cube():\n", " \"\"\"Load cube mesh and create a single handle Simplicits object. Returns (sim_obj, mesh).\"\"\"\n", " mesh_path = os.path.join(COMMON_DATA_DIR, \"meshes\")\n", " mesh = kaolin.io.import_mesh(\n", " mesh_path + \"/cube.obj\", triangulate=True\n", " ).to('cuda')\n", " mesh.vertices = kaolin.ops.pointcloud.center_points(\n", " mesh.vertices.unsqueeze(0), normalize=True\n", " ).squeeze(0) * 0.25\n", " orig_vertices = mesh.vertices.clone()\n", "\n", " uniform_pts = torch.rand(NUM_SAMPLES, 3, device='cuda') * (\n", " orig_vertices.max(dim=0).values - orig_vertices.min(dim=0).values\n", " ) + orig_vertices.min(dim=0).values\n", "\n", " physics_points = kaolin.physics.simplicits.PhysicsPoints(\n", " pts=uniform_pts,\n", " yms=SOFT_YOUNGS_MODULUS,\n", " prs=POISSON_RATIO,\n", " rhos=DENSITY,\n", " appx_vol=APPROX_VOLUME\n", " ).cuda()\n", "\n", " sim_obj = SimplicitsObject.create_rigid(\n", " physics_points=physics_points\n", " ) # Single handle objects acts as a affinely deformable at low yms. At higher stiffness it behaves rigidly.\n", " return sim_obj, mesh\n", "\n", "\n", "sim_obj, mesh = get_simplicits_cube()\n", "orig_vertices = mesh.vertices.clone()\n", "print(f\"Mesh: {mesh.vertices.shape[0]} vertices | Sim samples: {sim_obj.pts.shape[0]}\")" ] }, { "cell_type": "markdown", "id": "cell-articulation-md", "metadata": {}, "source": [ "## Robot Articulation\n", "\n", "`create_articulation(builder)` downloads the Franka Emika Panda URDF via\n", "`newton.utils.download_asset`, adds it to a `ModelBuilder`, sets the initial\n", "joint configuration, and defines the sequence of end-effector key poses used\n", "for the scripted manipulation task.\n", "\n", "Returns `(endeffector_id, endeffector_offset, robot_key_poses, targets,\n", "transition_duration, robot_key_poses_time)` for use by the IK controller." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-articulation", "metadata": {}, "outputs": [], "source": [ "def create_articulation(builder):\n", " \"\"\"Load Franka URDF, set initial pose, define key poses. Returns control metadata.\"\"\"\n", " asset_path = newton.utils.download_asset(\"franka_emika_panda\")\n", " builder.add_urdf(\n", " str(asset_path / \"urdf\" / \"fr3_franka_hand.urdf\"),\n", " xform=wp.transform(\n", " (-0.5, -0.5, -0.0),\n", " wp.quat_identity(),\n", " ),\n", " floating=False,\n", " scale=2,\n", " enable_self_collisions=False,\n", " collapse_fixed_joints=True,\n", " force_show_colliders=False,\n", " )\n", " builder.joint_q[:6] = [0.0, 0.0, 0.0, -1.59695, 0.0, 2.5307]\n", "\n", " robot_key_poses = np.array(\n", " [\n", " [2.5, 0.31, -0.60, 0.23, 1, 0.0, 0.0, 0.0, GRIPPER_CLAMP_OPEN],\n", " [2, 0.31, -0.60, 0.23, 1, 0.0, 0.0, 0.0, GRIPPER_CLAMP_CLOSE]\n", " ],\n", " dtype=np.float32,\n", " )\n", " targets = robot_key_poses[:, 1:]\n", " transition_duration = robot_key_poses[:, 0]\n", " robot_key_poses_time = np.cumsum(robot_key_poses[:, 0])\n", " endeffector_id = builder.body_count - 3\n", " endeffector_offset = wp.transform([0.0, 0.0, 0.22], wp.quat_identity())\n", " return endeffector_id, endeffector_offset, robot_key_poses, targets, transition_duration, robot_key_poses_time\n", "\n", "\n", "franka = ModelBuilder()\n", "endeffector_id, endeffector_offset, robot_key_poses, targets, transition_duration, robot_key_poses_time = \\\n", " create_articulation(franka)\n", "\n", "print(f\"Franka: {franka.body_count} bodies, {franka.joint_dof_count} DOFs, \"\n", " f\"endeffector_id={endeffector_id}\")" ] }, { "cell_type": "markdown", "id": "cell-builder-md", "metadata": {}, "source": [ "## Building the Combined Model\n", "\n", "`SimplicitsModelBuilder` holds both the soft-body Simplicits objects and the\n", "Franka robot articulation. After adding both, `finalize()` compiles all geometry\n", "and physics parameters into GPU arrays. `state_0` / `state_1` form the double\n", "buffer; `control` holds target joint inputs." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-builder", "metadata": {}, "outputs": [], "source": [ "simplicits_builder = SimplicitsModelBuilder()\n", "\n", "simplicits_builder.add_simplicits_object(\n", " sim_obj,\n", " num_qp=1000,\n", " init_transform=torch.tensor([\n", " [1, 0, 0, -0.5],\n", " [0, 1, 0, -0.5],\n", " [0, 0, 1, 1.5],\n", " [0, 0, 0, 1.0]\n", " ], dtype=torch.float32, device='cuda'),\n", " renderable_pts=orig_vertices.clone().detach()\n", ")\n", "simplicits_builder.add_simplicits_object(\n", " sim_obj,\n", " num_qp=1000,\n", " init_transform=torch.tensor([\n", " [1, 0, 0, -0.5],\n", " [0, 1, 0, -0.5],\n", " [0, 0, 1, 2.5],\n", " [0, 0, 0, 1.0]\n", " ], dtype=torch.float32, device='cuda'),\n", " renderable_pts=orig_vertices.clone().detach()\n", ")\n", "\n", "simplicits_builder.add_simplicits_collisions(\n", " collision_particle_radius=COLLISION_PARTICLE_RADIUS,\n", " detection_ratio=DETECTION_RATIO,\n", " impenetrable_barrier_ratio=IMPENETRABLE_BARRIER_RATIO,\n", " collision_penalty=COLLISION_PENALTY,\n", " max_contact_pairs=MAX_CONTACT_PAIRS,\n", " friction=COLLISION_FRICTION\n", ")\n", "\n", "# Add Franka robot and ground plane\n", "simplicits_builder.add_builder(franka)\n", "bodies_per_env = franka.body_count\n", "dof_q_per_env = franka.joint_coord_count\n", "dof_qd_per_env = franka.joint_dof_count\n", "\n", "simplicits_builder.add_ground_plane()\n", "\n", "# Finalize model\n", "model = simplicits_builder.finalize()\n", "model.soft_contact_ke = SOFT_CONTACT_KE\n", "model.soft_contact_kd = SOFT_CONTACT_KD\n", "model.soft_contact_mu = SELF_CONTACT_FRICTION\n", "\n", "state_0 = model.state()\n", "state_1 = model.state()\n", "control = model.control()\n", "target_joint_qd = wp.empty_like(state_0.joint_qd)\n", "\n", "print(f\"Model: {model.particle_count} particles, {model.body_count} bodies\")" ] }, { "cell_type": "markdown", "id": "cell-control-md", "metadata": {}, "source": [ "## Control Setup and IK\n", "\n", "`set_up_control()` allocates the Jacobian computation buffers and defines the\n", "`compute_body_out_kernel` that extracts the end-effector velocity for autodiff.\n", "\n", "`generate_control_joint_qd(state_in, sim_time)` computes target joint velocities\n", "via Jacobian pseudo-inverse IK plus a null-space term to maintain a reference posture.\n", "Gripper joints are controlled directly from the current key pose activation value." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-control", "metadata": {}, "outputs": [], "source": [ "def set_up_control():\n", " \"\"\"Allocate IK control buffers and define the end-effector velocity kernel.\"\"\"\n", " out_dim = 6\n", " in_dim = model.joint_dof_count\n", "\n", " def _onehot(i, n):\n", " return wp.array([1.0 if j == i else 0.0 for j in range(n)], dtype=float)\n", "\n", " Jacobian_one_hots = [_onehot(i, out_dim) for i in range(out_dim)]\n", "\n", " @wp.kernel\n", " def _compute_body_out(\n", " body_qd: wp.array(dtype=wp.spatial_vector),\n", " body_out_arr: wp.array(dtype=float)\n", " ):\n", " mv = transform_twist(wp.static(endeffector_offset), body_qd[wp.static(endeffector_id)])\n", " for i in range(6):\n", " body_out_arr[i] = mv[i]\n", "\n", " compute_body_out_kernel = _compute_body_out\n", " temp_state_for_jacobian = model.state(requires_grad=True)\n", " body_out = wp.empty(out_dim, dtype=float, requires_grad=True)\n", " J_flat = wp.empty(out_dim * in_dim, dtype=float)\n", " ee_delta = wp.empty(1, dtype=wp.spatial_vector)\n", " initial_pose = model.joint_q.numpy()\n", "\n", " return SimpleNamespace(\n", " one_hots=Jacobian_one_hots,\n", " kernel=compute_body_out_kernel,\n", " temp_state=temp_state_for_jacobian,\n", " body_out=body_out,\n", " J_flat=J_flat,\n", " ee_delta=ee_delta,\n", " initial_pose=initial_pose,\n", " )\n", "\n", "\n", "ik = set_up_control()\n", "\n", "\n", "def generate_control_joint_qd(state_in, sim_time: float):\n", " \"\"\"Compute target joint velocities via Jacobian IK + null-space posture control.\"\"\"\n", " t_mod = (\n", " sim_time\n", " if sim_time < robot_key_poses_time[-1]\n", " else sim_time % robot_key_poses_time[-1]\n", " )\n", " include_rotation = True\n", " current_interval = np.searchsorted(robot_key_poses_time, t_mod)\n", " target = targets[current_interval]\n", "\n", " wp.launch(\n", " compute_ee_delta,\n", " dim=1,\n", " inputs=[\n", " state_in.body_q,\n", " endeffector_offset,\n", " endeffector_id,\n", " bodies_per_env,\n", " wp.transform(*target[:7]),\n", " ],\n", " outputs=[ik.ee_delta],\n", " )\n", "\n", " # Compute Jacobian via autodiff\n", " state_in.joint_q.requires_grad = True\n", " state_in.joint_qd.requires_grad = True\n", " in_dim = model.joint_dof_count\n", " tape = wp.Tape()\n", " with tape:\n", " eval_fk(model, state_in.joint_q, state_in.joint_qd, ik.temp_state)\n", " wp.launch(\n", " ik.kernel, 1,\n", " inputs=[ik.temp_state.body_qd],\n", " outputs=[ik.body_out]\n", " )\n", " out_dim = 6\n", " for i in range(out_dim):\n", " tape.backward(grads={ik.body_out: ik.one_hots[i]})\n", " wp.copy(ik.J_flat[i * in_dim: (i + 1) * in_dim], state_in.joint_qd.grad)\n", " tape.zero()\n", "\n", " J = ik.J_flat.numpy().reshape(-1, in_dim)\n", " delta_target = ik.ee_delta.numpy()[0]\n", " J_inv = np.linalg.pinv(J)\n", "\n", " I = np.eye(J.shape[1], dtype=np.float32)\n", " N = I - J_inv @ J\n", "\n", " q = state_in.joint_q.numpy()\n", " q_des = q.copy()\n", " q_des[1:] = ik.initial_pose[1:]\n", " delta_q_null = NULL_SPACE_GAIN * (q_des - q)\n", "\n", " delta_q = J_inv @ delta_target + N @ delta_q_null\n", "\n", " # Gripper control\n", " delta_q[-2] = target[-1] * 0.04 - q[-2]\n", " delta_q[-1] = target[-1] * 0.04 - q[-1]\n", "\n", " target_joint_qd.assign(delta_q)" ] }, { "cell_type": "markdown", "id": "cell-solver-md", "metadata": {}, "source": [ "## Solvers\n", "\n", "Two solvers share the same model:\n", "\n", "- **`robot_solver`** (`SolverFeatherstone`) — articulated dynamics for the Franka arm.\n", " `model.particle_count` is temporarily set to 0 during its step so it skips soft-body\n", " particles. Gravity is set to zero during robot stepping (robot uses PD/IK control).\n", "- **`simplicits_solver`** (`SimplicitsSolver`) — Newton optimization for the soft cubes.\n", "\n", "Pre-allocated `gravity_zero` / `gravity_earth` arrays avoid CPU allocation in the loop." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-solver", "metadata": {}, "outputs": [], "source": [ "robot_solver = SolverFeatherstone(model, update_mass_matrix_interval=1)\n", "simplicits_solver = SimplicitsSolver(model)\n", "\n", "gravity_zero = wp.zeros(1, dtype=wp.vec3)\n", "gravity_earth = wp.array(wp.vec3(0.0, 0.0, -9.81), dtype=wp.vec3)\n", "\n", "# Evaluate initial forward kinematics\n", "newton.eval_fk(model, model.joint_q, model.joint_qd, state_0)\n", "\n", "sim_substeps = 1\n", "frame_dt = FRAME_DT\n", "sim_dt = frame_dt / sim_substeps\n", "\n", "print(\"Solvers ready\")" ] }, { "cell_type": "markdown", "id": "cell-simloop-md", "metadata": {}, "source": [ "## Simulation Loop\n", "\n", "Each call to `simulate()` first computes target joint velocities via IK, then\n", "runs `sim_substeps` physics steps:\n", "\n", "1. **Robot substep** — `model.particle_count=0` and `gravity=0` so the Featherstone\n", " solver only integrates the articulated chain. `shape_contact_pair_count=0` skips\n", " rigid–rigid contacts during this pass.\n", "2. **Simplicits substep** — collide + Newton solve for the two soft cubes against all\n", " shapes (including robot bodies).\n", "\n", "States are swapped after each substep." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-simloop", "metadata": {}, "outputs": [], "source": [ "sim_time = 0.0\n", "\n", "\n", "def simulate_step():\n", " global state_0, state_1\n", "\n", " state_0.clear_forces()\n", " state_1.clear_forces()\n", "\n", " # Robot step: disable particles and gravity\n", " particle_count = model.particle_count\n", " model.particle_count = 0\n", " model.gravity.assign(gravity_zero)\n", " model.shape_contact_pair_count = 0\n", "\n", " state_0.joint_qd.assign(target_joint_qd)\n", " robot_solver.step(state_0, state_1, control, None, sim_dt)\n", "\n", " # Restore and run Simplicits\n", " model.particle_count = particle_count\n", " model.gravity.assign(gravity_earth)\n", "\n", " contacts = model.collide(state_0)\n", " simplicits_solver.step(state_0, state_1, None, contacts, sim_dt)\n", "\n", " state_0, state_1 = state_1, state_0\n", "\n", "\n", "def simulate():\n", " global sim_time\n", " generate_control_joint_qd(state_0, sim_time)\n", " for _ in range(sim_substeps):\n", " simulate_step()\n", " sim_time += frame_dt" ] }, { "cell_type": "markdown", "id": "cell-viewer-md", "metadata": {}, "source": [ "## Visualization with k3d\n", "\n", "Two sets of k3d objects are rendered:\n", "\n", "- **`mesh_soft`** (blue) — the two Simplicits cubes combined into a single mesh.\n", " Updated each frame via `get_object_deformed_pts`.\n", "- **`mesh_robot`** (orange) — one combined mesh for the entire robot, built from\n", " `model.shape_source` after `finalize()`. `get_robot_mesh_verts` transforms each\n", " solid MESH shape from body frame to world frame using `body_q` each frame.\n", " Faces are static and cached once; only vertices are updated per frame.\n", "\n", "Play/Stop/Reset buttons control the background simulation thread." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-viewer", "metadata": {}, "outputs": [], "source": [ "orig_vertices_np = orig_vertices.cpu().numpy().astype(np.float32)\n", "orig_faces_np = mesh.faces.cpu().numpy().astype(np.uint32)\n", "num_verts = orig_vertices_np.shape[0]\n", "\n", "combined_faces_np = np.concatenate(\n", " [orig_faces_np, orig_faces_np + num_verts]\n", ").astype(np.uint32)\n", "\n", "# --- Robot mesh helpers ---\n", "\n", "\n", "def apply_transform(verts, tf_row):\n", " \"\"\"Apply a (tx, ty, tz, qx, qy, qz, qw) transform to an (N, 3) vertex array.\"\"\"\n", " trans = tf_row[:3]\n", " q = torch.from_numpy(tf_row[3:7].astype(np.float32)).unsqueeze(0)\n", " rot = kaolin.math.quat.rot33_from_quat(q)[0].numpy()\n", " return verts @ rot.T + trans\n", "\n", "\n", "def get_robot_mesh_verts(body_q_np):\n", " \"\"\"Collect all solid MESH shapes from model, transform to world, return combined verts+faces.\"\"\"\n", " shape_body = model.shape_body.numpy()\n", " shape_geo_src = model.shape_source # Python list, not a Warp array\n", " shape_geo_type = model.shape_type.numpy()\n", " shape_geo_scale = model.shape_scale.numpy()\n", " shape_transform = model.shape_transform.numpy()\n", " shape_is_solid = model.shape_is_solid.numpy()\n", "\n", " all_verts, all_faces, verts_offset = [], [], 0\n", "\n", " for i in range(len(shape_geo_src)):\n", " if shape_geo_type[i] != newton.GeoType.MESH or not shape_is_solid[i]:\n", " continue\n", " v = shape_geo_src[i].vertices.copy().astype(np.float32)\n", " f = shape_geo_src[i].indices.reshape(-1, 3).astype(np.uint32)\n", "\n", " q = torch.from_numpy(shape_transform[i][3:].astype(np.float32)).unsqueeze(0)\n", " R = kaolin.math.quat.rot33_from_quat(q)[0].numpy()\n", " p = shape_transform[i][:3]\n", " v = v * shape_geo_scale[i].reshape(1, -1) # per-axis scale\n", " v = (R @ v.T).T + p # shape-local → body frame\n", "\n", " if shape_body[i] >= 0:\n", " v = apply_transform(v, body_q_np[shape_body[i]]) # body → world\n", "\n", " all_verts.append(v)\n", " all_faces.append(f + verts_offset)\n", " verts_offset += len(v)\n", "\n", " return np.concatenate(all_verts, axis=0), np.concatenate(all_faces, axis=0)\n", "\n", "\n", "# --- Build k3d plot ---\n", "\n", "plot = k3d.plot(camera_auto_fit=False)\n", "plot.camera = [3, 3, 3, 0, 0, 0, 0, 0, 1]\n", "\n", "mesh_soft = k3d.mesh(\n", " np.concatenate([orig_vertices_np, orig_vertices_np]).astype(np.float32),\n", " combined_faces_np,\n", " color=0x3399ff\n", ")\n", "plot += mesh_soft\n", "\n", "# Robot: compute initial verts+faces, cache faces (static)\n", "init_robot_verts, robot_mesh_faces = get_robot_mesh_verts(state_0.body_q.numpy())\n", "mesh_robot = k3d.mesh(init_robot_verts, robot_mesh_faces, color=0xff6633)\n", "plot += mesh_robot\n", "\n", "# Floor plane at FLOOR_PLANE height\n", "_s = 3.0 # half-extent\n", "floor_verts = np.array([\n", " [-_s, -_s, FLOOR_PLANE],\n", " [ _s, -_s, FLOOR_PLANE],\n", " [ _s, _s, FLOOR_PLANE],\n", " [-_s, _s, FLOOR_PLANE],\n", "], dtype=np.float32)\n", "floor_faces = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.uint32)\n", "mesh_floor = k3d.mesh(floor_verts, floor_faces, color=0xcccccc, opacity=0.5)\n", "plot += mesh_floor\n", "\n", "plot.display()\n", "\n", "\n", "def update_vis():\n", " # Soft body cubes\n", " model.simplicits_scene.sim_z = state_0.sim_z\n", " v0 = model.simplicits_scene.get_object_deformed_pts(0, 'rendered')\n", " v1 = model.simplicits_scene.get_object_deformed_pts(1, 'rendered')\n", " mesh_soft.vertices = np.concatenate(\n", " [v0.cpu().numpy(), v1.cpu().numpy()]\n", " ).astype(np.float32)\n", " # Robot mesh (faces are static, only update vertices)\n", " robot_verts, _ = get_robot_mesh_verts(state_0.body_q.numpy())\n", " mesh_robot.vertices = robot_verts\n", "\n", "\n", "sim_running = [False]\n", "sim_thread = [None]\n", "\n", "\n", "def run_sim_loop():\n", " while sim_running[0]:\n", " simulate()\n", " update_vis()\n", "\n", "\n", "def on_play(b):\n", " if not sim_running[0]:\n", " sim_running[0] = True\n", " sim_thread[0] = threading.Thread(target=run_sim_loop, daemon=True)\n", " sim_thread[0].start()\n", "\n", "\n", "def on_stop(b):\n", " sim_running[0] = False\n", "\n", "\n", "def on_reset(b):\n", " global state_0, state_1, sim_time\n", " sim_running[0] = False\n", " if sim_thread[0] is not None:\n", " sim_thread[0].join()\n", " model.simplicits_scene.reset_scene()\n", " state_0 = model.state()\n", " state_1 = model.state()\n", " sim_time = 0.0\n", " update_vis()\n", "\n", "\n", "buttons = [Button(description=x) for x in [\"Play\", \"Stop\", \"Reset\"]]\n", "buttons[0].on_click(on_play)\n", "buttons[1].on_click(on_stop)\n", "buttons[2].on_click(on_reset)\n", "\n", "update_vis()\n", "display(VBox(buttons))" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }