cbb/Robot_Arm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(train): 首次正式预训练（21850 步闭环到达 0%）排查后修正——FK/平滑改加在配对去噪估计上（原加在随机噪声一步采样上=回归条件均值，多峰塌成单峰，实测候选标准差 0.030/幅度 0.151）；时间嵌入最高角频率 1000→50、MF 分支 fp32、p_neq 前 5000 步为 0（实测 du/dt 均方根 6–12 vs v 1.1）；增广 40%/3°/12 步收回；新增 dagger 阶段（策略闭环状态请老师重规划 ×2 轮微调）；开训前老师轨迹回放自检；评估加开环误差与候选多样性，best/早停按到达率+开环误差

Browse Source
This commit is contained in:
Hark
2026-10-18 15:56:39 +00:00
parent c2732fb3cd
commit 979a5814c1
19 changed files with 386 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/train.py
+59 -14
View File
@@ -2,7 +2,13 @@
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）：全程 60000 步 ≤ 3.5 h，即 samples_per_s ≥ 5000（3090 实测 6.2 k/s ≈ 2.8 h）；日志 eta_h 为剩余小时数。
每 eval_every 步做一次闭环验证（ml/sim/closed_loop.py，部署节奏），按 success 存 best.pt；连续 patience_evals 次不升即停。"""
每 eval_every 步做一次闭环验证（ml/sim/closed_loop.py，部署节奏），按 success 存 best.pt；连续 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 个候选的多样性），
塌缩或学偏当场可见；best / 早停按 success 为主、开环误差为次（success 同为 0 时也能分出好坏，不会 10000 步就误停）。"""
import argparse, json, math, os, sys, time
import torch
@@ -12,7 +18,7 @@ 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, one_step
from ml.model.meanflow import meanflow_loss
from ml.sim.closed_loop import evaluate
@@ -28,28 +34,61 @@ def q_from_feat(X):
return torch.atan2(s, c) / D2R
@torch.no_grad()
def open_loop(model, kin, vdata, stepv, fk_idx, n=1024, K=8):
"""开环指标（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='只载入权重、优化器与步数从 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
if a.init: tc['p_neq_warm'] = 0 # 微调起点已训过 r≠t，不再预热
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)
data = Samples(a.samples, device=dev, split='train')
val_tasks = [t for t in load_tasks(a.samples) if t['split'] == 'val'][:tc['eval_tasks']]
try: vdata = Samples(a.samples, 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)}, device {dev}', flush=True)
model = build(cfg).to(dev)
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
if a.init and not a.resume:
ck = torch.load(a.init, map_location=dev); model.load_state_dict(ck['model'])
print(f'微调：权重取自 {a.init}（第 {ck.get("step")} 步），优化器重置', flush=True)
if a.resume:
ck = torch.load(a.resume, map_location=dev); model.load_state_dict(ck['model'])
if 'opt' in ck: opt.load_state_dict(ck['opt'])
@@ -82,19 +121,23 @@ def main():
if model.head == 'wta':
l_mf, xh = wta_loss(model, mem, Y); mf_raw = l_mf.detach(); xh = xh.float()
else:
l_mf, mf_raw = meanflow_loss(model, mem, Y, p_neq=tc['p_neq'])
e = torch.randn_like(Y)
xh = one_step(model, mem, e).float()
p_neq = tc['p_neq'] if step >= tc.get('p_neq_warm', 5000) else 0.0
l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True)
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, s = model.heads(mem)
l_val = (v.float() - V).abs().mean(); l_sub = (s.float() - SG).abs().mean()
d2 = xh[:, 2:] - 2 * xh[:, 1:-1] + xh[:, :-2]
l_sm = d2.pow(2).mean()
l_fk = torch.zeros((), device=dev)
if tc['w_fk'] > 0:
q = q_from_feat(X)
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) * stepv, 1)[:, ks]
qt = q[:, None] + torch.cumsum(Y * stepv, 1)[:, ks]
qt = q[:, None] + torch.cumsum(Yf * stepv, 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()
@@ -119,12 +162,14 @@ def main():
model.eval()
# bs=64 + 归还缓存池：WDDM 卡满时 15.9 GiB 的评估峰值会被驱动换出到内存，eval 从 2 分钟变 1 小时
res, _ = evaluate(model, kin, val_tasks, dev, bs=64)
if vdata is not None: res.update(open_loop(model, kin, vdata, stepv, fk_idx))
torch.cuda.empty_cache()
res['step'] = step + 1; res['t'] = time.time()
print('EVAL', json.dumps(res), flush=True); log.write(json.dumps({'eval': res}) + '\n'); log.flush()
save('last.pt', step + 1)
if res['success'] > best + 1e-4:
best = res['success']; bad_evals = 0; save('best.pt', step + 1, {'eval': res})
sc = res['success'] - 1e-3 * min(res.get('ol_fk_mm', 0.0), 500.0) # 到达率为主，开环误差（mm）只在到达率打平时起作用
if sc > best + 1e-5:
best = sc; bad_evals = 0; save('best.pt', step + 1, {'eval': res})
else:
bad_evals += 1
if bad_evals >= tc['patience_evals']: print('闭环到达率不再上升，停止', flush=True); break
@@ -132,7 +177,7 @@ def main():
print('训练发散，last.pt / best.pt 保留发散前的存档', flush=True); return
save('last.pt', step + 1)
if a.no_eval: save('best.pt', step + 1)
print(f'done, best success {best:.4f}, {(time.time() - t0) / 3600:.2f} h', flush=True)
print(f'done, best score（到达率 − 0.001×开环 mm）{best:.4f}, {(time.time() - t0) / 3600:.2f} h', flush=True)
if __name__ == '__main__':