diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..55732f2d672afa8c8abdb64b98a052609f003acf 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +*.STL filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index fe07dfbee37e75038db4a56da0b3f4c1558214f3..d3b38a2ed8be8aac0d79082561f7ca42d5835949 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,52 @@ --- -title: Ardy Motion Api -emoji: 🔥 -colorFrom: pink -colorTo: pink +title: ARDY Motion Playground +emoji: 🕺 +colorFrom: gray +colorTo: purple sdk: gradio sdk_version: 6.20.0 -python_version: '3.13' app_file: app.py +short_description: Text-to-motion with an interactive 3D skeleton playground +python_version: "3.12" +startup_duration_timeout: 1h pinned: false --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# ARDY Motion Playground + +Interactive **text-to-motion** demo for [ARDY](https://research.nvidia.com/labs/sil/projects/ardy/) +(*Autoregressive Diffusion with Hybrid Representation for Interactive Human Motion Generation*, NVIDIA, SIGGRAPH 2026). + +Pick a **rig** (human or robot), type a natural-language prompt, generate a 3D motion clip, +then **orbit, scrub, and play** it in the built-in Three.js playground (play/pause, timeline +scrubber, speed control, loop, and a root-trajectory trail). + +- **Human rig:** `nvidia/ARDY-Core-RP-20FPS-Horizon40` — 27-joint "Core" skeleton at 20 fps. +- **Robot rig:** `nvidia/ARDY-G1-RP-25FPS-Horizon52` — 34-joint Unitree G1 humanoid robot at 25 fps + (same rig as the sibling Space `hugging-apps/ardy-g1-motion-generation`). +- **Text encoder:** LLM2Vec-Llama-3-8B using the public + `unsloth/llama-3-8b-Instruct` base plus the original McGill MNTP and supervised + adapters. No Meta-gated base-model access is required. +- **Hardware:** ZeroGPU. + +Post-processing (foot-skate cleanup, which needs a compiled C++ extension) is disabled in this +demo; generation uses the model's autoregressive path directly. + +Based on the [nv-tlabs/ardy](https://github.com/nv-tlabs/ardy) reference implementation +(`scripts/generate.py`). + +## Blender export + +Open **Blender export** to generate three downloadable files: + +- `ardy_motion.bvh` — directly importable with Blender's BVH importer. +- `ardy_motion.npz` — complete ARDY output including joint positions, local and + global rotations, root trajectory, and foot contacts. +- `metadata.json` — prompt, seed, FPS, skeleton hierarchy, and coordinate-system + information. + +The programmatic endpoint is `/generate_blender`. The existing `/generate` +endpoint and Three.js playground remain available. + +The base and adapter repositories can be overridden with the Space variables +`LLM2VEC_BASE_MODEL`, `LLM2VEC_MNTP_MODEL`, and `LLM2VEC_SUPERVISED_MODEL`. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..b60e0af5eeb7a45b20c7b62ce529603084287589 --- /dev/null +++ b/app.py @@ -0,0 +1,1113 @@ +import os + +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +# Load models locally (no external text-encoder service) and prefer the HF cache. +os.environ.setdefault("TEXT_ENCODER_MODE", "local") + +import base64 +import gzip +import json +import random +import sys +import tempfile +import time +import xml.etree.ElementTree as ET +from pathlib import Path + +import spaces # must precede torch / CUDA-touching imports +import torch +import numpy as np +import trimesh +import gradio as gr + +# The vendored ARDY package lives next to this file. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from ardy.model import load_model # noqa: E402 +from ardy.exports import write_bvh # noqa: E402 +from ardy.model.load_model import load_text_encoder # noqa: E402 +from ardy.motion_rep.tools import length_to_mask # noqa: E402 +from ardy.tools import seed_everything, to_numpy # noqa: E402 + +# Two rigs are offered in the playground: +# "human" -> ARDY-Core-RP-20FPS-Horizon40, 27-joint skeleton @ 20 fps +# "robot" -> ARDY-G1-RP-25FPS-Horizon52, 34-joint Unitree G1 robot @ 25 fps +# The robot rig matches the sibling Space hugging-apps/ardy-g1-motion-generation. +DEFAULT_RIG = "human" +MAX_SEED = 2**31 - 1 + +# ----------------------------------------------------------------------------- +# Model loading (module scope). +# +# ZeroGPU has no live GPU at startup: it intercepts .to("cuda") *placement* but +# NOT arbitrary CUDA compute. LLM2Vec's PEFT load runs a LoRA `merge_and_unload` +# (real matmuls), so we build everything on CPU first, then .to("cuda") — the +# placement call is what the ZeroGPU hijack packs to disk and streams into VRAM +# on the first @spaces.GPU request. +# ----------------------------------------------------------------------------- +# transformers / PEFT / safetensors infer their load device from +# torch.cuda.is_available(), which ZeroGPU reports True at startup even though no +# real GPU is attached — so a plain load tries to place weights on cuda and dies +# with "No CUDA GPUs are available". Force every loader onto CPU by masking +# is_available() during construction, then restore it and .to("cuda") (which the +# ZeroGPU hijack packs to disk + streams into VRAM on the first request). +# +# A single text encoder is built once and shared across both motion models +# (load_text_encoder is explicitly designed for this — see its docstring). +_real_cuda_available = torch.cuda.is_available +torch.cuda.is_available = lambda: False +try: + print("Loading ARDY text encoder (LLM2Vec-Llama-3-8B) on CPU…", flush=True) + _text_encoder = load_text_encoder(mode="local", device="cpu") + + print("Loading ARDY human (core) motion model on CPU…", flush=True) + MODEL_HUMAN = load_model("core", device="cpu", text_encoder=_text_encoder) + + print("Loading ARDY robot (G1) motion model on CPU…", flush=True) + MODEL_ROBOT = load_model("g1", device="cpu", text_encoder=_text_encoder) +finally: + torch.cuda.is_available = _real_cuda_available + +print("Moving models to CUDA (ZeroGPU-intercepted placement)…", flush=True) +_text_encoder = _text_encoder.to("cuda") +for _m in (MODEL_HUMAN, MODEL_ROBOT): + _m.to("cuda") + # self.device was captured at construction (device="cpu"); generation creates + # many tensors on it, so retarget it to cuda to match the moved weights. + _m.device = "cuda" + _m.eval() + + +def _rig_params(model): + fps = float(model.motion_rep.fps) + skeleton = model.skeleton + parents = skeleton.joint_parents.cpu().numpy().astype(int).tolist() + patch = model.num_frames_per_token + gen_horizon = model.gen_horizon_len + num_base_steps = int(model.diffusion.num_base_steps) + # History carried between autoregressive windows. The reference streaming demo + # keeps this SHORT so a new prompt takes effect within a window. + hist_crop = max(patch, (4 // patch) * patch) + root_idx = parents.index(-1) if -1 in parents else 0 + return { + "model": model, + "fps": fps, + "skeleton": skeleton, + "parents": parents, + "patch": patch, + "gen_horizon": gen_horizon, + "num_base_steps": num_base_steps, + "hist_crop": hist_crop, + "root_idx": root_idx, + } + + +RIGS = { + "human": _rig_params(MODEL_HUMAN), + "robot": _rig_params(MODEL_ROBOT), +} +# The largest base-step count across rigs bounds the diffusion-steps slider. +NUM_BASE_STEPS = max(r["num_base_steps"] for r in RIGS.values()) +for _name, _r in RIGS.items(): + print( + f"[{_name}] rig ready: {_r['skeleton'].nbjoints} joints, {_r['fps']} fps, " + f"horizon {_r['gen_horizon']}, patch {_r['patch']}, " + f"hist_crop {_r['hist_crop']}, base_steps {_r['num_base_steps']}", + flush=True, + ) + + +def _normalize_rig(rig) -> str: + """Map any UI/API rig value onto a valid RIGS key ('human' | 'robot').""" + if rig is None: + return DEFAULT_RIG + key = str(rig).strip().lower() + if key in RIGS: + return key + if key.startswith("hum") or "core" in key or "person" in key: + return "human" + if key.startswith("rob") or "g1" in key or "unitree" in key: + return "robot" + return DEFAULT_RIG + + +# ----------------------------------------------------------------------------- +# HUMAN rig: skinned body mesh (ARDY "CoreSkin" linear-blend skinning). +# +# The reference viz (ardy/viz/viser_utils.py) renders a smooth humanoid body by +# skinning a bind mesh with the per-frame *global* joint transforms: +# verts = CoreSkin.skin(global_rot_mats, posed_joints, rot_is_global=True) +# The browser holds the static skin data (bind vertices / faces / LBS +# indices+weights) and does the per-vertex blend, while the server sends only the +# tiny per-frame joint affine matrices +# A[f,j] = fk[f,j] @ bind_rig_transform_inv[j] (fk = [R_global | pos]) +# so the payload stays small (~0.1 MB/clip). +# ----------------------------------------------------------------------------- +_HUMAN_SKEL = RIGS["human"]["skeleton"] +_SKIN_PATH = Path(_HUMAN_SKEL.folder) / "skin_standard.npz" +_skin = np.load(_SKIN_PATH) +BIND_RIG_INV = np.linalg.inv( + np.asarray(_skin["bind_rig_transform"], dtype=np.float64) +).astype(np.float32) # [J, 4, 4] + + +def _build_skin_blob(): + """Pack the static skin data into one gzip+base64 blob (loaded once by the + browser). Layout: bind_vertices f32[V,3] | faces u32[F,3] | lbs_idx u8[V,W] + | lbs_wt f32[V,W].""" + bind_v = np.asarray(_skin["bind_vertices"], dtype=np.float32) + faces = np.asarray(_skin["faces"], dtype=np.uint32) + idx = np.asarray(_skin["lbs_indices"], dtype=np.uint8) + wt = np.asarray(_skin["lbs_weights"], dtype=np.float32) + raw = ( + np.ascontiguousarray(bind_v).tobytes() + + np.ascontiguousarray(faces).tobytes() + + np.ascontiguousarray(idx).tobytes() + + np.ascontiguousarray(wt).tobytes() + ) + meta = {"V": int(bind_v.shape[0]), "F": int(faces.shape[0]), "W": int(idx.shape[1])} + return base64.b64encode(gzip.compress(raw, 6)).decode("ascii"), meta + + +SKIN_B64, SKIN_META = _build_skin_blob() +print(f"CoreSkin ready: {SKIN_META['V']} verts / {SKIN_META['F']} faces, " + f"blob {len(SKIN_B64) // 1024} KB", flush=True) + + +def _joint_affines_human(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str: + """Per-frame joint affine matrices A = fk @ bind_rig_inv, base64 f32 [T,J,12]. + + global_rot_mats: [T, J, 3, 3]; posed_joints: [T, J, 3].""" + T, J = posed_joints.shape[:2] + fk = np.tile(np.eye(4, dtype=np.float32), (T, J, 1, 1)) + fk[..., :3, :3] = global_rot_mats.astype(np.float32) + fk[..., :3, 3] = posed_joints.astype(np.float32) + A = (fk @ BIND_RIG_INV)[..., :3, :] # [T, J, 3, 4] + A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32)) + return base64.b64encode(A.tobytes()).decode("ascii") + + +# ----------------------------------------------------------------------------- +# ROBOT rig: G1 robot mesh rig (rigid per-joint STL meshes). +# +# Unlike the human skeleton (rendered with one skinned body mesh via LBS), the +# Unitree G1 robot is rendered by attaching a rigid STL mesh to each articulated +# joint. This mirrors ardy/viz/g1_rig.py (G1MeshRig): each mesh has a local +# transform (geom_pos, geom_rot) relative to its joint, read from the MuJoCo +# g1.xml, plus a coordinate change from MuJoCo to ARDY axes. We precompute — for +# each mesh — its geometry PRE-TRANSFORMED into the joint-local frame +# (v' = geom_rot @ v + geom_pos), so at render time the browser just applies the +# per-frame joint transform: world_v = joint_pos + joint_rot @ v'. +# ----------------------------------------------------------------------------- +# G1 joint -> STL mesh mapping (mirrors ardy/viz/g1_rig.py G1_MESH_JOINT_MAP). +G1_MESH_JOINT_MAP = { + "pelvis_skel": ["pelvis.STL", "pelvis_contour_link.STL"], + "left_hip_pitch_skel": ["left_hip_pitch_link.STL"], + "left_hip_roll_skel": ["left_hip_roll_link.STL"], + "left_hip_yaw_skel": ["left_hip_yaw_link.STL"], + "left_knee_skel": ["left_knee_link.STL"], + "left_ankle_pitch_skel": ["left_ankle_pitch_link.STL"], + "left_ankle_roll_skel": ["left_ankle_roll_link.STL"], + "right_hip_pitch_skel": ["right_hip_pitch_link.STL"], + "right_hip_roll_skel": ["right_hip_roll_link.STL"], + "right_hip_yaw_skel": ["right_hip_yaw_link.STL"], + "right_knee_skel": ["right_knee_link.STL"], + "right_ankle_pitch_skel": ["right_ankle_pitch_link.STL"], + "right_ankle_roll_skel": ["right_ankle_roll_link.STL"], + "waist_yaw_skel": ["waist_yaw_link_rev_1_0.STL", "waist_yaw_link.STL"], + "waist_roll_skel": ["waist_roll_link_rev_1_0.STL", "waist_roll_link.STL"], + "waist_pitch_skel": [ + "torso_link_rev_1_0.STL", + "torso_link.STL", + "logo_link.STL", + "head_link.STL", + ], + "left_shoulder_pitch_skel": ["left_shoulder_pitch_link.STL"], + "left_shoulder_roll_skel": ["left_shoulder_roll_link.STL"], + "left_shoulder_yaw_skel": ["left_shoulder_yaw_link.STL"], + "left_elbow_skel": ["left_elbow_link.STL"], + "left_wrist_roll_skel": ["left_wrist_roll_link.STL"], + "left_wrist_pitch_skel": ["left_wrist_pitch_link.STL"], + "left_wrist_yaw_skel": ["left_wrist_yaw_link.STL", "left_rubber_hand.STL"], + "right_shoulder_pitch_skel": ["right_shoulder_pitch_link.STL"], + "right_shoulder_roll_skel": ["right_shoulder_roll_link.STL"], + "right_shoulder_yaw_skel": ["right_shoulder_yaw_link.STL"], + "right_elbow_skel": ["right_elbow_link.STL"], + "right_wrist_roll_skel": ["right_wrist_roll_link.STL"], + "right_wrist_pitch_skel": ["right_wrist_pitch_link.STL"], + "right_wrist_yaw_skel": ["right_wrist_yaw_link.STL", "right_rubber_hand.STL"], +} + +_ROBOT_SKEL = RIGS["robot"]["skeleton"] +_MUJOCO_TO_ARDY = np.array( + [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64 +) +_G1_SKEL_DIR = Path(_ROBOT_SKEL.folder) +_G1_MESH_DIR = _G1_SKEL_DIR / "meshes" / "g1" +_G1_XML = _G1_SKEL_DIR / "xml" / "g1.xml" + + +def _quat_wxyz_to_matrix(wxyz: np.ndarray) -> np.ndarray: + w, x, y, z = wxyz + n = np.sqrt(w * w + x * x + y * y + z * z) + if n < 1e-12: + return np.eye(3) + w, x, y, z = w / n, x / n, y / n, z / n + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ], + dtype=np.float64, + ) + + +def _mesh_local_transforms() -> dict: + """mesh_file -> (geom_pos[3], geom_rot[3x3]) in ARDY axes, parsed from g1.xml.""" + if not _G1_XML.exists(): + return {} + root = ET.parse(_G1_XML).getroot() + file_to_name = {} + for mesh in root.findall(".//asset/mesh"): + name, file = mesh.get("name"), mesh.get("file") + if name and file: + file_to_name[file] = name + name_to_tf = {} + for geom in root.findall(".//geom"): + name = geom.get("mesh") + if name is None: + continue + pos = geom.get("pos") + quat = geom.get("quat") + gp = np.zeros(3) if pos is None else np.array([float(v) for v in pos.split()]) + gr_ = np.eye(3) if quat is None else _quat_wxyz_to_matrix( + np.array([float(v) for v in quat.split()]) + ) + name_to_tf[name] = (gp, gr_) + out = {} + for file, name in file_to_name.items(): + gp, gr_ = name_to_tf.get(name, (np.zeros(3), np.eye(3))) + gp = _MUJOCO_TO_ARDY @ gp + gr_ = _MUJOCO_TO_ARDY @ gr_ @ _MUJOCO_TO_ARDY.T + out[file] = (gp, gr_) + return out + + +def _build_g1_mesh_blob(): + """Pack all rigid G1 meshes, pre-transformed into joint-local frame, into one + gzip+base64 blob. Returns (b64, meta). meta.parts lists per-mesh + {joint, v_off, v_cnt}. Layout: all verts f32[Vtot,3] then all faces + u32[Ftot,3] (face indices are GLOBAL into the concatenated vertex array).""" + skeleton = _ROBOT_SKEL + local_tf = _mesh_local_transforms() + all_v = [] + all_f = [] + parts = [] + v_cursor = 0 + for joint_name, mesh_files in G1_MESH_JOINT_MAP.items(): + if joint_name not in skeleton.bone_index: + continue + joint_idx = int(skeleton.bone_index[joint_name]) + for mesh_file in mesh_files: + mp = _G1_MESH_DIR / mesh_file + if not mp.exists(): + continue + mesh = trimesh.load_mesh(str(mp), process=True) + if isinstance(mesh, trimesh.Scene): + mesh = trimesh.util.concatenate(mesh.dump()) + verts = np.asarray(mesh.vertices, dtype=np.float64) @ _MUJOCO_TO_ARDY.T + faces = np.asarray(mesh.faces, dtype=np.int64) + gp, gr_ = local_tf.get(mesh_file, (np.zeros(3), np.eye(3))) + # Pre-apply the mesh's joint-local transform: v' = geom_rot @ v + geom_pos. + verts = (verts @ gr_.T) + gp + vcnt = verts.shape[0] + parts.append({"joint": joint_idx, "v_off": v_cursor, "v_cnt": vcnt}) + all_v.append(verts.astype(np.float32)) + all_f.append((faces + v_cursor).astype(np.uint32)) # global vertex indices + v_cursor += vcnt + V = np.concatenate(all_v, axis=0) if all_v else np.zeros((0, 3), np.float32) + F = np.concatenate(all_f, axis=0) if all_f else np.zeros((0, 3), np.uint32) + raw = np.ascontiguousarray(V).tobytes() + np.ascontiguousarray(F).tobytes() + meta = {"Vtot": int(V.shape[0]), "Ftot": int(F.shape[0]), "parts": parts} + b64 = base64.b64encode(gzip.compress(raw, 6)).decode("ascii") + return b64, meta + + +G1_MESH_B64, G1_MESH_META = _build_g1_mesh_blob() +print( + f"G1 rig ready: {len(G1_MESH_META['parts'])} meshes, " + f"{G1_MESH_META['Vtot']} verts / {G1_MESH_META['Ftot']} faces, " + f"blob {len(G1_MESH_B64) // 1024} KB", + flush=True, +) + + +def _joint_transforms_robot(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str: + """Per-frame joint affine matrices [R_global | pos], base64 f32 [T,J,12]. + + The browser applies world_v = pos + R_global @ v' per mesh (v' already in + joint-local frame).""" + T, J = posed_joints.shape[:2] + A = np.zeros((T, J, 3, 4), dtype=np.float32) + A[..., :3, :3] = global_rot_mats.astype(np.float32) + A[..., :3, 3] = posed_joints.astype(np.float32) + A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32)) + return base64.b64encode(A.tobytes()).decode("ascii") + + +# --- Autoregressive generation with a persistable latent state --------------- +# ARDY is autoregressive: it generates one `gen_horizon_len`-frame window at a +# time, conditioned on a history of previous frames. `autoregressive_step` is +# the streaming primitive — it returns the *normalized motion-feature tensor* +# for (history + new window), which can be fed straight back in as the next +# window's history. We thread that tensor to (a) fill a requested clip length +# and (b) CONTINUE a clip with a new prompt, exactly like the reference +# interactive demo. Persisting the tensor in a gr.State lets a second "Continue" +# click resume from where the first clip ended, on the same character. +def _generate_sequence(rig_key, prompt, num_new_frames, steps, cfg_weight, init_tensor): + """Run the AR loop for one prompt on the selected rig. `init_tensor`: + normalized feature tensor [1, Th, D] on cuda (prior motion to continue), or + None to start fresh. Returns the full normalized feature tensor.""" + r = RIGS[rig_key] + model = r["model"] + gen_horizon = r["gen_horizon"] + patch = r["patch"] + hist_crop = r["hist_crop"] + text_feat, text_pad_mask = model._encode_text([prompt]) + motion_tensor = init_tensor + target_new = max(1, int(np.ceil(num_new_frames / gen_horizon))) * gen_horizon + produced = 0 + while produced < target_new: + if motion_tensor is None: + history, hist_len = None, 0 + else: + hist_len = (min(motion_tensor.shape[1], hist_crop) // patch) * patch + history = motion_tensor[:, motion_tensor.shape[1] - hist_len:] if hist_len else None + hist_len = history.shape[1] if history is not None else 0 + samples = model.autoregressive_step( + num_frames=hist_len + gen_horizon, # exactly history + one window (no future) + num_denoising_steps=steps, + motion_mask=None, + observed_motion=None, + cfg_weight=float(cfg_weight), + texts=None, + text_feat=text_feat, + text_pad_mask=text_pad_mask, + init_history_sequence=history, + init_global_translation=None, # first window -> defaults (origin / +Z heading) + init_first_heading_angle=None, + ) + new_tail = samples[:, hist_len:] # the freshly generated window + motion_tensor = new_tail if motion_tensor is None else torch.cat([motion_tensor, new_tail], dim=1) + produced += new_tail.shape[1] + return motion_tensor + + +def _decode_motion(rig_key, motion_tensor): + r = RIGS[rig_key] + model = r["model"] + with torch.no_grad(): + out = to_numpy(model.motion_rep.inverse(motion_tensor, is_normalized=True)) + return { + key: np.asarray(value)[0] + for key, value in out.items() + } + + +def _pack_payload(rig_key, motion_tensor, prompt, seed): + """Decode a normalized feature tensor to the browser payload (per-frame joint + affines + root ground-track), tagged with the rig so the viewer loads the + correct skeleton/model.""" + r = RIGS[rig_key] + out = _decode_motion(rig_key, motion_tensor) + posed = out["posed_joints"] # [T, J, 3] global joint positions + grm = out["global_rot_mats"] # [T, J, 3, 3] global rotations + if rig_key == "robot": + affines = _joint_transforms_robot(grm, posed) + else: + affines = _joint_affines_human(grm, posed) + return { + "rig": rig_key, + "fps": r["fps"], + "num_frames": int(posed.shape[0]), + "num_joints": int(posed.shape[1]), + "affines": affines, # drives the browser rig + "root": np.round(posed[:, r["root_idx"], :].astype(np.float32), 4).tolist(), + "prompt": prompt, + "seed": seed, + } + + +def _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, init_np): + rig_key = _normalize_rig(rig) + r = RIGS[rig_key] + prompt = (prompt or "").strip() + if not prompt: + raise gr.Error("Please enter a text prompt describing the motion.") + if randomize_seed: + seed = random.randint(0, MAX_SEED) + seed = int(seed) + seed_everything(seed) + steps = max(1, min(int(diffusion_steps), r["num_base_steps"])) + num_new = max(r["patch"], int(round(float(duration) * r["fps"]))) + init_tensor = None if init_np is None else torch.from_numpy(init_np).to("cuda") + + t0 = time.perf_counter() + with torch.no_grad(): + full = _generate_sequence(rig_key, prompt, num_new, steps, cfg_weight, init_tensor) + payload = _pack_payload(rig_key, full, prompt, seed) + tag = "continue" if init_np is not None else "generate" + print(f"[{tag}:{rig_key}] '{prompt[:50]}' +{num_new}f -> {full.shape[1]}f total " + f"steps={steps} seed={seed} {time.perf_counter() - t0:.1f}s", flush=True) + return json.dumps(payload), seed, full.detach().cpu().numpy() + + +@spaces.GPU +def ui_generate(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS, + cfg_weight=2.0, seed=0, randomize_seed=True): + """Start a fresh clip (resets the running sequence).""" + return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None) + + +@spaces.GPU +def ui_continue(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS, + cfg_weight=2.0, seed=0, randomize_seed=True, state=None): + """Append a new action, continuing from the previous clip's final pose.""" + return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, state) + + +@spaces.GPU +def generate_motion(prompt: str, rig: str = DEFAULT_RIG, duration: float = 5.0, + diffusion_steps: int = NUM_BASE_STEPS, cfg_weight: float = 2.0, + seed: int = 0, randomize_seed: bool = True) -> tuple[str, int]: + """Generate a 3D motion clip from a text prompt with ARDY. + + Args: + prompt: Natural-language description of the motion (e.g. "a person walks in a circle"). + rig: Which character to animate — "human" (27-joint skeleton) or "robot" (Unitree G1). + duration: Length of the generated motion in seconds. + diffusion_steps: Number of denoising steps (1..num_base_steps). + cfg_weight: Classifier-free-guidance weight for the text prompt. + seed: Random seed for reproducibility. + randomize_seed: If True, ignore `seed` and draw a fresh random one. + + Returns: + A JSON string with the animated skeleton payload plus the seed used. + """ + payload_json, seed, _ = _core_generate( + rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None + ) + return payload_json, seed + + +@spaces.GPU +def generate_blender_motion( + prompt: str, + duration: float = 5.0, + diffusion_steps: int = NUM_BASE_STEPS, + cfg_weight: float = 2.0, + seed: int = 0, + randomize_seed: bool = True, +) -> tuple[str, str, str, int]: + """Generate a human motion and return BVH, NPZ, and JSON metadata files.""" + _, seed, motion_np = _core_generate( + "human", + prompt, + duration, + diffusion_steps, + cfg_weight, + seed, + randomize_seed, + None, + ) + motion_tensor = torch.from_numpy(motion_np).to("cuda") + out = _decode_motion("human", motion_tensor) + r = RIGS["human"] + skeleton = r["skeleton"] + joint_names = list(skeleton.bone_order_names) + parents = list(r["parents"]) + rest_positions = to_numpy(skeleton.neutral_joints) + + output_dir = Path(tempfile.mkdtemp(prefix="ardy_blender_")) + npz_path = output_dir / "ardy_motion.npz" + bvh_path = output_dir / "ardy_motion.bvh" + metadata_path = output_dir / "metadata.json" + + np.savez_compressed( + npz_path, + **out, + fps=np.asarray(r["fps"], dtype=np.float32), + text=np.asarray(prompt), + seed=np.asarray(seed, dtype=np.int64), + rig=np.asarray("human"), + joint_names=np.asarray(joint_names), + parents=np.asarray(parents, dtype=np.int32), + rest_positions=np.asarray(rest_positions, dtype=np.float32), + ) + write_bvh( + bvh_path, + joint_names=joint_names, + parents=parents, + rest_positions=rest_positions, + local_rot_mats=out["local_rot_mats"], + root_positions=out["root_positions"], + fps=r["fps"], + ) + metadata_path.write_text( + json.dumps( + { + "model": "nvidia/ARDY-Core-RP-20FPS-Horizon40", + "rig": "human", + "prompt": prompt, + "seed": seed, + "fps": r["fps"], + "frames": int(out["posed_joints"].shape[0]), + "joints": joint_names, + "parents": parents, + "coordinate_system": "Y-up, metric", + "post_processing": False, + }, + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + return str(bvh_path), str(npz_path), str(metadata_path), seed + + +# ----------------------------------------------------------------------------- +# Front-end: a self-contained Three.js playground, delivered as a Gradio-native +# custom HTML component (Gradio 6 `gr.HTML` templates + js_on_load). +# +# The motion JSON is carried as the component's own `value` prop. `js_on_load` +# imports three.js, builds the scene once, wires the controls, then registers a +# `watch('value', ...)` callback that Gradio fires whenever the component is +# updated as the output of a Python event (Generate button / Examples). +# +# Each payload is tagged with its `rig`. The viewer ships the static data for +# BOTH rigs (human skin blob + G1 rigid-mesh blob) and switches at load time: +# - "human": one skinned body mesh (linear-blend skinning in the browser). +# - "robot": rigid per-joint G1 STL meshes posed by the joint transforms. +# ----------------------------------------------------------------------------- + +PLAYER_TEMPLATE = """ +
+
+
Generate a motion to load it into the playground.
+
+
+ + + 0 / 0 + + + +
+
+
+""" + +# css_template rules are auto-scoped to this component by Gradio. +PLAYER_CSS_TEMPLATE = """ +.ardy-playground { width: 100%; } +.ardy-canvas-wrap { + position: relative; width: 100%; height: 480px; + border-radius: 12px; overflow: hidden; + background: #ffffff; + border: 1px solid #e5e7eb; +} +.ardy-canvas-wrap canvas { display:block; width:100% !important; height:100% !important; } +.ardy-hint { + position:absolute; top:50%; left:50%; transform:translate(-50%,-50%); + color:#98a2b3; font-size:14px; text-align:center; pointer-events:none; +} +.ardy-controls { + display:flex; align-items:center; gap:12px; flex-wrap:wrap; + margin-top:10px; padding:8px 4px; +} +.ardy-controls .ardy-btn { + background:#76B900; color:#fff; + border:none; border-radius:8px; padding:6px 16px; cursor:pointer; font-weight:600; +} +.ardy-scrub { flex:1; min-width:160px; accent-color:#76B900; } +.ardy-frame { font-variant-numeric: tabular-nums; color: var(--body-text-color); min-width:70px; } +.ardy-lbl { font-size:13px; color: var(--body-text-color); display:flex; align-items:center; gap:4px; } +.ardy-caption { margin-top:6px; font-size:13px; color:#667085; } +""" + +APP_CSS = """ +#col-container { max-width: 1200px; margin: 0 auto; } +.dark .gradio-container { color: var(--body-text-color); } +""" + +# Runs once, when the component first renders. `element`, `props`, and `watch` +# are injected by Gradio. We import three.js, decode the static skin data (human +# rig) and the static rigid-mesh data (robot rig), build the scene, and subscribe +# to value changes with `watch('value', ...)`. Each generated payload carries the +# per-frame joint affine matrices plus a `rig` tag; the viewer renders whichever +# rig the payload requests, swapping the on-screen mesh as needed. +PLAYER_JS_ON_LOAD = r""" +const root = element; +const q = (sel) => root.querySelector(sel); + +const state = { + ready:false, + THREE:null, OrbitControls:null, + renderer:null, scene:null, camera:null, controls:null, + trailLine:null, + data:null, verts:null, frame:0, playing:false, lastT:0, + speed:1, loop:true, trail:false, pending:null, + rig:null, // which rig mesh is currently mounted in the scene + // human rig + skin:null, humanMesh:null, humanGeom:null, + // robot rig + robotMesh:null, robotGeom:null, + robotBaseVerts:null, robotFaces:null, robotVtot:0, robotParts:null, +}; + +// --- binary helpers --------------------------------------------------------- +function b64ToBytes(b64){ + const bin = atob(b64); const out = new Uint8Array(bin.length); + for(let i=0;i{ + const w2 = wrap.clientWidth, h2 = wrap.clientHeight; + if(w2>0 && h2>0){ camera.aspect=w2/h2; camera.updateProjectionMatrix(); renderer.setSize(w2,h2); } + }).observe(wrap); + + animate(); +} + +// Mount the mesh for the requested rig (lazily built, then shown/hidden). Only +// one rig mesh is visible at a time; both share the scene once created. +function mountRig(rig){ + const THREE = state.THREE; + if(rig === "robot"){ + if(!state.robotMesh){ + const geom = new THREE.BufferGeometry(); + geom.setIndex(new THREE.BufferAttribute(state.robotFaces, 1)); + geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(state.robotVtot*3), 3)); + const mat = new THREE.MeshStandardMaterial({color:0xd7dde6, roughness:0.5, metalness:0.55}); + const mesh = new THREE.Mesh(geom, mat); + mesh.castShadow = true; mesh.frustumCulled = false; + state.scene.add(mesh); state.robotMesh = mesh; state.robotGeom = geom; + } + if(state.humanMesh) state.humanMesh.visible = false; + state.robotMesh.visible = true; + return state.robotGeom; + } else { + if(!state.humanMesh){ + const s = state.skin; + const geom = new THREE.BufferGeometry(); + geom.setIndex(new THREE.BufferAttribute(s.faces, 1)); + geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(s.V*3), 3)); + const mat = new THREE.MeshStandardMaterial({color:0x98bdff, roughness:0.85, metalness:0.0}); + const mesh = new THREE.Mesh(geom, mat); + mesh.castShadow = true; mesh.frustumCulled = false; + state.scene.add(mesh); state.humanMesh = mesh; state.humanGeom = geom; + } + if(state.robotMesh) state.robotMesh.visible = false; + state.humanMesh.visible = true; + return state.humanGeom; + } +} + +function activeGeom(){ + return (state.rig === "robot") ? state.robotGeom : state.humanGeom; +} + +function setFrame(f){ + if(!state.verts) return; + const T = state.data.num_frames; + f = Math.max(0, Math.min(T-1, f|0)); + state.frame = f; + const geom = activeGeom(); + if(!geom) return; + const pos = geom.getAttribute("position"); + pos.array.set(state.verts[f]); + pos.needsUpdate = true; + geom.computeVertexNormals(); + geom.computeBoundingSphere(); + updateTrail(); + const scrub = q(".ardy-scrub"); if(scrub) scrub.value = f; + const lbl = q(".ardy-frame"); if(lbl) lbl.textContent = (f+1)+" / "+T; +} + +function updateTrail(){ + const THREE = state.THREE; + if(!state.data || !state.data.root){ if(state.trailLine) state.trailLine.visible=false; return; } + if(!state.trail){ if(state.trailLine) state.trailLine.visible=false; return; } + const T = state.data.num_frames, root = state.data.root; + if(!state.trailLine){ + const g = new THREE.BufferGeometry(); + g.setAttribute("position", new THREE.BufferAttribute(new Float32Array(T*3),3)); + state.trailLine = new THREE.Line(g, new THREE.LineBasicMaterial({color:0xf59e0b})); + state.scene.add(state.trailLine); + } + state.trailLine.visible = true; + const attr = state.trailLine.geometry.getAttribute("position"); + for(let t=0;t= 1/Math.max(1e-3,fps)){ + state.lastT = now; + let nf = state.frame + 1; + if(nf >= state.data.num_frames){ + if(state.loop){ nf = 0; } else { nf = state.data.num_frames-1; setPlaying(false); } + } + setFrame(nf); + } + } + state.controls.update(); + state.renderer.render(state.scene, state.camera); +} + +function setPlaying(p){ + state.playing = p; + const b = q(".ardy-play"); + if(b) b.textContent = p ? "⏸ Pause" : "▶ Play"; + state.lastT = performance.now(); +} + +function loadData(data){ + initScene(); + const rig = (data.rig === "robot") ? "robot" : "human"; + state.rig = rig; + mountRig(rig); + state.data = data; + // Decode per-frame affines and pre-pose every frame (one-time cost per clip). + const T = data.num_frames, J = data.num_joints; + const Abytes = b64ToBytes(data.affines); + const A = new Float32Array(Abytes.buffer, Abytes.byteOffset, Abytes.byteLength/4); + state.verts = (rig === "robot") ? poseAllFrames(A, T, J) : skinAllFrames(A, T, J); + if(state.trailLine){ state.scene.remove(state.trailLine); state.trailLine.geometry.dispose(); state.trailLine=null; } + const scrub = q(".ardy-scrub"); if(scrub){ scrub.max = T-1; scrub.value = 0; } + const hint = q(".ardy-hint"); if(hint) hint.style.display = "none"; + const cap = q(".ardy-caption"); + if(cap) cap.textContent = '"' + data.prompt + '" · ' + (rig==="robot"?"robot":"human") + + ' · ' + T + ' frames @ ' + data.fps + ' fps · seed ' + data.seed; + root.dataset.ardyLoaded = "1"; + root.dataset.ardyRig = rig; + root.dataset.ardyFrames = String(T); + setFrame(0); + setPlaying(true); +} + +function applyPayload(payload){ + if(!payload) return; + if(!state.ready){ state.pending = payload; return; } // three.js / rigs still loading + try { loadData(JSON.parse(payload)); } + catch(e){ console.error("ARDY playground load error", e); } +} + +function wireControls(){ + const bind = (sel, ev, fn) => { + const el = q(sel); if(el && !el.dataset.wired){ el.dataset.wired="1"; el.addEventListener(ev, fn); } + }; + bind(".ardy-play", "click", ()=> setPlaying(!state.playing)); + bind(".ardy-scrub", "input", (e)=>{ setPlaying(false); setFrame(parseInt(e.target.value)); }); + bind(".ardy-speed", "change", (e)=>{ state.speed = parseFloat(e.target.value); }); + bind(".ardy-loop", "change", (e)=>{ state.loop = e.target.checked; }); + bind(".ardy-trail", "change", (e)=>{ state.trail = e.target.checked; updateTrail(); }); +} + +// Gradio-native hand-off: render whenever the component's value prop updates as +// the output of a Python event (Generate / Continue / Examples). +if (typeof watch === "function") { + watch("value", () => applyPayload(props.value)); +} + +// Import three.js (esm.sh, not jsDelivr: the OrbitControls addon has an internal +// bare `import ... from "three"` that a browser dynamic import() can't resolve +// without an import map; esm.sh rewrites it and dedupes three) and decode both +// rigs, then flush any value that already arrived. +Promise.all([ + import("https://esm.sh/three@0.160.0"), + import("https://esm.sh/three@0.160.0/examples/jsm/controls/OrbitControls.js"), + decodeSkin(), + decodeRig(), +]).then(([THREE, oc, skin, rig]) => { + state.THREE = THREE; + state.OrbitControls = oc.OrbitControls; + state.skin = skin; + state.robotBaseVerts = rig.baseVerts; + state.robotFaces = rig.faces; + state.robotVtot = rig.V; + state.robotParts = rig.parts; + state.ready = true; + wireControls(); + initScene(); + const start = state.pending || props.value; + if (start) applyPayload(start); +}).catch((e)=> console.error("ARDY viewer init failed", e)); +""" + +EXAMPLES = [ + ["A person walks forward confidently.", "human", 5.0], + ["A person walks in a circle.", "human", 6.0], + ["A person jumps up and down.", "human", 4.0], + ["The robot walks forward confidently.", "robot", 5.0], + ["The robot waves with the right hand.", "robot", 4.0], + ["The robot crouches down and then stands back up.", "robot", 5.0], +] + +with gr.Blocks() as demo: + with gr.Column(elem_id="col-container"): + gr.Markdown( + """ + # 🕺 ARDY Motion Playground + Interactive text-to-motion with **[ARDY](https://research.nvidia.com/labs/sil/projects/ardy/)** + (Autoregressive Diffusion with Hybrid Representation) by NVIDIA. + Pick a **rig** (human or robot), type a prompt and **Generate** a 3D motion clip, + then **orbit, scrub, and play** it below. + Chain actions with **Continue ▸** — the same character keeps going from where it stopped. + """ + ) + + # Running latent state (normalized feature tensor, CPU) — lets "Continue" + # resume the same character from the end of the previous clip. + seq_state = gr.State(None) + + prompt = gr.Textbox( + label="Motion prompt", + placeholder="e.g. a person walks in a circle then waves", + lines=2, + ) + rig = gr.Radio( + choices=[("🕺 Human", "human"), ("🤖 Robot (Unitree G1)", "robot")], + value=DEFAULT_RIG, + label="Rig", + ) + with gr.Row(): + run = gr.Button("Generate", variant="primary", scale=2) + cont = gr.Button("Continue ▸", variant="secondary", scale=1) + + # The interactive 3D playground — a Gradio-native custom HTML component. + # Its `value` (the motion JSON) is set directly by the Generate handler; + # `watch('value', ...)` in js_on_load renders it. Ship the static data for + # both rigs (human skin blob + G1 rigid-mesh blob) once, as a header + # prepended to js_on_load; the per-frame affines ride in each payload. + _rig_header = ( + f'const ARDY_SKIN_B64="{SKIN_B64}";\n' + f"const ARDY_SKIN_META={json.dumps(SKIN_META)};\n" + f'const G1_MESH_B64="{G1_MESH_B64}";\n' + f"const G1_MESH_META={json.dumps(G1_MESH_META)};\n" + ) + player = gr.HTML( + value="", + html_template=PLAYER_TEMPLATE, + css_template=PLAYER_CSS_TEMPLATE, + js_on_load=_rig_header + PLAYER_JS_ON_LOAD, + elem_id="ardy-player", + ) + + with gr.Accordion("Advanced settings", open=False): + duration = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Duration (seconds)") + diffusion_steps = gr.Slider( + 1, NUM_BASE_STEPS, value=NUM_BASE_STEPS, step=1, label="Diffusion steps" + ) + cfg_weight = gr.Slider(1.0, 6.0, value=2.0, step=0.5, label="Text guidance (CFG)") + with gr.Row(): + randomize_seed = gr.Checkbox(label="Randomize seed", value=True) + seed = gr.Number(label="Seed", value=0, precision=0) + + with gr.Accordion("Blender export", open=False): + gr.Markdown( + "Generate the human rig as **BVH**, full-fidelity **NPZ**, and " + "JSON metadata. Import the BVH directly in Blender." + ) + export_blender = gr.Button("Generate Blender files", variant="secondary") + with gr.Row(): + bvh_file = gr.File(label="Blender BVH") + npz_file = gr.File(label="ARDY NPZ") + metadata_file = gr.File(label="Metadata") + + gr.Examples( + examples=EXAMPLES, + inputs=[prompt, rig, duration], + outputs=[player, seed, seq_state], + fn=ui_generate, + cache_examples=False, + run_on_click=True, + ) + + gr.Markdown( + """ + Rigs: **ARDY-Core-RP-20FPS-Horizon40** (human, 27-joint skeleton @ 20 fps) + and **ARDY-G1-RP-25FPS-Horizon52** (Unitree G1 robot, 34-joint skeleton @ 25 fps). + Text encoder: LLM2Vec-Llama-3-8B. Post-processing (foot-skate cleanup) is disabled in this demo. + Motion is generated autoregressively; longer clips take longer. + **Generate** starts a new clip; **Continue ▸** keeps the same character going, + transitioning it into the new prompt (like the reference demo's prompt timeline). + """ + ) + + _gen_inputs = [prompt, rig, duration, diffusion_steps, cfg_weight, seed, randomize_seed] + # Generate starts fresh; Continue resumes from the running latent state. + # The payload is written straight into the player's `value`; its js_on_load + # `watch('value', ...)` renders it (loading the correct rig from the payload). + run.click(fn=ui_generate, inputs=_gen_inputs, + outputs=[player, seed, seq_state], api_name=False) + cont.click(fn=ui_continue, inputs=_gen_inputs + [seq_state], + outputs=[player, seed, seq_state], api_name=False) + export_blender.click( + fn=generate_blender_motion, + inputs=[prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed], + outputs=[bvh_file, npz_file, metadata_file, seed], + api_name=False, + ) + + # Clean single-shot endpoint for the HTTP API / MCP tool (no session state). + gr.api(generate_motion, api_name="generate") + gr.api(generate_blender_motion, api_name="generate_blender") + +demo.queue() + +if __name__ == "__main__": + # Gradio 6 moved theme/css from the Blocks constructor to launch(). The + # player's JS/CSS now live on the gr.HTML component itself (js_on_load / + # css_template), so no global `head=` script is needed. + demo.launch( + theme=gr.themes.Citrus(), + css=APP_CSS, + mcp_server=True, + ssr_mode=False, + ) diff --git a/ardy/__init__.py b/ardy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a6a5f2d78447f61515d5524a9d26422c820143cc --- /dev/null +++ b/ardy/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ARDY: text-to-motion generation.""" diff --git a/ardy/assets.py b/ardy/assets.py new file mode 100644 index 0000000000000000000000000000000000000000..91facad0faed373c9ed1ad4667980cf19788b093 --- /dev/null +++ b/ardy/assets.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +PACKAGE_ROOT = Path(__file__).resolve().parent +ASSETS_ROOT = PACKAGE_ROOT / "assets" +DEMO_ASSETS_ROOT = ASSETS_ROOT / "demo" +DEMO_EXAMPLES_ROOT = DEMO_ASSETS_ROOT / "examples" +SKELETONS_ROOT = ASSETS_ROOT / "skeletons" +SOMA_ASSETS_ROOT = ASSETS_ROOT / "SOMA" + + +def skeleton_asset_path(*parts: str) -> Path: + return SKELETONS_ROOT.joinpath(*parts) + + +def demo_asset_path(*parts: str) -> Path: + return DEMO_ASSETS_ROOT.joinpath(*parts) diff --git a/ardy/assets/skeletons/cskel27/joints.p b/ardy/assets/skeletons/cskel27/joints.p new file mode 100644 index 0000000000000000000000000000000000000000..0f72152f6bd511cf913864eea086f5f14e668b90 Binary files /dev/null and b/ardy/assets/skeletons/cskel27/joints.p differ diff --git a/ardy/assets/skeletons/cskel27/skin_standard.npz b/ardy/assets/skeletons/cskel27/skin_standard.npz new file mode 100644 index 0000000000000000000000000000000000000000..3d95f366379b0ece0ec9b0355672c8a115c4b526 --- /dev/null +++ b/ardy/assets/skeletons/cskel27/skin_standard.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ed2b1ea33a997d7777a59d74ec2312b52af2276331048c954441f5f44f6e58b +size 237170 diff --git a/ardy/assets/skeletons/cskel29/joints.p b/ardy/assets/skeletons/cskel29/joints.p new file mode 100644 index 0000000000000000000000000000000000000000..e97d757feeef5493fd724ad1073fb32dea922aa5 Binary files /dev/null and b/ardy/assets/skeletons/cskel29/joints.p differ diff --git a/ardy/assets/skeletons/cskel29/standard_t_pose_cskel29_joint_quat.p b/ardy/assets/skeletons/cskel29/standard_t_pose_cskel29_joint_quat.p new file mode 100644 index 0000000000000000000000000000000000000000..8cb8b2a3ac38a570921cdf7e1890a4432dc73b92 Binary files /dev/null and b/ardy/assets/skeletons/cskel29/standard_t_pose_cskel29_joint_quat.p differ diff --git a/ardy/assets/skeletons/g1skel34/joints.p b/ardy/assets/skeletons/g1skel34/joints.p new file mode 100644 index 0000000000000000000000000000000000000000..4d5853a6e1ac0e6f2f854cd1a8893293612fa3ac Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/joints.p differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/head_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/head_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..401f822751c33dab53c6cee77eb62890aca5cd8f --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/head_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:005fb67fbd3eff94aa8bf4a6e83238174e9f91b6721f7111594322f223724411 +size 932784 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..69de8490184afc698f633e34bb74e65351bd4689 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_pitch_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..cb69f65c6a8059fc6720cb393540f579b4d0c26e --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_ankle_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c4092af943141d4d9f74232f3cfa345afc6565f46a077793b8ae0e68b39dc33f +size 653384 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_elbow_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_elbow_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..1a96d99ba469960173129084ab6dd3bf8a732a71 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/left_elbow_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..4cf7475bf6bbaef47d5b3f7644a34c73b3d64037 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_pitch_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4725168105ee768ee31638ef22b53f6be2d7641bfd7cfefe803488d884776fa4 +size 181684 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..585f6040c38571177d55519417a0f46ca7b29381 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:91f25922ee8a7c3152790051bebad17b4d9cd243569c38fe340285ff93a97acf +size 192184 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..b46a74134f90c568cbdfb8484eeb83a0503b4504 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_hip_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a16d88aa6ddac8083aa7ad55ed317bea44b1fa003d314fba88965b7ed0f3b55b +size 296284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_knee_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_knee_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..2dcf84e4018ab9e4cff48cecdafe433601c48480 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_knee_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d92b9e3d3a636761150bb8025e32514c4602b91c7028d523ee42b3e632de477 +size 854884 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_rubber_hand.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_rubber_hand.STL new file mode 100644 index 0000000000000000000000000000000000000000..04a2fa22b7f60f2d00c64ec996b33a5219c2b633 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_rubber_hand.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cff2221a690fa69303f61fce68f2d155c1517b52efb6ca9262dd56e0bc6e70fe +size 2287484 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..926d9807f1cbf6ad0d41170904785eb30427c848 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_pitch_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f0d1cfd02fcf0d42f95e678eeca33da3afbcc366ffba5c052847773ec4f31d52 +size 176784 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..4c6840b93ea15d1ba7b4dac0f18f94655867cab7 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fb9df21687773522598dc384f1a2945c7519f11cbc8bd372a49170316d6eee88 +size 400284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..89b0e0661275ed832c674a46893239ad2e5bef85 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_shoulder_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1aa97e9748e924336567992181f78c7cd0652fd52a4afcca3df6b2ef6f9e712e +size 249184 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..82cc224a8e41251d879502f9809e31d0988ec7f9 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_pitch_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..2097ca3e5bbc0af985d150e5dcda4d8ba9bcf43f --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:edc387c9a0ba8c2237e9b296d32531426fabeb6f53e58df45c76106bca74148c +size 356184 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..692f4b07105fa0ec4ffd989dc60be9b2a33ac9c3 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/left_wrist_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:83f8fb3a726bf9613d65dd14f0f447cb918c3c95b3938042a0c9c09749267d3b +size 318684 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/logo_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/logo_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..6c25961c594c714ada9d32e922250d326c121f4a --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/logo_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8571a0a19bc4916fa55f91449f51d5fdefd751000054865a842449429d5f155b +size 243384 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis.STL new file mode 100644 index 0000000000000000000000000000000000000000..f98a88dbf6c7698791ba694e16cbcc180c40663c --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ba6bbc888e630550140d3c26763f10206da8c8bd30ed886b8ede41c61f57a31 +size 1060884 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis_contour_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis_contour_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..8025bc070cd65550f345e9e97bc44a8ff9740dea --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/pelvis_contour_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5cc5c2c7a312329e3feeb2b03d3fc09fc29705bd01864f6767e51be959662420 +size 1805184 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..e77d8a2fe1e5d56fac049833d254d6ffa4f6b350 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_pitch_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..3c38507b9060eda2c40e5e2bb1b174ef5bda0390 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_ankle_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4b66222ea56653e627711b56d0a8949b4920da5df091da0ceb343f54e884e3a5 +size 653784 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_elbow_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_elbow_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..f259e3812efb9985d6463c04d3e8a4b53793a699 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/right_elbow_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..064085fca108bfc2a5a365403ceb6cdac4d5cd5f --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_pitch_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4f3c99d7f4a7d34eadbef9461fc66e3486cb5442db1ec50c86317d459f1a9c6 +size 181284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..6544025eeb7dc7d09da549746ae28272e5bcf224 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4c254ef66a356f492947f360dd931965477b631e3fcc841f91ccc46d945d54f6 +size 192684 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..0ad7bee36419f85b490debe929957c6932df393f --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_hip_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e479c2936ca2057e9eb2f7dff6c189b7419d7b8484dea0b298cbb36a2a6aa668 +size 296284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_knee_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_knee_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..65e8a708a7e86f1f242cd32e8959a560ea66c2d7 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_knee_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63c4008449c9bbe701a6e2b557b7a252e90cf3a5abcf54cee46862b9a69f8ec8 +size 852284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_rubber_hand.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_rubber_hand.STL new file mode 100644 index 0000000000000000000000000000000000000000..58148fb9e7307176d43f71da139781a77b63d156 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_rubber_hand.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:99533b778bca6246144fa511bb9d4e555e075c641f2a0251e04372869cd99d67 +size 2192684 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..48a1c46009929341f2c312966acf81c3c733d7ec --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_pitch_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24cdb387e0128dfe602770a81c56cdce3a0181d34d039a11d1aaf8819b7b8c02 +size 176784 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..2a5d22f9e17a986ec8c958a6ff7668d65878a907 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:962b97c48f9ce9e8399f45dd9522e0865d19aa9fd299406b2d475a8fc4a53e81 +size 401884 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..0882a567fc1a11f0cf821182154551058b0517f3 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_shoulder_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0b76489271da0c72461a344c9ffb0f0c6e64f019ea5014c1624886c442a2fe5 +size 249984 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_pitch_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..da194543c40df9d492abb0553c06cc614a78caae Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_pitch_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..77d23a77ad6ab659884db4d540e0bd9fd0675d0f --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_roll_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a7ee9212ff5b94d6cb7f52bb1bbf3f352194d5b598acff74f4c77d340c5b344f +size 356084 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..77edbb4135bb5b9b64357a0226431fd665d794b2 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/right_wrist_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc9dece2d12509707e0057ba2e48df8f3d56db0c79410212963a25e8a50f61a6 +size 341484 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..e4fb87c4cc736c90970bb5f182981f0cf1544038 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e96d023f0368a4e3450b86ca5d4f10227d8141156a373e7da8cb3c93266523e0 +size 2232984 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link_rev_1_0.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link_rev_1_0.STL new file mode 100644 index 0000000000000000000000000000000000000000..836b9923b3edac7a1c06ecd2f11bfbbc0dedf7e0 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/torso_link_rev_1_0.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:11ddb46f2098efbbd8816b1d65893632d8e78be936376c7cdcd6771899ccc723 +size 2570584 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..65831abd2a64bc8c36e31016964413a3d2116725 Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link_rev_1_0.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link_rev_1_0.STL new file mode 100644 index 0000000000000000000000000000000000000000..a64f330c592582dc31cdf38f4d08cfff06681c5f Binary files /dev/null and b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_roll_link_rev_1_0.STL differ diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..c049debc4e0080a06ea95c2e7e9b74613cd42040 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2883f20e03f09b669b5b4ce10677ee6b5191c0934b584d7cbaef2d0662856ffb +size 336284 diff --git a/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link_rev_1_0.STL b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link_rev_1_0.STL new file mode 100644 index 0000000000000000000000000000000000000000..dc628fbdda04129fda29140243d148cbd2c81c83 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link_rev_1_0.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ec6db442b11f25eed898b5add07940c85d804f300de24dcbd264ccd8be7d554c +size 619984 diff --git a/ardy/assets/skeletons/g1skel34/xml/g1.xml b/ardy/assets/skeletons/g1skel34/xml/g1.xml new file mode 100644 index 0000000000000000000000000000000000000000..36231a5b97ab1784ad2e2c889b974b086639c120 --- /dev/null +++ b/ardy/assets/skeletons/g1skel34/xml/g1.xml @@ -0,0 +1,413 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + > + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/ardy/assets/skeletons/somaskel30/joints.p b/ardy/assets/skeletons/somaskel30/joints.p new file mode 100644 index 0000000000000000000000000000000000000000..9e319f6fa3cb4959ed3221040b0ae945ec9646c7 Binary files /dev/null and b/ardy/assets/skeletons/somaskel30/joints.p differ diff --git a/ardy/assets/skeletons/somaskel30/soma_base_fit_mhr_params.npz b/ardy/assets/skeletons/somaskel30/soma_base_fit_mhr_params.npz new file mode 100644 index 0000000000000000000000000000000000000000..25ea8aa2ee5d40a79c7cf2761b6b3a28bd7c1133 --- /dev/null +++ b/ardy/assets/skeletons/somaskel30/soma_base_fit_mhr_params.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:61f0b96c8386a19d44823faf4399816642b7e20691ed5b398c7402921090096e +size 1877 diff --git a/ardy/assets/skeletons/somaskel77/joints.p b/ardy/assets/skeletons/somaskel77/joints.p new file mode 100644 index 0000000000000000000000000000000000000000..8c5eb2e67e1314d8765c821ae8dd4e0ff1351a3b Binary files /dev/null and b/ardy/assets/skeletons/somaskel77/joints.p differ diff --git a/ardy/assets/skeletons/somaskel77/relaxed_hands_rest_pose.npy b/ardy/assets/skeletons/somaskel77/relaxed_hands_rest_pose.npy new file mode 100644 index 0000000000000000000000000000000000000000..74dc084f8e62e534a64e840141165cd56f6d85e5 --- /dev/null +++ b/ardy/assets/skeletons/somaskel77/relaxed_hands_rest_pose.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:64a3828e0d1ef1f1de8228c74eba8040c0810d898169f1892d8997a213b2b64c +size 2900 diff --git a/ardy/assets/skeletons/somaskel77/skin_standard.npz b/ardy/assets/skeletons/somaskel77/skin_standard.npz new file mode 100644 index 0000000000000000000000000000000000000000..4a82cd790d42a6a48a858b13bf8f6594bac30241 --- /dev/null +++ b/ardy/assets/skeletons/somaskel77/skin_standard.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:90ee2cdf50f168382a7dd7e9c88b118298aba0aca4b72022080c89c4bab0ceb2 +size 531434 diff --git a/ardy/assets/skeletons/somaskel77/standard_t_pose_global_offsets_rots.p b/ardy/assets/skeletons/somaskel77/standard_t_pose_global_offsets_rots.p new file mode 100644 index 0000000000000000000000000000000000000000..b6d68af1e73df0fc69c2cd9a59f2f6ed6aebb425 Binary files /dev/null and b/ardy/assets/skeletons/somaskel77/standard_t_pose_global_offsets_rots.p differ diff --git a/ardy/assets/skeletons/somaskel77/standard_t_pose_somaskel77_joint_quat.p b/ardy/assets/skeletons/somaskel77/standard_t_pose_somaskel77_joint_quat.p new file mode 100644 index 0000000000000000000000000000000000000000..5c5d808beae8a06e8a02b559250c55a38b9526af Binary files /dev/null and b/ardy/assets/skeletons/somaskel77/standard_t_pose_somaskel77_joint_quat.p differ diff --git a/ardy/constraints.py b/ardy/constraints.py new file mode 100644 index 0000000000000000000000000000000000000000..72f546d84d7e215bdb7115f1277e6c2373ed9620 --- /dev/null +++ b/ardy/constraints.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +from torch import Tensor + +from ardy.motion_rep.tools import compute_heading_angle +from ardy.skeleton import SkeletonBase + +from .geometry import axis_angle_to_matrix, matrix_to_axis_angle + + +def create_pairs(tensor_A, tensor_B): + pairs = torch.stack( + ( + tensor_A[:, None].expand(-1, len(tensor_B)), + tensor_B.expand(len(tensor_A), -1), + ), + dim=-1, + ).reshape(-1, 2) + return pairs + + +def compute_global_heading(global_joints_positions: Tensor, skeleton: SkeletonBase): + root_heading_angle = compute_heading_angle(global_joints_positions, skeleton) + global_root_heading = torch.stack([torch.cos(root_heading_angle), torch.sin(root_heading_angle)], dim=-1) + return global_root_heading + + +class Root2DConstraintSet: + name = "root2d" + + def __init__( + self, + skeleton: SkeletonBase, + frame_indices: Tensor, + root_2d: Tensor, + global_root_heading: Optional[Tensor] = None, + to_crop: bool = False, + ) -> None: + self.skeleton = skeleton + if to_crop: + root_2d = root_2d[frame_indices] + if global_root_heading is not None: + global_root_heading = global_root_heading[frame_indices] + else: + assert len(root_2d) == len(frame_indices), "The number of root 2d should be match the number of frames" + if global_root_heading is not None: + assert len(global_root_heading) == len(frame_indices), ( + "The number of global root heading should match the number of frames" + ) + self.root_2d = root_2d + self.global_root_heading = global_root_heading + self.frame_indices = frame_indices + + def update_constraints(self, data_dict: dict, index_dict: dict) -> None: + data_dict["root_2d"].append(self.root_2d) + index_dict["root_2d"].append(self.frame_indices) + if self.global_root_heading is not None: + # Convert heading angles to [cos, sin] format + # self.global_root_heading contains angles in radians + heading_cos_sin = torch.stack( + [ + torch.cos(self.global_root_heading), + torch.sin(self.global_root_heading), + ], + dim=-1, + ) + data_dict["global_root_heading"].append(heading_cos_sin) + index_dict["global_root_heading"].append(self.frame_indices) + + def crop_move(self, start: int, end: int): + mask = (self.frame_indices >= start) & (self.frame_indices < end) + return ( + Root2DConstraintSet( + self.skeleton, + self.frame_indices[mask] - start, + self.root_2d[mask], + self.global_root_heading[mask], + ) + if self.global_root_heading is not None + else Root2DConstraintSet(self.skeleton, self.frame_indices[mask] - start, self.root_2d[mask]) + ) + + def get_save_info(self): + info = { + "type": self.name, + "frame_indices": self.frame_indices, + "root_2d": self.root_2d, + } + if self.global_root_heading is not None: + info["global_root_heading"] = self.global_root_heading + return info + + @classmethod + def from_dict(cls, skeleton: SkeletonBase, dico: dict): + device = skeleton.device + root_2d_key = "root_2d" if "root_2d" in dico else "smooth_root_2d" + return cls( + skeleton, + frame_indices=torch.tensor(dico["frame_indices"]), + root_2d=torch.tensor(dico[root_2d_key], device=device), + global_root_heading=torch.tensor(dico["global_root_heading"]) if "global_root_heading" in dico else None, + ) + + +class FullBodyConstraintSet: + name = "fullbody" + + def __init__( + self, + skeleton: SkeletonBase, + frame_indices: Tensor, + global_joints_positions: Tensor, + global_joints_rots: Tensor, + root_2d: Optional[Tensor] = None, + to_crop: bool = False, + ): + self.skeleton = skeleton + self.frame_indices = frame_indices + + if to_crop: + global_joints_positions = global_joints_positions[frame_indices] + global_joints_rots = global_joints_rots[frame_indices] + if root_2d is not None: + root_2d = root_2d[frame_indices] + else: + assert len(global_joints_positions) == len(frame_indices), ( + "The number of global positions should be match the number of frames" + ) + assert len(global_joints_rots) == len(frame_indices), ( + "The number of global joint rotations should be match the number of frames" + ) + + if root_2d is not None: + assert len(root_2d) == len(frame_indices), ( + "The number of root 2d (if specified) should be match the number of frames" + ) + + if root_2d is None: + # substitute root 2d with the real root + root_2d = global_joints_positions[:, skeleton.root_idx, [0, 2]] + + # root y: from smooth or pelvis is the same + self.root_y_pos = global_joints_positions[:, skeleton.root_idx, 1] + + self.global_joints_positions = global_joints_positions + self.global_joints_rots = global_joints_rots + self.global_root_heading = compute_global_heading(global_joints_positions, skeleton) + self.root_2d = root_2d + + def update_constraints(self, data_dict, index_dict): + nbjoints = self.skeleton.nbjoints + indices_lst = create_pairs( + self.frame_indices, + torch.arange(nbjoints), + ) + data_dict["global_joints_positions"].append( + self.global_joints_positions.reshape(-1, 3) + ) # flatten the global positions + index_dict["global_joints_positions"].append(indices_lst) + + # global rotations are not used here + + # also constraint root 2d to get the same full body + # maybe keep storing the hips offset, if we smooth it ourselves + data_dict["root_2d"].append(self.root_2d) + index_dict["root_2d"].append(self.frame_indices) + + # constraint the y pos of the root + data_dict["root_y_pos"].append(self.root_y_pos) + index_dict["root_y_pos"].append(self.frame_indices) + + # constraint the global heading + data_dict["global_root_heading"].append(self.global_root_heading) + index_dict["global_root_heading"].append(self.frame_indices) + + def crop_move(self, start: int, end: int): + mask = (self.frame_indices >= start) & (self.frame_indices < end) + return FullBodyConstraintSet( + self.skeleton, + self.frame_indices[mask] - start, + self.global_joints_positions[mask], + self.global_joints_rots[mask], + self.root_2d[mask], + ) + + def get_save_info(self): + local_joints_rot = self.skeleton.global_rots_to_local_rots(self.global_joints_rots) + local_joints_rot = matrix_to_axis_angle(local_joints_rot) + + root_positions = self.global_joints_positions[:, self.skeleton.root_idx] + return { + "type": self.name, + "frame_indices": self.frame_indices, + "local_joints_rot": local_joints_rot, + "root_positions": root_positions, + "root_2d": self.root_2d, + } + + @classmethod + def from_dict(cls, skeleton: SkeletonBase, dico: dict): + frame_indices = torch.tensor(dico["frame_indices"]) + device = skeleton.device + global_joints_rots, global_joints_positions, _ = skeleton.fk( + axis_angle_to_matrix(torch.tensor(dico["local_joints_rot"], device=device)), + torch.tensor(dico["root_positions"], device=device), + ) + root_2d = None + if "root_2d" in dico: + root_2d = torch.tensor(dico["root_2d"], device=device) + elif "smooth_root_2d" in dico: + root_2d = torch.tensor(dico["smooth_root_2d"], device=device) + + return cls( + skeleton, + frame_indices=frame_indices, + global_joints_positions=global_joints_positions, + global_joints_rots=global_joints_rots, + root_2d=root_2d, + ) + + +class EndEffectorConstraintSet: + name = "end-effector" + + def __init__( + self, + skeleton: SkeletonBase, + frame_indices: Tensor, + global_joints_positions: Tensor, + global_joints_rots: Tensor, + root_2d: Optional[Tensor], + *, + joint_names: list[str], + to_crop: bool = False, + ) -> None: + self.skeleton = skeleton + self.frame_indices = frame_indices + self.joint_names = joint_names + + # joint_names are constant for all the frames + rot_joint_names, pos_joint_names = self.skeleton.expand_joint_names(self.joint_names) + # indexing works for motion_rep with smooth root only (contains pelvis index) + self.pos_indices = torch.tensor([self.skeleton.bone_index[jname] for jname in pos_joint_names]) + self.rot_indices = torch.tensor([self.skeleton.bone_index[jname] for jname in rot_joint_names]) + + if to_crop: + global_joints_positions = global_joints_positions[frame_indices] + global_joints_rots = global_joints_rots[frame_indices] + if root_2d is not None: + root_2d = root_2d[frame_indices] + else: + assert len(global_joints_positions) == len(frame_indices), ( + "The number of global positions should be match the number of frames" + ) + assert len(global_joints_rots) == len(frame_indices), ( + "The number of global joint rotations should be match the number of frames" + ) + if root_2d is not None: + assert len(root_2d) == len(frame_indices), ( + "The number of root 2d (if specified) should be match the number of frames" + ) + + if root_2d is None: + # substitute root 2d with the real root + root_2d = global_joints_positions[:, skeleton.root_idx, [0, 2]] + + # root y: from smooth or pelvis is the same + self.root_y_pos = global_joints_positions[:, skeleton.root_idx, 1] + + self.global_joints_positions = global_joints_positions + self.global_root_heading = compute_global_heading(global_joints_positions, skeleton) + self.global_joints_rots = global_joints_rots + self.root_2d = root_2d + + def update_constraints(self, data_dict, index_dict): + crop_frames_indexing = torch.arange(len(self.frame_indices)) + + # constraint positions + pos_indices_real = create_pairs( + self.frame_indices, + self.pos_indices, + ) + pos_indices_crop = create_pairs( + crop_frames_indexing, + self.pos_indices, + ) + data_dict["global_joints_positions"].append(self.global_joints_positions[tuple(pos_indices_crop.T)]) + index_dict["global_joints_positions"].append(pos_indices_real) + + # constraint rotations + rot_indices_real = create_pairs( + self.frame_indices, + self.rot_indices, + ) + rot_indices_crop = create_pairs( + crop_frames_indexing, + self.rot_indices, + ) + data_dict["global_joints_rots"].append(self.global_joints_rots[tuple(rot_indices_crop.T)]) + index_dict["global_joints_rots"].append(rot_indices_real) + + # also constraint root 2d to get the same full body + # maybe keep storing the hips offset, if we smooth it ourselves + data_dict["root_2d"].append(self.root_2d) + index_dict["root_2d"].append(self.frame_indices) + + # constraint the y pos of the root + data_dict["root_y_pos"].append(self.root_y_pos) + index_dict["root_y_pos"].append(self.frame_indices) + + # constraint the global heading + data_dict["global_root_heading"].append(self.global_root_heading) + index_dict["global_root_heading"].append(self.frame_indices) + + def crop_move(self, start: int, end: int): + mask = (self.frame_indices >= start) & (self.frame_indices < end) + + cls = type(self) + kwargs = {} + if not hasattr(cls, "joint_names"): + kwargs["joint_names"] = self.joint_names + + return cls( + self.skeleton, + self.frame_indices[mask] - start, + self.global_joints_positions[mask], + self.global_joints_rots[mask], + self.root_2d[mask], + **kwargs, + ) + + def get_save_info(self): + local_joints_rot = self.skeleton.global_rots_to_local_rots(self.global_joints_rots) + local_joints_rot = matrix_to_axis_angle(local_joints_rot) + + root_positions = self.global_joints_positions[:, self.skeleton.root_idx] + output = { + "type": self.name, + "frame_indices": self.frame_indices, + "local_joints_rot": local_joints_rot, + "root_positions": root_positions, + "root_2d": self.root_2d, + } + if not hasattr(self.__class__, "joint_names"): + # save the joint_names for this base class + # but not for children + output["joint_names"] = self.joint_names + return output + + @classmethod + def from_dict(cls, skeleton: SkeletonBase, dico: dict): + frame_indices = torch.tensor(dico["frame_indices"]) + device = skeleton.device + global_joints_rots, global_joints_positions, _ = skeleton.fk( + axis_angle_to_matrix(torch.tensor(dico["local_joints_rot"], device=device)), + torch.tensor(dico["root_positions"], device=device), + ) + root_2d = None + if "root_2d" in dico: + root_2d = torch.tensor(dico["root_2d"], device=device) + elif "smooth_root_2d" in dico: + root_2d = torch.tensor(dico["smooth_root_2d"], device=device) + + kwargs = {} + if not hasattr(cls, "joint_names"): + kwargs["joint_names"] = dico["joint_names"] + + return cls( + skeleton, + frame_indices=frame_indices, + global_joints_positions=global_joints_positions, + global_joints_rots=global_joints_rots, + root_2d=root_2d, + **kwargs, + ) + + +class LeftHandConstraintSet(EndEffectorConstraintSet): + name = "left-hand" + joint_names: list[str] = ["LeftHand", "Hips"] + + def __init__(self, *args, **kwargs: dict): + super().__init__(*args, joint_names=self.joint_names, **kwargs) + + +class RightHandConstraintSet(EndEffectorConstraintSet): + name = "right-hand" + joint_names: list[str] = ["RightHand", "Hips"] + + def __init__(self, *args, **kwargs: dict): + super().__init__(*args, joint_names=self.joint_names, **kwargs) + + +class LeftFootConstraintSet(EndEffectorConstraintSet): + name = "left-foot" + joint_names: list[str] = ["LeftFoot", "Hips"] + + def __init__(self, *args, **kwargs: dict): + super().__init__(*args, joint_names=self.joint_names, **kwargs) + + +class RightFootConstraintSet(EndEffectorConstraintSet): + name = "right-foot" + joint_names: list[str] = ["RightFoot", "Hips"] + + def __init__(self, *args, **kwargs: dict): + super().__init__(*args, joint_names=self.joint_names, **kwargs) + + +TYPE_TO_CLASS = { + "root2d": Root2DConstraintSet, + "fullbody": FullBodyConstraintSet, + "left-hand": LeftHandConstraintSet, + "right-hand": RightHandConstraintSet, + "left-foot": LeftFootConstraintSet, + "right-foot": RightFootConstraintSet, + "end-effector": EndEffectorConstraintSet, +} + + +def load_constraints_lst(path_or_data: str | list, skeleton: SkeletonBase): + from ardy.tools import load_json + + if isinstance(path_or_data, str): + saved = load_json(path_or_data) + else: + saved = path_or_data + + constraints_lst = [] + for el in saved: + cls = TYPE_TO_CLASS[el["type"]] + constraints_lst.append(cls.from_dict(skeleton, el)) + return constraints_lst + + +def save_constraints_lst(path: str, constraints_lst): + from ardy.tools import save_json + + if not constraints_lst: + print("The constraints lst is empty. Skip saving") + return + + to_save = [] + + def tensor_to_list(obj): + """Recursively convert tensors to lists for JSON serialization.""" + if isinstance(obj, Tensor): + return obj.cpu().tolist() + elif isinstance(obj, dict): + return {k: tensor_to_list(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [tensor_to_list(v) for v in obj] + else: + return obj + + for constraint in constraints_lst: + constraint_info = constraint.get_save_info() + # Convert all tensors to lists for JSON serialization + constraint_info = tensor_to_list(constraint_info) + to_save.append(constraint_info) + + save_json(path, to_save) + print(f"Saved constraints to {path}") + return to_save diff --git a/ardy/data_processing/__init__.py b/ardy/data_processing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..52a7a9daf028f654d6a02f65c99cc609a57b492f --- /dev/null +++ b/ardy/data_processing/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/ardy/data_processing/bvh.py b/ardy/data_processing/bvh.py new file mode 100644 index 0000000000000000000000000000000000000000..0a35bb951ee795e0e6d76768b75d1baca0ac34c9 --- /dev/null +++ b/ardy/data_processing/bvh.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import re +from typing import Optional + +import numpy as np + + +class BvhNode: + """Abstraction for bvh node.""" + + def __init__(self, value=[], parent=None): + self.value = value + self.children = [] + self.parent = parent + if self.parent: + self.parent.add_child(self) + + def add_child(self, item): + item.parent = self + self.children.append(item) + + def filter(self, key): + for child in self.children: + if child.value[0] == key: + yield child + + def __iter__(self): + for child in self.children: + yield child + + def __getitem__(self, key): + for child in self.children: + for index, item in enumerate(child.value): + if item == key: + if index + 1 >= len(child.value): + return None + else: + return child.value[index + 1 :] + raise IndexError("key {} not found".format(key)) + + def __repr__(self): + return str(" ".join(self.value)) + + @property + def name(self): + return self.value[1] + + +class Bvh: + """Abstraction for bvh.""" + + def __init__(self, data: str, backend: Optional[str] = "graph"): + """ + Args: + data (str): bvh file content + """ + self.data = data + self.root = BvhNode() + self.frames = [] + self.backend = backend + self.tokenize() + if self.backend == "np": + # cache important info for quick access later + self.build_data_array() + elif self.backend == "graph": + pass + else: + raise ValueError(f"Unknown backend for BVH loading: {backend}") + + def build_data_array(self): + joints = self.get_joints() + self.joint2idx = dict() + self.joint2channels = dict() + cur_idx = 0 + for joint in joints: + self.joint2idx[joint.value[1]] = cur_idx + cur_idx += int(joint["CHANNELS"][0]) + self.joint2channels[joint.value[1]] = joint["CHANNELS"][1:] + self.np_data_array = np.array(self.frames, dtype=np.float32) + + def tokenize(self): + first_round = [] + accumulator = "" + for char in self.data: + if char not in ("\n", "\r"): + accumulator += char + elif accumulator: + first_round.append(re.split("\\s+", accumulator.strip())) + accumulator = "" + node_stack = [self.root] + frame_time_found = False + node = None + for item in first_round: + if frame_time_found: + self.frames.append(item) + continue + key = item[0] + if key == "{": + node_stack.append(node) + elif key == "}": + node_stack.pop() + else: + node = BvhNode(item) + # print("new node: ", node, "\nparent: ", node_stack[-1]) + node_stack[-1].add_child(node) + if item[0] == "Frame" and item[1] == "Time:": + frame_time_found = True + + def search(self, *items): + found_nodes = [] + + def check_children(node): + if len(node.value) >= len(items): + failed = False + for index, item in enumerate(items): + if node.value[index] != item: + failed = True + break + if not failed: + found_nodes.append(node) + for child in node: + check_children(child) + + check_children(self.root) + return found_nodes + + def get_joints(self): + joints = [] + + def iterate_joints(joint): + joints.append(joint) + for child in joint.filter("JOINT"): + iterate_joints(child) + + iterate_joints(next(self.root.filter("ROOT"))) + return joints + + def get_joints_names(self): + joints = [] + + def iterate_joints(joint): + joints.append(joint.value[1]) + for child in joint.filter("JOINT"): + iterate_joints(child) + + iterate_joints(next(self.root.filter("ROOT"))) + return joints + + def joint_direct_children(self, name): + joint = self.get_joint(name) + return [child for child in joint.filter("JOINT")] + + def get_joint_index(self, name): + return self.get_joints().index(self.get_joint(name)) + + def get_joint(self, name): + found = self.search("ROOT", name) + if not found: + found = self.search("JOINT", name) + if found: + return found[0] + raise LookupError("joint not found") + + def joint_offset(self, name, idx=[0, 1, 2]): + joint = self.get_joint(name) + offset = joint["OFFSET"] + if len(offset) < max(idx): + return None + return (float(offset[idx[0]]), float(offset[idx[1]]), float(offset[idx[2]])) + + def joint_offset_rot(self, name): + return self.joint_offset(name, idx=[3, 4, 5]) + + def joint_channels(self, name): + if self.backend == "np": + return self.joint2channels[name] + else: + joint = self.get_joint(name) + return joint["CHANNELS"][1:] + + def get_joint_channels_index(self, joint_name): + if self.backend == "np": + return self.joint2idx[joint_name] + else: + index = 0 + for joint in self.get_joints(): + if joint.value[1] == joint_name: + return index + index += int(joint["CHANNELS"][0]) + raise LookupError("joint not found") + + def get_joint_channel_index(self, joint, channel): + channels = self.joint_channels(joint) + if channel in channels: + channel_index = channels.index(channel) + else: + raise ValueError(f"Channel {channel} not found in {channels}") + return channel_index + + def frame_joint_channel(self, frame_index, joint, channel, value=None): + """Get single frame data for on specific joint and one specific channel (e.g. Xrotation).""" + joint_index = self.get_joint_channels_index(joint) + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + return value + if self.backend == "np": + return self.np_data_array[frame_index, joint_index + channel_index] + else: + return float(self.frames[frame_index][joint_index + channel_index]) + + def frame_joint_channels(self, frame_index, joint, channels, value=None): + """Get single frame data for on specific joint from multiple specific channels (e.g. + Xrotation, Yrotation, Zrotation).""" + values = [] + joint_index = self.get_joint_channels_index(joint) + if self.backend == "np": + channel_idx = [self.get_joint_channel_index(joint, channel) for channel in channels] + channel_idx = np.array(channel_idx) + joint_index + values = self.np_data_array[frame_index, channel_idx] + else: + for channel in channels: + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + values.append(value) + else: + values.append(float(self.frames[frame_index][joint_index + channel_index])) + return values + + def frames_joint_channels(self, joint, channels, value=None): + """Get all frame data for one joint from multiple channels (e.g. Xrotation, Yrotation, + Zrotation).""" + joint_index = self.get_joint_channels_index(joint) + if self.backend == "np": + channel_idx = [self.get_joint_channel_index(joint, channel) for channel in channels] + channel_idx = np.array(channel_idx) + joint_index + all_frames = self.np_data_array[:, channel_idx] + else: + all_frames = [] + for frame in self.frames: + values = [] + for channel in channels: + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + values.append(value) + else: + values.append(float(frame[joint_index + channel_index])) + all_frames.append(values) + return all_frames + + def frames_joints_channels(self, joint_names, channels): + """Get all frames for all specified joints with one specified set of channels.""" + if self.backend != "np": + raise NotImplementedError("Only np backend is supported for this function") + joint_indices = [(joint_name, self.joint2idx[joint_name]) for joint_name in joint_names] + data_indices = [] + for joint_name, joint_idx in joint_indices: + channel_indices = [self.get_joint_channel_index(joint_name, channel) for channel in channels] + data_indices.extend([joint_idx + channel_idx for channel_idx in channel_indices]) + all_frames = self.np_data_array[:, data_indices] + all_frames = all_frames.reshape(-1, len(joint_names), len(channels)) + return all_frames + + def joint_parent(self, name): + joint = self.get_joint(name) + if joint.parent == self.root: + return None + return joint.parent + + def joint_parent_index(self, name): + joint = self.get_joint(name) + if joint.parent == self.root: + return -1 + return self.get_joints().index(joint.parent) + + @property + def nframes(self): + try: + return int(next(self.root.filter("Frames:")).value[1]) + except StopIteration: + raise LookupError("number of frames not found") + + @property + def frame_time(self): + try: + return float(next(self.root.filter("Frame")).value[2]) + except StopIteration: + raise LookupError("frame time not found") diff --git a/ardy/data_processing/skeleton.py b/ardy/data_processing/skeleton.py new file mode 100644 index 0000000000000000000000000000000000000000..9cd9b780eb875f4e4de8ef5acff666603b985e42 --- /dev/null +++ b/ardy/data_processing/skeleton.py @@ -0,0 +1,475 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional, Tuple + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from ardy.assets import skeleton_asset_path +from ardy.data_processing.bvh import Bvh + +_DEFAULT_G1_XML = str(skeleton_asset_path("g1skel34", "xml", "g1.xml")) + + +class Bone: + """Abstraction for a bone in character skeleton.""" + + def __init__(self): + # original bone info + self.id = None + self.name = None + self.orient = np.identity(3) + self.dof_index = [] + self.channels = [] # bvh only + self.lb = [] + self.ub = [] + self.parent = None + self.child = [] + + # asf specific + self.dir = np.zeros(3) + self.len = 0 + # bvh specific + self.offset = np.zeros(3) # default offset for position + self.offset_rot = None # rotation for custom nv bvh + + # inferred info + self.pos = np.zeros(3) + self.end = np.zeros(3) + + def __repr__(self): + return f"{self.name}" + + +class Skeleton: + """Abstraction for character skeleton.""" + + def __init__(self): + self.bones = [] + self.name2bone = {} + self.mass_scale = 1.0 + self.len_scale = 1.0 + self.dof_name = ["x", "y", "z"] + self.root = None + + def get_bones_names(self): + return [x.name for x in self.bones] + + def get_parent_indices(self): + parent_indices = [-1] * len(self.bones) + for bone in self.bones: + if bone.parent: + parent_indices[bone.id] = bone.parent.id + return parent_indices + + def get_neutral_joints(self): + joints = [] + for bone in self.bones: + joints.append(bone.pos) + joints = np.stack(joints, axis=0) + return joints + + def load_from_bvh(self, fname, exclude_bones=None, spec_channels=None): + if exclude_bones is None: + exclude_bones = {} + if spec_channels is None: + spec_channels = dict() + with open(fname) as f: + mocap = Bvh(f.read()) + + joint_names = list( + filter( + lambda x: all([t not in x for t in exclude_bones]), + mocap.get_joints_names(), + ) + ) + dof_ind = {"x": 0, "y": 1, "z": 2} + self.len_scale = 1.0 + self.root = Bone() + self.root.id = 0 + self.root.name = joint_names[0] + self.root.channels = mocap.joint_channels(self.root.name) + self.root.offset = np.array(mocap.joint_offset(self.root.name)) * self.len_scale + self.root.offset_rot = mocap.joint_offset_rot(self.root.name) + if self.root.offset_rot is not None: + self.root.offset_rot = np.array(self.root.offset_rot) + self.name2bone[self.root.name] = self.root + self.bones.append(self.root) + for i, joint in enumerate(joint_names[1:]): + bone = Bone() + bone.id = i + 1 + bone.name = joint + bone.channels = spec_channels[joint] if joint in spec_channels.keys() else mocap.joint_channels(joint) + bone.dof_index = [dof_ind[x[0].lower()] for x in bone.channels] + bone.offset = np.array(mocap.joint_offset(joint)) * self.len_scale + bone.offset_rot = mocap.joint_offset_rot(joint) + if bone.offset_rot is not None: + bone.offset_rot = np.array(bone.offset_rot) + bone.lb = [-180.0] * 3 + bone.ub = [180.0] * 3 + self.bones.append(bone) + self.name2bone[joint] = bone + + # for bone in self.bones: + # print(bone.name, bone.channels, bone.offset) + + for bone in self.bones[1:]: + parent_name = mocap.joint_parent(bone.name).name + if parent_name in self.name2bone.keys(): + bone_p = self.name2bone[parent_name] + bone_p.child.append(bone) + bone.parent = bone_p + + self.forward_bvh(self.root) + for bone in self.bones: + if len(bone.child) == 0: + child_vals = [str(node) for node in mocap.get_joint(bone.name).children] + if "End Site" in child_vals: + end_site_idx = child_vals.index("End Site") + end_site_offset = mocap.get_joint(bone.name).children[end_site_idx]["OFFSET"] + bone.end = bone.pos + np.array([float(x) for x in end_site_offset]) * self.len_scale + else: + pass + else: + bone.end = sum([bone_c.pos for bone_c in bone.child]) / len(bone.child) + + def load_from_xml(self, fname, exclude_bones=None, spec_channels=None): + if exclude_bones is None: + exclude_bones = {} + if spec_channels is None: + spec_channels = dict() + + import xml.etree.ElementTree as ET + + tree = ET.parse(fname) + root = tree.getroot() + + def map_bone_name(body_name): + if body_name == "torso_link": + return "waist_pitch_skel" + elif "_link" in body_name: + return body_name.replace("_link", "_skel") + else: + return body_name + "_skel" + + bodies = root.find("worldbody").findall(".//body") + name_to_body = {map_bone_name(body.get("name")): body for body in bodies} + joint_names = [body.get("name") for body in bodies] + + R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True) + x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True) + mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward + + def mujoco_joint_translation(body_name): + body = name_to_body[body_name] + if "pos" in body.attrib: + result = np.array([100 * float(x) for x in body.get("pos").strip().split(" ")]) + return mujoco_to_ardy.apply(result) + else: + return np.zeros(3) + + parent_map = { + map_bone_name(child.get("name")): map_bone_name(parent.get("name")) + for parent in root.iter() + if parent.get("name") + for child in parent + if child.get("name") + } + + self.len_scale = 1.0 + self.root = Bone() + self.root.id = 0 + self.root.name = map_bone_name(joint_names[0]) + self.root.offset = mujoco_joint_translation(self.root.name) + self.root.offset_rot = np.eye(3) + + self.name2bone[self.root.name] = self.root + self.bones.append(self.root) + for i, joint in enumerate(joint_names[1:]): + bone = Bone() + bone.id = i + 1 + bone.name = map_bone_name(joint) + bone.offset = mujoco_joint_translation(bone.name) + bone.offset_rot = np.eye(3) + self.bones.append(bone) + self.name2bone[bone.name] = bone + + for bone in self.bones[1:]: + parent_name = parent_map[bone.name] + if parent_name in self.name2bone.keys(): + bone_p = self.name2bone[parent_name] + bone_p.child.append(bone) + bone.parent = bone_p + + self.forward_bvh(self.root) + + def forward_bvh(self, bone): + if bone.parent: + bone.pos = bone.parent.pos + bone.offset + else: + bone.pos = bone.offset + for bone_c in bone.child: + self.forward_bvh(bone_c) + + +class SkeletonG1(Skeleton): + """Add functionality specific to the G1 including loading from XML specific to G1 and adding + dummy joints.""" + + def load_from_xml(self, fname, exclude_bones=None, spec_channels=None): + if exclude_bones is None: + exclude_bones = {} + if spec_channels is None: + spec_channels = dict() + + import xml.etree.ElementTree as ET + + tree = ET.parse(fname) + root = tree.getroot() + + def map_bone_name(body_name): + if body_name == "torso_link": + return "waist_pitch_skel" + elif "_link" in body_name: + return body_name.replace("_link", "_skel") + else: + return body_name + "_skel" + + bodies = root.find("worldbody").findall(".//body") + name_to_body = {map_bone_name(body.get("name")): body for body in bodies} + joint_names = [body.get("name") for body in bodies] + + R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True) + x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True) + mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward + + def mujoco_joint_translation(body_name): + body = name_to_body[body_name] + if "pos" in body.attrib: + result = np.array([100 * float(x) for x in body.get("pos").strip().split(" ")]) + return mujoco_to_ardy.apply(result) + else: + return np.zeros(3) + + parent_map = { + map_bone_name(child.get("name")): map_bone_name(parent.get("name")) + for parent in root.iter() + if parent.get("name") + for child in parent + if child.get("name") + } + + self.len_scale = 1.0 + self.root = Bone() + self.root.id = 0 + self.root.name = map_bone_name(joint_names[0]) + self.root.offset = mujoco_joint_translation(self.root.name) + self.root.offset_rot = np.eye(3) + + self.name2bone[self.root.name] = self.root + self.bones.append(self.root) + for i, joint in enumerate(joint_names[1:]): + bone = Bone() + bone.id = i + 1 + bone.name = map_bone_name(joint) + bone.offset = mujoco_joint_translation(bone.name) + bone.offset_rot = np.eye(3) + self.bones.append(bone) + self.name2bone[bone.name] = bone + + for bone in self.bones[1:]: + parent_name = parent_map[bone.name] + if parent_name in self.name2bone.keys(): + bone_p = self.name2bone[parent_name] + bone_p.child.append(bone) + bone.parent = bone_p + + self.forward_bvh(self.root) + + def add_joint(self, parent_name, bone_name, offset): + new_bone = Bone() + new_bone.name = bone_name + new_bone.parent = next((x for x in self.bones if x.name == parent_name)) + new_bone.parent.child.append(new_bone) + new_bone.offset = offset + new_bone.offset_rot = np.eye(3) + + parent_index = next((i for i, x in enumerate(self.bones) if x.name == parent_name)) + self.bones.insert(parent_index + 1, new_bone) + self.name2bone[bone_name] = new_bone + + +def load_bvh_animation( + fname: str, + skeleton: Skeleton, + rot_order: Optional[str] = "native", + backend: Optional[str] = "np", + return_quat: Optional[bool] = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Given a bvh file for one motion sequence, and Skeleton object containing the neutral + skeleton, return root translation and joint rotation matrics for the motion sequence. + + Args: + fname (str): full path of the bvh file + skeleton (Skeleton): Skeleton object initialized by the bvh file, containing the neutral skeleton + rot_order (str): one of "native", "ZXY", "YXZ", .... + return_quat (bool): if True, return joint rotation quaternions instead of rotation matrices + Returns: + torch.Tensor: [T, 3] root translations + torch.Tensor: [T, J, 3, 3] joint rotation matrix + """ + with open(fname) as f: + mocap = Bvh(f.read(), backend=backend) + + # assume all joints are same ordering, load in with native ordering + root_channels = mocap.joint_channels(skeleton.root.name) + pos_channels = [channel for channel in root_channels if channel.endswith("position")] + rot_channels = [channel for channel in root_channels if channel.endswith("rotation")] + + root_trans = np.array(mocap.frames_joint_channels(skeleton.root.name, pos_channels)) + + if backend == "np": + # NOTE: assumes rot channel ordering is the same for all joints + joint_eulers = mocap.frames_joints_channels(skeleton.get_bones_names(), rot_channels) + joint_eulers = np.deg2rad(joint_eulers) + elif backend == "graph": + joint_eulers = [] + for bone in skeleton.bones: + bone_channels = mocap.joint_channels(bone.name) + bone_rot_channels = [channel for channel in bone_channels if channel.endswith("rotation")] + assert bone_rot_channels == rot_channels, "Rotation channel ordering is not consistent across joints!" + # use native rotation order + euler = np.deg2rad(np.array(mocap.frames_joint_channels(bone.name, rot_channels))) + joint_eulers.append(euler) + joint_eulers = np.stack(joint_eulers, axis=1) + else: + raise ValueError(f"Unknown backend for BVH loading: {backend}") + + if rot_order == "native": + rot_order = "" + for axis in rot_channels: + rot_order += axis[0] + else: + # need to reorder dims + ordered_joint_eulers = [] + for axis in rot_order: + i = rot_channels.index(axis + "rotation") + ordered_joint_eulers.append(joint_eulers[..., i]) + joint_eulers = np.stack(ordered_joint_eulers, axis=-1) + + rotations = Rotation.from_euler(rot_order, joint_eulers.reshape(-1, 3)) + if return_quat: + joint_rots = rotations.as_quat(scalar_first=True).reshape(joint_eulers.shape[:-1] + (4,)) + else: + joint_rots = rotations.as_matrix().reshape(joint_eulers.shape[:-1] + (3, 3)) + + return root_trans, joint_rots + + +def load_csv_g1_animation( + fname: str, + skeleton: Skeleton, + csv_format: str = "retargeted_from_gmr", + xml_path: str = _DEFAULT_G1_XML, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Load a csv file that's in the format of mujoco qpos (32 joints, dof 29) return root + translation and joint rotation matrics for the motion sequence. + + Args: + fname (str): full path of the csv file + skeleton (Skeleton): Skeleton object initialized by the bvh file, containing the neutral skeleton + csv_format (str): format of the csv file ["retargeted_from_gmr", "retargeted_from_unitree"] + xml_path (str): full path of the xml file for the skeleton + Returns: + torch.Tensor: [T, 3] root translations + torch.Tensor: [T, J, 3, 3] joint rotation matrix + """ + assert csv_format in ["retargeted_from_gmr", "retargeted_from_unitree"] + import xml.etree.ElementTree as ET + + with open(fname) as f: + csv_raw_data = f.readlines() + if csv_format == "retargeted_from_gmr": + csv_raw_data = np.array( + [np.array([float(i) for i in row.split(",")[1:]]) for row in csv_raw_data[1:]] + ) # drop the title row, and the frame column (first column) + root_trans = csv_raw_data[:, :3] # cm since it will be converted back to m outside this func + xyz = csv_raw_data[:, 3:6] + root_rot = Rotation.from_euler("xyz", xyz, degrees=True) + + elif csv_format == "retargeted_from_unitree": + csv_raw_data = np.array([np.array([float(i) for i in row.split(",")]) for row in csv_raw_data]) + root_trans = ( + np.array(csv_raw_data[:, :3]) * 100 + ) # m -> cm since it will be converted back to m outside this func + root_rot = Rotation.from_quat(csv_raw_data[:, 3:7], scalar_first=False) # [x, y, z, w] is the order in csv file + + # there are two transformations from the mujoco coordinate to the ardy coordinate (both are right hand system): + # 1) the mujoco is x facing, and ardy is y facing in the mujoco coordinate (or z facing in the ardy system). + # 2) the source csv uses z up, and the ardy uses y up; + R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True) + x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True) + mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward + + # for root transform, we do the similar transformation for the root rotations. + root_trans = mujoco_to_ardy.apply(root_trans) + root_rot = mujoco_to_ardy * root_rot * mujoco_to_ardy.inv() + + # get the joint rotations from the csv file + rotations = Rotation.identity(csv_raw_data.shape[0] * len(skeleton.bones)) + joint_rots = rotations.as_matrix().reshape(csv_raw_data.shape[0], len(skeleton.bones), 3, 3) + joint_rots[:, 0] = root_rot.as_matrix() # the rotation of pelvis + + # load the mujoco xml file for joint axis; find all joint, get name and axis + tree = ET.parse(xml_path) + root = tree.getroot() + + xml_classes = [x for x in tree.findall(".//default") if "class" in x.attrib] + joint_axes = dict() + for xml_class in xml_classes: + j = xml_class.findall("joint") + if j: + joint_axes[xml_class.get("class")] = j[0].get("axis") + + parent_map = {child: parent for parent in root.iter() for child in parent} + + for joint_id_in_csv, joint in enumerate(root.find("worldbody").findall(".//joint")): # skip the base joint + # the order in the xml file is the order of the data shows up; note that in the xml file, the joint ends with + # "_joint", but the skeleton.get_bones_names() joints end with "_skel" + assert joint.get("name").endswith("_joint") + joint_name_in_skeleton = joint.get("name").replace("_joint", "_skel") + assert joint_name_in_skeleton in skeleton.get_bones_names() + axis_values = [float(x) for x in (joint.get("axis") or joint_axes[joint.get("class")]).split(" ")] + assert sum(axis_values) == 1 and all(x in [0, 1] for x in axis_values), ( + f"Invalid axis: {axis_values}: one and only one of the axis values should be 1; others should be 0s." + ) + + # the mapped axis in the ardy's g1 skeleton space is calculated as bones_axis = mujoco_to_ardy.apply(axis_values) + # [1, 0, 0] -> [0, 0, 1]; [0, 1, 0] -> [1, 0, 0]; [0, 0, 1] -> [0, 1, 0] + axis_in_ardy = ["x", "y", "z"][np.argmax(axis_values)] + if csv_format == "retargeted_from_gmr": + joint_dof = csv_raw_data[:, joint_id_in_csv + 6] * np.pi / 180 + elif csv_format == "retargeted_from_unitree": + joint_dof = csv_raw_data[:, joint_id_in_csv + 7] + else: + raise ValueError(f"Unknown csv format: {csv_format}") + + # Single-axis sequence needs angles of shape (N, 1) under SciPy's array-API + # from_euler, which reads the last axis as the axis count (not a bare (N,)). + joint_rot = Rotation.from_euler(axis_in_ardy, joint_dof[:, None], degrees=False) + + body = parent_map[joint] + if "quat" in body.attrib: + joint_extra_rot = Rotation.from_quat( + [float(x) for x in body.get("quat").strip().split(" ")], + scalar_first=True, + ) + joint_rot = joint_extra_rot * joint_rot + + joint_rots[:, skeleton.get_bones_names().index(joint_name_in_skeleton)] = ( + mujoco_to_ardy * joint_rot * mujoco_to_ardy.inv() + ).as_matrix() + + return root_trans, joint_rots diff --git a/ardy/exports/__init__.py b/ardy/exports/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..21558f429df2d353daeb738fbb7e59b36cb3bbae --- /dev/null +++ b/ardy/exports/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exporters for external motion formats (e.g. BVH and MuJoCo qpos CSV).""" + +from .bvh import write_bvh + +__all__ = ["write_bvh"] diff --git a/ardy/exports/bvh.py b/ardy/exports/bvh.py new file mode 100644 index 0000000000000000000000000000000000000000..99b7a5d77747a14cc49ff24046dc2506f5d67011 --- /dev/null +++ b/ardy/exports/bvh.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 +# SPDX-License-Identifier: Apache-2.0 +"""Small, dependency-light BVH exporter for ARDY motion outputs.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +import numpy as np +from scipy.spatial.transform import Rotation + + +def _children_from_parents(parents: Sequence[int]) -> list[list[int]]: + children = [[] for _ in parents] + roots = [] + for joint, parent in enumerate(parents): + if parent < 0: + roots.append(joint) + else: + children[parent].append(joint) + if len(roots) != 1: + raise ValueError(f"BVH requires exactly one root joint; found {len(roots)}") + return children + + +def _joint_offsets(rest_positions: np.ndarray, parents: Sequence[int]) -> np.ndarray: + offsets = np.asarray(rest_positions, dtype=np.float64).copy() + for joint, parent in enumerate(parents): + if parent < 0: + offsets[joint] = 0.0 + else: + offsets[joint] -= rest_positions[parent] + return offsets + + +def _leaf_offset(joint: int, offsets: np.ndarray, parents: Sequence[int]) -> np.ndarray: + direction = offsets[joint] + if np.linalg.norm(direction) < 1e-8: + parent = parents[joint] + direction = offsets[parent] if parent >= 0 else np.array([0.0, 0.05, 0.0]) + norm = np.linalg.norm(direction) + if norm < 1e-8: + return np.array([0.0, 0.05, 0.0]) + return direction / norm * 0.05 + + +def _fmt_vec(values: np.ndarray) -> str: + return " ".join(f"{float(value):.7f}" for value in values) + + +def write_bvh( + path: str | Path, + *, + joint_names: Sequence[str], + parents: Sequence[int], + rest_positions: np.ndarray, + local_rot_mats: np.ndarray, + root_positions: np.ndarray, + fps: float, +) -> Path: + """Write local joint rotations and root translation to a Blender-readable BVH. + + ARDY uses a Y-up metric coordinate system. The exported BVH preserves those + units and uses Z-X-Y Euler channels, which Blender's BVH importer supports. + """ + path = Path(path) + names = [str(name).replace(" ", "_") for name in joint_names] + parents = [int(parent) for parent in parents] + rotations = np.asarray(local_rot_mats, dtype=np.float64) + roots = np.asarray(root_positions, dtype=np.float64) + rest = np.asarray(rest_positions, dtype=np.float64) + + joint_count = len(names) + if rotations.ndim != 4 or rotations.shape[1:] != (joint_count, 3, 3): + raise ValueError( + "local_rot_mats must have shape [frames, joints, 3, 3]; " + f"received {rotations.shape}" + ) + if roots.shape != (rotations.shape[0], 3): + raise ValueError(f"root_positions must have shape [frames, 3]; received {roots.shape}") + if rest.shape != (joint_count, 3): + raise ValueError(f"rest_positions must have shape [joints, 3]; received {rest.shape}") + if fps <= 0: + raise ValueError("fps must be greater than zero") + + children = _children_from_parents(parents) + root = parents.index(-1) + offsets = _joint_offsets(rest, parents) + channel_order: list[int] = [] + lines = ["HIERARCHY"] + + def add_joint(joint: int, depth: int, is_root: bool = False) -> None: + indent = " " * depth + label = "ROOT" if is_root else "JOINT" + lines.append(f"{indent}{label} {names[joint]}") + lines.append(f"{indent}{{") + lines.append(f"{indent} OFFSET {_fmt_vec(offsets[joint])}") + if is_root: + lines.append( + f"{indent} CHANNELS 6 Xposition Yposition Zposition " + "Zrotation Xrotation Yrotation" + ) + else: + lines.append(f"{indent} CHANNELS 3 Zrotation Xrotation Yrotation") + channel_order.append(joint) + for child in children[joint]: + add_joint(child, depth + 1) + if not children[joint]: + lines.append(f"{indent} End Site") + lines.append(f"{indent} {{") + lines.append(f"{indent} OFFSET {_fmt_vec(_leaf_offset(joint, offsets, parents))}") + lines.append(f"{indent} }}") + lines.append(f"{indent}}}") + + add_joint(root, 0, is_root=True) + lines.extend( + [ + "MOTION", + f"Frames: {rotations.shape[0]}", + f"Frame Time: {1.0 / float(fps):.9f}", + ] + ) + + # scipy uses intrinsic rotations for uppercase sequences, matching BVH's + # channel interpretation. + eulers = Rotation.from_matrix(rotations.reshape(-1, 3, 3)).as_euler( + "ZXY", degrees=True + ) + eulers = eulers.reshape(rotations.shape[0], joint_count, 3) + for frame in range(rotations.shape[0]): + values: list[float] = [] + for index, joint in enumerate(channel_order): + if index == 0: + values.extend(roots[frame].tolist()) + values.extend(eulers[frame, joint].tolist()) + lines.append(" ".join(f"{value:.7f}" for value in values)) + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path diff --git a/ardy/exports/mujoco.py b/ardy/exports/mujoco.py new file mode 100644 index 0000000000000000000000000000000000000000..b83f0dc06fb3bb11d1820849b8e4dcf6c1078c64 --- /dev/null +++ b/ardy/exports/mujoco.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import xml.etree.ElementTree as ET +from typing import Optional + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from ardy.assets import skeleton_asset_path +from ardy.geometry import matrix_to_quaternion +from ardy.skeleton import SkeletonBase +from ardy.tools import ensure_batched, to_numpy, to_torch + +# Default G1 mujoco XML ships in the ardy package skeleton assets. +_DEFAULT_G1_XML = str(skeleton_asset_path("g1skel34", "xml", "g1.xml")) + + +class MujocoQposConverter(torch.nn.Module): + """Fast batch converter from our dictionary format to mujoco qpos with precomputed transforms. + + In mujoco, the coordination is z up and x forward, right handed + + features (30 joints): + root (pelvis, 7 = translation + rotation) + 29 dof joints (29) + + In ardy, the coordinate system is y up and z forward, right handed + features (34 joints): + root (pelvis) + (34 - 1) joints; among these joints, 4 are end-effector joints added by ardy. + """ + + def __init__( + self, + input_skeleton: SkeletonBase, + xml_path: str = _DEFAULT_G1_XML, + dead_joint_rotation_scheme: str = "dummy", + ): + """Initialize converter with precomputed transforms. + + Args: + xml_path: Path to the mujoco XML file containing joint definitions + dead_joint_rotation_scheme: Scheme for handling dead joints (end-effectors joints); + if "dummy", the dead joints's global rotations are set to identity matrix; + if "parent", the dead joints's global rotations are set to the parent's rotation. + """ + super().__init__() + self.xml_path = xml_path + self.skeleton = input_skeleton + self._prepare_transforms() + self._subtree_joints = {} + self._dead_joint_rotation_scheme = dead_joint_rotation_scheme + + def _prepare_transforms(self): + """Precompute all necessary transforms for efficient batch processing.""" + # Define coordinate transformations between mujoco and ardy space + # 1) R_zup_to_yup: rotation around x-axis by -90 degrees + # 2) x_forward_to_y_forward: rotation around z-axis by -90 degrees + # Combined transformation matrix: mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward + self.mujoco_to_ardy_matrix = torch.tensor( + [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=torch.float32 + ) + self.ardy_to_mujoco_matrix = self.mujoco_to_ardy_matrix.T # Inverse transformation: ardy_to_mujoco + + # Parse XML once and extract joint information + tree = ET.parse(self.xml_path) + root = tree.getroot() + + xml_classes = [x for x in tree.findall(".//default") if "class" in x.attrib] + joint_axes = dict() + for xml_class in xml_classes: + j = xml_class.findall("joint") + if j: + joint_axes[xml_class.get("class")] = j[0].get("axis") + + mujoco_hinge_joints = root.find("worldbody").findall(".//joint") # skip the base joint + self._mujoco_joint_axis_values_ardy_space = torch.zeros( + (len(mujoco_hinge_joints), 3), dtype=torch.float32 + ) # mujoco order but ardy space + self._mujoco_joint_axis_values_mujoco_space = torch.zeros( + (len(mujoco_hinge_joints), 3), dtype=torch.float32 + ) # mujoco order but mujoco space + + # for the below indices, mujoco_indices_to_ardy_indices does not include mujoco root (30 - 1 = 29 elements), + # while ardy_indices_to_mujoco_indices inclues the ardy root (32 elements). + self._mujoco_indices_to_ardy_indices = torch.zeros((len(mujoco_hinge_joints),), dtype=torch.int32) + self._ardy_indices_to_mujoco_indices = ( + torch.ones((self.skeleton.nbjoints,), dtype=torch.int32) * -1 + ) # -1 means not in the csv skeleton + + self._nb_joints_mujoco = len(mujoco_hinge_joints) + 1 + self._nb_joints_ardy = self.skeleton.nbjoints + self._mujoco_joint_including_root_parent_list = torch.full( + (len(mujoco_hinge_joints) + 1,), -1, dtype=torch.int32 + ) + self._mujoco_joint_including_root_list = ["pelvis_skel"] + + for joint_id_in_csv, joint in enumerate(mujoco_hinge_joints): + joint_name_in_skeleton = joint.get("name").replace("_joint", "_skel") + joint_parent_name_in_skeleton = self.skeleton.bone_parents[joint_name_in_skeleton] + + self._mujoco_joint_including_root_list.append(joint_name_in_skeleton) + self._mujoco_joint_including_root_parent_list[joint_id_in_csv + 1] = ( + self._mujoco_joint_including_root_list.index(joint_parent_name_in_skeleton) + ) + + joint_idx_in_ardy_skeleton = self.skeleton.bone_order_names.index(joint_name_in_skeleton) + axis_values = [float(x) for x in (joint.get("axis") or joint_axes[joint.get("class")]).split(" ")] + + # the mapped axis in ardy skeleton space is calculated as bones_axis = mujoco_to_ardy.apply(axis_values) + # [1, 0, 0] -> [0, 0, 1]; [0, 1, 0] -> [1, 0, 0]; [0, 0, 1] -> [0, 1, 0] + mujoco_joint_axis_mapping_ardy_space = [ + torch.tensor([0, 0, 1]), + torch.tensor([1, 0, 0]), + torch.tensor([0, 1, 0]), + ][np.argmax(axis_values)] + + self._mujoco_joint_axis_values_ardy_space[joint_id_in_csv] = mujoco_joint_axis_mapping_ardy_space + self._mujoco_joint_axis_values_mujoco_space[joint_id_in_csv] = torch.tensor(axis_values) + + self._mujoco_indices_to_ardy_indices[joint_id_in_csv] = joint_idx_in_ardy_skeleton + self._ardy_indices_to_mujoco_indices[joint_idx_in_ardy_skeleton] = joint_id_in_csv + 1 # +1 for the root + self._ardy_indices_to_mujoco_indices[0] = 0 # the root joint mapping + + # load the offset matrices from the xml + R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True) + x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True) + mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward + + self._rot_offsets_q2t = torch.zeros(len(self._ardy_indices_to_mujoco_indices), 3, 3, dtype=torch.float32) + self._rot_offsets_q2t[...] = torch.eye(3)[None] + + self._rot_offsets_f2q = torch.zeros(len(self._ardy_indices_to_mujoco_indices), 3, 3, dtype=torch.float32) + self._rot_offsets_f2q[...] = torch.eye(3)[None] + parent_map = {child: parent for parent in root.iter() for child in parent} + for i, joint in enumerate(mujoco_hinge_joints): + body = parent_map[joint] + if "quat" in body.attrib: + rot = Rotation.from_quat( + [float(x) for x in body.get("quat").strip().split(" ")], + scalar_first=True, + ) + idx = self._mujoco_indices_to_ardy_indices[i] + self._rot_offsets_q2t[idx] = torch.from_numpy(rot.as_matrix()) + rot = mujoco_to_ardy * rot * mujoco_to_ardy.inv() + self._rot_offsets_f2q[idx] = torch.from_numpy(rot.as_matrix().T) + + def dict_to_qpos( + self, + output: dict, + device: Optional[str] = None, + root_quat_w_first: bool = True, + numpy: bool = True, + ): + local_rot_mats = to_torch(output["local_rot_mats"], device) + root_positions = to_torch(output["root_positions"], device) + + qpos = self.to_qpos( + local_rot_mats, + root_positions, + root_quat_w_first=root_quat_w_first, + ) + if numpy: + qpos = to_numpy(qpos) + return qpos + + def save_csv(self, qpos: torch.Tensor | np.ndarray, csv_path): + # comment this + qpos = to_numpy(qpos) + shape = qpos.shape + if len(shape) == 2: + # only one motion: save it + np.savetxt(csv_path, qpos, delimiter=",") + if len(shape) == 3: + # batch of motions + if shape[0] == 1: + # if only one motion, just save it + np.savetxt(csv_path, qpos[0], delimiter=",") + else: + csv_path_base, ext = os.path.splitext(csv_path) + for i in range(shape[0]): + self.save_csv(qpos[i], csv_path_base + "_" + str(i).zfill(2) + ext) + + @ensure_batched(local_rot_mats=5, root_positions=3, lengths=1) + def to_qpos( + self, + local_rot_mats: torch.Tensor, + root_positions: torch.Tensor, + root_quat_w_first: bool = True, + ) -> torch.Tensor: + """Fast batch conversion from ARDY features to mujoco qpos format. + + Args: + local_rot_mats (torch.Tensor): [batch, numFrames, numJoints, 3, 3] + local joint rotation matrices in ARDY coordinates + root_positions (torch.Tensor): [batch, numFrames, 3] root joint + positions in ARDY coordinates + root_quat_w_first (bool): store the root quaternion as [w, x, y, z] + (mujoco convention) instead of [x, y, z, w] + + Returns: + torch.Tensor of shape [batch, numFrames, 36] containing mujoco qpos data: + - root_trans (3) + root_quat (4) + joint_dofs (29) = 36 columns + """ + + batch_size, num_frames, nb_joints = local_rot_mats.shape[:3] + device, dtype = local_rot_mats.device, local_rot_mats.dtype + + local_rot_mats = torch.matmul(self._rot_offsets_f2q.to(device), local_rot_mats) + + batch_size, num_frames = root_positions.shape[0], root_positions.shape[1] + + # Move precomputed matrices to the same device/dtype + ardy_to_mujoco_matrix = self.ardy_to_mujoco_matrix.to(device=device, dtype=dtype) + + # Initialize output tensor: [batch, numFrames, 36] + qpos = torch.zeros((batch_size, num_frames, 36), dtype=dtype, device=device) + + # Convert root translation: apply coordinate transformation + root_positions_mujoco = torch.matmul(ardy_to_mujoco_matrix[None, None, ...], root_positions[..., None]) + qpos[:, :, :3] = root_positions_mujoco.view(batch_size, num_frames, 3) + + # Convert root rotation: apply coordinate transformation to rotation matrix + root_rot = local_rot_mats[:, :, 0, :] # [batch, numFrames, 3, 3] + + # Apply coordinate transformation: R_mujoco = ardy_to_mujoco * R_ardy * ardy_to_mujoco^T + mujoco_to_ardy_matrix = ardy_to_mujoco_matrix.T + root_rot_mujoco = torch.matmul( + torch.matmul(ardy_to_mujoco_matrix[None, None, ...], root_rot), + mujoco_to_ardy_matrix[None, None, ...], + ) + root_rot_quat = matrix_to_quaternion(root_rot_mujoco) # [w, x, y, z] + if root_quat_w_first: + qpos[:, :, 3:7] = root_rot_quat[:, :, [0, 1, 2, 3]] # [w, x, y, z] + else: + qpos[:, :, 3:7] = root_rot_quat[:, :, [1, 2, 3, 0]] # [w, x, y, z] -> [x, y, z, w] + + # Convert joint DOFs using precomputed mappings + joint_rot_mujoco = local_rot_mats[ + :, :, self._mujoco_indices_to_ardy_indices, : + ] # mujoco joint order but ardy feature space + x_joint_dof = torch.atan2(joint_rot_mujoco[..., 2, 1], joint_rot_mujoco[..., 2, 2]) + y_joint_dof = torch.atan2(joint_rot_mujoco[..., 0, 2], joint_rot_mujoco[..., 0, 0]) + z_joint_dof = torch.atan2(joint_rot_mujoco[..., 1, 0], joint_rot_mujoco[..., 1, 1]) + xyz_joint_dofs = torch.stack([x_joint_dof, y_joint_dof, z_joint_dof], dim=-1) + joint_dofs = (xyz_joint_dofs * self._mujoco_joint_axis_values_ardy_space[None, None, :, :].to(device)).sum( + dim=-1 + ) + qpos[:, :, 7:] = joint_dofs + return qpos diff --git a/ardy/geometry.py b/ardy/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..9f0dced159d9c2efae8bf45e6ea2ec16712441d7 --- /dev/null +++ b/ardy/geometry.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn.functional as F + + +def angle_to_Y_rotation_matrix(angle): + cos, sin = torch.cos(angle), torch.sin(angle) + one, zero = torch.ones_like(angle), torch.zeros_like(angle) + mat = torch.stack((cos, zero, sin, zero, one, zero, -sin, zero, cos), -1) + mat = mat.reshape(angle.shape + (3, 3)) + return mat + + +def matrix_to_cont6d(matrix): + cont_6d = torch.concat([matrix[..., 0], matrix[..., 1]], dim=-1) + return cont_6d + + +def cont6d_to_matrix(cont6d): + assert cont6d.shape[-1] == 6, "The last dimension must be 6" + x_raw = cont6d[..., 0:3] + y_raw = cont6d[..., 3:6] + + x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True) + z = torch.cross(x, y_raw, dim=-1) + z = z / torch.norm(z, dim=-1, keepdim=True) + + y = torch.cross(z, x, dim=-1) + + x = x[..., None] + y = y[..., None] + z = z[..., None] + + mat = torch.cat([x, y, z], dim=-1) + return mat + + +def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert axis-angle to rotation matrix. + + Args: + axis_angle: (..., 3) axis-angle vectors (angle = norm, axis = normalized) + Returns: + rotmat: (..., 3, 3) rotation matrices + """ + eps = 1e-6 + angle = torch.norm(axis_angle, dim=-1, keepdim=True) # (..., 1) + axis = axis_angle / (angle + eps) + + x, y, z = axis.unbind(-1) + + zero = torch.zeros_like(x) + K = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).reshape(*axis.shape[:-1], 3, 3) + + eye = torch.eye(3, device=axis.device, dtype=axis.dtype) + eye = eye.expand(*axis.shape[:-1], 3, 3) + + sin = torch.sin(angle)[..., None] + cos = torch.cos(angle)[..., None] + + R = eye + sin * K + (1 - cos) * (K @ K) + return R + + +def matrix_to_axis_angle(R: torch.Tensor) -> torch.Tensor: + """Convert rotation matrix to axis-angle. + + Args: + R: (..., 3, 3) rotation matrices + Returns: + axis_angle: (..., 3) + """ + eps = 1e-6 + + trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + cos_angle = (trace - 1) / 2 + cos_angle = torch.clamp(cos_angle, -1 + eps, 1 - eps) + + angle = torch.acos(cos_angle) + + rx = R[..., 2, 1] - R[..., 1, 2] + ry = R[..., 0, 2] - R[..., 2, 0] + rz = R[..., 1, 0] - R[..., 0, 1] + axis = torch.stack([rx, ry, rz], dim=-1) + + sin_angle = torch.sin(angle) + axis = axis / (2 * sin_angle[..., None] + eps) + + return axis * angle[..., None] + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """Returns torch.sqrt(torch.max(0, x)) subgradient is zero where x is 0.""" + return torch.sqrt(x * (x > 0).to(x.dtype)) + + +def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(matrix.reshape(batch_dim + (9,)), dim=-1) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + quat_by_rijk = torch.stack( + [ + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + return ( + (F.one_hot(q_abs.argmax(dim=-1), num_classes=4)[..., None] * quat_candidates) + .sum(dim=-2) + .reshape(batch_dim + (4,)) + ) + + +def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: + """Convert rotations given as quaternions to rotation matrices. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) diff --git a/ardy/model/__init__.py b/ardy/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e5809193817bca24a5c05cc49c5db70e38aaa345 --- /dev/null +++ b/ardy/model/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ARDY model package: main model class, text encoders, and loading utilities.""" + +from .ardy_model import Ardy +from .llm2vec import LLM2VecEncoder +from .load_model import load_model +from .loading import ( + AVAILABLE_MODELS, + DEFAULT_MODEL, + DEFAULT_TEXT_ENCODER_URL, + MODEL_NAMES, + load_checkpoint_state_dict, +) + +# from .twostage_denoiser import TwostageDenoiser + +__all__ = [ + "Ardy", + "LLM2VecEncoder", + # "TwostageDenoiser", + "load_model", + "load_checkpoint_state_dict", + "AVAILABLE_MODELS", + "DEFAULT_MODEL", + "DEFAULT_TEXT_ENCODER_URL", + "MODEL_NAMES", +] diff --git a/ardy/model/ardy_model.py b/ardy/model/ardy_model.py new file mode 100644 index 0000000000000000000000000000000000000000..0f974c05ccb4151bd1af282ee416c8393d176042 --- /dev/null +++ b/ardy/model/ardy_model.py @@ -0,0 +1,808 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Any, List, Optional, Tuple, Union + +import numpy as np +import torch +from torch import nn +from tqdm.auto import tqdm + +from ardy.model.cfg import AutoLatentClassifierFreeGuidedModel +from ardy.model.diffusion import DDIMSampler, Diffusion +from ardy.model.latent_utils import HybridMotionConverter + +log = logging.getLogger(__name__) + + +def get_three_mask_from_len(history_len, generation_len, future_len, num_frames, device): + indices = torch.arange(num_frames, device=device)[None, :] # (1, num_frames) + history_mask = indices < history_len[:, None] + generation_start = history_len + generation_end = generation_start + generation_len + generation_mask = (indices < generation_end[:, None]) & (indices >= generation_start[:, None]) + future_mask = indices >= generation_end[:, None] + return history_mask, generation_mask, future_mask + + +def translate_normalized_root_motion(root_motion, translation, motion_rep): + """The y coords of the input translation are zero.""" + assert (translation[:, 1] == 0).all(), "the y coords of the input translation are zero" + + root_motion = motion_rep.global_root_stats.unnormalize(root_motion) + root_motion[:, :, motion_rep.slice_dict["root_pos"]] = ( + root_motion[:, :, motion_rep.slice_dict["root_pos"]] + translation[:, None, :] + ) + root_motion = motion_rep.global_root_stats.normalize(root_motion) + return root_motion + + +class Ardy(nn.Module): + """Helper class for test time.""" + + def __init__( + self, + denoiser: nn.Module, + autoencoder: nn.Module, + gen_horizon_len: int, + num_base_steps: int, + text_encoder: Optional[Any], + device: Optional[Union[str, torch.device]] = None, + cfg_type: Optional[str] = "regular", + ): + super().__init__() + + if cfg_type is None: + cfg_type = "nocfg" + # eval mode + self.denoiser = AutoLatentClassifierFreeGuidedModel(denoiser.eval(), cfg_type=cfg_type) + + self.motion_rep = denoiser.motion_rep + self.skeleton = self.motion_rep.skeleton + self.gen_horizon_len = gen_horizon_len + + self.autoencoder = autoencoder + + self.diffusion = Diffusion(num_base_steps=num_base_steps) + self.sampler = DDIMSampler(self.diffusion) + self.text_encoder = text_encoder + self.hybrid = HybridMotionConverter.from_model(self) + + self.num_frames_per_token = self.autoencoder.num_frames_per_token + + self.device = device + self.to(device) + + def train(self, mode: bool): + self.denoiser.train(mode) + return self + + def eval(self): + self.denoiser.eval() + return self + + def set_autoencoder(self, autoencoder: nn.Module) -> None: + """Replace the autoencoder everywhere it is referenced. + + ``self.autoencoder`` and ``self.hybrid.autoencoder`` are separate + references: the hybrid converter captured the autoencoder at + construction time (see ``HybridMotionConverter.from_model``), and + detokenization during generation goes through the hybrid converter. + Rebinding only ``self.autoencoder`` (e.g. to swap in a TRT decoder) + therefore has no effect on generation. Updating both keeps them in sync + so the swap actually takes effect. + """ + self.autoencoder = autoencoder + self.hybrid.autoencoder = autoencoder + + def compile_denoiser( + self, + backend: str = "inductor", + mode: str = "default", + dynamic: bool = False, + ) -> None: + """Compile the denoiser with torch.compile for faster inference. + + Must be called after loading checkpoint weights, as torch.compile wraps + the module in a way that may interfere with checkpoint loading. + + Args: + backend: torch.compile backend ("torch_tensorrt" or "inductor"). + mode: Compilation mode ("reduce-overhead" uses CUDA graphs). + dynamic: Whether to allow dynamic tensor shapes without recompilation. + """ + log.info( + "Compiling denoiser with torch.compile (backend=%s, mode=%s, dynamic=%s)...", + backend, + mode, + dynamic, + ) + compile_kwargs: dict = { + "backend": backend, + "mode": mode, + "dynamic": dynamic, + } + if backend == "torch_tensorrt": + try: + import torch_tensorrt # noqa: F401 — registers the backend + except (ModuleNotFoundError, OSError) as e: + log.warning( + "torch_tensorrt unavailable (%s), falling back to inductor backend.", + e, + ) + backend = "inductor" + compile_kwargs["backend"] = backend + else: + # torch_tensorrt uses options instead of mode + del compile_kwargs["mode"] + compile_kwargs["options"] = { + "enabled_precisions": {torch.float32}, + "min_block_size": 3, + "use_python_runtime": True, + # Workaround: torch_tensorrt cat converter crashes on + # FakeTensor frozen params (tries .numpy() on them). + # Force cat to run in PyTorch fallback. + "torch_executed_ops": {"torch.ops.aten.cat.default"}, + } + self.denoiser = torch.compile(self.denoiser, **compile_kwargs) + log.info("Denoiser compiled.") + + def warmup( + self, + num_tokens: int = 3, + num_text_tokens: int = 32, + num_iterations: int = 3, + ) -> None: + """Run the denoiser with dummy inputs to trigger compilation. + + Call this after compile_denoiser() to avoid compilation latency during + the first real inference call. Multiple iterations are used because + torch_tensorrt may refine its engine across the first few calls. + + Args: + num_tokens: Number of motion tokens for warmup (history + generation). + num_text_tokens: Number of text tokens for warmup. + num_iterations: Number of warmup forward passes. + """ + log.info("Warming up compiled denoiser (%d iterations)...", num_iterations) + device = self.device + # The dimension attributes live directly on the inner denoiser, reached + # via denoiser.model (the module wrapped by the CFG wrapper). Going through + # .model preserves attribute access even after torch.compile. + model = self.denoiser.model + nfpt = model.num_frames_per_token + num_frames = num_tokens * nfpt + dim_token = model.nframe_root_dim + model.latent_embedding_dim + motion_rep_dim = model.motion_rep.motion_rep_dim + + with torch.no_grad(): + for i in range(num_iterations): + self.denoiser( + # self.denoiser is the CFG wrapper; it requires the text and + # constraint guidance weights as its first two arguments (see + # denoising_step). Values are arbitrary for warmup. + cfg_weight_text=torch.ones(1, device=device), + cfg_weight_cstr=torch.zeros(1, device=device), + x=torch.randn(1, num_tokens, dim_token, device=device), + history_len=torch.tensor([nfpt], device=device), + generation_len=torch.tensor([num_frames - nfpt], device=device), + future_len=torch.tensor([0], device=device), + history_mask=torch.cat( + [ + torch.ones(1, nfpt, dtype=torch.bool, device=device), + torch.zeros(1, num_frames - nfpt, dtype=torch.bool, device=device), + ], + dim=1, + ), + generation_mask=torch.cat( + [ + torch.zeros(1, nfpt, dtype=torch.bool, device=device), + torch.ones(1, num_frames - nfpt, dtype=torch.bool, device=device), + ], + dim=1, + ), + future_mask=torch.zeros(1, num_frames, dtype=torch.bool, device=device), + history_token_mask=torch.cat( + [ + torch.ones(1, 1, dtype=torch.bool, device=device), + torch.zeros(1, num_tokens - 1, dtype=torch.bool, device=device), + ], + dim=1, + ), + generation_token_mask=torch.cat( + [ + torch.zeros(1, 1, dtype=torch.bool, device=device), + torch.ones(1, num_tokens - 1, dtype=torch.bool, device=device), + ], + dim=1, + ), + future_token_mask=torch.zeros(1, num_tokens, dtype=torch.bool, device=device), + text_feat=torch.randn(1, num_text_tokens, model.llm_shape[-1], device=device), + text_feat_pad_mask=torch.ones(1, num_text_tokens, dtype=torch.bool, device=device), + timesteps=torch.tensor([0], device=device), + first_heading_angle=torch.zeros(1, device=device, dtype=torch.float32), + motion_mask=torch.zeros(1, num_frames, motion_rep_dim, device=device), + observed_motion=torch.zeros(1, num_frames, motion_rep_dim, device=device), + ) + log.info("Warmup iteration %d/%d complete.", i + 1, num_iterations) + log.info("Warmup complete.") + + def denoising_step( + self, + x: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_pad_mask: torch.Tensor, + t: torch.Tensor, + first_heading_angle: Optional[torch.Tensor], + motion_mask: torch.Tensor, + observed_motion: torch.Tensor, + num_denoising_steps: torch.Tensor, + cfg_weight: Union[float, Tuple[float, float]], + target_motion: Optional[torch.Tensor] = None, + cfg_type: Optional[str] = None, + ) -> torch.Tensor: + """Single denoising step. + + Returns: + torch.Tensor: [B, T, D] noisy motion input to t-1 + """ + # subsample timesteps + # NOTE: do this at every step due to ONNX export, i.e. num_samp_stepsmay change dynamically when + # running onnx version so need to account for that. + num_denoising_steps = num_denoising_steps[0] + use_timesteps, map_tensor = self.diffusion.space_timesteps(num_denoising_steps) + self.diffusion.calc_diffusion_vars(use_timesteps) + + # first compute initial clean prediction from denoiser + t_map = map_tensor[t] + + # self.denoiser is the CFG wrapper (PyTorch or TRTCFGDenoiser); both take the + # text and constraint guidance weights as two tensors. This is the single + # place the combined cfg_weight (text, constraint) is turned into tensors. + if isinstance(cfg_weight, (tuple, list)): + w_text, w_cstr = cfg_weight + else: + w_text, w_cstr = cfg_weight, 0.0 + cfg_weight_text = torch.tensor([w_text], device=x.device, dtype=torch.float32) + cfg_weight_cstr = torch.tensor([w_cstr], device=x.device, dtype=torch.float32) + + with torch.inference_mode(): + token_seq_pred_clean = self.denoiser( + cfg_weight_text, + cfg_weight_cstr, + x, + history_len, + generation_len, + future_len, + history_mask, + generation_mask, + future_mask, + history_token_mask, + generation_token_mask, + future_token_mask, + text_feat, + text_pad_mask, + t_map, + first_heading_angle, + motion_mask, + observed_motion, + ) + + # sampler computes next step noisy motion + batch_size, num_token_dim = x.shape[0], x.shape[2] + num_generation_tokens = self.gen_horizon_len // self.num_frames_per_token + generation_token_t = x[generation_token_mask].reshape(batch_size, num_generation_tokens, num_token_dim) + generation_token_clean = token_seq_pred_clean[generation_token_mask].reshape( + batch_size, num_generation_tokens, num_token_dim + ) + generation_token_tm1 = self.sampler(generation_token_t, generation_token_clean, t) + + # Clone: x is still needed by the caller; the masked write below must + # not alias it. + xm1 = x.clone() + xm1[generation_token_mask] = generation_token_tm1.reshape(-1, num_token_dim) + return xm1 + + def _encode_text(self, texts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode text prompts into features and a padding mask.""" + device = self.device + log.info("Encoding text...") + text_feat, text_length = self.text_encoder(texts) + text_feat = text_feat.to(device) + + # handle empty string (set to zero) + empty_text_mask = [len(text.strip()) == 0 for text in texts] + text_feat[empty_text_mask] = 0 + + # Create the pad mask for the text + batch_size, maxlen = text_feat.shape[:2] + tensor_text_length = torch.tensor(text_length, device=device) + tensor_text_length[empty_text_mask] = 0 + text_pad_mask = torch.arange(maxlen, device=device).expand(batch_size, maxlen) < tensor_text_length[:, None] + return text_feat, text_pad_mask + + def _recenter_history( + self, + history_sequence: torch.Tensor, + center_frame_index: torch.Tensor, + requantize: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Recenter the hybrid history around ``center_frame_index``. + + Returns the recentered hybrid history, the center position, and the recentered (explicit) + root motion. + """ + root_motion, latent_body_motion = self.hybrid.get_root_and_latent_body_motion_from_hybrid( + history_sequence, + ) + new_root_motion, center_pos = self.motion_rep.recenter_root_motion( + root_motion, + center_frame_index, + is_normalized=True, + to_normalize=True, + return_center_pos=True, + ) + # quantize the latent body motion if the autoencoder uses quantization + if requantize and self.autoencoder.encode_with_quantization: + latent_body_motion = self.autoencoder.requantize( + latent_body_motion, + ) + # combine back to hybrid motion + history_sequence = self.hybrid.get_hybrid_motion_from_root_and_latent_body_motion( + new_root_motion, + latent_body_motion, + ) + return history_sequence, center_pos, new_root_motion + + def _encode_init_history( + self, + init_history_sequence: torch.Tensor, + batch_size: int, + crop_history_length: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode an explicit init history into a recentered hybrid sequence. + + Returns the hybrid history sequence, the initial global translation, and the first-frame + heading angle. + """ + init_history_len = init_history_sequence.shape[1] + history_sequence, history_pad_mask = self.hybrid.get_hybrid_motion_from_explicit( + motion=init_history_sequence, + motion_len=torch.ones(batch_size, device=self.device, dtype=torch.long) * init_history_len, + motion_pad_mask=torch.ones(batch_size, init_history_len, device=self.device, dtype=torch.bool), + ) + # transl_center_mode == "last_history": + center_frame_index = torch.ones(batch_size, device=self.device, dtype=torch.long) * ( + init_history_len - 1 + ) # [B,] + + history_sequence, center_pos, _ = self._recenter_history(history_sequence, center_frame_index, requantize=False) + # update global translation + global_transl = center_pos.clone() + + # set first frame heading angle + heading_angles = self.motion_rep.get_root_heading_angle(self.motion_rep.unnormalize(init_history_sequence)) + first_heading_angle = heading_angles[:, 0] + return history_sequence, global_transl, first_heading_angle + + def _generate_window( + self, + history_sequence: Optional[torch.Tensor], + global_transl: torch.Tensor, + history_start_frame: int, + history_end_frame: int, + total_frames: int, + text_feat: torch.Tensor, + text_pad_mask: torch.Tensor, + first_heading_angle: Optional[torch.Tensor], + motion_mask: Optional[torch.Tensor], + observed_motion: Optional[torch.Tensor], + num_denoising_steps: torch.Tensor, + cfg_weight: Union[float, Tuple[float, float]], + indices: List[int], + progress_bar=tqdm, + target_motion: Optional[torch.Tensor] = None, + cfg_type: Optional[str] = None, + ) -> torch.Tensor: + """Generate a single window of ``gen_horizon_len`` frames. + + Builds the history/generation/future masks and conditioning, denoises a block of pure noise, + and appends the newly generated tokens (in hybrid latent form) to ``history_sequence``. + Returns the updated history. + """ + device = self.device + batch_size = text_feat.shape[0] + num_frames_per_token = self.num_frames_per_token + latent_embedding_dim = self.denoiser.latent_embedding_dim + nframe_root_dim = self.denoiser.nframe_root_dim + gen_horizon_len = self.gen_horizon_len + num_generation_tokens = gen_horizon_len // num_frames_per_token + generation_len = torch.ones(batch_size, device=device, dtype=torch.long) * gen_horizon_len + + history_token_end = (history_end_frame - history_start_frame) // num_frames_per_token + generation_token_end = history_token_end + num_generation_tokens + history_len = torch.ones(batch_size, device=device, dtype=torch.long) * history_token_end * num_frames_per_token + future_len = torch.ones(batch_size, device=device, dtype=torch.long) * ( + total_frames - gen_horizon_len - history_end_frame + ) + + history_mask, generation_mask, future_mask = get_three_mask_from_len( + history_len, + generation_len, + future_len, + total_frames - history_start_frame, + device, + ) + # update observed motion roots with global translation + if motion_mask is not None: + cur_motion_mask = motion_mask[:, history_start_frame:] * ( + ~history_mask[:, :, None] + ) # only allow motion mask for non-history frames + cur_observed_root_motion = self.motion_rep.extract_root(observed_motion[:, history_start_frame:]) + cur_observed_body_motion = self.motion_rep.extract_body(observed_motion[:, history_start_frame:]) + + translated_observed_root_motion = translate_normalized_root_motion( + cur_observed_root_motion, -global_transl, self.motion_rep + ) + cur_observed_motion = self.motion_rep.concat_root_body( + translated_observed_root_motion, cur_observed_body_motion + ) + + cur_observed_motion = ( + cur_observed_motion * cur_motion_mask + ) # mask out the unobserved frames roots/joints that are non-zero due to translation + else: + cur_motion_mask = None + cur_observed_motion = None + + history_token_mask, generation_token_mask, future_token_mask = self.hybrid.convert_frame_mask_to_token_mask( + history_mask, + generation_mask, + future_mask, + cur_motion_mask, + ) + + # init pure noise x_T and x + shape = ( + batch_size, + num_generation_tokens, + nframe_root_dim + latent_embedding_dim, + ) + x_t = torch.randn(shape, device=device) + x = torch.zeros( + batch_size, + (total_frames - history_start_frame) // num_frames_per_token, + nframe_root_dim + latent_embedding_dim, + device=device, + ) + if history_token_end > 0: + x[:, :history_token_end] = history_sequence[:, -history_token_end:] + x[:, history_token_end:generation_token_end] = x_t + + for i in progress_bar(indices): + t = torch.tensor([i] * x_t.size(0), device=device) + with torch.no_grad(): + x = self.denoising_step( + x, + history_len, + generation_len, + future_len, + history_mask, + generation_mask, + future_mask, + history_token_mask, + generation_token_mask, + future_token_mask, + text_feat, + text_pad_mask, + t, + first_heading_angle, + cur_motion_mask, + cur_observed_motion, + num_denoising_steps, + cfg_weight, + target_motion, + cfg_type=cfg_type, + ) + + # update the full history sequence + if history_sequence is None: + history_sequence = x[:, :generation_token_end] + else: + history_sequence = torch.cat( + [ + history_sequence, + x[:, history_token_end:generation_token_end], + ], + dim=1, + ) + return history_sequence + + def __call__( + self, + texts: List[str], + num_frames: int, + num_denoising_steps: int, + pad_mask: torch.Tensor, + first_heading_angle: Optional[torch.Tensor], + motion_mask: torch.Tensor, + observed_motion: torch.Tensor, + target_motion: Optional[torch.Tensor] = None, + cfg_weight: Optional[float] = 2.0, + text_feat: Optional[torch.Tensor] = None, + text_pad_mask: Optional[torch.Tensor] = None, + cfg_type: Optional[str] = None, + progress_bar=tqdm, + return_text_embeddings_x_dict: bool = False, + crop_history_length: Optional[int] = None, + init_history_sequence: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Sample full denoising loop. + + Args: + texts (List[str]): batch of text prompts to use for sampling (if text_feat is not passed in) + """ + if text_feat is None: + assert text_pad_mask is None + text_feat, text_pad_mask = self._encode_text(texts) + + batch_size = text_feat.shape[0] + num_frames_per_token = self.num_frames_per_token + motion_pad_mask = pad_mask + motion_len = motion_pad_mask.sum(dim=-1) # [batch_size], exact length of the motion output + motion_len_pad = ( + torch.ceil(motion_len / num_frames_per_token).int() * num_frames_per_token + ) # padded to multiple of num_frames_per_token, used for decoding + + gen_horizon_len = self.gen_horizon_len + init_history_len = 0 if init_history_sequence is None else init_history_sequence.shape[1] + num_autoregressive_steps = int(np.ceil((num_frames - init_history_len) / gen_horizon_len)) + num_gen_frames = num_autoregressive_steps * gen_horizon_len + init_history_len + motion_pad_mask = torch.arange(num_gen_frames, device=self.device) < motion_len_pad[:, None] + if num_frames < num_gen_frames: + pad_num_frames = num_gen_frames - num_frames + if motion_mask is not None: + motion_mask = torch.cat( + [ + motion_mask, + torch.zeros( + batch_size, + pad_num_frames, + motion_mask.shape[-1], + device=self.device, + ), + ], + dim=1, + ) + if observed_motion is not None: + observed_motion = torch.cat( + [ + observed_motion, + torch.zeros( + batch_size, + pad_num_frames, + observed_motion.shape[-1], + device=self.device, + ), + ], + dim=1, + ) + + generation_len = torch.ones(batch_size, device=self.device, dtype=torch.long) * gen_horizon_len + assert crop_history_length is None or crop_history_length % num_frames_per_token == 0, ( + "crop_history_length should be a multiple of num_frames_per_token" + ) + + if init_history_sequence is not None: + assert first_heading_angle is None + assert crop_history_length is None, ( + "not supporting combination of crop_history_length and init_history_sequence" + ) + history_sequence, global_transl, first_heading_angle = self._encode_init_history( + init_history_sequence, batch_size, crop_history_length + ) + else: + history_sequence = None + global_transl = torch.zeros( + (batch_size, self.motion_rep.nfeats_dict["root_pos"]), + device=self.device, + ) + + # sample loop + indices = list(range(num_denoising_steps))[::-1] + num_denoising_steps = torch.tensor( + [num_denoising_steps], device=self.device + ) # this and t need to be tensor for onnx export + # init diffusion with correct num steps before looping + use_timesteps = self.diffusion.space_timesteps(num_denoising_steps[0])[0] + self.diffusion.calc_diffusion_vars(use_timesteps) + + for auto_step in range(num_autoregressive_steps): + history_end_frame = auto_step * gen_horizon_len + init_history_len + history_start_frame = ( + 0 if crop_history_length is None else max(0, history_end_frame - crop_history_length) + ) # set history start frame to 0 if no crop, otherwise set to the first frame of the cropped history, the start frame should be non-negative + + history_sequence = self._generate_window( + history_sequence, + global_transl, + history_start_frame, + history_end_frame, + num_gen_frames, + text_feat, + text_pad_mask, + first_heading_angle, + motion_mask, + observed_motion, + num_denoising_steps, + cfg_weight, + indices, + progress_bar=progress_bar, + target_motion=target_motion, + cfg_type=cfg_type, + ) + + # recenter the history sequence as specified + # transl_center_mode == "last_history": + center_frame_index = (history_end_frame + generation_len - 1).clamp(min=0) # [B,] + history_sequence, center_pos, new_root_motion = self._recenter_history( + history_sequence, center_frame_index, requantize=True + ) + + # update global translation + global_transl = global_transl + center_pos + + # update first heading angle if history cropping is applied + if crop_history_length is not None: + # new_root_motion holds root features only — unnormalize with the + # root-slice stats, not the full-feature stats. + heading_angles = self.motion_rep.get_root_heading_angle( + self.motion_rep.global_root_stats.unnormalize(new_root_motion), + ) + full_history_len = new_root_motion.shape[1] + cropped_history_start_frame = max(0, full_history_len - crop_history_length) + first_heading_angle = heading_angles[:, cropped_history_start_frame] + + root_motion, latent_body_motion = self.hybrid.get_root_and_latent_body_motion_from_hybrid( + history_sequence, + ) + new_root_motion = translate_normalized_root_motion(root_motion, global_transl, self.motion_rep) + history_sequence = self.hybrid.get_hybrid_motion_from_root_and_latent_body_motion( + new_root_motion, + latent_body_motion, + ) + if crop_history_length is not None: + motion_output = self.hybrid.get_explicit_motion_from_hybrid_autoregressive( + history_sequence, + motion_pad_mask, + motion_len_pad, + motion_mask=motion_mask, + crop_history_length=crop_history_length, + ) + else: + motion_output = self.hybrid.get_explicit_motion_from_hybrid( + history_sequence, + motion_pad_mask, + motion_len_pad, + motion_mask=motion_mask, + ) + motion_pred = motion_output[:, :num_frames] + + if return_text_embeddings_x_dict: + x_dict = {"x": text_feat, "mask": text_pad_mask} + return motion_pred, x_dict + return motion_pred + + def autoregressive_step( + self, + num_frames: int, + num_denoising_steps: int, + motion_mask: torch.Tensor | None, + observed_motion: torch.Tensor | None, + cfg_weight: Optional[float] = 2.0, + texts: Optional[List[str]] = None, + text_feat: Optional[torch.Tensor] = None, + text_pad_mask: Optional[torch.Tensor] = None, + init_history_sequence: Optional[torch.Tensor] = None, + init_global_translation: Optional[torch.Tensor] = None, # [B, 3] + init_first_heading_angle: Optional[torch.Tensor] = None, # [B,] + ) -> torch.Tensor: + """Perform a single autoregressive generation step. + + Returns: + torch.Tensor: motions in the generated window + """ + # Encode text if not provided + if text_feat is None: + assert text_pad_mask is None + text_feat, text_pad_mask = self._encode_text(texts) + + batch_size = text_feat.shape[0] + num_frames_per_token = self.num_frames_per_token + gen_horizon_len = self.gen_horizon_len + + assert num_frames % num_frames_per_token == 0, "num_frames should be a multiple of num_frames_per_token" + + # Init diffusion with correct num steps + indices = list(range(num_denoising_steps))[::-1] + num_denoising_steps_tensor = torch.tensor([num_denoising_steps], device=self.device) + use_timesteps = self.diffusion.space_timesteps(num_denoising_steps_tensor[0])[0] + self.diffusion.calc_diffusion_vars(use_timesteps) + + # process init history sequence + init_history_len = 0 if init_history_sequence is None else init_history_sequence.shape[1] + if init_history_sequence is not None: + history_sequence, global_transl, first_heading_angle = self._encode_init_history( + init_history_sequence, batch_size + ) + else: + history_sequence = None + global_transl = ( + init_global_translation + if init_global_translation is not None + else torch.zeros( + (batch_size, self.motion_rep.nfeats_dict["root_pos"]), + device=self.device, + ) + ) + first_heading_angle = ( + init_first_heading_angle + if init_first_heading_angle is not None + else torch.zeros(batch_size, device=self.device) + ) + + # Generate a single window on top of the (optional) history. + history_end_frame = init_history_len + history_start_frame = 0 + history_sequence = self._generate_window( + history_sequence, + global_transl, + history_start_frame, + history_end_frame, + num_frames, + text_feat, + text_pad_mask, + first_heading_angle, + motion_mask, + observed_motion, + num_denoising_steps_tensor, + cfg_weight, + indices, + progress_bar=lambda iterable: iterable, + target_motion=None, + cfg_type=None, + ) + + root_motion, latent_body_motion = self.hybrid.get_root_and_latent_body_motion_from_hybrid(history_sequence) + new_root_motion = translate_normalized_root_motion(root_motion, global_transl, self.motion_rep) + # Quantize the latent body motion if the autoencoder uses quantization + if self.autoencoder.encode_with_quantization: + latent_body_motion = self.autoencoder.requantize(latent_body_motion) + # Combine back to hybrid motion + history_sequence = self.hybrid.get_hybrid_motion_from_root_and_latent_body_motion( + new_root_motion, + latent_body_motion, + ) + + generated_motion_len = init_history_len + gen_horizon_len + motion_pad_mask = torch.ones(batch_size, generated_motion_len, device=self.device, dtype=torch.bool) + motion_len = ( + torch.ones(batch_size, device=self.device, dtype=torch.long) * generated_motion_len + ) # length of generated motions + motion_output = self.hybrid.get_explicit_motion_from_hybrid( + history_sequence, + motion_pad_mask, + motion_len, + motion_mask=motion_mask, + ) + return motion_output diff --git a/ardy/model/auto_latent_twostage_denoiser.py b/ardy/model/auto_latent_twostage_denoiser.py new file mode 100644 index 0000000000000000000000000000000000000000..62cd4d94b419cee075f00cb3e085d3f6d7d0662b --- /dev/null +++ b/ardy/model/auto_latent_twostage_denoiser.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +from typing import Optional + +import torch +from einops import rearrange +from torch import nn + +from ardy.model.latent_utils import HybridMotionConverter + +from .backbone import TransformerEncoderBlock +from .loading import load_checkpoint_state_dict + + +def sparsify_token_seq(token_seq, token_mask, token_index): + """Sparsify the token sequence to only include the tokens at positions that contain at least one + valid token in the batch. + + The token index will then not be continuous. + Args: + token_seq (torch.Tensor): [B, num_tokens, dim_token] token sequence + token_mask (torch.Tensor): [B, num_tokens] boolean token mask + token_index (torch.Tensor): [B, num_tokens] integer token index + Returns: + token_seq (torch.Tensor): [B, num_valid_tokens, dim_token] token sequence + token_mask (torch.Tensor): [B, num_valid_tokens] boolean token mask + token_index (torch.Tensor): [B, num_valid_tokens] integer token index + index_mask (torch.Tensor): [num_tokens,] boolean mask, if any token at this index is valid in the batch + """ + index_mask = token_mask.any(dim=0) # [num_tokens,], if any token at this index is valid in the batch + # slice the token sequence, token_mask and token index to only include the valid tokens + token_seq = token_seq[:, index_mask] + token_mask = token_mask[:, index_mask] + token_index = token_index[:, index_mask] + return token_seq, token_mask, token_index, index_mask + + +class AutoLatentTwostageDenoiser(nn.Module): + """Two-stage denoiser: first predicts global root features, then body features conditioned on + local root.""" + + def __init__( + self, + motion_rep, + motion_mask_mode, + num_frames_per_token, + nframe_root_dim, + latent_embedding_dim, + # other params + ckpt_path: Optional[str] = None, + sparsify_token_seq: bool = False, + trt_compatible: bool = True, + **kwargs, + ): + """Build root and body transformer blocks; optionally load checkpoint from ckpt_path.""" + assert not (sparsify_token_seq and trt_compatible), "Should choose between sparsification or trt compatibility." + + super().__init__() + self.motion_rep = motion_rep + self.motion_mask_mode = motion_mask_mode + + self.num_frames_per_token = num_frames_per_token + self.nframe_root_dim = nframe_root_dim + + self.latent_embedding_dim = latent_embedding_dim + self.latent_dim = kwargs["latent_dim"] + self.positional_encoding_mode = kwargs["positional_encoding_mode"] + self.llm_shape = kwargs["llm_shape"] + + self.sparsify_token_seq = sparsify_token_seq + + print(f"sparsify_token_seq: {self.sparsify_token_seq}") + + local_root_dim = motion_rep.local_root_dim + self.nframe_local_root_dim = local_root_dim * self.num_frames_per_token + global_root_hybrid_token_dim = self.nframe_root_dim + self.latent_embedding_dim + + local_root_hybrid_token_dim = ( + self.nframe_local_root_dim + self.latent_embedding_dim + ) # body stage always takes in local root info for motion (but still the global mask) + + if motion_mask_mode == "concat": + # stage 1: root pred only + self.global_root_hybrid_proj = nn.Linear(global_root_hybrid_token_dim, self.latent_dim) + global_root_hybrid_constraints_dim = ( + global_root_hybrid_token_dim + + (motion_rep.motion_rep_dim + motion_rep.body_dim) * self.num_frames_per_token + ) + self.global_root_hybrid_constraints_proj = nn.Linear( + global_root_hybrid_constraints_dim, self.latent_dim + ) # generation tokens with constraints infilling and concat + # stage 2: body pred only + self.local_root_hybrid_proj = nn.Linear(local_root_hybrid_token_dim, self.latent_dim) + local_root_hybrid_constraints_dim = ( + local_root_hybrid_token_dim + + (motion_rep.motion_rep_dim + motion_rep.body_dim) * self.num_frames_per_token + ) # body stage always takes in local root info for motion (but still the global mask) + self.local_root_hybrid_constraints_proj = nn.Linear( + local_root_hybrid_constraints_dim, self.latent_dim + ) # generation tokens with constraints infilling and concat + # explicit future constraints projection for both stages + self.future_constraints_proj = nn.Linear( + motion_rep.motion_rep_dim * 2 * self.num_frames_per_token, + self.latent_dim, + ) + else: + assert motion_mask_mode is None + self.global_root_hybrid_proj = nn.Linear(global_root_hybrid_token_dim, self.latent_dim) + self.local_root_hybrid_proj = nn.Linear(local_root_hybrid_token_dim, self.latent_dim) + + root_output_dim = self.nframe_root_dim + self.root_model = TransformerEncoderBlock( + add_input_proj=False, + input_dim=-1, # for compatibility + output_dim=root_output_dim, + skeleton=self.motion_rep.skeleton, + **kwargs, + ) + + body_output_dim = self.latent_embedding_dim + self.body_model = TransformerEncoderBlock( + add_input_proj=False, + input_dim=-1, # for compatibility + output_dim=body_output_dim, + skeleton=self.motion_rep.skeleton, + **kwargs, + ) + + self.trt_compatible = trt_compatible + self.hybrid = HybridMotionConverter(denoiser=self) + + if ckpt_path: + self.load_ckpt(ckpt_path) + + def load_ckpt(self, ckpt_path: str) -> None: + """Load checkpoint from path; state dict keys are stripped of 'denoiser.backbone.' + prefix.""" + state_dict = load_checkpoint_state_dict(ckpt_path) + state_dict = { + key.replace("denoiser.", ""): val for key, val in state_dict.items() if key.startswith("denoiser") + } + state_dict = {key.replace("backbone.", ""): val for key, val in state_dict.items()} + self.load_state_dict(state_dict) + + def forward( + self, + x: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: torch.Tensor, + first_heading_angle: Optional[torch.Tensor] = None, + motion_mask: Optional[torch.Tensor] = None, + observed_motion: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Args: + x (torch.Tensor): [B, num_tokens, dim_token] current noisy hybrid motion + history_len (torch.Tensor): [B,] history length + generation_len (torch.Tensor): [B,] generation length + future_len (torch.Tensor): [B,] future length. history_len + generation_len + future length == motion_len, but not necessarily equal to num_frames which are padded to the same length + history_mask (torch.Tensor): [B, num_frames] attention mask, positions with True are allowed to attend, False are not + generation_mask (torch.Tensor): [B, num_frames] attention mask, positions with True are allowed to attend, False are not + future_mask (torch.Tensor): [B, num_frames] attention mask, positions with True are allowed to attend, False are not + history_token_mask (torch.Tensor): [B, num_tokens] attention mask, positions with True are allowed to attend, False are not + generation_token_mask (torch.Tensor): [B, num_tokens] attention mask, positions with True are allowed to attend, False are not + future_token_mask (torch.Tensor): [B, num_tokens] attention mask, positions with True are allowed to attend, False are not + text_feat (torch.Tensor): [B, max_text_len, llm_dim] embedded text prompts + text_feat_pad_mask (torch.Tensor): [B, max_text_len] attention mask, positions with True are allowed to attend, False are not + timesteps (torch.Tensor): [B,] current denoising step + motion_mask (torch.Tensor): [B, num_frames, dim_motion] attention mask, positions with True are allowed to attend, False are not + observed_motion (torch.Tensor): [B, num_frames, dim_motion] observed motion + + Returns: + torch.Tensor: same size as input x + """ + bs, num_tokens, dim_token = x.shape + num_frames = num_tokens * self.num_frames_per_token + global_root_motion, latent_body_motion = self.hybrid.get_root_and_latent_body_motion_from_hybrid(x) + token_index = torch.arange(num_tokens, device=x.device)[None, :].repeat(bs, 1) # [B, num_tokens] + if self.positional_encoding_mode == "learned_prefix_zero_at_first_generation": + origin_index = history_len // self.num_frames_per_token + token_index = token_index - origin_index[:, None] + + if self.motion_mask_mode is not None: + if motion_mask is None or observed_motion is None: + motion_shape = (bs, num_frames, self.motion_rep.motion_rep_dim) + motion_mask = torch.zeros(motion_shape, device=x.device) + observed_motion = torch.zeros(motion_shape, device=x.device) + + if self.motion_mask_mode == "concat": + observed_root_motion = self.motion_rep.extract_root(observed_motion) + observed_body_motion = self.motion_rep.extract_body(observed_motion) + root_motion_mask = self.motion_rep.extract_root(motion_mask) + # infill root constraints from observed motion and motion mask. the root motion mask will always be false for the history tokens, so the history root motion will not be changed. + global_root_motion = ( + global_root_motion * (1 - root_motion_mask) + observed_root_motion * root_motion_mask + ) + + x_infilled_extended = torch.cat( + [ + latent_body_motion, + global_root_motion.reshape((bs, num_tokens, self.nframe_root_dim)), + observed_body_motion.reshape(bs, num_tokens, -1), + motion_mask.reshape(bs, num_tokens, -1), + ], + dim=-1, + ) + + if self.trt_compatible: + # Dense masking instead of fancy indexing (TRT-compatible) + generation_token_proj = ( + self.global_root_hybrid_constraints_proj(x_infilled_extended) + * generation_token_mask[:, :, None] + ) + history_token_proj = self.global_root_hybrid_proj(x) * history_token_mask[:, :, None] + future_token_proj = ( + self.future_constraints_proj(x_infilled_extended[:, :, self.latent_embedding_dim :]) + * future_token_mask[:, :, None] + ) + root_stage_input = generation_token_proj + history_token_proj + future_token_proj + root_stage_pad_mask = history_token_mask | generation_token_mask | future_token_mask + else: + generation_token_extended = x_infilled_extended[ + generation_token_mask + ] # [num_generation_tokens, D], float32 + generation_token_proj = self.global_root_hybrid_constraints_proj( + generation_token_extended + ) # [num_generation_tokens, latent_dim for transformer], float16, note the dtype can change + history_token = x[history_token_mask] # [num_history_tokens, D] + history_token_proj = self.global_root_hybrid_proj(history_token) + future_constraints_input = x_infilled_extended[future_token_mask] + future_constraints_input = future_constraints_input[ + :, self.latent_embedding_dim : + ] # [num_future_tokens, D], extract the part after the latent body embedding + future_token_proj = self.future_constraints_proj(future_constraints_input) + dtype = generation_token_proj.dtype + device = generation_token_proj.device + root_stage_input = torch.zeros( + bs, num_tokens, self.latent_dim, dtype=dtype, device=device + ) # [B, num_tokens, latent_dim for transformer] + root_stage_input[generation_token_mask] = generation_token_proj + root_stage_input[history_token_mask] = history_token_proj + root_stage_input[future_token_mask] = future_token_proj + root_stage_pad_mask = history_token_mask | generation_token_mask | future_token_mask + else: + raise NotImplementedError(f"This motion mask mode ({self.motion_mask_mode}) is not supported.") + else: + root_stage_input = self.global_root_hybrid_proj(x) + root_stage_pad_mask = history_token_mask | generation_token_mask + root_stage_token_index = token_index + if self.sparsify_token_seq: + ( + root_stage_input, + root_stage_pad_mask, + root_stage_token_index, + root_stage_index_mask, + ) = sparsify_token_seq(root_stage_input, root_stage_pad_mask, root_stage_token_index) + + # Stage 1: predict root motion in global + global_root_motion_pred = self.root_model( + root_stage_input, + root_stage_pad_mask, + text_feat, + text_feat_pad_mask, + timesteps, + first_heading_angle, + root_stage_token_index, + ) # [B, num_tokens, num_frames_per_token * 5] + if self.sparsify_token_seq: # convert back to dense tensor + dense_tensor = torch.zeros( + bs, + num_tokens, + self.nframe_root_dim, + dtype=global_root_motion_pred.dtype, + device=global_root_motion_pred.device, + ) + dense_tensor[:, root_stage_index_mask, :] = global_root_motion_pred + global_root_motion_pred = dense_tensor + global_root_motion_pred = rearrange( + global_root_motion_pred, + "b t (f d) -> b (t f) d", + f=self.num_frames_per_token, + ) + global_root_motion_pred = ( + global_root_motion_pred * generation_mask[:, :, None] + global_root_motion * history_mask[:, :, None] + ) + lengths = history_len + generation_len + + # Convert root pred to local rep + # At test-time want to allow gradient through for guidance + convert_ctx = torch.no_grad() if self.training else contextlib.nullcontext() + with convert_ctx: + local_root_motion = self.motion_rep.global_root_to_local_root( + global_root_motion_pred, + normalized=True, + lengths=lengths, + ) + if self.training: + local_root_motion = local_root_motion.detach() + + # concatenate the predicted local root with the body motion + local_root_motion = rearrange( + local_root_motion, + "b (t f) d -> b t (f d)", + f=self.num_frames_per_token, + ) + x_new = torch.cat([local_root_motion, latent_body_motion], dim=-1) + + if self.motion_mask_mode == "concat": + x_new_extended = torch.cat( + [ + x_new, + observed_body_motion.reshape(bs, num_tokens, -1), + motion_mask.reshape(bs, num_tokens, -1), + ], + dim=-1, + ) + if self.trt_compatible: + generation_token_proj = ( + self.local_root_hybrid_constraints_proj(x_new_extended) * generation_token_mask[:, :, None] + ) + history_token_proj = self.local_root_hybrid_proj(x_new) * history_token_mask[:, :, None] + body_stage_input = generation_token_proj + history_token_proj + future_token_proj + body_stage_pad_mask = history_token_mask | generation_token_mask | future_token_mask + else: + generation_token_extended = x_new_extended[generation_token_mask] # [num_generation_tokens, D] + generation_token_proj = self.local_root_hybrid_constraints_proj( # noqa + generation_token_extended + ) + history_token = x_new[history_token_mask] # [num_history_tokens, D] + history_token_proj = self.local_root_hybrid_proj(history_token) + # reuse the future_token_proj from body stage, because the future constraints are explicit and sparse, we can not convert to local root + dtype = generation_token_proj.dtype + device = generation_token_proj.device + body_stage_input = torch.zeros( + bs, num_tokens, self.latent_dim, dtype=dtype, device=device + ) # [B, num_tokens, latent_dim for transformer] + body_stage_input[generation_token_mask] = generation_token_proj + body_stage_input[history_token_mask] = history_token_proj + body_stage_input[future_token_mask] = future_token_proj + body_stage_pad_mask = history_token_mask | generation_token_mask | future_token_mask + else: + assert self.motion_mask_mode is None + body_stage_input = self.local_root_hybrid_proj(x_new) + body_stage_pad_mask = history_token_mask | generation_token_mask + body_stage_token_index = token_index + if self.sparsify_token_seq: + ( + body_stage_input, + body_stage_pad_mask, + body_stage_token_index, + body_stage_index_mask, + ) = sparsify_token_seq(body_stage_input, body_stage_pad_mask, body_stage_token_index) + + # Stage 2: predict local body motion based on local root + latent_body_motion_pred = self.body_model( + body_stage_input, + body_stage_pad_mask, + text_feat, + text_feat_pad_mask, + timesteps, + first_heading_angle, + body_stage_token_index, + ) + if self.sparsify_token_seq: # convert back to dense tensor + dense_tensor = torch.zeros( + bs, + num_tokens, + self.latent_embedding_dim, + dtype=latent_body_motion_pred.dtype, + device=latent_body_motion_pred.device, + ) + dense_tensor[:, body_stage_index_mask, :] = latent_body_motion_pred + latent_body_motion_pred = dense_tensor + latent_body_motion_pred = ( + latent_body_motion_pred * generation_token_mask[:, :, None] + + latent_body_motion * history_token_mask[:, :, None] + ) # fuse history and generation tokens + + # concatenate the predicted local body with the predicted root + output = torch.cat( + [ + global_root_motion_pred.reshape(bs, num_tokens, self.nframe_root_dim), + latent_body_motion_pred, + ], + dim=-1, + ) + return output diff --git a/ardy/model/autoencoder/__init__.py b/ardy/model/autoencoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6736106d50a3622383d94f546e49aca20d916f55 --- /dev/null +++ b/ardy/model/autoencoder/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Motion autoencoder modules (FSQ quantizer and transformer encoder/decoder).""" diff --git a/ardy/model/autoencoder/fsq.py b/ardy/model/autoencoder/fsq.py new file mode 100644 index 0000000000000000000000000000000000000000..71af60bf9efa1b4583de904ef0e13139b52e8f80 --- /dev/null +++ b/ardy/model/autoencoder/fsq.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +import torch as t +from torch import nn +from vector_quantize_pytorch import FSQ + +from ardy.model.loading import load_checkpoint_state_dict +from ardy.motion_rep import MotionRepBase +from ardy.motion_rep.stats import Stats + +from .transformer import ( + DoubleCondDecoderTransformer, + EncoderTransformer, +) + + +def round_ste(z): + """Round with straight through gradients.""" + zhat = z.round() + return z + (zhat - z).detach() + + +class FSQVAETransformer(nn.Module): + ALLOWED_FEATURE_MODE = ["pose", "body", "root"] + + def __init__( + self, + motion_rep: MotionRepBase, + feature_mode: list, + # fsq config + num_fsq_levels: int = 5, + fsq_level_list: int | list[int] = [16, 16, 16, 16, 16], + # encoding config + encode_with_normalization: bool = True, + encode_with_quantization: bool = True, + # network config + latent_embedding_dim: int = 128, # number of the output embedding dimension + num_frames_per_token: int = 4, + latent_dim: int = 512, + ff_size: int = 1024, + num_layers: int = 8, + num_heads: int = 4, + activation: str = "gelu", + dropout: float = 0.1, + pe_dropout: float = 0.1, + norm_first: bool = False, + causal_encoder: bool = False, + causal_decoder: bool = False, + ckpt_path: Optional[str] = None, + **kwargs, + ): + super().__init__() + self._motion_rep = motion_rep + self._num_frames_per_token = num_frames_per_token + self._latent_embedding_dim = num_fsq_levels + + self.encode_with_normalization = encode_with_normalization + self.encode_with_quantization = encode_with_quantization + + assert len(feature_mode) == 4 and np.all([f in self.ALLOWED_FEATURE_MODE for f in feature_mode]) + ( + self.encoder_input_feature_mode, # body + self.decoder_output_feature_mode, # pose + self.decoder_target_cond_feature_mode, # body + self.decoder_external_cond_feature_mode, # root + ) = feature_mode + + # always use the local root + dim_dict = { + "pose": motion_rep.local_motion_rep_dim, # local root + body + "body": motion_rep.body_dim, + "root": motion_rep.local_root_dim, + } + encoder_state_dim = dim_dict[self.encoder_input_feature_mode] + decoder_state_dim = dim_dict[self.decoder_output_feature_mode] + decoder_target_cond_dim = dim_dict[self.decoder_target_cond_feature_mode] + decoder_external_cond_dim = dim_dict[self.decoder_external_cond_feature_mode] + + latent_embedding_dim = num_fsq_levels + self.encoder = EncoderTransformer( + encoder_state_dim, + latent_embedding_dim, + num_frames_per_token, + latent_dim, + num_heads, + ff_size, + dropout, + activation, + norm_first, + num_layers, + pe_dropout, + causal_encoder, + ) + self.decoder = DoubleCondDecoderTransformer( + latent_embedding_dim, + decoder_state_dim, + num_frames_per_token, + latent_dim, + num_heads, + ff_size, + dropout, + activation, + norm_first, + num_layers, + pe_dropout, + causal_decoder, + target_cond_dim=decoder_target_cond_dim, + external_cond_dim=decoder_external_cond_dim, + ) + + if isinstance(fsq_level_list, int): + fsq_level_list = [fsq_level_list] * num_fsq_levels + assert len(fsq_level_list) == num_fsq_levels, ( + f"fsq_level_list must have num_fsq_levels={num_fsq_levels} entries, got {len(fsq_level_list)}" + ) + self.quantizer = FSQ(levels=fsq_level_list, return_indices=False) + + if ckpt_path: + self.load_ckpt(ckpt_path) + self.eval_and_freeze() + + # make forward by default used for detokenization + self.forward = self.detokenize + + def eval_and_freeze(self): + self.eval() + # freeze the autoencoder + for param in self.parameters(): + param.requires_grad = False + + def load_ckpt(self, ckpt_path: str) -> None: + """Load checkpoint from path; state dict keys are stripped of 'denoiser.backbone.' prefix. + + Then load the stats. + """ + # Load ckpt + state_dict = load_checkpoint_state_dict(ckpt_path) + state_dict = { + key.removeprefix("pose_net."): val for key, val in state_dict.items() if key.startswith("pose_net") + } + self.load_state_dict(state_dict) + + ckpt_path = str(ckpt_path) + autoencoder_stats_dir = Path(ckpt_path).parent / "stats" + self.post_quantization_stats = Stats(folder=autoencoder_stats_dir / "post_quantization", load=True) + self.pre_quantization_stats = Stats(folder=autoencoder_stats_dir / "pre_quantization", load=True) + + @property + def motion_rep(self): + return self._motion_rep + + @property + def num_frames_per_token(self): + return self._num_frames_per_token + + def extract_feature( + self, + x: t.Tensor, + feature: str = "", + lengths: Optional[torch.Tensor] = None, + ): + """@brief: extract the full pose / root or body from the full local features""" + assert feature in self.ALLOWED_FEATURE_MODE + + if feature == "body": + return self.motion_rep.extract_body(x) + + global_root = self.motion_rep.extract_root(x) + local_root = self.motion_rep.global_root_to_local_root(global_root, normalized=True, lengths=lengths) + + if feature == "root": + return local_root + + # pose + assert feature == "pose" + body = self.motion_rep.extract_body(x) + local_features = torch.cat([local_root, body], dim=-1) + return local_features + + def tokenize( + self, + x: t.Tensor, + motion_pad_mask: t.BoolTensor, + ): + """@brief: get the embeddings of the input motion + @param x: local poses [batch_size, numFrames, feat_dim] + @param motion_pad_mask: [batch_size, numFrames], 1 means valid frames + @return embeddings: [batch_size, num_tokens, feat_dim] + """ + # create lengths + x_in = self.extract_feature(x, self.encoder_input_feature_mode, lengths=motion_pad_mask.sum(1)) + x_encoder = self.encoder(x_in, motion_pad_mask) + + if self.encode_with_quantization: + x_quantized, indices = self.quantizer(x_encoder) + token_embeddings = x_quantized + else: + token_embeddings = x_encoder + + if self.encode_with_normalization: + # mean = self.stats['mean_after_quantization'] if encode_with_quantization else self.stats['mean_before_quantization'] + # std = self.stats['std_after_quantization'] if encode_with_quantization else self.stats['std_before_quantization'] + # token_embeddings = (token_embeddings - mean) / std + stats = self.post_quantization_stats if self.encode_with_quantization else self.pre_quantization_stats + token_embeddings = stats.normalize(token_embeddings) + + return token_embeddings + + def detokenize( + self, + token_embeddings: t.Tensor, + external_cond: t.Tensor = None, + motion_pad_mask: t.BoolTensor = None, + ): + """@brief: get the original motion from the token embeddings + @param token_embeddings: [batch_size, num_tokens, feat_dim] + @param motion_pad_mask: [batch_size, numFrames], 1 means valid frames + @return x: [batch_size, numFrames, feat_dim] + """ + if self.encode_with_normalization: + # mean = self.stats['mean_after_quantization'] if encode_with_quantization else self.stats['mean_before_quantization'] + # std = self.stats['std_after_quantization'] if encode_with_quantization else self.stats['std_before_quantization'] + # token_embeddings = token_embeddings * std + mean + stats = self.post_quantization_stats if self.encode_with_quantization else self.pre_quantization_stats + token_embeddings = stats.unnormalize(token_embeddings) + + if not self.encode_with_quantization: # quantization not applied at encoding + token_embeddings, indices = self.quantizer(token_embeddings) + else: # redo the rounding operation do ensure quantization, but we can not directly call quantizer since it include additional operation apart from rounding + half_width = self.quantizer._levels // 2 + token_embeddings = round_ste(token_embeddings.clamp(-1, 1) * half_width) / half_width + + x_decoder = self.decoder( + token_embeddings, + external_cond=external_cond, + motion_pad_mask=motion_pad_mask, + ) + # produce a dict instead of the raw output + output = {} + if self.decoder_output_feature_mode == "pose": + output["root"] = x_decoder[..., : self.motion_rep.local_root_dim] + output["body"] = x_decoder[..., self.motion_rep.local_root_dim :] + elif self.decoder_output_feature_mode == "body": + output["body"] = x_decoder + elif self.decoder_output_feature_mode == "root": + output["root"] = x_decoder + else: + raise NotImplementedError + return output + + def requantize( + self, + token_embeddings: t.Tensor, + ): + """@brief: requantize the token embeddings to the discrete values + @param token_embeddings: [batch_size, num_tokens, feat_dim] + @return requantized_token_embeddings: [batch_size, num_tokens, feat_dim] + """ + assert self.encode_with_quantization, "Only support encode_with_quantization=True" + + if self.encode_with_normalization: # unnormalize the token embeddings + stats = self.post_quantization_stats + token_embeddings = stats.unnormalize(token_embeddings) + + half_width = self.quantizer._levels // 2 + token_embeddings = round_ste(token_embeddings.clamp(-1, 1) * half_width) / half_width + + if self.encode_with_normalization: # normalize the token embeddings + stats = self.post_quantization_stats + token_embeddings = stats.normalize(token_embeddings) + + return token_embeddings diff --git a/ardy/model/autoencoder/transformer.py b/ardy/model/autoencoder/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..6205676d468c169475b9179dfcad1b20bbebe1de --- /dev/null +++ b/ardy/model/autoencoder/transformer.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + + +class PositionalEncoding(nn.Module): + """Non-learned positional encoding.""" + + def __init__( + self, + d_model: int, + dropout: Optional[float] = 0.1, + max_len: Optional[int] = 5000, + ): + """ + Args: + d_model (int): input dim + dropout (Optional[float] = 0.1): dropout probability on output + max_len (Optional[int] = 5000): maximum sequence length + """ + super(PositionalEncoding, self).__init__() + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + + # Note: have to replace torch.exp() and math.log() with torch.pow() + # due to MKL exp() and ln() throws floating point exceptions on certain CPUs + # see corresponding commit and MR + div_term = torch.pow(10000.0, -torch.arange(0, d_model, 2).float() / d_model) + # div_term = torch.exp( + # torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model) + # ) + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) # [1, T, D] + + self.register_buffer("pe", pe, persistent=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply positional encoding to input sequence. + + Args: + x (torch.Tensor): [B, T, D] input motion sequence + + Returns: + torch.Tensor: [B, T, D] input motion with PE added to it (and optionally dropout) + """ + x = x + self.pe[:, : x.shape[1], :] + return self.dropout(x) + + +class EncoderTransformer(nn.Module): + def __init__( + self, + input_dim: int, + output_dim: int, + num_frames_per_token: int, + latent_dim: int, + num_heads: int, + ff_size: int, + dropout: float, + activation: str, + norm_first: bool, + num_layers: int, + pe_dropout: float, + is_causal: bool = True, + ): + super().__init__() + self.is_causal = is_causal + self.num_frames_per_token = num_frames_per_token + self.latent_dim = latent_dim + self.num_heads = num_heads + self.ff_size = ff_size + self.dropout = dropout + self.activation = activation + self.norm_first = norm_first + self.num_layers = num_layers + self.pe_dropout = pe_dropout + + self.input_proj = nn.Linear(input_dim * num_frames_per_token, latent_dim) + self.sequence_pos_encoder = PositionalEncoding(latent_dim, pe_dropout) + trans_enc_layer = nn.TransformerEncoderLayer( + d_model=latent_dim, + nhead=num_heads, + dim_feedforward=ff_size, + dropout=dropout, + activation=activation, + batch_first=True, + norm_first=norm_first, + ) + # technically we don't need to disable nested tensors here, just + # need to make sure flash attention is not used. One way is to set + # torch.backends.mha.set_fastpath_enabled(False) but this only + # works on torch>=2.3.0. We should test how much perf is lost without + # using nested tensors. + self.seqTransEncoder = nn.TransformerEncoder(trans_enc_layer, num_layers=num_layers, enable_nested_tensor=False) + self.output_proj = nn.Linear(latent_dim, output_dim) + + def forward(self, x, motion_pad_mask: Optional[torch.BoolTensor] = None): + """Tranformer-based encoder. + + Args: + x (torch.Tensor): motions frames [batch_size, numFrames, latent_dim] + motion_pad_mask (Optional[torch.BoolTensor]): shape -> [batch, num_frames] (dtype=bool) + + Returns: + torch.Tensor: latent embeddings [batch_size, num_tokens, latent_dim] + """ + # reshape motion frames to [batch_size, num_tokens, num_frames_per_token * latent_dim] + x = rearrange(x, "b (t f) d -> b t (f d)", f=self.num_frames_per_token) + x = self.input_proj(x) # input projection + x = self.sequence_pos_encoder(x) # positional encoding + + attention_mask = ( + torch.nn.Transformer.generate_square_subsequent_mask(x.shape[1], device=x.device, dtype=torch.bool) + if self.is_causal + else None + ) # create causal mask + if motion_pad_mask is not None: + motion_pad_mask = rearrange(motion_pad_mask, "b (t f) -> b t f", f=self.num_frames_per_token) + # `> 0.5` casts to bool so the downstream `~` traces to ONNX `Not` + # on a BOOL tensor (TRT rejects `Not` on FLOAT). Accepts a float mask + # (TRT I/O convention) or a bool mask (eager); semantics are identical. + token_pad_mask = (motion_pad_mask > 0.5).all(dim=-1) # [B, num_tokens] + # negate the mask to get the padding mask, True means not allowed in attention + src_key_padding_mask = ~token_pad_mask + else: + src_key_padding_mask = None + x = self.seqTransEncoder( + x, mask=attention_mask, src_key_padding_mask=src_key_padding_mask + ) # transformer encoder layers + output_embeddings = self.output_proj(x) # output latent embedding [batch_size, num_tokens, output_dim] + + return output_embeddings + + +class DoubleCondDecoderTransformer(nn.Module): + def __init__( + self, + input_dim: int, + output_dim: int, + num_frames_per_token: int, + latent_dim: int, + num_heads: int, + ff_size: int, + dropout: float, + activation: str, + norm_first: bool, + num_layers: int, + pe_dropout: float, + is_causal: bool = True, + target_cond_dim: int = -1, + external_cond_dim: int = -1, + ): + super().__init__() + self.is_causal = is_causal + self.num_frames_per_token = num_frames_per_token + self.latent_dim = latent_dim + self.num_heads = num_heads + self.ff_size = ff_size + self.dropout = dropout + self.activation = activation + self.norm_first = norm_first + self.num_layers = num_layers + self.pe_dropout = pe_dropout + + # target and external condition configs + self._target_cond_dim, self._external_cond_dim = ( + target_cond_dim, + external_cond_dim, + ) + self._HAS_EXTERNAL_COND, self._HAS_TARGET_COND = ( + external_cond_dim > 0, + target_cond_dim > 0, + ) + + self.input_proj = nn.Linear(input_dim, latent_dim) + self.sequence_pos_encoder = PositionalEncoding(latent_dim, pe_dropout) + trans_enc_layer = nn.TransformerEncoderLayer( + d_model=latent_dim, + nhead=num_heads, + dim_feedforward=ff_size, + dropout=dropout, + activation=activation, + batch_first=True, + norm_first=norm_first, + ) + # technically we don't need to disable nested tensors here, just + # need to make sure flash attention is not used. One way is to set + # torch.backends.mha.set_fastpath_enabled(False) but this only + # works on torch>=2.3.0. We should test how much perf is lost without + # using nested tensors. + self.seqTransEncoder = nn.TransformerEncoder(trans_enc_layer, num_layers=num_layers, enable_nested_tensor=False) + """Step 2: external cond embeddings. + + At each layer, the external embedding will be merged with the hidden state. The external + condition is dense and expected to be always available in each frame. + """ + if self._HAS_EXTERNAL_COND: + external_cond_blocks = [] + cond_feat_dim = num_frames_per_token * self._external_cond_dim + external_cond_blocks.append(nn.Linear(cond_feat_dim + latent_dim, latent_dim)) + external_cond_blocks.append(nn.ReLU()) + self.external_cond_blocks = nn.Sequential(*external_cond_blocks) + """Step 3: target cond embeddings. + + The target condition is sparse and expected to be available in certain frames. we replace + the hidden state with the target condition embeddings for given frames. In earlier layers + where each position corresponds to multiple frames, we reshape the hidden states to map to + each frame position (see @forward method) + """ + if self._HAS_TARGET_COND: + target_cond_blocks = [] + assert latent_dim % num_frames_per_token == 0, ( + "latent_dim % num_frames_per_token needs to 0 so that hidden can be split for each frame in earlier layers." + ) + target_cond_blocks.append(nn.Linear(self._target_cond_dim, int(latent_dim / num_frames_per_token))) + target_cond_blocks.append(nn.ReLU()) + self.target_cond_blocks = nn.Sequential(*target_cond_blocks) + + # output projection + self.output_proj = nn.Linear(latent_dim, output_dim * num_frames_per_token) + + def forward( + self, + x: torch.Tensor, + external_cond: torch.Tensor = None, + target_cond: Optional[torch.Tensor] = None, + has_target_cond: Optional[torch.Tensor] = None, + motion_pad_mask: Optional[torch.BoolTensor] = None, + ): + """@brief: the decoder could take the external condition and target condition as input. + @params x: shape -> [batch, num_tokens, latent_dim] + @params external_cond: shape -> [batch, num_frames, external_cond_dim] + @params target_cond: shape -> [batch, num_frames, target_cond_dim] + @params has_target_cond: shape -> [batch, num_frames] (dtype=bool) + @params motion_pad_mask: shape -> [batch, num_frames] (dtype=bool) + """ + batch_size = x.shape[0] + num_tokens = x.shape[1] + num_frames_per_token = self.num_frames_per_token + num_frames = num_tokens * num_frames_per_token + + h = self.input_proj(x) # [batch, num_tokens, latent_dim] + + if (not self._HAS_TARGET_COND) or target_cond is None or has_target_cond is None: # no target condition + pass + else: # replace parts of tokens with target condition + h_target_cond = self.target_cond_blocks(target_cond) # [batch, num_frames, latent_dim] + h = h.reshape([batch_size, num_frames, self.latent_dim // self.num_frames_per_token]) + h = torch.where(has_target_cond[:, :, None], h_target_cond, h) + h = h.reshape([batch_size, num_tokens, self.latent_dim]) + + if self._HAS_EXTERNAL_COND: # concatenate external condition + assert external_cond is not None and external_cond.shape[1] == num_frames + external_cond = external_cond.reshape( + [batch_size, num_tokens, -1] + ) # [batch, num_tokens, external_cond_dim * num_frames_per_token] + h = torch.cat( + [h, external_cond], dim=-1 + ) # [batch, num_tokens, latent_dim + external_cond_dim * num_frames_per_token] + h = self.external_cond_blocks(h) # [batch, num_tokens, latent_dim] + + h = self.sequence_pos_encoder(h) # positional encoding + attention_mask = ( + torch.nn.Transformer.generate_square_subsequent_mask(h.shape[1], device=h.device, dtype=torch.bool) + if self.is_causal + else None + ) # create causal mask + if motion_pad_mask is not None: + motion_pad_mask = rearrange(motion_pad_mask, "b (t f) -> b t f", f=self.num_frames_per_token) + # `> 0.5` casts to bool so the downstream `~` traces to ONNX `Not` + # on a BOOL tensor (TRT rejects `Not` on FLOAT). Accepts a float mask + # (TRT I/O convention) or a bool mask (eager); semantics are identical. + token_pad_mask = (motion_pad_mask > 0.5).all(dim=-1) # [B, num_tokens] + # negate the mask to get the padding mask, True means not allowed in attention + src_key_padding_mask = ~token_pad_mask + else: + src_key_padding_mask = None + h = self.seqTransEncoder( + h, mask=attention_mask, src_key_padding_mask=src_key_padding_mask + ) # transformer encoder layers + h = self.output_proj(h) # output projection + h = rearrange(h, "b t (f d) -> b (t f) d", f=self.num_frames_per_token) # [batch, num_frames, output_dim] + + return h diff --git a/ardy/model/backbone.py b/ardy/model/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..287b103cb42391e751bdd0e902beba6935ed7df4 --- /dev/null +++ b/ardy/model/backbone.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import logging +from typing import Optional, Union + +import torch +import torch.nn.functional as F +from omegaconf import ListConfig +from pydantic.dataclasses import dataclass +from torch import Tensor, nn + +from ardy.tools import validate + +log = logging.getLogger(__name__) + + +def pad_x_and_mask_to_fixed_size(x: Tensor, mask: Tensor, size: int): + """Pad a feature vector x and the mask to always have the same size. + + Args: + x (torch.Tensor): [B, T, D] + mask (torch.Tensor): [B, T] + size (int) + Returns: + torch.Tensor: [B, size, D] + torch.Tensor: [B, size] + """ + cur_max_size = x.shape[1] + + if cur_max_size == size: + # already padded to this size, probably in the collate function + return x, mask + + if cur_max_size > size: + # This issue should have been handled in the collate function + # useful as a check for test time + log.warn("The size of the tensor is larger than the maximum size. Cropping the input..") + return x[:, :size], mask[:, :size] + + # Pad with zeros along the time dimension (torch.compile-compatible) + pad_len = size - cur_max_size + new_x = torch.nn.functional.pad(x, (0, 0, 0, pad_len)) # pad last-but-one dim + new_mask = torch.nn.functional.pad(mask, (0, pad_len)) # pad last dim + return new_x, new_mask + + +def _get_activation_fn(activation): + """Resolve an activation the same way nn.TransformerEncoderLayer does.""" + if callable(activation): + return activation + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + raise ValueError(f"Unsupported activation: {activation!r} (expected 'relu' or 'gelu')") + + +class SDPATransformerEncoderLayer(nn.Module): + """Drop-in replacement for nn.TransformerEncoderLayer using SDPA. + + Computes the identical function as nn.TransformerEncoderLayer and uses the same submodule names + (``self_attn`` / ``linear1`` / ``linear2`` / ``norm1`` / ``norm2``), so state_dicts are + interchangeable and existing checkpoints load unchanged. The only difference is that self- + attention runs through ``F.scaled_dot_product_attention`` instead of the nn.MultiheadAttention / + BetterTransformer fast path -- which is friendly to torch.compile / CUDA graphs / ONNX export. + + Self-attention only (query == key == value); only ``batch_first=True`` is supported (the + convention used throughout this repo). + """ + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation="relu", + layer_norm_eps: float = 1e-5, + batch_first: bool = True, + norm_first: bool = False, + bias: bool = True, + ) -> None: + super().__init__() + if d_model % nhead != 0: + raise ValueError(f"d_model ({d_model}) must be divisible by nhead ({nhead})") + if not batch_first: + raise NotImplementedError("Only batch_first=True is supported") + + self.d_model = d_model + self.nhead = nhead + self.head_dim = d_model // nhead + self.norm_first = norm_first + self.attn_dropout_p = dropout + + # Parameter container with the SAME names as nn.TransformerEncoderLayer's + # self_attn (in_proj_weight, in_proj_bias, out_proj.{weight,bias}). We reuse + # its parameters but run attention via SDPA rather than calling its forward(). + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, bias=bias, batch_first=True) + + # Feed-forward block: identical structure/names to nn.TransformerEncoderLayer. + self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias) + + self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias) + self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + + def _build_attn_mask(self, src, src_key_padding_mask, src_mask): + """Merge masks into a single additive float mask, as MHA does internally. + + Output broadcasts to [B, nhead, L, S]. Bool masks use the PyTorch convention (True == + ignore); float masks are additive. + """ + if src_key_padding_mask is None and src_mask is None: + return None + bs, seq_len, _ = src.shape + attn_mask = torch.zeros(bs, 1, 1, seq_len, dtype=src.dtype, device=src.device) + if src_key_padding_mask is not None: + if src_key_padding_mask.dtype == torch.bool: + attn_mask = attn_mask.masked_fill(src_key_padding_mask[:, None, None, :], float("-inf")) + else: + attn_mask = attn_mask + src_key_padding_mask[:, None, None, :] + if src_mask is not None: + if src_mask.dtype == torch.bool: + add = torch.zeros_like(src_mask, dtype=src.dtype).masked_fill(src_mask, float("-inf")) + else: + add = src_mask.to(src.dtype) + # [L, S] broadcasts against [B, 1, 1, S] -> [B, 1, L, S] + attn_mask = attn_mask + add + return attn_mask + + def _sa_block(self, x, attn_mask): + bs, seq_len, _ = x.shape + qkv = F.linear(x, self.self_attn.in_proj_weight, self.self_attn.in_proj_bias) + q, k, v = qkv.chunk(3, dim=-1) + q = q.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2) + k = k.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2) + v = v.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2) + dropout_p = self.attn_dropout_p if self.training else 0.0 + attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p) + attn = attn.transpose(1, 2).reshape(bs, seq_len, self.d_model) + attn = self.self_attn.out_proj(attn) + return self.dropout1(attn) + + def _ff_block(self, x): + x = self.linear2(self.dropout(self.activation(self.linear1(x)))) + return self.dropout2(x) + + def forward(self, src, src_mask=None, src_key_padding_mask=None, is_causal=False): + if is_causal: + raise NotImplementedError("is_causal is not supported") + attn_mask = self._build_attn_mask(src, src_key_padding_mask, src_mask) + x = src + if self.norm_first: + x = x + self._sa_block(self.norm1(x), attn_mask) + x = x + self._ff_block(self.norm2(x)) + else: + x = self.norm1(x + self._sa_block(x, attn_mask)) + x = self.norm2(x + self._ff_block(x)) + return x + + +class SDPATransformerEncoder(nn.Module): + """Drop-in replacement for nn.TransformerEncoder built from SDPA layers. + + Same state_dict layout as nn.TransformerEncoder (``layers.{i}.<...>``), so existing checkpoints + load unchanged. + """ + + def __init__(self, encoder_layer, num_layers, norm=None): + super().__init__() + self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)]) + self.num_layers = num_layers + self.norm = norm + + def forward(self, src, mask=None, src_key_padding_mask=None, is_causal=None): + out = src + for layer in self.layers: + out = layer(out, src_mask=mask, src_key_padding_mask=src_key_padding_mask) + if self.norm is not None: + out = self.norm(out) + return out + + +@dataclass(frozen=True, config=dict(extra="forbid", arbitrary_types_allowed=True)) +class TransformerEncoderBlockConfig: + """Configuration for the transformer encoder backbone.""" + + # input features dimension + input_dim: int + # output features dimension + output_dim: int + + # skeleton object + skeleton: object + + # dimension of the text embeddings + llm_shape: Union[list[int], ListConfig] + + # mask the text or not + use_text_mask: bool + + # latent dimension of the model + latent_dim: int + # dimension of the feedforward network in transformer + ff_size: int + # num layers in transformer + num_layers: int + # num heads in transformer + num_heads: int + # activation in transformer + activation: str + # dropout rate for the transformer + dropout: float + # dropout rate for the positional embeddings + pe_dropout: float + # use norm first or not + norm_first: bool = False + + # Input first heading angle + input_first_heading_angle: bool = False + + # auto latent model + add_input_proj: bool = True + positional_encoding_mode: str = "default" + + +class TransformerEncoderBlock(nn.Module): + @validate(TransformerEncoderBlockConfig, save_args=True, super_init=True) + def __init__(self, conf): + self.nbjoints = self.skeleton.nbjoints + llm_dim = self.llm_shape[-1] + self.embed_text = nn.Linear(llm_dim, self.latent_dim) + + # maximum number of tokens + self.num_text_tokens = self.llm_shape[0] + + self.sequence_pos_encoder = PositionalEncoding(self.latent_dim, self.pe_dropout) + self.embed_timestep = TimestepEmbedder(self.latent_dim, self.sequence_pos_encoder) + + if self.add_input_proj: + self.input_linear = nn.Linear(self.input_dim, self.latent_dim) + else: + self.input_linear = nn.Identity() + self.output_linear = nn.Linear(self.latent_dim, self.output_dim) + + if self.input_first_heading_angle: + self.linear_first_heading_angle = nn.Linear(2, self.latent_dim) + + if self.positional_encoding_mode == "learned_prefix_zero_at_first_generation": + prefix_length = self.num_text_tokens + 1 # text tokens + diffusion step token + if self.input_first_heading_angle: + prefix_length += 1 + self.prefix_length = prefix_length + self.learned_prefix_embedding = LearnedPositionalEncoding( + self.latent_dim, self.pe_dropout, max_len=prefix_length + ) + self.motion_token_embedding = PositionalEncodingNegativeIndex(self.latent_dim, self.pe_dropout) + elif self.positional_encoding_mode == "default": + pass + else: + raise ValueError(f"Invalid positional encoding mode: {self.positional_encoding_mode}") + + trans_enc_layer = SDPATransformerEncoderLayer( + d_model=self.latent_dim, + nhead=self.num_heads, + dim_feedforward=self.ff_size, + dropout=self.dropout, + activation=self.activation, + batch_first=True, + norm_first=self.norm_first, + ) + self.seqTransEncoder = SDPATransformerEncoder( + trans_enc_layer, + num_layers=self.num_layers, + ) + + def forward( + self, + x: Tensor, + x_pad_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: Tensor, + first_heading_angle: Optional[Tensor] = None, + token_index: Optional[Tensor] = None, + ) -> Tensor: + """ + Args: + x (torch.Tensor): [B, T, dim_motion] current noisy motion + x_pad_mask (torch.Tensor): [B, T] attention mask, positions with True are allowed to attend, False are not + text_feat (torch.Tensor): [B, max_text_len, llm_dim] embedded text prompts + text_feat_pad_mask (torch.Tensor): [B, max_text_len] attention mask, positions with True are allowed to attend, False are not + timesteps (torch.Tensor): [B,] current denoising step + token_index (torch.Tensor): [B,] token index for positional encoding, can be negative if the indices are centered at first generation token. When the future constraints are sparse, indices are not continuous. + Returns: + torch.Tensor: [B, T, output_dim] + """ + batch_size = x.shape[0] + x = self.input_linear(x) # [B, T, D] + + # Pad the text tokens + mask to always have the same size == self.num_text_tokens + # done here if it was not done in the collate function + if self.num_text_tokens is not None: + text_feat, text_feat_pad_mask = pad_x_and_mask_to_fixed_size( + text_feat, + text_feat_pad_mask, + self.num_text_tokens, + ) + + # Encode the text features and the time information. + # The text encoder may run in a different precision (e.g. bfloat16) + # than the denoiser (float32), so align the dtype before projecting to + # avoid "mat1 and mat2 must have the same dtype" errors. + emb_text = self.embed_text(text_feat.to(self.embed_text.weight.dtype)) # [B, max_text_len, D] + emb_time = self.embed_timestep(timesteps) # [B, 1, D] + + # Create mask for the time information + time_mask = torch.ones((batch_size, 1), dtype=bool, device=x.device) + + # Create the prefix features (text, time, etc): [B, max_text_len*repeat_text_token_num + 1 + etc] + prefix_feats = torch.cat((emb_text, emb_time), axis=1) + + # Behavior from old code: not use text mask -> True for all the tokens + if not self.use_text_mask: + # text_feat_pad_mask = torch.ones_like(text_feat_pad_mask) + text_feat_pad_mask = torch.ones( + (batch_size, emb_text.shape[1]), + dtype=torch.bool, + device=x.device, + ) + + prefix_mask = torch.cat((text_feat_pad_mask, time_mask), axis=1) + + # add the input first heading angle + if self.input_first_heading_angle: + assert first_heading_angle is not None, "The first heading angle is mandatory for this model" + # cos(angle) / sin(angle) + first_heading_angle_feats = torch.stack( + [ + torch.cos(first_heading_angle), + torch.sin(first_heading_angle), + ], + axis=-1, + ) + + first_heading_angle_feats = self.linear_first_heading_angle(first_heading_angle_feats) + first_heading_angle_feats = first_heading_angle_feats[:, None] # for cat + first_heading_angle_mask = torch.ones( + (batch_size, 1), + dtype=bool, + device=x.device, + ) + prefix_feats = torch.cat((prefix_feats, first_heading_angle_feats), axis=1) + prefix_mask = torch.cat((prefix_mask, first_heading_angle_mask), axis=1) + + # compute the number of prefix features + pose_start_ind = prefix_feats.shape[1] + + if self.positional_encoding_mode == "default": # prefix-prepended style + # Concatenate prefix and x: [B, len(prefix) + T, D] + xseq = torch.cat((prefix_feats, x), axis=1) + # Add positional encoding + xseq = self.sequence_pos_encoder(xseq) + elif self.positional_encoding_mode == "learned_prefix_zero_at_first_generation": + # apply learned positional encoding to the prefix features + prefix_feats_pe = self.learned_prefix_embedding(prefix_feats) + x_pe = self.motion_token_embedding(x, token_index) + # Concatenate prefix and x: [B, len(prefix) + T, D] + xseq = torch.cat((prefix_feats_pe, x_pe), axis=1) + + # Concatenate the masks and negate them: [B, len(prefix) + T] + src_key_padding_mask = ~torch.cat((prefix_mask, x_pad_mask), axis=1) + + # Input to the transformer and keep the motion indexes + output = self.seqTransEncoder( + xseq, + src_key_padding_mask=src_key_padding_mask, + ) + output = output[:, pose_start_ind:] # [B, T, D] + output = self.output_linear(output) # [B, T, OD] + return output + + +class PositionalEncoding(nn.Module): + """Non-learned positional encoding.""" + + def __init__( + self, + d_model: int, + dropout: Optional[float] = 0.1, + max_len: Optional[int] = 5000, + ): + """ + Args: + d_model (int): input dim + dropout (Optional[float] = 0.1): dropout probability on output + max_len (Optional[int] = 5000): maximum sequence length + """ + super(PositionalEncoding, self).__init__() + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + + # Note: have to replace torch.exp() and math.log() with torch.pow() + # due to MKL exp() and ln() throws floating point exceptions on certain CPUs + # see corresponding commit and MR + div_term = torch.pow(10000.0, -torch.arange(0, d_model, 2).float() / d_model) + # div_term = torch.exp( + # torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model) + # ) + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) # [1, T, D] + + self.register_buffer("pe", pe, persistent=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply positional encoding to input sequence. + + Args: + x (torch.Tensor): [B, T, D] input motion sequence + + Returns: + torch.Tensor: [B, T, D] input motion with PE added to it (and optionally dropout) + """ + x = x + self.pe[:, : x.shape[1], :] + return self.dropout(x) + + +class TimestepEmbedder(nn.Module): + """Encoder for diffusion step.""" + + def __init__(self, latent_dim: int, sequence_pos_encoder: PositionalEncoding): + """ + Args: + latent_dim (int): dim to encode to + sequence_pos_encoder (PositionalEncoding): the PE to use on timesteps + """ + super().__init__() + self.latent_dim = latent_dim + self.sequence_pos_encoder = sequence_pos_encoder + + time_embed_dim = self.latent_dim + self.time_embed = nn.Sequential( + nn.Linear(self.latent_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + """Embed timesteps by adding PE then going through linear layers. + + Args: + timesteps (torch.Tensor): [B] + + Returns: + torch.Tensor: [B, 1, D] + """ + return self.time_embed(F.embedding(timesteps.int(), self.sequence_pos_encoder.pe.squeeze(0))).unsqueeze(1) + + +class LearnedPositionalEncoding(nn.Module): + def __init__( + self, + d_model, + dropout: Optional[float] = 0.1, + max_len=5000, + ): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + self.embedding = nn.Embedding(max_len, d_model) + self.max_len = max_len + self.d_model = d_model + + def forward(self, x): + assert x.shape[1] <= self.max_len, f"Input length {x.shape[1]} is greater than max length {self.max_len}" + assert x.shape[2] == self.d_model, f"Input dimension {x.shape[2]} is not equal to d_model {self.d_model}" + assert x.ndim == 3, f"Input dimension {x.ndim} is not 3" + + positions = torch.arange(0, x.shape[1], device=x.device, dtype=torch.int32).unsqueeze(0) # [1, T] + x = x + self.embedding(positions) + return self.dropout(x) + + +class PositionalEncodingNegativeIndex(nn.Module): + """Non-learned positional encoding. + + The input indices can be negative. + """ + + def __init__( + self, + d_model: int, + dropout: Optional[float] = 0.1, + max_len: Optional[int] = 5000, + ): + """ + Args: + d_model (int): input dim + dropout (Optional[float] = 0.1): dropout probability on output + max_len (Optional[int] = 5000): maximum absolute index value, e.g. if max_len is 5000, the index can be in (-5000, 5000) + """ + super().__init__() + self.max_len = max_len + self.d_model = d_model + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + + # Note: have to replace torch.exp() and math.log() with torch.pow() + # due to MKL exp() and ln() throws floating point exceptions on certain CPUs + # see corresponding commit and MR + div_term = torch.pow(10000.0, -torch.arange(0, d_model, 2).float() / d_model) + # div_term = torch.exp( + # torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model) + # ) + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + + pe_negative = torch.zeros(max_len - 1, d_model) + pe_negative[:, 0::2] = torch.sin(-position[1:] * div_term) + pe_negative[:, 1::2] = torch.cos(-position[1:] * div_term) + + # reverse the pe_negative and concatenate with pe + pe_negative = torch.flip(pe_negative, dims=[0]) + pe = torch.cat([pe, pe_negative], dim=0) # [2T-1, D] + + self.register_buffer("pe", pe, persistent=False) + + def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Apply positional encoding to input sequence. + + Args: + x (torch.Tensor): [B, T, D] input motion sequence + index (torch.Tensor): [B, T] index for each position, can be negative + Returns: + torch.Tensor: [B, T, D] input motion with PE added to it (and optionally dropout) + """ + assert index.abs().max() < self.max_len, f"Index {index.abs().max()} is greater than max length {self.max_len}" + + # Convert negative indices to positive offsets into the pe buffer for tensorrt compatibility. + # pe layout: [0..max_len-1, reversed_negative(max_len..2*max_len-2)] + # Python negative indexing: pe[-k] == pe[len - k] + safe_index = torch.where(index >= 0, index, index + self.pe.shape[0]) + positional_encoding = F.embedding(safe_index.int(), self.pe) # [B, T, D] + x = x + positional_encoding + return self.dropout(x) diff --git a/ardy/model/cfg.py b/ardy/model/cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..2fa6f63b023503f8698cec26973edfddabc08644 --- /dev/null +++ b/ardy/model/cfg.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import torch.nn as nn + +CFG_TYPES = ["nocfg", "regular", "separated"] + + +def AutoLatentClassifierFreeGuidedModel(model: nn.Module, cfg_type: Optional[str] = "separated"): + if cfg_type == "nocfg": + return AutoLatentClassifierFreeGuidedModelNoCFG(model) + elif cfg_type == "regular": + return AutoLatentClassifierFreeGuidedModelRegular(model) + elif cfg_type == "separated": + return AutoLatentClassifierFreeGuidedModelSeparated(model) + raise ValueError(f"Unknown cfg_type {cfg_type!r}. Available: {CFG_TYPES}") + + +class AutoLatentClassifierFreeGuidedModelSeparated(nn.Module): + """Wrapper around denoiser to use classifier-free guidance at sampling time.""" + + def __init__(self, model: nn.Module): + """ + Args: + model (nn.Module): the denoiser to wrap in CFG + """ + super().__init__() + self.model = model + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward( + self, + cfg_weight_text: torch.Tensor, + cfg_weight_cstr: torch.Tensor, + x: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: torch.Tensor, + first_heading_angle: torch.Tensor, + motion_mask: torch.Tensor, + observed_motion: torch.Tensor, + ) -> torch.Tensor: + """ + Args: + cfg_weight (float): guidance weight float or tuple of floats with (text, constraint) weights if using separated cfg + x (torch.Tensor): [B, T, dim_motion] current noisy motion + x_pad_mask (torch.Tensor): [B, T] attention mask, positions with True are allowed to attend, False are not + text_feat (torch.Tensor): [B, max_text_len, llm_dim] embedded text prompts + text_feat_pad_mask (torch.Tensor): [B, max_text_len] attention mask, positions with True are allowed to attend, False are not + timesteps (torch.Tensor): [B,] current denoising step + motion_mask + observed_motion + + Returns: + torch.Tensor: same size as input x + """ + # ── CFG separated batching (B=1 → B=3) ── + # Pass 0: text-only (real text, zero constraints) + # Pass 1: constraint-only (zero text, real constraints) + # Pass 2: unconditional (zero text, zero constraints) + x_3 = torch.cat([x, x, x], dim=0) + history_len_3 = torch.cat([history_len, history_len, history_len], dim=0) + generation_len_3 = torch.cat([generation_len, generation_len, generation_len], dim=0) + future_len_3 = torch.cat([future_len, future_len, future_len], dim=0) + history_mask_3 = torch.cat([history_mask, history_mask, history_mask], dim=0) + generation_mask_3 = torch.cat([generation_mask, generation_mask, generation_mask], dim=0) + future_mask_3 = torch.cat([future_mask, future_mask, future_mask], dim=0) + history_token_mask_3 = torch.cat([history_token_mask, history_token_mask, history_token_mask], dim=0) + generation_token_mask_3 = torch.cat( + [generation_token_mask, generation_token_mask, generation_token_mask], + dim=0, + ) + future_token_mask_3 = torch.cat( + [ + 0 * future_token_mask, + future_token_mask, + 0 * future_token_mask, + ], + dim=0, + ) + text_feat_3 = torch.cat([text_feat, 0 * text_feat, 0 * text_feat], dim=0) + text_feat_pad_mask_3 = torch.cat( + [ + text_feat_pad_mask, + 0 * text_feat_pad_mask, + 0 * text_feat_pad_mask, + ], + dim=0, + ) + timesteps_3 = torch.cat([timesteps, timesteps, timesteps], dim=0) + # motion_mask / observed_motion are None for constraint-free generation + # and first_heading_angle may be None; pass None through (the inner + # denoiser fills in zeros for None constraints), matching the Regular + # variant. Only pass 1 (constraint-only) carries the real constraints. + first_heading_angle_3 = ( + torch.cat( + [first_heading_angle, first_heading_angle, first_heading_angle], + dim=0, + ) + if first_heading_angle is not None + else None + ) + motion_mask_3 = ( + torch.cat([0 * motion_mask, motion_mask, 0 * motion_mask], dim=0) if motion_mask is not None else None + ) + observed_motion_3 = ( + torch.cat([0 * observed_motion, observed_motion, 0 * observed_motion], dim=0) + if observed_motion is not None + else None + ) + + out_3 = self.model( + x=x_3, + history_len=history_len_3, + generation_len=generation_len_3, + future_len=future_len_3, + history_mask=history_mask_3 > 0.5, + generation_mask=generation_mask_3 > 0.5, + future_mask=future_mask_3 > 0.5, + history_token_mask=history_token_mask_3 > 0.5, + generation_token_mask=generation_token_mask_3 > 0.5, + future_token_mask=future_token_mask_3 > 0.5, + text_feat=text_feat_3, + text_feat_pad_mask=text_feat_pad_mask_3 > 0.5, + timesteps=timesteps_3, + first_heading_angle=first_heading_angle_3, + motion_mask=motion_mask_3, + observed_motion=observed_motion_3, + ) + + out_text, out_cstr, out_uncond = torch.chunk(out_3, 3, dim=0) + return out_uncond + cfg_weight_text * (out_text - out_uncond) + cfg_weight_cstr * (out_cstr - out_uncond) + + +class AutoLatentClassifierFreeGuidedModelRegular(nn.Module): + """Regular (single-weight) classifier-free guidance at sampling time.""" + + def __init__(self, model: nn.Module): + super().__init__() + self.model = model + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward( + self, + cfg_weight_text: torch.Tensor, + cfg_weight_cstr: torch.Tensor, + x: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: torch.Tensor, + first_heading_angle: torch.Tensor = None, + motion_mask: torch.Tensor = None, + observed_motion: torch.Tensor = None, + ) -> torch.Tensor: + """Regular CFG: one conditional (real text + constraints) pass and one unconditional pass (B + -> 2B). + + Uses cfg_weight_text as the single guidance weight; cfg_weight_cstr is accepted for a + uniform API but unused here. + """ + # Pass 0: conditional, Pass 1: unconditional (zero text + zero constraints) + text_feat = torch.cat([text_feat, 0 * text_feat], dim=0) + if motion_mask is not None: + motion_mask = torch.cat([motion_mask, 0 * motion_mask], dim=0) + if observed_motion is not None: + observed_motion = torch.cat([observed_motion, 0 * observed_motion], dim=0) + if first_heading_angle is not None: + first_heading_angle = torch.cat([first_heading_angle, first_heading_angle], dim=0) + + out_cond_uncond = self.model( + torch.cat([x, x], dim=0), + torch.cat([history_len, history_len], dim=0), + torch.cat([generation_len, generation_len], dim=0), + torch.cat([future_len, future_len], dim=0), + torch.cat([history_mask, history_mask], dim=0), + torch.cat([generation_mask, generation_mask], dim=0), + torch.cat([future_mask, future_mask], dim=0), + torch.cat([history_token_mask, history_token_mask], dim=0), + torch.cat([generation_token_mask, generation_token_mask], dim=0), + torch.cat([future_token_mask, False * future_token_mask], dim=0) if future_token_mask is not None else None, + text_feat, + torch.cat([text_feat_pad_mask, False * text_feat_pad_mask], dim=0), + torch.cat([timesteps, timesteps], dim=0), + first_heading_angle=first_heading_angle, + motion_mask=motion_mask, + observed_motion=observed_motion, + ) + + out, out_uncond = torch.chunk(out_cond_uncond, 2) + return out_uncond + cfg_weight_text * (out - out_uncond) + + +class AutoLatentClassifierFreeGuidedModelNoCFG(nn.Module): + """No classifier-free guidance: a single conditional denoiser pass.""" + + def __init__(self, model: nn.Module): + super().__init__() + self.model = model + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward( + self, + cfg_weight_text: torch.Tensor, + cfg_weight_cstr: torch.Tensor, + x: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: torch.Tensor, + first_heading_angle: torch.Tensor = None, + motion_mask: torch.Tensor = None, + observed_motion: torch.Tensor = None, + ) -> torch.Tensor: + """No guidance: run the denoiser once. + + cfg_weight_text / cfg_weight_cstr are accepted for a uniform API but unused. + """ + return self.model( + x, + history_len, + generation_len, + future_len, + history_mask, + generation_mask, + future_mask, + history_token_mask, + generation_token_mask, + future_token_mask, + text_feat, + text_feat_pad_mask, + timesteps, + first_heading_angle=first_heading_angle, + motion_mask=motion_mask, + observed_motion=observed_motion, + ) diff --git a/ardy/model/diffusion.py b/ardy/model/diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..69035ae0c482a602e3cfa8691c8f99c7450c0768 --- /dev/null +++ b/ardy/model/diffusion.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Diffusion process and DDIM sampling for motion generation.""" + +import math +from typing import Optional, Tuple + +import torch +from torch import nn + + +def get_beta_schedule( + num_diffusion_timesteps: int, + max_beta: Optional[float] = 0.999, +) -> torch.Tensor: + """Get cosine beta schedule.""" + + def alpha_bar(t): + return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return torch.tensor(betas, dtype=torch.float) + + +class Diffusion(torch.nn.Module): + """Cosine-schedule diffusion process: betas, alphas, and DDIM step mapping.""" + + def __init__(self, num_base_steps: int): + """Set up cosine beta schedule and precompute diffusion variables for num_base_steps.""" + super().__init__() + self.num_base_steps = num_base_steps + betas_base = get_beta_schedule(self.num_base_steps) + self.register_buffer("betas_base", betas_base, persistent=False) + alphas_cumprod_base = torch.cumprod(1.0 - self.betas_base, dim=0) + self.register_buffer("alphas_cumprod_base", alphas_cumprod_base, persistent=False) + use_timesteps, _ = self.space_timesteps(self.num_base_steps) + self.calc_diffusion_vars(use_timesteps) + + def extra_repr(self) -> str: + return f"num_base_steps={self.num_base_steps}" + + @property + def device(self): + return self.betas_base.device + + def space_timesteps(self, num_denoising_steps: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (use_timesteps, map_tensor) for a subsampled denoising schedule of + num_denoising_steps.""" + nsteps_train = self.num_base_steps + frac_stride = (nsteps_train - 1) / max(1, num_denoising_steps - 1) + use_timesteps = torch.round(torch.arange(nsteps_train, device=self.device) * frac_stride).to(torch.long) + use_timesteps = torch.clamp(use_timesteps, max=nsteps_train - 1) + map_tensor = torch.arange(nsteps_train, device=self.device, dtype=torch.long)[use_timesteps] + return use_timesteps, map_tensor + + def calc_diffusion_vars(self, use_timesteps: torch.Tensor) -> None: + """Update buffers (betas, alphas, alphas_cumprod, etc.) for the given subsampled + timesteps.""" + alphas_cumprod = self.alphas_cumprod_base[use_timesteps] + last_alpha_cumprod = torch.cat([torch.tensor([1.0]).to(alphas_cumprod), alphas_cumprod[:-1]]) + betas = 1.0 - alphas_cumprod / last_alpha_cumprod + self.register_buffer("betas", betas, persistent=False) + + alphas = 1.0 - self.betas + self.register_buffer("alphas", alphas, persistent=False) + alphas_cumprod = torch.cumprod(self.alphas, dim=0) + alphas_cumprod = torch.clamp(alphas_cumprod, min=1e-9) + self.register_buffer("alphas_cumprod", alphas_cumprod, persistent=False) + + alphas_cumprod_prev = torch.cat([torch.tensor([1.0]).to(self.alphas_cumprod), self.alphas_cumprod[:-1]]) + self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev, persistent=False) + + sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod) + self.register_buffer("sqrt_recip_alphas_cumprod", sqrt_recip_alphas_cumprod, persistent=False) + + sqrt_recipm1_alphas_cumprod = torch.rsqrt(self.alphas_cumprod / (1.0 - self.alphas_cumprod)) + self.register_buffer("sqrt_recipm1_alphas_cumprod", sqrt_recipm1_alphas_cumprod, persistent=False) + + posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + self.register_buffer("posterior_variance", posterior_variance, persistent=False) + + sqrt_alphas_cumprod = torch.rsqrt(1.0 / self.alphas_cumprod) + self.register_buffer("sqrt_alphas_cumprod", sqrt_alphas_cumprod, persistent=False) + + sqrt_one_minus_alphas_cumprod = torch.rsqrt(1.0 / (1.0 - self.alphas_cumprod)) + self.register_buffer( + "sqrt_one_minus_alphas_cumprod", + sqrt_one_minus_alphas_cumprod, + persistent=False, + ) + + def q_sample( + self, + x_start: torch.Tensor, + t: torch.Tensor, + noise: torch.Tensor = None, + ): + if noise is None: + noise = torch.randn_like(x_start) + assert noise.shape == x_start.shape + + xt = ( + self.sqrt_alphas_cumprod[t, None, None] * x_start + + self.sqrt_one_minus_alphas_cumprod[t, None, None] * noise + ) + return xt + + +class DDIMSampler(nn.Module): + """Deterministic DDIM sampler (eta = 0).""" + + def __init__(self, diffusion: Diffusion): + super().__init__() + self.diffusion = diffusion + + def __call__( + self, + x_t: torch.Tensor, + pred_xstart: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + # Assumes the diffusion buffers are current for the active schedule + # (denoising_step refreshes them every step before calling). + eps = ( + self.diffusion.sqrt_recip_alphas_cumprod[t, None, None] * x_t - pred_xstart + ) / self.diffusion.sqrt_recipm1_alphas_cumprod[t, None, None] + alpha_bar_prev = self.diffusion.alphas_cumprod_prev[t, None, None] + x = pred_xstart * torch.sqrt(alpha_bar_prev) + torch.sqrt(1 - alpha_bar_prev) * eps + return x diff --git a/ardy/model/latent_utils.py b/ardy/model/latent_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bcd2957a77ed4f62cc90dee9e555b2e49c8fc008 --- /dev/null +++ b/ardy/model/latent_utils.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared utility functions for latent motion models. + +These functions were originally in latent_ardy_train.py and are used by auto_latent training and +test-time code. +""" + +import logging +from functools import cached_property + +import torch + +log = logging.getLogger(__name__) + + +class HybridMotionConverter: + """Converts between explicit motion and the hybrid (root + latent token) rep. + + Bundles the ``denoiser``, ``autoencoder`` and ``motion_rep`` that every conversion needs so call + sites don't have to thread them through manually. Derived attributes (``num_frames_per_token``, + ``motion_rep``) are resolved lazily and raise a clear error if the required object is missing, + instead of half-initializing. + """ + + def __init__(self, *, denoiser=None, autoencoder=None, motion_rep=None, gen_horizon_len=None): + if denoiser is None and autoencoder is None: + raise ValueError("HybridMotionConverter needs a denoiser and/or an autoencoder") + self.denoiser = denoiser + self.autoencoder = autoencoder + self._motion_rep = motion_rep + self.gen_horizon_len = gen_horizon_len + + @classmethod + def from_model(cls, ardy_model): + """Build from an object exposing ``denoiser``/``autoencoder``/``motion_rep``.""" + return cls( + denoiser=ardy_model.denoiser, + autoencoder=ardy_model.autoencoder, + motion_rep=ardy_model.motion_rep, + gen_horizon_len=getattr(ardy_model, "gen_horizon_len", None), + ) + + @cached_property + def motion_rep(self): + rep = self._motion_rep + if rep is None and self.denoiser is not None: + rep = self.denoiser.motion_rep + if rep is None and self.autoencoder is not None: + rep = self.autoencoder.motion_rep + if rep is None: + raise ValueError("No motion_rep available on denoiser or autoencoder") + return rep + + @cached_property + def num_frames_per_token(self): + if self.autoencoder is not None: + return self.autoencoder.num_frames_per_token + return self.denoiser.num_frames_per_token + + def get_num_frames_from_hybrid(self, hybrid_motion): + num_tokens = hybrid_motion.shape[1] + num_frames = self.num_frames_per_token * num_tokens + return num_frames + + def get_root_and_latent_body_motion_from_hybrid(self, hybrid_motion): + num_frames = self.get_num_frames_from_hybrid(hybrid_motion) + bs = hybrid_motion.shape[0] + + root_motion = hybrid_motion[ + :, :, : self.denoiser.nframe_root_dim + ] # (B, num_frames // num_frames_per_token, dim_root * num_frames_per_token) + root_motion = root_motion.reshape(bs, num_frames, self.motion_rep.motion_root_dim) # (B, num_frames, dim_root) + latent_body_motion = hybrid_motion[:, :, self.denoiser.nframe_root_dim :] # (B, num_tokens, dim_latent_body) + return root_motion, latent_body_motion + + def get_explicit_motion_from_hybrid( + self, + hybrid_motion, + motion_pad_mask, + motion_len, + motion_mask=None, + ): + # Separate the root and the latents + global_root_motion, latent_body_motion = self.get_root_and_latent_body_motion_from_hybrid(hybrid_motion) + # Compute the local root as a condition + local_root_motion = self.motion_rep.global_root_to_local_root( + global_root_motion, normalized=True, lengths=motion_len + ) + + # Detokenize the tokens + output = self.autoencoder.detokenize( + latent_body_motion, + external_cond=local_root_motion, + motion_pad_mask=motion_pad_mask, + ) + decoded_body_motion = output["body"] + # Construct back the motion features + motion = self.motion_rep.concat_root_body(global_root_motion, decoded_body_motion) + return motion + + def get_explicit_motion_from_hybrid_autoregressive( + self, + hybrid_motion, + motion_pad_mask, + motion_len, + motion_mask=None, + crop_history_length=None, + ): + if self.gen_horizon_len is None: + raise ValueError("gen_horizon_len is required for autoregressive decoding (see from_model)") + gen_horizon_len = self.gen_horizon_len + num_frames_per_token = self.num_frames_per_token + num_generation_tokens = gen_horizon_len // num_frames_per_token + token_len = hybrid_motion.shape[1] + results = [] + for token_idx in range(0, token_len, num_generation_tokens): + generation_token_end = min(token_idx + num_generation_tokens, token_len) + generation_token_start = token_idx + history_token_start = ( + max(0, token_idx - crop_history_length // num_frames_per_token) + if crop_history_length is not None + else 0 + ) + generation_frame_start = generation_token_start * num_frames_per_token + generation_frame_end = generation_token_end * num_frames_per_token + history_frame_start = history_token_start * num_frames_per_token + explicit_motion = self.get_explicit_motion_from_hybrid( + hybrid_motion[:, history_token_start:generation_token_end], + motion_pad_mask[:, history_frame_start:generation_frame_end], + motion_pad_mask[:, history_frame_start:generation_frame_end].sum(dim=-1), + motion_mask=motion_mask[:, history_frame_start:generation_frame_end] + if motion_mask is not None + else None, + ) + results.append( + explicit_motion[ + :, + generation_frame_start - history_frame_start : generation_frame_end - history_frame_start, + ] + ) + motion = torch.cat(results, dim=1) + return motion + + def get_hybrid_motion_from_root_and_latent_body_motion( + self, + root_motion, + latent_body_motion, + ): + bs = root_motion.shape[0] + nframe_root_dim = self.denoiser.nframe_root_dim + num_tokens = latent_body_motion.shape[1] + hybrid_motion = torch.cat( + [root_motion.reshape(bs, num_tokens, nframe_root_dim), latent_body_motion], + dim=-1, + ) + return hybrid_motion + + def get_hybrid_motion_from_explicit( + self, + motion, + motion_len, + motion_pad_mask, + ): + token_embedding = self.autoencoder.tokenize(motion, motion_pad_mask=motion_pad_mask) # (B, T, D) + + # ensure masked tokens are set to 0 + bs, num_frames = motion_pad_mask.shape + num_frames_per_token = self.num_frames_per_token + + motion_token_mask = motion_pad_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all( + dim=-1 + ) # (B, num_frames // num_frames_per_token) + token_embedding = token_embedding * motion_token_mask[:, :, None] + + root = self.motion_rep.extract_root(motion) + root_reshape = root.reshape( + bs, + num_frames // num_frames_per_token, + self.motion_rep.motion_root_dim * num_frames_per_token, + ) # (B, num_frames // num_frames_per_token, dim_root * num_frames_per_token) + hybrid_motion = torch.cat( + [root_reshape, token_embedding], dim=-1 + ) # (B, num_frames // num_frames_per_token, dim_hybrid) + hybrid_motion_pad_mask = motion_pad_mask.reshape( + bs, num_frames // num_frames_per_token, num_frames_per_token + ).all(dim=-1) # (B, num_frames // num_frames_per_token) + + return hybrid_motion, hybrid_motion_pad_mask + + def convert_frame_mask_to_token_mask( + self, + history_mask, + generation_mask, + future_mask, + motion_mask=None, + ): + num_frames_per_token = self.num_frames_per_token + bs, num_frames = history_mask.shape[0], history_mask.shape[1] + history_token_mask = history_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all( + dim=-1 + ) # (B, num_frames // num_frames_per_token) + generation_token_mask = generation_mask.reshape( + bs, num_frames // num_frames_per_token, num_frames_per_token + ).all(dim=-1) # (B, num_frames // num_frames_per_token) + future_token_mask = torch.zeros_like(history_token_mask) + if motion_mask is not None: + future_token_mask = future_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all( + dim=-1 + ) # (B, num_frames // num_frames_per_token) + has_observation = motion_mask.any(dim=-1) # if any feature is observed, the token is valid + has_observation = has_observation.reshape( + bs, num_frames // num_frames_per_token, num_frames_per_token + ).any( + dim=-1 + ) # (B, num_frames // num_frames_per_token), future constraints are sparse, any frame with constraints make the token valid + future_token_mask = future_token_mask & has_observation + + return history_token_mask, generation_token_mask, future_token_mask diff --git a/ardy/model/llm2vec/README.md b/ardy/model/llm2vec/README.md new file mode 100644 index 0000000000000000000000000000000000000000..450e995c0aa9cdb0f1117aa33cb30a30d84a41c0 --- /dev/null +++ b/ardy/model/llm2vec/README.md @@ -0,0 +1 @@ +This is a patched version of the original [LLM2Vec](https://github.com/McGill-NLP/llm2vec) codebase so that `McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised` works with `transformers==5.0.0rc3`. diff --git a/ardy/model/llm2vec/__init__.py b/ardy/model/llm2vec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..53d7a92811d75bf5e772e1eb2d5b18ce1b63eef4 --- /dev/null +++ b/ardy/model/llm2vec/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LLM2Vec text encoder and wrapper for ARDY.""" + +from .llm2vec import LLM2Vec +from .llm2vec_wrapper import LLM2VecEncoder + +__all__ = [ + "LLM2Vec", + "LLM2VecEncoder", +] diff --git a/ardy/model/llm2vec/llm2vec.py b/ardy/model/llm2vec/llm2vec.py new file mode 100644 index 0000000000000000000000000000000000000000..3fa009a66602eb9c7c1259620a7e51405afe8a1a --- /dev/null +++ b/ardy/model/llm2vec/llm2vec.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import os +from functools import partial +from typing import Dict, List, Optional, Union + +import numpy as np +import torch +import torch.multiprocessing as mp +from peft import PeftModel +from torch import Tensor, device, nn +from tqdm.autonotebook import tqdm, trange +from transformers import ( + AutoConfig, + AutoModel, + AutoTokenizer, + GemmaConfig, + LlamaConfig, + MistralConfig, + PretrainedConfig, + Qwen2Config, +) + +logger = logging.getLogger(__name__) + + +def batch_to_device(batch, target_device: device): + """Send a pytorch batch to a device (CPU/GPU)""" + for key in batch: + if isinstance(batch[key], Tensor): + batch[key] = batch[key].to(target_device) + return batch + + +class LLM2Vec(nn.Module): + def __init__( + self, + model: AutoModel, + tokenizer: AutoTokenizer, + prompt_model_family: Optional[str] = None, + pooling_mode: str = "mean", + max_length: int = 512, + doc_max_length: int = 400, + skip_instruction: bool = True, + ): + super().__init__() + self.model = model + self.tokenizer = tokenizer + self.prompt_model_family = prompt_model_family + self.pooling_mode = pooling_mode + self.skip_instruction = skip_instruction + self.max_length = max_length + self.doc_max_length = doc_max_length + self.config = model.config + + @classmethod + def _get_model_class(cls, config_class_name, enable_bidirectional): + if not enable_bidirectional: + return AutoModel + if config_class_name == "MistralConfig": + from .models.bidirectional_mistral import MistralBiModel + + return MistralBiModel + elif config_class_name == "LlamaConfig": + from .models.bidirectional_llama import LlamaBiModel + + return LlamaBiModel + elif config_class_name == "GemmaConfig": + from .models.bidirectional_gemma import GemmaBiModel + + return GemmaBiModel + elif config_class_name == "Qwen2Config": + from .models.bidirectional_qwen2 import Qwen2BiModel + + return Qwen2BiModel + else: + raise ValueError(f"{config_class_name} is not supported yet with bidirectional models.") + + @classmethod + def from_pretrained( + cls, + base_model_name_or_path, + mntp_model_name_or_path=None, + peft_model_name_or_path=None, + prompt_model_family=None, + merge_peft=False, + enable_bidirectional=True, + **kwargs, + ): + # pop out encoder args + keys = ["pooling_mode", "max_length", "doc_max_length", "skip_instruction"] + encoder_args = {key: kwargs.pop(key, None) for key in keys if kwargs.get(key) is not None} + + tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path) + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "left" + + config = AutoConfig.from_pretrained(base_model_name_or_path) + config_class_name = config.__class__.__name__ + + model_class = cls._get_model_class(config_class_name, enable_bidirectional=enable_bidirectional) + + model = model_class.from_pretrained(base_model_name_or_path, **kwargs) + + if os.path.isdir(base_model_name_or_path) and os.path.exists(f"{base_model_name_or_path}/config.json"): + with open(f"{base_model_name_or_path}/config.json", "r") as fIn: + config_dict = json.load(fIn) + config = PretrainedConfig.from_dict(config_dict) + model.config._name_or_path = config._name_or_path + + # For special case where config.json and adapter weights are in the same directory + if hasattr(model, "peft_config"): + model = PeftModel.from_pretrained( + model, + base_model_name_or_path, + ) + model = model.merge_and_unload() + + if mntp_model_name_or_path is not None: + model = PeftModel.from_pretrained( + model, + mntp_model_name_or_path, + ) + model = model.merge_and_unload() + + if peft_model_name_or_path is not None: + model = PeftModel.from_pretrained( + model, + peft_model_name_or_path, + ) + if merge_peft: + model = model.merge_and_unload() + + config = {} + config_addr = peft_model_name_or_path if peft_model_name_or_path is not None else base_model_name_or_path + if os.path.exists(f"{config_addr}/llm2vec_config.json"): + with open(f"{config_addr}/llm2vec_config.json", "r") as fIn: + llm2vec_config = json.load(fIn) + config.update(llm2vec_config) + + for key, value in encoder_args.items(): + config[key] = value + + return cls( + model=model, + tokenizer=tokenizer, + prompt_model_family=prompt_model_family, + **config, + ) + + def prepare_for_tokenization(self, text): + if ( + self.prompt_model_family == "llama3-instruct" + or self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B-Instruct" + ): + text = "<|start_header_id|>user<|end_header_id|>\n\n" + text.strip() + "<|eot_id|>" + return text + if self.model.config._name_or_path in [ + "mistralai/Mistral-7B-Instruct-v0.2", + "meta-llama/Llama-2-7b-chat-hf", + ]: + text = "[INST] " + text.strip() + " [/INST]" + if self.model.config._name_or_path in [ + "google/gemma-2-9b-it", + ]: + text = "user\n" + text.strip() + "" + if self.model.config._name_or_path in [ + "Qwen/Qwen2-1.5B-Instruct", + "Qwen/Qwen2-7B-Instruct", + ]: + text = "<|im_start|>user\n" + text.strip() + "<|im_end|>" + if self.pooling_mode == "eos_token": + if self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B": + text = text.strip() + "<|end_of_text|>" + elif isinstance(self.model.config, LlamaConfig) or isinstance(self.model.config, MistralConfig): + text = text.strip() + " " + elif isinstance(self.model.config, GemmaConfig): + text = text.strip() + "" + elif isinstance(self.model.config, Qwen2Config): + text = text.strip() + "<|endoftext|>" + return text + + def tokenize(self, texts): + texts_2 = [] + original_texts = [] + for text in texts: + t = text.split("!@#$%^&*()") + texts_2.append(t[1] if len(t) > 1 else "") + original_texts.append("".join(t)) + + original = self.tokenizer( + original_texts, + return_tensors="pt", + padding=True, + truncation=True, + max_length=self.max_length, + ) + embed_mask = None + for t_i, t in enumerate(texts_2): + ids = self.tokenizer( + [t], + return_tensors="pt", + padding=True, + truncation=True, + max_length=self.max_length, + add_special_tokens=False, + ) + if embed_mask is None: + e_m = torch.zeros_like(original["attention_mask"][t_i]) + if len(ids["input_ids"][0]) > 0: + e_m[-len(ids["input_ids"][0]) :] = torch.ones(len(ids["input_ids"][0])) + embed_mask = e_m.unsqueeze(0) + else: + e_m = torch.zeros_like(original["attention_mask"][t_i]) + if len(ids["input_ids"][0]) > 0: + e_m[-len(ids["input_ids"][0]) :] = torch.ones(len(ids["input_ids"][0])) + embed_mask = torch.cat((embed_mask, e_m.unsqueeze(0)), dim=0) + + original["embed_mask"] = embed_mask + return original + + def _skip_instruction(self, sentence_feature): + assert sentence_feature["attention_mask"].shape == sentence_feature["embed_mask"].shape + sentence_feature["attention_mask"] = sentence_feature["embed_mask"] + + def forward(self, sentence_feature: Dict[str, Tensor]): + embed_mask = None + if "embed_mask" in sentence_feature: + embed_mask = sentence_feature.pop("embed_mask") + reps = self.model(**sentence_feature) + sentence_feature["embed_mask"] = embed_mask + + return self.get_pooling(sentence_feature, reps.last_hidden_state) + + def get_pooling(self, features, last_hidden_states): # All models padded from left + assert self.tokenizer.padding_side == "left", "Pooling modes are implemented for padding from left." + if self.skip_instruction: + self._skip_instruction(features) + seq_lengths = features["attention_mask"].sum(dim=-1) + if self.pooling_mode == "mean": + return torch.stack( + [last_hidden_states[i, -length:, :].mean(dim=0) for i, length in enumerate(seq_lengths)], + dim=0, + ) + elif self.pooling_mode == "weighted_mean": + bs, l, _ = last_hidden_states.shape + complete_weights = torch.zeros(bs, l, device=last_hidden_states.device) + for i, seq_l in enumerate(seq_lengths): + if seq_l > 0: + complete_weights[i, -seq_l:] = torch.arange(seq_l) + 1 + complete_weights[i] /= torch.clamp(complete_weights[i].sum(), min=1e-9) + return torch.sum(last_hidden_states * complete_weights.unsqueeze(-1), dim=1) + elif self.pooling_mode == "eos_token" or self.pooling_mode == "last_token": + return last_hidden_states[:, -1] + elif self.pooling_mode == "bos_token": + return last_hidden_states[features["input_ids"] == self.tokenizer.bos_token_id] + else: + raise ValueError(f"{self.pooling_mode} is not implemented yet.") + + def _convert_to_str(self, instruction, text): + tokenized_q = self.tokenizer( + text, + return_tensors="pt", + padding=True, + truncation=True, + max_length=self.max_length, + add_special_tokens=False, + ) + tokenized_q_length = len(tokenized_q["input_ids"][0]) + + while tokenized_q_length > self.doc_max_length: + reduction_ratio = self.doc_max_length / tokenized_q_length + reduced_length = int(len(text.split()) * reduction_ratio) + text = " ".join(text.split()[:reduced_length]) + tokenized_q = self.tokenizer( + text, + return_tensors="pt", + padding=True, + truncation=True, + max_length=self.max_length, + add_special_tokens=False, + ) + tokenized_q_length = len(tokenized_q["input_ids"][0]) + + return f"{instruction.strip()} !@#$%^&*(){text}" if instruction else f"!@#$%^&*(){text}" + + def encode( + self, + sentences: Union[str, List[str]], + batch_size: int = 32, + show_progress_bar: bool = True, + convert_to_numpy: bool = False, + convert_to_tensor: bool = False, + device: Optional[str] = None, + ): + """Encode a list of sentences to their respective embeddings. + + The sentences can be a list of strings or a string. Args: + sentences: sentence or sentences to encode. + batch_size: batch size for turning sentence tokens into embeddings. + show_progress_bar: whether to show progress bars during encoding steps. + convert_to_numpy: If true, return numpy arrays instead of torch tensors. + convert_to_tensor: If true, return torch tensors (default). + device: torch backend device identifier (e.g., 'cuda', 'cpu','mps' etc.). If not specified, + the default is to use cuda when available, otherwise cpu. Note that only the choice of 'cuda' supports + multiprocessing as currently implemented. + + Returns: embeddings of the sentences. Embeddings are detached and always on the CPU (see _encode implementation). + """ + if isinstance(sentences[0], str) and isinstance(sentences[-1], int): + sentences = [sentences] + # required for MEDI version of MTEB + if isinstance(sentences[0], str): + sentences = [[""] + [sentence] for sentence in sentences] + + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + concatenated_input_texts = [] + for sentence in sentences: + assert isinstance(sentence[0], str) + assert isinstance(sentence[1], str) + concatenated_input_texts.append(self._convert_to_str(sentence[0], sentence[1])) + sentences = concatenated_input_texts + + self.eval() + + if convert_to_tensor: + convert_to_numpy = False + + length_sorted_idx = np.argsort([-self._text_length(sen) for sen in sentences]) + sentences_sorted = [sentences[idx] for idx in length_sorted_idx] + all_embeddings = [] + + if torch.cuda.device_count() <= 1: + # This branch also support mps devices + self.to(device) + for start_index in trange( + 0, + len(sentences), + batch_size, + desc="Batches", + disable=not show_progress_bar, + ): + sentences_batch = sentences_sorted[start_index : start_index + batch_size] + embeddings = self._encode(sentences_batch, device=device, convert_to_numpy=convert_to_numpy) + all_embeddings.append(embeddings) + else: + num_proc = torch.cuda.device_count() + cuda_compatible_multiprocess = mp.get_context("spawn") + with cuda_compatible_multiprocess.Pool(num_proc) as p: + sentences_batches = [ + sentences_sorted[start_index : start_index + batch_size] + for start_index in range(0, len(sentences), batch_size) + ] + + progress_bar = tqdm( + total=len(sentences_batches), + desc="Batches", + disable=not show_progress_bar, + ) + results = [] + + def update(*args): + progress_bar.update() + + for batch in sentences_batches: + results.append( + p.apply_async( + self._encode, + args=(batch, None, convert_to_numpy, True), + callback=update, + ) + ) + + all_embeddings = [result.get() for result in results] + progress_bar.close() + + all_embeddings = torch.cat(all_embeddings, dim=0) + all_embeddings = all_embeddings[np.argsort(length_sorted_idx)] + all_embeddings = all_embeddings.to(torch.float32) + if convert_to_numpy: + all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings]) + return all_embeddings + + def save(self, output_path, merge_before_save=False, save_config=True): + if merge_before_save and isinstance(self.model, PeftModel): + self.model = self.model.merge_and_unload() + # Fixes the issue of saving - https://huggingface.co/McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-unsup-simcse/discussions/1 + if hasattr(self.model, "_hf_peft_config_loaded"): + self.model._hf_peft_config_loaded = False + + self.model.save_pretrained(output_path) + self.tokenizer.save_pretrained(output_path) + + llm2vec_config = { + "pooling_mode": self.pooling_mode, + "max_length": self.max_length, + "doc_max_length": self.doc_max_length, + "skip_instruction": self.skip_instruction, + } + + if save_config: + os.makedirs(output_path, exist_ok=True) + with open(f"{output_path}/llm2vec_config.json", "w") as fOut: + json.dump(llm2vec_config, fOut, indent=4) + + def _encode( + self, + sentences_batch, + device: Optional[str] = None, + convert_to_numpy: bool = False, + multiprocessing=False, + ): + if multiprocessing: + # multiprocessing only supports CUDA devices at this time, so we ignore the value of device + # and use cuda:rank for the device + rank = mp.current_process()._identity[0] + if device is None and torch.cuda.is_available(): + device = f"cuda:{rank % torch.cuda.device_count()}" + + self.to(device) + features = self.tokenize([self.prepare_for_tokenization(sentence) for sentence in sentences_batch]) + features = batch_to_device(features, device) + + with torch.no_grad(): + embeddings = self.forward(features) + embeddings = embeddings.detach() + embeddings = embeddings.cpu() + + return embeddings + + def _text_length(self, text: Union[List[int], List[List[int]]]): + """Help function to get the length for the input text. + + Text can be either a string (which means a single text) a list of ints (which means a single + tokenized text), or a tuple of list of ints (representing several text inputs to the model). + """ + if ( + isinstance(text, str) or (isinstance(text, list) and isinstance(text[0], int)) or len(text) == 0 + ): # Single text, list of ints, or empty + return len(text) + if isinstance(text, dict): # {key: value} case + return len(next(iter(text.values()))) + elif not hasattr(text, "__len__"): # Object has no len() method + return 1 + else: + return sum([len(t) for t in text]) + + def resize_token_embeddings( + self, + new_num_tokens: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + ) -> nn.Embedding: + return self.model.resize_token_embeddings(new_num_tokens=new_num_tokens, pad_to_multiple_of=pad_to_multiple_of) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs) diff --git a/ardy/model/llm2vec/llm2vec_wrapper.py b/ardy/model/llm2vec/llm2vec_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..c4a82d26ab2fb4665fc1ca84d0df0fadea22193d --- /dev/null +++ b/ardy/model/llm2vec/llm2vec_wrapper.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LLM2Vec encoder wrapper for ARDY text conditioning.""" + +import os + +import numpy as np +import torch + +from .llm2vec import LLM2Vec + + +class LLM2VecEncoder: + """LLM2Vec text embeddings.""" + + def __init__( + self, + base_model_name_or_path: str, + peft_model_name_or_path: str, + dtype: str, + llm_dim: int, + device: str = "auto", + mntp_model_name_or_path: str | None = None, + prompt_model_family: str | None = None, + ) -> None: + torch_dtype = getattr(torch, dtype) + self.llm_dim = llm_dim + + cache_dir = os.environ.get("HUGGINGFACE_CACHE_DIR") + + if "TEXT_ENCODERS_DIR" in os.environ: + base_model_name_or_path = os.path.join(os.environ["TEXT_ENCODERS_DIR"], base_model_name_or_path) + if mntp_model_name_or_path is not None: + mntp_model_name_or_path = os.path.join( + os.environ["TEXT_ENCODERS_DIR"], + mntp_model_name_or_path, + ) + peft_model_name_or_path = os.path.join(os.environ["TEXT_ENCODERS_DIR"], peft_model_name_or_path) + + self.model = LLM2Vec.from_pretrained( + base_model_name_or_path=base_model_name_or_path, + mntp_model_name_or_path=mntp_model_name_or_path, + peft_model_name_or_path=peft_model_name_or_path, + prompt_model_family=prompt_model_family, + torch_dtype=torch_dtype, + cache_dir=cache_dir, + ) + + env_device = os.environ.get("TEXT_ENCODER_DEVICE") + if env_device: + device = env_device + if device == "auto": + device = "cuda" if torch.cuda.is_available() else "cpu" + self._device = device + if device is not None: + self.model = self.model.to(device) + + self.model.eval() + for p in self.model.parameters(): + p.requires_grad = False + + def to(self, device: torch.device | str | None = None, dtype: torch.dtype | None = None): + if device is not None and dtype is not None: + self.model = self.model.to(device=device, dtype=dtype) + elif device is not None: + self.model = self.model.to(device) + elif dtype is not None: + self.model = self.model.to(dtype=dtype) + if device is not None: + self._device = str(device) if not isinstance(device, str) else device + return self + + def eval(self): + self.model.eval() + return self + + def get_device(self): + return self.model.model.device + + def __call__(self, text: list[str] | str): + is_string = False + if isinstance(text, str): + text = [text] + is_string = True + + with torch.no_grad(): + encoded_text = self.model.encode( + text, + # IMPORTANT: different batch sizes unexpectedly change the output embeddings, so we always set it to 1 + # here for repeatability no matter how many texts are being encoded. This + # is a fundamental issue with transformers, and is especially bad at lower + # precisions (https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535) + # note: this is an internal batch size used by llm2vec - the text list can still be of arbitrary length. + batch_size=1, + show_progress_bar=False, + device=self._device, + ) + + assert len(encoded_text.shape) + assert self.llm_dim == encoded_text.shape[-1] + + encoded_text = encoded_text[:, None] + lengths = np.ones(len(encoded_text), dtype=int).tolist() + + if is_string: + encoded_text = encoded_text[0] + lengths = lengths[0] + + encoded_text = torch.tensor(encoded_text).to(self._device) + return encoded_text, lengths diff --git a/ardy/model/llm2vec/models/__init__.py b/ardy/model/llm2vec/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd5bbbf7881c1d8a258493735d2f334dc791097 --- /dev/null +++ b/ardy/model/llm2vec/models/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +# from .bidirectional_gemma import GemmaBiForMNTP, GemmaBiModel +# from .bidirectional_llama import LlamaBiForMNTP, LlamaBiModel +# from .bidirectional_mistral import MistralBiForMNTP, MistralBiModel +# from .bidirectional_qwen2 import Qwen2BiForMNTP, Qwen2BiModel diff --git a/ardy/model/llm2vec/models/attn_mask_utils.py b/ardy/model/llm2vec/models/attn_mask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eaa2752cdd7c22dbf1e164aa2b355b498586f02f --- /dev/null +++ b/ardy/model/llm2vec/models/attn_mask_utils.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Optional, Tuple, Union + +import torch +from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D + mask of shape `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter( + is_causal=False, sliding_window=sliding_window + ) # is_causal=True in original implementation + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + key_value_length=key_value_length, + dtype=inputs_embeds.dtype, + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], + input_shape[-1], + key_value_length, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """Prepares the correct `attn_mask` argument to be used by + `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the + cases `query_length == 1` and `key_value_length == query_length`, and rely instead on SDPA + `is_causal` argument to use causal/non-causal masks, allowing to dispatch to the flash attention + kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter( + is_causal=False, sliding_window=sliding_window + ) # is_causal=True in original implementation + + key_value_length = input_shape[-1] + past_key_values_length + batch_size, query_length = input_shape + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400). + is_tracing = ( + torch.jit.is_tracing() + or isinstance(inputs_embeds, torch.fx.Proxy) + or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling()) + ) + + if attention_mask is not None: + # 4d mask is passed through + if len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype) + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + return attention_mask + + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + attention_mask = None + elif key_value_length == query_length: + attention_mask = None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + pass + elif query_length > 1 and key_value_length != query_length: + # See the comment above (https://github.com/pytorch/pytorch/issues/108108). + # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`. + attention_mask = True + elif is_tracing: + raise ValueError( + 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.' + ) + + if attention_mask is None: + expanded_4d_mask = None + elif attention_mask is True: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], + input_shape[-1], + key_value_length, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # Attend to all tokens in masked rows from the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + if not is_tracing and expanded_4d_mask.device.type == "cuda": + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, min_dtype=torch.finfo(inputs_embeds.dtype).min + ) + + return expanded_4d_mask diff --git a/ardy/model/llm2vec/models/bidirectional_llama.py b/ardy/model/llm2vec/models/bidirectional_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..2255f4072457a16f1bc7769a7671106fcc696e9d --- /dev/null +++ b/ardy/model/llm2vec/models/bidirectional_llama.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from peft import PeftModel +from torch import nn +from transformers import LlamaConfig, LlamaForCausalLM, LlamaModel, LlamaPreTrainedModel +from transformers.cache_utils import Cache, StaticCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.models.llama.modeling_llama import ( + LlamaAttention, + LlamaDecoderLayer, + # LlamaFlashAttention2, + LlamaMLP, + LlamaRMSNorm, + LlamaRotaryEmbedding, + # LlamaSdpaAttention, +) +from transformers.utils import logging + +from .utils import is_transformers_attn_greater_or_equal_4_43_1 + +logger = logging.get_logger(__name__) + + +class ModifiedLlamaAttention(LlamaAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.is_causal = False + + +# class ModifiedLlamaFlashAttention2(LlamaFlashAttention2): +# def __init__(self, *args, **kwargs): +# super().__init__(*args, **kwargs) +# self.is_causal = False + + +# class ModifiedLlamaSdpaAttention(LlamaSdpaAttention): +# def __init__(self, *args, **kwargs): +# super().__init__(*args, **kwargs) +# self.is_causal = False + + +# LLAMA_ATTENTION_CLASSES = { +# "eager": ModifiedLlamaAttention, +# "flash_attention_2": ModifiedLlamaFlashAttention2, +# "sdpa": ModifiedLlamaSdpaAttention, +# } + + +class ModifiedLlamaDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: LlamaConfig, layer_idx: int): + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + + self.self_attn = ModifiedLlamaAttention(config=config, layer_idx=layer_idx) + # self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation]( + # config=config, layer_idx=layer_idx + # ) + + self.mlp = LlamaMLP(config) + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class LlamaBiModel(LlamaModel): + _no_split_modules = ["ModifiedLlamaDecoderLayer"] + + def __init__(self, config: LlamaConfig): + if not is_transformers_attn_greater_or_equal_4_43_1(): + raise ValueError( + "The current implementation of LlamaEncoderModel follows modeling_llama.py of transformers version >= 4.43.1" + ) + LlamaPreTrainedModel.__init__(self, config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [ModifiedLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LlamaRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def _update_causal_mask( + self, + attention_mask, + input_tensor, + cache_position, + past_key_values: Cache, + output_attentions: bool, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + # if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + # if AttentionMaskConverter._ignore_causal_mask_sdpa( + # attention_mask, + # inputs_embeds=input_tensor, + # past_key_values_length=past_seen_tokens, + # is_training=self.training, + # ): + # return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + causal_mask = torch.zeros( + (sequence_length, target_length), dtype=dtype, device=device + ) # in original implementation - torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) + # Commenting out next 2 lines to disable causal masking + # if sequence_length != 1: + # causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.dim() == 2: + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0) + causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype) + elif attention_mask.dim() == 4: + # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with + # cache. In that case, the 4D attention mask attends to the newest tokens only. + if attention_mask.shape[-2] < cache_position[0] + sequence_length: + offset = cache_position[0] + else: + offset = 0 + mask_shape = attention_mask.shape + mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype + causal_mask[ + : mask_shape[0], + : mask_shape[1], + offset : mask_shape[2] + offset, + : mask_shape[3], + ] = mask_slice + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + +class LlamaBiForMNTP(LlamaForCausalLM): + def __init__(self, config): + LlamaPreTrainedModel.__init__(self, config) + self.model = LlamaBiModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + # getter for PEFT model + def get_model_for_peft(self): + return self.model + + # setter for PEFT model + def set_model_for_peft(self, model: PeftModel): + self.model = model + + # save the PEFT model + def save_peft_model(self, path): + self.model.save_pretrained(path) diff --git a/ardy/model/llm2vec/models/utils.py b/ardy/model/llm2vec/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3f6fbe59bea002d7ecd29de59ac83e41b85260b9 --- /dev/null +++ b/ardy/model/llm2vec/models/utils.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.metadata + +from packaging import version +from transformers.utils.import_utils import _is_package_available + + +def is_transformers_attn_greater_or_equal_4_43_1(): + if not _is_package_available("transformers"): + return False + + return version.parse(importlib.metadata.version("transformers")) >= version.parse("4.43.1") diff --git a/ardy/model/load_model.py b/ardy/model/load_model.py new file mode 100644 index 0000000000000000000000000000000000000000..9c4c7aeea6a510d9df9862581f18931f9dcc4665 --- /dev/null +++ b/ardy/model/load_model.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Load Ardy diffusion models from local checkpoints or Hugging Face.""" + +from pathlib import Path +from typing import Optional + +import torch +from huggingface_hub import snapshot_download +from omegaconf import OmegaConf + +from .loading import ( + DEFAULT_MODEL, + DEFAULT_TEXT_ENCODER_URL, + get_env_var, + instantiate_from_dict, +) +from .registry import hf_repo_id, resolve_model_name + +DEFAULT_TEXT_ENCODER = "llm2vec" +TEXT_ENCODER_PRESETS = { + "llm2vec": { + "target": "ardy.model.LLM2VecEncoder", + "kwargs": { + # The original McGill adapter points at Meta's gated repository. + # Load the same Llama 3.0 architecture from an ungated mirror, then + # apply the original MNTP and supervised adapters explicitly. + "base_model_name_or_path": get_env_var( + "LLM2VEC_BASE_MODEL", + "unsloth/llama-3-8b-Instruct", + ), + "mntp_model_name_or_path": get_env_var( + "LLM2VEC_MNTP_MODEL", + "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp", + ), + "peft_model_name_or_path": get_env_var( + "LLM2VEC_SUPERVISED_MODEL", + "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised", + ), + "prompt_model_family": "llama3-instruct", + "dtype": "bfloat16", + "llm_dim": 4096, + "device": "auto", + }, + } +} + + +def _download_from_hf(full_name: str) -> Path: + """Download a released model from Hugging Face; returns the local snapshot dir. + + With LOCAL_CACHE=true, tries the local HF cache first and falls back online. + """ + repo_id = hf_repo_id(full_name) + local_cache = get_env_var("LOCAL_CACHE", "False").lower() == "true" + if local_cache: + try: + return Path(snapshot_download(repo_id=repo_id, local_files_only=True)) + except Exception: + pass # cache miss -> download online below + return Path(snapshot_download(repo_id=repo_id)) + + +def _build_api_text_encoder_conf(text_encoder_url: str) -> dict: + return { + "_target_": "ardy.model.text_encoder_api.TextEncoderAPI", + "url": text_encoder_url, + } + + +def _build_local_text_encoder_conf(text_encoder_fp32: bool = False) -> dict: + text_encoder_name = get_env_var("TEXT_ENCODER", DEFAULT_TEXT_ENCODER) + if text_encoder_name not in TEXT_ENCODER_PRESETS: + available = ", ".join(sorted(TEXT_ENCODER_PRESETS)) + raise ValueError(f"Unknown TEXT_ENCODER='{text_encoder_name}'. Available: {available}") + + preset = TEXT_ENCODER_PRESETS[text_encoder_name] + # Copy before overriding so the shared preset dict is never mutated. + kwargs = dict(preset["kwargs"]) + if text_encoder_fp32: + kwargs["dtype"] = "float32" + return { + "_target_": preset["target"], + **kwargs, + } + + +def _select_text_encoder_conf( + text_encoder_url: str, + text_encoder_fp32: bool = False, + mode: str = "auto", +) -> tuple: + """Return ``(conf, probe)``: the selected encoder conf plus the already-instantiated encoder + when auto-mode probing built one (reused by the caller so the API client is not constructed + twice). + + ``mode`` is resolved and validated by load_text_encoder: + - "api": force TextEncoderAPI + - "local": force local LLM2VecEncoder + - "auto": try API first, fallback to local if unreachable + """ + if mode == "local": + return _build_local_text_encoder_conf(text_encoder_fp32), None + if mode == "api": + return _build_api_text_encoder_conf(text_encoder_url), None + + api_conf = _build_api_text_encoder_conf(text_encoder_url) + try: + text_encoder = instantiate_from_dict(api_conf) + # Probe availability early so inference doesn't fail later. + text_encoder(["healthcheck"]) + return api_conf, text_encoder + except Exception as error: + print( + "Text encoder service is unreachable, falling back to local LLM2Vec " + f"encoder. ({type(error).__name__}: {error})" + ) + return _build_local_text_encoder_conf(text_encoder_fp32), None + + +def load_text_encoder( + mode: Optional[str] = None, + url: Optional[str] = None, + fp32: bool = False, + device: Optional[str] = None, +): + """Select and instantiate a text encoder, ready for inference. + + This is the single place that owns text-encoder selection + instantiation, + so it can be built once and reused across multiple models (e.g. core / g1 / + soma) by passing the result into ``load_model(..., text_encoder=...)``. + + Args: + mode: Backend selection ("auto"/"api"/"local"). When None, falls back + to the TEXT_ENCODER_MODE env var (default "auto"). + url: Remote service URL. When None, falls back to the TEXT_ENCODER_URL + env var. + fp32: Use float32 instead of the default bfloat16. + device: Target device. When None, uses cuda if available else cpu. + + Returns: + The instantiated text encoder placed on ``device``. + """ + if mode is None: + mode = get_env_var("TEXT_ENCODER_MODE", "auto") + mode = str(mode).lower() + if mode not in ("auto", "api", "local"): + raise ValueError( + f"Unknown text-encoder mode {mode!r}. Choose 'auto', 'api' or 'local'; " + "to load a model without a text encoder, pass text_encoder=False to " + "load_model()." + ) + + resolved_url = url or get_env_var("TEXT_ENCODER_URL", DEFAULT_TEXT_ENCODER_URL) + print( + f"Setting up text encoder (mode={mode}); first run may take a while...", + flush=True, + ) + conf, text_encoder = _select_text_encoder_conf(resolved_url, fp32, mode=mode) + if text_encoder is None: + # Placement is handled below via .to(); drop any device kwarg the preset + # may carry so it doesn't conflict (e.g. accelerate device_map="auto"). + conf = dict(conf) + conf.pop("device", None) + text_encoder = instantiate_from_dict(conf) + + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float32 if fp32 else torch.bfloat16 + return text_encoder.to(device=device, dtype=dtype) + + +def load_model( + modelname=None, + device=None, + eval_mode: bool = True, + default_family: Optional[str] = None, + text_encoder=None, + text_encoder_fp32: bool = False, + text_encoder_mode: Optional[str] = None, + text_encoder_url: Optional[str] = None, + return_config: bool = False, + checkpoints_dir: Optional[str] = None, +): + """Load a released Ardy model. + + ``modelname`` may be a short key ("core"/"g1"/"soma") or a full folder name + ("Ardy-Core-RP-20FPS-Horizon40"). If a local checkpoints dir is given (via + ``checkpoints_dir`` or the ``CHECKPOINTS_DIR`` env var) the model is loaded + from ``/``; otherwise it is downloaded from HF. + + Args: + modelname: Short key or full folder name; uses DEFAULT_MODEL if None. + device: Target device for the model (e.g. 'cuda', 'cpu'). + eval_mode: If True, set model to eval mode. + default_family: Ignored (kept for call-site compatibility). + text_encoder: Pre-built text encoder to reuse, or False to load the + model without any text encoder (model.text_encoder is left as + None). When None (the default), one is built via + ``load_text_encoder``. + text_encoder_fp32: If True, uses fp32 for the text encoder rather than default bfloat16. + text_encoder_mode: Backend selection ("auto"/"api"/"local"). When None, + falls back to the TEXT_ENCODER_MODE env var (default "auto"). + Ignored unless ``text_encoder`` is None. + text_encoder_url: URL of the remote text-encoder service. When None, + falls back to the TEXT_ENCODER_URL env var. + checkpoints_dir: Local dir holding released model folders. When None, + falls back to the CHECKPOINTS_DIR env var; if neither is set the + model is downloaded from Hugging Face. + + Returns: + Loaded model in eval mode, or (model, confg) if return_config is true + + Raises: + ValueError: If modelname cannot be resolved to a released model. + FileNotFoundError: If config.yaml is missing in the model folder. + """ + if modelname is None: + modelname = DEFAULT_MODEL + + # Local dir if CHECKPOINTS_DIR is set (arg or env), otherwise download from HF. + checkpoints_dir = checkpoints_dir or get_env_var("CHECKPOINTS_DIR") + # Resolve after checkpoints_dir so local-only folders (beyond the three + # released models) are accepted when loading from a local dir. + full_name = resolve_model_name(modelname, checkpoints_dir=checkpoints_dir) + if checkpoints_dir: + model_path = Path(checkpoints_dir) / full_name + if not model_path.exists(): + raise FileNotFoundError(f"Model {full_name!r} not found under CHECKPOINTS_DIR {checkpoints_dir!r}.") + else: + model_path = _download_from_hf(full_name) + + model_config_path = model_path / "config.yaml" + if not model_config_path.exists(): + raise FileNotFoundError(f"The model folder exists but config.yaml is missing: {model_config_path}") + + model_conf = OmegaConf.load(model_config_path) + + # Resolve the text encoder: False means load the model without one, a + # pre-built instance is reused as-is, and None (the default) builds one + # here via load_text_encoder (which resolves a None mode through the + # TEXT_ENCODER_MODE env var). Identity checks, not truthiness: False and + # None mean different things. + if text_encoder is False: + text_encoder = None + elif text_encoder is None: + text_encoder = load_text_encoder( + mode=text_encoder_mode, + url=text_encoder_url, + fp32=text_encoder_fp32, + device=device, + ) + + runtime_conf = OmegaConf.create({"checkpoint_dir": str(model_path)}) + model_cfg = OmegaConf.to_container(OmegaConf.merge(model_conf, runtime_conf), resolve=True) + model_cfg.pop("checkpoint_dir", None) + # The text encoder is attached after construction (or left as None for + # text_encoder=False), so prevent Hydra from instantiating one during + # construction. + model_cfg["text_encoder"] = None + + model = instantiate_from_dict(model_cfg, overrides={"device": device}) + if text_encoder is not None: + model.text_encoder = text_encoder + + if eval_mode: + model = model.eval() + if return_config: + return model, model_cfg + return model diff --git a/ardy/model/loading.py b/ardy/model/loading.py new file mode 100644 index 0000000000000000000000000000000000000000..0df962156bd151a3dcc93e323b73e0bc822ed3e6 --- /dev/null +++ b/ardy/model/loading.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model loading utilities: checkpoints, registry, env, and Hydra-based instantiation.""" + +import os +from pathlib import Path +from typing import Any, Dict, Optional, Union + +import torch +from hydra.utils import instantiate +from omegaconf import OmegaConf +from safetensors.torch import load_file as load_safetensors + +from .registry import ( + ARDY_MODELS, + AVAILABLE_MODELS, + DEFAULT_MODEL, + DEFAULT_TEXT_ENCODER_URL, + MODEL_NAMES, + TMR_MODELS, +) + + +def get_env_var(name: str, default: Optional[str] = None) -> Optional[str]: + """Return environment variable value, or default if unset/empty.""" + return os.environ.get(name) or default + + +def instantiate_from_dict( + cfg: Dict[str, Any], + overrides: Optional[Dict[str, Any]] = None, +): + """Instantiate an object from a config dict (e.g. from OmegaConf.to_container). + + The dict must contain _target_ with a fully qualified class path. Nested configs are + instantiated recursively. + """ + if overrides: + cfg = {**cfg, **overrides} + conf = OmegaConf.create(cfg) + return instantiate(conf) + + +def load_checkpoint_state_dict(ckpt_path: Union[str, Path]) -> dict: + """Load a state dict from a checkpoint file. + + If the checkpoint is a dict with a 'state_dict' key (e.g. PyTorch Lightning), + that is returned; otherwise the whole checkpoint is treated as the state dict. + + Args: + ckpt_path: Path to the checkpoint file. + + Returns: + state_dict suitable for model.load_state_dict(). + """ + ckpt_path = str(ckpt_path) + + if ckpt_path.endswith(".safetensors"): + state_dict = load_safetensors(ckpt_path) + else: + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) + if isinstance(checkpoint, dict) and "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + elif isinstance(checkpoint, dict): + state_dict = checkpoint + else: + raise ValueError(f"Unsupported checkpoint format: {ckpt_path}") + return {key: val.detach().cpu() for key, val in state_dict.items()} + + +__all__ = [ + "get_env_var", + "instantiate_from_dict", + "ARDY_MODELS", + "TMR_MODELS", + "AVAILABLE_MODELS", + "MODEL_NAMES", + "DEFAULT_MODEL", + "DEFAULT_TEXT_ENCODER_URL", + "load_checkpoint_state_dict", +] diff --git a/ardy/model/registry.py b/ardy/model/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..a532e0c468729e06efb7895b245356edf486c670 --- /dev/null +++ b/ardy/model/registry.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model registry: map nicknames to released model folders / Hugging Face repos. + +Released models are organized by skeleton and generation horizon (in frames). +``load_model`` accepts: + +- a skeleton nickname (``"core"``, ``"g1"``, ``"soma"``) — resolves to that + skeleton's default horizon (see ``DEFAULT_HORIZON``), +- a skeleton+horizon nickname (``"core8"``, ``"g152"``, ``"soma60"``), +- the full folder / repo name (``"ARDY-SOMA-RP-30FPS-Horizon60"``). +""" + +import os +import re + +# Hugging Face org that hosts the released models. +HF_ORG = "nvidia" + +# skeleton -> generation horizon (frames) -> released folder name +# (the folder name is also the HF repo name under HF_ORG). +MODELS_BY_SKELETON = { + "core": { + 40: "ARDY-Core-RP-20FPS-Horizon40", + 8: "ARDY-Core-RP-20FPS-Horizon8", + }, + "g1": { + 52: "ARDY-G1-RP-25FPS-Horizon52", + 8: "ARDY-G1-RP-25FPS-Horizon8", + }, +} + +# Horizon a bare skeleton nickname resolves to ("core" -> "core40"). +DEFAULT_HORIZON = {"core": 40, "g1": 52, "soma": 60} + +# nickname -> released folder name: "core8"/"core40"/... plus the bare +# skeleton names, which map to their DEFAULT_HORIZON variant. +MODELS = { + f"{skeleton}{horizon}": folder + for skeleton, by_horizon in MODELS_BY_SKELETON.items() + for horizon, folder in by_horizon.items() +} +MODELS.update({skeleton: MODELS_BY_SKELETON[skeleton][DEFAULT_HORIZON[skeleton]] for skeleton in MODELS_BY_SKELETON}) + +DEFAULT_MODEL = "core" +DEFAULT_TEXT_ENCODER_URL = "http://127.0.0.1:9550/" + +# --- Aliases kept for imports elsewhere (ardy.model.loading re-exports these) -- +# nickname -> HF repo id ("org/name") +MODEL_NAMES = {key: f"{HF_ORG}/{name}" for key, name in MODELS.items()} +# a modelname is valid if it is a nickname or a full folder name +AVAILABLE_MODELS = list(MODELS) + list(dict.fromkeys(MODELS.values())) +ARDY_MODELS = list(MODELS) +TMR_MODELS: list[str] = [] + +# Released-style folder name, e.g. "ARDY-Core-RP-20FPS-Horizon40". +_NAME_PATTERN = re.compile(r"ardy-(core|g1|soma)-.*horizon(\d+)$", re.IGNORECASE) + + +def parse_model_name(folder: str): + """``(skeleton, horizon)`` parsed from a released-style folder name. + + Returns e.g. ``("core", 40)`` for ``"ARDY-Core-RP-20FPS-Horizon40"`` (case-insensitive), or + ``None`` when the name does not follow the released naming scheme (e.g. a local training-run + folder). + """ + m = _NAME_PATTERN.match(folder) + if not m: + return None + return m.group(1).lower(), int(m.group(2)) + + +def resolve_model_name(name: str, default_family=None, checkpoints_dir=None) -> str: + """Return the released folder / repo name for a nickname or full name. + + Accepts a nickname (``"soma"``, ``"core8"``), the full folder name (``"ARDY-SOMA-RP-30FPS- + Horizon60"``, case-insensitive), or a full HF repo id (``"nvidia/ARDY-SOMA-RP-30FPS- + Horizon60"``). ``default_family`` is ignored (kept for call-site compatibility). + + When ``checkpoints_dir`` is given, the valid model set is whatever folders live there — not just + the released models — so a name matching a folder in it is accepted as-is (nicknames still + resolve via the registry). + """ + if name in MODELS: + return MODELS[name] + # Full folder name, optionally HF-org-prefixed; match case-insensitively + # and return the canonical casing (HF resolves either way, but local + # folder lookups are case-sensitive). + bare = name.split("/", 1)[1] if "/" in name else name + canonical = {folder.lower(): folder for folder in MODELS.values()} + if bare.lower() in canonical: + return canonical[bare.lower()] + if checkpoints_dir and os.path.isdir(os.path.join(checkpoints_dir, name)): + return name + raise ValueError( + f"Unknown model {name!r}. Choose a nickname {list(MODELS)} " + f"or a full name {list(dict.fromkeys(MODELS.values()))}." + ) + + +def hf_repo_id(full_name: str) -> str: + """Hugging Face repo id for a resolved full model name.""" + return f"{HF_ORG}/{full_name}" diff --git a/ardy/model/text_encoder_api.py b/ardy/model/text_encoder_api.py new file mode 100644 index 0000000000000000000000000000000000000000..e94dc27e713351b22cc5ab585f7395347edccd22 --- /dev/null +++ b/ardy/model/text_encoder_api.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Remote text encoder API client (Gradio) for motion generation.""" + +import logging + +import numpy as np +import torch +from gradio_client import Client + +# Suppress the [httpx] logs (GET requests) +logging.getLogger("httpx").setLevel(logging.WARNING) + +# Suppress internal gradio_client logs +logging.getLogger("gradio_client").setLevel(logging.WARNING) + + +class TextEncoderAPI: + """Text encoder API client for motion generation.""" + + def __init__(self, url: str): + self.client = Client(url, verbose=False) + self.device = "cpu" + self.dtype = torch.float + + def _create_np_random_name(self): + import uuid + + return str(uuid.uuid4()) + ".npy" + + def to(self, device=None, dtype=None): + if device is not None: + self.device = device + if dtype is not None: + self.dtype = dtype + return self + + def __call__(self, texts): + """Encode text prompts into tensors. + + Args: + texts (str | list[str]): text prompts to encode + + Returns: + tuple[torch.Tensor, list[int]]: encoded text tensors and their lengths + """ + if isinstance(texts, str): + texts = [texts] + if not texts: + raise ValueError("texts must contain at least one prompt") + + tensors = [] + lengths = [] + for text in texts: + filename = self._create_np_random_name() + + result = self.client.predict( + text=text, + filename=filename, + api_name="/DemoWrapper", + ) + path = result[0]["value"] + tensor = np.load(path) + length = tensor.shape[0] + + tensors.append(tensor) + lengths.append(length) + + padded_tensor = np.zeros((len(lengths), max(lengths), tensors[0].shape[-1]), dtype=tensors[0].dtype) + for idx, (tensor, length) in enumerate(zip(tensors, lengths)): + padded_tensor[idx, :length] = tensor + + padded_tensor = torch.from_numpy(padded_tensor) + padded_tensor = padded_tensor.to(device=self.device, dtype=self.dtype) + return padded_tensor, lengths diff --git a/ardy/model/trt.py b/ardy/model/trt.py new file mode 100644 index 0000000000000000000000000000000000000000..dd6f0afb1fccd9035eecadefddd5fa1ec53d502f --- /dev/null +++ b/ardy/model/trt.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TensorRT engine runtime wrappers with PyTorch-compatible interfaces. + +Provides drop-in replacements for the denoiser wrapper and autoencoder decoder that run inference +via pre-built TRT engines instead of PyTorch. +""" + +import logging +from typing import Dict, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +log = logging.getLogger(__name__) + +try: + import tensorrt as trt +except ImportError: + trt = None + + +class TRTEngine: + """Loads a serialized TensorRT engine and runs inference. + + Manages device memory for engine I/O bindings and provides a simple ``infer(output_shapes, + **inputs) -> dict[str, Tensor]`` interface. + """ + + def __init__(self, engine_path: str) -> None: + if trt is None: + raise ImportError("tensorrt is required for TRT inference. Install with: pip install tensorrt") + self.logger = trt.Logger(trt.Logger.WARNING) + engine_path = str(engine_path) + log.info("Loading TRT engine from %s", engine_path) + with open(engine_path, "rb") as f: + runtime = trt.Runtime(self.logger) + self.engine = runtime.deserialize_cuda_engine(f.read()) + if self.engine is None: + raise RuntimeError( + f"Failed to deserialize TRT engine: {engine_path}\n" + f"The engine was built with a different TensorRT version. " + f"Rebuild it with the currently installed version " + f"({trt.__version__}) by running scripts/export_onnx.py." + ) + self.context = self.engine.create_execution_context() + + # Catalogue bindings and their static shapes + self.input_names = [] + self.output_names = [] + self.binding_dtypes = {} + self.input_shapes = {} + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + mode = self.engine.get_tensor_mode(name) + dtype = trt.nptype(self.engine.get_tensor_dtype(name)) + shape = self.engine.get_tensor_shape(name) + self.binding_dtypes[name] = dtype + if mode == trt.TensorIOMode.INPUT: + self.input_names.append(name) + self.input_shapes[name] = tuple(shape) + else: + self.output_names.append(name) + log.info( + "TRT engine loaded: %d inputs (%s), %d outputs (%s)", + len(self.input_names), + ", ".join(self.input_names), + len(self.output_names), + ", ".join(self.output_names), + ) + + def infer( + self, + output_shapes: Optional[Dict[str, Tuple[int, ...]]] = None, + stream: Optional[torch.cuda.Stream] = None, + **inputs: torch.Tensor, + ) -> dict: + """Run inference with the given named inputs. + + Args: + output_shapes: Optional expected shape per output tensor. Any output + not listed here (or all of them, if None) has its shape resolved + from the execution context after the input shapes are set — which + works for graphs with multiple or auto-named outputs. + stream: CUDA stream to use. If None, uses the current stream. + **inputs: Named input tensors (must be contiguous, on CUDA). + + Returns: + Dictionary mapping output tensor names to GPU tensors. + """ + if stream is None: + stream = torch.cuda.current_stream() + output_shapes = output_shapes or {} + + # Set input shapes and addresses + for name in self.input_names: + tensor = inputs[name] + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + inputs[name] = tensor + self.context.set_input_shape(name, tuple(tensor.shape)) + self.context.set_tensor_address(name, tensor.data_ptr()) + + # Allocate outputs. Use the caller-provided shape when given, otherwise + # ask the context for the shape it resolved from the input shapes above. + outputs = {} + for name in self.output_names: + shape = output_shapes.get(name) + if shape is None: + shape = tuple(self.context.get_tensor_shape(name)) + np_dtype = self.binding_dtypes[name] + torch_dtype = torch.from_numpy(np.array([], dtype=np_dtype)).dtype + out = torch.empty(shape, dtype=torch_dtype, device="cuda") + outputs[name] = out + self.context.set_tensor_address(name, out.data_ptr()) + + # Execute + self.context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + return outputs + + +def _pad_to(tensor: torch.Tensor, target: int, dim: int) -> torch.Tensor: + """Zero-pad ``tensor`` along ``dim`` to ``target`` size.""" + cur = tensor.shape[dim] + if cur >= target: + return tensor + pad_sizes = [0] * (2 * tensor.ndim) + # F.pad uses reversed dim order: last dim first + pad_idx = 2 * (tensor.ndim - 1 - dim) + pad_sizes[pad_idx + 1] = target - cur + return F.pad(tensor, pad_sizes) + + +class TRTCFGDenoiser(nn.Module): + """Drop-in replacement for AutoLatentClassifierFreeGuidedModel using a TRT engine. + + The TRT engine includes both separated CFG batching and the denoiser. It takes B=1 inputs + + cfg_weights and returns the CFG-combined B=1 output. + + Because the ONNX graph has fixed attention reshapes, all inputs are padded to the trace-time + sizes (``num_tokens`` and ``num_text_tokens``) before calling the engine, and the output is + sliced back to the original size. + """ + + def __init__( + self, + engine_path: str, + denoiser: nn.Module, + num_frames_per_token: int, + num_tokens: int, + num_text_tokens: int, + ) -> None: + """ + Args: + engine_path: Path to the serialized CFG+denoiser TRT engine. + denoiser: The original denoiser (or infer wrapper) for attribute access. + num_tokens: The fixed num_tokens the ONNX was traced with. + num_text_tokens: The fixed num_text_tokens the ONNX was traced with. + """ + super().__init__() + self._engine = TRTEngine(engine_path) + self._denoiser = denoiser + self._num_tokens = num_tokens + self._num_text_tokens = num_text_tokens + self._nfpt = num_frames_per_token + self._num_frames = num_tokens * self._nfpt + self._dim_token = denoiser.nframe_root_dim + denoiser.latent_embedding_dim + self._motion_rep_dim = denoiser.motion_rep.motion_rep_dim + # Expose model attribute so denoising_step can access backbone dims + self.model = denoiser + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self._denoiser, name) + + def forward( + self, + cfg_weight_text: torch.Tensor, + cfg_weight_cstr: torch.Tensor, + token_seq_t: torch.Tensor, + history_len: torch.Tensor, + generation_len: torch.Tensor, + future_len: torch.Tensor, + history_mask: torch.Tensor, + generation_mask: torch.Tensor, + future_mask: torch.Tensor, + history_token_mask: torch.Tensor, + generation_token_mask: torch.Tensor, + future_token_mask: torch.Tensor, + text_feat: torch.Tensor, + text_feat_pad_mask: torch.Tensor, + timesteps: torch.Tensor, + first_heading_angle: Optional[torch.Tensor] = None, + motion_mask: Optional[torch.Tensor] = None, + observed_motion: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + actual_num_tokens = token_seq_t.shape[1] + if actual_num_tokens > self._num_tokens: + # TensorRT's Python API only logs oversize setInputShape errors and + # then runs with the engine's stale shapes, returning a truncated + # output that crashes far downstream — so refuse loudly here. + raise ValueError( + f"Generation window of {actual_num_tokens} tokens " + f"({actual_num_tokens * self._nfpt} frames) exceeds the " + f"{self._num_tokens}-token ({self._num_frames}-frame) capacity " + f"this TRT engine was built with. Reduce the window " + f"(history/future crop) or rebuild the engine with more tokens." + ) + T = self._num_tokens # padded token count + F_pad = self._num_frames # padded frame count + T_text = self._num_text_tokens + device = token_seq_t.device + mrep = self._motion_rep_dim + + if first_heading_angle is None: + first_heading_angle = torch.zeros(1, device=device, dtype=torch.float32) + if motion_mask is None: + motion_mask = torch.zeros(1, F_pad, mrep, device=device) + if observed_motion is None: + observed_motion = torch.zeros(1, F_pad, mrep, device=device) + + # Pad all inputs to the fixed trace-time sizes + x_pad = _pad_to(token_seq_t, T, dim=1) + hmask_pad = _pad_to(history_mask, F_pad, dim=1) + gmask_pad = _pad_to(generation_mask, F_pad, dim=1) + fmask_pad = _pad_to(future_mask, F_pad, dim=1) + htmask_pad = _pad_to(history_token_mask, T, dim=1) + gtmask_pad = _pad_to(generation_token_mask, T, dim=1) + ftmask_pad = _pad_to(future_token_mask, T, dim=1) + text_pad = _pad_to(text_feat, T_text, dim=1) + tpmask_pad = _pad_to(text_feat_pad_mask, T_text, dim=1) + mm_pad = _pad_to(motion_mask, F_pad, dim=1) + om_pad = _pad_to(observed_motion, F_pad, dim=1) + + output_shapes = {"output": (1, T, self._dim_token)} + results = self._engine.infer( + output_shapes, + cfg_weight_text=cfg_weight_text.float(), + cfg_weight_cstr=cfg_weight_cstr.float(), + x=x_pad.float(), + history_len=history_len.long(), + generation_len=generation_len.long(), + future_len=future_len.long(), + history_mask=hmask_pad.float(), + generation_mask=gmask_pad.float(), + future_mask=fmask_pad.float(), + history_token_mask=htmask_pad.float(), + generation_token_mask=gtmask_pad.float(), + future_token_mask=ftmask_pad.float(), + text_feat=text_pad.float(), + text_feat_pad_mask=tpmask_pad.float(), + timesteps=timesteps.long(), + first_heading_angle=first_heading_angle.float(), + motion_mask=mm_pad.float(), + observed_motion=om_pad.float(), + ) + # Slice back to actual size + return results["output"][:, :actual_num_tokens, :] + + +# FSQVAETransformer.detokenize splits the raw decoder output into these dict +# keys (in this order) depending on decoder_output_feature_mode. The decoder +# ONNX therefore has one graph output per key, in the same order, so the engine +# outputs line up with these positionally. Kept in sync with fsq.py. +_DECODER_OUTPUT_KEYS = { + "pose": ("root", "body"), + "body": ("body",), + "root": ("root",), +} + + +class TRTDecoder(nn.Module): + """Drop-in replacement for the autoencoder decoder using a TRT engine. + + Pads inputs to the trace-time sizes (num_tokens / num_frames) to match the fixed attention + reshapes in the ONNX graph, runs the engine, and slices each output back to the real frame + count. Returns one tensor per graph output, in graph order (``detokenize`` splits its result + into a dict, so the decoder graph has one output per dict entry — e.g. root and body). + """ + + def __init__( + self, + engine_path: str, + decoder_output_dim: Optional[int] = None, + num_tokens: int = 8, + num_frames_per_token: int = 4, + ) -> None: + super().__init__() + self._engine = TRTEngine(engine_path) + # decoder_output_dim is accepted for backward compatibility but no longer + # needed: output shapes are resolved from the engine at inference time. + self._num_tokens = num_tokens + self._num_frames = num_tokens * num_frames_per_token + + def forward( + self, + latent_tokens: torch.Tensor, + external_cond: torch.Tensor, + motion_pad_mask: torch.Tensor, + ) -> list: + actual_num_frames = motion_pad_mask.shape[1] + T = self._num_tokens + F_pad = self._num_frames + + # Pad to trace-time sizes + lt_pad = _pad_to(latent_tokens, T, dim=1) + ec_pad = _pad_to(external_cond, F_pad, dim=1) + mp_pad = _pad_to(motion_pad_mask, F_pad, dim=1) + + results = self._engine.infer( + latent_tokens=lt_pad.float(), + external_cond=ec_pad.float(), + motion_pad_mask=mp_pad.float(), + ) + # One tensor per graph output, in graph order, sliced back to actual size. + return [results[name][:, :actual_num_frames, :] for name in self._engine.output_names] + + +class TRTAutoencoder(nn.Module): + """Wraps the original autoencoder but replaces detokenize() with a TRT decoder.""" + + def __init__(self, trt_decoder: TRTDecoder, autoencoder: nn.Module) -> None: + super().__init__() + self._trt_decoder = trt_decoder + self._autoencoder = autoencoder + # The decoder graph emits one output per detokenize() dict key; recover + # those keys (in order) so detokenize() below returns the same dict the + # PyTorch autoencoder does. + mode = autoencoder.decoder_output_feature_mode + if mode not in _DECODER_OUTPUT_KEYS: + raise ValueError( + f"Unsupported decoder_output_feature_mode {mode!r}; expected one of {list(_DECODER_OUTPUT_KEYS)}." + ) + self._output_keys = _DECODER_OUTPUT_KEYS[mode] + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self._autoencoder, name) + + def detokenize( + self, + token_embeddings: torch.Tensor, + external_cond: torch.Tensor = None, + motion_pad_mask: torch.Tensor = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + """Route through the TRT decoder, returning the same dict as the PyTorch ``detokenize`` + (e.g. ``{"root": ..., "body": ...}``).""" + outputs = self._trt_decoder(token_embeddings, external_cond, motion_pad_mask) + if len(outputs) != len(self._output_keys): + raise RuntimeError( + f"Decoder engine produced {len(outputs)} outputs but " + f"decoder_output_feature_mode implies keys {self._output_keys}." + ) + return dict(zip(self._output_keys, outputs)) diff --git a/ardy/motion_rep/__init__.py b/ardy/motion_rep/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8570f8d73ee375f405f85696897be6697dc8a8b9 --- /dev/null +++ b/ardy/motion_rep/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Motion representation utilities.""" + +from .reps import ArdyMotionRep, MotionRepBase + +__all__ = [ + "ArdyMotionRep", + "MotionRepBase", +] diff --git a/ardy/motion_rep/conditioning.py b/ardy/motion_rep/conditioning.py new file mode 100644 index 0000000000000000000000000000000000000000..c7c64da24dd143b7684b16a96429ab41ed15db3f --- /dev/null +++ b/ardy/motion_rep/conditioning.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constraint conditioning: build index and data dicts from constraint sets for the denoiser.""" + +from collections import defaultdict + +import torch + + +def build_condition_dicts(constraints_lst: list): + index_dict = defaultdict(list) + data_dict = defaultdict(list) + for constraint in constraints_lst: + constraint.update_constraints(data_dict, index_dict) + return index_dict, data_dict + + +def get_unique_index_and_data(indices_lst, data): + # unique + sort them by t + indices_unique, inverse = torch.unique(indices_lst, dim=0, return_inverse=True) + # pick first value for each unique (t, j) + first_idx = torch.zeros(indices_unique.size(0), dtype=torch.long, device=inverse.device) + first_idx.scatter_(0, inverse, torch.arange(len(inverse), device=inverse.device)) + assert (indices_lst[first_idx] == indices_unique).all() + # get the data + indices_lst = indices_lst[first_idx] + data = data[first_idx] + return indices_lst, data diff --git a/ardy/motion_rep/feet.py b/ardy/motion_rep/feet.py new file mode 100644 index 0000000000000000000000000000000000000000..2765eb170c661de7293813ee26edfab4f238301b --- /dev/null +++ b/ardy/motion_rep/feet.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def foot_detect_from_pos_and_vel( + positions: torch.Tensor, + velocity: torch.Tensor, + skeleton, + vel_thres: float, + height_thresh: float, +) -> torch.Tensor: + """Compute foot contact labels using heuristics combining joint height and velocities. + + Args: + positions (torch.Tensor): [X, T, J, 3] global joint positions + velocity (torch.Tensor): [X, T, J, 3] velocities (already padded correctly), already multiplied by 1 / dt + skeleton: skeleton providing left/right foot joint indices + vel_thres (float): threshold for joint velocity + height_thresh (float): threshold for joint height + + Returns: + torch.Tensor: [X, T, 4] contact labels (left heel, left toe, right heel, + right toe), 1 for foot plant + """ + device = positions.device + fid_l = skeleton.left_foot_joint_idx + fid_r = skeleton.right_foot_joint_idx + + velfactor, heightfactor = ( + torch.tensor([vel_thres, vel_thres], device=device), + torch.tensor([height_thresh, height_thresh], device=device), + ) + + feet_l_v = torch.linalg.norm(velocity[:, :, fid_l], axis=-1) + feet_l_h = positions[:, :, fid_l, 1] + + feet_l = torch.logical_and( + feet_l_v < velfactor, + feet_l_h < heightfactor, + ).to(positions.dtype) + + feet_r_v = torch.linalg.norm(velocity[:, :, fid_r], axis=-1) + feet_r_h = positions[:, :, fid_r, 1] + + feet_r = torch.logical_and( + feet_r_v < velfactor, + feet_r_h < heightfactor, + ).to(positions.dtype) + + foot_contacts = torch.cat((feet_l, feet_r), axis=-1) + return foot_contacts diff --git a/ardy/motion_rep/reps/__init__.py b/ardy/motion_rep/reps/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3871970958b88c86301a19e3af325ff3e247328e --- /dev/null +++ b/ardy/motion_rep/reps/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .ardy_motionrep import ArdyMotionRep +from .base import MotionRepBase + +__all__ = ["ArdyMotionRep", "MotionRepBase"] diff --git a/ardy/motion_rep/reps/ardy_motionrep.py b/ardy/motion_rep/reps/ardy_motionrep.py new file mode 100644 index 0000000000000000000000000000000000000000..1d8112671764109f4aa097ea1f559282a6144dbb --- /dev/null +++ b/ardy/motion_rep/reps/ardy_motionrep.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ardy motion representation for autoregressive inference.""" + +from typing import Optional + +import einops +import torch +from torch import Tensor + +from ...geometry import cont6d_to_matrix, matrix_to_cont6d +from ...skeleton.kinematics import fk +from ...skeleton.transforms import global_rots_to_local_rots +from ...tools import ensure_batched, to_numpy +from ..conditioning import get_unique_index_and_data +from ..feet import foot_detect_from_pos_and_vel +from ..tools import ( + RotateFeatures, + compute_heading_angle, + compute_vel_xyz, +) +from .base import MotionRepBase + + +class ArdyMotionRep(MotionRepBase): + """Global root / global joint representation used by Ardy inference. + + Feature layout: + - ``root_pos``: root position ``[x, y, z]``. + - ``global_root_heading``: root heading as ``[cos(theta), sin(theta)]``. + - ``local_joints_positions``: non-root joints in root-local coordinates. + - ``global_rot_data``: global joint rotations in 6D representation. + - ``velocities``: global joint velocities. + - ``foot_contacts``: four foot contact channels. + """ + + def __init__( + self, + skeleton, + fps, + stats_path: Optional[str] = None, + stats=None, + name: Optional[str] = None, + **kwargs, + ): + # `stats`, `name`, and **kwargs let ArdyMotionRep be built straight from the + # training config via Hydra instantiate(), which passes a (core) skeleton, a Stats + # object (stats=...), a `name`, and possibly extra keys. See _ensure_ardy_skeleton. + skeleton = self._ensure_ardy_skeleton(skeleton) + assert skeleton.root_idx == 0, "ArdyMotionRep assumes the skeleton root index is 0." + self.name = name if name is not None else f"{skeleton.name}_dual_root_global_joints" + # Stats object (with a .folder) -> reuse ardy's single-folder sliced stats loader. + if stats_path is None and stats is not None: + stats_path = getattr(stats, "folder", None) + nbjoints = skeleton.nbjoints + + self.size_dict = { + "root_pos": torch.Size([3]), + "global_root_heading": torch.Size([2]), + "local_joints_positions": torch.Size([nbjoints - 1, 3]), # removed the pelvis joint + "global_rot_data": torch.Size([nbjoints, 6]), + "velocities": torch.Size([nbjoints, 3]), + "foot_contacts": torch.Size([4]), + } + self.last_root_feature = "global_root_heading" + self.local_root_size_dict = { + "local_root_rot_vel": torch.Size([1]), + "local_root_vel": torch.Size([2]), + "global_root_y": torch.Size([1]), + } + super().__init__(skeleton, fps, stats_path) + + @staticmethod + def _ensure_ardy_skeleton(skeleton): + """Return an ardy.motion_rep skeleton. + + ArdyMotionRep relies on ardy.motion_rep FK/geometry, which require an ardy.motion_rep + skeleton. When built from the (core) training config the loader passes a core skeleton, so + rebuild the matching ardy skeleton from the same folder. + """ + from ardy.skeleton import SkeletonBase as ArdySkeletonBase + + if isinstance(skeleton, ArdySkeletonBase): + return skeleton + + from ardy.skeleton import ( + CoreSkeleton27, + G1Skeleton34, + SOMASkeleton30, + SOMASkeleton77, + ) + + skel_by_njoints = { + 27: CoreSkeleton27, + 34: G1Skeleton34, + 30: SOMASkeleton30, + 77: SOMASkeleton77, + } + nbjoints = skeleton.nbjoints + if nbjoints not in skel_by_njoints: + raise ValueError(f"No ardy.motion_rep skeleton for nbjoints={nbjoints} (known: {sorted(skel_by_njoints)}).") + try: + device = next(skeleton.buffers()).device + except (StopIteration, AttributeError): + device = "cpu" + return skel_by_njoints[nbjoints]( + folder=skeleton.folder, + load=True, + t_pose=getattr(skeleton, "t_pose", None), + ).to(device) + + def recenter_root_motion( + self, + root_motion: torch.Tensor, + center_frame_index: torch.Tensor, + is_normalized: bool, + to_normalize: bool, + return_center_pos: bool = False, + ): + """Translate root x/z so a selected frame becomes the local origin.""" + if is_normalized: + root_motion = self.global_root_stats.unnormalize(root_motion) + + batch_idx = torch.arange(root_motion.shape[0], device=root_motion.device) + center_pos = root_motion[batch_idx, center_frame_index.long(), :3].clone() + center_pos[:, 1] = 0 + root_motion = root_motion.clone() + root_motion[:, :, [0, 2]] -= center_pos[:, None, [0, 2]] + + if to_normalize: + root_motion = self.global_root_stats.normalize(root_motion) + if return_center_pos: + return root_motion, center_pos + return root_motion + + @ensure_batched(local_joint_rots=5, root_positions=3, lengths=1) + def __call__( + self, + local_joint_rots: torch.Tensor, + root_positions: torch.Tensor, + to_normalize: bool, + to_canonicalize: bool = False, + lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Convert local rotations and root positions into Ardy features.""" + device = local_joint_rots.device + if lengths is None: + assert local_joint_rots.shape[0] == 1, "If lengths is not provided, the input should not be batched." + lengths = torch.tensor([local_joint_rots.shape[1]], device=device) + + global_joint_rots, global_positions, local_joints_positions_origin_is_pelvis = fk( + local_joint_rots, + root_positions, + self.skeleton, + ) + + root_heading_angle = compute_heading_angle(global_positions, self.skeleton) + global_root_heading = torch.stack([torch.cos(root_heading_angle), torch.sin(root_heading_angle)], dim=-1) + + ground_offset = torch.zeros_like(root_positions) + ground_offset[..., 1] = root_positions[..., 1] + # remove the pelvis joint (root_idx == 0) and align onto the ground + local_joints_positions = local_joints_positions_origin_is_pelvis[:, :, 1:] + ground_offset[:, :, None] + + velocities = compute_vel_xyz(global_positions, self.fps, lengths=lengths) + foot_contacts = foot_detect_from_pos_and_vel(global_positions, velocities, self.skeleton, 0.15, 0.10) + global_rot_data = matrix_to_cont6d(global_joint_rots) + + features, _ = einops.pack( + [ + root_positions, + global_root_heading, + local_joints_positions, + global_rot_data, + velocities, + foot_contacts, + ], + "batch time *", + ) + + assert features.shape[-1] == self.motion_rep_dim + + if to_canonicalize: + features = self.canonicalize(features, normalized=False) + + if to_normalize: + features = self.normalize(features) + return features + + @ensure_batched(features=3, angle=1) + def rotate(self, features: torch.Tensor, angle: torch.Tensor): + """Rotate root/joint positional and rotational features by heading.""" + # assume it is not normalized + bs = features.shape[0] + device = features.device + [ + root_pos, + global_root_heading, + local_joints_positions, + global_rot_data, + velocities, + foot_contacts, + ] = einops.unpack(features, self.ps, "batch time *") + + if not isinstance(angle, torch.Tensor): + angle = torch.tensor(angle, device=device) + if len(angle.shape) == 0: + angle = angle.repeat(bs) + + RF = RotateFeatures(angle) + new_features, _ = einops.pack( + [ + RF.rotate_positions(root_pos), + RF.rotate_2d_positions(global_root_heading), + RF.rotate_positions(local_joints_positions), + RF.rotate_6d_rotations(global_rot_data), + RF.rotate_positions(velocities), + foot_contacts, + ], + "batch time *", + ) + return new_features + + @ensure_batched(features=3, translation_2d=2) + def translate_2d(self, features: torch.Tensor, translation_2d: torch.Tensor) -> torch.Tensor: + """Translate root planar position by ``(dx, dz)``.""" + bs = features.shape[0] + if len(translation_2d.shape) == 1: + translation_2d = translation_2d.repeat(bs, 1) + + new_features = features.clone() + new_root_pos = new_features[:, :, self.slice_dict["root_pos"]] + new_root_pos[:, :, 0] += translation_2d[:, [0]] + new_root_pos[:, :, 2] += translation_2d[:, [1]] + return new_features + + @ensure_batched(features=3) + def inverse( + self, + features: torch.Tensor, + is_normalized: bool, + posed_joints_from="rotations", + return_numpy: bool = False, + ) -> dict: + """Decode Ardy features into motion tensors.""" + assert posed_joints_from in ["rotations", "positions"], "posed_joints_from should be rotations or positions" + + if is_normalized: + features = self.unnormalize(features) + + [ + root_positions, + global_root_heading, + local_joints_positions, + global_rot_data, + velocities, + foot_contacts, + ] = einops.unpack(features, self.ps, "batch time *") + + global_rot_mats = cont6d_to_matrix(global_rot_data) + local_rot_mats = global_rots_to_local_rots(global_rot_mats, self.skeleton) + + if posed_joints_from == "rotations": + _, posed_joints, _ = fk(local_rot_mats, root_positions, self.skeleton) + else: + dummy_root = torch.zeros_like(local_joints_positions[:, :, [0]]) + posed_joints = torch.cat([dummy_root, local_joints_positions], dim=2) + posed_joints[..., 0] += root_positions[..., None, 0] + posed_joints[..., 2] += root_positions[..., None, 2] + + output_tensor_dict = { + "local_rot_mats": local_rot_mats, + "global_rot_mats": global_rot_mats, + "posed_joints": posed_joints, + "root_positions": root_positions, + "smooth_root_pos": root_positions, + "foot_contacts": foot_contacts > 0.5, + "global_root_heading": global_root_heading, + } + if return_numpy: + return to_numpy(output_tensor_dict) + return output_tensor_dict + + def create_conditions( + self, + index_dict: dict, + data_dict: dict, + length: int, + to_normalize: bool, + device: str, + ): + observed_motion = torch.zeros(length, self.motion_rep_dim, device=device) + motion_mask = torch.zeros(length, self.motion_rep_dim, dtype=bool, device=device) + + self._fill_root_2d_constraints(observed_motion, motion_mask, index_dict, data_dict, device) + self._fill_global_heading_constraints(observed_motion, motion_mask, index_dict, data_dict, device) + self._fill_root_y_constraints(observed_motion, motion_mask, index_dict, data_dict, device) + self._fill_global_rotation_constraints(observed_motion, motion_mask, index_dict, data_dict, device) + self._fill_global_position_constraints(observed_motion, motion_mask, index_dict, data_dict, device) + + motion_mask = motion_mask.float() + if to_normalize: + observed_motion = self.normalize(observed_motion) * motion_mask + return observed_motion, motion_mask + + def _cat_indices(self, values, device): + indices = torch.cat([torch.tensor(x) if not isinstance(x, Tensor) else x for x in values]) + return indices.to(device=device, dtype=torch.long) + + def _fill_root_2d_constraints(self, observed_motion, motion_mask, index_dict, data_dict, device): + fname = "root_2d" if index_dict.get("root_2d") else "smooth_root_2d" + if fname not in index_dict or not index_dict[fname]: + return + + indices = self._cat_indices(index_dict[fname], device) + indices, root_pos_2d = get_unique_index_and_data(indices, torch.cat(data_dict[fname]).to(device)) + f_sliced = observed_motion[:, self.slice_dict["root_pos"]] + f_sliced[indices, 0] = root_pos_2d[:, 0] + f_sliced[indices, 2] = root_pos_2d[:, 1] + m_sliced = motion_mask[:, self.slice_dict["root_pos"]] + m_sliced[indices, 0] = True + m_sliced[indices, 2] = True + + def _fill_global_heading_constraints(self, observed_motion, motion_mask, index_dict, data_dict, device): + fname = "global_root_heading" + if fname not in index_dict or not index_dict[fname]: + return + + indices = self._cat_indices(index_dict[fname], device) + indices, global_root_heading = get_unique_index_and_data(indices, torch.cat(data_dict[fname]).to(device)) + f_sliced = observed_motion[:, self.slice_dict[fname]] + f_sliced[indices] = global_root_heading + m_sliced = motion_mask[:, self.slice_dict[fname]] + m_sliced[indices] = True + + def _fill_root_y_constraints(self, observed_motion, motion_mask, index_dict, data_dict, device): + fname = "root_y_pos" + if fname not in index_dict or not index_dict[fname]: + return + + indices = self._cat_indices(index_dict[fname], device) + indices, root_y_pos = get_unique_index_and_data(indices, torch.cat(data_dict[fname]).to(device)) + root_y_pos = root_y_pos.reshape(-1) + f_sliced = observed_motion[:, self.slice_dict["root_pos"]] + f_sliced[indices, 1] = root_y_pos + m_sliced = motion_mask[:, self.slice_dict["root_pos"]] + m_sliced[indices, 1] = True + + def _fill_global_rotation_constraints(self, observed_motion, motion_mask, index_dict, data_dict, device): + fname = "global_joints_rots" + if fname not in index_dict or not index_dict[fname]: + return + + indices_lst = self._cat_indices(index_dict[fname], device) + indices_lst, global_joints_rots = get_unique_index_and_data(indices_lst, torch.cat(data_dict[fname]).to(device)) + global_rot_data = matrix_to_cont6d(global_joints_rots) + + f_sliced = observed_motion[:, self.slice_dict["global_rot_data"]] + masking = torch.zeros(len(f_sliced) * self.nbjoints, 6, device=device, dtype=bool) + masking[indices_lst.T[0] * self.nbjoints + indices_lst.T[1]] = True + masking = masking.reshape(len(f_sliced), self.nbjoints * 6) + f_sliced[masking] = global_rot_data.flatten() + m_sliced = motion_mask[:, self.slice_dict["global_rot_data"]] + m_sliced[masking] = True + + def _fill_global_position_constraints(self, observed_motion, motion_mask, index_dict, data_dict, device): + fname = "global_joints_positions" + if fname not in index_dict or not index_dict[fname]: + return + + indices_lst = self._cat_indices(index_dict[fname], device) + indices_lst, global_joints_positions = get_unique_index_and_data( + indices_lst, + torch.cat(data_dict[fname]).to(device), + ) + + time_indices = indices_lst[:, 0].contiguous() + unique_times = time_indices.unique().contiguous() + value_indices = torch.searchsorted(unique_times, time_indices) + hips_mask = indices_lst[:, 1] == self.skeleton.root_idx + assert hips_mask.sum() == len(unique_times) + assert (indices_lst[hips_mask, 0] == unique_times).all() + + root_positions = global_joints_positions[hips_mask][value_indices].clone() + root_positions_y = root_positions[:, 1].clone() + + root_test = motion_mask[time_indices, self.slice_dict["root_pos"]] + if not root_test[:, [0, 2]].all(): + raise ValueError("For constraining global positions, root 2D should also be constrained.") + + ground_offset = torch.zeros_like(root_positions) + ground_offset[:, 1] = root_positions_y + local_joints_positions = global_joints_positions - root_positions + ground_offset + + f_sliced = observed_motion[:, self.slice_dict["local_joints_positions"]] + masking = torch.zeros(len(f_sliced) * (self.nbjoints - 1), 3, device=device, dtype=bool) + non_root_mask = ~hips_mask + indices_lst_no_root = indices_lst[non_root_mask] + local_joints_positions_no_root = local_joints_positions[non_root_mask] + masking[indices_lst_no_root[:, 0] * (self.nbjoints - 1) + (indices_lst_no_root[:, 1] - 1)] = True + masking = masking.reshape(len(f_sliced), (self.nbjoints - 1) * 3) + f_sliced[masking] = local_joints_positions_no_root.flatten() + m_sliced = motion_mask[:, self.slice_dict["local_joints_positions"]] + m_sliced[masking] = True diff --git a/ardy/motion_rep/reps/base.py b/ardy/motion_rep/reps/base.py new file mode 100644 index 0000000000000000000000000000000000000000..fff37da4adf011ac5034e7b0e4c27aa6347e72b8 --- /dev/null +++ b/ardy/motion_rep/reps/base.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Base motion representation: feature layout, normalization, and conditioning helpers.""" + +from typing import Optional + +import einops +import numpy as np +import torch +from einops import repeat + +from ...tools import ensure_batched +from ..conditioning import build_condition_dicts +from ..stats import Stats +from ..tools import compute_vel_angle, compute_vel_xyz + + +class MotionRepBase: + """Base class for motion representations used in generation and conditioning. + + Subclasses define: + - ``size_dict``: feature blocks and their shapes, + - ``last_root_feature``: last entry of the root block, + - ``local_root_size_dict``: local-root feature layout, + and implement transform-specific methods such as ``__call__``, ``inverse``, + ``rotate``, ``translate_2d`` and ``create_conditions``. + """ + + def __init__( + self, + skeleton, + fps, + stats_path: Optional[str] = None, + ): + """Initialize feature slicing metadata and optional normalization stats.""" + self.skeleton = skeleton + self.fps = fps + self.nbjoints = skeleton.nbjoints + + self.feature_names = list(self.size_dict.keys()) + self.ps = list(self.size_dict.values()) + self.nfeats_dict = {key: val.numel() for key, val in self.size_dict.items()} + feats_cumsum = np.cumsum([0] + list(self.nfeats_dict.values())).tolist() + self.slice_dict = {key: slice(feats_cumsum[i], feats_cumsum[i + 1]) for i, key in enumerate(self.feature_names)} + + self.motion_rep_dim = sum(self.nfeats_dict.values()) + self.root_slice = slice(0, self.slice_dict[self.last_root_feature].stop) + self.body_slice = slice(self.root_slice.stop, self.motion_rep_dim) + self.body_dim = self.body_slice.stop - self.body_slice.start + self.global_root_dim = self.motion_root_dim = self.root_slice.stop + self.local_root_dim = sum(val.numel() for val in self.local_root_size_dict.values()) + + self.local_motion_rep_dim = self.local_root_dim + self.body_dim + self.global_motion_rep_dim = self.motion_rep_dim + + # Stats ship as a single mean.npy / std.npy pair covering the full + # feature vector; the per-part stats below are slices of it. + if stats_path: + all_stats = Stats(stats_path) + self.local_root_stats = all_stats.sliced( + indices=(self.global_root_dim + np.arange(self.local_root_dim)).tolist(), + ) + self.body_stats = all_stats.sliced( + indices=(self.global_root_dim + self.local_root_dim + np.arange(self.body_dim)).tolist(), + ) + self.global_root_stats = all_stats.sliced( + indices=(0 + np.arange(self.global_root_dim)).tolist(), + ) + self.stats = all_stats.sliced( + indices=(0 + np.arange(self.global_root_dim)).tolist() + + (self.global_root_dim + self.local_root_dim + np.arange(self.body_dim)).tolist(), + ) + + def get_root_pos(self, features: torch.Tensor, fallback_to_smooth: bool = True): + """Extract root positions from a feature tensor. + + Supports both ``root_pos`` and ``smooth_root_pos`` representations. + """ + if "root_pos" in self.slice_dict: + return features[..., self.slice_dict["root_pos"]] + + if "smooth_root_pos" not in self.slice_dict: + raise TypeError("This motion rep should have either a root_pos or smooth_root_pos field") + + if fallback_to_smooth: + return features[:, :, self.slice_dict["smooth_root_pos"]] + + # else compute the root pos from the smooth root and local joints offset + smooth_root_pos = features[:, :, self.slice_dict["smooth_root_pos"]].clone() + local_joints_positions_flatten = features[..., self.slice_dict["local_joints_positions"]] + hips_offset = local_joints_positions_flatten[..., self.skeleton.root_idx : self.skeleton.root_idx + 3] + root_pos = torch.stack( + [ + smooth_root_pos[..., 0] + hips_offset[..., 0], + smooth_root_pos[..., 1], + smooth_root_pos[..., 2] + hips_offset[..., 2], + ], + axis=-1, + ) + return root_pos + + @ensure_batched(root_features=3, lengths=1) + def global_root_to_local_root( + self, + root_features: torch.Tensor, + normalized: bool, + lengths: Optional[torch.Tensor], + ): + """Convert global root features to local-root motion features. + + Args: + root_features: Root feature tensor containing root position and + global heading, shaped ``[B, T, D_root]``. + normalized: Whether ``root_features`` are normalized. + lengths: Optional valid lengths per sequence. + + Returns: + Tensor ``[B, T, 4]`` with local root rotational velocity, planar + velocity, and global root height. + """ + if normalized: + root_features = self.global_root_stats.unnormalize(root_features) + + [root_pos, global_root_heading] = einops.unpack(root_features, self.ps[:2], "batch time *") + cos, sin = global_root_heading.unbind(-1) + heading_angle = torch.arctan2(sin, cos) + + local_root_rot_vel = compute_vel_angle(heading_angle, self.fps, lengths=lengths) + local_root_vel = compute_vel_xyz( + root_pos[..., None, :], + self.fps, + lengths=lengths, + )[..., 0, [0, 2]] + global_root_y = root_pos[..., 1] + local_root_motion = torch.cat( + [ + local_root_rot_vel[..., None], + local_root_vel, + global_root_y[..., None], + ], + axis=-1, + ) + + if normalized: + local_root_motion = self.local_root_stats.normalize(local_root_motion) + return local_root_motion + + def get_root_heading_angle(self, features: torch.Tensor) -> torch.Tensor: + """Compute root heading angle from cosine/sine heading features.""" + global_root_heading = features[:, :, self.slice_dict["global_root_heading"]] + cos, sin = global_root_heading.unbind(-1) + return torch.arctan2(sin, cos) + + @ensure_batched(features=3) + def rotate_to( + self, + features: torch.Tensor, + target_angle: torch.Tensor, + return_delta_angle=False, + ): + """Rotate each sequence so frame-0 heading matches ``target_angle``.""" + # rotate so that the first frame angle is the target + # it put the motion_rep to the angle + current_first_angle = self.get_root_heading_angle(features)[:, 0] + delta_angle = target_angle - current_first_angle + rotated_features = self.rotate(features, delta_angle) + if return_delta_angle: + return rotated_features, delta_angle + return rotated_features + + @ensure_batched(features=3) + def rotate_to_zero( + self, + features: torch.Tensor, + return_delta_angle=False, + ): + """Rotate each sequence so frame-0 heading becomes zero.""" + target_angle = torch.zeros(len(features), device=features.device) + return self.rotate_to(features, target_angle, return_delta_angle=return_delta_angle) + + @ensure_batched(features=3) + def randomize_first_heading( + self, + features: torch.Tensor, + return_delta_angle=False, + ) -> torch.Tensor: + """Rotate each sequence to a random frame-0 heading.""" + target_heading_angle = torch.rand(features.shape[0], device=features.device) * 2 * np.pi + return self.rotate_to( + features, + target_heading_angle, + return_delta_angle=return_delta_angle, + ) + + @ensure_batched(features=3, target_2d_pos=2) + def translate_2d_to( + self, + features: torch.Tensor, + target_2d_pos: torch.Tensor, + return_delta_pos: bool = False, + ) -> torch.Tensor: + """Translate each sequence so frame-0 root ``(x, z)`` matches a target.""" + root_pos = self.get_root_pos(features) + current_first_2d_pos = root_pos[:, 0, [0, 2]].clone() + delta_2d_pos = target_2d_pos - current_first_2d_pos + translated_features = self.translate_2d(features, delta_2d_pos) + if return_delta_pos: + return translated_features, delta_2d_pos + return translated_features + + @ensure_batched(features=3) + def translate_2d_to_zero( + self, + features: torch.Tensor, + return_delta_pos: bool = False, + ) -> torch.Tensor: + """Translate each sequence so frame-0 root ``(x, z)`` is at the origin.""" + target_2d_pos = torch.zeros(len(features), 2, device=features.device) + return self.translate_2d_to(features, target_2d_pos, return_delta_pos=return_delta_pos) + + @ensure_batched(features=3) + def canonicalize(self, features: torch.Tensor, normalized: bool = False): + """Canonicalize heading and planar position at frame 0.""" + if normalized: + features = self.unnormalize(features) + rotated_features = self.rotate_to_zero(features) + canonicalized_features = self.translate_2d_to_zero(rotated_features) + if normalized: + canonicalized_features = self.normalize(canonicalized_features) + return canonicalized_features + + def normalize(self, features): + """Normalize features.""" + return self.stats.normalize(features) + + def unnormalize(self, features): + """Undo feature normalization.""" + return self.stats.unnormalize(features) + + def extract_root(self, features: torch.Tensor) -> torch.Tensor: + """Extract normalized/global root features from a full motion tensor.""" + return features[..., self.root_slice] + + def extract_body(self, features: torch.Tensor) -> torch.Tensor: + """Extract body features from a full motion tensor.""" + return features[..., self.body_slice] + + def concat_root_body(self, root: torch.Tensor, body: torch.Tensor) -> torch.Tensor: + """Concatenate root and body features into a full Ardy feature tensor.""" + return torch.cat([root, body], dim=-1) + + def create_conditions_from_constraints( + self, + constraints_lst: list, + length: int, + to_normalize: bool, + device: str, + ): + """Create a conditioning tensor and mask from constraint objects.""" + index_dict, data_dict = build_condition_dicts(constraints_lst) + return self.create_conditions(index_dict, data_dict, length, to_normalize, device) + + def create_conditions_from_constraints_batched( + self, + constraints_lst: list | list[list], + lengths: torch.Tensor, + to_normalize: bool, + device: str, + ): + """Batched version of ``create_conditions_from_constraints``. + + Supports either one shared constraint list for all batch elements, or a per-sample list of + constraint lists. + """ + num_samples = len(lengths) + if not constraints_lst or not isinstance(constraints_lst[0], list): + # If no constraints, or constraints are shared across the batch, + # build once and repeat. + observed_motion, motion_mask = self.create_conditions_from_constraints( + constraints_lst, int(lengths.max()), to_normalize, device + ) + observed_motion = repeat(observed_motion, "t d -> b t d", b=num_samples) + motion_mask = repeat(motion_mask, "t d -> b t d", b=num_samples) + return observed_motion, motion_mask + + length = int(lengths.max()) + observed_motion_lst = [] + motion_mask_lst = [] + for constraints_lst_el in constraints_lst: + observed_motion, motion_mask = self.create_conditions_from_constraints( + constraints_lst_el, + length, + to_normalize, + device, + ) + observed_motion_lst.append(observed_motion) + motion_mask_lst.append(motion_mask) + observed_motion = torch.stack(observed_motion_lst, axis=0) + motion_mask = torch.stack(motion_mask_lst, axis=0) + return observed_motion, motion_mask diff --git a/ardy/motion_rep/smooth_root.py b/ardy/motion_rep/smooth_root.py new file mode 100644 index 0000000000000000000000000000000000000000..2049cc0c9079c0ec4ec6ccf4b0c14be1bf8fb183 --- /dev/null +++ b/ardy/motion_rep/smooth_root.py @@ -0,0 +1,463 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import numpy as np +import torch +from scipy import sparse +from scipy.sparse.linalg import splu + + +class TrajectorySmoother: + """A class for modifying trajectories to hit specific values at specific frames while respecting + soft constraints. + + This class modifies a trajectory to hit specific values at specific + frames, while respecting the following soft constraints: + * Preserve the original positions + * Bring the accelerations as close to zero as possible + + The weights of the position soft constraints are specified in pos_weight. + + This is posed as a minimization problem: + E(x) = pos_weight * |x - x_orig|^2 + + |A x|^2 + + where you minimize E(x) subject to specified values at indices where + "mask" is equal to 1. A is a matrix that computes the N-2 accelerations + associated with frames n-1, n and n+1. + """ + + """ + min f(x) + g(z) + s.t. I x - z = 0 + + x --> argmin_x (f(x) + p/2 ||I x - z + u||^2) + z --> argmin_z (g(z) + p/2 ||I x - z + u||^2) + u --> u + I x - z + + f(x) = pos_weight * |x - x_orig|^2 + |A x|^2 + g(z) = inf if any(|z-t| > margin) else 0 + + x minimization: + + E(x) = wp/2 * |x - x_orig|^2 + 1/2 |A x|^2 + p/2 |I x - z + u|^2 + E(x) = wp/2 * (x - x_orig)^T (x - x_orig) + x^T A^T A x + p/2 (I x - z + u)^T (I x - z + u) + E(x) = wp/2 * (x^T x - x_orig^T x - x^T x_orig + x_orig^T x_orig) + + 1/2 x^T A^T A x + + p/2 (x^T I^T I x - x^T I^T z + x^T I^T u - z^T I x + z^T z - z^T u + u^T I x - u^T z + u^T u) + + argmin E(x) = argmin [ + wp/2 * (x^T x - 2 x^T x_orig) + + 1/2 x^T A^T A x + + p/2 (x^T I^T I x + 2 x^T I^T (u - z) ) + ] + = argmin [ + x^T wp/2 * I * x - wp * x^T x_orig + + 1/2 x^T A^T A x + + x^T p/2 I^T I x + p x^T I^T (u - z) + ] + = argmin [ + 1/2 x^T (wp * I + A^T A + p I^T I) x - x^T (wp * x_orig + p I^T (z - u)) + ] + + x = (wp * I + A^T A + p I)^-1 (wp * x_orig + p (z - u)) + + + """ + + def __init__( + self, + margins, + pos_weight=0.0, + loop=False, + admm_iters=100, + alpha_overrelax=1.0, + circle_project=False, + ): + """Initialize the TrajectorySmoother. + + Args: + margins: Array of margin values for each frame. + margins[i] < 0: unconstrained + margins[i] == 0: pinned on this frame + margins[i] > 0: can deviate within the margin + pos_weight: Weight for position preservation + loop: Whether the trajectory should loop + admm_iters: Number of ADMM iterations + """ + self.pos_weight = pos_weight + self.admm_iters = admm_iters + self.alpha_overrelax = alpha_overrelax + self.circle_project = circle_project + N = len(margins) + + # Store margin information as numpy arrays + self.margin_vals = margins + + # Build acceleration matrix A + a_data = [] + a_rows = [] + a_cols = [] + + for i in range(1, N - 1): + scale = 1.0 + a_data.extend([-scale, 2.0 * scale, -scale]) + a_rows.extend([i, i, i]) + a_cols.extend([i - 1, i, i + 1]) + + if loop: + # Add periodic accelerations + scale = 1.0 + a_data.extend([-scale, 2.0 * scale, -scale]) + a_rows.extend([0, 0, 0]) + a_cols.extend([N - 1, 0, 1]) + + scale = 1.0 + a_data.extend([-scale, 2.0 * scale, -scale]) + a_rows.extend([N - 1, N - 1, N - 1]) + a_cols.extend([N - 2, N - 1, 0]) + + A = sparse.csr_matrix((a_data, (a_rows, a_cols)), shape=(N, N)) + + # Build identity matrix + identity_matrix = sparse.eye(N) + + # Build system matrix M + M = pos_weight * identity_matrix + A.T @ A + + # Calculate ADMM step size + diag_max = max(abs(M.diagonal())) + self.admm_stepsize = 0.25 * np.sqrt(diag_max) + + M = M + self.admm_stepsize * identity_matrix + self.system_lu = splu(M.tocsc()) + + def smooth(self, targets, x0): + """Interpolate between reference positions while satisfying constraints. + + Args: + observations: Target positions for constrained frames (numpy array) + ref_positions: Reference positions defining original shape + (numpy array) + + Returns: + Interpolated positions (numpy array) + """ + x_target = targets.copy() + x = x0.copy() + z = np.zeros_like(x) + u = np.zeros_like(x) + + for _ in range(self.admm_iters): + self.z_update(z, x, x_target, u) + self.u_update(u, x, z) + self.x_update(x, z, u, x_target) + + return x + + def x_update(self, x, z, u, x_t): + """Update x in the ADMM iteration.""" + # x = (wp * I + A^T A + p I)^-1 (wp * x_orig + p (z - u)) + r = self.pos_weight * x_t + self.admm_stepsize * (z - u) + x[:] = self.system_lu.solve(r) + + def z_update(self, z, x, z_t, u): + """Update z in the ADMM iteration using vectorized operations.""" + # Compute the difference from target for all margin locations at once + z[:] = x + u - z_t + + # Check if we need to project back to margin + z_diff_norms = np.linalg.norm(z, axis=1) + mask = z_diff_norms > self.margin_vals + if np.any(mask): + scale_factors = self.margin_vals[mask] / z_diff_norms[mask] + z[mask] *= scale_factors[:, np.newaxis] + + # Add back the target + z[:] += z_t + + if self.circle_project: + z[:] = z / (np.linalg.norm(z, axis=1, keepdims=True) + 1.0e-6) + + def u_update(self, u, x, z): + """Update u in the ADMM iteration using vectorized operations.""" + u[:] += self.alpha_overrelax * (x - z) + + +def smooth_signal(x, margins, pos_weight=0, alpha_overrelax=1.8, admm_iters=500, circle_project=False): + x_smoothed = x.copy() + x_smoothed[:] = x.mean(axis=0, keepdims=True) + + # smooth the signal, multigrid style by starting out coarse, + # doubling the resolution and repeating until we're at the full + # resolution, using the previous result as the initial guess. + levels = int(math.floor(math.log2(len(x)))) + levels = max(levels - 4, 1) + + stepsize = 2**levels + while True: + # smooth signals at this level: + num_steps = len(x_smoothed[::stepsize]) + smoother = TrajectorySmoother( + margins=margins[::stepsize], + pos_weight=pos_weight, + alpha_overrelax=alpha_overrelax, + admm_iters=admm_iters, + circle_project=circle_project, + ) + x_smoothed[::stepsize] = smoother.smooth(x[::stepsize], x_smoothed[::stepsize]) + + # interpolate to next level: + next_stepsize = stepsize // 2 + num_interleaved = len(x_smoothed[next_stepsize::stepsize]) + if num_interleaved == num_steps: + # linearly extrapolate the last value if we have to: + x_smoothed[next_stepsize::stepsize][-1] = ( + x_smoothed[::stepsize][-1] + (x_smoothed[::stepsize][-1] - x_smoothed[::stepsize][-2]) / 2 + ) + num_interleaved = num_interleaved - 1 + + # linearly interpolate the remaining values: + x_smoothed[next_stepsize::stepsize][:num_interleaved] = ( + x_smoothed[::stepsize][:-1] + x_smoothed[::stepsize][1:] + ) / 2 + + if stepsize == 1: + break + + stepsize //= 2 + + return x_smoothed + + +def get_smooth_root_pos(hip_translations): + root_translations_xz = hip_translations[..., [0, 2]] + root_translations_y = hip_translations[..., [1]] + + batch_size, nframes = root_translations_xz.shape[:2] + margins = np.full(root_translations_xz.shape[1], 0.06) + + root_translations_smoothed_xz = [] + for batch in range(batch_size): + root_translations_smoothed_xz.append( + smooth_signal(root_translations_xz[batch].detach().cpu().numpy(), margins)[None] + ) + + root_translations_smoothed_xz = torch.tensor(np.concatenate(root_translations_smoothed_xz)) + + root_translations = torch.cat( + [ + root_translations_smoothed_xz.to(root_translations_y.device), + root_translations_y, + ], + dim=-1, + )[..., [0, 2, 1]] + + return root_translations + + +def smooth_trajectory_gpu( + positions: torch.Tensor, + margin: float = 0.06, + smoothness_weight: float = 1.0, + position_weight: float = 0.01, + num_iters: int = 50, +) -> torch.Tensor: + """GPU-accelerated trajectory smoothing using gradient descent. + + Args: + positions: [N, D] tensor of positions to smooth + margin: Maximum allowed deviation from original positions + smoothness_weight: Weight for smoothness term (minimizes acceleration) + position_weight: Weight for position preservation term + num_iters: Number of optimization iterations + + Returns: + Smoothed positions [N, D] + """ + device = positions.device + n_frames = positions.shape[0] + + if n_frames < 3: + return positions.clone() + + # Initialize smoothed positions with original + smoothed = positions.clone().requires_grad_(True) + optimizer = torch.optim.Adam([smoothed], lr=0.01) + + for _ in range(num_iters): + optimizer.zero_grad() + + # Smoothness loss: minimize second derivative (acceleration) + vel = smoothed[1:] - smoothed[:-1] + acc = vel[1:] - vel[:-1] + smoothness_loss = smoothness_weight * torch.sum(acc**2) + + # Position preservation loss (soft constraint) + position_loss = position_weight * torch.sum((smoothed - positions) ** 2) + + # Total loss + loss = smoothness_loss + position_loss + + loss.backward() + optimizer.step() + + # Project back to margin constraints + with torch.no_grad(): + diff = smoothed - positions + diff_norm = torch.norm(diff, dim=1, keepdim=True) + mask = diff_norm > margin + if mask.any(): + scale = margin / (diff_norm[mask] + 1e-8) + smoothed[mask] = positions[mask] + diff[mask] * scale + + return smoothed.detach() + + +def smooth_trajectory_gpu_fast( + positions: torch.Tensor, + margin: float = 0.06, + smoothness_weight: float = 100.0, + position_weight: float = 0.1, +) -> torch.Tensor: + """ + Fast GPU-accelerated trajectory smoothing using direct least squares solution. + No iterative optimization - solves in one shot. + + Args: + positions: [N, D] tensor of positions to smooth + margin: Maximum allowed deviation from original positions + smoothness_weight: Weight for smoothness term (minimizes acceleration) + position_weight: Weight for position preservation term + + Returns: + Smoothed positions [N, D] + """ + device = positions.device + dtype = positions.dtype + n_frames = positions.shape[0] + n_dims = positions.shape[1] + + if n_frames < 3: + return positions.clone() + + # Build acceleration matrix A for second derivative + # A[i] computes: x[i-1] - 2*x[i] + x[i+1] + A = torch.zeros((n_frames - 2, n_frames), device=device, dtype=dtype) + for i in range(n_frames - 2): + A[i, i] = 1.0 + A[i, i + 1] = -2.0 + A[i, i + 2] = 1.0 + + # Solve: minimize smoothness_weight * ||Ax||^2 + position_weight * ||x - x_orig||^2 + # Solution: x = (smoothness_weight * A^T A + position_weight * I)^(-1) * (position_weight * x_orig) + + ATA = A.T @ A # [N, N] + M = smoothness_weight * ATA + position_weight * torch.eye(n_frames, device=device, dtype=dtype) + + # Solve for each dimension independently + smoothed = torch.zeros_like(positions) + for d in range(n_dims): + b = position_weight * positions[:, d] # [N] + # Use torch.linalg.solve for numerical stability + smoothed[:, d] = torch.linalg.solve(M, b) + + # Apply margin constraints + diff = smoothed - positions + diff_norm = torch.norm(diff, dim=1) # [N] - no keepdim to avoid shape mismatch + mask = diff_norm > margin # [N] boolean mask + if mask.any(): + scale = margin / (diff_norm[mask] + 1e-8) # [M] where M = number of True in mask + # Need to broadcast scale [M] to match diff[mask] which is [M, D] + smoothed[mask] = positions[mask] + diff[mask] * scale[:, None] + + return smoothed + + +def get_smooth_root_pos_gpu(hip_translations: torch.Tensor) -> torch.Tensor: + """GPU-accelerated version of get_smooth_root_pos. + + Args: + hip_translations: [B, T, 3] root positions + + Returns: + Smoothed root positions [B, T, 3] + """ + device = hip_translations.device + root_translations_xz = hip_translations[..., [0, 2]] # [B, T, 2] + root_translations_y = hip_translations[..., [1]] # [B, T, 1] + + batch_size = root_translations_xz.shape[0] + + # Smooth XZ components + root_translations_smoothed_xz = [] + for batch_idx in range(batch_size): + smoothed_xz = smooth_trajectory_gpu( + root_translations_xz[batch_idx], # [T, 2] + margin=0.06, + smoothness_weight=1.0, + position_weight=0.01, + num_iters=50, + ) + root_translations_smoothed_xz.append(smoothed_xz[None]) # [1, T, 2] + + root_translations_smoothed_xz = torch.cat(root_translations_smoothed_xz, dim=0) # [B, T, 2] + + # Combine smoothed XZ with original Y + root_translations = torch.cat( + [ + root_translations_smoothed_xz, + root_translations_y, + ], + dim=-1, + ) # [B, T, 3] with order [x, z, y] + + # Reorder to [x, y, z] + root_translations = root_translations[..., [0, 2, 1]] + + return root_translations + + +def get_smooth_root_pos_gpu_fast(hip_translations: torch.Tensor) -> torch.Tensor: + """Fast GPU-accelerated version of get_smooth_root_pos using direct solver. + + Args: + hip_translations: [B, T, 3] root positions + + Returns: + Smoothed root positions [B, T, 3] + """ + device = hip_translations.device + root_translations_xz = hip_translations[..., [0, 2]] # [B, T, 2] + root_translations_y = hip_translations[..., [1]] # [B, T, 1] + + batch_size = root_translations_xz.shape[0] + + # Smooth XZ components using fast method + root_translations_smoothed_xz = [] + for batch_idx in range(batch_size): + smoothed_xz = smooth_trajectory_gpu_fast( + root_translations_xz[batch_idx], # [T, 2] + margin=0.06, + smoothness_weight=100.0, + position_weight=0.1, + ) + root_translations_smoothed_xz.append(smoothed_xz[None]) # [1, T, 2] + + root_translations_smoothed_xz = torch.cat(root_translations_smoothed_xz, dim=0) # [B, T, 2] + + # Combine smoothed XZ with original Y + root_translations = torch.cat( + [ + root_translations_smoothed_xz, + root_translations_y, + ], + dim=-1, + ) # [B, T, 3] with order [x, z, y] + + # Reorder to [x, y, z] + root_translations = root_translations[..., [0, 2, 1]] + + return root_translations diff --git a/ardy/motion_rep/stats.py b/ardy/motion_rep/stats.py new file mode 100644 index 0000000000000000000000000000000000000000..393a0a19b29e6bc6a0a0bb70cffe99f3b5eab7e2 --- /dev/null +++ b/ardy/motion_rep/stats.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Feature normalization statistics (mean/std) for motion representations.""" + +import logging +import os +from typing import Optional + +import numpy as np +import torch + +log = logging.getLogger(__name__) + + +class Stats(torch.nn.Module): + """Utility module for feature normalization statistics. + + Normalization follows: + ``(data - mean) / sqrt(std**2 + eps)`` + """ + + def __init__( + self, + folder: Optional[str] = None, + load: bool = True, + load_as_fp32: bool = True, + eps=1e-05, + ): + super().__init__() + self.folder = folder + self.eps = eps + self.load_as_fp32 = load_as_fp32 + if folder is not None and load: + self.load() + + def sliced(self, indices): + """Return a new ``Stats`` object containing selected feature indices.""" + new_stats = Stats(folder=self.folder, load=False, eps=self.eps) + new_stats.register_from_tensors( + self.mean[..., indices].clone(), + self.std[..., indices].clone(), + ) + return new_stats + + def load(self): + """Load ``mean.npy`` and ``std.npy`` from ``self.folder``.""" + mean_path = os.path.join(self.folder, "mean.npy") + std_path = os.path.join(self.folder, "std.npy") + if not os.path.exists(mean_path) or not os.path.exists(std_path): + raise FileNotFoundError( + f"Missing stats files in '{self.folder}'. Expected:\n" + f" - {mean_path}\n" + f" - {std_path}\n\n" + "These files ship alongside the model checkpoint. Re-download the " + "checkpoint, or check that the checkpoint directory you passed is " + "complete." + ) + + mean = torch.from_numpy(np.load(mean_path)) + std = torch.from_numpy(np.load(std_path)) + if self.load_as_fp32: + mean = mean.to(dtype=torch.float32) + std = std.to(dtype=torch.float32) + self.register_from_tensors(mean, std) + + def register_from_tensors(self, mean: torch.Tensor, std: torch.Tensor): + """Register mean/std tensors as non-persistent buffers.""" + self.register_buffer("mean", mean, persistent=False) + self.register_buffer("std", std, persistent=False) + # Precomputed eps-adjusted denominator so normalize/unnormalize skip + # the per-call sqrt. + self.register_buffer("std_eps", torch.sqrt(std**2 + self.eps), persistent=False) + + def normalize(self, data: torch.Tensor) -> torch.Tensor: + """Normalize data using the stored statistics.""" + mean = self.mean.to(device=data.device, dtype=data.dtype) + std_eps = self.std_eps.to(device=data.device, dtype=data.dtype) + return (data - mean) / std_eps + + def unnormalize(self, data: torch.Tensor) -> torch.Tensor: + """Undo normalization using the stored statistics.""" + mean = self.mean.to(device=data.device, dtype=data.dtype) + std_eps = self.std_eps.to(device=data.device, dtype=data.dtype) + return data * std_eps + mean + + def is_loaded(self): + """Return whether statistics are currently available.""" + return hasattr(self, "mean") + + def get_dim(self): + """Return feature dimensionality.""" + return self.mean.shape[0] + + def save( + self, + folder: Optional[str] = None, + mean: Optional[torch.Tensor] = None, + std: Optional[torch.Tensor] = None, + ): + """Save statistics to ``folder`` as ``mean.npy`` and ``std.npy``.""" + if folder is None: + folder = self.folder + if folder is None: + raise ValueError("No folder to save stats") + + if mean is None and std is None: + try: + mean = self.mean.cpu().numpy() + std = self.std.cpu().numpy() + except AttributeError: + raise ValueError("Stats were not loaded") + + # don't override stats folder + os.makedirs(folder, exist_ok=False) + + np.save(os.path.join(folder, "mean.npy"), mean) + np.save(os.path.join(folder, "std.npy"), std) + + def __eq__(self, other): + return (self.mean.cpu() == other.mean.cpu()).all() and (self.std.cpu() == other.std.cpu()).all() + + # should define a hash value for pytorch, as we defined __eq__ + def __hash__(self): + # Convert mean and std to bytes for a consistent hash value + mean_hash = hash(self.mean.detach().cpu().numpy().tobytes()) + std_hash = hash(self.std.detach().cpu().numpy().tobytes()) + return hash((mean_hash, std_hash)) + + def __repr__(self): + return f'Stats(folder="{self.folder}")' diff --git a/ardy/motion_rep/tools.py b/ardy/motion_rep/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..fcfe089458568182184393fcbee11deec8d1d4e4 --- /dev/null +++ b/ardy/motion_rep/tools.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import List, Optional, Union + +import einops +import torch + +from ..geometry import cont6d_to_matrix, matrix_to_cont6d +from ..tools import ensure_batched + + +def diff_angles(angles, fps: float): + """Computes differences between angles. + + Args: + angles (Tensor): [..., T] the batched sequences of rotation angles in radians. + + Returns: + Tensor: [..., T-1] the difference between consecutive angles + """ + cos = torch.cos(angles) + sin = torch.sin(angles) + + cos_diff = cos[..., 1:] * cos[..., :-1] + sin[..., 1:] * sin[..., :-1] + sin_diff = sin[..., 1:] * cos[..., :-1] - cos[..., 1:] * sin[..., :-1] + + # should be close to angles.diff() but more robust + # multiply by fps = 1 / dt + angles_diff = fps * torch.arctan2(sin_diff, cos_diff) + return angles_diff + + +@ensure_batched(positions=4, lengths=1) +def compute_vel_xyz( + positions: torch.Tensor, + fps: float, + lengths: Optional[torch.Tensor] = None, +): + """Compute the velocities from positions: dx/dt. + + Works with batches. The last velocity is duplicated to keep the same size. + Args: + positions (torch.Tensor): [..., T, J, 3] xyz positions of a human skeleton + fps (float): frame per seconds + lengths (Optional[torch.Tensor]): [...] size of each input batched. If not provided, positions should not be batched + + Returns: + velocity (torch.Tensor): [..., T, J, 3] velocities computed from the positions + """ + device = positions.device + + if lengths is None: + assert positions.shape[0] == 1, "If lenghts is not provided, the input should not be batched." + lengths = torch.tensor([len(positions)], device=device) + + # compute velocities with fps + velocity = fps * (positions[:, 1:] - positions[:, :-1]) + # pading the velocity vector + vel_pad = torch.zeros_like(velocity[:, 0]) + velocity, _ = einops.pack([velocity, vel_pad], "batch * nbjoints dim") + + # repeat the last velocities + # with special care for different lengths with batches + # Use gather/scatter instead of in-place advanced indexing (TRT-compatible) + nj, nd = velocity.shape[2], velocity.shape[3] + # clamp: for length-1 sequences src would be -1 (invalid for gather); + # index 0 is the zero pad, so the copy is a harmless no-op there + src_idx = (lengths - 2).clamp(min=0).long().view(-1, 1, 1, 1).expand(-1, 1, nj, nd) + dst_idx = (lengths - 1).long().view(-1, 1, 1, 1).expand(-1, 1, nj, nd) + velocity = velocity.scatter(1, dst_idx, torch.gather(velocity, 1, src_idx)) + return velocity + + +@ensure_batched(root_rot_angles=2, lengths=1) +def compute_vel_angle( + root_rot_angles: torch.Tensor, + fps: float, + lengths: Optional[torch.Tensor] = None, +): + """Compute the local root rotation velocity: dtheta/dt. + + Args: + root_rot_angles (torch.Tensor): [..., T] rotation angle (in radian) + fps (float): frame per seconds + lengths (Optional[torch.Tensor]): [...] size of each input batched. If not provided, root_rot_angles should not be batched + + Returns: + local_root_rot_vel (torch.Tensor): [..., T] local root rotation velocity (in radian/s) + """ + device = root_rot_angles.device + if lengths is None: + assert root_rot_angles.shape[0] == 1, "If lenghts is not provided, the input should not be batched." + lengths = torch.tensor([len(root_rot_angles)], device=device) + + local_root_rot_vel = diff_angles(root_rot_angles, fps) + pad_rot_vel_angles = torch.zeros_like(root_rot_angles[:, 0]) + local_root_rot_vel, _ = einops.pack( + [local_root_rot_vel, pad_rot_vel_angles], + "batch *", + ) + # repeat the last rotation angle + # with special care for different lengths with batches + # Use gather/scatter instead of in-place advanced indexing (TRT-compatible) + # clamp: for length-1 sequences src would be -1 (invalid for gather); + # index 0 is the zero pad, so the copy is a harmless no-op there + src_idx = (lengths - 2).clamp(min=0).unsqueeze(-1).long() # [B, 1] + dst_idx = (lengths - 1).unsqueeze(-1).long() # [B, 1] + local_root_rot_vel = local_root_rot_vel.scatter(1, dst_idx, torch.gather(local_root_rot_vel, 1, src_idx)) + return local_root_rot_vel + + +@ensure_batched(posed_joints=4) +def compute_heading_angle(posed_joints: torch.Tensor, skeleton): + """Compute the heading direction from the joint positions, by looking at the hip vector. + + Args: + posed_joints (torch.Tensor): [B, T, J, 3] global positions + skeleton (SkeletonBase): skeleton of the human, used to find location of hips + Returns: + heading (torch.Tensor): [B] heading angle + """ + # compute root heading for the sequence from hip positions + r_hip, l_hip = skeleton.hip_joint_idx + diff = posed_joints[:, :, r_hip] - posed_joints[:, :, l_hip] + heading_angle = torch.atan2(diff[..., 2], -diff[..., 0]) + return heading_angle + + +def length_to_mask( + length: Union[torch.Tensor, List], + max_len: Optional[int] = None, + device=None, +) -> torch.Tensor: + if isinstance(length, list): + if device is None: + device = "cpu" + length = torch.tensor(length, device=device) + + if device is not None: + assert device == length.device + device = length.device + + if max_len is None: + max_len = max(length) + + mask = torch.arange(max_len, device=device).expand(len(length), max_len) < length.unsqueeze(1) + return mask + + +class RotateFeatures: + """Helper that applies a global heading rotation to motion features.""" + + def __init__(self, angle: torch.Tensor): + """Precompute 2D and 3D rotation matrices for a batch of angles. + + Args: + angle: Rotation angle(s) in radians, shaped ``[B]``. + """ + self.angle = angle + + cos, sin = torch.cos(angle), torch.sin(angle) + one, zero = torch.ones_like(angle), torch.zeros_like(angle) + + # 2D rotation transposed (sin are -sin) + self.corrective_mat_2d_T = torch.stack((cos, sin, -sin, cos), -1).reshape(angle.shape + (2, 2)) + # 3D rotation on Y axis + self.corrective_mat_Y = torch.stack((cos, zero, sin, zero, one, zero, -sin, zero, cos), -1).reshape( + angle.shape + (3, 3) + ) + self.corrective_mat_Y_T = self.corrective_mat_Y.transpose(-2, -1).contiguous() + + def rotate_positions(self, positions: torch.Tensor): + """Rotate 3D positions around the Y axis.""" + return positions @ self.corrective_mat_Y_T + + def rotate_2d_positions(self, positions_2d: torch.Tensor): + """Rotate 2D ``(x, z)`` vectors in the ground plane.""" + return positions_2d @ self.corrective_mat_2d_T + + def rotate_rotations(self, rotations: torch.Tensor): + """Left-multiply global rotation matrices by the heading correction.""" + # "Rotate" the global rotations + # which means add an extra Y rotation after the transform + # so at the left R' = R_y R + # (since we use the convention x' = R x) + # "bik,btdkj->btdij" + B, T, J = rotations.shape[:3] + BTJ = B * T * J + return ( + self.corrective_mat_Y[:, None, None].expand(B, T, J, 3, 3).reshape(BTJ, 3, 3) @ rotations.reshape(BTJ, 3, 3) + ).reshape(B, T, J, 3, 3) + + def rotate_6d_rotations(self, rotations_6d: torch.Tensor): + """Rotate 6D rotation features via matrix conversion.""" + return matrix_to_cont6d(self.rotate_rotations(cont6d_to_matrix(rotations_6d))) diff --git a/ardy/postprocess.py b/ardy/postprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..a8b2f14919a85b724ef7116be92379a607d863f5 --- /dev/null +++ b/ardy/postprocess.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Post-processing utilities for motion generation output.""" + +from types import SimpleNamespace +from typing import Dict, List, Optional, Tuple + +import numpy as np +import torch + +from .constraints import ( + EndEffectorConstraintSet, + FullBodyConstraintSet, + Root2DConstraintSet, +) +from .geometry import matrix_to_quaternion, quaternion_to_matrix +from .skeleton import G1Skeleton34, SkeletonBase, fk + + +def extract_input_motion_from_constraints( + constraint_lst: List, + skeleton: SkeletonBase, + num_frames: int, + num_joints: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Extract hip translations and local rotations from constraints for postprocessing. + + Args: + constraint_lst: List of constraints (FullBodyConstraintSet, EndEffectorConstraintSet, etc.) + skeleton: Skeleton instance + num_frames: Total number of frames in the motion + num_joints: Number of joints + + Returns: + Tuple of (hip_translations_input, rotations_input): + - hip_translations_input: Hip translations, shape (T, 3) + - rotations_input: Local joint rotations as quaternions, shape (T, J, 4) + """ + # Initialize with zeros for all frames + hip_translations_input = torch.zeros(num_frames, 3) + rotations_input = torch.zeros(num_frames, num_joints, 4) + rotations_input[..., 0] = 1.0 # Initialize as identity quaternions (w=1, x=y=z=0) + + def _match_hip_dtype(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to(device=hip_translations_input.device, dtype=hip_translations_input.dtype) + + def _match_rot_dtype(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to(device=rotations_input.device, dtype=rotations_input.dtype) + + if not constraint_lst: + return hip_translations_input, rotations_input + + # Sort constraints to ensure FullBodyConstraintSet is processed last. + # This gives fullbody constraints the last say on hip translations. + sorted_constraints = sorted(constraint_lst, key=lambda c: isinstance(c, FullBodyConstraintSet)) + + for constraint in sorted_constraints: + frame_indices = constraint.frame_indices + + # Filter any constraint frames that exceed num_frames. + if isinstance(frame_indices, torch.Tensor): + valid_mask = frame_indices < num_frames + if valid_mask.sum() == 0: + continue + frame_indices = frame_indices[valid_mask] + else: + valid_positions = [i for i, idx in enumerate(frame_indices) if idx < num_frames] + if not valid_positions: + continue + frame_indices = [frame_indices[i] for i in valid_positions] + + # Handle Root2DConstraintSet separately - only assign root_2d at xz dimensions + if isinstance(constraint, Root2DConstraintSet): + root_2d = getattr(constraint, "root_2d", None) # (K, 2) + if isinstance(frame_indices, torch.Tensor): + root_2d = root_2d[valid_mask] + else: + root_2d = root_2d[valid_positions] + root_2d = _match_hip_dtype(root_2d) # (K, 2) where K = len(frame_indices) + hip_translations_input[frame_indices, 0] = root_2d[:, 0] # x + hip_translations_input[frame_indices, 2] = root_2d[:, 1] # z + continue + elif isinstance(constraint, FullBodyConstraintSet) or isinstance(constraint, EndEffectorConstraintSet): + global_rots = constraint.global_joints_rots # (K, J, 3, 3) where K = len(frame_indices) + global_positions = constraint.global_joints_positions # (K, J, 3) + if isinstance(frame_indices, torch.Tensor): + global_rots = global_rots[valid_mask] + global_positions = global_positions[valid_mask] + root_2d = constraint.root_2d[valid_mask] + else: + global_rots = global_rots[valid_positions] + global_positions = global_positions[valid_positions] + root_2d = constraint.root_2d[valid_positions] + + root_positions = global_positions[:, skeleton.root_idx] # (K, 3) + # For EE constraints that do not include Hips, overwrite xz with root_2d. + # If Hips is constrained, we want to keep its exact root trajectory. + if isinstance(constraint, EndEffectorConstraintSet) and "Hips" not in getattr( + constraint, "joint_names", [] + ): + root_positions[:, 0] = root_2d[:, 0] # x + root_positions[:, 2] = root_2d[:, 1] # z + + local_rot_mats = skeleton.global_rots_to_local_rots(global_rots) # (K, J, 3, 3) + local_rot_quats = matrix_to_quaternion(local_rot_mats) # (K, J, 4) + + hip_translations_input[frame_indices] = _match_hip_dtype(root_positions) + rotations_input[frame_indices] = _match_rot_dtype(local_rot_quats) + else: + NotImplementedError(f"Constraint {constraint.name} is not supported") + + return hip_translations_input, rotations_input + + +def create_working_rig_from_skeleton( + skeleton: SkeletonBase, above_ground_offset: float = 0.007 +) -> List[SimpleNamespace]: + """Create the working rig as a list of SimpleNamespace objects from skeleton. + + Args: + skeleton: SkeletonBase instance with bone_order_names, neutral_joints, joint_parents + above_ground_offset: Additional offset to position the rig slightly above ground + Returns: + List of SimpleNamespace objects representing the working rig + """ + working_rig_joints = [] + + joint_names = skeleton.bone_order_names + neutral_positions = skeleton.neutral_joints.cpu().numpy() + parent_indices = skeleton.joint_parents.cpu().numpy() + + if isinstance(skeleton, G1Skeleton34): + retarget_map = { + skeleton.bone_order_names[skeleton.root_idx]: "Hips", + skeleton.left_hand_joint_names[0]: "LeftHand", + skeleton.right_hand_joint_names[0]: "RightHand", + skeleton.left_foot_joint_names[0]: "LeftFoot", + skeleton.right_foot_joint_names[0]: "RightFoot", + } + else: + retarget_map = { + "Hips": "Hips", + "Head": "Head", + "LeftHand": "LeftHand", + "RightHand": "RightHand", + "LeftFoot": "LeftFoot", + "RightFoot": "RightFoot", + } + + for i, joint_name in enumerate(joint_names): + parent_name = None if parent_indices[i] == -1 else joint_names[parent_indices[i]] + + # Calculate local translation relative to parent + if parent_indices[i] == -1: + # Move the rig so that the lowest point (toe) is at ground level (y=0), + # plus a small offset to position the rig slightly above ground + toe_height = neutral_positions[:, 1].min() # lowest y-coordinate (toe) + local_translation = ( + neutral_positions[i] + np.array([0.0, -toe_height + above_ground_offset, 0.0]) + ).tolist() + else: + parent_idx = parent_indices[i] + parent_position = neutral_positions[parent_idx] + joint_position = neutral_positions[i] + local_translation = (joint_position - parent_position).tolist() + + # Default rotation (identity quaternion: x=0, y=0, z=0, w=1) + default_rotation = [0.0, 0.0, 0.0, 1.0] + + joint_info = SimpleNamespace( + name=joint_name, + parent=parent_name, + t_pose_rotation=default_rotation, + t_pose_translation=local_translation, + retarget_tag=retarget_map.get(joint_name), + ) + + working_rig_joints.append(joint_info) + + return working_rig_joints + + +def post_process_motion( + local_rot_mats: torch.Tensor, + root_positions: torch.Tensor, + contacts: torch.Tensor, + skeleton: SkeletonBase, + constraint_lst: Optional[List] = None, + contact_threshold: float = 0.5, + root_margin: float = 0.04, +) -> Dict[str, torch.Tensor]: + """Post-process generated motion to reduce foot skating and improve quality. + + Args: + local_rot_mats: Local joint rotation matrices, shape (B, T, J, 3, 3) + root_positions: Root joint positions, shape (B, T, 3) + contacts: Foot contact labels, shape (B, T, num_contacts) + skeleton: Skeleton instance + constraint_lst: Optional list of constraints (or list of lists of constraints for batched inference)(FullBodyConstraintSet, etc.) + contact_threshold: Threshold for foot contact detection + root_margin: Margin for root position correction + + Returns: + Dictionary with corrected motion data: + - local_rot_mats: Corrected local rotation matrices (B, T, J, 3, 3) + - root_positions: Corrected root positions (B, T, 3) + - posed_joints: Corrected global joint positions (B, T, J, 3) + - global_rot_mats: Corrected global rotation matrices (B, T, J, 3, 3) + """ + # Ensure batch dimension + assert local_rot_mats.dim() == 5, "local_rot_mats should be 5D, make sure to include the batch dimension" + + batch_size, num_frames, num_joints = local_rot_mats.shape[:3] + + def _build_constraint_masks_dict(constraints: List) -> Dict[str, torch.Tensor]: + out = { + key: torch.zeros(num_frames, dtype=torch.float32) + for key in [ + "FullBody", + "LeftFoot", + "RightFoot", + "LeftHand", + "RightHand", + "Root", + ] + } + for constraint in constraints: + frame_indices = constraint.frame_indices + if isinstance(frame_indices, torch.Tensor): + frame_indices = frame_indices[frame_indices < num_frames] + if frame_indices.numel() == 0: + continue + else: + frame_indices = [idx for idx in frame_indices if idx < num_frames] + if not frame_indices: + continue + if constraint.name == "fullbody": + out["FullBody"][frame_indices] = 1.0 + elif constraint.name == "left-foot": + out["LeftFoot"][frame_indices] = 1.0 + elif constraint.name == "right-foot": + out["RightFoot"][frame_indices] = 1.0 + elif constraint.name == "left-hand": + out["LeftHand"][frame_indices] = 1.0 + elif constraint.name == "right-hand": + out["RightHand"][frame_indices] = 1.0 + elif constraint.name == "root2d": + out["Root"][frame_indices] = 1.0 + return out + + # Create constraint masks from constraint_lst (one dict per batch item when batched). + batched_constraints = bool(constraint_lst) and isinstance(constraint_lst[0], list) + if batched_constraints: + constraint_masks_dict_lst = [_build_constraint_masks_dict(constraint_lst[b]) for b in range(batch_size)] + else: + constraint_masks_dict = ( + _build_constraint_masks_dict(constraint_lst) + if constraint_lst + else { + key: torch.zeros(num_frames, dtype=torch.float32) + for key in [ + "FullBody", + "LeftFoot", + "RightFoot", + "LeftHand", + "RightHand", + "Root", + ] + } + ) + + # Create working rig + working_rig = create_working_rig_from_skeleton(skeleton) + has_double_ankle_joints = isinstance(skeleton, G1Skeleton34) + + # Prepare input tensors. The generated motion will be modified in place. Clone first. + neutral_joints_pelvis_offset = skeleton.neutral_joints[0].cpu().clone() + hip_translations_corrected = root_positions.cpu().clone() + rotations_corrected = matrix_to_quaternion(local_rot_mats).cpu().clone() # (B, T, J, 4) + contacts = contacts.cpu() + + # Extract input motion (target keyframes) from constraints for each batch + # For constrained keyframes, use the original motion from constraints + # For non-constrained frames, zeros are used + hip_translations_input = torch.zeros(batch_size, num_frames, 3) + rotations_input = torch.zeros(batch_size, num_frames, num_joints, 4) + rotations_input[..., 0] = 1.0 # Initialize as identity quaternions (w=1, x=y=z=0) + + if constraint_lst: + for b in range(batch_size): + # Get constraints for this batch item (if batched) or use the same list + constraints_lst_el = ( + constraint_lst[b] + if isinstance( + constraint_lst[0], list + ) # when the constraint_list is in batch format, each item in a list is a constraintlist for one sample + else constraint_lst # single constraint list shared for all samples in the batch + ) + hip_translations_input[b], rotations_input[b] = extract_input_motion_from_constraints( + constraints_lst_el, + skeleton, + num_frames, + num_joints, + ) + + # Call the motion correction for each batch (optional package). + try: + from motion_correction import motion_postprocess + except ImportError as e: + raise RuntimeError( + "Motion correction is required for this postprocessing path but the " + "motion_correction package is not installed. Install with: pip install -e ." + ) from e + for b in range(batch_size): + motion_postprocess.correct_motion( + hip_translations_corrected[b : b + 1], + rotations_corrected[b : b + 1], + contacts[b : b + 1], + hip_translations_input[b : b + 1], + rotations_input[b : b + 1], + constraint_masks_dict_lst[b] if batched_constraints else constraint_masks_dict, + contact_threshold, + root_margin, + working_rig, + has_double_ankle_joints, + ) + + local_rot_mats_corrected = quaternion_to_matrix(rotations_corrected) + + # Compute posed joints using FK + device = local_rot_mats.device + global_rot_mats, posed_joints, _ = fk( + local_rot_mats_corrected.to(device), + hip_translations_corrected.to(device), + skeleton, + ) + + result = { + "local_rot_mats": local_rot_mats_corrected.to(device), + "root_positions": hip_translations_corrected.to(device), + "posed_joints": posed_joints, + "global_rot_mats": global_rot_mats, + } + + return result diff --git a/ardy/skeleton/__init__.py b/ardy/skeleton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e22d2e26385d16e86b62d576c552047cb697dda7 --- /dev/null +++ b/ardy/skeleton/__init__.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Skeleton definitions and utilities used across ARDY.""" + +from .base import SkeletonBase +from .definitions import ( + CoreSkeleton27, + G1Skeleton34, + SOMASkeleton30, + SOMASkeleton77, +) +from .kinematics import batch_rigid_transform, fk +from .registry import build_skeleton +from .transforms import global_rots_to_local_rots, to_standard_tpose + +__all__ = [ + "SkeletonBase", + "G1Skeleton34", + "SOMASkeleton30", + "SOMASkeleton77", + "CoreSkeleton27", + "batch_rigid_transform", + "fk", + "build_skeleton", + "global_rots_to_local_rots", + "to_standard_tpose", +] diff --git a/ardy/skeleton/base.py b/ardy/skeleton/base.py new file mode 100644 index 0000000000000000000000000000000000000000..81937e5b64e838b89e6834fe41f8f75e7bf515d2 --- /dev/null +++ b/ardy/skeleton/base.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Base skeleton class: hierarchy, joint metadata, and helpers for kinematics and motion.""" + +from pathlib import Path +from typing import Optional + +import torch + +from ardy.assets import skeleton_asset_path + +from .kinematics import fk +from .transforms import ( + from_standard_tpose, + global_rots_to_local_rots, + to_standard_tpose, +) + + +class SkeletonBase(torch.nn.Module): + """Base class that stores a skeleton hierarchy and helper metadata. + + Subclasses define the static joint layout (joint names and parent links) and semantic groups + (feet, hands, hips). This class builds index mappings, parent tensors, and convenience helpers + used by kinematics, constraints, and motion conversion utilities. + """ + + # these should be defined in the subclass + name = None + bone_order_names_with_parents = None + bone_order_names_no_root = None + root_idx = None + foot_joint_names = None + foot_joint_idx = None + hip_joint_names = None # in order [right, left] + hip_joint_idx = None # in order [right, left] + + def __init__( + self, + folder: Optional[str] = None, + name: Optional[str] = None, + load: bool = True, + **kwargs, # to catch addition args in configs + ): + """Initialize a skeleton instance and optional neutral-pose assets. + + Args: + folder: Folder containing serialized skeleton assets (for example + `joints.p` and optional `standard_t_pose_global_offsets_rots.p`). + name: Optional runtime name used to validate subclass compatibility. + load: Whether to load tensor assets from `folder`. + **kwargs: Unused extra config keys kept for config compatibility. + """ + super().__init__() + + if name is not None: + # Check that the name is not too far from the actual skeleton class name + assert self.name in name + self.name = name + + if folder is None: + # Take the skeleton asset folder of the repo from the name + # in case we don't override it + folder = str(skeleton_asset_path(self.name)) + + self.folder = folder + + self.dim = len(self.bone_order_names_with_parents) + + if load and folder is not None: + pfolder = Path(folder) + neutral_joints = torch.load(pfolder / "joints.p").squeeze() + self.register_buffer("neutral_joints", neutral_joints, persistent=False) + + if (pfolder / "bvh_joints.p").exists(): + bvh_neutral_joints = torch.load(pfolder / "bvh_joints.p").squeeze() + self.register_buffer("bvh_neutral_joints", bvh_neutral_joints, persistent=False) + + global_offset_path = pfolder / "standard_t_pose_global_offsets_rots.p" + if global_offset_path.exists(): + global_rot_offsets = torch.load(global_offset_path).squeeze() + self.register_buffer("global_rot_offsets", global_rot_offsets, persistent=False) + + self.bone_order_names = [x for x, y in self.bone_order_names_with_parents] + + self.bone_parents = dict(self.bone_order_names_with_parents) + self.bone_index = {x: idx for idx, x in enumerate(self.bone_order_names)} + self.bone_order_names_index = self.bone_index + + # create the parents tensor on the fly + joint_parents = torch.tensor( + [-1 if (y := self.bone_parents[x]) is None else self.bone_index[y] for x in self.bone_order_names] + ) + self.register_buffer("joint_parents", joint_parents, persistent=False) + + self.nbjoints = len(self.bone_order_names) + + # check lengths + assert self.nbjoints == len(self.joint_parents) + if "neutral_joints" in self.__dict__: + assert self.nbjoints == len(self.neutral_joints) + + root_indices = torch.where(joint_parents == -1)[0] + assert len(root_indices) == 1 # should be one root only + self.root_idx = root_indices[0].item() + + if "neutral_joints" in self.__dict__: + assert (self.neutral_joints[0] == 0).all() + + # remove the root + self.bone_order_names_no_root = ( + self.bone_order_names[: self.root_idx] + self.bone_order_names[self.root_idx + 1 :] + ) + + self.foot_joint_names = self.left_foot_joint_names + self.right_foot_joint_names + self.foot_joint_names_index = {x: idx for idx, x in enumerate(self.foot_joint_names)} + + self.left_foot_joint_idx = [ + self.bone_order_names.index(foot_joint) for foot_joint in self.left_foot_joint_names + ] + + self.right_foot_joint_idx = [ + self.bone_order_names.index(foot_joint) for foot_joint in self.right_foot_joint_names + ] + + self.foot_joint_idx = self.left_foot_joint_idx + self.right_foot_joint_idx + + self.hip_joint_idx = [self.bone_order_names.index(hip_joint) for hip_joint in self.hip_joint_names] + + def expand_joint_names(self, joint_names): + """Expand base EE names [LeftFoot, RightFoot, LeftHand, RightHand] actual joint names to + constrain position and rotations. + + Args: + joint_names: list of list of base EE names to constrain + + Returns: + rot_joint_names: list of list of joint names to constrain rotations + pos_joint_names: list of list of joint names to constrain positions + """ + base_ee = ["LeftFoot", "RightFoot", "LeftHand", "RightHand", "Hips"] + + pelvis_name = self.bone_order_names[self.root_idx] + + base_pos_names = [ + self.left_foot_joint_names, + self.right_foot_joint_names, + self.left_hand_joint_names, + self.right_hand_joint_names, + [pelvis_name], + ] + # base of each chain + base_rot_names = [ + self.left_foot_joint_names[:1], + self.right_foot_joint_names[:1], + self.left_hand_joint_names[:1], + self.right_hand_joint_names[:1], + [pelvis_name], + ] + rot_joint_names = [] + pos_joint_names = [] + # loop through each EE joint group to constrain in the current keyframe + for jname in joint_names: + idx = base_ee.index(jname) + rot_joint_names += base_rot_names[idx] + pos_joint_names += base_pos_names[idx] + return rot_joint_names, pos_joint_names + + def expand_joint_names_batched(self, joint_names): + """Expand base EE names [LeftFoot, RightFoot, LeftHand, RightHand] actual joint names to + constrain position and rotations. + + Args: + joint_names: list of list of base EE names to constrain + + Returns: + rot_joint_names: list of list of joint names to constrain rotations + pos_joint_names: list of list of joint names to constrain positions + """ + base_ee = ["LeftFoot", "RightFoot", "LeftHand", "RightHand", "Hips"] + + pelvis_name = self.bone_order_names[self.root_idx] + + base_pos_names = [ + self.left_foot_joint_names, + self.right_foot_joint_names, + self.left_hand_joint_names, + self.right_hand_joint_names, + [pelvis_name], + ] + # base of each chain + base_rot_names = [ + self.left_foot_joint_names[:1], + self.right_foot_joint_names[:1], + self.left_hand_joint_names[:1], + self.right_hand_joint_names[:1], + [pelvis_name], + ] + # loop through each keyframe + rot_joint_names = [] + pos_joint_names = [] + for key_joint_names in joint_names: + key_rot_names = [] + key_pos_names = [] + # loop through each EE joint group to constrain in the current keyframe + for jname in key_joint_names: + idx = base_ee.index(jname) + key_rot_names += base_rot_names[idx] + key_pos_names += base_pos_names[idx] + rot_joint_names.append(key_rot_names) + pos_joint_names.append(key_pos_names) + return rot_joint_names, pos_joint_names + + def __repr__(self): + if self.folder is None: + return f"{self.__class__.__name__}()" + return f'{self.__class__.__name__}(folder="{self.folder}")' + + @property + def device(self): + """Device where neutral-joint buffers are stored. + + Returns 'cpu' if neutral_joints is not present. + """ + if getattr(self, "neutral_joints", None) is None: + return "cpu" + return self.neutral_joints.device + + def fk(self, local_joint_rots: torch.Tensor, root_positions: torch.Tensor): + """Run forward kinematics for this skeleton layout. + + Args: + local_joint_rots: Local joint rotation matrices with shape + `(..., J, 3, 3)`. + root_positions: Root translations with shape `(..., 3)`. + + Returns: + Tuple of `(global_joint_rots, posed_joints, posed_joints_norootpos)`. + """ + global_joint_rots, posed_joints, posed_joints_norootpos = fk(local_joint_rots, root_positions, self) + return global_joint_rots, posed_joints, posed_joints_norootpos + + def to_standard_tpose(self, local_rot_mats: torch.Tensor): + """Convert local rotations into the skeleton's standard T-pose frame.""" + return to_standard_tpose(local_rot_mats, self) + + def from_standard_tpose(self, local_rot_mats: torch.Tensor): + """Convert local rotations from the skeleton's standard T-pose frame.""" + return from_standard_tpose(local_rot_mats, self) + + def global_rots_to_local_rots(self, global_joint_rots: torch.Tensor): + """Convert global joint rotations to local rotations for this hierarchy.""" + return global_rots_to_local_rots(global_joint_rots, self) + + def get_skel_slice(self, skeleton: "SkeletonBase"): + """Build index mapping from another skeleton into this skeleton order. + + Args: + skeleton: Source skeleton whose joint order is used by input tensors. + + Returns: + A list of source indices ordered as `self.bone_order_names`. + + Raises: + ValueError: If at least one required joint is missing from `skeleton`. + """ + try: + skel_slice = [skeleton.bone_index[x] for x in self.bone_order_names] + except KeyError: + raise ValueError("The current skeleton contain joints that are not in the input") + return skel_slice diff --git a/ardy/skeleton/bvh.py b/ardy/skeleton/bvh.py new file mode 100644 index 0000000000000000000000000000000000000000..6df00182183789459441ec594c656ba785de3dc7 --- /dev/null +++ b/ardy/skeleton/bvh.py @@ -0,0 +1,578 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BVH parsing utilities and skeleton/animation conversion helpers.""" + +import re +from typing import Optional, Tuple + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + + +class BvhNode: + """Lightweight tree node used to represent parsed BVH hierarchy lines.""" + + def __init__(self, value=[], parent=None): + """Create a node from tokenized BVH line values.""" + self.value = value + self.children = [] + self.parent = parent + if self.parent: + self.parent.add_child(self) + + def add_child(self, item): + """Attach a child node and set its parent reference.""" + item.parent = self + self.children.append(item) + + def filter(self, key): + """Yield direct children whose first token matches `key`.""" + for child in self.children: + if child.value[0] == key: + yield child + + def __iter__(self): + for child in self.children: + yield child + + def __getitem__(self, key): + """Return all tokens following `key` from the first matching child node.""" + for child in self.children: + for index, item in enumerate(child.value): + if item == key: + if index + 1 >= len(child.value): + return None + else: + return child.value[index + 1 :] + raise IndexError("key {} not found".format(key)) + + def __repr__(self): + return str(" ".join(self.value)) + + @property + def name(self): + """Joint name for `ROOT`/`JOINT` entries.""" + return self.value[1] + + +class Bvh: + """Parsed BVH file with hierarchy graph and per-frame channel values.""" + + def __init__(self, data: str, backend: Optional[str] = "graph"): + """ + Args: + data: Raw BVH file content. + backend: Parsing mode. `"graph"` keeps list-based frame storage, + while `"np"` precomputes a NumPy array and index caches. + """ + self.data = data + self.root = BvhNode() + self.frames = [] + self.backend = backend + self.tokenize() + if self.backend == "np": + # cache important info for quick access later + self.build_data_array() + elif self.backend == "graph": + pass + else: + raise ValueError(f"Unknown backend for BVH loading: {backend}") + + def build_data_array(self): + """Build cached channel indices and contiguous frame data for `"np"` backend.""" + joints = self.get_joints() + self.joint2idx = dict() + self.joint2channels = dict() + cur_idx = 0 + for joint in joints: + self.joint2idx[joint.value[1]] = cur_idx + cur_idx += int(joint["CHANNELS"][0]) + self.joint2channels[joint.value[1]] = joint["CHANNELS"][1:] + self.np_data_array = np.array(self.frames, dtype=np.float32) + + def tokenize(self): + """Tokenize BVH text and populate hierarchy plus frame values.""" + first_round = [] + accumulator = "" + for char in self.data: + if char not in ("\n", "\r"): + accumulator += char + elif accumulator: + first_round.append(re.split("\\s+", accumulator.strip())) + accumulator = "" + node_stack = [self.root] + frame_time_found = False + node = None + for item in first_round: + if frame_time_found: + self.frames.append(item) + continue + key = item[0] + if key == "{": + node_stack.append(node) + elif key == "}": + node_stack.pop() + else: + node = BvhNode(item) + # print("new node: ", node, "\nparent: ", node_stack[-1]) + node_stack[-1].add_child(node) + if item[0] == "Frame" and item[1] == "Time:": + frame_time_found = True + + def search(self, *items): + """Depth-first search for nodes matching a prefix of tokens.""" + found_nodes = [] + + def check_children(node): + if len(node.value) >= len(items): + failed = False + for index, item in enumerate(items): + if node.value[index] != item: + failed = True + break + if not failed: + found_nodes.append(node) + for child in node: + check_children(child) + + check_children(self.root) + return found_nodes + + def get_joints(self): + """Return all `ROOT`/`JOINT` hierarchy joints in BVH traversal order.""" + joints = [] + + def iterate_joints(joint): + joints.append(joint) + for child in joint.filter("JOINT"): + iterate_joints(child) + + iterate_joints(next(self.root.filter("ROOT"))) + return joints + + def get_joints_names(self): + """Return joint names in the same order as :meth:`get_joints`.""" + joints = [] + + def iterate_joints(joint): + joints.append(joint.value[1]) + for child in joint.filter("JOINT"): + iterate_joints(child) + + iterate_joints(next(self.root.filter("ROOT"))) + return joints + + def joint_direct_children(self, name): + """Return direct child joints of the given joint name.""" + joint = self.get_joint(name) + return [child for child in joint.filter("JOINT")] + + def get_joint_index(self, name): + """Return hierarchy index of the named joint.""" + return self.get_joints().index(self.get_joint(name)) + + def get_joint(self, name): + """Return hierarchy node for a joint name.""" + found = self.search("ROOT", name) + if not found: + found = self.search("JOINT", name) + if found: + return found[0] + raise LookupError("joint not found") + + def joint_offset(self, name, idx=[0, 1, 2]): + """Return selected `OFFSET` components for a joint.""" + joint = self.get_joint(name) + offset = joint["OFFSET"] + if len(offset) < max(idx): + return None + return (float(offset[idx[0]]), float(offset[idx[1]]), float(offset[idx[2]])) + + def joint_offset_rot(self, name): + """Return optional rotational offset components from custom BVH files.""" + return self.joint_offset(name, idx=[3, 4, 5]) + + def joint_channels(self, name): + """Return channel names declared for a joint.""" + if self.backend == "np": + return self.joint2channels[name] + else: + joint = self.get_joint(name) + return joint["CHANNELS"][1:] + + def get_joint_channels_index(self, joint_name): + """Return the flattened starting channel index for one joint.""" + if self.backend == "np": + return self.joint2idx[joint_name] + else: + index = 0 + for joint in self.get_joints(): + if joint.value[1] == joint_name: + return index + index += int(joint["CHANNELS"][0]) + raise LookupError("joint not found") + + def get_joint_channel_index(self, joint, channel): + """Return per-joint channel offset for a specific channel name.""" + channels = self.joint_channels(joint) + if channel in channels: + channel_index = channels.index(channel) + else: + raise ValueError(f"Channel {channel} not found in {channels}") + return channel_index + + def frame_joint_channel(self, frame_index, joint, channel, value=None): + """Return one channel value for one joint at one frame index.""" + joint_index = self.get_joint_channels_index(joint) + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + return value + if self.backend == "np": + return self.np_data_array[frame_index, joint_index + channel_index] + else: + return float(self.frames[frame_index][joint_index + channel_index]) + + def frame_joint_channels(self, frame_index, joint, channels, value=None): + """Get single frame data for on specific joint from multiple specific channels (e.g. + Xrotation, Yrotation, Zrotation).""" + values = [] + joint_index = self.get_joint_channels_index(joint) + if self.backend == "np": + channel_idx = [self.get_joint_channel_index(joint, channel) for channel in channels] + channel_idx = np.array(channel_idx) + joint_index + values = self.np_data_array[frame_index, channel_idx] + else: + for channel in channels: + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + values.append(value) + else: + values.append(float(self.frames[frame_index][joint_index + channel_index])) + return values + + def frames_joint_channels(self, joint, channels, value=None): + """Get all frame data for one joint from multiple channels (e.g. Xrotation, Yrotation, + Zrotation).""" + joint_index = self.get_joint_channels_index(joint) + if self.backend == "np": + channel_idx = [self.get_joint_channel_index(joint, channel) for channel in channels] + channel_idx = np.array(channel_idx) + joint_index + all_frames = self.np_data_array[:, channel_idx] + else: + all_frames = [] + for frame in self.frames: + values = [] + for channel in channels: + channel_index = self.get_joint_channel_index(joint, channel) + if channel_index == -1 and value is not None: + values.append(value) + else: + values.append(float(frame[joint_index + channel_index])) + all_frames.append(values) + return all_frames + + def frames_joints_channels(self, joint_names, channels): + """Get all frames for all specified joints with one specified set of channels.""" + if self.backend != "np": + raise NotImplementedError("Only np backend is supported for this function") + joint_indices = [(joint_name, self.joint2idx[joint_name]) for joint_name in joint_names] + data_indices = [] + for joint_name, joint_idx in joint_indices: + channel_indices = [self.get_joint_channel_index(joint_name, channel) for channel in channels] + data_indices.extend([joint_idx + channel_idx for channel_idx in channel_indices]) + all_frames = self.np_data_array[:, data_indices] + all_frames = all_frames.reshape(-1, len(joint_names), len(channels)) + return all_frames + + def joint_parent(self, name): + """Return parent joint node, or `None` for the root.""" + joint = self.get_joint(name) + if joint.parent == self.root: + return None + return joint.parent + + def joint_parent_index(self, name): + """Return parent joint index, or `-1` for the root.""" + joint = self.get_joint(name) + if joint.parent == self.root: + return -1 + return self.get_joints().index(joint.parent) + + @property + def nframes(self): + """Number of motion frames declared in the BVH header.""" + try: + return int(next(self.root.filter("Frames:")).value[1]) + except StopIteration: + raise LookupError("number of frames not found") + + @property + def frame_time(self): + """Frame duration in seconds declared in the BVH header.""" + try: + return float(next(self.root.filter("Frame")).value[2]) + except StopIteration: + raise LookupError("frame time not found") + + +class Bone: + """Container for one skeleton bone and its kinematic metadata.""" + + def __init__(self): + # original bone info + self.id = None + self.name = None + self.orient = np.identity(3) + self.dof_index = [] + self.channels = [] # bvh only + self.lb = [] + self.ub = [] + self.parent = None + self.child = [] + + # asf specific + self.dir = np.zeros(3) + self.len = 0 + # bvh specific + self.offset = np.zeros(3) # default offset for position + self.offset_rot = None # rotation for custom nv bvh + + # inferred info + self.pos = np.zeros(3) + self.end = np.zeros(3) + + def __repr__(self): + return f"{self.name}" + + +class SkeletonBvh: + """Skeleton structure reconstructed from BVH hierarchy metadata.""" + + def __init__(self): + self.bones = [] + self.name2bone = {} + self.mass_scale = 1.0 + self.len_scale = 1.0 + self.dof_name = ["x", "y", "z"] + self.root = None + + def get_bones_names(self): + """Return bone names in skeleton order.""" + return [x.name for x in self.bones] + + def get_parent_indices(self): + """Return parent index array aligned with `self.bones`.""" + parent_indices = [-1] * len(self.bones) + for bone in self.bones: + if bone.parent: + parent_indices[bone.id] = bone.parent.id + return parent_indices + + def get_neutral_joints(self): + """Return neutral/rest joint positions as a NumPy array `(J, 3)`.""" + joints = [] + for bone in self.bones: + joints.append(bone.pos) + joints = np.stack(joints, axis=0) + return joints + + def load_from_bvh(self, fname, exclude_bones=None, spec_channels=None, mocap=None): + """Load skeleton hierarchy and rest offsets from a BVH file. + + Args: + fname: Path to a BVH file (ignored when *mocap* is given). + exclude_bones: Bone-name substrings to ignore while constructing the + skeleton. + spec_channels: Optional per-joint channel overrides. + mocap: Pre-parsed :class:`Bvh` object. When provided the file is + not re-read from disk. + """ + if exclude_bones is None: + exclude_bones = {} + if spec_channels is None: + spec_channels = dict() + if mocap is None: + with open(fname) as f: + mocap = Bvh(f.read()) + + joint_names = list( + filter( + lambda x: all([t not in x for t in exclude_bones]), + mocap.get_joints_names(), + ) + ) + dof_ind = {"x": 0, "y": 1, "z": 2} + self.len_scale = 1.0 + self.root = Bone() + self.root.id = 0 + self.root.name = joint_names[0] + self.root.channels = mocap.joint_channels(self.root.name) + self.root.offset = np.array(mocap.joint_offset(self.root.name)) * self.len_scale + self.root.offset_rot = mocap.joint_offset_rot(self.root.name) + if self.root.offset_rot is not None: + self.root.offset_rot = np.array(self.root.offset_rot) + # self.root.offset = np.zeros_like(self.root.offset) # TODO: remove this + self.name2bone[self.root.name] = self.root + self.bones.append(self.root) + for i, joint in enumerate(joint_names[1:]): + bone = Bone() + bone.id = i + 1 + bone.name = joint + bone.channels = spec_channels[joint] if joint in spec_channels.keys() else mocap.joint_channels(joint) + bone.dof_index = [dof_ind[x[0].lower()] for x in bone.channels] + bone.offset = np.array(mocap.joint_offset(joint)) * self.len_scale + bone.offset_rot = mocap.joint_offset_rot(joint) + if bone.offset_rot is not None: + bone.offset_rot = np.array(bone.offset_rot) + bone.lb = [-180.0] * 3 + bone.ub = [180.0] * 3 + self.bones.append(bone) + self.name2bone[joint] = bone + + # for bone in self.bones: + # print(bone.name, bone.channels, bone.offset) + + for bone in self.bones[1:]: + parent_name = mocap.joint_parent(bone.name).name + if parent_name in self.name2bone.keys(): + bone_p = self.name2bone[parent_name] + bone_p.child.append(bone) + bone.parent = bone_p + + self.forward_bvh(self.root) + for bone in self.bones: + if len(bone.child) == 0: + child_vals = [str(node) for node in mocap.get_joint(bone.name).children] + if "End Site" in child_vals: + end_site_idx = child_vals.index("End Site") + end_site_offset = mocap.get_joint(bone.name).children[end_site_idx]["OFFSET"] + bone.end = bone.pos + np.array([float(x) for x in end_site_offset]) * self.len_scale + else: + pass + else: + bone.end = sum([bone_c.pos for bone_c in bone.child]) / len(bone.child) + + def forward_bvh(self, bone): + """Recursively accumulate absolute joint positions from local offsets.""" + if bone.parent: + bone.pos = bone.parent.pos + bone.offset + else: + bone.pos = bone.offset + for bone_c in bone.child: + self.forward_bvh(bone_c) + + +def load_bvh_animation( + fname: str, + skeleton: SkeletonBvh, + rot_order: Optional[str] = "native", + backend: Optional[str] = "np", + return_quat: Optional[bool] = False, + mocap: Optional["Bvh"] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Load motion channels from BVH into root translations and joint rotations. + + Args: + fname: Full path to the BVH file (ignored when *mocap* is given). + skeleton: Parsed neutral skeleton built from compatible BVH hierarchy. + rot_order: Euler order to use for conversion (`"native"` keeps BVH order). + backend: BVH parser backend (`"np"` or `"graph"`). + return_quat: If `True`, return quaternions instead of rotation matrices. + mocap: Pre-parsed :class:`Bvh` object. When provided the file is + not re-read from disk. + + Returns: + Root translations `(T, 3)` and joint rotations `(T, J, 3, 3)` or + `(T, J, 4)` when `return_quat=True`. + """ + if mocap is None: + with open(fname) as f: + mocap = Bvh(f.read(), backend=backend) + + # assume all joints are same ordering, load in with native ordering + root_channels = mocap.joint_channels(skeleton.root.name) + pos_channels = [channel for channel in root_channels if channel.endswith("position")] + rot_channels = [channel for channel in root_channels if channel.endswith("rotation")] + + root_trans = np.array(mocap.frames_joint_channels(skeleton.root.name, pos_channels)) + + effective_backend = mocap.backend + if effective_backend == "np": + # NOTE: assumes rot channel ordering is the same for all joints + joint_eulers = mocap.frames_joints_channels(skeleton.get_bones_names(), rot_channels) + joint_eulers = np.deg2rad(joint_eulers) + elif effective_backend == "graph": + joint_eulers = [] + for bone in skeleton.bones: + bone_channels = mocap.joint_channels(bone.name) + bone_rot_channels = [channel for channel in bone_channels if channel.endswith("rotation")] + assert bone_rot_channels == rot_channels, "Rotation channel ordering is not consistent across joints!" + # use native rotation order + euler = np.deg2rad(np.array(mocap.frames_joint_channels(bone.name, rot_channels))) + joint_eulers.append(euler) + joint_eulers = np.stack(joint_eulers, axis=1) + else: + raise ValueError(f"Unknown backend for BVH loading: {effective_backend}") + + if rot_order == "native": + rot_order = "" + for axis in rot_channels: + rot_order += axis[0] + else: + # need to reorder dims + ordered_joint_eulers = [] + for axis in rot_order: + i = rot_channels.index(axis + "rotation") + ordered_joint_eulers.append(joint_eulers[..., i]) + joint_eulers = np.stack(ordered_joint_eulers, axis=-1) + + rotations = Rotation.from_euler(rot_order, joint_eulers.reshape(-1, 3)) + if return_quat: + joint_rots = rotations.as_quat(scalar_first=True).reshape(joint_eulers.shape[:-1] + (4,)) + else: + joint_rots = rotations.as_matrix().reshape(joint_eulers.shape[:-1] + (3, 3)) + + return root_trans, joint_rots + + +def parse_bvh_motion(file_path_input: str, parse_neutral_joints: bool = False): + """Parse a BVH motion into tensors used by ARDY motion pipelines. + + Args: + file_path_input: Path to input BVH file. + parse_neutral_joints: If `True`, also return neutral joints in meters. + + Returns: + ``(local_rot_mats, root_trans, fps)`` or + ``(local_rot_mats, root_trans, fps, neutral_joints)`` when requested. + """ + with open(file_path_input) as f: + mocap = Bvh(f.read(), backend="np") + + fps = 1.0 / mocap.frame_time + + skeletonBVH = SkeletonBvh() + exclude_bones = {"Root"} + skeletonBVH.load_from_bvh(file_path_input, exclude_bones=exclude_bones, mocap=mocap) + + root_trans, local_rot_mats = load_bvh_animation(file_path_input, skeletonBVH, mocap=mocap) + root_trans *= 0.01 # unit change: cm -> m + root_trans = torch.tensor(root_trans) + local_rot_mats = torch.tensor(local_rot_mats) + + # Don't parse neutral_joints here + # it is not actually needed right now: + # the skeleton is always the same, and saved in the folder + # carefull: the one saved in the folder it relative to the standard t_pose + # whereas the parsed one is not + if not parse_neutral_joints: + return local_rot_mats, root_trans, fps + + neutral_joints = skeletonBVH.get_neutral_joints() + neutral_joints *= 0.01 # unit change: cm -> m + # remove the root position of the skeleton + # (it is already "included" in the root_translation) + root_idx = 0 + neutral_joints = torch.tensor(neutral_joints - neutral_joints[root_idx]) + return local_rot_mats, root_trans, fps, neutral_joints diff --git a/ardy/skeleton/definitions.py b/ardy/skeleton/definitions.py new file mode 100644 index 0000000000000000000000000000000000000000..2805604260844658658a72be9ef39a3556811d25 --- /dev/null +++ b/ardy/skeleton/definitions.py @@ -0,0 +1,381 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Concrete skeleton definitions: SOMA, G1, and Core with joint names and hierarchy.""" + +from pathlib import Path + +import numpy as np +import torch + +from ..tools import ensure_batched +from .base import SkeletonBase + + +class SOMASkeleton77(SkeletonBase): + """High-detail 77-joint SOMA skeleton with full finger and toe chains.""" + + name = "somaskel77" + + right_foot_joint_names = [ + "RightFoot", + "RightToeBase", + "RightToeEnd", + ] # in order of chain + left_foot_joint_names = [ + "LeftFoot", + "LeftToeBase", + "LeftToeEnd", + ] # in order of chain + right_hand_joint_names = [ + "RightHand", + "RightHandThumb1", + "RightHandThumb2", + "RightHandThumb3", + "RightHandThumbEnd", + "RightHandIndex1", + "RightHandIndex2", + "RightHandIndex3", + "RightHandIndex4", + "RightHandIndexEnd", + "RightHandMiddle1", + "RightHandMiddle2", + "RightHandMiddle3", + "RightHandMiddle4", + "RightHandMiddleEnd", + "RightHandRing1", + "RightHandRing2", + "RightHandRing3", + "RightHandRing4", + "RightHandRingEnd", + "RightHandPinky1", + "RightHandPinky2", + "RightHandPinky3", + "RightHandPinky4", + "RightHandPinkyEnd", + ] # in order of chain + left_hand_joint_names = [ + "LeftHand", + "LeftHandThumb1", + "LeftHandThumb2", + "LeftHandThumb3", + "LeftHandThumbEnd", + "LeftHandIndex1", + "LeftHandIndex2", + "LeftHandIndex3", + "LeftHandIndex4", + "LeftHandIndexEnd", + "LeftHandMiddle1", + "LeftHandMiddle2", + "LeftHandMiddle3", + "LeftHandMiddle4", + "LeftHandMiddleEnd", + "LeftHandRing1", + "LeftHandRing2", + "LeftHandRing3", + "LeftHandRing4", + "LeftHandRingEnd", + "LeftHandPinky1", + "LeftHandPinky2", + "LeftHandPinky3", + "LeftHandPinky4", + "LeftHandPinkyEnd", + ] # in order of chain + + hip_joint_names = ["RightLeg", "LeftLeg"] # in order [right, left] + + bone_order_names_with_parents = [ + ("Hips", None), + ("Spine1", "Hips"), + ("Spine2", "Spine1"), + ("Chest", "Spine2"), + ("Neck1", "Chest"), + ("Neck2", "Neck1"), + ("Head", "Neck2"), + ("HeadEnd", "Head"), + ("Jaw", "Head"), + ("LeftEye", "Head"), + ("RightEye", "Head"), + ("LeftShoulder", "Chest"), + ("LeftArm", "LeftShoulder"), + ("LeftForeArm", "LeftArm"), + ("LeftHand", "LeftForeArm"), + ("LeftHandThumb1", "LeftHand"), + ("LeftHandThumb2", "LeftHandThumb1"), + ("LeftHandThumb3", "LeftHandThumb2"), + ("LeftHandThumbEnd", "LeftHandThumb3"), + ("LeftHandIndex1", "LeftHand"), + ("LeftHandIndex2", "LeftHandIndex1"), + ("LeftHandIndex3", "LeftHandIndex2"), + ("LeftHandIndex4", "LeftHandIndex3"), + ("LeftHandIndexEnd", "LeftHandIndex4"), + ("LeftHandMiddle1", "LeftHand"), + ("LeftHandMiddle2", "LeftHandMiddle1"), + ("LeftHandMiddle3", "LeftHandMiddle2"), + ("LeftHandMiddle4", "LeftHandMiddle3"), + ("LeftHandMiddleEnd", "LeftHandMiddle4"), + ("LeftHandRing1", "LeftHand"), + ("LeftHandRing2", "LeftHandRing1"), + ("LeftHandRing3", "LeftHandRing2"), + ("LeftHandRing4", "LeftHandRing3"), + ("LeftHandRingEnd", "LeftHandRing4"), + ("LeftHandPinky1", "LeftHand"), + ("LeftHandPinky2", "LeftHandPinky1"), + ("LeftHandPinky3", "LeftHandPinky2"), + ("LeftHandPinky4", "LeftHandPinky3"), + ("LeftHandPinkyEnd", "LeftHandPinky4"), + ("RightShoulder", "Chest"), + ("RightArm", "RightShoulder"), + ("RightForeArm", "RightArm"), + ("RightHand", "RightForeArm"), + ("RightHandThumb1", "RightHand"), + ("RightHandThumb2", "RightHandThumb1"), + ("RightHandThumb3", "RightHandThumb2"), + ("RightHandThumbEnd", "RightHandThumb3"), + ("RightHandIndex1", "RightHand"), + ("RightHandIndex2", "RightHandIndex1"), + ("RightHandIndex3", "RightHandIndex2"), + ("RightHandIndex4", "RightHandIndex3"), + ("RightHandIndexEnd", "RightHandIndex4"), + ("RightHandMiddle1", "RightHand"), + ("RightHandMiddle2", "RightHandMiddle1"), + ("RightHandMiddle3", "RightHandMiddle2"), + ("RightHandMiddle4", "RightHandMiddle3"), + ("RightHandMiddleEnd", "RightHandMiddle4"), + ("RightHandRing1", "RightHand"), + ("RightHandRing2", "RightHandRing1"), + ("RightHandRing3", "RightHandRing2"), + ("RightHandRing4", "RightHandRing3"), + ("RightHandRingEnd", "RightHandRing4"), + ("RightHandPinky1", "RightHand"), + ("RightHandPinky2", "RightHandPinky1"), + ("RightHandPinky3", "RightHandPinky2"), + ("RightHandPinky4", "RightHandPinky3"), + ("RightHandPinkyEnd", "RightHandPinky4"), + ("LeftLeg", "Hips"), + ("LeftShin", "LeftLeg"), + ("LeftFoot", "LeftShin"), + ("LeftToeBase", "LeftFoot"), + ("LeftToeEnd", "LeftToeBase"), + ("RightLeg", "Hips"), + ("RightShin", "RightLeg"), + ("RightFoot", "RightShin"), + ("RightToeBase", "RightFoot"), + ("RightToeEnd", "RightToeBase"), + ] + + @property + def relaxed_hands_rest_pose(self): + # lazy loading + if hasattr(self, "_relaxed_hands_rest_pose"): + return self._relaxed_hands_rest_pose + + relaxed_hands_pose_path = Path(self.folder) / "relaxed_hands_rest_pose.npy" + relaxed_hands_rest_pose = torch.from_numpy(np.load(relaxed_hands_pose_path)).squeeze() + self.register_buffer( + "_relaxed_hands_rest_pose", + relaxed_hands_rest_pose, + persistent=False, + ) + return self._relaxed_hands_rest_pose + + +class SOMASkeleton30(SkeletonBase): + """Compact 30-joint SOMA variant with reduced hand and end-effector detail.""" + + name = "somaskel30" + + right_foot_joint_names = [ + "RightFoot", + "RightToeBase", + ] # in order of chain + left_foot_joint_names = [ + "LeftFoot", + "LeftToeBase", + ] # in order of chain + right_hand_joint_names = [ + "RightHand", + "RightHandMiddleEnd", + ] # in order of chain + left_hand_joint_names = [ + "LeftHand", + "LeftHandMiddleEnd", + ] # in order of chain + + hip_joint_names = ["RightLeg", "LeftLeg"] # in order [right, left] + + bone_order_names_with_parents = [ + ("Hips", None), + ("Spine1", "Hips"), + ("Spine2", "Spine1"), + ("Chest", "Spine2"), + ("Neck1", "Chest"), + ("Neck2", "Neck1"), + ("Head", "Neck2"), + ("Jaw", "Head"), + ("LeftEye", "Head"), + ("RightEye", "Head"), + ("LeftShoulder", "Chest"), + ("LeftArm", "LeftShoulder"), + ("LeftForeArm", "LeftArm"), + ("LeftHand", "LeftForeArm"), + ("LeftHandThumbEnd", "LeftHand"), + ("LeftHandMiddleEnd", "LeftHand"), + ("RightShoulder", "Chest"), + ("RightArm", "RightShoulder"), + ("RightForeArm", "RightArm"), + ("RightHand", "RightForeArm"), + ("RightHandThumbEnd", "RightHand"), + ("RightHandMiddleEnd", "RightHand"), + ("LeftLeg", "Hips"), + ("LeftShin", "LeftLeg"), + ("LeftFoot", "LeftShin"), + ("LeftToeBase", "LeftFoot"), + ("RightLeg", "Hips"), + ("RightShin", "RightLeg"), + ("RightFoot", "RightShin"), + ("RightToeBase", "RightFoot"), + ] + + @property + def somaskel77(self): + # lazy loading + if not hasattr(self, "_somaskel77"): + self._somaskel77 = SOMASkeleton77() + return self._somaskel77 + + @ensure_batched(local_joint_rots_subset=4) + def to_SOMASkeleton77(self, local_joint_rots_subset: torch.Tensor): + # Converting from 30-joint to 77-joint to have relaxed hands + + device = local_joint_rots_subset.device + nF = len(local_joint_rots_subset) + local_joint_rots_mats = self.somaskel77.relaxed_hands_rest_pose.clone().to(device).repeat(nF, 1, 1, 1) + + skel_slice = self.get_skel_slice(self.somaskel77) + local_joint_rots_mats[:, skel_slice] = local_joint_rots_subset + return local_joint_rots_mats + + @ensure_batched(local_joint_rots_full=4) # [BT, J, 3, 3] + def from_SOMASkeleton77(self, local_joint_rots_full: torch.Tensor) -> torch.Tensor: + """Extract the 30-joint subset from 77-joint local rotation data.""" + skel_slice = self.get_skel_slice(self.somaskel77) + return local_joint_rots_full[:, skel_slice] + + def output_to_SOMASkeleton77(self, output: dict) -> dict: + """Convert model output dict from somaskel30 to somaskel77. + + Expands local_rot_mats to 77 joints, re-runs FK for global_rot_mats and posed_joints. Foot + contacts are expanded from 4 channels to 6 (toe-end copies toe-base contact). + """ + local_rot_mats_77 = self.to_SOMASkeleton77(output["local_rot_mats"]) + root_positions = output["root_positions"] + global_rot_mats_77, posed_joints_77, _ = self.somaskel77.fk(local_rot_mats_77, root_positions) + out_77 = dict(output) + out_77["local_rot_mats"] = local_rot_mats_77 + out_77["global_rot_mats"] = global_rot_mats_77 + out_77["posed_joints"] = posed_joints_77 + + if "foot_contacts" in output: + fc = output["foot_contacts"] # [..., 4]: [L_heel, L_toe, R_heel, R_toe] + # -> [..., 6]: [L_heel, L_toe, L_toe_end, R_heel, R_toe, R_toe_end] + out_77["foot_contacts"] = torch.cat([fc[..., :2], fc[..., 1:2], fc[..., 2:4], fc[..., 3:4]], dim=-1) + + return out_77 + + +class G1Skeleton34(SkeletonBase): + """Unitree G1 skeleton with 32 articulated joints plus 2 toe endpoints.""" + + name = "g1skel34" + right_foot_joint_names = ["right_ankle_roll_skel", "right_toe_base"] + left_foot_joint_names = ["left_ankle_roll_skel", "left_toe_base"] + right_hand_joint_names = ["right_wrist_yaw_skel", "right_hand_roll_skel"] + left_hand_joint_names = ["left_wrist_yaw_skel", "left_hand_roll_skel"] + + hip_joint_names = [ + "right_hip_pitch_skel", + "left_hip_pitch_skel", + ] # used to calculate root orientation, only need 1 pair of hip joints + + bone_order_names_with_parents = [ + ("pelvis_skel", None), + ("left_hip_pitch_skel", "pelvis_skel"), + ("left_hip_roll_skel", "left_hip_pitch_skel"), + ("left_hip_yaw_skel", "left_hip_roll_skel"), + ("left_knee_skel", "left_hip_yaw_skel"), + ("left_ankle_pitch_skel", "left_knee_skel"), + ("left_ankle_roll_skel", "left_ankle_pitch_skel"), + ("left_toe_base", "left_ankle_roll_skel"), + ("right_hip_pitch_skel", "pelvis_skel"), + ("right_hip_roll_skel", "right_hip_pitch_skel"), + ("right_hip_yaw_skel", "right_hip_roll_skel"), + ("right_knee_skel", "right_hip_yaw_skel"), + ("right_ankle_pitch_skel", "right_knee_skel"), + ("right_ankle_roll_skel", "right_ankle_pitch_skel"), + ("right_toe_base", "right_ankle_roll_skel"), + ("waist_yaw_skel", "pelvis_skel"), + ("waist_roll_skel", "waist_yaw_skel"), + ("waist_pitch_skel", "waist_roll_skel"), + ("left_shoulder_pitch_skel", "waist_pitch_skel"), + ("left_shoulder_roll_skel", "left_shoulder_pitch_skel"), + ("left_shoulder_yaw_skel", "left_shoulder_roll_skel"), + ("left_elbow_skel", "left_shoulder_yaw_skel"), + ("left_wrist_roll_skel", "left_elbow_skel"), + ("left_wrist_pitch_skel", "left_wrist_roll_skel"), + ("left_wrist_yaw_skel", "left_wrist_pitch_skel"), + ("left_hand_roll_skel", "left_wrist_yaw_skel"), + ("right_shoulder_pitch_skel", "waist_pitch_skel"), + ("right_shoulder_roll_skel", "right_shoulder_pitch_skel"), + ("right_shoulder_yaw_skel", "right_shoulder_roll_skel"), + ("right_elbow_skel", "right_shoulder_yaw_skel"), + ("right_wrist_roll_skel", "right_elbow_skel"), + ("right_wrist_pitch_skel", "right_wrist_roll_skel"), + ("right_wrist_yaw_skel", "right_wrist_pitch_skel"), + ("right_hand_roll_skel", "right_wrist_yaw_skel"), + ] + + +class CoreSkeleton27(SkeletonBase): + """Core Skeleton with 27 joints.""" + + name = "cskel27" + right_foot_joint_names = ["RightFoot", "RightToeBase"] # in order of chain + left_foot_joint_names = ["LeftFoot", "LeftToeBase"] # in order of chain + right_hand_joint_names = ["RightHand", "RightHandEnd"] # in order of chain + left_hand_joint_names = ["LeftHand", "LeftHandEnd"] # in order of chain + hip_joint_names = ["RightUpLeg", "LeftUpLeg"] # in order [right, left] + + bone_order_names_with_parents = [ + ("Hips", None), + ## + ("Spine", "Hips"), + ("Spine1", "Spine"), + ("Spine2", "Spine1"), + ("Spine3", "Spine2"), + ("Neck", "Spine3"), + ("Head", "Neck"), + ## + ("RightShoulder", "Spine3"), + ("RightArm", "RightShoulder"), + ("RightForeArm", "RightArm"), + ("RightHand", "RightForeArm"), + ("RightHandEnd", "RightHand"), + ("RightHandThumb1", "RightHand"), + ## + ("LeftShoulder", "Spine3"), + ("LeftArm", "LeftShoulder"), + ("LeftForeArm", "LeftArm"), + ("LeftHand", "LeftForeArm"), + ("LeftHandEnd", "LeftHand"), + ("LeftHandThumb1", "LeftHand"), + ## + ("RightUpLeg", "Hips"), + ("RightLeg", "RightUpLeg"), + ("RightFoot", "RightLeg"), + ("RightToeBase", "RightFoot"), + ## + ("LeftUpLeg", "Hips"), + ("LeftLeg", "LeftUpLeg"), + ("LeftFoot", "LeftLeg"), + ("LeftToeBase", "LeftFoot"), + ] diff --git a/ardy/skeleton/kinematics.py b/ardy/skeleton/kinematics.py new file mode 100644 index 0000000000000000000000000000000000000000..e8c77cf4fd5b4108a438d093c1883a1d74fee5b5 --- /dev/null +++ b/ardy/skeleton/kinematics.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Forward-kinematics primitives for articulated skeletons.""" + +from typing import List + +import einops +import torch +import torch.nn.functional as F + +from ..tools import ensure_batched + + +@ensure_batched(local_joint_rots=4, root_positions=2) +def fk( + local_joint_rots: torch.Tensor, + root_positions: torch.Tensor, + skeleton, + root_positions_is_global: bool = True, +): + """Compute global joint rotations and positions from local rotations. + + Args: + local_joint_rots: Local rotation matrices with shape `(..., J, 3, 3)`. + root_positions: Root translations with shape `(..., 3)`. + skeleton: Skeleton object exposing `neutral_joints`, `joint_parents`, and + `root_idx`. + root_positions_is_global: If `True`, neutral joints are recentered so root + translations are interpreted in world space. + + Returns: + Tuple `(global_joint_rots, posed_joints, posed_joints_norootpos)`. + """ + device = local_joint_rots.device + dtype = local_joint_rots.dtype + + neutral_joints = skeleton.neutral_joints.to(device=device, dtype=dtype) + + if root_positions_is_global is True: + # Removing the pelvis offset from the neutral joints + # as the root positions does not depends on the pelvis offset of the skeleton + pelvis_offset = neutral_joints[skeleton.root_idx] + neutral_joints = neutral_joints - pelvis_offset + + # compute joint position and global rotations + joints = einops.repeat( + neutral_joints, + "j k -> b j k", + b=len(local_joint_rots), + ) + posed_joints_norootpos, global_joint_rots = batch_rigid_transform( + local_joint_rots, + joints, + skeleton.joint_parents, + skeleton.root_idx, + ) + # if root_positions_is_global is True: + # posed_joints_norootpos always start at zero + # otherwise it could start with the pelvis offset + + posed_joints = posed_joints_norootpos + root_positions[:, None] + return global_joint_rots, posed_joints, posed_joints_norootpos + + +def compute_idx_levels(parents): + """Group joint indices by hierarchy depth for level-wise FK updates. + + Args: + parents: Parent index tensor of shape `(J,)` with root parent `-1`. + + Returns: + List of index tensors, where each tensor contains joints at one depth. + """ + idx_levs = [[]] + lev_dicts = {0: -1} + for i in range(1, parents.shape[0]): + assert int(parents[i]) in lev_dicts + lev = lev_dicts[int(parents[i])] + 1 + if lev + 1 > len(idx_levs): + idx_levs.append([]) + idx_levs[lev].append(int(i)) + lev_dicts[int(i)] = lev + idx_levs = [torch.tensor(x).long() for x in idx_levs] + return idx_levs + + +def batch_rigid_transform(rot_mats, joints, parents, root_idx): + """Perform batch rigid transformation on a skeletal structure. + + Args: + rot_mats: Local rotation matrices for each joint: (B, J, 3, 3) + joints: Initial joint positions: (B, J, 3) + parents: Tensor indicating the parent of each joint: (J,) + root_idx (int): index of the root + + Returns: + Transformed joint positions after applying forward kinematics. + """ + # Compute the hierarchical levels of joints based on their parent relationships + idx_levs = compute_idx_levels(parents) + + # Apply forward kinematics to transform the joints + return forward_kinematics(rot_mats, joints, parents, idx_levs, root_idx) + + +@torch.jit.script +def transform_mat(R, t): + """Creates a batch of transformation matrices. + + Args: + - R: Bx3x3 array of a batch of rotation matrices + - t: Bx3x1 array of a batch of translation vectors + Returns: + - T: Bx4x4 Transformation matrix + """ + # No padding left or right, only add an extra row + return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1.0)], dim=2) + + +@torch.jit.script +def forward_kinematics( + rot_mats, + joints, + parents: torch.Tensor, + idx_levs: List[torch.Tensor], + root_idx: int, +): + """Perform forward kinematics to compute posed joints and global rotation matrices. + + Args: + rot_mats: Local rotation matrices for each joint: (B, J, 3, 3) + joints: Initial joint positions: (B, J, 3) + parents: Tensor indicating the parent of each joint: (J,) + idx_levs: Tensors of joint indices grouped by depth in the kinematic tree. + root_idx (int): index of the root + Returns: + Posed joints: (B, J, 3) + Global rotation matrices: (B, J, 3, 3) + """ + # Add an extra dimension to joints + joints = torch.unsqueeze(joints, dim=-1) + + # Compute relative joint positions + rel_joints = joints.clone() + + mask_no_root = torch.ones(joints.shape[1], dtype=torch.bool) + mask_no_root[root_idx] = False + rel_joints[:, mask_no_root] -= joints[:, parents[mask_no_root]].clone() + + # Compute initial transformation matrices + # (B, J + 1, 4, 4) + transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3), rel_joints.reshape(-1, 3, 1)).reshape( + -1, joints.shape[1], 4, 4 + ) + + # Initialize the root transformation matrices + transforms = torch.zeros_like(transforms_mat) + transforms[:, root_idx] = transforms_mat[:, root_idx] + + # Compute global transformations level by level + for indices in idx_levs: + curr_res = torch.matmul(transforms[:, parents[indices]], transforms_mat[:, indices]) + transforms[:, indices] = curr_res + + # Extract posed joint positions from the transformation matrices + posed_joints = transforms[:, :, :3, 3] + + # Extract global rotation matrices from the transformation matrices + global_rot_mat = transforms[:, :, :3, :3] + + return posed_joints, global_rot_mat diff --git a/ardy/skeleton/registry.py b/ardy/skeleton/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..12f1e2314806beeaa5661fd2c56f3f771442c17b --- /dev/null +++ b/ardy/skeleton/registry.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Factory helpers for building predefined skeleton variants.""" + +from pathlib import Path + +from ardy.assets import SKELETONS_ROOT + +from .definitions import ( + CoreSkeleton27, + G1Skeleton34, + SOMASkeleton30, + SOMASkeleton77, +) + + +def build_skeleton(nbjoints: int, assets_folder: str | Path = SKELETONS_ROOT): + """Instantiate a known skeleton class from its joint count. + + Supported joint counts: 30 (SOMA compact), 34 (G1), 77 (SOMA full). + + Args: + nbjoints: Number of joints expected in the skeleton representation. + assets_folder: Base skeleton-assets directory containing per-skeleton subfolders. + + Returns: + A configured `SkeletonBase` subclass instance. + + Raises: + ValueError: If `nbjoints` does not match a registered skeleton. + """ + assets_folder = Path(assets_folder) + if nbjoints == 34: + return G1Skeleton34(assets_folder / "g1skel34") + elif nbjoints == 30: + return SOMASkeleton30(assets_folder / "somaskel30") + elif nbjoints == 77: + return SOMASkeleton77(assets_folder / "somaskel77") + elif nbjoints == 27: + return CoreSkeleton27(assets_folder / "cskel27") + else: + raise ValueError("This skeleton is not recognized.") diff --git a/ardy/skeleton/transforms.py b/ardy/skeleton/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..5f45b0df206d76f025469ba9220f1acab0226829 --- /dev/null +++ b/ardy/skeleton/transforms.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rotation-space conversion utilities for skeleton motion data.""" + +import einops +import torch + +from ..tools import ensure_batched +from .kinematics import batch_rigid_transform + + +def global_rots_to_local_rots(global_joint_rots: torch.Tensor, skeleton): + """Convert global rotations to local rotations using a skeleton hierarchy. + + Args: + global_joint_rots: Global rotation matrices with shape `(..., J, 3, 3)`. + skeleton: Skeleton object exposing `joint_parents` and `root_idx`. + + Returns: + Local rotation matrices with the same leading shape as the input. + """ + # Doing big batch + global_joint_mats, ps = einops.pack( + [global_joint_rots], + "* nbjoints dim1 dim2", + ) + + # obtain back the local rotations from the new global rotations. + # skeleton.joint_parents is a registered buffer that follows the skeleton's + # device (e.g. CUDA for a loaded model), while the input may be on CPU (e.g. + # motion parsed on CPU for caching); align the index/eye to the data device. + device = global_joint_mats.device + parent_rot_mats = global_joint_mats[:, skeleton.joint_parents.to(device)] + + parent_rot_mats[:, skeleton.root_idx] = torch.eye(3, device=device, dtype=parent_rot_mats.dtype) # the root joint + parent_rot_mats_inv = parent_rot_mats.transpose(2, 3) + local_rot_mats = torch.einsum( + "T N m n, T N n o -> T N m o", + parent_rot_mats_inv, + global_joint_mats, + ) + [local_rot_mats] = einops.unpack(local_rot_mats, ps, "* nbjoints dim1 dim2") + return local_rot_mats + + +@ensure_batched(local_rot_mats=4) +def change_tpose(local_rot_mats: torch.Tensor, global_rot_offsets: torch.Tensor, skeleton): + """Re-express local rotations in another t_pose based on the global rotation offsets. + + Args: + local_rot_mats: Local rotation matrices with shape `(..., J, 3, 3)`. + global_rot_offsets: Global rotation offsets with shape `(..., J, 3, 3)`. + skeleton: Skeleton object exposing `joint_parents`, + `root_idx`, and `nbjoints`. + + Returns: + Tuple `(new_local_rot_mats, new_global_rot_mats)` in the standard frame. + """ + device, dtype = local_rot_mats.device, local_rot_mats.dtype + global_rot_offsets = global_rot_offsets.to(device=device, dtype=dtype) + + root_idx = skeleton.root_idx + joint_parents = skeleton.joint_parents + # These are dummy joint positions, will not be used + neutral_joints = torch.ones((len(local_rot_mats), skeleton.nbjoints, 3), device=device, dtype=dtype) + + # get the old joint rotations in the same global space as the t-pose + # Note: the neutral joints we use here doesn't matter, because we are only using the global rotation outputs + _, global_rot_mats = batch_rigid_transform(local_rot_mats, neutral_joints, joint_parents, root_idx) # (T, N, 3, 3) + + # compute the desired joint rotations in the frame of the new t-pose + new_global_rot_mats = torch.einsum("T N m n, N o n -> T N m o", global_rot_mats, global_rot_offsets) + # convert back to local rotations + new_local_rot_mats = global_rots_to_local_rots(new_global_rot_mats, skeleton) + return new_local_rot_mats, new_global_rot_mats + + +@ensure_batched(local_rot_mats=4) +def to_standard_tpose(local_rot_mats: torch.Tensor, skeleton): + """Re-express local rotations in the skeleton's standard T-pose convention. + + Args: + local_rot_mats: Local rotation matrices with shape `(..., J, 3, 3)`. + skeleton: Skeleton object exposing `global_rot_offsets`, `joint_parents`, + `root_idx`, and `nbjoints`. + + Returns: + Tuple `(new_local_rot_mats, new_global_rot_mats)` in the standard frame. + """ + global_rot_offsets = skeleton.global_rot_offsets + return change_tpose(local_rot_mats, global_rot_offsets, skeleton) + + +@ensure_batched(local_rot_mats=4) +def from_standard_tpose(local_rot_mats: torch.Tensor, skeleton): + """Re-express local rotations from the skeleton's standard T-pose convention to the original + formulation. + + Args: + local_rot_mats: Local rotation matrices with shape `(..., J, 3, 3)`. + skeleton: Skeleton object exposing `global_rot_offsets`, `joint_parents`, + `root_idx`, and `nbjoints`. + + Returns: + Tuple `(new_local_rot_mats, new_global_rot_mats)` in the standard frame. + """ + global_rot_offsets = skeleton.global_rot_offsets + global_rot_offsets_T = global_rot_offsets.mT # do the inverse transform + return change_tpose(local_rot_mats, global_rot_offsets_T, skeleton) diff --git a/ardy/tools.py b/ardy/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..25c4b91dc73a143a3714f830c9d735826c6f9bcb --- /dev/null +++ b/ardy/tools.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import json +import math +import random +from collections.abc import Mapping, Sequence +from functools import wraps +from math import prod +from pathlib import Path +from typing import Any, Callable, Mapping, Optional, ParamSpec, TypeVar, Union + +import numpy as np +import torch + + +def validate(validator, save_args: bool = False, super_init: bool = False): + """Create a decorator function for validating user inputs. + + Args: + validator: the function to validate (pydantic dataclass) + save (bool): save all the attributes to the obj [args[0]] + super_init (bool): init parent with no arguments (usefull for using save on a nn.Module) + + Returns: + decorator: the decorator function + """ + + def decorator(func): + @wraps(func) + def validated_func(*args, **kwargs): + conf = validator(**kwargs) + + if save_args: + assert len(args) != 0 + obj = args[0] + + if super_init: + # init the parent module + super(type(obj), obj).__init__() + + for key, val in conf.__dict__.items(): + setattr(obj, key, val) + return func(*args, conf) + + return validated_func + + return decorator + + +# Type alias for clarity +Tensor = Any + +P = ParamSpec("P") +R = TypeVar("R") + + +def ensure_batched(**spec: int) -> Callable[[Callable[P, R]], Callable[P, R]]: + """Decorator to flatten complex batch dimensions. + + Fixes included: + 1. Handles 1D tensors (tail_ndim=0) correctly without slicing errors. + 2. Skips .reshape() if the input is already purely flat (Optimization). + """ + if not spec: + raise ValueError("At least one argument spec must be provided.") + + canonical_name, canonical_ndim = next(iter(spec.items())) + + def decorator(fn: Callable[P, R]) -> Callable[P, R]: + sig = inspect.signature(fn) + + @wraps(fn) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + + # --- 1. CANONICAL ARGUMENT --- + if canonical_name not in bound.arguments: + raise TypeError(f"Missing canonical argument '{canonical_name}'.") + + x0 = bound.arguments[canonical_name] + if x0 is None: + raise ValueError(f"Canonical '{canonical_name}' cannot be None.") + + # Calculate split between Batch dims and Feature dims + expected_tail_dims = canonical_ndim - 1 # e.g. 3 - 1 = 2 (Sequence, Feat) + + # Validation + if x0.ndim < expected_tail_dims: + raise ValueError(f"'{canonical_name}' ndim={x0.ndim} < expected {expected_tail_dims} tail dims.") + + # --- LOGIC FIX 1: Handle 0 tail dims correctly --- + if expected_tail_dims == 0: + orig_batch_shape = x0.shape + tail_shape = () + else: + orig_batch_shape = x0.shape[:-expected_tail_dims] + tail_shape = x0.shape[-expected_tail_dims:] + + # Calculate flattened batch size + # If orig_batch_shape is () (scalar input), size is 1. + B_flat = prod(orig_batch_shape) if orig_batch_shape else 1 + + # Determine if we added a fake batch dim (unbatched input) + is_unbatched_input = len(orig_batch_shape) == 0 + + # --- LOGIC FIX 2: Skip reshape if already flat (Optimization) --- + # If batch shape is already 1D (e.g. [2]), we don't need to reshape [2, 140, 5] -> [2, 140, 5] + is_already_flat = len(orig_batch_shape) == 1 + + if is_unbatched_input: + # (H, W) -> (1, H, W) + x0_batched = x0.reshape(1, *tail_shape) + elif is_already_flat: + # (B, H, W) -> Keep as is + x0_batched = x0 + else: + # (B1, B2, H, W) -> (B1*B2, H, W) + x0_batched = x0.reshape(B_flat, *tail_shape) + + bound.arguments[canonical_name] = x0_batched + + # --- 2. OTHER ARGUMENTS --- + for name, target_ndim in list(spec.items())[1:]: + val = bound.arguments.get(name, None) + if val is None: + continue + + arg_tail_dims = target_ndim - 1 # e.g. for lengths=1, tail=0 + + # Validate + if val.ndim < arg_tail_dims: + raise ValueError(f"'{name}' ndim={val.ndim} too small.") + + # --- Get Batch Shape (With 0-tail fix) --- + if arg_tail_dims == 0: + val_batch_shape = val.shape + val_tail_shape = () + else: + val_batch_shape = val.shape[:-arg_tail_dims] + val_tail_shape = val.shape[-arg_tail_dims:] + + # --- Check Mismatch --- + # Unbatched inputs must match unbatched canonical + if len(val_batch_shape) == 0: + if not is_unbatched_input: + raise ValueError(f"'{name}' is unbatched but canonical is batched.") + val_batched = val.reshape(1, *val_tail_shape) + else: + # Batched inputs must match canonical batch shape EXACTLY + if val_batch_shape != orig_batch_shape: + raise ValueError( + f"Batch dimensions mismatch! '{canonical_name}' has {orig_batch_shape}, " + f"but '{name}' has {val_batch_shape}." + ) + + # Optimization: Don't reshape if already flat + if is_already_flat: + val_batched = val + else: + val_batched = val.reshape(B_flat, *val_tail_shape) + + bound.arguments[name] = val_batched + + # --- 3. EXECUTION --- + out = fn(**bound.arguments) + + # --- 4. RESTORE --- + def restore(obj): + if isinstance(obj, Mapping): + return {k: restore(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(restore(x) for x in obj) + + if hasattr(obj, "shape"): + if obj.ndim == 0: + return obj + + # Verify batch dimension exists and wasn't reduced + if obj.shape[0] != B_flat: + return obj + + # If input was simple (B, ...), return simple (B, ...) + if is_already_flat: + return obj + + rest = obj.shape[1:] + + if is_unbatched_input: + return obj.reshape(*rest) + + return obj.reshape(*orig_batch_shape, *rest) + return obj + + return restore(out) + + return wrapper + + return decorator + + +def to_numpy(obj): + if isinstance(obj, Mapping): + return {k: to_numpy(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(to_numpy(x) for x in obj) + if isinstance(obj, torch.Tensor): + return obj.cpu().numpy() + return obj + + +def to_torch(obj, device=None, dtype=None): + if isinstance(obj, Mapping): + return {k: to_torch(v, device, dtype) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(to_torch(x, device, dtype) for x in obj) + if isinstance(obj, np.ndarray): + obj = torch.from_numpy(obj) + if isinstance(obj, torch.Tensor): + if dtype is not None: + obj = obj.to(dtype=dtype) + if device is None: + return obj + return obj.to(device) + return obj + + +def seed_everything(seed: int, deterministic: bool = False) -> None: + """Seed all random number generators.""" + random.seed(seed) # for Python random module. + np.random.seed(seed) # for NumPy. + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + if deterministic: + torch.backends.cudnn.deterministic = True # for deterministic behavior. + torch.backends.cudnn.benchmark = False # if you want to make the behavior deterministic. + + +def load_json(path: Union[str, Path]) -> Any: + """Load a JSON file and return its contents. + + Args: + path (str | Path): Path to the JSON file. + + Returns: + Any: Parsed JSON content (dict, list, etc.). + + Raises: + FileNotFoundError: If the file does not exist. + ValueError: If the file is not valid JSON. + """ + path = Path(path) + + if not path.exists(): + raise FileNotFoundError(f"JSON file not found: {path}") + + try: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in file {path}: {e}") from e + + +def save_json(path: Union[str, Path], data: Any) -> None: + """Save data to a JSON file. + + Args: + path (str | Path): Path to the JSON file. + data (Any): Data to save (must be JSON serializable). + + Raises: + ValueError: If the data is not JSON serializable. + """ + path = Path(path) + + # Create parent directories if they don't exist + path.parent.mkdir(parents=True, exist_ok=True) + + try: + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + except (TypeError, ValueError) as e: + raise ValueError(f"Data is not JSON serializable: {e}") from e diff --git a/ardy/viz/__init__.py b/ardy/viz/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2431d663a14af04b57cde832144facd59b3935a7 --- /dev/null +++ b/ardy/viz/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Visualization utilities for viser-based rendering of skeletons and skinned meshes.""" diff --git a/ardy/viz/core_skin.py b/ardy/viz/core_skin.py new file mode 100644 index 0000000000000000000000000000000000000000..28878e82f3c5211b51981659bc8eede597b80b10 --- /dev/null +++ b/ardy/viz/core_skin.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import numpy as np +import torch + +from ardy.skeleton.kinematics import batch_rigid_transform + +SKIN_NAME = "skin_standard.npz" + + +class CoreSkin: + def __init__(self, skeleton): + self.skeleton = skeleton + skin_data_path = Path(skeleton.folder) / SKIN_NAME + + assert skeleton.neutral_joints is not None, "CoreSkeleton27 must have neutral joints instantiated" + + device = skeleton.neutral_joints.device + + # bind_rig_transform: [R, 4, 4] + # bind_vertices: [V, 3] + # faces: [F, 3] + # lbs indices, lbs weights: [V, W] (W = max (num joints vertice is related to), in our case W=5) + skin_data = np.load(skin_data_path) + bind_rig_np = np.array(skin_data["bind_rig_transform"], dtype=np.float32) + self.bind_rig_transform = torch.from_numpy(bind_rig_np).to(device=device, dtype=torch.float) + # Precompute the inverse in numpy to avoid torch lazy evaluation issues + bind_rig_inv_np = np.linalg.inv(bind_rig_np) + self.bind_rig_transform_inv = torch.from_numpy(bind_rig_inv_np).to(device=device, dtype=torch.float) + self.bind_vertices = torch.tensor(skin_data["bind_vertices"], device=device, dtype=torch.float) + self.faces = torch.tensor(skin_data["faces"], device=device, dtype=torch.long) + self.lbs_indices = torch.tensor(skin_data["lbs_indices"], device=device, dtype=torch.long) + self.lbs_weights = torch.tensor(skin_data["lbs_weights"], device=device, dtype=torch.float) + + # double check the rig matches expected skeleton + rig_joint_names = list(skin_data["rig_joint_names"]) # list(str) : [R] + for sname, rname in zip(self.skeleton.bone_order_names, rig_joint_names): + if sname != rname: + raise ValueError(f"MISMATCH in skinnging rig: expected='{sname}' vs rig='{rname}'") + + def lbs(self, posed_transform): + bind_rig_transform_inv = self.bind_rig_transform_inv + bind_vertices = self.bind_vertices + lbs_weights = self.lbs_weights + # posed_transform: [B, F, J, 4, 4] or [B, J, 4, 4] or [J, 4, 4] + # unsqueeze to match posed_transform dim + for _ in range(posed_transform.dim() - 3): + bind_rig_transform_inv = bind_rig_transform_inv.unsqueeze(0) + bind_vertices = bind_vertices.unsqueeze(0) + lbs_weights = lbs_weights.unsqueeze(0) + # bind_rig_transform_inv: [..., R, 4, 4] + # bind_vertices: [..., V, 3] + # lbs_weights: [..., V, W] + + affine_mat = (posed_transform @ bind_rig_transform_inv)[..., :3, :] # [..., J, 3, 4] + vs = ( + affine_mat[..., self.lbs_indices, :, :] + @ torch.concat([bind_vertices, torch.ones_like(bind_vertices[..., 0:1])], dim=-1)[..., None, :, None] + ) # [..., V, W, 3, 1] + ws = lbs_weights[..., None, None] + resv = (vs * ws).sum(dim=-3).squeeze(-1) # [..., V, 3] + return resv + + def skin(self, joint_rotmat, joint_pos, rot_is_global=False): + """ + joint_rotmat: [T, J, 3, 3] local or global joint rotation matrices + joint_pos: [T, J, 3] global joint positions + rot_is_global: bool, if True, joint_rotmat is global rotation matrices, otherwise it is local rotation matrices and FK is performed internally + """ + nF, nJ = joint_pos.shape[:2] + device = joint_rotmat.device + + # prepare full transformation matrices + fk_transform = torch.eye(4, device=device)[None, None].repeat(nF, nJ, 1, 1) + fk_transform[..., :3, 3] = joint_pos + if rot_is_global: + fk_transform[..., :3, :3] = joint_rotmat + else: + neutral_joints_seq = self.skeleton.neutral_joints[None].repeat((nF, 1, 1)).to(device) + # FK to get the global rotations + _, global_joint_rotmat = batch_rigid_transform( + joint_rotmat, + neutral_joints_seq, + self.skeleton.joint_parents.to(device), + self.skeleton.root_idx, + ) + fk_transform[..., :3, :3] = global_joint_rotmat + + vertices = self.lbs(fk_transform) + return vertices diff --git a/ardy/viz/g1_rig.py b/ardy/viz/g1_rig.py new file mode 100644 index 0000000000000000000000000000000000000000..d8917dd3661051b687598e712aefd4195526ef22 --- /dev/null +++ b/ardy/viz/g1_rig.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import xml.etree.ElementTree as ET +from typing import Optional, Tuple + +import numpy as np +import trimesh +import viser +import viser.transforms as tf + +from ardy.skeleton import G1Skeleton34 + +G1_MESH_JOINT_MAP = { + "pelvis_skel": ["pelvis.STL", "pelvis_contour_link.STL"], + "left_hip_pitch_skel": ["left_hip_pitch_link.STL"], + "left_hip_roll_skel": ["left_hip_roll_link.STL"], + "left_hip_yaw_skel": ["left_hip_yaw_link.STL"], + "left_knee_skel": ["left_knee_link.STL"], + "left_ankle_pitch_skel": ["left_ankle_pitch_link.STL"], + "left_ankle_roll_skel": ["left_ankle_roll_link.STL"], + "right_hip_pitch_skel": ["right_hip_pitch_link.STL"], + "right_hip_roll_skel": ["right_hip_roll_link.STL"], + "right_hip_yaw_skel": ["right_hip_yaw_link.STL"], + "right_knee_skel": ["right_knee_link.STL"], + "right_ankle_pitch_skel": ["right_ankle_pitch_link.STL"], + "right_ankle_roll_skel": ["right_ankle_roll_link.STL"], + "waist_yaw_skel": ["waist_yaw_link_rev_1_0.STL", "waist_yaw_link.STL"], + "waist_roll_skel": ["waist_roll_link_rev_1_0.STL", "waist_roll_link.STL"], + "waist_pitch_skel": [ + "torso_link_rev_1_0.STL", + "torso_link.STL", + "logo_link.STL", + "head_link.STL", + ], + "left_shoulder_pitch_skel": ["left_shoulder_pitch_link.STL"], + "left_shoulder_roll_skel": ["left_shoulder_roll_link.STL"], + "left_shoulder_yaw_skel": ["left_shoulder_yaw_link.STL"], + "left_elbow_skel": ["left_elbow_link.STL"], + "left_wrist_roll_skel": ["left_wrist_roll_link.STL"], + "left_wrist_pitch_skel": ["left_wrist_pitch_link.STL"], + "left_wrist_yaw_skel": ["left_wrist_yaw_link.STL", "left_rubber_hand.STL"], + "right_shoulder_pitch_skel": ["right_shoulder_pitch_link.STL"], + "right_shoulder_roll_skel": ["right_shoulder_roll_link.STL"], + "right_shoulder_yaw_skel": ["right_shoulder_yaw_link.STL"], + "right_elbow_skel": ["right_elbow_link.STL"], + "right_wrist_roll_skel": ["right_wrist_roll_link.STL"], + "right_wrist_pitch_skel": ["right_wrist_pitch_link.STL"], + "right_wrist_yaw_skel": ["right_wrist_yaw_link.STL", "right_rubber_hand.STL"], +} + + +class G1MeshRig: + _mesh_geom_cache: dict[str, dict[str, tuple[np.ndarray, np.ndarray]]] = {} + _mesh_transform_cache: dict[str, dict[str, tuple[np.ndarray, np.ndarray]]] = {} + + def __init__( + self, + name: str, + server: viser.ViserServer | viser.ClientHandle, + skeleton: G1Skeleton34, + mesh_dir: str, + color: Tuple[int, int, int], + ): + self.server = server + self.skeleton = skeleton + self.mesh_dir = mesh_dir + self.color = color + self.mesh_handles: list[viser.SceneHandle] = [] + self.mesh_items: list[dict[str, object]] = [] + self._defer_initial_visibility = True + self._mujoco_to_ardy = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64) + self._mesh_local_transforms = self._get_mesh_local_transforms() + + for joint_name, mesh_files in G1_MESH_JOINT_MAP.items(): + if joint_name not in self.skeleton.bone_index: + continue + joint_idx = self.skeleton.bone_index[joint_name] + for mesh_file in mesh_files: + mesh_path = os.path.join(self.mesh_dir, mesh_file) + if not os.path.exists(mesh_path): + continue + vertices, faces = self._get_mesh_geometry(mesh_file, mesh_path) + if vertices is None: + continue + handle = self.server.scene.add_mesh_simple( + f"/{name}/g1_mesh/{os.path.splitext(mesh_file)[0]}", + vertices=vertices, + faces=faces, + opacity=None, + color=self.color, + wireframe=False, + visible=not self._defer_initial_visibility, + ) + self.mesh_handles.append(handle) + geom_pos, geom_rot = self._mesh_local_transforms.get( + mesh_file, + (np.zeros(3, dtype=np.float64), np.eye(3, dtype=np.float64)), + ) + self.mesh_items.append( + { + "handle": handle, + "joint_idx": joint_idx, + "geom_pos": geom_pos, + "geom_rot": geom_rot, + } + ) + + if self._defer_initial_visibility: + for handle in self.mesh_handles: + handle.visible = True + + def _get_mesh_geometry(self, mesh_file: str, mesh_path: str) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]: + cache_key = self.mesh_dir + cached = self._mesh_geom_cache.get(cache_key) + if cached is not None and mesh_file in cached: + vertices, faces = cached[mesh_file] + return vertices.copy(), faces.copy() + + mesh = trimesh.load_mesh(mesh_path, process=True) + if isinstance(mesh, trimesh.Scene): + mesh = trimesh.util.concatenate(mesh.dump()) + vertices = mesh.vertices @ self._mujoco_to_ardy.T + faces = mesh.faces + + if cache_key not in self._mesh_geom_cache: + self._mesh_geom_cache[cache_key] = {} + self._mesh_geom_cache[cache_key][mesh_file] = (vertices, faces) + return vertices.copy(), faces.copy() + + def _get_mesh_local_transforms(self) -> dict[str, tuple[np.ndarray, np.ndarray]]: + cached = self._mesh_transform_cache.get(self.mesh_dir) + if cached is not None: + return {mesh_file: (pos.copy(), rot.copy()) for mesh_file, (pos, rot) in cached.items()} + + xml_path = os.path.abspath(os.path.join(self.mesh_dir, "..", "..", "xml", "g1.xml")) + if not os.path.exists(xml_path): + return {} + tree = ET.parse(xml_path) + root = tree.getroot() + + mesh_file_to_mesh_name = {} + for mesh in root.findall(".//asset/mesh"): + mesh_name = mesh.get("name") + mesh_file = mesh.get("file") + if mesh_name and mesh_file: + mesh_file_to_mesh_name[mesh_file] = mesh_name + + mesh_name_to_transform = {} + for geom in root.findall(".//geom"): + mesh_name = geom.get("mesh") + if mesh_name is None: + continue + pos = geom.get("pos") + quat = geom.get("quat") + if pos is None: + geom_pos = np.zeros(3, dtype=np.float64) + else: + geom_pos = np.array([float(x) for x in pos.split()], dtype=np.float64) + if quat is None: + geom_rot = np.eye(3, dtype=np.float64) + else: + wxyz = np.array([float(x) for x in quat.split()], dtype=np.float64) + geom_rot = tf.SO3(wxyz=wxyz).as_matrix() + mesh_name_to_transform[mesh_name] = (geom_pos, geom_rot) + + mesh_file_transforms = {} + for mesh_file, mesh_name in mesh_file_to_mesh_name.items(): + geom_pos, geom_rot = mesh_name_to_transform.get( + mesh_name, + (np.zeros(3, dtype=np.float64), np.eye(3, dtype=np.float64)), + ) + geom_pos = self._mujoco_to_ardy @ geom_pos + geom_rot = self._mujoco_to_ardy @ geom_rot @ self._mujoco_to_ardy.T + mesh_file_transforms[mesh_file] = (geom_pos, geom_rot) + + self._mesh_transform_cache[self.mesh_dir] = { + mesh_file: (pos.copy(), rot.copy()) for mesh_file, (pos, rot) in mesh_file_transforms.items() + } + return mesh_file_transforms + + def set_visibility(self, visible: bool) -> None: + for handle in self.mesh_handles: + handle.visible = visible + + def set_opacity(self, opacity: float) -> None: + for handle in self.mesh_handles: + handle.opacity = opacity + + def set_wireframe(self, wireframe: bool) -> None: + for handle in self.mesh_handles: + handle.wireframe = wireframe + + def set_color(self, color: Tuple[int, int, int]) -> None: + self.color = color + for handle in self.mesh_handles: + handle.color = color + + def set_pose(self, joints_pos: np.ndarray, joints_rot: np.ndarray) -> None: + for item in self.mesh_items: + handle = item["handle"] + joint_idx = item["joint_idx"] + geom_pos = item["geom_pos"] + geom_rot = item["geom_rot"] + + joint_pos = joints_pos[joint_idx] + joint_rot = joints_rot[joint_idx] + mesh_pos = joint_pos + joint_rot @ geom_pos + mesh_rot = joint_rot @ geom_rot + + handle.position = mesh_pos + handle.wxyz = tf.SO3.from_matrix(mesh_rot).wxyz + + def clear(self) -> None: + for handle in self.mesh_handles: + self.server.scene.remove_by_name(handle.name) diff --git a/ardy/viz/soma_skin.py b/ardy/viz/soma_skin.py new file mode 100644 index 0000000000000000000000000000000000000000..0fcb43617bf95fbed41a7ebc3cd489043a5f7c8f --- /dev/null +++ b/ardy/viz/soma_skin.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Soma skeleton skinning for visualization (SOMASkeleton30 / SOMASkeleton77).""" + +from pathlib import Path + +import numpy as np +import torch + +from ardy.assets import skeleton_asset_path +from ardy.skeleton import SOMASkeleton30, SOMASkeleton77 +from ardy.skeleton.kinematics import batch_rigid_transform +from ardy.skeleton.transforms import global_rots_to_local_rots + +SKEL_PATH = "somaskel77" +SKIN_NAME = "skin_standard.npz" +# Fallback path: skins ship with the ardy package under ardy/assets/skeletons/. +ASSETS_SKIN_PATH = skeleton_asset_path(SKEL_PATH, SKIN_NAME) + + +class SOMASkin: + def __init__(self, skeleton): + self.skeleton_input = skeleton + assert isinstance(skeleton, (SOMASkeleton30, SOMASkeleton77)), ( + "SOMASkin only supports SOMASkeleton30 or SOMASkeleton77" + ) + assert skeleton.neutral_joints is not None, "The skeleton must have neutral joints instantiated" + device = skeleton.neutral_joints.device + + # Resolve skin data path: try skeleton folder parent first, then packaged ardy/assets + skel_path = Path(skeleton.folder).parent / SKEL_PATH + skin_data_path = skel_path / SKIN_NAME + if not skin_data_path.exists(): + skin_data_path = ASSETS_SKIN_PATH + skel_path = skin_data_path.parent + + # the skin is always the 77-joint skeleton + self.skeleton_skin = SOMASkeleton77(folder=str(skel_path), load=True).to(device) + + skin_data = np.load(skin_data_path) + bind_rig_np = np.array(skin_data["bind_rig_transform"], dtype=np.float32) + self.bind_rig_transform = torch.from_numpy(bind_rig_np).to(device=device, dtype=torch.float) + bind_rig_inv_np = np.linalg.inv(bind_rig_np) + self.bind_rig_transform_inv = torch.from_numpy(bind_rig_inv_np).to(device=device, dtype=torch.float) + self.bind_vertices = torch.tensor(skin_data["bind_vertices"], device=device, dtype=torch.float) + self.faces = torch.tensor(skin_data["faces"], device=device, dtype=torch.long) + self.lbs_indices = torch.tensor(skin_data["lbs_indices"], device=device, dtype=torch.long) + self.lbs_weights = torch.tensor(skin_data["lbs_weights"], device=device, dtype=torch.float) + + # double check the rig matches expected skeleton + rig_joint_names = list(skin_data["rig_joint_names"]) + for sname, rname in zip(self.skeleton_skin.bone_order_names, rig_joint_names): + if sname != rname: + raise ValueError(f"MISMATCH in skinning rig: expected='{sname}' vs rig='{rname}'") + + def lbs(self, posed_transform): + bind_rig_transform_inv = self.bind_rig_transform_inv + bind_vertices = self.bind_vertices + lbs_weights = self.lbs_weights + for _ in range(posed_transform.dim() - 3): + bind_rig_transform_inv = bind_rig_transform_inv.unsqueeze(0) + bind_vertices = bind_vertices.unsqueeze(0) + lbs_weights = lbs_weights.unsqueeze(0) + + affine_mat = (posed_transform @ bind_rig_transform_inv)[..., :3, :] + vs = ( + affine_mat[..., self.lbs_indices, :, :] + @ torch.concat([bind_vertices, torch.ones_like(bind_vertices[..., 0:1])], dim=-1)[..., None, :, None] + ) + ws = lbs_weights[..., None, None] + resv = (vs * ws).sum(dim=-3).squeeze(-1) + return resv + + def skin(self, joint_rotmat, joint_pos, rot_is_global=False): + """ + joint_rotmat: [T, J, 3, 3] local or global joint rotation matrices + joint_pos: [T, J, 3] global joint positions + rot_is_global: bool, if True, joint_rotmat is global rotation matrices + """ + nF, nJ = joint_pos.shape[:2] + device = joint_rotmat.device + + if nJ != self.skeleton_skin.nbjoints: + assert nJ == 30, "SOMASkin currently only supports 30-joint or 77-joint skeletons" + + # make sure we have local joint rotations + if rot_is_global: + local_joint_rots_mats_subset = global_rots_to_local_rots(joint_rotmat, self.skeleton_input) + else: + local_joint_rots_mats_subset = joint_rotmat + + local_joint_rots_mats = self.skeleton_input.to_SOMASkeleton77(local_joint_rots_mats_subset) + + # FK to get the global joint pos and rot + neutral_joints_seq = self.skeleton_skin.neutral_joints[None].repeat((nF, 1, 1)).to(device) + new_joint_pos, joint_rotmat = batch_rigid_transform( + local_joint_rots_mats, + neutral_joints_seq, + self.skeleton_skin.joint_parents.to(device), + self.skeleton_skin.root_idx, + ) + joint_pos = new_joint_pos + joint_pos[:, self.skeleton_input.root_idx : self.skeleton_input.root_idx + 1] + nJ = self.skeleton_skin.nbjoints + rot_is_global = True + + # prepare full transformation matrices + fk_transform = torch.eye(4, device=device)[None, None].repeat(nF, nJ, 1, 1) + fk_transform[..., :3, 3] = joint_pos + if rot_is_global: + fk_transform[..., :3, :3] = joint_rotmat + else: + neutral_joints_seq = self.skeleton_skin.neutral_joints[None].repeat((nF, 1, 1)).to(device) + _, global_joint_rotmat = batch_rigid_transform( + joint_rotmat, + neutral_joints_seq, + self.skeleton_skin.joint_parents.to(device), + self.skeleton_skin.root_idx, + ) + fk_transform[..., :3, :3] = global_joint_rotmat + + vertices = self.lbs(fk_transform) + return vertices diff --git a/ardy/viz/viser_utils.py b/ardy/viz/viser_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..250947db90b07504ab89778d3ff24f92fe79b43c --- /dev/null +++ b/ardy/viz/viser_utils.py @@ -0,0 +1,2604 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: I001 +import os +import xml.etree.ElementTree as ET +from typing import List, Optional, Tuple, Union +import time + +import numpy as np +import torch +import trimesh +import viser +import viser.transforms as tf + +from scipy.interpolate import interp1d + +from ardy.assets import skeleton_asset_path +from ardy.skeleton.kinematics import batch_rigid_transform +from ardy.skeleton import ( + CoreSkeleton27, + G1Skeleton34, + SOMASkeleton30, + SOMASkeleton77, + SkeletonBase, +) +from ardy.motion_rep.smooth_root import get_smooth_root_pos +from ardy.skeleton.transforms import global_rots_to_local_rots +from ardy.tools import to_numpy, to_torch +from ardy.viz.core_skin import CoreSkin +from ardy.viz.g1_rig import G1MeshRig +from ardy.viz.soma_skin import SOMASkin + + +_G1_JOINT_AXIS_INDEX_CACHE: Optional[dict[str, int]] = None + + +def _get_g1_joint_axis_indices() -> dict[str, int]: + """Return a map from G1 joint names to a single rotation axis index.""" + global _G1_JOINT_AXIS_INDEX_CACHE + if _G1_JOINT_AXIS_INDEX_CACHE is not None: + return _G1_JOINT_AXIS_INDEX_CACHE + + xml_path = str(skeleton_asset_path("g1skel34", "xml", "g1.xml")) + if not os.path.exists(xml_path): + _G1_JOINT_AXIS_INDEX_CACHE = {} + return _G1_JOINT_AXIS_INDEX_CACHE + + tree = ET.parse(xml_path) + root = tree.getroot() + + joint_axes = {} + for xml_class in tree.findall(".//default"): + if "class" not in xml_class.attrib: + continue + joint_nodes = xml_class.findall("joint") + if joint_nodes: + joint_axes[xml_class.get("class")] = joint_nodes[0].get("axis") + + # mujoco (z-up, x-forward) -> ardy (y-up, z-forward) + mujoco_to_ardy = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64) + axis_indices_by_name: dict[str, int] = {} + for joint in root.find("worldbody").findall(".//joint"): + axis_str = joint.get("axis") or joint_axes.get(joint.get("class")) + if axis_str is None: + continue + axis_vals = np.array([float(x) for x in axis_str.split()], dtype=np.float64) + if not np.any(axis_vals): + continue + axis_ardy = mujoco_to_ardy @ axis_vals + axis_idx = int(np.argmax(np.abs(axis_ardy))) + axis_indices_by_name[joint.get("name").replace("_joint", "_skel")] = axis_idx + + _G1_JOINT_AXIS_INDEX_CACHE = axis_indices_by_name + return _G1_JOINT_AXIS_INDEX_CACHE + + +def _skew(v: np.ndarray) -> np.ndarray: + """Skew-symmetric matrix for cross products: skew(v) @ x == np.cross(v, x).""" + vx, vy, vz = float(v[0]), float(v[1]), float(v[2]) + return np.array([[0.0, -vz, vy], [vz, 0.0, -vx], [-vy, vx, 0.0]], dtype=np.float64) + + +def _rotation_matrix_from_two_vec(v_from: np.ndarray, v_to: np.ndarray, eps: float = 1e-8) -> np.ndarray: + """Return R such that R @ v_from ~= v_to (both treated as 3D vectors). + + Uses a Rodrigues-style construction, with special handling for near-parallel and near-opposite + vectors for numerical stability. + """ + a = np.asarray(v_from, dtype=np.float64).reshape(3) + b = np.asarray(v_to, dtype=np.float64).reshape(3) + na = np.linalg.norm(a) + nb = np.linalg.norm(b) + if na < eps or nb < eps: + return np.eye(3, dtype=np.float64) + a = a / na + b = b / nb + + c = float(np.clip(np.dot(a, b), -1.0, 1.0)) # cos(theta) + if c > 1.0 - eps: + return np.eye(3, dtype=np.float64) + if c < -1.0 + eps: + # 180 deg rotation about any axis orthogonal to a: + # R = -I + 2 * uu^T, where u is a unit axis orthogonal to a. + axis_seed = np.array([1.0, 0.0, 0.0], dtype=np.float64) + if abs(float(np.dot(a, axis_seed))) > 0.9: + axis_seed = np.array([0.0, 1.0, 0.0], dtype=np.float64) + u = np.cross(a, axis_seed) + u = u / np.linalg.norm(u).clip(min=eps) + return -np.eye(3, dtype=np.float64) + 2.0 * np.outer(u, u) + + v = np.cross(a, b) + s2 = float(np.dot(v, v)) # ||v||^2 == sin^2(theta) + K = _skew(v) + # R = I + K + K^2 * ((1 - c) / s^2) + return np.eye(3, dtype=np.float64) + K + (K @ K) * ((1.0 - c) / s2) + + +# TODO: should things in here by kept on cpu with numpy to avoid latency when interacting with UI? +# the downside is we need torch/GPU for FK + +# Cache arrow meshes to avoid recreating them repeatedly +_CACHED_ARROW_BASE = None +_CACHED_ARROW_HEAD = None + + +def _get_cached_arrow_meshes(): + """Get cached arrow base and head meshes (vertices and faces only).""" + global _CACHED_ARROW_BASE, _CACHED_ARROW_HEAD + if _CACHED_ARROW_BASE is None: + arrow_base_mesh = trimesh.creation.cylinder(radius=0.01, height=0.2) + arrow_head_mesh = trimesh.creation.cone(radius=0.03, height=0.05) + _CACHED_ARROW_BASE = { + "vertices": arrow_base_mesh.vertices.copy(), + "faces": arrow_base_mesh.faces.copy(), + } + _CACHED_ARROW_HEAD = { + "vertices": arrow_head_mesh.vertices.copy(), + "faces": arrow_head_mesh.faces.copy(), + } + return _CACHED_ARROW_BASE, _CACHED_ARROW_HEAD + + +class WaypointMesh: + def __init__( + self, + name: str, + server: viser.ViserServer, + position: np.ndarray, + heading: Optional[np.ndarray] = None, + color: Optional[Tuple[int, int, int]] = (255, 0, 0), + add_annulus: bool = True, + ): + self.server = server + self.color = color + self.base_name = name # Store base name to remove parent folder later + + sphere = trimesh.creation.icosphere(subdivisions=3, radius=0.025) + + z_to_y_up = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]) + + self.sphere = self.server.scene.add_mesh_simple( + name=f"{name}/sphere", + vertices=sphere.vertices, + faces=sphere.faces, + position=position, + color=color, + ) + + if add_annulus: + annulus = trimesh.creation.annulus(r_min=0.1, r_max=0.2, height=0.005) + annulus_vertices = annulus.vertices @ z_to_y_up + self.annulus = self.server.scene.add_mesh_simple( + name=f"{name}/annulus", + vertices=annulus_vertices, + faces=annulus.faces, + position=position, + color=color, + ) + else: + self.annulus = None + + self.arrow_base = None + self.arrow_head = None + if heading is not None: + assert heading.shape == (2,), "Heading must be a 2D vector" + heading_norm = heading / np.linalg.norm(heading) + heading_scaled = 0.2 * heading_norm + heading_3d = np.array([heading_scaled[0], 0, heading_scaled[1]]) + + # Calculate rotation to align Y-axis (default cylinder/cone orientation) with heading + # Rotation angle around Y-axis + angle = np.arctan2(heading_norm[0], heading_norm[1]) # heading = [cos, sin] -> angle + cos_a, sin_a = np.cos(angle), np.sin(angle) + # Rotation matrix around Y-axis + rot_y = np.array([[cos_a, 0, sin_a], [0, 1, 0], [-sin_a, 0, cos_a]]) + + # Use cached arrow meshes + arrow_base_cache, arrow_head_cache = _get_cached_arrow_meshes() + + # Rotate vertices to align with heading direction + arrow_base_vertices = arrow_base_cache["vertices"] @ rot_y.T + arrow_head_vertices = arrow_head_cache["vertices"] @ rot_y.T + + self.arrow_base = self.server.scene.add_mesh_simple( + name=f"{name}/arrow_base", + vertices=arrow_base_vertices, + faces=arrow_base_cache["faces"], + position=position + (heading_3d / 2), + color=color, + ) + self.arrow_head = self.server.scene.add_mesh_simple( + name=f"{name}/arrow_head", + vertices=arrow_head_vertices, + faces=arrow_head_cache["faces"], + position=position + heading_3d, + color=color, + ) + + def update_position(self, position: np.ndarray, heading: Optional[np.ndarray] = None): + self.sphere.position = position + if self.annulus is not None: + self.annulus.position = position + if heading is not None: + assert heading.shape == (2,), "Heading must be a 2D vector" + heading_norm = heading / np.linalg.norm(heading) + heading_scaled = 0.2 * heading_norm + heading_3d = np.array([heading_scaled[0], 0, heading_scaled[1]]) + + # Calculate rotation to align Y-axis with heading + angle = np.arctan2(heading_norm[0], heading_norm[1]) + cos_a, sin_a = np.cos(angle), np.sin(angle) + rot_y = np.array([[cos_a, 0, sin_a], [0, 1, 0], [-sin_a, 0, cos_a]]) + + # Remove old arrow meshes if they exist + if self.arrow_base is not None: + self.server.scene.remove_by_name(self.arrow_base.name) + if self.arrow_head is not None: + self.server.scene.remove_by_name(self.arrow_head.name) + + # Use cached arrow meshes + arrow_base_cache, arrow_head_cache = _get_cached_arrow_meshes() + + # Rotate vertices to align with heading direction + arrow_base_vertices = arrow_base_cache["vertices"] @ rot_y.T + arrow_head_vertices = arrow_head_cache["vertices"] @ rot_y.T + + # Get base name from sphere (e.g., "/waypoint_0/sphere" -> "/waypoint_0") + base_name = self.sphere.name.rsplit("/", 1)[0] + + self.arrow_base = self.server.scene.add_mesh_simple( + name=f"{base_name}/arrow_base", + vertices=arrow_base_vertices, + faces=arrow_base_cache["faces"], + position=position + (heading_3d / 2), + color=self.color, + ) + self.arrow_head = self.server.scene.add_mesh_simple( + name=f"{base_name}/arrow_head", + vertices=arrow_head_vertices, + faces=arrow_head_cache["faces"], + position=position + heading_3d, + color=self.color, + ) + + def clear(self): + """Remove all waypoint meshes and the parent folder from the scene.""" + # self.server.scene.remove_by_name(self.sphere.name) + # if self.annulus is not None: + # self.server.scene.remove_by_name(self.annulus.name) + # if self.arrow_base is not None: + # self.server.scene.remove_by_name(self.arrow_base.name) + # if self.arrow_head is not None: + # self.server.scene.remove_by_name(self.arrow_head.name) + # Remove the parent folder to prevent empty folders in scene tree + try: + # print(f"Removing parent folder {self.base_name}") + self.server.scene.remove_by_name(self.base_name) + except Exception as e: + print(f"Error removing parent folder {self.base_name}: {e}") + pass # Parent folder might have already been removed + + +class VelocityArrowMesh: + """Visualizes root velocity as an arrow (line segment + cone head).""" + + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + color: tuple = (50, 150, 255), # Default blue color (RGB) + ): + """Initialize velocity arrow visualization. + + Args: + name: str, base name for the arrow components + server: viser.ViserServer, server to add the arrow to + skeleton: SkeletonBase, skeleton to get root index from + color: tuple, RGB color tuple (0-255), default is blue + """ + self.name = name + self.server = server + self.skeleton = skeleton + self.color = color + + # Arrow components + self.arrow_line = None # Line segment + self.arrow_cone = None # Cone head + self.should_show = False # Track if arrow should be visible based on velocity magnitude + + def update( + self, + root_velocity: Optional[Union[np.ndarray, torch.Tensor]], + root_pos: Union[np.ndarray, torch.Tensor], + visible: bool = True, + ): + """Update the velocity arrow visualization. + + Args: + root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s + root_pos: Union[np.ndarray, torch.Tensor], [3] root position + visible: bool, whether the arrow should be visible (controlled by skeleton visibility) + """ + if root_velocity is None: + # Hide arrow if no velocity provided + self.should_show = False + if self.arrow_line is not None: + self.arrow_line.visible = False + self.arrow_cone.visible = False + return + + # Convert to numpy if tensor + if isinstance(root_velocity, torch.Tensor): + root_velocity = root_velocity.detach().cpu().numpy() + if isinstance(root_pos, torch.Tensor): + root_pos = root_pos.detach().cpu().numpy() + + # Project velocity to XZ plane + velocity_xz = np.array([root_velocity[0], 0.0, root_velocity[2]]) + velocity_magnitude = np.linalg.norm(velocity_xz) + + # Only show arrow if velocity is significant (> 0.1 m/s) + if velocity_magnitude <= 0.1: + self.should_show = False + if self.arrow_line is not None: + self.arrow_line.visible = False + self.arrow_cone.visible = False + return + + # Arrow should be shown + self.should_show = True + + # Calculate arrow geometry + root_2d_pos = np.array([root_pos[0], 0.0, root_pos[2]]) # Project to ground + velocity_dir = velocity_xz / velocity_magnitude + arrow_length = velocity_magnitude / 4.0 + arrow_end = root_2d_pos + velocity_dir * arrow_length + + # Calculate rotation quaternion for arrow head + from_vec = np.array([0.0, 0.0, 1.0]) + to_vec = velocity_dir + rot_mat = _rotation_matrix_from_two_vec(from_vec, to_vec) + quat = tf.SO3.from_matrix(rot_mat).wxyz + + # Create or update arrow components + if self.arrow_line is None: + # Create line segment + self.arrow_line = self.server.scene.add_line_segments( + name=f"{self.name}/velocity_line", + points=np.array([[root_2d_pos, arrow_end]]), + colors=self.color, + line_width=3.0, + ) + + # Create cone head + arrow_head = trimesh.creation.cone(radius=0.04, height=0.1) + self.arrow_cone = self.server.scene.add_mesh_simple( + name=f"{self.name}/velocity_cone", + vertices=arrow_head.vertices, + faces=arrow_head.faces, + color=self.color, + position=arrow_end, + wxyz=quat, + ) + else: + # Update existing components atomically + # Hide both during update + old_line_visible = self.arrow_line.visible + old_cone_visible = self.arrow_cone.visible + self.arrow_line.visible = False + self.arrow_cone.visible = False + + # Update geometry + self.arrow_line.points = np.array([[root_2d_pos, arrow_end]]) + self.arrow_cone.position = arrow_end - 0.05 * velocity_dir + self.arrow_cone.wxyz = quat + + # Restore visibility + self.arrow_line.visible = old_line_visible + self.arrow_cone.visible = old_cone_visible + + # Set visibility based on skeleton visibility and should_show + self.arrow_line.visible = visible and self.should_show + self.arrow_cone.visible = visible and self.should_show + + def set_visibility(self, visible: bool): + """Set visibility of the velocity arrow.""" + if self.arrow_line is not None: + self.arrow_line.visible = visible and self.should_show + if self.arrow_cone is not None: + self.arrow_cone.visible = visible and self.should_show + + def clear(self): + """Remove the velocity arrow from the scene.""" + if self.arrow_line is not None: + self.server.scene.remove_by_name(self.arrow_line.name) + self.arrow_line = None + if self.arrow_cone is not None: + self.server.scene.remove_by_name(self.arrow_cone.name) + self.arrow_cone = None + + +class SkeletonMesh: + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + joint_color: Optional[Tuple[float, float, float] | np.ndarray] = ( + 255, + 235, + 0, + ), + bone_color: Optional[Tuple[float, float, float] | np.ndarray] = ( + 27, + 106, + 0, + ), + starting_joints_pos: Optional[torch.Tensor] = None, + show_root_2d_projection: bool = False, + ): + """ + name: str, name of the skeleton mesh + server: viser.ViserServer, server to add the skeleton mesh to + skeleton: SkeletonBase, skeleton to visualize (must be a CoreSkeleton27) + joint_color: Optional[Tuple[float, float, float] | np.ndarray], color of the joints, either (3,) or (J, 3) + bone_color: Optional[Tuple[float, float, float] | np.ndarray], color of the bones, either (3,) or (J-1, 3) + starting_joints_pos: Optional[torch.Tensor], starting joint positions (if None, will use neutral pose) + show_root_2d_projection: bool, whether to show the 2D root projection as a blue sphere + """ + self.server = server + self.skeleton = skeleton + self.show_root_2d_projection = show_root_2d_projection + joint_mesh = trimesh.creation.icosphere(subdivisions=3, radius=0.02) + bone_mesh = trimesh.creation.cylinder(radius=0.01, height=1.0) + + init_joints_pos = skeleton.neutral_joints.clone() + self.num_joints = init_joints_pos.shape[0] + num_bones = self.num_joints - 1 + non_root_bones = [ + joint_name + for joint_name, parent_name in self.skeleton.bone_order_names_with_parents + if parent_name is not None + ] + self.bone_to_idx = {bone_name: idx for idx, bone_name in enumerate(non_root_bones)} + + # initialize meshes + init_joints_wxyzs = np.concatenate([np.ones((self.num_joints, 1)), np.zeros((self.num_joints, 3))], axis=1) + if isinstance(joint_color, tuple): + self.joint_colors = np.full((self.num_joints, 3), joint_color) + elif isinstance(joint_color, np.ndarray): + assert joint_color.shape == ( + self.num_joints, + 3, + ), "Joint colors must be (J, 3)" + self.joint_colors = joint_color + self.joints_batched_mesh = server.scene.add_batched_meshes_simple( + f"{name}/joints", + vertices=joint_mesh.vertices, + faces=joint_mesh.faces, + batched_wxyzs=init_joints_wxyzs, + batched_positions=np.zeros((self.num_joints, 3)), + batched_scales=np.ones((self.num_joints, 3)), + batched_colors=self.joint_colors, + ) + init_bones_wxyzs = np.concatenate([np.ones((num_bones, 1)), np.zeros((num_bones, 3))], axis=1) + if isinstance(bone_color, tuple): + bone_color = np.full((num_bones, 3), bone_color) + elif isinstance(bone_color, np.ndarray): + assert bone_color.shape == (num_bones, 3), "Bone colors must be (J-1, 3)" + bone_color = bone_color + self.bones_batched_mesh = server.scene.add_batched_meshes_simple( + f"{name}/bones", + vertices=bone_mesh.vertices, + faces=bone_mesh.faces, + batched_wxyzs=init_bones_wxyzs, + batched_positions=np.zeros((num_bones, 3)), + batched_scales=np.ones((num_bones, 3)), + batched_colors=bone_color, + ) + + # Initialize 2D root projection sphere (blue) + if self.show_root_2d_projection: + root_2d_mesh = trimesh.creation.icosphere(subdivisions=2, radius=0.03) + self.root_2d_sphere = server.scene.add_mesh_simple( + f"{name}/root_2d_projection", + vertices=root_2d_mesh.vertices, + faces=root_2d_mesh.faces, + color=(50, 150, 255), # Blue color + position=(0, 0, 0), + ) + else: + self.root_2d_sphere = None + + # Initialize root velocity arrow (visualize XZ projection) + self.velocity_arrow_mesh = VelocityArrowMesh( + name=name, + server=server, + skeleton=skeleton, + ) + + # used if precomputed meshes are used + self.mesh_info_cache = None + + if starting_joints_pos is not None: + self.set_pose(starting_joints_pos) + else: + # set them to neutral pose + min_height = init_joints_pos[:, 1].min().item() + init_joints_pos[:, 1] -= min_height # move to be on ground + self.set_pose(init_joints_pos) + + def compute_single_pose(self, joints_pos: np.ndarray): + """Compute the mesh for a single frame. + + joints_pos: [J, 3] global joint positions + """ + # compute bone transforms + new_batched_positions = np.zeros((self.skeleton.nbjoints - 1, 3)) + new_batched_wxyzs = np.zeros((self.skeleton.nbjoints - 1, 4)) + new_batched_scales = np.ones((self.skeleton.nbjoints - 1, 3)) + for joint_name, parent_name in self.skeleton.bone_order_names_with_parents: + if parent_name is None: + continue + joint_idx = self.skeleton.bone_index[joint_name] + parent_idx = self.skeleton.bone_index[parent_name] + joint_pos = joints_pos[joint_idx] + parent_pos = joints_pos[parent_idx] + + bone_pos = (joint_pos + parent_pos) / 2.0 + bone_scale = np.linalg.norm(joint_pos - parent_pos) + if bone_scale < 1e-8: + bone_wxyz = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64) + else: + bone_dir = (joint_pos - parent_pos) / bone_scale + R = _rotation_matrix_from_two_vec(np.array([0.0, 0.0, 1.0], dtype=np.float64), bone_dir) + bone_wxyz = tf.SO3.from_matrix(R).wxyz + + bone_idx = self.bone_to_idx[joint_name] + new_batched_positions[bone_idx] = bone_pos + new_batched_wxyzs[bone_idx] = bone_wxyz + new_batched_scales[bone_idx] = np.array([1.0, 1.0, bone_scale], dtype=float) + + return new_batched_positions, new_batched_wxyzs, new_batched_scales + + def precompute_mesh_info(self, joints_pos: torch.Tensor): + """Precompute the meshes for all frames at once. + + joints_pos: [T, J, 3] global joint positions + """ + joints_pos = joints_pos.cpu().numpy() + # compute bone transforms + num_frames = joints_pos.shape[0] + self.mesh_info_cache = { + "positions": np.zeros((num_frames, self.skeleton.nbjoints - 1, 3)), + "wxyzs": np.zeros((num_frames, self.skeleton.nbjoints - 1, 4)), + "scales": np.ones((num_frames, self.skeleton.nbjoints - 1, 3)), + } + for i in range(num_frames): + new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose(joints_pos[i]) + self.mesh_info_cache["positions"][i] = new_batched_positions + self.mesh_info_cache["wxyzs"][i] = new_batched_wxyzs + self.mesh_info_cache["scales"][i] = new_batched_scales + + def update_mesh_info_cache(self, joints_pos: torch.Tensor, frame_idx: int): + """Update the mesh info cache for the given frame. + + joints_pos: [J, 3] global joint positions + """ + assert self.mesh_info_cache is not None + new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose( + joints_pos.cpu().numpy() + ) + self.mesh_info_cache["positions"][frame_idx] = new_batched_positions + self.mesh_info_cache["wxyzs"][frame_idx] = new_batched_wxyzs + self.mesh_info_cache["scales"][frame_idx] = new_batched_scales + + def set_pose( + self, + joints_pos: torch.Tensor, + foot_contacts: Optional[torch.Tensor] = None, + frame_idx: Optional[int] = None, + root_velocity: Optional[Union[np.ndarray, torch.Tensor]] = None, + ): + """ + joints_pos: [J, 3] global joint positions + foot_contacts: [4] contact labels for left heel/toe and right heel/toe, 1 for in contact + frame_idx: int, index of the frame to set the pose for (only needed if precomputed mesh info is used) + root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s + """ + self.cur_joints_pos = joints_pos + joints_pos = joints_pos.cpu().numpy() + + if self.mesh_info_cache is not None: + assert frame_idx is not None + new_batched_positions = self.mesh_info_cache["positions"][frame_idx] + new_batched_wxyzs = self.mesh_info_cache["wxyzs"][frame_idx] + new_batched_scales = self.mesh_info_cache["scales"][frame_idx] + else: + new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose(joints_pos) + + # update meshes + self.bones_batched_mesh.batched_positions = new_batched_positions + self.bones_batched_mesh.batched_wxyzs = new_batched_wxyzs + self.bones_batched_mesh.batched_scales = new_batched_scales + # directly set joint positions + self.joints_batched_mesh.batched_positions = joints_pos + + # update 2D root projection sphere + if self.root_2d_sphere is not None: + root_pos = joints_pos[self.skeleton.root_idx] + root_2d_pos = np.array([root_pos[0], 0.0, root_pos[2]]) + self.root_2d_sphere.position = root_2d_pos + + # Update root velocity arrow visualization (XZ projection) + root_pos = joints_pos[self.skeleton.root_idx] + skeleton_visible = self.joints_batched_mesh.visible + self.velocity_arrow_mesh.update( + root_velocity=root_velocity, + root_pos=root_pos, + visible=skeleton_visible, + ) + + # update colors for foot contacts + if foot_contacts is not None: + cur_joint_colors = self.joint_colors.copy() + foot_contacts = foot_contacts.bool().cpu().numpy().astype(bool) + foot_joints = np.array(self.skeleton.foot_joint_idx, dtype=int) + contact_idx = foot_joints[foot_contacts] + cur_joint_colors[contact_idx] = (160, 32, 240) + self.joints_batched_mesh.batched_colors = cur_joint_colors + else: + self.joints_batched_mesh.batched_colors = self.joint_colors + + def set_visibility(self, visible: bool): + self.joints_batched_mesh.visible = visible + self.bones_batched_mesh.visible = visible + if self.root_2d_sphere is not None: + self.root_2d_sphere.visible = visible + # Update velocity arrow visibility + self.velocity_arrow_mesh.set_visibility(visible) + + def get_pose(self) -> np.ndarray: + return self.cur_joints_pos + + def clear(self): + names = [mesh.name for mesh in [self.joints_batched_mesh, self.bones_batched_mesh]] + for name in names: + self.server.scene.remove_by_name(name) + if self.root_2d_sphere is not None: + self.server.scene.remove_by_name(self.root_2d_sphere.name) + # Clear velocity arrow + self.velocity_arrow_mesh.clear() + + +LIGHT_THEME = dict( + mesh=(152, 189, 255), # (90, 200, 255) - original viser blue +) + +DARK_THEME = dict( + mesh=(60, 85, 130), +) + + +class Character: + def __init__( + self, + name: str, + server: viser.ViserServer | viser.ClientHandle, + skeleton: SkeletonBase, + create_skeleton_mesh: bool = True, + create_skinned_mesh: bool = True, + visible_skeleton: bool = False, + visible_skinned_mesh: bool = True, + skinned_mesh_opacity: float = 1.0, + show_foot_contacts: bool = True, + dark_mode: bool = False, + mesh_mode: str = "core_skin", + g1_mesh_dir: Optional[str] = None, + show_root_2d_projection: bool = False, + ): + self.server = server + self.name = name + self.skeleton = skeleton + self.cur_joints_pos = None + self.cur_joints_rot = None + self.cur_foot_contacts = None + + self.skeleton_mesh = None + self.show_foot_contacts = show_foot_contacts + if create_skeleton_mesh: + self.skeleton_mesh = SkeletonMesh( + f"/{name}/skeleton", + server, + skeleton, + show_root_2d_projection=show_root_2d_projection, + ) + # init with default rest pose + self.cur_joints_pos = self.skeleton_mesh.get_pose() + self.skeleton_mesh.set_visibility(visible_skeleton) + + self.skinned_mesh = None + self.g1_mesh_rig = None + self.skin = None + self.mesh_mode = mesh_mode + if create_skinned_mesh: + if isinstance(self.skeleton, CoreSkeleton27) and mesh_mode == "core_skin": + self.skin = CoreSkin(self.skeleton) + self.skinned_mesh = server.scene.add_mesh_simple( + f"/{name}/simple_skinned", + vertices=self.skin.bind_vertices.cpu().numpy(), + faces=self.skin.faces.cpu().numpy(), + opacity=None, + color=LIGHT_THEME["mesh"] if not dark_mode else DARK_THEME["mesh"], + wireframe=False, + visible=False, + ) + self.skinned_verts_cache = None + + bind_pos = self.skeleton.neutral_joints.clone() + min_height = bind_pos[:, 1].min().item() + bind_pos[:, 1] -= min_height # move to be on ground + bind_rotmat = torch.eye(3, device=bind_pos.device).repeat(bind_pos.shape[0], 1, 1) + self.set_pose(bind_pos, bind_rotmat) + self.skinned_mesh.visible = True # avoid blinking + self.set_skinned_mesh_visibility(visible_skinned_mesh) + self.set_skinned_mesh_opacity(skinned_mesh_opacity) + elif isinstance(self.skeleton, (SOMASkeleton30, SOMASkeleton77)) and mesh_mode == "soma_skin": + self.skin = SOMASkin(self.skeleton) + self.skinned_mesh = server.scene.add_mesh_simple( + f"/{name}/simple_skinned", + vertices=self.skin.bind_vertices.cpu().numpy(), + faces=self.skin.faces.cpu().numpy(), + opacity=None, + color=LIGHT_THEME["mesh"] if not dark_mode else DARK_THEME["mesh"], + wireframe=False, + visible=False, + ) + self.skinned_verts_cache = None + + bind_pos = self.skeleton.neutral_joints.clone() + min_height = bind_pos[:, 1].min().item() + bind_pos[:, 1] -= min_height # move to be on ground + bind_rotmat = torch.eye(3, device=bind_pos.device).repeat(bind_pos.shape[0], 1, 1) + self.set_pose(bind_pos, bind_rotmat) + self.skinned_mesh.visible = True # avoid blinking + self.set_skinned_mesh_visibility(visible_skinned_mesh) + self.set_skinned_mesh_opacity(skinned_mesh_opacity) + elif isinstance(self.skeleton, G1Skeleton34) and mesh_mode == "g1_stl": + if g1_mesh_dir is None: + g1_mesh_dir = os.path.join( + os.path.dirname(__file__), + "..", + "assets", + "skeletons", + "g1skel34", + "meshes", + "g1", + ) + g1_mesh_dir = os.path.abspath(g1_mesh_dir) + if not os.path.exists(g1_mesh_dir): + print(f"G1 mesh directory not found: {g1_mesh_dir}") + self.g1_mesh_rig = G1MeshRig( + name, + server, + self.skeleton, + g1_mesh_dir, + DARK_THEME["mesh"] if dark_mode else LIGHT_THEME["mesh"], + ) + init_joints_pos = self.skeleton.neutral_joints.clone() + min_height = init_joints_pos[:, 1].min().item() + init_joints_pos[:, 1] -= min_height # move to be on ground + init_joints_rot = torch.eye(3, device=init_joints_pos.device).repeat(init_joints_pos.shape[0], 1, 1) + self.set_pose(init_joints_pos, init_joints_rot) + self.set_skinned_mesh_visibility(visible_skinned_mesh) + self.set_skinned_mesh_opacity(skinned_mesh_opacity) + else: + raise ValueError( + "Unsupported mesh mode for skeleton type: " + f"{type(self.skeleton).__name__} with mesh_mode={mesh_mode}" + ) + + def change_theme(self, is_dark_mode): + color = DARK_THEME["mesh"] if is_dark_mode else LIGHT_THEME["mesh"] + if self.skinned_mesh is not None: + self.skinned_mesh.color = color + if self.g1_mesh_rig is not None: + self.g1_mesh_rig.set_color(color) + + def set_skeleton_visibility(self, visible: bool): + if self.skeleton_mesh is not None: + self.skeleton_mesh.set_visibility(visible) + + def set_show_foot_contacts(self, show: bool): + self.show_foot_contacts = show + + def set_skinned_mesh_visibility(self, visible: bool): + if self.skinned_mesh is not None: + self.skinned_mesh.visible = visible + if self.g1_mesh_rig is not None: + self.g1_mesh_rig.set_visibility(visible) + + def set_skinned_mesh_opacity(self, opacity: float): + if self.skinned_mesh is not None: + self.skinned_mesh.opacity = opacity + if self.g1_mesh_rig is not None: + self.g1_mesh_rig.set_opacity(opacity) + + def set_skinned_mesh_wireframe(self, wireframe: bool): + if self.skinned_mesh is not None: + self.skinned_mesh.wireframe = wireframe + if self.g1_mesh_rig is not None: + self.g1_mesh_rig.set_wireframe(wireframe) + + def precompute_skinning(self, joints_pos: torch.Tensor, joints_rot: torch.Tensor): + """If using simple skinning, we can precompute the skinning for all frames at once. + + joints_pos: [T, J, 3] global joint positions + joints_rot: [T, J, 3, 3] global joint rotation matrices + """ + assert self.skin is not None + self.skinned_verts_cache = self.skin.skin(joints_rot, joints_pos, rot_is_global=True).cpu().numpy() + + def update_skinning_cache(self, joints_pos: torch.Tensor, joints_rot: torch.Tensor, frame_idx: int): + """Update the skinning cache for the given frame. + + joints_pos: [J, 3] global joint positions + joints_rot: [J, 3, 3] global joint rotation matrices + frame_idx: int, index of the frame to update the cache for + """ + if self.skinned_verts_cache is None: + return + + new_skinned_verts = self.skin.skin(joints_rot[None], joints_pos[None], rot_is_global=True)[0].cpu().numpy() + self.skinned_verts_cache[frame_idx] = new_skinned_verts + + def set_pose( + self, + joints_pos: torch.Tensor, + joints_rot: torch.Tensor, + foot_contacts: Optional[torch.Tensor] = None, + frame_idx: Optional[int] = None, + root_velocity: Optional[Union[np.ndarray, torch.Tensor]] = None, + ): + """ + joints_pos: [J, 3] global joint positions + joints_rot: [J, 3, 3] global joint rotation matrices + frame_idx: int, index of the frame to set the pose for (only needed if precomputed skinning is used) + root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s + """ + if self.skeleton_mesh is not None: + cur_foot_contacts = foot_contacts if self.show_foot_contacts else None + self.skeleton_mesh.set_pose( + joints_pos, + foot_contacts=cur_foot_contacts, + frame_idx=frame_idx, + root_velocity=root_velocity, + ) + self.cur_foot_contacts = cur_foot_contacts + + if self.skinned_mesh is not None: + if self.skinned_verts_cache is not None: + assert frame_idx is not None + skinned_verts = self.skinned_verts_cache[frame_idx] + else: + skinned_verts = self.skin.skin(joints_rot[None], joints_pos[None], rot_is_global=True)[0].cpu().numpy() + + # update the vertices + self.skinned_mesh.vertices = skinned_verts + if self.g1_mesh_rig is not None: + joints_pos_np = joints_pos.detach().cpu().numpy() + joints_rot_np = joints_rot.detach().cpu().numpy() + self.g1_mesh_rig.set_pose(joints_pos_np, joints_rot_np) + + self.cur_joints_pos = joints_pos + self.cur_joints_rot = joints_rot + + def get_pose(self) -> torch.Tensor: + return self.cur_joints_pos, self.cur_joints_rot + + def clear(self): + if self.skeleton_mesh is not None: + self.skeleton_mesh.clear() + if self.skinned_mesh is not None: + self.server.scene.remove_by_name(self.skinned_mesh.name) + if self.g1_mesh_rig is not None: + self.g1_mesh_rig.clear() + + +class CharacterMotion: + def __init__( + self, + character: Character, + joints_pos: torch.Tensor, + joints_rot: torch.Tensor, + foot_contacts: Optional[torch.Tensor] = None, + ): + self.character = character + self.server = character.server + self.skeleton = character.skeleton + self.name = character.name + + # [T, J, 3] global joint positions + self.joints_pos = joints_pos + # [T, J, 3, 3] global joint rotation matrices + self.joints_rot = joints_rot + assert joints_pos.shape[0] == joints_rot.shape[0] + # keep track of local rots as well for convenience during pose editing + self.joints_local_rot = global_rots_to_local_rots(joints_rot, self.skeleton) + + self.length = joints_pos.shape[0] + self.cur_frame_idx = None + + self.foot_contacts = foot_contacts + if foot_contacts is not None: + assert foot_contacts.shape[0] == self.length + + self.precompute_mesh_info() + + # gizmos for pose editing + self.root_translation_gizmo = None + self.updating_root_translation_gizmo = False + self.joint_gizmos = None + self.updating_joint_gizmos = False + + def precompute_mesh_info(self): + if self.character.skeleton_mesh is not None: + print("Caching skeleton mesh info...") + self.character.skeleton_mesh.precompute_mesh_info(self.joints_pos) + if self.character.skinned_mesh is not None: + print("Caching skinning info...") + self.character.precompute_skinning(self.joints_pos, self.joints_rot) + + def set_frame(self, idx: int): + """Sets the pose of the character to the given frame index.""" + idx = min(idx, self.length - 1) # clamp to last frame + cur_foot_contacts = self.foot_contacts[idx] if self.foot_contacts is not None else None + self.character.set_pose( + self.joints_pos[idx], + self.joints_rot[idx], + frame_idx=idx, + foot_contacts=cur_foot_contacts, + ) + self.cur_frame_idx = idx + + # update gizmos if frame has changed due to playback + cur_root_pos = self.joints_pos[self.cur_frame_idx, self.skeleton.root_idx].clone() + cur_root_pos[1] = 0.0 + if self.root_translation_gizmo is not None and not self.updating_root_translation_gizmo: + self.root_translation_gizmo.position = cur_root_pos.cpu().numpy() + if self.joint_gizmos is not None: + for i, joint_gizmo in enumerate(self.joint_gizmos): + if not self.updating_joint_gizmos: + joint_gizmo.position = self.joints_pos[self.cur_frame_idx, i].cpu().numpy() + joint_gizmo.wxyz = tf.SO3.from_matrix( + self.joints_local_rot[self.cur_frame_idx, i].cpu().numpy() + ).wxyz + + def update_pose_at_frame( + self, + frame_idx: int, + joints_pos: Optional[torch.Tensor] = None, + joints_rot: Optional[torch.Tensor] = None, + joints_local_rot: Optional[torch.Tensor] = None, + foot_contacts: Optional[torch.Tensor] = None, + ): + """Overwrites one or more of the pose components at the given frame. + + If only a subset of joints_pos, joints_rot, or joints_local_rot are provided, the other + components will be updated with FK. + """ + if joints_pos is not None: + joints_pos = to_torch(joints_pos, device=self.joints_pos.device, dtype=self.joints_pos.dtype) + self.joints_pos[frame_idx] = joints_pos + if joints_local_rot is None and joints_rot is None: + raise NotImplementedError("No IK to update joint rotations accordingly.") + if joints_rot is not None: + joints_rot = to_torch(joints_rot, device=self.joints_rot.device, dtype=self.joints_rot.dtype) + self.joints_rot[frame_idx] = joints_rot + if joints_local_rot is None: + # update local rots from global rots + self.joints_local_rot[frame_idx] = global_rots_to_local_rots(joints_rot, self.skeleton) + if joints_pos is None: + # need to update with FK + new_posed_joints, _ = batch_rigid_transform( + self.joints_local_rot[frame_idx : frame_idx + 1], + self.skeleton.neutral_joints[None].to(self.joints_local_rot.device), + self.skeleton.joint_parents.to(self.joints_local_rot.device), + self.skeleton.root_idx, + ) + new_posed_joints = ( + new_posed_joints[0] + + self.joints_pos[frame_idx, self.skeleton.root_idx : self.skeleton.root_idx + 1] + - self.skeleton.neutral_joints[[self.skeleton.root_idx]] + ) + self.joints_pos[frame_idx] = new_posed_joints + if joints_local_rot is not None: + joints_local_rot = to_torch(joints_local_rot, device=self.joints_local_rot.device).to( + dtype=self.joints_local_rot.dtype + ) + self.joints_local_rot[frame_idx] = joints_local_rot + if joints_rot is None or joints_pos is None: + # need to update with FK + new_posed_joints, new_global_rots = batch_rigid_transform( + self.joints_local_rot[frame_idx : frame_idx + 1], + self.skeleton.neutral_joints[None].to(self.joints_local_rot.device), + self.skeleton.joint_parents.to(self.joints_local_rot.device), + self.skeleton.root_idx, + ) + new_posed_joints = ( + new_posed_joints[0] + + self.joints_pos[frame_idx, self.skeleton.root_idx : self.skeleton.root_idx + 1] + - self.skeleton.neutral_joints[[self.skeleton.root_idx]] + ) + if joints_rot is None: + self.joints_rot[frame_idx] = new_global_rots[0] + if joints_pos is None: + self.joints_pos[frame_idx] = new_posed_joints + if foot_contacts is not None: + foot_contacts = to_torch(foot_contacts, device=self.foot_contacts.device).to(dtype=self.foot_contacts.dtype) + self.foot_contacts[frame_idx] = foot_contacts + + if self.character.skeleton_mesh is not None: + self.character.skeleton_mesh.update_mesh_info_cache(self.joints_pos[frame_idx], frame_idx) + if self.character.skinned_mesh is not None: + self.character.update_skinning_cache(self.joints_pos[frame_idx], self.joints_rot[frame_idx], frame_idx) + + def clear(self): + self.character.clear() + + # + # Editing helpers + # + def get_current_projected_root_pos(self) -> np.ndarray: + """Get the projected root position on the ground at the current frame.""" + root_pos = self.joints_pos[self.cur_frame_idx, self.skeleton.root_idx].clone() + root_pos[1] = 0.0 + return to_numpy(root_pos) + + def get_projected_root_pos(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray: + """If requested frames are out of range, simply pads with the last frame to get expected + length.""" + if end_frame_idx is None: + expected_len = 1 + else: + expected_len = end_frame_idx - start_frame_idx + 1 + if start_frame_idx >= self.length: + start_frame_idx = self.length - 1 + if end_frame_idx is None or expected_len == 1: + root_pos = self.joints_pos[start_frame_idx, self.skeleton.root_idx].clone() + root_pos[1] = 0.0 + return to_numpy(root_pos) + else: + if end_frame_idx >= self.length: + end_frame_idx = self.length = 1 + root_pos = self.joints_pos[start_frame_idx : end_frame_idx + 1, self.skeleton.root_idx].clone() + root_pos[:, 1] = 0.0 + if root_pos.shape[0] < expected_len: + # pad with the last root position + root_pos = torch.cat( + [ + root_pos, + root_pos[-1:].repeat(expected_len - root_pos.shape[0], 1), + ], + dim=0, + ) + return to_numpy(root_pos) + + def set_projected_root_pos_path( + self, + root_pos_path: np.ndarray | torch.Tensor, + min_frame_idx: int = None, + max_frame_idx: int = None, + ): + """Sets the projected root position path for the character motion. + + Can set only a subset of the path by providing min_frame_idx and max_frame_idx. If not provided, will set the full + path. + + Args: + root_pos_path: torch.Tensor, [T, 2] projected root positions + min_frame_idx: int, optional, minimum frame index to set the path at + max_frame_idx: int, optional, maximum frame index to set the path at + """ + if min_frame_idx is not None or max_frame_idx is not None: + assert min_frame_idx is not None and max_frame_idx is not None, ( + "min_frame_idx and max_frame_idx must be provided if setting path at specific frames" + ) + if min_frame_idx >= self.length: + # both are out of bounds + return + max_frame_idx = min(max_frame_idx, self.length - 1) + root_pos_path = root_pos_path[min_frame_idx : max_frame_idx + 1] + else: + assert root_pos_path.shape[0] == self.length + min_frame_idx = 0 + max_frame_idx = self.length - 1 + + cur_joints_pos = self.joints_pos.clone()[min_frame_idx : max_frame_idx + 1] + root_pos_tensor = to_torch(root_pos_path, device=cur_joints_pos.device, dtype=cur_joints_pos.dtype) + diff = root_pos_tensor - cur_joints_pos[:, self.skeleton.root_idx, [0, 2]] + cur_joints_pos[:, :, [0, 2]] += diff.unsqueeze(1) + for frame_idx in range(min_frame_idx, max_frame_idx + 1): + rel_idx = frame_idx - min_frame_idx + self.update_pose_at_frame( + frame_idx, + joints_pos=cur_joints_pos[rel_idx], + joints_rot=self.joints_rot[frame_idx], + joints_local_rot=self.joints_local_rot[frame_idx], + ) + # update immediately to show changes + self.set_frame(self.cur_frame_idx) + + def get_joints_pos(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray: + """If requested frames are out of range, simply pads with the last frame to get expected + length.""" + if end_frame_idx is None: + expected_len = 1 + else: + expected_len = end_frame_idx - start_frame_idx + 1 + if start_frame_idx >= self.length: + start_frame_idx = self.length - 1 + if end_frame_idx is None or expected_len == 1: + return to_numpy(self.joints_pos[start_frame_idx].clone()) + else: + if end_frame_idx >= self.length: + end_frame_idx = self.length - 1 + return_joints_pos = self.joints_pos[start_frame_idx : end_frame_idx + 1].clone() + if return_joints_pos.shape[0] < expected_len: + # pad with the last pose + return_joints_pos = torch.cat( + [ + return_joints_pos, + return_joints_pos[-1:].repeat(expected_len - return_joints_pos.shape[0], 1, 1), + ], + dim=0, + ) + return to_numpy(return_joints_pos) + + def get_joints_rot(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray: + """If requested frames are out of range, simply pads with the last frame to get expected + length.""" + if end_frame_idx is None: + expected_len = 1 + else: + expected_len = end_frame_idx - start_frame_idx + 1 + if start_frame_idx >= self.length: + start_frame_idx = self.length - 1 + if end_frame_idx is None or expected_len == 1: + return to_numpy(self.joints_rot[start_frame_idx].clone()) + else: + if end_frame_idx >= self.length: + end_frame_idx = self.length - 1 + return_joints_rot = self.joints_rot[start_frame_idx : end_frame_idx + 1].clone() + if return_joints_rot.shape[0] < expected_len: + # pad with the last pose + return_joints_rot = torch.cat( + [ + return_joints_rot, + return_joints_rot[-1:].repeat(expected_len - return_joints_rot.shape[0], 1, 1, 1), + ], + dim=0, + ) + return to_numpy(return_joints_rot) + + def get_current_joints_pos(self) -> torch.Tensor: + return self.joints_pos[self.cur_frame_idx].clone() + + def get_current_joints_rot(self) -> torch.Tensor: + return self.joints_rot[self.cur_frame_idx].clone() + + def add_root_translation_gizmo(self, constraints: dict): + """Create and initialize gizmo to control the root translation.""" + # TODO: could also allow rotation around y-axis + self.root_translation_gizmo = self.server.scene.add_transform_controls( + f"/{self.name}/gizmo_root_translation", + scale=0.5, + line_width=2.5, + active_axes=(True, False, True), # only allow translation on xz plane + disable_axes=False, + disable_sliders=False, + disable_rotations=True, + depth_test=False, # render even when occluded + ) + init_position = self.get_current_projected_root_pos() + self.root_translation_gizmo.position = init_position + + @self.root_translation_gizmo.on_update + def _(_): + self.updating_root_translation_gizmo = True + # translate to gizmo position + new_root_pos = to_torch( + self.root_translation_gizmo.position, + device=self.joints_pos.device, + ).to(dtype=self.joints_pos.dtype) + cur_joints_pos = self.joints_pos[self.cur_frame_idx].clone() + root_diff = new_root_pos - cur_joints_pos[self.skeleton.root_idx] + root_diff[1] = 0.0 # don't change height + cur_joints_pos += root_diff[None] + self.update_pose_at_frame( + self.cur_frame_idx, + joints_pos=cur_joints_pos, + joints_rot=self.joints_rot[self.cur_frame_idx], + joints_local_rot=self.joints_local_rot[self.cur_frame_idx], + ) + + self.updating_root_translation_gizmo = False + # update immediately to show user changes + self.set_frame(self.cur_frame_idx) + # update the 2D waypoint constraints as well if there is one + if "2D Root" in constraints: + root_2d_contraints = constraints["2D Root"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in root_2d_contraints.keyframes: + for keyframe_id in root_2d_contraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + root_2d_contraints.add_keyframe( + keyframe_id, + frame_idx, + root_pos=new_root_pos, + exists_ok=True, + ) + if "Full-Body" in constraints: + full_body_constraints = constraints["Full-Body"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in full_body_constraints.keyframes: + current_dict = full_body_constraints.keyframes[frame_idx] + for keyframe_id in full_body_constraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + full_body_constraints.add_keyframe( + keyframe_id, + frame_idx, + joints_pos=cur_joints_pos, + joints_rot=current_dict["joints_rot"], + exists_ok=True, + ) + if "End-Effectors" in constraints: + end_effector_constraints = constraints["End-Effectors"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in end_effector_constraints.keyframes: + current_dict = end_effector_constraints.keyframes[frame_idx] + for keyframe_id, _ in end_effector_constraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + end_effector_constraints.add_keyframe( + keyframe_id, + frame_idx, + joints_pos=cur_joints_pos, + joints_rot=current_dict["joints_rot"], + joint_names=current_dict["joint_names"], + end_effector_type=current_dict["end_effector_type"], + exists_ok=True, + ) + + def add_joint_gizmos(self, constraints: dict): + self.joint_gizmos = [] + + joint_axis_indices = None + hidden_gizmo_joints = None + if isinstance(self.skeleton, G1Skeleton34): + joint_axis_indices = _get_g1_joint_axis_indices() + hidden_gizmo_joints = set( + self.skeleton.left_hand_joint_names + + self.skeleton.right_hand_joint_names + + self.skeleton.left_foot_joint_names + + self.skeleton.right_foot_joint_names + ) + elif isinstance(self.skeleton, CoreSkeleton27): + hidden_gizmo_joints = { + "RightHandThumb1", + "RightHandEnd", + "LeftHandThumb1", + "LeftHandEnd", + } + + joints_wxyzs = tf.SO3.from_matrix(self.joints_local_rot[self.cur_frame_idx].cpu().numpy()).wxyz + for joint_idx in range(self.skeleton.nbjoints): + disable_axes = True # by default, only rotation controls + disable_sliders = True + if joint_idx == self.skeleton.root_idx: + disable_axes = False # allow translation for root + disable_sliders = False + active_axes = (True, True, True) + if joint_axis_indices is not None: + joint_name = self.skeleton.bone_order_names[joint_idx] + axis_idx = joint_axis_indices.get(joint_name) + if axis_idx is not None: + # PivotControls shows rotation handles when a plane is active. + # To allow rotation about one axis, enable the other two axes. + active_axes = ( + axis_idx != 0, + axis_idx != 1, + axis_idx != 2, + ) + joint_visible = True + if hidden_gizmo_joints is not None: + joint_name = self.skeleton.bone_order_names[joint_idx] + joint_visible = joint_name not in hidden_gizmo_joints + cur_joint_gizmo = self.server.scene.add_transform_controls( + f"/{self.name}/gizmo_joint_{joint_idx}", + scale=0.075, + line_width=4.0, + active_axes=active_axes, + disable_axes=disable_axes, + disable_sliders=disable_sliders, + disable_rotations=False, + depth_test=False, # render even when occluded + position=self.joints_pos[self.cur_frame_idx, joint_idx].cpu().numpy(), + wxyz=joints_wxyzs[joint_idx], + visible=joint_visible, + ) + self.joint_gizmos.append(cur_joint_gizmo) + + def set_callback_in_closure(i: int) -> None: + @cur_joint_gizmo.on_update + def _(_) -> None: + self.updating_joint_gizmos = True + new_local_joint_rots = self.joints_local_rot[self.cur_frame_idx].clone() + new_local_rot = tf.SO3(self.joint_gizmos[i].wxyz) + new_local_rot_mat_np = new_local_rot.as_matrix() + if joint_axis_indices is not None: + joint_name = self.skeleton.bone_order_names[i] + axis_idx = joint_axis_indices.get(joint_name) + if axis_idx is not None: + rotvec = new_local_rot.log() + axis = np.zeros(3, dtype=np.float64) + axis[axis_idx] = 1.0 + new_local_rot_mat_np = tf.SO3.exp(rotvec[axis_idx] * axis).as_matrix() + new_local_rot_mat = torch.tensor(new_local_rot_mat_np).to(new_local_joint_rots.device) + new_local_joint_rots[i] = new_local_rot_mat + + self.update_pose_at_frame( + self.cur_frame_idx, + joints_local_rot=new_local_joint_rots, + ) + + # handle root translation separately + cur_joints_pos = self.joints_pos[self.cur_frame_idx].clone() + if i == self.skeleton.root_idx: + new_root_pos = to_torch( + self.joint_gizmos[i].position, + device=self.joints_pos.device, + ).to(dtype=self.joints_pos.dtype) + root_diff = new_root_pos - self.joints_pos[self.cur_frame_idx, i] + if torch.norm(root_diff) > 1e-3: + # the root translation has been changed + # translate to gizmo position + cur_joints_pos += root_diff[None] + self.update_pose_at_frame( + self.cur_frame_idx, + joints_pos=cur_joints_pos, + joints_rot=self.joints_rot[self.cur_frame_idx], + joints_local_rot=self.joints_local_rot[self.cur_frame_idx], + ) + + self.updating_joint_gizmos = False + + # update immediately to show user changes + self.set_frame(self.cur_frame_idx) + + if i == self.skeleton.root_idx: + # update the 2D waypoint constraints as well if there is one + if "2D Root" in constraints: + root_2d_contraints = constraints["2D Root"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in root_2d_contraints.keyframes: + new_root_pos[1] = 0.0 # force y to 0 + for keyframe_id in root_2d_contraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + root_2d_contraints.add_keyframe( + keyframe_id, + frame_idx, + root_pos=new_root_pos, + exists_ok=True, + ) + + if "Full-Body" in constraints: + full_body_constraints = constraints["Full-Body"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in full_body_constraints.keyframes: + for keyframe_id in full_body_constraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + full_body_constraints.add_keyframe( + keyframe_id, + frame_idx, + joints_pos=self.joints_pos[frame_idx], + joints_rot=self.joints_rot[frame_idx], + exists_ok=True, + ) + if "End-Effectors" in constraints: + end_effector_constraints = constraints["End-Effectors"] + # if there is a constraint at that frame, we want to update it + frame_idx = self.cur_frame_idx + if frame_idx in end_effector_constraints.keyframes: + current_dict = end_effector_constraints.keyframes[frame_idx] + for keyframe_id, _ in end_effector_constraints.frame2keyid[frame_idx]: + # add will modify the existing constraint + end_effector_constraints.add_keyframe( + keyframe_id, + frame_idx, + joints_pos=self.joints_pos[frame_idx], + joints_rot=self.joints_rot[frame_idx], + joint_names=current_dict["joint_names"], + end_effector_type=current_dict["end_effector_type"], + exists_ok=True, + ) + + set_callback_in_closure(joint_idx) + + def clear_all_gizmos(self): + self.updating_root_translation_gizmo = True + self.updating_joint_gizmos = True + if self.root_translation_gizmo is not None: + self.server.scene.remove_by_name(self.root_translation_gizmo.name) + self.root_translation_gizmo = None + if self.joint_gizmos is not None: + for joint_gizmo in self.joint_gizmos: + self.server.scene.remove_by_name(joint_gizmo.name) + self.joint_gizmos = None + self.updating_root_translation_gizmo = False + self.updating_joint_gizmos = False + + +# +# Constraint classes +# + + +class ConstraintSet: + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + display_name: Optional[str] = None, + ): + self.name = name + self.server = server + self.skeleton = skeleton + self.display_name = display_name if display_name is not None else name + + self.keyframes = dict() # frame_idx -> poses + self.frame2keyid = dict() # frame_idx -> list of keyframe ids at this frame + self.scene_elements = dict() # frame_idx -> meshes, labels, etc. + self.interval_labels = dict() # (start_frame_idx, end_frame_idx) -> interval_label + self.labels_visible = True + + def set_label_visibility(self, visible: bool) -> None: + """Show or hide constraint labels without deleting them.""" + self.labels_visible = visible + for scene_data in self.scene_elements.values(): + label = scene_data.get("label") + if label is not None: + label.visible = visible + for interval_label in self.interval_labels.values(): + interval_label.visible = visible + + def add_keyframe(self, keyframe_id: str, frame_idx: int, pose_data: torch.Tensor): + """Adds a single keyframe at the given frame with the given pose data. + + Args: + keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx. + frame_idx: int, frame index to add the keyframe at + pose_data: torch.Tensor, e.g. full-body pose, EE pose, 2D root pose, etc. + """ + raise NotImplementedError("Subclasses must implement this method") + + def add_interval( + self, + interval_id: str, + start_frame_idx: int, + end_frame_idx: int, + pose_seq_data: torch.Tensor, + ): + """Adds a keyframe interval between the given start and end frames with the given pose data. + + Args: + interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx. + start_frame_idx: int, start frame index of the interval + end_frame_idx: int, end frame index of the interval + pose_seq_data: torch.Tensor, data for constrained interval, e.g. full-body poses, EE poses, 2D root poses, etc. + """ + raise NotImplementedError("Subclasses must implement this method") + + def _add_interval_label(self, start_frame_idx: int, end_frame_idx: int): + """ + Adds an interval label between the given start and end frames + Args: + start_frame_idx: int, start frame index of the interval + end_frame_idx: int, end frame index of the interval + """ + mid = int((start_frame_idx + end_frame_idx) / 2) + interval_label_pos = self._get_label_pos(mid) + interval_label = self.server.scene.add_label( + name=f"/{self.name}/interval_label_{start_frame_idx}_{end_frame_idx}", + text=f"{self.display_name} @ [{start_frame_idx}, {end_frame_idx}]", + position=interval_label_pos, + font_size_mode="screen", + font_screen_scale=0.7, + anchor="center-center", + ) + interval_label.visible = self.labels_visible + self.interval_labels[(start_frame_idx, end_frame_idx)] = interval_label + + def remove_keyframe(self, keyframe_id: str, frame_idx: int): + """ + Removes a keyframe at the given frame + Args: + keyframe_id: str, id for the keyframe to remove + frame_idx: int, frame index to remove the keyframe at + """ + raise NotImplementedError("Subclasses must implement this method") + + def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int): + """ + Removes an interval between the given start and end frames + Args: + interval_id: str, id for the interval to remove + start_frame_idx: int, start frame index of the interval + end_frame_idx: int, end frame index of the interval + """ + raise NotImplementedError("Subclasses must implement this method") + + def _get_label_pos(self, frame_idx: int): + """ + Returns the position of where to place the displayed label for the given frame index + Args: + frame_idx: int, frame index to get the label position for + """ + raise NotImplementedError("Subclasses must implement this method") + + def _remove_interval_and_update_label(self, interval_id: str, start_frame_idx: int, end_frame_idx: int): + """ + Removes an interval between the given start and end frames and updates the label + Args: + start_frame_idx: int, start frame index of the interval + end_frame_idx: int, end frame index of the interval + """ + for frame_idx in range(start_frame_idx, end_frame_idx + 1): + self.remove_keyframe(interval_id, frame_idx) + + # Update interval labels that overlap with the removed range + intervals_to_update = [] + for (interval_start, interval_end), label in list(self.interval_labels.items()): + # Check if intervals overlap + if interval_start <= end_frame_idx and interval_end >= start_frame_idx: + intervals_to_update.append((interval_start, interval_end, label)) + + for interval_start, interval_end, label in intervals_to_update: + # Remove old label from scene and dict + self.server.scene.remove_by_name(label.name) + del self.interval_labels[(interval_start, interval_end)] + + new_start, new_end = update_interval(interval_start, interval_end, start_frame_idx, end_frame_idx) + + if new_start is None or new_end is None: + continue + + # Create updated label with new range + if new_start <= new_end: + # Position label at midpoint - these keyframes are guaranteed to exist + # since the new range is outside the removal range + mid_frame = (new_start + new_end) // 2 + label_pos = self._get_label_pos(mid_frame) + new_label = self.server.scene.add_label( + name=f"/{self.name}/interval_label_{new_start}_{new_end}", + text=f"{self.display_name} @ [{new_start}, {new_end}]", + position=label_pos, + font_size_mode="screen", + font_screen_scale=0.7, + anchor="center-center", + ) + new_label.visible = self.labels_visible + self.interval_labels[(new_start, new_end)] = new_label + + def get_constraint_info(self, device: Optional[str] = None): + """Returns constraint information for generation (torch) or UI (numpy).""" + raise NotImplementedError("Subclasses must implement this method") + + def get_frame_idx(self): + """Returns all constrained frame indices in the set.""" + return [frame_idx for frame_idx in list(self.keyframes.keys())] + + def set_keyframe_visibility(self, keyframe_idx: int, visible: bool): + """Sets the visibility of scene elements at the given keyframe index. + + Args: + keyframe_idx: int, keyframe index to set visibility for + visible: bool, whether to make the elements visible + """ + raise NotImplementedError("Subclasses must implement this method") + + def clear(self, frame_idx: Optional[int] = None): + """ + Clears all keyframes and intervals from the constraint set + Args: + frame_idx: int, sing frame index to clear if given + """ + raise NotImplementedError("Subclasses must implement this method") + + +def build_constraint_set_table_markdown(constraint_list: List[ConstraintSet]): + markdown = "| Track | Frame Num |\n" + markdown += "|------|----------|\n" + + # Sort constraints by frame_idx + for constraint in constraint_list: + frame_info = constraint.get_frame_idx() + if len(frame_info) > 0: + frame_info = ", ".join([str(frame) for frame in sorted(frame_info)]) + else: + frame_info = "-" + markdown += f"| {constraint.display_name} | {frame_info} |\n" + + return markdown + + +def update_interval(interval_start, interval_end, start_frame_idx, end_frame_idx): + """Updates an interval after removing the range from start_frame_idx to end_frame_idx.""" + # Calculate new range after removing [start_frame_idx, end_frame_idx] + # Case 1: Removal fully contains the interval -> delete entirely + if start_frame_idx <= interval_start and end_frame_idx >= interval_end: + return None, None # Already removed, don't recreate + # Case 2: Removal is at the start of interval -> shrink from start + elif start_frame_idx <= interval_start and end_frame_idx < interval_end: + new_start = end_frame_idx + 1 + new_end = interval_end + # Case 3: Removal is at the end of interval -> shrink from end + elif start_frame_idx > interval_start and end_frame_idx >= interval_end: + new_start = interval_start + new_end = start_frame_idx - 1 + # Case 4: Removal is in the middle -> keep the larger portion + else: # start_frame_idx > interval_start and end_frame_idx < interval_end + left_size = start_frame_idx - interval_start + right_size = interval_end - end_frame_idx + if left_size >= right_size: + new_start = interval_start + new_end = start_frame_idx - 1 + else: + new_start = end_frame_idx + 1 + new_end = interval_end + return new_start, new_end + + +class FullbodyKeyframeSet(ConstraintSet): + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + display_name: Optional[str] = None, + ): + super().__init__(name, server, skeleton, display_name=display_name) + + def add_keyframe( + self, + keyframe_id: str, + frame_idx: int, + joints_pos: torch.Tensor | np.ndarray, + joints_rot: torch.Tensor | np.ndarray, + viz_label: bool = True, + exists_ok: bool = False, + ): + """Adds a single full-body keyframe at the given frame or updates the existing one at this + frame. Note if a keyframe already exists at this frame, it will be updated to the given + pose. + + Args: + keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx. + frame_idx: int, frame index to add the keyframe at + joints_pos: torch.Tensor, [J, 3] joints positions to add the keyframe at + """ + # create/update scene elements + if frame_idx in self.keyframes: + skeleton_mesh = self.scene_elements[frame_idx]["skeleton_mesh"] + skeleton_mesh.set_pose(to_torch(joints_pos)) + if viz_label and "label" in self.scene_elements[frame_idx]: + label = self.scene_elements[frame_idx]["label"] + label.position = to_numpy(joints_pos)[self.skeleton.root_idx] + label.visible = self.labels_visible + else: + # create skeleton to visualize the full-body constraint + skeleton_mesh = SkeletonMesh( + f"/{self.name}/skeleton_{frame_idx}", + self.server, + self.skeleton, + joint_color=(255, 235, 0), + bone_color=(255, 0, 0), + starting_joints_pos=to_torch(joints_pos), + ) + self.scene_elements[frame_idx] = { + "skeleton_mesh": skeleton_mesh, + } + if viz_label: + label = self.server.scene.add_label( + name=f"/{self.name}/label_{frame_idx}", + text=f"{self.display_name} @ {frame_idx}", + position=to_numpy(joints_pos)[self.skeleton.root_idx], + font_size_mode="screen", + font_screen_scale=0.7, + anchor="center-center", + ) + label.visible = self.labels_visible + self.scene_elements[frame_idx]["label"] = label + + # set/update data + self.keyframes[frame_idx] = { + "joints_pos": to_numpy(joints_pos), + "joints_rot": to_numpy(joints_rot), + } + + if frame_idx not in self.frame2keyid: + self.frame2keyid[frame_idx] = [] + + if keyframe_id in self.frame2keyid[frame_idx]: + if not exists_ok: + raise AssertionError("keyframe_id already exists in this frame!") + else: + self.frame2keyid[frame_idx].append(keyframe_id) + + def add_interval( + self, + interval_id: str, + start_frame_idx: int, + end_frame_idx: int, + joints_pos: torch.Tensor, + joints_rot: torch.Tensor, + ): + """Adds a full-body keyframe interval between the given start and end frames. + + Args: + start_frame_idx: int, start frame index of the interval + end_frame_idx: int, end frame index of the interval + joints_pos: torch.Tensor, [T, J, 3] joints positions within the interval + """ + assert joints_pos.shape[0] == end_frame_idx - start_frame_idx + 1 + for frame_idx in range(start_frame_idx, end_frame_idx + 1): + rel_idx = frame_idx - start_frame_idx + self.add_keyframe( + interval_id, + frame_idx, + joints_pos[rel_idx], + joints_rot[rel_idx], + viz_label=False, + ) + + # add separate interval label + self._add_interval_label(start_frame_idx, end_frame_idx) + + def remove_keyframe(self, keyframe_id: str, frame_idx: int): + if frame_idx not in self.keyframes: + return + if keyframe_id not in self.frame2keyid[frame_idx]: + return + self.frame2keyid[frame_idx].remove(keyframe_id) + if len(self.frame2keyid[frame_idx]) == 0: + del self.frame2keyid[frame_idx] + self.clear(frame_idx) + + def _get_label_pos(self, frame_idx: int): + return self.keyframes[frame_idx]["joints_pos"][self.skeleton.root_idx] + + def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int): + self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx) + + def get_constraint_info(self, device: Optional[str] = None): + all_joints_pos = [] + all_joints_rot = [] + for v in self.keyframes.values(): + joints_pos = to_torch(v["joints_pos"], device=device) + joints_rot = to_torch(v["joints_rot"], device=device) + if len(joints_pos.shape) == 2: + all_joints_pos.append(joints_pos[None]) + else: + all_joints_pos.append(joints_pos) + if len(joints_rot.shape) == 3: + all_joints_rot.append(joints_rot[None]) + else: + all_joints_rot.append(joints_rot) + + all_joints_pos = torch.cat(all_joints_pos, dim=0) if len(all_joints_pos) > 0 else None + all_joints_rot = torch.cat(all_joints_rot, dim=0) if len(all_joints_rot) > 0 else None + + return { + "frame_idx": self.get_frame_idx(), + "joints_pos": all_joints_pos, + "joints_rot": all_joints_rot, + } + + def set_keyframe_visibility(self, keyframe_idx: int, visible: bool): + """Sets the visibility of scene elements at the given keyframe index.""" + if keyframe_idx not in self.scene_elements: + return + + scene_elements = self.scene_elements[keyframe_idx] + if "skeleton_mesh" in scene_elements: + skeleton_mesh = scene_elements["skeleton_mesh"] + if hasattr(skeleton_mesh, "joints_batched_mesh"): + skeleton_mesh.joints_batched_mesh.visible = visible + if hasattr(skeleton_mesh, "bones_batched_mesh"): + skeleton_mesh.bones_batched_mesh.visible = visible + + if "label" in scene_elements: + label = scene_elements["label"] + if hasattr(label, "visible"): + label.visible = visible + + def clear(self, frame_idx: Optional[int] = None): + frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx] + for fidx in frame_idx_list: + self.scene_elements[fidx]["skeleton_mesh"].clear() + if "ee_rotation_axes" in self.scene_elements[fidx]: + self.server.scene.remove_by_name(self.scene_elements[fidx]["ee_rotation_axes"].name) + if "label" in self.scene_elements[fidx]: + self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name) + + self.keyframes.pop(fidx) + self.scene_elements.pop(fidx) + + if frame_idx is None: + # clear all interval labels if clearing all keyframes + for interval_label in list(self.interval_labels.values()): + self.server.scene.remove_by_name(interval_label.name) + self.interval_labels.clear() + + +class EEJointsKeyframeSet(ConstraintSet): + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + display_name: Optional[str] = None, + ): + super().__init__(name, server, skeleton, display_name=display_name) + + # frame_idx -> list of (keyframe_id, joint_names) at this frame + self.frame2keyid = dict() + + def create_scene_elements( + self, + frame_idx: int, + joints_pos: torch.Tensor | np.ndarray, + joints_rot: Optional[torch.Tensor | np.ndarray], + joint_names: List[str], + viz_label: bool = True, + ): + # create skeleton to visualize the full-body constraint + ee_joint_indices = [] + ee_gizmo_indices = [] + constrained_bone_idx = [] + for joint_name in joint_names: + if joint_name == "Hips": + continue + elif joint_name in ["LeftHand", "RightHand", "LeftFoot", "RightFoot"]: + expanded_joint_names = { + "LeftHand": self.skeleton.left_hand_joint_names, + "RightHand": self.skeleton.right_hand_joint_names, + "LeftFoot": self.skeleton.left_foot_joint_names, + "RightFoot": self.skeleton.right_foot_joint_names, + }[joint_name] + ee_joint_indices.extend([self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names]) + if len(expanded_joint_names) > 1: + ee_gizmo_indices.extend( + [self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names[:-1]] + ) + constrained_bone_idx.extend( + [self.skeleton.bone_order_names_index[joint] - 1 for joint in expanded_joint_names[1:]] + ) + else: + raise ValueError(f"Invalid joint name: {joint_name}") + + # de-duplicate while preserving order + ee_joint_indices = list(dict.fromkeys(ee_joint_indices)) + ee_gizmo_indices = list(dict.fromkeys(ee_gizmo_indices)) + constrained_bone_idx = list(dict.fromkeys(constrained_bone_idx)) + + constrained_idx = np.array([self.skeleton.root_idx] + ee_joint_indices, dtype=np.intp) + constrained_bone_idx = np.array(constrained_bone_idx, dtype=np.intp) + + # create skeleton to visualize the full-body constraint + joint_color = np.full((self.skeleton.nbjoints, 3), (220, 220, 220)) + bone_color = np.full((self.skeleton.nbjoints - 1, 3), (220, 220, 220)) + # color constrained joints and bones red + joint_color[constrained_idx] = (255, 0, 0) + if len(constrained_bone_idx) > 0: + bone_color[constrained_bone_idx] = (255, 0, 0) + skeleton_mesh = SkeletonMesh( + f"/{self.name}/skeleton_{frame_idx}", + self.server, + self.skeleton, + joint_color=joint_color, + bone_color=bone_color, + starting_joints_pos=to_torch(joints_pos), + ) + + self.scene_elements[frame_idx] = { + "skeleton_mesh": skeleton_mesh, + } + joints_pos_np = to_numpy(joints_pos) + joints_rot_np = to_numpy(joints_rot) if joints_rot is not None else None + if joints_rot_np is not None and len(ee_gizmo_indices) > 0: + ee_axes = self.server.scene.add_batched_axes( + f"/{self.name}/ee_rot_axes_{frame_idx}", + batched_wxyzs=tf.SO3.from_matrix(joints_rot_np[ee_gizmo_indices]).wxyz, + batched_positions=joints_pos_np[ee_gizmo_indices], + axes_length=0.07, + axes_radius=0.007, + ) + self.scene_elements[frame_idx]["ee_rotation_axes"] = ee_axes + if viz_label: + label = self.server.scene.add_label( + name=f"/{self.name}/label_{frame_idx}", + text=f"{self.display_name} @ {frame_idx}", + position=joints_pos_np[self.skeleton.root_idx] + np.array([0.0, 0.05, 0.0]), + font_size_mode="screen", + font_screen_scale=0.7, + anchor="bottom-center", + ) + label.visible = self.labels_visible + self.scene_elements[frame_idx]["label"] = label + + def add_keyframe( + self, + keyframe_id: str, + frame_idx: int, + joints_pos: torch.Tensor | np.ndarray, + joints_rot: torch.Tensor | np.ndarray, + joint_names: List[str], + end_effector_type: str, + viz_label: bool = True, + exists_ok: bool = False, + ): + """Adds a single EE keyframe at the given frame or updates the existing one at this frame. + + Args: + keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx. + frame_idx: int, frame index to add the keyframe at + joints_pos: torch.Tensor, [J, 3] joints positions to add the keyframe at + joints_rot: torch.Tensor, [J, 3, 3] joints rotation matrices to add the keyframe at + joint_names: List[str], names of the joints to add the keyframe at + """ + need_create_viz = True + joint_names_input = joint_names + + if not isinstance(end_effector_type, set): + end_effector_type = set([end_effector_type]) + + # create/update scene elements + if frame_idx in self.keyframes: + if joint_names != self.keyframes[frame_idx]["joint_names"]: + # merge together with existing constraint if needed + joint_names = set(joint_names) + joint_names.update(set(self.keyframes[frame_idx]["joint_names"])) + joint_names = list(joint_names) + end_effector_type.update(self.keyframes[frame_idx]["end_effector_type"]) + # need to re-create viz elements + self.clear(frame_idx) + else: + need_create_viz = False + # overwrite the pose with the latest one + skeleton_mesh = self.scene_elements[frame_idx]["skeleton_mesh"] + skeleton_mesh.set_pose(to_torch(joints_pos)) + if "ee_rotation_axes" in self.scene_elements[frame_idx]: + ee_gizmo_indices = [] + for joint_name in joint_names: + if joint_name == "Hips": + continue + elif joint_name in [ + "LeftHand", + "RightHand", + "LeftFoot", + "RightFoot", + ]: + expanded_joint_names = { + "LeftHand": self.skeleton.left_hand_joint_names, + "RightHand": self.skeleton.right_hand_joint_names, + "LeftFoot": self.skeleton.left_foot_joint_names, + "RightFoot": self.skeleton.right_foot_joint_names, + }[joint_name] + if len(expanded_joint_names) > 1: + ee_gizmo_indices.extend( + [self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names[:-1]] + ) + else: + raise ValueError(f"Invalid joint name: {joint_name}") + ee_gizmo_indices = list(dict.fromkeys(ee_gizmo_indices)) + if len(ee_gizmo_indices) > 0: + ee_axes = self.scene_elements[frame_idx]["ee_rotation_axes"] + joints_pos_np = to_numpy(joints_pos) + joints_rot_np = to_numpy(joints_rot) + ee_axes.batched_positions = joints_pos_np[ee_gizmo_indices] + ee_axes.batched_wxyzs = tf.SO3.from_matrix(joints_rot_np[ee_gizmo_indices]).wxyz + if viz_label and "label" in self.scene_elements[frame_idx]: + label = self.scene_elements[frame_idx]["label"] + label.position = to_numpy(joints_pos)[self.skeleton.root_idx] + label.visible = self.labels_visible + + if need_create_viz: + self.create_scene_elements(frame_idx, joints_pos, joints_rot, joint_names, viz_label=viz_label) + + # set/update data + self.keyframes[frame_idx] = { + "joints_pos": to_numpy(joints_pos), + "joints_rot": to_numpy(joints_rot), + "joint_names": joint_names, + "end_effector_type": end_effector_type, + } + + if frame_idx not in self.frame2keyid: + self.frame2keyid[frame_idx] = [] + + known_keyframe_ids = {k: idx for idx, (k, _) in enumerate(self.frame2keyid[frame_idx])} + + if keyframe_id in known_keyframe_ids.keys(): + if not exists_ok: + raise AssertionError("keyframe_id already exists in this frame!") + idx = known_keyframe_ids[keyframe_id] + # override previous exisiting keyframe + self.frame2keyid[frame_idx][idx] = (keyframe_id, joint_names_input) + else: + # track which subset of joints are constrained by this keyframe_id + self.frame2keyid[frame_idx].append((keyframe_id, joint_names_input)) + + def add_interval( + self, + interval_id: str, + start_frame_idx: int, + end_frame_idx: int, + joints_pos: torch.Tensor | np.ndarray, + joints_rot: torch.Tensor | np.ndarray, + joint_names: List[str], + end_effector_type: str, + ): + """Adds an interval of EE keyframes at the given frame or updates the existing one at this + frame. + + Args: + interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx. + start_frame_idx: int, start frame index to add the interval at + end_frame_idx: int, end frame index to add the interval at + joints_pos: torch.Tensor, [T, J, 3] joints positions to add the interval at + joints_rot: torch.Tensor, [T, J, 3, 3] joints rotation matrices to add the interval at + joint_names: List[str], names of the joints to add for the entire interval + """ + num_frames = end_frame_idx - start_frame_idx + 1 + joints_pos_np = to_numpy(joints_pos) + joints_rot_np = to_numpy(joints_rot) + assert joints_pos_np.shape[0] == num_frames + assert joints_rot_np.shape[0] == num_frames + + for frame_idx in range(start_frame_idx, end_frame_idx + 1): + rel_idx = frame_idx - start_frame_idx + self.add_keyframe( + interval_id, + frame_idx, + joints_pos_np[rel_idx], + joints_rot_np[rel_idx], + joint_names, + end_effector_type, + viz_label=False, + ) + self._add_interval_label(start_frame_idx, end_frame_idx) + + def remove_keyframe(self, keyframe_id: str, frame_idx: int): + """Removes a keyframe at the given frame or updates the existing one at this frame by + removing the specified joints. + + Args: + keyframe_id: str, id for the keyframe to remove. This determines which joints to remove. + frame_idx: int, frame index to remove the keyframe at + """ + if frame_idx not in self.keyframes: + return + + remaining_joint_names = set() + delete_idx = None + for i, (keyid, joint_names) in enumerate(self.frame2keyid[frame_idx]): + if keyid == keyframe_id: + delete_idx = i + else: + remaining_joint_names.update(joint_names) + if delete_idx is None: + # this keyframe_id is not in the specified frame + return + + self.frame2keyid[frame_idx].pop(delete_idx) + if len(remaining_joint_names) == 0: + # no more keyframes in this frame, clear the frame + del self.frame2keyid[frame_idx] + self.clear(frame_idx) + return + + # only deleting part of keyframe (potentially some subset of joints) + # delete the old visualization and add a new one with the updated joint set + new_joint_names = list(remaining_joint_names) + self.clear(frame_idx, scene_elements_only=True) + joints_pos = self.keyframes[frame_idx]["joints_pos"] + joints_rot = self.keyframes[frame_idx]["joints_rot"] + self.create_scene_elements(frame_idx, joints_pos, joints_rot, new_joint_names) + self.keyframes[frame_idx]["joint_names"] = new_joint_names + + def _get_label_pos(self, frame_idx: int): + return self.keyframes[frame_idx]["joints_pos"][self.skeleton.root_idx] + + def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int): + self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx) + + def get_constraint_info(self, device: Optional[str] = None): + all_joints_pos = [] + all_joints_rot = [] + all_joints_names = [] + all_end_effector_type = [] + for v in self.keyframes.values(): + joints_pos = to_torch(v["joints_pos"], device=device) + joints_rot = to_torch(v["joints_rot"], device=device) + if len(joints_pos.shape) == 2: + all_joints_pos.append(joints_pos[None]) + else: + all_joints_pos.append(joints_pos) + if len(joints_rot.shape) == 3: + all_joints_rot.append(joints_rot[None]) + else: + all_joints_rot.append(joints_rot) + all_joints_names.append(v["joint_names"]) + all_end_effector_type.append(v["end_effector_type"]) + + all_joints_pos = torch.cat(all_joints_pos, dim=0) if len(all_joints_pos) > 0 else None + all_joints_rot = torch.cat(all_joints_rot, dim=0) if len(all_joints_rot) > 0 else None + + return { + "frame_idx": self.get_frame_idx(), + "joints_pos": all_joints_pos, + "joints_rot": all_joints_rot, + "joint_names": all_joints_names, + "end_effector_type": all_end_effector_type, + } + + def set_keyframe_visibility(self, keyframe_idx: int, visible: bool, show_rotation_axes: bool = True): + """Sets the visibility of scene elements at the given keyframe index.""" + if keyframe_idx not in self.scene_elements: + return + + scene_elements = self.scene_elements[keyframe_idx] + if "skeleton_mesh" in scene_elements: + skeleton_mesh = scene_elements["skeleton_mesh"] + if hasattr(skeleton_mesh, "joints_batched_mesh"): + skeleton_mesh.joints_batched_mesh.visible = visible + if hasattr(skeleton_mesh, "bones_batched_mesh"): + skeleton_mesh.bones_batched_mesh.visible = visible + + if "ee_rotation_axes" in scene_elements: + ee_axes = scene_elements["ee_rotation_axes"] + if hasattr(ee_axes, "visible"): + ee_axes.visible = visible and show_rotation_axes + + if "label" in scene_elements: + label = scene_elements["label"] + if hasattr(label, "visible"): + label.visible = visible + + def clear(self, frame_idx: Optional[int] = None, scene_elements_only: bool = False): + frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx] + for fidx in frame_idx_list: + self.scene_elements[fidx]["skeleton_mesh"].clear() + if "ee_rotation_axes" in self.scene_elements[fidx]: + self.server.scene.remove_by_name(self.scene_elements[fidx]["ee_rotation_axes"].name) + if "label" in self.scene_elements[fidx]: + self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name) + self.scene_elements.pop(fidx) + if not scene_elements_only: + self.keyframes.pop(fidx) + + if frame_idx is None: + # clear all interval labels if clearing all keyframes + for interval_label in list(self.interval_labels.values()): + self.server.scene.remove_by_name(interval_label.name) + self.interval_labels.clear() + + +class RootKeyframe2DSet(ConstraintSet): + def __init__( + self, + name: str, + server: viser.ViserServer, + skeleton: SkeletonBase, + display_name: Optional[str] = None, + ): + super().__init__(name, server, skeleton, display_name=display_name) + self.dense_path = False + self.smooth_path = True + self.line_segments = None # visualization of dense path + # Cache for interpolated path + self._cached_t = None + self._cached_path3d = None + # Store root headings for each keyframe + self.root_headings = {} # frame_idx -> heading (float) + + def add_keyframe( + self, + keyframe_id: str, + frame_idx: int, + root_pos: torch.Tensor | np.ndarray, + global_root_heading: Optional[float] = None, + viz_label: bool = True, + update_path: bool = True, + exists_ok: bool = False, + add_annulus: bool = True, + ): + """Adds a single 2D root keyframe at the given frame or updates the existing one at this + frame. + + Args: + keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx. + frame_idx: int, frame index to add the keyframe at + root_pos: torch.Tensor, [3] root position to add the keyframe at, y entry (index 1) should be 0 + viz_label: bool, whether to visualize the label for the keyframe + """ + root_pos_np = to_numpy(root_pos) + + # Convert heading angle to 2D direction vector for visualization + heading_2d = None + if global_root_heading is not None: + # Heading is in radians, convert to [x, z] direction vector as [cos, sin] + heading_2d = np.array([np.cos(global_root_heading), np.sin(global_root_heading)]) + + if frame_idx in self.keyframes: + waypoint = self.scene_elements[frame_idx]["waypoint"] + waypoint.update_position(root_pos_np, heading=heading_2d) + if viz_label and "label" in self.scene_elements[frame_idx]: + label = self.scene_elements[frame_idx]["label"] + label.position = root_pos.cpu().numpy() + label.visible = self.labels_visible + else: + waypoint = WaypointMesh( + f"/{self.name}/{keyframe_id}", + self.server, + position=root_pos_np, + heading=heading_2d, + add_annulus=add_annulus, + ) + self.scene_elements[frame_idx] = { + "waypoint": waypoint, + } + if viz_label: + label = self.server.scene.add_label( + name=f"/{self.name}/label_{frame_idx}", + text=f"{self.display_name} @ {frame_idx}", + position=root_pos_np, + font_size_mode="screen", + font_screen_scale=0.7, + anchor="bottom-left", + ) + label.visible = self.labels_visible + self.scene_elements[frame_idx]["label"] = label + + # set/update data + self.keyframes[frame_idx] = root_pos_np + if global_root_heading is not None: + self.root_headings[frame_idx] = global_root_heading + if frame_idx not in self.frame2keyid: + self.frame2keyid[frame_idx] = [] + + if keyframe_id in self.frame2keyid[frame_idx]: + if not exists_ok: + raise AssertionError("keyframe_id already exists in this frame!") + else: + self.frame2keyid[frame_idx].append(keyframe_id) + + # need to update path visualization + if self.dense_path and self.line_segments is None: + # visualize dense path with line segments + self.line_segments = self.server.scene.add_line_segments( + name=f"/{self.name}/line_segments", + points=np.zeros((1, 2, 3)), + colors=(255, 0, 0), + line_width=5.0, + ) + if update_path: + self.update_line_segments() + + def add_interval( + self, + interval_id: str, + start_frame_idx: int, + end_frame_idx: int, + root_pos: torch.Tensor | np.ndarray, + add_annulus: bool = True, + ): + """Adds an interval of 2D root keyframes between the given start and end frames. + + Args: + interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx. + start_frame_idx: int, start frame index to add the interval at + end_frame_idx: int, end frame index to add the interval at + root_pos: torch.Tensor, [T, 3] root positions to add the interval at + """ + root_pos_np = to_numpy(root_pos) + assert root_pos_np.shape[0] == end_frame_idx - start_frame_idx + 1 + for frame_idx in range(start_frame_idx, end_frame_idx + 1): + rel_idx = frame_idx - start_frame_idx + self.add_keyframe( + f"{interval_id}/{frame_idx}", + frame_idx, + root_pos_np[rel_idx], + viz_label=False, + update_path=False, + add_annulus=add_annulus, + ) + self._add_interval_label(start_frame_idx, end_frame_idx) + if self.line_segments is not None: + self.update_line_segments() + + def set_smooth_path(self, smooth_path: bool): + self.smooth_path = smooth_path + if self.line_segments is not None: + self.update_line_segments() + + def set_dense_path(self, dense_path: bool): + """If dense_path is True, will make the path dense by interpolated between added keyframes. + + Args: + dense_path: bool, whether to make the path dense + """ + self.dense_path = dense_path + if self.dense_path: + # visualize dense path with line segments + self.line_segments = self.server.scene.add_line_segments( + name=f"/{self.name}/line_segments", + points=np.zeros((1, 2, 3)), + colors=(255, 0, 0), + line_width=5.0, + ) + self.update_line_segments() + else: + if self.line_segments is not None: + self.server.scene.remove_by_name(self.line_segments.name) + self.line_segments = None + + # TODO: accelerate interp1d. Note that np and torch interp gives different results, much less smooth + def interpolate_path(self, t: np.ndarray): + """Interpolates the path between the given frame indices. + + Args: + t: np.ndarray, frame indices to interpolate at + """ + cur_info = self._get_sparse_constraint_info() + frame_idx = cur_info["frame_idx"] + all_root_pos = cur_info["root_pos"] + + x = all_root_pos[:, 0] + z = all_root_pos[:, 2] + + kind = "linear" + # if self.smooth_path and len(frame_idx) >= 3: + # kind = "quadratic" + + interp_x = interp1d(frame_idx, x, kind=kind) + interp_z = interp1d(frame_idx, z, kind=kind) + + x_new = interp_x(t) + z_new = interp_z(t) + + path3d = np.stack([x_new, np.zeros_like(x_new), z_new], axis=1) + + if self.smooth_path and len(frame_idx) >= 3: + start_time = time.time() + path3d = get_smooth_root_pos(torch.from_numpy(path3d[None]))[0].numpy() + print(f"Time taken to smooth path: {time.time() - start_time} seconds") + return path3d + + def update_line_segments(self, frame_idx: int = 0): + if len(self.keyframes) < 2: + return + + t = np.array(sorted(self.get_frame_idx())) + if self.smooth_path: + # more points for smoothed curve + # t = np.linspace(t[0], t[-1], 100) + t = np.arange(max(t[0], frame_idx), t[-1] + 1) + + path3d = self.interpolate_path(t) + + # Cache the computed t and path3d for use in get_constraint_info + self._cached_t = t + self._cached_path3d = path3d + + points = np.zeros((len(path3d) - 1, 2, 3)) + points[:, 0] = path3d[:-1] + points[:, 1] = path3d[1:] + + self.line_segments.points = points + + def remove_keyframe(self, keyframe_id: str, frame_idx: int): + if frame_idx not in self.keyframes: + return + if keyframe_id not in self.frame2keyid[frame_idx]: + return + self.frame2keyid[frame_idx].remove(keyframe_id) + if len(self.frame2keyid[frame_idx]) == 0: + del self.frame2keyid[frame_idx] + # Also remove heading if it exists + if frame_idx in self.root_headings: + del self.root_headings[frame_idx] + self.clear(frame_idx) + if self.line_segments is not None: + self.update_line_segments() + + def _get_label_pos(self, frame_idx: int): + return self.keyframes[frame_idx] + + def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int): + self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx) + + def _get_sparse_constraint_info(self): + all_root_pos = [] + all_root_headings = [] + frame_indices = self.get_frame_idx() + + for frame_idx in frame_indices: + v = self.keyframes[frame_idx] + # Handle both numpy arrays and torch tensors + if isinstance(v, torch.Tensor): + if len(v.shape) == 1: + all_root_pos.append(v.unsqueeze(0).cpu().numpy()) + else: + all_root_pos.append(v.cpu().numpy()) + else: # numpy array + if len(v.shape) == 1: + all_root_pos.append(v[np.newaxis, :]) + else: + all_root_pos.append(v) + # Get heading if it exists + if frame_idx in self.root_headings: + all_root_headings.append(self.root_headings[frame_idx]) + + if len(all_root_pos) > 0: + all_root_pos = np.concatenate(all_root_pos, axis=0) + else: + all_root_pos = None + + result = { + "frame_idx": frame_indices, + "root_pos": all_root_pos, + } + + if len(all_root_headings) > 0 and len(all_root_headings) == len(frame_indices): + result["global_root_heading"] = torch.tensor(all_root_headings, dtype=torch.float32) + + return result + + def get_constraint_info(self): + if not self.dense_path or len(self.keyframes) == 0: + return self._get_sparse_constraint_info() + else: + # Use cached path if available, otherwise compute it + if self._cached_t is not None and self._cached_path3d is not None: + t = self._cached_t + path3d = self._cached_path3d + else: + frame_idx_list = self.get_frame_idx() + min_frame_idx = min(frame_idx_list) + max_frame_idx = max(frame_idx_list) + t = np.arange(min_frame_idx, max_frame_idx + 1) + path3d = self.interpolate_path(t) + + result = { + "frame_idx": t.tolist(), + "root_pos": torch.tensor(path3d, dtype=torch.float32), + } + + # Interpolate headings if any exist + if len(self.root_headings) > 0: + frame_idx_list = self.get_frame_idx() + # Check if we have headings for all keyframes + if all(fi in self.root_headings for fi in frame_idx_list): + # Interpolate headings for dense path + heading_values = [self.root_headings[fi] for fi in sorted(frame_idx_list)] + heading_frames = sorted(frame_idx_list) + + # Linear interpolation of headings + interpolated_headings = [] + for frame in t: + if frame in self.root_headings: + interpolated_headings.append(self.root_headings[frame]) + else: + # Find surrounding keyframes and interpolate + left_frame = max([f for f in heading_frames if f <= frame], default=None) + right_frame = min([f for f in heading_frames if f >= frame], default=None) + + if left_frame is not None and right_frame is not None and left_frame != right_frame: + # Linear interpolation + alpha = (frame - left_frame) / (right_frame - left_frame) + interp_heading = (1 - alpha) * self.root_headings[ + left_frame + ] + alpha * self.root_headings[right_frame] + interpolated_headings.append(interp_heading) + elif left_frame is not None: + interpolated_headings.append(self.root_headings[left_frame]) + elif right_frame is not None: + interpolated_headings.append(self.root_headings[right_frame]) + + if len(interpolated_headings) == len(t): + result["global_root_heading"] = torch.tensor(interpolated_headings, dtype=torch.float32) + + return result + + def set_keyframe_visibility(self, keyframe_idx: int, visible: bool): + """Sets the visibility of scene elements at the given keyframe index.""" + if keyframe_idx not in self.scene_elements: + return + + scene_elements = self.scene_elements[keyframe_idx] + if "waypoint" in scene_elements: + waypoint = scene_elements["waypoint"] + # WaypointMesh has sphere, annulus, arrow_base, arrow_head components + if hasattr(waypoint, "sphere"): + waypoint.sphere.visible = visible + if hasattr(waypoint, "annulus") and waypoint.annulus is not None: + waypoint.annulus.visible = visible + if hasattr(waypoint, "arrow_base") and waypoint.arrow_base is not None: + waypoint.arrow_base.visible = visible + if hasattr(waypoint, "arrow_head") and waypoint.arrow_head is not None: + waypoint.arrow_head.visible = visible + + if "label" in scene_elements: + label = scene_elements["label"] + if hasattr(label, "visible"): + label.visible = visible + + def set_interval_labels_visibility(self, frame_idx: int): + # set interval labels visibility + for interval in self.interval_labels.keys(): + interval_label = self.interval_labels[interval] + start_frame_idx, end_frame_idx = interval + visibility = frame_idx <= end_frame_idx + interval_label.visible = visibility + + def clear(self, frame_idx: Optional[int] = None): + frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx] + for fidx in frame_idx_list: + self.scene_elements[fidx]["waypoint"].clear() + if "label" in self.scene_elements[fidx]: + self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name) + + self.keyframes.pop(fidx) + self.scene_elements.pop(fidx) + # Also clear heading if it exists + if fidx in self.root_headings: + self.root_headings.pop(fidx) + + if frame_idx is None: + # clear all interval labels if clearing all keyframes + for interval_label in list(self.interval_labels.values()): + self.server.scene.remove_by_name(interval_label.name) + self.interval_labels.clear() + + # clear line segments if turning off dense path + if self.line_segments is not None: + self.server.scene.remove_by_name(self.line_segments.name) + self.line_segments = None + + # Clear all headings when clearing everything + self.root_headings.clear() + + # Invalidate cache when clearing all keyframes + self._cached_t = None + self._cached_path3d = None + else: + # Invalidate cache when clearing a specific keyframe + self._cached_t = None + self._cached_path3d = None + + +def load_example_cases(examples_base_dir): + example_dirs = os.listdir(examples_base_dir) + example_names = sorted([dir for dir in example_dirs if os.path.isdir(os.path.join(examples_base_dir, dir))]) + example_dict = {name: os.path.join(examples_base_dir, name) for name in example_names} + return example_dict diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..54656adf21118dc848d6082ddaa659d48148c9e1 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,13 @@ +hydra-core>=1.3 +omegaconf>=2.3 +pydantic>=2.0 +numpy>=1.23,<2 +scipy>=1.10 +transformers==5.8.1 +peft>=0.19 +einops>=0.7 +vector-quantize-pytorch>=1.25,<1.26 +safetensors>=0.4 +trimesh>=4.0 +tqdm>=4.0 +packaging>=21.0 diff --git a/tests/test_bvh.py b/tests/test_bvh.py new file mode 100644 index 0000000000000000000000000000000000000000..b1cdb8797f3cb4af8b07fb0f2fda032b3b3d36a6 --- /dev/null +++ b/tests/test_bvh.py @@ -0,0 +1,46 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from ardy.exports.bvh import write_bvh + + +class BvhExportTests(unittest.TestCase): + def test_writes_blender_readable_hierarchy_and_motion(self): + rotations = np.repeat(np.eye(3)[None, None], 4, axis=0) + rotations = np.repeat(rotations, 2, axis=1) + roots = np.array( + [ + [0.0, 0.0, 0.0], + [0.1, 0.0, 0.0], + [0.2, 0.0, 0.0], + [0.3, 0.0, 0.0], + ] + ) + + with tempfile.TemporaryDirectory() as folder: + output = Path(folder) / "motion.bvh" + write_bvh( + output, + joint_names=["Hips", "Spine"], + parents=[-1, 0], + rest_positions=np.array([[0.0, 0.0, 0.0], [0.0, 0.5, 0.0]]), + local_rot_mats=rotations, + root_positions=roots, + fps=20, + ) + text = output.read_text(encoding="utf-8") + + self.assertIn("ROOT Hips", text) + self.assertIn("JOINT Spine", text) + self.assertIn("Frames: 4", text) + self.assertIn("Frame Time: 0.050000000", text) + motion_lines = text.split("Frame Time: 0.050000000\n", 1)[1].strip().splitlines() + self.assertEqual(len(motion_lines), 4) + self.assertTrue(motion_lines[-1].startswith("0.3000000 0.0000000 0.0000000")) + + +if __name__ == "__main__": + unittest.main()