diff --git a/scripts/eval-t2m-posture.py b/scripts/eval-t2m-posture.py new file mode 100644 index 00000000..c35c976a --- /dev/null +++ b/scripts/eval-t2m-posture.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +# ruff: noqa: E702, E741 +"""Score a t2m ONNX model on the POSTURE metrics that gate shipping (#837). + +ONE-TIME OFFLINE dev tool — NOT shipped. + +The v6/v6.1 quality bar was tracked by hand in ~/t2m_v6/EVAL_NOTES.md. This +script recomputes those exact numbers so any two models (and the training +data itself) are directly comparable: + + spineY mean world up-component of the spine aim (v6.1 shipped: 0.99) + headY ... of the neck/head aim (v6.1 shipped: 0.90) + armY signed up-component per upper arm, WORST (v6.1 shipped: -0.98) + fwd/side foot travel ratio, locomotion only (v6.1 shipped: 2.29, + broken v6: 1.2, bar 2.0) + +Metrics are computed on the model's own canonical output (the same quantity +prep-t2m-v6.py gates the training windows on), so "model vs data" is an +apples-to-apples read. Reports best-of-N per action the way the shipped +MotionGenerator picks a candidate, plus the mean, so sample VARIANCE (the +v6 walk failure mode) is visible rather than hidden by a lucky draw. + +Usage: + python3 scripts/eval-t2m-posture.py --model ~/t2m_v62/flow/t2m.onnx \ + --vocab ~/t2m_v62/flow/t2m-vocab.json [--data ~/t2m_v62/t2m_v62.npz] \ + [--actions walk,run,jump] [--samples 16] +""" +import argparse +import importlib.util +import json +import os + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def load_module(name, fname): + spec = importlib.util.spec_from_file_location(name, os.path.join(HERE, fname)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +prep5 = load_module("prep5", "prep-t2m-v5.py") +prep6 = load_module("prep6", "prep-t2m-v6.py") + +J = 22 +D_CANON = prep5.D_CANON +LOCOMOTION = prep6.LOCOMOTION + + +def aim_y(w, roles): + """Mean world up-component of the first VALID role's canonical aim.""" + for r in roles: + return prep6.mean_dir_y(w, r) + return float("nan") + + +def torso_world_up(w): + """World up-component of the HIP->HEAD axis, via the same FK the retarget + uses. This is what catches a whole-body forward PITCH about the root. + + `spineY` (the per-joint aim's up-component) does NOT: the spine can aim + "up" within its own canonical frame while the entire body is rotated + toward horizontal. A render-verified diving sample scored spineY 0.97 + yet was pitched near-horizontal, which is why this term exists. + + Returns the mean over frames of dot(normalize(head - hip), +Y). + 1.0 = perfectly upright, 0.0 = horizontal, <0 = inverted/head-down. + """ + T = len(w) + PAR = prep6.PAR + + def pos(role): + p = np.zeros((T, 3), np.float32) + r = role + while PAR[r] >= 0: + p = p + prep6.qrot(w[:, PAR[r]], np.broadcast_to(D_CANON[r], (T, 3))) + r = PAR[r] + return p + + axis = pos(5) - pos(0) # head - hip + n = np.linalg.norm(axis, axis=-1, keepdims=True) + axis = axis / (n + 1e-12) + return float(axis[:, 1].mean()) + + +def ankle_spread(w): + """Mean world distance between the two ankles (roles 17/21), in unit bone + lengths, over the window. + + Catches the COLLAPSED-LEGS failure that the aim-based metrics miss: a + render-verified bad draw showed both legs fused into one tapering mass + while torsoUp/spineY/fwd_side all still looked healthy. A real stride + separates the ankles for most of the cycle. + """ + T = len(w) + PAR = prep6.PAR + + def pos(role): + p = np.zeros((T, 3), np.float32) + r = role + while PAR[r] >= 0: + p = p + prep6.qrot(w[:, PAR[r]], np.broadcast_to(D_CANON[r], (T, 3))) + r = PAR[r] + return p + + d = np.linalg.norm(pos(17) - pos(21), axis=-1) + return float(d.mean()) + + +def score_window(action, w): + """Posture metrics for one [T,J,4] canonical window.""" + spine = aim_y(w, (0,)) + chest = aim_y(w, (2,)) + head = aim_y(w, (4,)) + arms = [prep6.mean_dir_y(w, r) for r in (7, 11)] + out = { + "torsoUp": torso_world_up(w), + "ankleSpread": ankle_spread(w), + "spineY": spine, + "chestY": chest, + "headY": head, + "armY_worst": max(arms), # arms hang => strongly negative; worst = least negative + "armY_L": arms[0], + "armY_R": arms[1], + } + if action in LOCOMOTION: + out["fwd_side"] = prep6.foot_travel_ratio(w) + return out + + +def quat_from_motion(motion): + """MotionGenerator output [T,220] -> canonical quats [T,J,4] (x,y,z,w).""" + m = np.asarray(motion, np.float32).reshape(motion.shape[0], J, 10) + q = m[:, :, 3:7] + return q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-12) + + +def fmt(d): + keys = ["torsoUp", "ankleSpread", "spineY", "headY", "armY_worst", + "fwd_side"] + return " ".join(f"{k}={d[k]:+.3f}" for k in keys if k in d) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--vocab", default="") + ap.add_argument("--data", default="", help="npz cache — also score the DATA") + ap.add_argument("--actions", default="", help="comma list; default = all vocab") + ap.add_argument("--samples", type=int, default=16) + ap.add_argument("--seed", type=int, default=0) + a = ap.parse_args() + + vocab_path = a.vocab or os.path.join(os.path.dirname(a.model), "t2m-vocab.json") + vj = json.load(open(os.path.expanduser(vocab_path))) + vocab = vj["vocab"] if isinstance(vj, dict) and "vocab" in vj else vj + if isinstance(vocab, dict): + vocab = vocab.get("actions", []) + print(f"vocab({len(vocab)}): {vocab}") + + import onnxruntime as ort + so = ort.SessionOptions() + so.log_severity_level = 3 + sess = ort.InferenceSession(os.path.expanduser(a.model), so, + providers=["CPUExecutionProvider"]) + inp = {i.name: i.shape for i in sess.get_inputs()} + print("model inputs:", inp) + tok_name = next(n for n in inp if "tok" in n.lower()) + seed_name = next((n for n in inp if "seed" in n.lower()), None) + zdim = None + if seed_name is not None: + zdim = int(inp[seed_name][-1]) + + actions = [s for s in a.actions.split(",") if s] or list(vocab) + rng = np.random.default_rng(a.seed) + + # ---- data reference (what the model is trying to match) ---- + if a.data: + z = np.load(os.path.expanduser(a.data), allow_pickle=False) + mo, tk = z["mo"], z["tk"] + dvocab = [str(s) for s in z["vocab"]] + print("\n=== TRAINING DATA (reference) ===") + for act in actions: + if act not in dvocab: + continue + idx = np.nonzero(tk[:, dvocab.index(act)])[0] + if not len(idx): + continue + pick = idx[:64] + ss = [score_window(act, mo[i]) for i in pick] + agg = {k: float(np.mean([s[k] for s in ss])) for k in ss[0]} + print(f" {act:10s} n={len(idx):6d} {fmt(agg)}") + + # ---- model ---- + print(f"\n=== MODEL ({a.samples} samples/action) ===") + print(f"{'action':10s} {'best-of-N':>44s} | {'mean':>44s}") + for act in actions: + if act not in vocab: + print(f" {act:10s} NOT IN VOCAB") + continue + t = np.zeros((1, len(vocab)), np.float32) + t[0, vocab.index(act)] = 1.0 + scores = [] + for _ in range(a.samples): + feeds = {tok_name: t} + if seed_name is not None: + feeds[seed_name] = (rng.standard_normal((1, zdim)) * 0.5).astype(np.float32) + out = sess.run(None, feeds)[0][0] + scores.append(score_window(act, quat_from_motion(out))) + # rank the way the shipped scorer does: upright + arms hanging + def rank(s): + # torsoUp dominates: a pitched-forward body must never win best-of-N + r = 3.0 * s["torsoUp"] + s["spineY"] + s["headY"] - s["armY_worst"] + if "fwd_side" in s: + r += min(s["fwd_side"], 4.0) + return r + best = max(scores, key=rank) + mean = {k: float(np.mean([s[k] for s in scores])) for k in scores[0]} + print(f" {act:10s} {fmt(best)} | {fmt(mean)}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval-t2m-refdist.py b/scripts/eval-t2m-refdist.py new file mode 100644 index 00000000..8201de17 --- /dev/null +++ b/scripts/eval-t2m-refdist.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# ruff: noqa: E702, E741 +"""Score a t2m model by DISTANCE TO REAL REFERENCE CLIPS (#837). + +ONE-TIME OFFLINE dev tool — NOT shipped. + +Why this exists. The hand-designed property metrics (travel direction, +contralateral phase, torso twist, limb amplitude, joint speed, gait +periodicity) can each be driven to match real motion while the render still +collapses: six independent hinges leave corners that satisfy all of them and +still look wrong, and five of those metrics turned out to have bugs that made +them reward the very defect they were added to catch. + +Real motion is a joint distribution, so compare against it directly. For each +generated clip, find its nearest real reference window in the canonical +representation and report that distance. One number, computed against ground +truth, that cannot be gamed by satisfying a property in isolation. + +Distance is per-frame mean geodesic quaternion angle over the valid canonical +joints, minimised over reference windows AND over a cyclic time shift of the +generated clip (phase is not a defect, so it should not be penalised). + +Usage: + python3 scripts/eval-t2m-refdist.py --model ~/t2m_v76/flow \ + [--refs "Walk.fbx=walk,Running.fbx=run"] [--samples 12] +""" +import argparse +import importlib.util +import json +import os +import subprocess +import tempfile + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) +DOWNLOADS = os.path.expanduser("~/Downloads") + + +def load_module(name, fname): + spec = importlib.util.spec_from_file_location(name, os.path.join(HERE, fname)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +prep5 = load_module("prep5", "prep-t2m-v5.py") + +J = 22 + + +def canon_from_fbx(path, cache_dir): + """Dump a real clip to canonical quats via the shipping extractor.""" + base = os.path.basename(path).replace(" ", "_") + ".canon.json" + out = os.path.join(cache_dir, base) + if not os.path.exists(out): + subprocess.run(["qtmesh", "anim", path, "--dump-canonical", out], + capture_output=True, text=True) + if not os.path.exists(out): + return None + d = json.load(open(out)) + clips = d.get("clips", []) + if not clips: + return None + c = clips[0] + cq, valid = prep5.canonicalize(np.asarray(c["quats"], np.float32), + c["restWorld"], c["restDir"]) + return cq, valid + + +def windows_of(cq, T=60, stride=5): + """Reference windows, cycle-extending short clips. + + The real clips are short — Walk.fbx is exactly 2.0 s (one 60-frame window) + and Running.fbx only 0.7 s — so a plain slice yields one or zero windows and + the distance has almost nothing to match against. Loop the clip and take + phase-shifted windows so every phase of the cycle is represented. + """ + n = len(cq) + if n < 2: + return [] + reps = int(np.ceil((T * 2) / n)) + 1 + ext = np.concatenate([cq] * reps, axis=0) + outs = [] + limit = min(len(ext) - T, max(n, T)) + for s in range(0, limit + 1, stride): + outs.append(ext[s:s + T]) + return outs + + +def geodesic_dist(a, b, valid): + """Mean per-frame per-joint geodesic angle (radians) between two clips.""" + dot = np.abs((a * b).sum(-1)).clip(0.0, 1.0) + ang = 2.0 * np.arccos(dot) # [T,J] + m = valid[None, :] + return float((ang * m).sum() / max(m.sum() * ang.shape[0], 1e-6)) + + +def best_distance(gen, refs, valid): + """Min over reference windows and over cyclic shifts of `gen`.""" + T = len(gen) + best = float("inf") + for r in refs: + for sh in range(0, T, 4): + g = np.roll(gen, sh, axis=0) + d = geodesic_dist(g, r, valid) + if d < best: + best = d + return best + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, help="dir with t2m.onnx + vocab") + ap.add_argument("--refs", default="Walk.fbx=walk,Running.fbx=run") + ap.add_argument("--samples", type=int, default=12) + ap.add_argument("--seed", type=int, default=11) + a = ap.parse_args() + + cache = os.path.join(tempfile.gettempdir(), "t2m_refcache") + os.makedirs(cache, exist_ok=True) + + refs = {} + valid_all = np.ones(J, np.float32) + for spec in a.refs.split(","): + if "=" not in spec: + continue + fname, action = spec.split("=", 1) + path = os.path.join(DOWNLOADS, fname.strip()) + if not os.path.exists(path): + print(f" (missing reference {fname})") + continue + got = canon_from_fbx(path, cache) + if got is None: + print(f" (could not extract {fname})") + continue + cq, valid = got + refs.setdefault(action.strip(), []).extend(windows_of(cq)) + valid_all = np.minimum(valid_all, valid) + if not refs: + raise SystemExit("no references extracted") + print(f"references: " + ", ".join(f"{k}({len(v)} windows)" + for k, v in refs.items())) + + mp = os.path.expanduser(a.model) + import onnxruntime as ort + so = ort.SessionOptions() + so.log_severity_level = 3 + sess = ort.InferenceSession(os.path.join(mp, "t2m.onnx"), so, + providers=["CPUExecutionProvider"]) + vocab = json.load(open(os.path.join(mp, "t2m-vocab.json")))["vocab"] + zd = int(sess.get_inputs()[1].shape[-1]) + rng = np.random.default_rng(a.seed) + + # self-distance floor: how close do DIFFERENT real windows get to each + # other? That is the best any model could plausibly score. + for act, rws in refs.items(): + if len(rws) > 1: + ds = [geodesic_dist(rws[i], rws[j], valid_all) + for i in range(len(rws)) for j in range(len(rws)) if i != j] + print(f" {act}: real-vs-real distance min {min(ds):.4f} " + f"median {np.median(ds):.4f} rad") + + print(f"\n{'action':8s} {'refDist(best)':>14s} {'refDist(mean)':>14s}") + for act, rws in refs.items(): + if act not in vocab: + print(f" {act:8s} not in vocab") + continue + t = np.zeros((1, len(vocab)), np.float32) + t[0, vocab.index(act)] = 1.0 + ds = [] + for _ in range(a.samples): + seed = (rng.standard_normal((1, zd)) * 0.5).astype(np.float32) + out = sess.run(None, {"tokens": t, "seed": seed})[0][0] + q = out.reshape(-1, J, 10)[:, :, 3:7] + q = q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-12) + ds.append(best_distance(q, rws, valid_all)) + print(f" {act:8s} {min(ds):14.4f} {np.mean(ds):14.4f} rad " + f"({np.degrees(np.mean(ds)):.1f} deg/joint)") + + +if __name__ == "__main__": + main() diff --git a/scripts/export-t2m-from-ckpt.py b/scripts/export-t2m-from-ckpt.py new file mode 100644 index 00000000..aa599d84 --- /dev/null +++ b/scripts/export-t2m-from-ckpt.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +# ruff: noqa: E702, E741 +"""Export a t2m ONNX from a MID-RUN training checkpoint (#837). + +ONE-TIME OFFLINE dev tool — NOT shipped. + +train-t2m-flow-v5.py only exports ONNX after its LAST epoch, so a 400-epoch +(~9h) run is unobservable until it finishes. This reuses the trainer's own +FlowDiT/Sampler definitions and export block against `/ckpt.pt`, so a +run in progress can be scored with eval-t2m-posture.py at any milestone +(the v6 notes show walk quality still climbing at ep200 → ep400, which is +exactly what you want to watch rather than assume). + +Reads the arch dims from the checkpoint's own tensor shapes, so it does not +need to be told --dim/--layers. + +Usage: + python3 scripts/export-t2m-from-ckpt.py --ckpt ~/t2m_v62/flow/ckpt.pt \ + --data ~/t2m_v62/t2m_v62.npz --out ~/t2m_v62/eval_ep100 [--steps 24] +""" +import argparse +import importlib.util +import json +import os + +import numpy as np +import torch + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def load_module(name, fname): + spec = importlib.util.spec_from_file_location(name, os.path.join(HERE, fname)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +tr = load_module("tr", "train-t2m-flow-v5.py") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt", required=True) + ap.add_argument("--data", required=True, help="npz (for vocab + canonRestDir)") + ap.add_argument("--out", required=True) + ap.add_argument("--steps", type=int, default=24) + # 1.0 = the value the SHIPPED v8.0 model was exported with. A mismatch + # here silently changes generated motion and invalidates checkpoint + # comparisons against it, so keep all three defaults (this, the trainer + # flag, and Sampler.__init__) on 1.0. + ap.add_argument("--guidance", type=float, default=1.0) + a = ap.parse_args() + + z = np.load(os.path.expanduser(a.data), allow_pickle=False) + vocab = [str(s) for s in z["vocab"]] + fps = int(z["fps"]) if "fps" in z else 30 + canon_rd = z["canonRestDir"] + T = int(z["mo"].shape[1]) + V = len(vocab) + + ck = torch.load(os.path.expanduser(a.ckpt), map_location="cpu", + weights_only=True) + sd = ck.get("net", ck) + epoch = ck.get("epoch", "?") + + # infer arch from tensor shapes rather than trusting flags + dim = None + for k, v in sd.items(): + if k.endswith("inp.weight") or ("inp" in k and v.ndim == 2): + dim = v.shape[0] + break + if dim is None: + dim = next(v.shape[-1] for v in sd.values() if v.ndim == 2) + layers = 1 + max( + (int(k.split(".")[1]) for k in sd if k.startswith("blocks.")), + default=0) + print(f"ckpt epoch={epoch} dim={dim} layers={layers} T={T} V={V}") + + net = tr.FlowDiT(V, T, dim=dim, layers=layers) + net.load_state_dict(sd) + net.eval() + + samp = tr.Sampler(net, V, T, a.steps, a.guidance).eval() + C6 = tr.C6 + J = tr.J + Z = T * C6 + tokens = torch.zeros(1, V); tokens[0, 0] = 1.0 + seed = torch.randn(1, Z) * 0.5 + + os.makedirs(os.path.expanduser(a.out), exist_ok=True) + onnx_path = os.path.join(os.path.expanduser(a.out), "t2m.onnx") + torch.onnx.export(samp, (tokens, seed), onnx_path, + input_names=["tokens", "seed"], + output_names=["motion"], opset_version=17, + dynamo=False) + vj = { + "vocab": vocab, "Z": Z, "T": T, "C": J * 10, "J": J, + "fps": fps, "frame": "world", "version": f"v5-flow-ep{epoch}", + "flowSteps": a.steps, + "restWorld": [[0.0, 0.0, 0.0, 1.0]] * J, + "restDir": [[float(v) for v in row] for row in canon_rd], + } + with open(os.path.join(os.path.expanduser(a.out), "t2m-vocab.json"), "w") as f: + json.dump(vj, f) + print(f"exported {onnx_path} ({os.path.getsize(onnx_path)/1e6:.1f} MB) " + f"+ vocab (epoch {epoch})") + + +if __name__ == "__main__": + main() diff --git a/scripts/pick-t2m-checkpoint.py b/scripts/pick-t2m-checkpoint.py new file mode 100644 index 00000000..ad30c6f8 --- /dev/null +++ b/scripts/pick-t2m-checkpoint.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Score archived t2m checkpoints on the actions a USER validated (#837). + +ONE-TIME OFFLINE dev tool — NOT shipped. + +Picking "the last epoch" is wrong: walk peaked and then regressed in earlier +runs, and the metrics disagree with each other (ep75 has the best walk refDist +0.269 while ep30 has the best facing across the validated actions). So score +every archived checkpoint on BOTH, restricted to the actions a human actually +confirmed are usable, and choose deliberately. + +USER_GOOD below is real user feedback on v7.8 ep60, not a guess. Keep it in sync +when new feedback arrives — and note the caveat recorded in EVAL_NOTES: the +shoulder-line facing metric measures FACING only and does NOT capture limb +articulation, so it ranks actions unreliably. It informs the choice between +checkpoints of the SAME model; it does not replace looking at renders. +""" +import argparse +import glob +import importlib.util +import json +import os +import tempfile + +import numpy as np +import onnxruntime as ort +import torch + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _load(name, fname): + spec = importlib.util.spec_from_file_location(name, os.path.join(HERE, fname)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +rd = _load("rd", "eval-t2m-refdist.py") +p6 = _load("p6", "prep-t2m-v6.py") +D=rd.prep5.D_CANON; PAR=p6.PAR +def pos(role,w): + T=len(w); p=np.zeros((T,3),np.float32); r=role + while PAR[r]>=0: + p=p+p6.qrot(w[:,PAR[r]],np.broadcast_to(D[r],(T,3))); r=PAR[r] + return p +def lat_err(q): + ls=p6.qrot(q[:,11],np.broadcast_to(D[11],(len(q),3))); rs=p6.qrot(q[:,7],np.broadcast_to(D[7],(len(q),3))) + v=ls-rs; v/=(np.linalg.norm(v,axis=-1,keepdims=True)+1e-9) + return float(np.sqrt(v[:,1]**2+v[:,2]**2).mean()) +_tr = _load("tr", "train-t2m-flow-v5.py") +_DIRS=torch.tensor(_tr.DIR_CANON,dtype=torch.float32) +def legchain(q): + """The user's reported 'second knee': bend in the ankle instead of the knee. + + Real Walk.fbx scores 0.014 and the training data 0.000, while v7.8 ep60 — + the checkpoint the user described as having 'almost a flat knee, slightly + backwards, then a second knee below' — scores 0.222. So this is the metric + that tracks the one defect the user named on the walk, and no other metric + here captures it (they all measure facing, phase or stride, never WHICH + joint bends). + """ + return float(_tr.leg_chain_penalty(torch.from_numpy(q[None]).float(),_DIRS).item()) +def scissor(q): + la,ra=pos(17,q),pos(21,q); return float(np.abs(la[:,2]-ra[:,2]).max()) +_ap = argparse.ArgumentParser(description=__doc__) +_ap.add_argument("--roots", default=os.path.expanduser("~/t2m_v78"), + help="comma-separated checkpoint roots to scan (each is " + "globbed for e*/ and ck*/ subdirs holding t2m.onnx)") +_ap.add_argument("--ref", default=os.path.expanduser("~/Downloads/Walk.fbx"), + help="real reference walk clip to score against") +_a = _ap.parse_args() + +cache=os.path.join(tempfile.gettempdir(),"t2m_refcache") +_ref = os.path.expanduser(_a.ref) +if not os.path.exists(_ref): + raise SystemExit(f"reference clip not found: {_ref} (pass --ref)") +cqw,vw=rd.canon_from_fbx(_ref,cache) +refs=rd.windows_of(cqw) +# actions the USER confirmed as good/usable on ep60 — these are what matter +USER_GOOD=["walk","run","wave","cough","death","pickup","attack","crouch", + "salute","working","shake","punch"] +so=ort.SessionOptions(); so.log_severity_level=3 +print(f"{'ckpt':10s} {'walk refD':>10s} {'walk latE':>10s} {'walk sciss':>11s} " + f"{'walk legchn':>12s} {'good-act latE':>14s}") +print(" (lower) (real .483) (real ~1.0) (real .014) (lower)") +rows=[] +cands = [] +for _root in _a.roots.split(","): + _root = os.path.expanduser(_root.strip()) + cands += glob.glob(os.path.join(_root, "e*")) + glob.glob(os.path.join(_root, "ck*")) +cands = sorted(set(cands)) +if not cands: + raise SystemExit(f"no checkpoint dirs under {_a.roots} (pass --roots)") +for d in cands: + f=os.path.join(d,"t2m.onnx") + if not os.path.exists(f): continue + s=ort.InferenceSession(f,so,providers=["CPUExecutionProvider"]) + vv=json.load(open(os.path.join(d,"t2m-vocab.json")))["vocab"]; zd=int(s.get_inputs()[1].shape[-1]) + def run(act,n): + rng=np.random.default_rng(7); t=np.zeros((1,len(vv)),np.float32) + if act not in vv: return None + t[0,vv.index(act)]=1.0 + R=[] + for _ in range(n): + out=s.run(None,{"tokens":t,"seed":(rng.standard_normal((1,zd))*0.5).astype(np.float32)})[0][0] + q=out.reshape(-1,22,10)[:,:,3:7]; q=q/np.linalg.norm(q,axis=-1,keepdims=True) + R.append(q) + return R + wq=run("walk",24) + wd=min(rd.best_distance(q,refs,vw) for q in wq) + wl=float(np.mean([lat_err(q) for q in wq])); ws=float(np.mean([scissor(q) for q in wq])) + wlc=float(np.mean([legchain(q) for q in wq])) + ls=[] + for a in USER_GOOD: + qs=run(a,8) + if qs: ls.append(np.mean([lat_err(q) for q in qs])) + gl=float(np.mean(ls)) + print(f"{os.path.basename(d):10s} {wd:10.3f} {wl:10.3f} {ws:11.3f} " + f"{wlc:12.3f} {gl:14.3f}") + rows.append((gl,wd,wlc,os.path.basename(d))) +print("\nreal walk: refD floor 0.437 | latErr 0.483 | scissor 0.652 | legchain 0.014") +rows.sort() +print(f"BEST by good-action facing: {rows[0][3]}") +print(f"BEST by walk refDist: {sorted(rows,key=lambda r:r[1])[0][3]}") +print(f"BEST by knee/ankle (the user's reported walk defect): " + f"{sorted(rows,key=lambda r:r[2])[0][3]}") diff --git a/scripts/prep-t2m-v5.py b/scripts/prep-t2m-v5.py index e7a387f2..ae0ae6a9 100644 --- a/scripts/prep-t2m-v5.py +++ b/scripts/prep-t2m-v5.py @@ -79,6 +79,19 @@ def load_module(name, fname): # FIXED canonical T-pose bone directions D (canonical axes: +Y up, +Z fwd, # +X left — the CMU frame convention every dump is normalized into). +# ROLE ORDER (AnimationMerger.cpp kParentCanon comment — do NOT re-derive this +# from the "X=left" axis note, which describes the AXES, not the role order): +# 0-5 hip, abdomen, chest, neck, neck1, head +# 6-9 RIGHT collar, shoulder, elbow, hand <-- RIGHT side, so -X +# 10-13 LEFT collar, shoulder, elbow, hand <-- LEFT side, so +X +# 14-17 RIGHT buttock, hip, knee, foot +# 18-21 LEFT buttock, hip, knee, foot +# Confirmed independently by compensateCanonicalHandedness(), which reads +# `lx = worldXForCanon(19)` (LEFT leg) vs `rx = worldXForCanon(15)` (RIGHT leg). +# A previous "chirality fix" swapped 6-9 with 10-13 on the false premise that +# 6-9 were the left arm; that INVERTED an already-correct table and produced a +# ~119 deg shoulder-vs-hip yaw in the cache (real gait counter-rotates 10-20), +# i.e. the anatomically twisted torso the user saw. Reverted. D_CANON = np.array([ [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [-1, 0, 0], [-1, 0, 0], [-1, 0, 0], [-1, 0, 0], diff --git a/scripts/prep-t2m-v6.py b/scripts/prep-t2m-v6.py index 85417c67..3a1598e2 100644 --- a/scripts/prep-t2m-v6.py +++ b/scripts/prep-t2m-v6.py @@ -32,7 +32,10 @@ import importlib.util import json import os +import re +import subprocess import sys +import tempfile import numpy as np @@ -58,8 +61,39 @@ def load_module(name, fname): MIRROR_PERM[a], MIRROR_PERM[b] = b, a LOCOMOTION = {"walk", "run", "march"} +# actions whose real-world joint speed exceeds the walk-tuned energy ceiling +FAST_ACTIONS = {"run", "march", "jump", "kick", "punch", "boxing", "attack", + "throw", "dance"} +# minimum dot(travel, body-forward) for a locomotion window to be kept +MIN_TRAVEL_FORWARD = 0.15 +# minimum ankle-scissor autocorrelation for a locomotion window (see +# gait_periodicity); real Mixamo Walk.fbx scores 0.968 +MIN_PERIODICITY = 0.6 +# real Mixamo clips used as the ground-truth shape for each action +REFERENCE_CLIPS = {"walk": "Walk.fbx", "run": "Running.fbx"} +# max geodesic distance (rad/joint) from the reference. Calibrated from the +# data: at 0.8 the pre-period-gate cache keeps 149 walk / 16 run windows, and +# real DIFFERENT motions sit 0.66 (jump) to 1.19 (punch) away, so 0.8 admits +# only genuinely walk-shaped motion. +MAX_REFERENCE_DISTANCE = 0.8 +# plausible gait-cycle lag window in frames at 30 fps (0.67 .. 1.0 s) +GAIT_LAG_MIN = 20 +GAIT_LAG_MAX = 30 +# max tolerated BACKWARD bend (degrees) for the knee/elbow hinges; small +# positive slack absorbs canonicalisation noise near full extension +MAX_HYPEREXTEND_KNEE = 20.0 +MAX_HYPEREXTEND_ELBOW = 35.0 HORIZONTAL_OK = {"death", "crawl", "roll", "swim", "fall", "sleep", "sit"} +# Source-exclusion (v6.2): some corpus characters carry a deliberately +# NON-HUMAN gait that the flow model faithfully reproduces on human prompts. +# Zombies shamble (bent spine, dragging/splayed stride, arms forward) and were +# 27% of the curated library incl. 5 walks + 1 run; "fruit"/produce characters +# (avocado etc.) are legless blobs whose canonical limbs are placeholders. +# Both poison locomotion. Matched case-insensitively against the clip's source +# string (asset title + animation name). +DEFAULT_EXCLUDE = r"zombie|avacado|avocado|fruit|banana|melon|undead|ghoul" + def qrot(q, v): qv = q[..., :3] @@ -101,13 +135,287 @@ def pos(role): return fwd / (side + 1e-6) +def travel_forward(w): + """dot(direction of travel, body's own forward). >0 = moving forward. + + Travel is inferred from the STANCE foot: a planted foot drifts backward + relative to a body moving forward, so travel = -velocity of the slower + (planted) foot per frame. Forward comes from the shoulder line + (left - right) x up. + + Some source clips genuinely walk BACKWARD, and they formed a real second + mode in the cache (37% of walk windows, mean -0.810). Flow matching then + samples both modes, so the model walked backward about half the time + (#837, user-reported "moving backwards"). Gating on this collapses the + conditional to one mode. + """ + # NB roles 7/11 are RIGHT/LEFT respectively (D_CANON[7] = -X, D_CANON[11] + # = +X), so `side` below is right->left... which is why the cross product + # with +Y yields the FORWARD axis with these operands in this order. A + # review flagged the old `lsh, rsh = fk_pos(w, 7), fk_pos(w, 11)` naming as + # a swapped binding and proposed exchanging them; that would INVERT the sign + # and make the gate reject every genuine forward walk. Verified empirically: + # the real Mixamo walk scores +0.581 and 100% of cached walk/run windows are + # positive as written. Renamed to match reality rather than changing it. + rsh, lsh = fk_pos(w, 7), fk_pos(w, 11) + side = rsh - lsh + side = side / (np.linalg.norm(side, axis=-1, keepdims=True) + 1e-9) + up = np.broadcast_to(np.array([0, 1, 0], np.float32), side.shape) + f = np.cross(side, up) + f = (f / (np.linalg.norm(f, axis=-1, keepdims=True) + 1e-9)).mean(0) + la, ra = fk_pos(w, 17), fk_pos(w, 21) + vl, vr = np.diff(la, axis=0), np.diff(ra, axis=0) + sl, sr = np.linalg.norm(vl, axis=-1), np.linalg.norm(vr, axis=-1) + v = np.where((sl < sr)[:, None], vl, vr) + t = -v.mean(0) + n = np.linalg.norm(t) + if n < 1e-6: + return 0.0 + fn = np.linalg.norm(f) + if fn < 1e-6: + return 0.0 + return float(np.dot(t / n, f / fn)) + + +def fk_pos(w, role): + """World position of `role` over the window via the canonical chain.""" + T = len(w) + p = np.zeros((T, 3), np.float32) + r = role + while PAR[r] >= 0: + p = p + qrot(w[:, PAR[r]], np.broadcast_to(D_CANON[r], (T, 3))) + r = PAR[r] + return p + + +# ---- reference-distance gate (v7.6) ---- +# The decisive measurement of this whole effort: scored against real Mixamo +# clips, the corpus's own "walk" windows sit a median 1.73 rad from Walk.fbx — +# FURTHER than a real punch (1.19), a real dance (0.91) or a real run (0.83). +# Real-vs-real is 0.03..0.44. So the corpus does not contain Mixamo-style +# walking, and every model trained on it faithfully reproduced that: the model +# measured 1.89 against its data's 1.73. No loss, gate or sampling change can +# close a gap that lives in the data. But a walk-LIKE subset exists (149 windows +# within 0.8 rad), so gate on distance to the real clips and train on that. +_REF_CACHE = {} +_REF_WARNED = set() + + +def _reference_windows(action, T): + """Cycle-extended canonical windows of the real reference clip, or None.""" + if action in _REF_CACHE: + return _REF_CACHE[action] + fname = REFERENCE_CLIPS.get(action) + if not fname: + _REF_CACHE[action] = None + return None + path = os.path.expanduser(os.path.join("~/Downloads", fname)) + if not os.path.exists(path): + _REF_CACHE[action] = None + return None + cache_dir = os.path.join(tempfile.gettempdir(), "t2m_refcache") + os.makedirs(cache_dir, exist_ok=True) + out = os.path.join(cache_dir, + os.path.basename(path).replace(" ", "_") + ".canon.json") + if not os.path.exists(out): + subprocess.run(["qtmesh", "anim", path, "--dump-canonical", out], + capture_output=True, text=True) + if not os.path.exists(out): + _REF_CACHE[action] = None + return None + try: + clips = json.load(open(out)).get("clips", []) + c = clips[0] + cq, valid = prep5.canonicalize(np.asarray(c["quats"], np.float32), + c["restWorld"], c["restDir"]) + except Exception: + _REF_CACHE[action] = None + return None + n = len(cq) + reps = int(np.ceil((T * 2) / max(n, 1))) + 1 + ext = np.concatenate([cq] * reps, axis=0) + wins = [ext[s:s + T] for s in range(0, min(len(ext) - T, max(n, T)) + 1, 5)] + _REF_CACHE[action] = (wins, valid) + return _REF_CACHE[action] + + +def reference_distance(action, w): + """Min geodesic distance to the real reference clip; inf if no reference. + + Minimised over reference windows and over cyclic shifts of `w`, so phase + is not penalised — only the shape of the motion. + """ + got = _reference_windows(action, len(w)) + if not got: + return float("inf") + wins, valid = got + m = valid[None, :] + denom = max(float(m.sum()) * w.shape[0], 1e-6) + best = float("inf") + for sh in range(0, len(w), 6): + g = np.roll(w, sh, axis=0) + for r in wins: + dot = np.abs((g * r).sum(-1)).clip(0.0, 1.0) + d = float((2.0 * np.arccos(dot) * m).sum() / denom) + if d < best: + best = d + return best + + +def gait_periodicity(w): + """Strongest autocorrelation of the ankle-scissor signal at lag >= 8. + + THE quality the other gates cannot see. A real gait is strongly cyclic — + real Mixamo Walk.fbx scores 0.968 — but every other metric here (energy, + stride ratio, travel, uprightness) is an average or a correlation that + NON-cyclic twitching satisfies just as well. The v6.8 model measured + healthy on all of them yet rendered as trembling in place, and its + periodicity was 0.26 against a data median of 0.446. + + 1.0 = perfect cycle, 0 = no repeat structure. + """ + la = fk_pos(w, 17)[:, 2] + ra = fk_pos(w, 21)[:, 2] + sig = la - ra + n = len(sig) + best = -1.0 + # Search only PLAUSIBLE GAIT CADENCES. A lag floor of 8 frames (0.27 s) was + # a serious bug: the corpus's median best-lag came out at exactly 8, i.e. + # the gate was accepting high-frequency WOBBLE as "periodic" and the model + # dutifully learned it — the trembling the user reported. Restricted to + # 20..30 frames (0.67..1.0 s at 30 fps, one full stride) the real Mixamo + # walk scores 0.968 while the corpus median is -0.379: the data is + # ANTI-correlated at true gait cadence. + # Pearson correlation computed PER OVERLAP. Normalising by the whole + # window's std while averaging over the (n - lag)-sample overlap is not a + # correlation and is not bounded: it scored a model at 1.131, above the real + # walk's 0.968, because the signal happened to be larger inside the overlap. + # An unbounded score also lets a single large arc pass as a "cycle". + for lag in range(GAIT_LAG_MIN, min(GAIT_LAG_MAX + 1, n // 2 + 1)): + a = sig[:-lag] + b = sig[lag:] + a = a - a.mean() + b = b - b.mean() + da, db = float(a.std()), float(b.std()) + if da < 1e-6 or db < 1e-6: + continue + c = float((a * b).mean() / (da * db)) + if c > best: + best = c + return max(best, 0.0) + + +def joint_hinge_signs(w): + """(worst knee, worst elbow) hinge angle, degrees. NOT CURRENTLY GATED ON. + + Requested as a guard against knees/elbows bending backwards, but it does + not work in the canonical frame and is left here documented rather than + silently enabled. Two sign conventions were tried and both FAILED against + real Mixamo clips: + - sign from the lateral axis (shoulder-left minus shoulder-right): that + axis flips between rigs, so 90% of the corpus read POSITIVE while the + real clip read negative — the gate would have dropped 90% of the data + (100% of march) on a convention mismatch alone. + - sign from "the shin swings backward relative to the thigh": false, + because mid-stride the shin legitimately swings AHEAD of the thigh; the + real walk then scored 75.7 and the real run 131.5, i.e. rejected. + An unsigned bound cannot work either: a real knee flexes 7..76 deg on a + walk and to 131 deg on a run, so "bent 75 forward" and "bent 75 backward" + are the same magnitude. A correct guard needs each rig's own bind-pose + hinge AXIS, which the canonical representation does not carry — it would + have to be computed in AnimationMerger against the target skeleton at + retarget time, not here. + + Knees and elbows are HINGES — they bend one way only. A negative value is + natural flexion in this frame and a positive one is HYPEREXTENSION (the + joint bending backwards), which reads as a broken limb. + + Convention measured on real Mixamo clips: knees run -75..-7 deg on a walk + and to -131 deg on a run (always negative = flexion); elbows sit near zero + while the arms hang. So the guard is on POSITIVE knee/elbow angles. + """ + def wdir(role): + d = qrot(w[:, role], np.broadcast_to(D_CANON[role], (len(w), 3))) + return d / (np.linalg.norm(d, axis=-1, keepdims=True) + 1e-9) + + # RIG-INDEPENDENT sign. A first cut took the sign from the lateral axis + # (shoulder-left minus shoulder-right), but that axis flips between rigs: + # 90% of the corpus scored a POSITIVE worst-knee while the real Mixamo clip + # scored negative, so the gate would have rejected 90% of the data (100% of + # march) purely from a convention mismatch. + # + # Anatomy gives a frame-free reference instead: a knee flexes so the SHIN + # swings BACKWARD relative to the thigh, i.e. the lower segment gains a + # component OPPOSITE the body's forward direction. Same for the forearm at + # the elbow. Forward comes from the shoulder line crossed with up, which is + # sign-stable because it is defined by the canonical +Y, not by which + # shoulder is "left". + up = np.array([0, 1, 0], np.float32) + lsh = qrot(w[:, 11], np.broadcast_to(D_CANON[11], (len(w), 3))) + rsh = qrot(w[:, 7], np.broadcast_to(D_CANON[7], (len(w), 3))) + lat = lsh - rsh + lat = lat / (np.linalg.norm(lat, axis=-1, keepdims=True) + 1e-9) + fwd = np.cross(lat, np.broadcast_to(up, lat.shape)) + fwd = fwd / (np.linalg.norm(fwd, axis=-1, keepdims=True) + 1e-9) + + def hyperextension(upper, lower): + """Degrees the hinge opens the WRONG way (0 when flexing naturally).""" + u, l = wdir(upper), wdir(lower) + ang = np.degrees(np.arccos(np.clip((u * l).sum(-1), -1, 1))) + # component of the lower segment along +forward, relative to the upper: + # flexion moves it backward (negative), hyperextension forward. + rel = ((l - u) * fwd).sum(-1) + # only count the angle when the joint opens forward + return np.where(rel > 0.05, ang, 0.0) + + knees = max(float(hyperextension(19, 20).max()), + float(hyperextension(15, 16).max())) + elbows = max(float(hyperextension(12, 13).max()), + float(hyperextension(8, 9).max())) + return knees, elbows + + +def skip_ankle_check(action): + """Actions where a large ankle angle is legitimate (kneeling, ground work).""" + return action in HORIZONTAL_OK or action in ("sit", "crouch", "pray", + "kick", "climb") + + +def dir_of(w, role): + """World aim of `role` over the window, unit-normalised.""" + d = qrot(w[:, role], np.broadcast_to(D_CANON[role], (len(w), 3))) + return d / (np.linalg.norm(d, axis=-1, keepdims=True) + 1e-9) + + def window_quality(action, w, valid): """True when the window meets the library curation bar.""" - # energy band — mean joint rotation speed (rad/frame) + # Energy band — mean joint rotation speed (rad/frame). The upper bound is + # ACTION-AWARE: a real Mixamo run measures 0.187 and a real walk 0.059, so + # the single 0.11 ceiling (tuned on walk-dominated data) systematically + # excluded genuine running — 5 of the 12 curated `run` clips died on it + # (0.13-0.22), which is why run trained on 24-36 windows and rendered as a + # gentle walk. Fast actions get headroom; walk keeps the tight bound. dq = np.abs((w[1:] * w[:-1]).sum(-1)).clip(0, 1) e = float((2 * np.arccos(dq)).mean()) - if not (0.004 <= e <= 0.11): + hi = 0.26 if action in FAST_ACTIONS else 0.11 + if not (0.004 <= e <= hi): return False + # ANATOMICAL ankle bound. The foot cannot fold perpendicular to the shin, + # yet 100% of the curated `march` windows measured a 104.5 deg ankle bend + # (knee a healthy 39.4 deg) — the source march templates are simply broken, + # which is exactly the "march is a twisted mess" the user reported. Left + # ungated it also poisons the trainer's leg-chain loss, since that term is + # masked to ALL locomotion and would be pulled toward this geometry. + # Real walk ankle 9.7 deg, real run 21.7 deg; 70 deg is generous headroom + # for a genuine high-knee march or kick while still rejecting a fold. + if not skip_ankle_check(action): + for knee, foot in ((20, 21), (16, 17)): + if valid[knee] and valid[foot]: + kd = dir_of(w, knee) + fd = dir_of(w, foot) + cos = np.clip((kd * fd).sum(-1), -1.0, 1.0) + if float(np.degrees(np.arccos(cos)).mean()) > 70.0: + return False if action in HORIZONTAL_OK: return True floor = 0.7 if action in LOCOMOTION else 0.5 @@ -129,13 +437,75 @@ def window_quality(action, w, valid): # sideways. 59% of raw CMU walk windows are splayed/side-stepping # (measured fwd/side < 1.5) — the model faithfully learned that # majority and walked sideways. Require a clean front-to-back stride. - if valid[17] and valid[21] and foot_travel_ratio(w) < 2.0: + # Stride-directionality gate: catches splayed/sideways walking. EXEMPT + # march — marching lifts the knees in place, so its fore/aft travel is + # legitimately low (CMU march windows: median ratio 1.56, only 36% clear + # 2.0), and applying a walk gate to it discarded 64% of the data and + # left march with 136 windows. + if (action != "march" and valid[17] and valid[21] + and foot_travel_ratio(w) < 2.0): return False + # REFERENCE-DISTANCE gate. Supersedes the periodicity gate, which was + # both too strict (it discarded genuinely walk-like windows: 149 survive + # at 0.8 rad here versus 15 in the periodicity-gated cache) and unable + # to see whether the motion actually resembles a walk. Actions with no + # reference clip are unaffected. + if action in REFERENCE_CLIPS: + rdst = reference_distance(action, w) + if not np.isfinite(rdst): + # No reference could be extracted (missing file / dump failure). + # Fail OPEN rather than silently rejecting every window of this + # action, and say so once — a silent total rejection would look + # like a data problem instead of a setup problem. + if action not in _REF_WARNED: + _REF_WARNED.add(action) + print(f"WARNING: no reference clip for '{action}' " + f"({REFERENCE_CLIPS[action]}) — distance gate SKIPPED", + flush=True) + elif rdst > MAX_REFERENCE_DISTANCE: + return False + # Travel-direction gate (v6.4): drop windows that move BACKWARD + # relative to the body's own forward. See travel_forward(). + # + # The threshold is RELAXED for the data-poor actions. At 0.15 the gate + # left run with 24 windows and march 132 (from 192/308), and the ep45 + # renders showed run and march too gentle to distinguish from a walk — + # the sampler cannot manufacture variety from 24 windows. Walk has 2704 + # and can afford the strict gate; for run/march, merely NOT going + # backward (> 0) is the useful signal. + if valid[7] and valid[11] and valid[17] and valid[21]: + floor = MIN_TRAVEL_FORWARD if action == "walk" else 0.0 + if travel_forward(w) < floor: + return False return True def mirror(w, valid): - """Sagittal mirror: reflect each quat (x,-y,-z,w) and swap L/R roles.""" + """Sagittal (left<->right) mirror: reflect about the X=0 plane and swap + L/R roles. + + A reflection about the plane with unit normal n maps a quaternion + (v, w) -> (-(v - 2(v.n)n), w) — i.e. negate the two components + PERPENDICULAR to n and keep the one along it. The sagittal plane's normal + is X (canonical X = left), so the correct sagittal mirror negates Y and Z + and KEEPS X... which is what this did. But that reflection also flips the + FORWARD (Z) axis, so every mirrored window travelled the opposite way and + the augmentation manufactured a backward-locomotion mode for free + (measured: 37% of walk windows travel backward, a real mode at -0.810). + Flow matching then samples both modes and the model walks backwards about + half the time (#837, user-reported). + + A true left<->right mirror reflects about the YZ plane (negating the LATERAL + X axis) and swaps the L/R joint roles. Conjugating a rotation by that + reflection, S @ R(q) @ S with S = diag(-1, 1, 1), is exactly q -> (x, -y, + -z, w), which is what this does. + + NB an earlier revision "fixed" this to (-x, y, z, -w). That is the NEGATION + of the same quaternion, and q and -q are the SAME rotation, so it changed + nothing — verified numerically against the conjugation above. The transform + below was correct all along; the note is kept so the no-op is not + reintroduced as a fix. + """ m = w * np.array([1, -1, -1, 1], np.float32) return m[:, MIRROR_PERM], valid[MIRROR_PERM] @@ -166,8 +536,26 @@ def main(): ap.add_argument("--T", type=int, default=60) ap.add_argument("--min-roles", type=int, default=12) ap.add_argument("--min-action-windows", type=int, default=16) + ap.add_argument("--library-repeat", type=int, default=1, + help="emit each curated-library take N times (#837). The " + "template clips face correctly and measure near real " + "motion; repeating them stops the far larger CMU " + "corpus from drowning them.") + ap.add_argument("--exclude-sources", default=DEFAULT_EXCLUDE, + help="regex; clips whose source matches are DROPPED " + "(default: non-human gaits — zombies, produce). " + "Pass '' to disable.") a = ap.parse_args() + excl = re.compile(a.exclude_sources, re.I) if a.exclude_sources else None + dropped_src = [0] + + def excluded(src): + if excl is None or not src or not excl.search(str(src)): + return False + dropped_src[0] += 1 + return True + T = a.T mo, msk, acts = [], [], [] gated = [0] @@ -179,8 +567,19 @@ def add(action, w, valid): def windows(action, cq, valid): nF = cq.shape[0] + # Admit clips shorter than T/2 WHEN THEY LOOP. Game run cycles are short + # loops: 11 of the 12 curated `run` clips are 20-27 frames, and the old + # T//2 (30-frame) floor discarded them before the cycle-repeat below + # could use them — leaving run with 24 training windows and a render too + # gentle to tell apart from a walk. Measured: 10 of those 11 loop + # cleanly (cyc <= 0.24, most exactly 0.000), so repeating them is sound. + # Non-looping short clips are still rejected (their repeat would jump). if nF < T // 2: - return + if nF < 16: + return + dq0 = np.abs((cq[-1] * cq[0]).sum(-1)).clip(0, 1) + if float((2 * np.arccos(dq0)).mean()) >= 0.25: + return if nF < T: dq = np.abs((cq[-1] * cq[0]).sum(-1)).clip(0, 1) cyc = float((2 * np.arccos(dq)).mean()) < 0.25 @@ -195,38 +594,81 @@ def windows(action, cq, valid): gated[0] += 1 continue add(action, w, valid) + # Augmented copies must clear the SAME gates as the base window. + # They previously bypassed window_quality entirely, so 18% of the + # "periodicity-gated" walk windows were actually below the 0.6 bar + # (min 0.200) — retime() resamples the window and can break the + # cycle at its edges, which is exactly the property being gated on. mw, mv = mirror(w, valid) - add(action, mw, mv) + if window_quality(action, mw, mv): + add(action, mw, mv) for f in (0.85, 1.15): - add(action, retime(w, f), valid) + rw_ = retime(w, f) + if window_quality(action, rw_, valid): + add(action, rw_, valid) if a.corpus: n = 0 - for action, cq, valid, _src in prep5.corpus_clips( + for action, cq, valid, src in prep5.corpus_clips( os.path.expanduser(a.corpus), a.min_roles): + if excluded(src): + continue windows(action, cq, valid) n += 1 print(f"corpus: {n} clips → {len(mo)} windows (cum)") if a.library: lib = json.load(open(os.path.expanduser(a.library))) n = 0 + dropped_roles = [0] for c in lib.get("clips", []): rw, rd = c.get("restWorld"), c.get("restDir") if not rw or not rd: continue + if excluded(c.get("source", "")): + continue cq, valid = prep5.canonicalize( np.asarray(c["quats"], np.float32), rw, rd) - windows(c["action"], cq, valid) + # --min-roles was only ever applied to the CORPUS loader, so with + # --library-repeat making the library the dominant source the flag + # was effectively inert (another silently-dead knob, same class as + # the amplitude term). It matters most here: the `hey` take resolves + # only 9 of 22 roles with ALL FOUR knee/ankle roles invalid, so the + # anatomical ankle gate cannot even run on it and the model has no + # leg data to learn from — it invents legs. Sorting the curated + # takes by valid-role count reproduces the user's verdict: `hey` 9 + # and `confession` 14 are the two worst and the two they called + # "not that good", while everything they rated GOOD has >=16. + if int(valid.sum()) < a.min_roles: + dropped_roles[0] += 1 + continue + # The curated TEMPLATE clips are the only source that both renders + # with correct facing on real rigs and measures near real motion: + # refDist to the real Mixamo walk is 0.303-0.462 for the template + # walks (the real-vs-real floor is 0.437 and the corpus median is + # 2.07), and their shoulder-line lateral error is 0.519 vs the + # corpus's 0.665 and real motion's 0.483. --library-repeat emits + # each take that many times so this good data is not drowned by the + # much larger CMU corpus. + for _ in range(max(1, a.library_repeat)): + windows(c["action"], cq, valid) n += 1 - print(f"library: {n} takes → {len(mo)} windows (cum)") + print(f"library: {n} takes x{max(1, a.library_repeat)} " + f"→ {len(mo)} windows (cum)") + if dropped_roles[0]: + print(f"library: dropped {dropped_roles[0]} takes below " + f"--min-roles {a.min_roles}") if a.bvh and a.index: n = 0 - for action, cq, valid, _src in prep5.cmu_clips(a.bvh, a.index): + for action, cq, valid, src in prep5.cmu_clips(a.bvh, a.index): + if excluded(src): + continue windows(action, cq, valid) n += 1 print(f"cmu: {n} trials → {len(mo)} windows (cum)") print(f"quality gates dropped {gated[0]} base windows") + print(f"source exclusion dropped {dropped_src[0]} clips " + f"(pattern: {a.exclude_sources or None})") if not mo: sys.exit("no windows") from collections import Counter @@ -241,9 +683,17 @@ def windows(action, cq, valid): print(f"windows: {mo.shape[0]} vocab({len(vocab)}):", {w: int(tk[:, vocab.index(w)].sum()) for w in vocab}) os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) + # Stamp the build command into the cache. The v7.8 npz recorded only + # fps/vocab/canonRestDir, so reproducing it later meant inferring the flags + # from action-overlap against the library files — slow and error-prone. np.savez_compressed(a.out, mo=mo, msk=msk, tk=tk, vocab=np.array(vocab), fps=FPS, - canonRestDir=D_CANON) + canonRestDir=D_CANON, + # Plain unicode arrays, NOT dtype=object: an object + # array forces allow_pickle=True on every reader, and + # a crafted npz could then execute code on load. + buildArgv=np.array(sys.argv, dtype=np.str_), + buildFlags=np.array(json.dumps(vars(a)), dtype=np.str_)) print(f"wrote {a.out} mo{mo.shape}") diff --git a/scripts/train-t2m-flow-v5.py b/scripts/train-t2m-flow-v5.py index ba6c145e..abc87d7a 100644 --- a/scripts/train-t2m-flow-v5.py +++ b/scripts/train-t2m-flow-v5.py @@ -93,6 +93,289 @@ def d6_to_quat(d): return F.normalize(q, dim=-1, eps=1e-6) +# ---------- gait-phase supervision (#837 orientation correctness) ---------- +# Canonical parent chain + rest directions, mirroring prep-t2m-v5/v6 exactly. +PAR_CANON = [-1, 0, 1, 2, 3, 4, 2, 6, 7, 8, 2, 10, 11, 12, + 0, 14, 15, 16, 0, 18, 19, 20] +DIR_CANON = [ # keep in sync with prep-t2m-v5.D_CANON + [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], + [-1, 0, 0], [-1, 0, 0], [-1, 0, 0], [-1, 0, 0], # roles 6-9 = RIGHT arm + [1, 0, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0], # roles 10-13 = LEFT arm + [0, -1, 0], [0, -1, 0], [0, -1, 0], [0, 0, 1], + [0, -1, 0], [0, -1, 0], [0, -1, 0], [0, 0, 1], +] +# roles whose Z trajectory defines the gait phase +L_ANKLE, R_ANKLE, L_WRIST, R_WRIST = 17, 21, 9, 13 + + +def qrot_t(q, v): + """Rotate v[...,3] by quat q[...,4] (x,y,z,w). Differentiable.""" + qv, qw = q[..., :3], q[..., 3:] + uv = torch.cross(qv, v, dim=-1) + uuv = torch.cross(qv, uv, dim=-1) + return v + 2.0 * (qw * uv + uuv) + + +def fk_pos(q, role, dirs): + """World position of `role` by walking the canonical parent chain. + + q: [B,T,J,4]. Mirrors the numpy FK in prep-t2m-v6/eval-t2m-posture, so the + quantity supervised here is the same one the metrics and the retarget read. + """ + p = torch.zeros(q.shape[0], q.shape[1], 3, device=q.device, dtype=q.dtype) + r = role + while PAR_CANON[r] >= 0: + d = dirs[r].expand(q.shape[0], q.shape[1], 3) + p = p + qrot_t(q[:, :, PAR_CANON[r]], d) + r = PAR_CANON[r] + return p + + +def _centred_z(p): + z = p[..., 2] + return z - z.mean(dim=1, keepdim=True) + + +def travel_forward_torch(q, dirs): + """dot(direction of travel, body forward), differentiable, per sample. + + Same quantity as prep-t2m-v6.travel_forward and the eval: travel is + inferred from the STANCE (slower) foot, which drifts backward while the + body moves forward. Forward is (role7 - role11) x up, where role 7 is the + RIGHT shoulder and role 11 the LEFT (D_CANON[7] = -X, D_CANON[11] = +X) — + see the note in prep-t2m-v6.travel_forward: this operand order is what + yields the forward axis, and exchanging them inverts the sign. + + Even with a 100%-forward training set the model sat at ~50% forward + (measured at ep40 with foot motion at 73% of the data's magnitude, so this + is a real direction ambiguity and not measurement noise). Flow matching has + no term tying the emitted gait to a travel direction, so this supervises it + directly — the same reasoning as the gait-phase hinge. + """ + rsh = fk_pos(q, 7, dirs) # role 7 = RIGHT shoulder + lsh = fk_pos(q, 11, dirs) # role 11 = LEFT shoulder + side = F.normalize(rsh - lsh, dim=-1, eps=1e-6) + up = torch.zeros_like(side) + up[..., 1] = 1.0 + fwd = F.normalize(torch.cross(side, up, dim=-1), dim=-1, eps=1e-6) + fwd = F.normalize(fwd.mean(dim=1), dim=-1, eps=1e-6) # [B,3] + + la = fk_pos(q, 17, dirs) + ra = fk_pos(q, 21, dirs) + vl = la[:, 1:] - la[:, :-1] + vr = ra[:, 1:] - ra[:, :-1] + sl = vl.norm(dim=-1, keepdim=True) + sr = vr.norm(dim=-1, keepdim=True) + v = torch.where(sl < sr, vl, vr) # stance foot + travel = -v.mean(dim=1) # [B,3] + travel = F.normalize(travel, dim=-1, eps=1e-6) + return (travel * fwd).sum(-1) # [B] in [-1,1] + + +def amplitude_excess(q, dirs, lo_scale=None): + """How far the ankle/wrist fore-aft excursion EXCEEDS human range. + + The phase + travel hinges reward motion and nothing bounded it, so the + model over-drove the limbs: measured stride 3.46 and armSwing 3.73 against + real Mixamo walk values of 0.765 and 1.718 - rendering as a knee-to-chest + exaggerated march instead of a walk. This penalises only the EXCESS above + generous ceilings, so normal motion is untouched. + """ + def pk(role): + z = fk_pos(q, role, dirs)[..., 2] + return z.max(dim=1).values - z.min(dim=1).values + + stride = 0.5 * (pk(17) + pk(21)) + swing = 0.5 * (pk(9) + pk(13)) + # TWO-SIDED, like the speed band. A ceiling-only amplitude term contributes + # exactly 0.0 once the model is under it, so nothing defends the limbs from + # COLLAPSING: measured at v6.7 ep45 the jitter term was 80% of the guard + # loss (0.1446/0.181) while amp was 0.0000, and armSwing/stride fell + # 1.481/1.126 -> 0.894/0.690, i.e. BELOW the real walk's 1.718/0.765. The + # limbs were collateral damage of jitter pushing speed down. Real-walk + # values anchor the floors (stride 0.765, swing 1.718); floors are set just + # under them and weighted 2x so shrinking is punished harder than growing. + # Ceilings sit above the ALL-ACTION p99 (stride 3.33, swing 3.54) so real + # data is never penalised. The original 1.6 stride ceiling was BELOW the + # data's p95 of 2.969 — it was clipping legitimate motion across march, + # climb and kick, which contributed to the amplitude collapse. + over = (stride - 3.5).clamp_min(0.0) + (swing - 3.7).clamp_min(0.0) + # Floors calibrated to the TRAINING data's 5th percentile (walk stride 0.741, + # swing 0.809) — NOT to the reference Mixamo clip, whose swing (1.718) sits + # well above this corpus's median (1.052). Floors above the data would + # penalise legitimate windows: at swing floor 1.50, 67% of real walk windows + # scored a penalty. + st_floor = 0.70 * (1.0 if lo_scale is None else lo_scale) + sw_floor = 0.75 * (1.0 if lo_scale is None else lo_scale) + under = ((st_floor - stride).clamp_min(0.0) + + (sw_floor - swing).clamp_min(0.0)) * 2.0 + return over + under + + +def jitter_excess(q, hi=None, lo=None): + """Per-frame angular speed ABOVE the training band = high-frequency jitter. + + THE defect that made every amplitude metric lie. The ep90 model emitted a + mean joint speed of 0.266 rad/frame against the data's 0.0557 (band + 0.032-0.099), and its worst joint (role 21) ran at 1.416 rad/frame — about + 81 deg per frame, physically impossible. The retarget's smoothing then + damped it to 0.041, so the RENDER was nearly static while the canonical + ankle "scissor" measured LARGER than the real walk: the metrics were + reading jitter amplitude, not stride. + + Penalise the mean per-joint speed above a ceiling just past the data band, + plus the worst single joint, which is where the impossible spikes live. + """ + dot = (q[:, 1:] * q[:, :-1]).sum(-1).abs().clamp(0.0, 1.0) + speed = 2.0 * torch.acos(dot.clamp(max=1.0 - 1e-7)) # [B,T-1,J] + per_joint = speed.mean(dim=1) # [B,J] + # TWO-SIDED band, not a ceiling. A pure ceiling removes the jitter AND the + # motion with it: at ep30 speed fell 0.353 -> 0.227 and the render went + # smooth but nearly STATIC (legs barely separating, contra dropped + # 0.345 -> 0.284). Jitter and stride were entangled, so the objective has + # to pull the speed TOWARD the data band (walk mean 0.056, real Walk.fbx + # 0.059) rather than merely below a ceiling. + m = per_joint.mean(dim=-1) + # The band is PER-SAMPLE so it can be action-aware. A single band is + # action-blind, and after the run/march data recovery the cache's overall + # energy rose 0.0429 -> 0.0550 (p90 0.078 -> 0.113) while WALK-only energy + # stayed at 0.0510 — the faster fast-action windows pull the shared model + # and walk inherits the speed, which is what degraded the v7.2 walk render. + hi_t = 0.10 if hi is None else hi + lo_t = 0.045 if lo is None else lo + mean_excess = ((m - hi_t).clamp_min(0.0) + + (lo_t - m).clamp_min(0.0) * 4.0) + worst_excess = (per_joint.max(dim=-1).values - 0.45).clamp_min(0.0) + return mean_excess + worst_excess + + +def gait_aperiodicity(q, dirs): + """1 - (best autocorrelation of the ankle-scissor signal), differentiable. + + The property NO other term can see. A real Mixamo walk autocorrelates at + 0.968; the v6.8 model measured 0.264 and rendered as trembling in place + while scoring healthy on travel, contra, twist, amplitude and speed — all + of which are averages or correlations that non-cyclic twitching satisfies. + + Gating the DATA to >= 0.6 periodicity was not enough on its own: the model + still emitted 0.247-0.293, i.e. no better than before. Same lesson as travel + direction — flow matching does not inherit a property just because the data + has it. If it is not in the loss, it is not in the output. + + Autocorrelation is evaluated at a fixed lag set spanning plausible gait + periods (10..29 frames at 30 fps = 0.33..0.97 s) and the best is taken, so + the term is agnostic to cadence. + """ + la = fk_pos(q, 17, dirs)[..., 2] + ra = fk_pos(q, 21, dirs)[..., 2] + sig = la - ra + T = sig.shape[1] + best = torch.full((sig.shape[0],), -1.0, device=q.device, dtype=sig.dtype) + # Lag range must MATCH prep-t2m-v6.gait_periodicity (8 .. T//2), otherwise + # the loss disagrees with the gate that selected the data: a narrower + # 10..29 window scored the >=0.6-periodic training set at 0.43-0.51 + # aperiodicity because it missed the slower cadences. + # Match prep-t2m-v6's GAIT_LAG_MIN/MAX. Searching from lag 8 rewarded + # high-frequency wobble instead of a stride (see gait_periodicity). + # Pearson per overlap (bounded), matching prep-t2m-v6.gait_periodicity. + for lag in range(20, min(31, max(21, T // 2 + 1))): + a = sig[:, :-lag] + b = sig[:, lag:] + a = a - a.mean(dim=1, keepdim=True) + b = b - b.mean(dim=1, keepdim=True) + c = ((a * b).mean(dim=1) + / ((a.std(dim=1) * b.std(dim=1)) + 1e-4)) + best = torch.maximum(best, c) + return (1.0 - best).clamp_min(0.0) + + +def leg_chain_penalty(q, dirs): + """Keep the leg's bend in the KNEE, not the ankle. + + User report on the v7.8 walk: "almost like a flat knee, slightly backwards, + then a second knee below on the correct orientation". Measured cause — the + bend is in the wrong joint: + + thigh->shin (knee) real 38.4 deg model 24.9 deg (UNDER-bent) + shin->foot (ankle) real 9.7 deg model 31.5 deg (3x OVER-bent) + + The over-bent ankle is what reads as a false second knee. Penalise ankle + flexion above a generous ceiling and knee flexion below a floor, both + anchored on the real-walk values, so the bend moves back up the chain. + """ + def wdir(role): + d = qrot_t(q[:, :, role], dirs[role].expand(q.shape[0], q.shape[1], 3)) + return F.normalize(d, dim=-1, eps=1e-6) + + pen = 0.0 + for hip, knee, foot in ((19, 20, 21), (15, 16, 17)): + th, sh, ft = wdir(hip), wdir(knee), wdir(foot) + knee_ang = torch.acos(((th * sh).sum(-1)).clamp(-1 + 1e-6, 1 - 1e-6)) + ankle_ang = torch.acos(((sh * ft).sum(-1)).clamp(-1 + 1e-6, 1 - 1e-6)) + # radians: real walk knee 0.67, ankle 0.17 + pen = pen + (ankle_ang.mean(dim=1) - 0.35).clamp_min(0.0) + pen = pen + (0.45 - knee_ang.mean(dim=1)).clamp_min(0.0) + return pen * 0.5 + + +def spine_twist_penalty(q): + """Mean |hip->chest yaw| in radians — an anatomy guard for the phase loss. + + The phase term alone is satisfiable by a DEGENERATE shortcut: rotate the + torso ~180 deg so the existing swing reads as contralateral. Measured at + --phase-weight 0.15 without this guard: hip->chest twist went to 170 deg + (data 6.3 deg, v6.2 6.7 deg) — the user-visible "walk got all twisted". + Penalising the yaw closes that escape route so the only way to earn the + phase reward is to actually fix limb timing. + """ + def conj(a): + return torch.cat([-a[..., :3], a[..., 3:]], -1) + + def mul(a, b): + ax, ay, az, aw = a.unbind(-1) + bx, by, bz, bw = b.unbind(-1) + return torch.stack([aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + aw * bw - ax * bx - ay * by - az * bz], -1) + + rel = mul(conj(q[:, :, 0]), q[:, :, 2]) # hip -> chest + yaw = 2.0 * torch.atan2(rel[..., 1], rel[..., 3]) + # wrap to [-pi, pi] so a 180 deg cheat is maximally penalised + yaw = torch.atan2(torch.sin(yaw), torch.cos(yaw)) + return yaw.abs().mean(dim=1) + + +def gait_phase_corr(q, dirs): + """Per-sample contralateral gait correlation, in [-1, 1]. + + A correct walk pairs LEFT leg forward with RIGHT arm forward, so + corr(Lankle_z, Rwrist_z) > 0 and corr(Lankle_z, Lwrist_z) < 0. The + training data measures +0.63 / -0.58 (92% correct sign) but the v6.1 and + v6.2 models emit only 67% / 75% correct — samples with the arms swinging + on the wrong side, which renders as "limbs moving as if facing backwards" + (user-reported). Nothing in the flow-matching loss constrained this, so + this term supervises it directly. + + Returns the mean of (contralateral - ipsilateral)/2, averaged over both + body sides so the term is itself mirror-symmetric. + """ + la = _centred_z(fk_pos(q, L_ANKLE, dirs)) + ra = _centred_z(fk_pos(q, R_ANKLE, dirs)) + lw = _centred_z(fk_pos(q, L_WRIST, dirs)) + rw = _centred_z(fk_pos(q, R_WRIST, dirs)) + + def corr(x, y): + xs = x / (x.std(dim=1, keepdim=True) + 1e-4) + ys = y / (y.std(dim=1, keepdim=True) + 1e-4) + return (xs * ys).mean(dim=1) + + # both sides, so the objective cannot prefer one chirality + contra = 0.5 * (corr(la, rw) + corr(ra, lw)) + ipsi = 0.5 * (corr(la, lw) + corr(ra, rw)) + return 0.5 * (contra - ipsi) + + # ---------- model ---------- class Block(nn.Module): def __init__(self, dim, heads): @@ -147,7 +430,11 @@ def forward(self, x, t, tok): class Sampler(nn.Module): """Euler flow sampler UNROLLED for ONNX export — MotionGenerator contract.""" - def __init__(self, net, V, T, steps, guidance=2.0): + # guidance default 1.0 — the value the shipped v8.0 model was trained + # and exported with. This used to default to 2.0 while both CLI flags + # defaulted to 1.0, so a caller constructing a Sampler directly got + # different motion than the CLI produced. + def __init__(self, net, V, T, steps, guidance=1.0): super().__init__() self.net, self.V, self.T, self.steps = net, V, T, steps self.guidance = guidance @@ -188,12 +475,63 @@ def main(): ap.add_argument("--guidance", type=float, default=1.0, help="CFG scale baked into the exported sampler") ap.add_argument("--device", default="mps") + ap.add_argument("--balance-power", type=float, default=1.0, + help="class-balance exponent: 1.0 = inverse-frequency " + "(equal per action), 0.5 = sqrt-tempered (favours " + "the data-rich actions like walk), 0.0 = raw") + ap.add_argument("--phase-weight", type=float, default=0.0, + help="weight of the contralateral gait-phase loss " + "(#837 orientation correctness). 0 = off. Applied to " + "LOCOMOTION actions only, on the reconstructed clean " + "sample. 0.05-0.2 is a sane range.") + ap.add_argument("--travel-weight", type=float, default=0.0, + help="weight of the travel-DIRECTION hinge: the body must " + "move along its own forward axis (#837). Locomotion " + "rows only. 0 = off; 0.1-0.3 is a sane range.") + ap.add_argument("--jitter-weight", type=float, default=0.0, + help="penalise per-frame angular speed above the data band " + "(#837). Applied to ALL actions — jitter is never " + "wanted. 0 = off; 1.0 is a sane start.") + ap.add_argument("--loco-boost", type=float, default=1.0, + help="multiply walk/run/march sampling weight (#837). The " + "periodicity gate leaves them ~7%% of batches, so the " + "gait losses barely reach any rows. 6.0 gives ~33%%.") + ap.add_argument("--period-weight", type=float, default=0.0, + help="reward a periodic GAIT CYCLE on locomotion rows " + "(#837). Data gating alone does not produce one. " + "0 = off; 0.5 is a sane start.") + ap.add_argument("--legchain-weight", type=float, default=0.0, + help="keep the leg bend in the KNEE not the ankle (#837): " + "an over-bent ankle reads as a false second knee. " + "0 = off; 0.3 is a sane start.") + ap.add_argument("--amp-weight", type=float, default=0.0, + help="penalise limb excursion ABOVE human range (#837): " + "the phase/travel hinges reward motion and otherwise " + "over-drive the legs into a knee-to-chest march. " + "0 = off; 0.3 is a sane start.") + ap.add_argument("--twist-weight", type=float, default=0.5, + help="weight of the hip->chest yaw penalty that stops the " + "phase loss from cheating by rotating the torso 180 " + "deg. Only active when --phase-weight > 0.") + ap.add_argument("--phase-actions", default="walk,run,march", + help="comma list the phase loss applies to") ap.add_argument("--resume", action="store_true", help="resume from /ckpt.pt (long runs survive " "sleep/restarts)") a = ap.parse_args() - d = np.load(a.data, allow_pickle=True) + # The leg-chain/amplitude/travel/period terms live inside the phase-loss + # block (they all need the reconstructed clean sample q_hat), so with + # --phase-weight 0 they are silently DEAD. That class of bug already cost + # several runs via the amplitude term, so refuse rather than no-op. + if a.phase_weight <= 0 and max(a.legchain_weight, a.amp_weight, + a.travel_weight, a.period_weight) > 0: + raise SystemExit( + "--legchain/amp/travel/period-weight require --phase-weight > 0 " + "(they share its reconstructed-sample block)") + + + d = np.load(a.data, allow_pickle=False) mo, msk, tk = d["mo"], d["msk"], d["tk"] vocab = [str(w) for w in d["vocab"]] fps = int(d["fps"]) @@ -209,9 +547,60 @@ def main(): m6 = torch.from_numpy(msk).repeat_interleave(D6, -1) # [N,C6] tok = torch.from_numpy(tk) - # class-balanced sampling — walk is 64% of windows (v4 lesson) + # gait-phase loss plumbing (#837): a per-action gate + the canonical rest + # directions as a device tensor for the differentiable FK. + dirs_t = torch.tensor(DIR_CANON, dtype=torch.float32, device=dev) + phase_acts = {w for w in a.phase_actions.split(",") if w} + loco_idx = [i for i, w in enumerate(vocab) if w in phase_acts] + loco_mask_v = torch.zeros(V, device=dev) + for i in loco_idx: + loco_mask_v[i] = 1.0 + # actions whose real joint speed exceeds the walk band (mirrors + # prep-t2m-v6.FAST_ACTIONS) + # Measured mean-speed medians on the v72 cache: run 0.110, dance 0.095, + # punch 0.064, walk 0.055, march 0.050, jump 0.042. Only run and dance + # genuinely need a raised ceiling; march/jump are NOT fast by this metric. + fast_names = {"run", "dance"} + fast_mask_v = torch.zeros(V, device=dev) + for i, w in enumerate(vocab): + if w in fast_names: + fast_mask_v[i] = 1.0 + if a.phase_weight > 0: + if a.legchain_weight > 0: + print(f"leg-chain loss ON (w={a.legchain_weight}) — bend stays in " + f"the knee, not the ankle") + print(f"gait-phase loss ON (w={a.phase_weight}) for " + f"{[vocab[i] for i in loco_idx]}", flush=True) + + # Class-balanced sampling — walk dominates the window count (v4 lesson). + # `--balance-power` tempers it: 1.0 = pure inverse-frequency (every action + # drawn equally often), 0.5 = sqrt-tempered, 0.0 = raw distribution. + # + # Pure inverse-frequency badly starves the actions that matter most. On the + # v6.2 cache walk is 30.85% of windows but only 4.35% of draws (0.14x its + # share) while 16-window curiosities like `confession` get 67x oversampling + # — equal capacity spent memorising 16 windows as on 7688 walk windows. + # Measured effect: walk fwd/side plateaued at ~2.6 by ep167 while the loss + # kept falling. sqrt-tempering gives walk 16.1% and still leaves run 2.8%. freq = tk.sum(0) - w = (tk @ (1.0 / np.maximum(freq, 1.0))).astype(np.float64) + w = (tk @ (1.0 / np.maximum(freq, 1.0) ** a.balance_power)).astype(np.float64) + # LOCOMOTION BOOST. The periodicity gate shrank walk/run/march to 467/39/41 + # windows of 19693 while the other 21 actions kept theirs, so locomotion was + # only 7.1% of sampled batches (~18 rows of 256) — the gait losses reach + # only those rows, so they were weak in practice no matter their weight. + # balance-power alone caps locomotion at 12.5% (3 of 24 actions at bp=1.0), + # which is still far too little for the three actions that matter here. + if a.loco_boost > 1.0: + loco_names = {"walk", "run", "march"} + li = [i for i, n in enumerate(vocab) if n in loco_names] + if li: + boost = np.ones(len(vocab), np.float64) + for i in li: + boost[i] = a.loco_boost + w = w * (tk @ boost) + sh = w[(tk[:, li].sum(1) > 0)].sum() / w.sum() + print(f"locomotion sampling share: {100 * sh:.1f}% " + f"(boost x{a.loco_boost})", flush=True) sampler = torch.utils.data.WeightedRandomSampler( torch.from_numpy(w), num_samples=N, replacement=True) ds = torch.utils.data.TensorDataset(x1, m6, tok) @@ -232,18 +621,46 @@ def main(): # epoch — no pickled objects — so load safely (no code execution). ck = torch.load(ckpt_path, map_location=dev, weights_only=True) net.load_state_dict(ck["net"]) - opt.load_state_dict(ck["opt"]) - sched.load_state_dict(ck["sched"]) - start_ep = ck["epoch"] + 1 - print(f"resumed from epoch {start_ep}", flush=True) + # A WARM START (weights transplanted from a run with a different vocab) + # deliberately carries no optimizer/scheduler state: their moment + # buffers are shaped for the OLD act_emb and would be stale or + # shape-mismatched. Restore them only when present. + if "opt" in ck and "sched" in ck: + opt.load_state_dict(ck["opt"]) + sched.load_state_dict(ck["sched"]) + start_ep = ck["epoch"] + 1 + print(f"resumed from epoch {start_ep}", flush=True) + else: + start_ep = ck.get("epoch", 0) + print(f"WARM START from transplanted weights at epoch {start_ep} " + f"(fresh optimizer/scheduler)", flush=True) + + # `range(start_ep, a.epochs)` is EMPTY when --epochs is at or below the + # resume point, and the export block below would then happily write an + # ONNX of untrained (or, on a warm start, merely transplanted) weights — + # silently, and indistinguishable from a real result. Refuse instead. + if start_ep >= a.epochs: + raise SystemExit( + f"--epochs {a.epochs} is not beyond the checkpoint's epoch " + f"{start_ep}: there is nothing to train. Raise --epochs to export " + f"a trained model, or use export-t2m-from-ckpt.py to export the " + f"checkpoint as-is.") for ep in range(start_ep, a.epochs): - tot, nb = 0.0, 0 + # Accumulate the epoch loss ON DEVICE. A per-batch `loss.item()` is a + # GPU->CPU sync (`_local_scalar_dense_mps` -> `waitUntilCompleted`) + # that drains the whole MPS queue ~97x/epoch; on long runs the MPS + # allocator degrades until each sync parks the process in + # uninterruptible wait and throughput collapses ~11x (observed at + # ep84: 84s/epoch -> ~16min/epoch). One sync per EPOCH instead. + tot = torch.zeros((), device=dev) + nb = 0 for xb, mb, tb in dl: xb, mb, tb = xb.to(dev), mb.to(dev), tb.to(dev) x0 = torch.randn_like(xb) # classifier-free guidance: drop the action condition 10% of the # time so the sampler can extrapolate cond vs uncond at export. + tb_raw = tb drop = (torch.rand(tb.shape[0], device=dev) < 0.1).float() tb = tb * (1.0 - drop)[:, None] t = torch.rand(xb.shape[0], device=dev) @@ -252,13 +669,87 @@ def main(): tgt = xb - x0 mask = mb[:, None, :] # [B,1,C6] loss = ((v - tgt) ** 2 * mask).sum() / mask.sum() / T + if a.jitter_weight > 0: + q_all = d6_to_quat( + (xt + (1.0 - t[:, None, None]) * v).reshape(-1, T, J, D6)) + # per-sample band: fast actions may legitimately move ~2.5x a + # walk (real Mixamo run 0.187 vs walk 0.059) + # Bands are set from the DATA's per-action p95, measured on the + # v72 cache. Note march's mean speed (median 0.050) is close to + # WALK's (0.055) — marching in place is not fast by this measure, + # so lumping it with run and raising its FLOOR to 0.10 penalised + # 96% of real march windows. Only genuinely fast actions (run, + # dance) get the raised ceiling, and the floor stays low for all. + is_fast = (tb_raw * fast_mask_v).sum(-1).clamp(0, 1) + hi = 0.10 + is_fast * 0.12 # walk 0.10, fast 0.22 + lo = torch.full_like(hi, 0.035) # low floor for every action + loss = loss + a.jitter_weight * jitter_excess( + q_all, hi=hi, lo=lo).mean() + if a.phase_weight > 0: + # Flow matching predicts a VELOCITY, so reconstruct the clean + # sample the model implies at this t before measuring gait: + # x_t = (1-t)x0 + t*x1 and v* = x1 - x0 => x1 = x_t + (1-t)v + x1_hat = xt + (1.0 - t[:, None, None]) * v + q_hat = d6_to_quat(x1_hat.reshape(-1, T, J, D6)) + ph = gait_phase_corr(q_hat, dirs_t) # [B], want -> +1 + # gate to locomotion rows only (tb is the CFG-dropped token, so + # use the pre-drop labels for the gate) + g = (tb_raw * loco_mask_v).sum(-1).clamp(0, 1) + denom = g.sum().clamp_min(1.0) + # hinge: only penalise below a firm-but-not-saturating target, + # so well-phased samples are left alone + # anatomy guard: free below ~20 deg of hip->chest yaw (the + # data sits at ~6 deg), then linearly penalised. + tw = spine_twist_penalty(q_hat) + tw_pen = (tw - 0.35).clamp_min(0.0) + loss = loss + a.phase_weight * ( + (g * (0.6 - ph).clamp_min(0.0)).sum() / denom) + loss = loss + a.twist_weight * ( + (g * tw_pen).sum() / denom) + if a.period_weight > 0: + # WALK rows only — see prep-t2m-v6's periodicity gate: the + # real Mixamo run has no positive autocorrelation at any + # lag, so demanding it of run/march asks for something real + # running does not have. + walk_j = vocab.index("walk") if "walk" in vocab else -1 + if walk_j >= 0: + gp = tb_raw[:, walk_j].clamp(0, 1) + ap = gait_aperiodicity(q_hat, dirs_t) + loss = loss + a.period_weight * ( + (gp * ap).sum() / gp.sum().clamp_min(1.0)) + if a.legchain_weight > 0: + lc = leg_chain_penalty(q_hat, dirs_t) + loss = loss + a.legchain_weight * ( + (g * lc).sum() / denom) + if a.amp_weight > 0: + # Floors are calibrated to WALK's data p5; march/run have + # legitimately larger excursion (march stride p5 1.174 vs + # walk 0.741), so a global floor penalised 81% of real march + # windows. Apply the floor only to walk rows; the others + # keep the over-drive ceiling. + walk_i = vocab.index("walk") if "walk" in vocab else -1 + if walk_i >= 0: + gw = tb_raw[:, walk_i].clamp(0, 1) + else: + gw = torch.zeros_like(g) + amp_ceil = amplitude_excess(q_hat, dirs_t, lo_scale=0.0) + amp_full = amplitude_excess(q_hat, dirs_t) + amp = amp_ceil + gw * (amp_full - amp_ceil) + loss = loss + a.amp_weight * ( + (g * amp).sum() / denom) + if a.travel_weight > 0: + tv = travel_forward_torch(q_hat, dirs_t) + # hinge toward the data's level (~+0.75); no reward above it + loss = loss + a.travel_weight * ( + (g * (0.5 - tv).clamp_min(0.0)).sum() / denom) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0) opt.step() sched.step() - tot += loss.item(); nb += 1 - print(f"ep {ep + 1}/{a.epochs} loss {tot / max(1, nb):.4f}", flush=True) + tot += loss.detach(); nb += 1 + print(f"ep {ep + 1}/{a.epochs} " + f"loss {tot.item() / max(1, nb):.4f}", flush=True) torch.save({"net": net.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(), "epoch": ep}, ckpt_path) diff --git a/src/MotionGenerator.cpp b/src/MotionGenerator.cpp index 8e862d8a..d654b446 100644 --- a/src/MotionGenerator.cpp +++ b/src/MotionGenerator.cpp @@ -42,7 +42,11 @@ const Syn kSynonyms[] = { {"walking", "walk"}, {"stroll", "walk"}, {"leap", "jump"}, {"hop", "jump"}, {"dancing", "dance"}, {"punching", "punch"}, {"kicking", "kick"}, - {"waving", "wave"}, {"marching", "march"}, {"climbing", "climb"}, + // "marching" -> walk: the march training data was broken (ankle folded + // ~106 deg) and was dropped from the vocab in #837, and no curated march + // clip exists either, so route it to the nearest good gait. + {"waving", "wave"}, {"marching", "walk"}, {"march", "walk"}, + {"climbing", "climb"}, {"idle", "idle"}, {"stand", "idle"}, {"standing", "idle"}, {"sitting", "sit"}, {"seat", "sit"}, {"throwing", "throw"}, {"toss", "throw"}, {"pitch", "throw"}, diff --git a/src/MotionLibrary.cpp b/src/MotionLibrary.cpp index c84b4125..74949544 100644 --- a/src/MotionLibrary.cpp +++ b/src/MotionLibrary.cpp @@ -35,7 +35,14 @@ const Syn kSynonyms[] = { {"walking", "walk"}, {"stroll", "walk"}, {"step", "walk"}, {"jumping", "jump"}, {"hop", "jump"}, {"leap", "jump"}, {"dancing", "dance"}, {"spin", "dance"}, - {"marching", "march"}, + // "march"/"marching" resolve to WALK, not to a `march` action. The curated + // library has never contained a march clip, and the t2m model's march was + // trained on leftover CMU corpus whose ankle folds ~106 deg (the foot + // perpendicular to the shin in 100% of windows) — it rendered as a twisted + // mess and was dropped from the model vocab in #837. With no clip and no + // vocab entry, "marching" resolved to nothing and returned a hard error, so + // point it at the nearest good action instead: a marching gait IS a walk. + {"marching", "walk"}, {"march", "walk"}, {"kicking", "kick"}, {"punching", "punch"}, {"strike", "punch"}, {"hit", "punch"}, {"waving", "wave"}, {"greet", "wave"}, {"hello", "wave"},