File size: 6,387 Bytes
eedc275 acdece1 eedc275 acdece1 eedc275 acdece1 eedc275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""Probe: which input protocol makes DreamZero (dreamzero/vla) produce sane actions on T-Rex data?
A single current frame at an arbitrary anchor (what openloop_compare.py did)
B sequential causal rollout from frame 0 with teacher-forced observations: chunk 0 = frame 0,
chunk j = frames [a-23, a-16, a-8, a] with a = 24*j (the deployment path of inference_trex.py)
C gt_video_action_pred: 25 real frames [a .. a+24] as (teacher-forced) video
For every protocol the 24 predicted joint actions are scored against the recorded targets
action[a .. a+23]: RMSE per group at k=0 and over the chunk, plus the copy-current-state baseline.
"""
import os, sys, time, json
from pathlib import Path
import numpy as np, torch
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "scripts"))
import inference_trex as it
import torch.distributed as dist
from groot.vla.data.schema import EmbodimentTag
from groot.vla.model.n1_5.sim_policy import GrootSimPolicy
G = {"left_arm": slice(0, 7), "left_hand": slice(7, 29), "right_arm": slice(29, 36), "right_hand": slice(36, 58)}
KEYS = ("action.left_arm", "action.left_hand", "action.right_arm", "action.right_hand")
DIM = {"action.left_arm": 7, "action.left_hand": 22, "action.right_arm": 7, "action.right_hand": 22}
def act58(result):
a = result.act
return np.concatenate([np.asarray(a[k], dtype=np.float64).reshape(-1, DIM[k]) for k in KEYS], -1)
def score(pred, gt, state):
out = {}
n = min(len(pred), len(gt))
for g, s in G.items():
out[g] = (float(np.sqrt(((pred[0, s] - gt[0, s]) ** 2).mean())), float(np.sqrt(((pred[:n, s] - gt[:n, s]) ** 2).mean())),
float(np.sqrt(((state[s] - gt[:n, s]) ** 2).mean()))) # k0, chunk, copy-state baseline (chunk)
return out
def fmt(sc):
return " ".join(f"{g}: k0 {v[0]:.3f} / chunk {v[1]:.3f} (hold {v[2]:.3f})" for g, v in sc.items())
def reset(policy):
ah = policy.trained_model.action_head
ah.current_start_frame = 0
ah.language = None
def main():
root = sys.argv[1] if len(sys.argv) > 1 else "data/trex_small_force"
episodes = [int(x) for x in (sys.argv[2] if len(sys.argv) > 2 else "0,45").split(",")]
chunks = [int(x) for x in (sys.argv[3] if len(sys.argv) > 3 else "1,2,3,4,5").split(",")]
if not dist.is_initialized():
os.environ.setdefault("MASTER_ADDR", "localhost"); os.environ.setdefault("MASTER_PORT", str(29600 + os.getpid() % 300))
dist.init_process_group(backend="gloo", world_size=1, rank=0)
policy = GrootSimPolicy(embodiment_tag=EmbodimentTag.TREX,
model_path="checkpoints/dreamzero_trex_wan22_lora/checkpoint-16000", device="cuda:0")
H, W = it.get_expected_video_resolution(policy)
ah = policy.trained_model.action_head
if os.environ.get("DIT_ALL_STEPS") == "1" and hasattr(ah, "dit_step_mask"):
ah.dit_step_mask = [True] * len(ah.dit_step_mask)
print(f"DIT_ALL_STEPS: running all {len(ah.dit_step_mask)} denoising steps", flush=True)
if os.environ.get("PROBE_ONLY_A") == "1":
chunks_only_a = True
else:
chunks_only_a = False
print(f"video {H}x{W} num_frames={getattr(ah,'num_frames',None)} num_frame_per_block={getattr(ah,'num_frame_per_block',None)} action_horizon={getattr(ah,'action_horizon',None)}", flush=True)
results = {}
for ep in episodes:
episode = it.TrexEpisode(root, ep)
task = episode.get_task(0)
actions = np.stack([np.array(episode.table.column("action")[i].as_py()) for i in range(episode.length)])
states = np.stack([episode.get_state(i) for i in range(episode.length)])
chunks_ep = [j for j in chunks if 24 * j + 24 <= episode.length]
# ---- A: single frame at the chunk boundary
for j in chunks_ep:
a = 24 * j
reset(policy)
with torch.inference_mode():
res, _ = policy.lazy_joint_forward_causal(it.Batch(obs=it.build_obs(episode, [a], task, H, W)))
sc = score(act58(res), actions[a:a + 24], states[a])
results[(ep, j, "A_single")] = sc; print(f"ep{ep} chunk{j} A_single {fmt(sc)}", flush=True)
# ---- B: sequential causal rollout from frame 0 (teacher-forced frames)
if chunks_only_a:
continue
reset(policy)
sched = it.build_chunk_schedule(episode.length, max(chunks_ep) + 1)
for j, frames in enumerate(sched):
with torch.inference_mode():
res, _ = policy.lazy_joint_forward_causal(it.Batch(obs=it.build_obs(episode, frames, task, H, W)))
a = frames[-1]
if j in chunks_ep or j == 0:
sc = score(act58(res), actions[a:a + 24], states[a])
results[(ep, j, "B_sequential")] = sc; print(f"ep{ep} chunk{j} B_seq frames={frames} {fmt(sc)} start_frame={ah.current_start_frame}", flush=True)
# ---- C: teacher-forced future video
for j in chunks_ep[:2]:
a = 24 * j
try:
reset(policy)
obs = it.build_obs(episode, list(range(a, a + 25)), task, H, W)
with torch.inference_mode():
res, _ = policy.gt_video_action_pred(it.Batch(obs=obs))
sc = score(act58(res), actions[a:a + 24], states[a])
results[(ep, j, "C_gtvideo")] = sc; print(f"ep{ep} chunk{j} C_gtvideo {fmt(sc)}", flush=True)
except Exception as e:
print(f"ep{ep} chunk{j} C_gtvideo FAILED: {type(e).__name__}: {str(e)[:200]}", flush=True)
break
print("\n==== mean over probed chunks (k0 RMSE / chunk RMSE), rad")
for proto in ("A_single", "B_sequential", "C_gtvideo"):
rows = [v for (e, j, p), v in results.items() if p == proto and j > 0]
if not rows: continue
print(f"{proto:13s} " + " ".join(f"{g}: {np.mean([r[g][0] for r in rows]):.3f}/{np.mean([r[g][1] for r in rows]):.3f}" for g in G))
rows = [v for (e, j, p), v in results.items() if p == "A_single"]
print(f"{'hold-state':13s} " + " ".join(f"{g}: {np.mean([r[g][2] for r in rows]):.3f}" for g in G))
json.dump({f"{e}|{j}|{p}": v for (e, j, p), v in results.items()}, open(os.environ.get("PROBE_OUT","ablation_runs/openloop_compare/dz_probe.json"), "w"), indent=1)
if __name__ == "__main__":
main()
|