145 lines
8.4 KiB
Python
145 lines
8.4 KiB
Python
"""老师轨迹回放自检（train 阶段开训前自动跑，约 1–3 分钟）：
|
||
python -m ml.eval.replay_teacher --n 512 --out ml/runs/replay_teacher.json
|
||
两部分，都过才 pass（pipeline.sh 读 pass 字段与指纹，不过就停 rc 5）：
|
||
1. 运动学口径（2026-10-18）：gen 写出的老师 50 Hz 轨迹（keep=true、cat≠D）直接放进 PyTorch 运动学 / 碰撞 / 到达判据：
|
||
到达率 ≥ 0.95、碰撞率 ≤ 0.01。验证的是「仿真口径（kin.json、间隙、到达门槛）与 JS 认证过的数据一致」。
|
||
2. 滚动口径（2026-10-19，审查 review_closedloop F1/D1、review_data_actions C1）：上面那条完全不经过 rollout()，
|
||
队列 / 前缀 / 提交 3 步 / qe / 标签 clamp / DLS / done / 时限上的错位它都发现不了。这里让老师当策略（closed_loop.TeacherOracle：
|
||
每次规划从当前时间下标取老师 16 步增量，按标签口径 clamp）原样走 rollout()（select='first'、use_sub=False）：
|
||
到达（时限内）≥ 0.95、碰撞 ≤ 0.01、DLS 接管前与老师轨迹的最大关节偏差 p99 ≤ 0.1°（shard 轨迹按 1e-3° 取整，标签 clamp 会截掉取整的余量）。
|
||
输出另带 fp（当前 kin.json）、data_fp（shard 里记录的碰撞模型指纹）、ver=2；pipeline 三者都核对。
|
||
load_teacher / pad_traj / oracle_raw / oracle_summary 供 ml/eval/diag_cl.py 的 H1 段复用。"""
|
||
import argparse, glob, json, math, os, sys
|
||
import torch
|
||
|
||
from ml.common.kin_torch import Kin
|
||
from ml.sim.closed_loop import goal_err, rollout, TeacherOracle, DT
|
||
|
||
DEV_TOL_DEG = 0.1
|
||
|
||
|
||
def load_teacher(shards, ids):
|
||
"""按任务 id 从 shard jsonl 取老师记录（traj、path、exec_s…）。只解析 id 命中的行。"""
|
||
want, out = set(ids), {}
|
||
for f in sorted(glob.glob(shards)):
|
||
for line in open(f, encoding='utf-8'):
|
||
k = line.find('"id":"')
|
||
if k < 0: continue
|
||
rid = line[k + 6:line.find('"', k + 6)]
|
||
if rid in want and rid not in out:
|
||
r = json.loads(line)
|
||
if r.get('traj'): out[rid] = r
|
||
if len(out) == len(want): break
|
||
return out
|
||
|
||
|
||
def pad_traj(trajs, device, extra=20):
|
||
"""[B,Tmax+extra,6]：各题老师轨迹末帧重复补齐（oracle 在轨迹末尾之后给零增量）。"""
|
||
Tm = max(len(t) for t in trajs) + extra
|
||
out = torch.zeros(len(trajs), Tm, 6, device=device)
|
||
for i, t in enumerate(trajs):
|
||
x = torch.tensor(t, dtype=torch.float32, device=device)
|
||
out[i, :len(t)] = x; out[i, len(t):] = x[-1]
|
||
return out
|
||
|
||
|
||
def oracle_raw(per, TR, lens, tasks):
|
||
"""oracle 滚动的逐题量：DLS 接管前（且到达前、老师轨迹结束前）与老师的最大关节偏差 mdev（度）、时限内到达 ok、
|
||
时限内碰撞 col、是否计入到达统计（cat≠D）、完成时刻 − 老师时长 dt（s）、是否用过 DLS、2 s 目标误差比。"""
|
||
tr = per['traj'].to(TR.device); B, Tn, _ = tr.shape
|
||
L = min(Tn, TR.shape[1])
|
||
dev = (tr[:, :L] - TR[:, :L]).abs().max(-1).values # [B,L]
|
||
k = torch.arange(L, device=TR.device)[None]
|
||
kd = torch.tensor(lens, device=TR.device)[:, None] - 1
|
||
dls = per['dls_it'].to(TR.device)[:, None]
|
||
td = per['t_done'].to(TR.device)
|
||
kdone = torch.where(torch.isfinite(td), (td / DT).round(), torch.full_like(td, 1e9)).long()[:, None]
|
||
valid = (k <= kd) & ((dls < 0) | (k <= dls)) & (k <= kdone)
|
||
ex = torch.tensor([t.get('exec_s') or float('nan') for t in tasks])
|
||
return {'mdev': torch.where(valid, dev, torch.zeros_like(dev)).max(-1).values.cpu(), 'ok': per['ok'], 'col': per['col'] > 0,
|
||
'reach_m': torch.tensor([t.get('cat') != 'D' for t in tasks]), 'dt': per['t_done'] - ex,
|
||
'dls': per['dls_it'] >= 0, 'r2': per['goal_ratio_2s']}
|
||
|
||
|
||
def _q(v, p):
|
||
v = v[torch.isfinite(v)]
|
||
return float(v.kthvalue(max(1, min(len(v), int(math.ceil(p * len(v)))))).values) if len(v) else None
|
||
|
||
|
||
def oracle_summary(raws):
|
||
c = {k: torch.cat([r[k] for r in raws]) for k in raws[0]}
|
||
m = c['reach_m']; ok = c['ok']
|
||
dt = c['dt'][ok & m]
|
||
return {'n': int(len(m)), 'n_reach_tasks': int(m.sum()),
|
||
'reach': float(ok[m].float().mean()) if m.any() else None,
|
||
'collision': float(c['col'][m].float().mean()) if m.any() else None,
|
||
'D': {'n': int((~m).sum()), 'collision': float(c['col'][~m].float().mean()) if (~m).any() else None},
|
||
'max_dev_deg_before_dls': {'p50': _q(c['mdev'], .5), 'p99': _q(c['mdev'], .99), 'max': float(c['mdev'].max())},
|
||
'dls_used': float(c['dls'].float().mean()),
|
||
'finish_minus_teacher_s': {'p50': _q(dt, .5), 'p90': _q(dt, .9), 'max': float(dt.max()) if len(dt) else None},
|
||
'goal_ratio_2s_p50': _q(c['r2'], .5)}
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--shards', default='ml/data/shard_*.jsonl'); ap.add_argument('--n', type=int, default=512)
|
||
ap.add_argument('--samples', default='ml/data/samples', help='读打包 meta 的 act_scale（oracle 按同一标签口径 clamp）；没有就按 1.0')
|
||
ap.add_argument('--bs', type=int, default=128)
|
||
ap.add_argument('--out', default='ml/runs/replay_teacher.json')
|
||
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
|
||
a = ap.parse_args()
|
||
kin = Kin(device=a.device)
|
||
recs = []
|
||
for f in sorted(glob.glob(a.shards)):
|
||
for line in open(f, encoding='utf-8'):
|
||
r = json.loads(line)
|
||
if r.get('keep') and r.get('traj') and r.get('cat') != 'D': recs.append(r)
|
||
if len(recs) >= a.n: break
|
||
if len(recs) >= a.n: break
|
||
if not recs: sys.exit('没有可回放的老师轨迹（先跑 gen）')
|
||
# ---- 1. 运动学口径
|
||
reach = col = 0; worst_gap = []
|
||
for r in recs:
|
||
q = torch.tensor(r['traj'], dtype=torch.float32, device=a.device)
|
||
mid = 0.5 * (q[1:] + q[:-1])
|
||
g = kin.gaps(torch.cat([q, mid], 0), 0.5)['gm'].min(-1).values
|
||
hit = bool((g <= 0).any()); col += hit; worst_gap.append(g.min().item())
|
||
T = torch.tensor([r['T']], dtype=torch.float32, device=a.device)
|
||
R = torch.tensor([r['R'] if r.get('R') else [0.] * 9], dtype=torch.float32, device=a.device)
|
||
e, oe = goal_err(kin, q[-1:], T, R)
|
||
reach += bool((e < 1).item() and (oe < 0.5).item())
|
||
n = len(recs)
|
||
kin_ok = reach / n >= 0.95 and col / n <= 0.01
|
||
# ---- 2. 滚动口径：老师 oracle 走完整 rollout()
|
||
try:
|
||
from ml.data import read_meta
|
||
act_scale = float(read_meta(a.samples).get('act_scale', 1.0)) if os.path.isdir(a.samples) else 1.0
|
||
except SystemExit:
|
||
act_scale = 1.0
|
||
tasks = [{'q0': r['q0'], 'T': r['T'], 'R': r.get('R'), 'exec_s': r.get('exec_s') or (len(r['traj']) - 1) * DT,
|
||
'travel': r.get('travel_deg'), 'cat': r.get('cat')} for r in recs]
|
||
raws = []
|
||
for i in range(0, n, a.bs):
|
||
tb, rb = tasks[i:i + a.bs], recs[i:i + a.bs]
|
||
TR = pad_traj([r['traj'] for r in rb], a.device)
|
||
_, per = rollout(None, kin, tb, a.device, K=1, select='first', use_sub=False, return_traj=True, seed=i,
|
||
oracle=TeacherOracle(TR, act_scale=act_scale))
|
||
raws.append(oracle_raw(per, TR, [len(r['traj']) for r in rb], tb))
|
||
orc = oracle_summary(raws); orc['act_scale'] = act_scale
|
||
orc_ok = (orc['reach'] or 0) >= 0.95 and (orc['collision'] or 0) <= 0.01 and orc['max_dev_deg_before_dls']['p99'] <= DEV_TOL_DEG
|
||
data_fps = {r.get('fp') for r in recs}
|
||
data_fp = data_fps.pop() if len(data_fps) == 1 else 'mixed'
|
||
res = {'ver': 2, 'n': n, 'reach': reach / n, 'collision': col / n, 'min_gap_mm_p10': sorted(worst_gap)[n // 10], 'kin_pass': kin_ok,
|
||
'oracle': orc, 'oracle_pass': orc_ok, 'fp': kin.fp, 'data_fp': data_fp,
|
||
'pass': bool(kin_ok and orc_ok and data_fp in (None, kin.fp))}
|
||
os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
|
||
json.dump(res, open(a.out, 'w', encoding='utf-8'), indent=1, ensure_ascii=False)
|
||
print('老师轨迹回放：到达 %.1f%%、碰撞 %.1f%%（%d 条）；oracle 走 rollout：到达 %.1f%%、碰撞 %.1f%%、DLS 前最大偏差 p99 %.4f°；指纹 %s' % (
|
||
100 * res['reach'], 100 * res['collision'], n, 100 * orc['reach'], 100 * orc['collision'],
|
||
orc['max_dev_deg_before_dls']['p99'], '一致' if data_fp in (None, kin.fp) else '不符'), flush=True)
|
||
sys.exit(0 if res['pass'] else 5)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|