初始化项目，由ModelHub XC社区提供模型
Model: EmpathicRobotics/vla-1.7b-qwen3-v2 Source: Original Platform
This commit is contained in:
ModelHub XC
309
pipeline_pose/phase5_adaptive_pchip.py
Normal file
309
pipeline_pose/phase5_adaptive_pchip.py
Normal file
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
Phase 5 — Adaptive PCHIP per-joint tokenizer.
|
||||
|
||||
For each window from Phase 4 (WINDOW_FRAMES frames -- 8 originally, 24 as of
|
||||
2026-07-22, see REPORT.md #38), each of the 17 joints gets an independent
|
||||
PCHIP compression with adaptive control-point count based on per-joint
|
||||
curvature over the WHOLE window (all WINDOW_FRAMES frames considered, not a
|
||||
fixed sub-grid):
|
||||
|
||||
Tier 2 (2 CPs: start + end) — max curvature < tau_low
|
||||
Tier 4 (4 CPs: start + end + top-2 interior) — tau_low <= curvature < tau_high
|
||||
Tier MAX_CPS (MAX_CPS CPs: start+end+top-(MAX_CPS-2) interior, chosen by
|
||||
curvature out of ALL WINDOW_FRAMES candidates) — curvature >= tau_high
|
||||
|
||||
MAX_CPS is fixed at 8 regardless of WINDOW_FRAMES -- widening the window
|
||||
gives the top tier more candidate positions to pick its best 8 from, not more
|
||||
tokens/joint. See MAX_CPS's own docstring below for why this matters.
|
||||
|
||||
Token stream per window (WINDOW_FRAMES=24 example; t values now range over
|
||||
however many of the window's real frame indices got chosen as CPs, e.g.
|
||||
t_0 and t_23 for tier 2, or t_0/t_5/t_14/t_23 for tier 4 if frames 5 and 14
|
||||
had the highest curvature):
|
||||
<fps_30>
|
||||
<pelvis> <pelvis_t_0> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N>
|
||||
<pelvis_t_23> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N> </pelvis>
|
||||
<r_hip> <r_hip_t_0> <r_hip_x_N> ... </r_hip>
|
||||
...
|
||||
|
||||
Quantization: [-2.0 m, +2.0 m] -> [0, 255] (precision ~15.7 mm)
|
||||
Time tokens: real frame index within the window, 0 to WINDOW_FRAMES-1 --
|
||||
requires <{joint}_t_N> tokens up to WINDOW_FRAMES-1 in the
|
||||
tokenizer vocab (only 0-7 existed before 2026-07-22).
|
||||
|
||||
Input: outputs/yolo_cleaned_30fps/{video_id}_cleaned.jsonl
|
||||
Output: outputs/agent_tokens_adaptive/{video_id}_tokens.jsonl
|
||||
Each line: {"video_id", "window_id", "fps", "token_str", "cp_counts"}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
# ── Constants ─────────────────────────────────────────────────────────────────
|
||||
|
||||
TARGET_FPS = 30
|
||||
WINDOW_FRAMES = 8
|
||||
N_JOINTS = 17
|
||||
COORD_RANGE = 2.0
|
||||
STRIDE = 8
|
||||
# 2026-07-22 (REPORT.md #38): cap on control points per joint, independent of
|
||||
# WINDOW_FRAMES. Before this change WINDOW_FRAMES==MAX_CPS==8 always (the top
|
||||
# tier was literally "use every frame"), so widening the window to 24 frames
|
||||
# would have silently tripled worst-case tokens/joint (24 CPs instead of 8)
|
||||
# with no code change needed to trigger it. Keeping MAX_CPS fixed at 8 means
|
||||
# the top tier now means "pick the best 8 of WINDOW_FRAMES candidates by
|
||||
# curvature" instead of "use all of them" -- same worst-case token cost as
|
||||
# before, but the 8 chosen points can be anywhere in the (now wider) window
|
||||
# instead of forced onto a fixed 8-slot grid. This is the whole point of
|
||||
# "Option 2" (dense pose, no subsampling before curve-fitting) agreed with
|
||||
# the user: Phase 3 keeps every real frame, Phase 5 decides freely which
|
||||
# frames matter most.
|
||||
MAX_CPS = 8
|
||||
|
||||
TAU_LOW = 0.005
|
||||
TAU_HIGH = 0.05
|
||||
|
||||
JOINT_NAMES = [
|
||||
"pelvis", "r_hip", "r_knee", "r_ankle",
|
||||
"l_hip", "l_knee", "l_ankle",
|
||||
"spine", "thorax", "nose", "head_top",
|
||||
"l_shoulder", "l_elbow", "l_wrist",
|
||||
"r_shoulder", "r_elbow", "r_wrist",
|
||||
]
|
||||
|
||||
# ── Quantization ──────────────────────────────────────────────────────────────
|
||||
|
||||
def quantize(v: float) -> int:
|
||||
return int(np.clip(round((v + COORD_RANGE) / (2.0 * COORD_RANGE) * 255), 0, 255))
|
||||
|
||||
|
||||
def dequantize(n: int) -> float:
|
||||
return n / 255.0 * (2.0 * COORD_RANGE) - COORD_RANGE
|
||||
|
||||
|
||||
# ── Per-joint adaptive CP selection ───────────────────────────────────────────
|
||||
|
||||
def joint_curvature(trajectory: np.ndarray) -> float:
|
||||
"""Max curvature (acceleration norm) for a single joint trajectory (8, 3)."""
|
||||
if trajectory.shape[0] < 3:
|
||||
return 0.0
|
||||
vel = np.diff(trajectory, axis=0)
|
||||
acc = np.diff(vel, axis=0)
|
||||
return float(np.max(np.linalg.norm(acc, axis=1)))
|
||||
|
||||
|
||||
def select_cp_indices(trajectory: np.ndarray, tau_low: float, tau_high: float) -> np.ndarray:
|
||||
"""Choose which frame indices become control points for one joint.
|
||||
|
||||
Tier sizes are fixed at 2 / 4 / MAX_CPS regardless of how many frames are
|
||||
in the window (see MAX_CPS's docstring) -- the top tier picks the
|
||||
MAX_CPS-2 highest-curvature *interior* frames out of every candidate in
|
||||
the window, not a fixed grid position. With WINDOW_FRAMES==MAX_CPS==8
|
||||
(the original config) this is exactly equivalent to the old
|
||||
`np.arange(WINDOW_FRAMES)` behavior, since "top 6 of 6 interior
|
||||
candidates" is all of them.
|
||||
"""
|
||||
curv = joint_curvature(trajectory)
|
||||
n_frames = trajectory.shape[0]
|
||||
|
||||
if curv < tau_low:
|
||||
return np.array([0, n_frames - 1])
|
||||
|
||||
n_interior = 2 if curv < tau_high else (MAX_CPS - 2)
|
||||
n_interior = min(n_interior, max(n_frames - 2, 0))
|
||||
|
||||
vel = np.diff(trajectory, axis=0)
|
||||
acc = np.diff(vel, axis=0)
|
||||
acc_norms = np.linalg.norm(acc, axis=1) # (n_frames-2,)
|
||||
|
||||
# acc[i] corresponds to frame i+1 (second derivative offset)
|
||||
interior_curv = np.zeros(n_frames)
|
||||
for i in range(len(acc_norms)):
|
||||
interior_curv[i + 1] = acc_norms[i]
|
||||
|
||||
# Exclude endpoints (already included), pick top-n_interior interior frames
|
||||
interior_curv[0] = -1.0
|
||||
interior_curv[-1] = -1.0
|
||||
top_n = np.argsort(interior_curv)[-n_interior:] if n_interior > 0 else np.array([], dtype=int)
|
||||
|
||||
indices = np.unique(np.sort(np.concatenate(([0], top_n, [n_frames - 1]))))
|
||||
return indices.astype(int)
|
||||
|
||||
|
||||
# ── Token builder ─────────────────────────────────────────────────────────────
|
||||
|
||||
def build_token_str(
|
||||
states: np.ndarray,
|
||||
fps: int = TARGET_FPS,
|
||||
tau_low: float = TAU_LOW,
|
||||
tau_high: float = TAU_HIGH,
|
||||
) -> tuple:
|
||||
"""
|
||||
states : (8, 17, 3) float32, root-centred metric coordinates
|
||||
Returns (token_str, cp_counts_dict)
|
||||
"""
|
||||
parts = [f"<fps_{fps}>"]
|
||||
cp_counts = {}
|
||||
|
||||
for j in range(N_JOINTS):
|
||||
name = JOINT_NAMES[j]
|
||||
trajectory = states[:, j, :] # (8, 3)
|
||||
cp_idx = select_cp_indices(trajectory, tau_low, tau_high)
|
||||
cp_counts[name] = len(cp_idx)
|
||||
|
||||
parts.append(f"<{name}>")
|
||||
for fi in cp_idx:
|
||||
x, y, z = trajectory[fi]
|
||||
parts.append(f"<{name}_t_{fi}>")
|
||||
parts.append(f"<{name}_x_{quantize(x)}>")
|
||||
parts.append(f"<{name}_y_{quantize(y)}>")
|
||||
parts.append(f"<{name}_z_{quantize(z)}>")
|
||||
parts.append(f"</{name}>")
|
||||
|
||||
return " ".join(parts), cp_counts
|
||||
|
||||
|
||||
# ── Per-file processing ──────────────────────────────────────────────────────
|
||||
|
||||
def process_file(
|
||||
input_path: str,
|
||||
output_jsonl: str,
|
||||
video_id: str,
|
||||
stride: int = STRIDE,
|
||||
tau_low: float = TAU_LOW,
|
||||
tau_high: float = TAU_HIGH,
|
||||
) -> int:
|
||||
records = []
|
||||
|
||||
with open(input_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
data = json.loads(line)
|
||||
window_id = int(data["window_id"])
|
||||
|
||||
if window_id % stride != 0:
|
||||
continue
|
||||
|
||||
states = np.array(data["states"], dtype=np.float32)
|
||||
|
||||
if states.shape != (WINDOW_FRAMES, N_JOINTS, 3):
|
||||
continue
|
||||
if np.isnan(states).any():
|
||||
continue
|
||||
|
||||
token_str, cp_counts = build_token_str(states, TARGET_FPS, tau_low, tau_high)
|
||||
|
||||
records.append({
|
||||
"video_id": video_id,
|
||||
"window_id": window_id,
|
||||
"fps": TARGET_FPS,
|
||||
"token_str": token_str,
|
||||
"cp_counts": cp_counts,
|
||||
})
|
||||
|
||||
if not records:
|
||||
return 0
|
||||
|
||||
tmp = output_jsonl + ".tmp"
|
||||
os.makedirs(os.path.dirname(output_jsonl), exist_ok=True)
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
for rec in records:
|
||||
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
os.replace(tmp, output_jsonl)
|
||||
|
||||
return len(records)
|
||||
|
||||
|
||||
# ── Entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(
|
||||
description="Phase 5 — Adaptive PCHIP per-joint tokenizer."
|
||||
)
|
||||
p.add_argument("--input-dir", required=True,
|
||||
help="Directory with *_cleaned.jsonl from Phase 4.")
|
||||
p.add_argument("--output-dir", required=True,
|
||||
help="Directory to write *_tokens.jsonl files.")
|
||||
p.add_argument("--stride", type=int, default=STRIDE,
|
||||
help=f"Keep windows where window_id %% stride == 0. Default: {STRIDE}")
|
||||
p.add_argument("--tau-low", type=float, default=TAU_LOW,
|
||||
help=f"Curvature threshold for 2-CP tier. Default: {TAU_LOW}")
|
||||
p.add_argument("--tau-high", type=float, default=TAU_HIGH,
|
||||
help=f"Curvature threshold for 8-CP tier. Default: {TAU_HIGH}")
|
||||
p.add_argument("--file-list", default=None,
|
||||
help="Optional text file listing specific *_cleaned.jsonl paths.")
|
||||
p.add_argument("--window-frames", type=int, default=WINDOW_FRAMES,
|
||||
help=f"Must match Phase 3/4's --window-size. Default: {WINDOW_FRAMES}. "
|
||||
f"2026-07-22: use 24 to match the wider cosmos chunk window -- "
|
||||
f"see REPORT.md #38. MAX_CPS (top-tier control-point count) is "
|
||||
f"NOT tied to this and stays 8 either way.")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global WINDOW_FRAMES
|
||||
args = parse_args()
|
||||
WINDOW_FRAMES = args.window_frames
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.file_list:
|
||||
with open(args.file_list) as f:
|
||||
all_files = [l.strip() for l in f if l.strip()]
|
||||
else:
|
||||
all_files = sorted(glob.glob(os.path.join(args.input_dir, "*_cleaned.jsonl")))
|
||||
|
||||
task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
|
||||
num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 1))
|
||||
|
||||
my_files = [f for i, f in enumerate(all_files) if i % num_tasks == task_id]
|
||||
total = len(my_files)
|
||||
|
||||
print(f"\n[Worker {task_id}/{num_tasks}] {total} files to process.")
|
||||
print("=" * 60)
|
||||
|
||||
processed = skipped = empty = 0
|
||||
tier_counts = {2: 0, 4: 0, 8: 0}
|
||||
|
||||
for idx, input_path in enumerate(my_files, start=1):
|
||||
base = os.path.basename(input_path)
|
||||
video_id = base[: -len("_cleaned.jsonl")]
|
||||
|
||||
out_jsonl = os.path.join(args.output_dir, f"{video_id}_tokens.jsonl")
|
||||
|
||||
if os.path.exists(out_jsonl):
|
||||
skipped += 1
|
||||
print(f"[{idx}/{total}] {video_id} — already done", end="\r")
|
||||
continue
|
||||
|
||||
try:
|
||||
n = process_file(
|
||||
input_path, out_jsonl, video_id,
|
||||
stride=args.stride,
|
||||
tau_low=args.tau_low,
|
||||
tau_high=args.tau_high,
|
||||
)
|
||||
if n > 0:
|
||||
processed += 1
|
||||
pct = (processed + skipped + empty) / total * 100
|
||||
print(f"[{idx}/{total}] {pct:.1f}% | {video_id} — {n} windows")
|
||||
else:
|
||||
empty += 1
|
||||
except Exception as e:
|
||||
print(f"[{idx}/{total}] ERROR {video_id} — {e}")
|
||||
for p in (out_jsonl + ".tmp",):
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"[Worker {task_id}] done — processed: {processed}, skipped: {skipped}, empty: {empty}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user