"""预训练（方案 §6）：MeanFlow 一步头 + FK 损失 + 平滑 + 价值头 + 子目标头。 python -m ml.train --samples ml/data/samples --out ml/runs/base [--config ml/configs/base.json] 吞吐摸底：python -m ml.train --samples ml/data/samples --out ml/runs/bench_base --steps 2000 --no-eval 通过线（方案 R10）：samples_per_s ≥ 5000；日志 eta_h 为剩余小时数。 每 eval_every 步用 EMA 权重做一次闭环验证（ml/sim/closed_loop.py，部署节奏），按选模分数存 best.pt；min_step 之后连续 patience_evals 次不升即停。 2026-10-18 首次正式预训练（21850 步，闭环到达 0–0.2%、穿柱 77–97%）后改： - FK / 平滑损失不再加在 one_step(随机噪声) 上（那样等于对任意噪声回归同一个 Y → 条件均值，多峰头塌成单峰）， 改加在 MF 分支里 r==t 且 t ≤ fk_tmax 的样本的配对去噪估计 x̂ = z − t·u 上（见 ml/model/meanflow.py）； - p_neq 前 p_neq_warm 步为 0（纯 flow matching 先把瞬时速度场学稳，再开 r≠t 的自举目标）； - 每次评估加开环指标（val 样本上一步采样的 16 步末端 FK 误差 mm、K 个候选的多样性）， 塌缩或学偏当场可见。 2026-10-19（审查 review_model_train P0-2/P0-3、review_data_actions C3）： - EMA（train.ema，默认 0.999，前期按 (1+s)/(10+s) 放宽）：评估、best/last 里的 'ema' 都是 EMA 权重；下游一律经 ckpt.load_policy 取 ck.get('ema', ck['model'])； - 选模：p_neq 预热期内及预热后第一次评估之前不存 best.pt（那时一步映射 u(z,0,1) 还没训过）； 分数 = success + 0.05·(1 − collision_rate) + 0.05·(1 − min(1, goal_ratio_2s)) − 1e-4·min(ol_best_mm, 500)，success 全 0 时闭环指标仍能排序； 早停只在 min_step = max(预热 + 2·eval_every, steps/2) 之后计数； - 余弦学习率按真实预算走完（configs 的 steps 改为 24000，原 60000 从来跑不到，14k 步早停时学习率还有 0.87×）； - MeanFlow 采样：p_neq 0.5、p_edge 0.1（固定 (r=0,t=1) 的部署点，见 meanflow.sample_rt）； - 子目标头只看 sub 特征（68:75）置零的输入：sub 存在时 SG[0:6] 与输入 sd 是同一个量，原来 80% 的样本可以直接抄（标签泄露）； 运行时子目标头只在冷启动 sf=0 时用，这样训练与使用同口径。价值头仍用完整输入（稳态推理时带 sub）； - 条件输入增广 aug_pre_drop / aug_pre_sigma / aug_dq_sigma（见 ml/data.Samples.augment）； - 动作尺度 act_scale 取自数据（ml/data.read_meta），存进 checkpoint；续训 / 微调起点与数据不一致直接报错。 - DAgger 微调（--init）：只有起点 checkpoint 已走过自己的 p_neq 预热时才把预热置 0；起点优先取 EMA 权重。""" import argparse, copy, json, math, os, sys, time import torch from ml.common.kin_torch import Kin from ml.common.features import STEP, D2R from ml.common.ckpt import Saver, atomic_save, rng_state, set_rng_state from ml.common import vram from ml.data import Samples, load_tasks from ml.model.policy import build, wta_loss from ml.model.meanflow import meanflow_loss from ml.sim.closed_loop import evaluate def fk_points(kin, q, idx): """选定胶囊端点 + TCP 的世界坐标 [B, 2*len(idx)+1, 3]""" xf = kin.cum_xf(q) a, b = kin.caps_world(q, xf) return torch.cat([a[:, idx], b[:, idx], kin.fk(q, xf)[:, None]], 1) def q_from_feat(X): s, c = X[:, 0:12:2], X[:, 1:12:2] return torch.atan2(s, c) / D2R @torch.no_grad() def open_loop(model, kin, vdata, stepv, fk_idx, n=1024, K=8): """stepv：每单位动作对应的度数（= act_scale·STEP）。""" """开环指标（val 样本，老师状态上）： ol_fk_mm 候选 0 的 16 步末端与老师 Y 末端的胶囊点平均距离（mm）； ol_best_mm K 个候选里最接近老师的那条的同一误差（多峰时它应明显小于 ol_fk_mm）； ol_div K 个候选间标准差 / 动作平均幅度（塌成单峰时接近 0；首次正式预训练 last.pt 实测 0.20）。""" g = torch.Generator(device=vdata.X.device); g.manual_seed(123) X, Y, _, _ = vdata.batch(min(n, vdata.n), gen=g) act, _, _ = model.sample_k(X, K) act = act.float().clamp(-1, 1) q = q_from_feat(X) qt = q + (Y * stepv).sum(1) qh = q[:, None] + (act * stepv).sum(2) # [B,K,6] pt = fk_points(kin, qt, fk_idx)[:, None] ph = fk_points(kin, qh.reshape(-1, 6), fk_idx).view(X.shape[0], K, -1, 3) err = (ph - pt).norm(dim=-1).mean(-1) # [B,K] mm div = (act.std(1).mean() / act.abs().mean().clamp(min=1e-6)).item() return {'ol_fk_mm': err[:, 0].mean().item(), 'ol_best_mm': err.min(1).values.mean().item(), 'ol_div': div} def main(): ap = argparse.ArgumentParser() ap.add_argument('--samples', nargs='+', default=['ml/data/samples']); ap.add_argument('--out', default='ml/runs/base') ap.add_argument('--config', default='ml/configs/base.json'); ap.add_argument('--steps', type=int, default=None) ap.add_argument('--resume', default=None); ap.add_argument('--no-eval', action='store_true') ap.add_argument('--init', default=None, help='只载入权重（优先 EMA）、优化器与步数从 0 开始（DAgger 微调用）') ap.add_argument('--lr', type=float, default=None); ap.add_argument('--warmup', type=int, default=None) ap.add_argument('--stop-at', type=int, default=0, help='自检用：跑到第 N 步按停止请求处理（存 last.pt、rc 143）') ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--seed', type=int, default=0) a = ap.parse_args() cfg = json.load(open(a.config)); tc = cfg['train'] if a.steps: tc['steps'] = a.steps if a.lr: tc['lr'] = a.lr if a.warmup is not None: tc['warmup'] = a.warmup init_ck = None if a.init: init_ck = torch.load(a.init, map_location='cpu') # 2026-10-19（审查 P0-2）：原来 --init 一律置 p_neq_warm=0，前提「起点已训过 r≠t」；best.pt 落在预热期（第 4000 步）时前提不成立， # DAgger 一上来就用没训过的场做 JVP 自举。现在只有起点已走过它自己的预热才置 0。 init_warm = (init_ck.get('cfg') or {}).get('train', {}).get('p_neq_warm', 5000) if init_ck.get('step', 0) >= init_warm: tc['p_neq_warm'] = 0 print(f'微调起点第 {init_ck.get("step")} 步（其预热 {init_warm} 步）→ 本次 p_neq_warm = {tc.get("p_neq_warm", 5000)}', flush=True) os.makedirs(a.out, exist_ok=True) torch.manual_seed(a.seed) dev = a.device vram.limit() if dev.startswith('cuda'): torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True # MF 分支改 fp32 后用 TF32 保吞吐 kin = Kin(device=dev) # 2026-10-19（审查 F11）：val 只取主数据目录（第一个 --samples）；DAgger 目录的样本全部当 train data = Samples(a.samples, device=dev, split='train', val_dirs=a.samples[:1]) act_scale = data.act_scale val_tasks = [t for t in load_tasks(a.samples[:1]) if t['split'] == 'val'][:tc['eval_tasks']] try: vdata = Samples(a.samples[:1], device=dev, split='val') except SystemExit: vdata = None if vdata is not None and vdata.n == 0: vdata = None print(f'train samples {data.n}, val tasks {len(val_tasks)}, act_scale {act_scale}, data_ver {data.data_ver}, device {dev}', flush=True) model = build(cfg).to(dev); model.act_scale = act_scale print(f'params {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M', flush=True) opt = torch.optim.AdamW(model.parameters(), lr=tc['lr'], betas=tuple(tc['betas']), weight_decay=tc['wd']) step0, best, bad_evals = 0, -1.0, 0 def check_scale(ck, src): s0 = float(ck.get('act_scale', 1.0)) if abs(s0 - act_scale) > 1e-9: raise SystemExit(f'{src} 的 act_scale={s0} 与样本 act_scale={act_scale} 不一致：标签尺度变了，须从头训练（FRESH_FROM=prep）') if init_ck is not None and not a.resume: check_scale(init_ck, a.init) model.load_state_dict(init_ck.get('ema') if init_ck.get('ema') is not None else init_ck['model']) print(f'微调：权重取自 {a.init}（{"EMA" if init_ck.get("ema") is not None else "原始"}权重），优化器重置', flush=True) ema = copy.deepcopy(model).eval(); ema.act_scale = act_scale for p_ in ema.parameters(): p_.requires_grad_(False) if a.resume: ck = torch.load(a.resume, map_location=dev); check_scale(ck, a.resume) model.load_state_dict(ck['model']) ema.load_state_dict(ck['ema'] if ck.get('ema') is not None else ck['model']) if 'opt' in ck: opt.load_state_dict(ck['opt']) step0, best, bad_evals = ck.get('step', 0), ck.get('best', -1.0), ck.get('bad_evals', 0) if set_rng_state(ck.get('rng')): print('续训：随机数状态已恢复（batch 顺序与不断点一致）', flush=True) print(f'续训：从第 {step0} 步继续（共 {tc["steps"]} 步），当前最好分数 {best:.4f}', flush=True) del init_ck total = tc['steps'] lr_at = lambda s: tc['lr'] * (s / tc['warmup'] if s < tc['warmup'] else 0.5 * (1 + math.cos(math.pi * (s - tc['warmup']) / max(1, total - tc['warmup'])))) amp = dev.startswith('cuda') and tc.get('amp') == 'bf16' stepv = torch.tensor(STEP, device=dev) astep = stepv * act_scale # 动作 → 度 fk_idx = torch.linspace(0, kin.cap_a.shape[0] - 1, tc['fk_caps'], device=dev).long() ema_d = float(tc.get('ema', 0.999)) warm = 0 if model.head == 'wta' else int(tc.get('p_neq_warm', 5000)) sel_from = warm + tc['eval_every'] # 第一次可以存 best.pt 的步数：预热结束后的第一次评估 min_step = max(warm + 2 * tc['eval_every'], total // 2) # 早停计数起点 log = open(os.path.join(a.out, 'log.jsonl'), 'a' if a.resume else 'w') # 新训练覆盖旧日志，续训才追加 if not val_tasks and not a.no_eval: print('没有 val 任务（样本太少），本次跳过闭环验证', flush=True); a.no_eval = True t0, seen, diverged = time.time(), 0, False saver = Saver(tc.get('ckpt_min', 10)) # 2026-10-18：每 10 分钟 + 收到停止请求时存 last.pt（原子写） def save(name, step, extra=None): d = {'model': model.state_dict(), 'ema': ema.state_dict(), 'opt': opt.state_dict(), 'cfg': cfg, 'step': step, 'best': best, 'bad_evals': bad_evals, 'fp': kin.fp, 'act_scale': act_scale, 'data_ver': data.data_ver} if name == 'last.pt': d['rng'] = rng_state() if extra: d.update(extra) atomic_save(d, os.path.join(a.out, name)) if name == 'last.pt': saver.mark() step = step0 - 1 # 续训时已训满：循环一次不跑，下面照常收尾 for step in range(step0, total): for gpar in opt.param_groups: gpar['lr'] = lr_at(step) X, Y, V, SG = data.batch(tc['batch'], aug=tc) model.train() with torch.autocast('cuda', dtype=torch.bfloat16, enabled=amp): mem = model.encode(X) if model.head == 'wta': l_mf, xh = wta_loss(model, mem, Y); mf_raw = l_mf.detach(); xh = xh.float() else: p_neq = tc['p_neq'] if step >= warm else 0.0 l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True, p_edge=tc.get('p_edge', 0.0)) sel = ((aux['r'] == aux['t']) & (aux['t'] <= tc.get('fk_tmax', 0.3))).nonzero(as_tuple=True)[0] tt = aux['t'][sel][:, None, None] xh = aux['z'][sel] - tt * aux['u'][sel] # 配对去噪估计（fp32） v, _ = model.heads(mem) Xs = X.clone(); Xs[:, 68:75] = 0 # 子目标头：sub 特征置零（切断抄输入的捷径，见文件头） _, s = model.heads(model.encode(Xs)) l_val = (v.float() - V).abs().mean(); l_sub = (s.float() - SG).abs().mean() Xf, Yf = (X, Y) if model.head == 'wta' else (X[sel], Y[sel]) l_sm = torch.zeros((), device=dev); l_fk = torch.zeros((), device=dev) if xh.shape[0] > 0: d2 = xh[:, 2:] - 2 * xh[:, 1:-1] + xh[:, :-2] l_sm = d2.pow(2).mean() if tc['w_fk'] > 0 and xh.shape[0] > 0: q = q_from_feat(Xf) ks = torch.tensor([3, 7, 11, 15], device=dev) qh = q[:, None] + torch.cumsum(xh.clamp(-1, 1) * astep, 1)[:, ks] qt = q[:, None] + torch.cumsum(Yf * astep, 1)[:, ks] B4 = qh.shape[0] * 4 ph = fk_points(kin, qh.reshape(B4, 6), fk_idx); pt = fk_points(kin, qt.reshape(B4, 6), fk_idx) l_fk = ((ph - pt).norm(dim=-1) / 10).mean() loss = tc['w_mf'] * l_mf + tc['w_fk'] * l_fk + tc['w_smooth'] * l_sm + tc['w_value'] * l_val + tc['w_sub'] * l_sub opt.zero_grad(set_to_none=True) loss.backward() gn = torch.nn.utils.clip_grad_norm_(model.parameters(), tc['grad_clip']) if not torch.isfinite(loss): print('非有限损失 → 训练发散（M5 门禁：触发 WTA 回退条件）', flush=True); save('diverged.pt', step); diverged = True; break opt.step(); seen += X.shape[0] with torch.no_grad(): # EMA；前期按 (1+s)/(10+s) 放宽，避免随机初值占权重 dk = min(ema_d, (1 + step) / (10 + step)) for pe, pm in zip(ema.parameters(), model.parameters()): pe.lerp_(pm.detach(), 1 - dk) if saver.stop() or (a.stop_at and step + 1 >= a.stop_at): save('last.pt', step + 1); log.close() print(f'收到停止请求：已存档 last.pt（第 {step + 1} 步），下次运行自动续训', flush=True); sys.exit(143) if saver.due(): save('last.pt', step + 1) if step % tc['log_every'] == 0: el = time.time() - t0 rec = {'step': step, 'total': total, 't': time.time(), 'loss': loss.item(), 'mf': mf_raw.item(), 'fk_cm': l_fk.item(), 'smooth': l_sm.item(), 'value': l_val.item(), 'sub': l_sub.item(), 'gnorm': gn.item(), 'lr': lr_at(step), 'samples_per_s': seen / max(el, 1e-6), 'eta_h': (total - step - 1) * tc['batch'] / max(seen / max(el, 1e-6), 1e-6) / 3600, 'vram': vram.peak_gib()} print(json.dumps(rec), flush=True); log.write(json.dumps(rec) + '\n'); log.flush() if not a.no_eval and (step + 1) % tc['eval_every'] == 0: # bs=64 + 归还缓存池：WDDM 卡满时 15.9 GiB 的评估峰值会被驱动换出到内存，eval 从 2 分钟变 1 小时 res, _ = evaluate(ema, kin, val_tasks, dev, bs=64) if vdata is not None: res.update(open_loop(ema, kin, vdata, astep, fk_idx)) torch.cuda.empty_cache() sc = res['success'] + 0.05 * (1 - res['collision_rate']) + 0.05 * (1 - min(1.0, res.get('goal_ratio_2s', 1.0))) \ - 1e-4 * min(res.get('ol_best_mm', 0.0), 500.0) res['step'] = step + 1; res['t'] = time.time(); res['score'] = sc; res['weights'] = 'ema' print('EVAL', json.dumps(res), flush=True); log.write(json.dumps({'eval': res}) + '\n'); log.flush() if step + 1 >= sel_from and sc > best + 1e-5: best = sc; bad_evals = 0; save('best.pt', step + 1, {'eval': res}) elif step + 1 >= min_step: bad_evals += 1 if bad_evals >= tc['patience_evals']: save('last.pt', step + 1); print('闭环分数不再上升，停止', flush=True); break save('last.pt', step + 1) if diverged: # 发散的权重不能覆盖 last.pt / best.pt：last.pt 保留发散前最后一次存档，M5 门禁会因此判不过 → 自动 WTA print('训练发散，last.pt / best.pt 保留发散前的存档', flush=True); return save('last.pt', step + 1) if a.no_eval or best <= -1.0: save('best.pt', step + 1) # 没评估过 / 总步数不到 sel_from：best = last print(f'done, best score {best:.4f}, {(time.time() - t0) / 3600:.2f} h', flush=True) if __name__ == '__main__': main()