269 lines
9.1 KiB
Python
269 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Decode adaptive PCHIP agent tokens into 3D pose trajectories.
|
|
|
|
Takes the raw token string that the VLA model produces and reconstructs
|
|
the full (n_frames, 17, 3) skeleton trajectory via PCHIP interpolation.
|
|
|
|
Usage:
|
|
# Decode tokens from stdin
|
|
echo "<fps_30> <pelvis> <pelvis_t_0> ..." | python decode_agent_tokens.py
|
|
|
|
# Decode from a text file
|
|
python decode_agent_tokens.py --input generated_tokens.txt
|
|
|
|
# Decode and save JSON
|
|
python decode_agent_tokens.py --input tokens.txt --output poses.json
|
|
|
|
Token format (one 8-frame window):
|
|
<fps_30>
|
|
<pelvis> <pelvis_t_0> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N>
|
|
<pelvis_t_7> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N> </pelvis>
|
|
<r_hip> <r_hip_t_0> <r_hip_x_N> <r_hip_y_N> <r_hip_z_N> ... </r_hip>
|
|
...17 joints...
|
|
|
|
Dequantization: coord = N / 255.0 * 4.0 - 2.0 (metres, range [-2, 2])
|
|
Time tokens: frame index 0-7 within the 8-frame window
|
|
Reconstruction: PCHIP interpolation over control points -> 8 frames
|
|
|
|
H36M joint order (17 joints):
|
|
0 pelvis 4 l_hip 8 thorax 12 l_elbow 16 r_wrist
|
|
1 r_hip 5 l_knee 9 nose 13 l_wrist
|
|
2 r_knee 6 l_ankle 10 head_top 14 r_shoulder
|
|
3 r_ankle 7 spine 11 l_shoulder 15 r_elbow
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import sys
|
|
|
|
import numpy as np
|
|
from scipy.interpolate import PchipInterpolator
|
|
|
|
WINDOW_FRAMES = 8
|
|
COORD_RANGE = 2.0
|
|
|
|
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",
|
|
]
|
|
JOINT_INDEX = {name: i for i, name in enumerate(JOINT_NAMES)}
|
|
N_JOINTS = len(JOINT_NAMES)
|
|
|
|
|
|
def dequantize(n: int) -> float:
|
|
return n / 255.0 * (2.0 * COORD_RANGE) - COORD_RANGE
|
|
|
|
|
|
def parse_window(tokens: list[str]) -> dict:
|
|
"""Parse tokens for a single 8-frame window into per-joint control points.
|
|
|
|
Args:
|
|
tokens: list of token strings like ['<fps_30>', '<pelvis>', '<pelvis_t_0>', ...]
|
|
|
|
Returns:
|
|
dict with fps (int) and joints (dict mapping joint name to
|
|
t_indices ndarray and cp_coords ndarray of shape (n_cp, 3)).
|
|
"""
|
|
fps = 30
|
|
if tokens and tokens[0].startswith("<fps_"):
|
|
fps = int(re.match(r"<fps_(\d+)>", tokens[0]).group(1))
|
|
tokens = tokens[1:]
|
|
|
|
joints = {}
|
|
i = 0
|
|
while i < len(tokens):
|
|
m = re.match(r"^<([a-z_]+)>$", tokens[i])
|
|
if not m or m.group(1) not in JOINT_INDEX:
|
|
i += 1
|
|
continue
|
|
|
|
name = m.group(1)
|
|
close = f"</{name}>"
|
|
i += 1
|
|
t_indices = []
|
|
coords = []
|
|
|
|
while i < len(tokens) and tokens[i] != close:
|
|
tm = re.match(rf"<{name}_t_(\d+)>$", tokens[i])
|
|
if tm and i + 3 < len(tokens):
|
|
t_indices.append(int(tm.group(1)))
|
|
xm = re.match(rf"<{name}_x_(\d+)>$", tokens[i + 1])
|
|
ym = re.match(rf"<{name}_y_(\d+)>$", tokens[i + 2])
|
|
zm = re.match(rf"<{name}_z_(\d+)>$", tokens[i + 3])
|
|
if xm and ym and zm:
|
|
coords.append([
|
|
dequantize(int(xm.group(1))),
|
|
dequantize(int(ym.group(1))),
|
|
dequantize(int(zm.group(1))),
|
|
])
|
|
i += 4
|
|
continue
|
|
i += 1
|
|
|
|
if i < len(tokens) and tokens[i] == close:
|
|
i += 1
|
|
|
|
if t_indices:
|
|
joints[name] = {
|
|
"t_indices": np.array(t_indices, dtype=int),
|
|
"cp_coords": np.array(coords, dtype=np.float32),
|
|
}
|
|
|
|
return {"fps": fps, "joints": joints}
|
|
|
|
|
|
def reconstruct(parsed: dict) -> np.ndarray:
|
|
"""PCHIP-interpolate sparse control points into a full per-frame trajectory.
|
|
|
|
Window length is inferred from the data itself (max t_index + 1 seen
|
|
across all joints) rather than a fixed constant -- 2026-07-22 (REPORT.md
|
|
#38): windows are now 24 frames (t up to 23), but this decoder is also
|
|
used on older 8-frame data (t up to 7), so it must handle both without
|
|
being told which convention a given token string uses.
|
|
|
|
Returns ndarray of shape (window_frames, 17, 3) in metres, root-centred.
|
|
"""
|
|
max_t = WINDOW_FRAMES - 1
|
|
for jdata in parsed["joints"].values():
|
|
if len(jdata["t_indices"]):
|
|
max_t = max(max_t, int(jdata["t_indices"].max()))
|
|
window_frames = max_t + 1
|
|
|
|
t_out = np.arange(window_frames, dtype=np.float64)
|
|
traj = np.zeros((window_frames, N_JOINTS, 3), dtype=np.float32)
|
|
|
|
for name, jdata in parsed["joints"].items():
|
|
j = JOINT_INDEX[name]
|
|
t_cp = jdata["t_indices"].astype(np.float64)
|
|
cp = jdata["cp_coords"]
|
|
|
|
if len(t_cp) < 2:
|
|
traj[:, j, :] = cp[0]
|
|
continue
|
|
|
|
for d in range(3):
|
|
traj[:, j, d] = PchipInterpolator(t_cp, cp[:, d])(t_out)
|
|
|
|
return traj
|
|
|
|
|
|
def decode(token_str: str) -> list[np.ndarray]:
|
|
"""Decode a token string into a list of (8, 17, 3) trajectories.
|
|
|
|
Handles both single windows and multiple consecutive windows.
|
|
"""
|
|
all_tokens = re.findall(r"<[^>]+>", token_str)
|
|
if not all_tokens:
|
|
return []
|
|
|
|
# Split on <fps_N> boundaries — each is one window
|
|
window_starts = [i for i, t in enumerate(all_tokens) if t.startswith("<fps_")]
|
|
|
|
if not window_starts:
|
|
parsed = parse_window(all_tokens)
|
|
return [reconstruct(parsed)]
|
|
|
|
trajectories = []
|
|
for wi, start in enumerate(window_starts):
|
|
end = window_starts[wi + 1] if wi + 1 < len(window_starts) else len(all_tokens)
|
|
parsed = parse_window(all_tokens[start:end])
|
|
trajectories.append(reconstruct(parsed))
|
|
|
|
return trajectories
|
|
|
|
|
|
def to_json(trajectories: list[np.ndarray], fps: int = 30) -> dict:
|
|
"""Convert decoded trajectories to a JSON-serialisable dict.
|
|
|
|
Per-window frame count is read from each trajectory's own shape (not a
|
|
fixed constant) -- see reconstruct()'s docstring, 2026-07-22."""
|
|
windows = []
|
|
cum_frames = 0
|
|
for i, traj in enumerate(trajectories):
|
|
n_frames_this = traj.shape[0]
|
|
motion = np.linalg.norm(traj[-1] - traj[0], axis=-1)
|
|
top_movers = sorted(
|
|
[(JOINT_NAMES[j], round(float(motion[j]), 4)) for j in range(N_JOINTS)],
|
|
key=lambda x: x[1], reverse=True,
|
|
)
|
|
n_missing = sum(1 for name in JOINT_NAMES if name not in
|
|
{JOINT_NAMES[j] for j in range(N_JOINTS) if np.any(traj[:, j, :] != 0)})
|
|
|
|
windows.append({
|
|
"window": i,
|
|
"time_sec": round(cum_frames / fps, 4),
|
|
"trajectory": traj.tolist(),
|
|
"value_range_m": [round(float(traj.min()), 4), round(float(traj.max()), 4)],
|
|
"top_movers": top_movers[:5],
|
|
"joints_all_zero": n_missing,
|
|
})
|
|
cum_frames += n_frames_this
|
|
|
|
total_frames = sum(traj.shape[0] for traj in trajectories)
|
|
return {
|
|
"n_windows": len(trajectories),
|
|
"total_frames": total_frames,
|
|
"duration_sec": round(total_frames / fps, 4),
|
|
"shape": [len(trajectories)] + list(trajectories[0].shape) if trajectories else [],
|
|
"value_range_m": [round(float(min(t.min() for t in trajectories)), 4),
|
|
round(float(max(t.max() for t in trajectories)), 4)] if trajectories else [0, 0],
|
|
"joint_names": JOINT_NAMES,
|
|
"windows": windows,
|
|
}
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description="Decode agent tokens to 3D poses.")
|
|
p.add_argument("--input", "-i", default=None,
|
|
help="File containing agent tokens (default: read stdin)")
|
|
p.add_argument("--output", "-o", default=None,
|
|
help="Save decoded poses to JSON file")
|
|
args = p.parse_args()
|
|
|
|
if args.input:
|
|
with open(args.input, "r") as f:
|
|
token_str = f.read()
|
|
else:
|
|
token_str = sys.stdin.read()
|
|
|
|
token_str = token_str.strip()
|
|
if not token_str:
|
|
print("No tokens provided.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
trajectories = decode(token_str)
|
|
if not trajectories:
|
|
print("Could not parse any agent windows from input.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
result = to_json(trajectories)
|
|
|
|
print(f"Decoded {result['n_windows']} windows "
|
|
f"({result['total_frames']} frames, {result['duration_sec']}s)")
|
|
print(f"Shape: {result['shape']} (windows, frames, joints, xyz)")
|
|
print(f"Value range: {result['value_range_m']} m")
|
|
|
|
for w in result["windows"][:3]:
|
|
print(f"\n Window {w['window']} (t={w['time_sec']}s):")
|
|
if w["joints_all_zero"] > 0:
|
|
print(f" WARNING: {w['joints_all_zero']} joints are all-zero (missing)")
|
|
print(f" Top movers: ", end="")
|
|
print(", ".join(f"{name} {d:.3f}m" for name, d in w["top_movers"]))
|
|
|
|
if result["n_windows"] > 3:
|
|
print(f"\n ... {result['n_windows'] - 3} more windows")
|
|
|
|
if args.output:
|
|
with open(args.output, "w") as f:
|
|
json.dump(result, f, indent=2)
|
|
print(f"\nSaved to: {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|