Robot_Arm/ml/apg/apg_finetune.py

"""APG 可微微调（方案 §7.1）：按部署节奏可微展开策略，直接对"走得到、不碰、平滑"求梯度；不需要老师。
  python -m ml.apg.apg_finetune --init ml/runs/base/best.pt --samples ml/data/samples --out ml/runs/apg [--pool ml/runs/mine/pool.jsonl]
每条展开固定一个噪声 z（确定性可微映射）；每次推理提交 3 步（已提交前缀 3 步，与部署一致）；
截断反传：每 seg_plans 次推理（=3·seg_plans 步）反传一次并截断状态梯度。
损失：段末到达项 + Σ碰撞铰链 ReLU(margin_train − 间隙)（逐胶囊对，端点+中点）+ 平滑（速度²+加速度²）+ 行程 + BC 锚定（MeanFlow 损失）。
2026-10-18：展开长度不再定长 60 次推理（3.6 s，会截掉老师时长 > 3.6 s 的绕远题），改为覆盖本批 exec_s p90×1.2，上限 plans_max；
  "到达"仍是每 seg_plans 段末的距离项（无到达即停），属有意设计：段末项给出持续的接近梯度。
  margin_train 3 mm：与训练数据过滤线一致（认证最小间隙 < 3 mm 的老师轨迹不进训练集），碰撞项不与 BC 标签对抗；
  运行时投影护栏独立保持 5 mm。端点与中点两个铰链取均值，与只查端点同量纲，w_col 无需重标。"""
import argparse, json, math, os, random, sys, time
import torch

from ml.common.kin_torch import Kin
from ml.common.features import features, STEP, H, wrap180
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, tasks_to_tensors, goal_err
from ml.common.ckpt import Saver, atomic_save, rng_state, set_rng_state
from ml.common import vram


def load_pool(paths):
    out = []
    for p in paths or []:
        if os.path.exists(p):
            out += [json.loads(l) for l in open(p, encoding='utf-8') if l.strip()]
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--init', required=True); ap.add_argument('--samples', nargs='+', default=['ml/data/samples'])
    ap.add_argument('--pool', nargs='*', default=[]); ap.add_argument('--out', default='ml/runs/apg')
    ap.add_argument('--config', default='ml/configs/apg.json'); ap.add_argument('--steps', type=int, default=None)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--stop-at', type=int, default=0, help='自检用：跑到第 N 次按停止请求处理'); ap.add_argument('--resume', default=None, help='从 <out>/last.pt 续跑（2026-10-18 第五轮，审查 N-11：几小时的 APG 断了不再从头来）')
    a = ap.parse_args()
    cfg = json.load(open(a.config)); C = cfg['apg']
    if a.steps: C['steps'] = a.steps
    dev = a.device; torch.manual_seed(a.seed); random.seed(a.seed)
    vram.limit()
    os.makedirs(a.out, exist_ok=True)
    ck = torch.load(a.init, map_location=dev)
    mcfg = ck['cfg']
    model = build(mcfg).to(dev); model.load_state_dict(ck['model'])
    kin = Kin(device=dev)
    data = Samples(a.samples, device=dev, split='train')
    all_tasks = load_tasks(a.samples)
    train_tasks = [t for t in all_tasks if t['split'] == 'train'] + load_pool(a.pool)
    val_tasks = [t for t in all_tasks if t['split'] == 'val'][:C['eval_tasks']]
    w = [C['cat_weights'].get(t.get('cat', 'A'), 1.0) * (2.0 if t.get('hard') else 1.0) for t in train_tasks]
    opt = torch.optim.AdamW(model.parameters(), lr=C['lr'], betas=(0.9, 0.95), weight_decay=0.0)
    amp = dev.startswith('cuda') and C.get('amp') == 'bf16'
    step = torch.tensor(STEP, device=dev)
    it0 = 0
    if a.resume and os.path.exists(a.resume):
        rk = torch.load(a.resume, map_location=dev); model.load_state_dict(rk['model'])
        if 'opt' in rk: opt.load_state_dict(rk['opt'])
        it0 = rk.get('step', 0)
        if set_rng_state(rk.get('rng')): print('续跑：随机数状态已恢复', flush=True)
        bp = os.path.join(a.out, 'best.pt')
        best = torch.load(bp, map_location='cpu').get('eval', {}).get('success', -1.0) if os.path.exists(bp) else -1.0
        log = open(os.path.join(a.out, 'log.jsonl'), 'a')
        print(f'续跑：从第 {it0} 步继续，当前最好 success {best:.4f}', flush=True)
    else:
        log = open(os.path.join(a.out, 'log.jsonl'), 'w')
        base, _ = evaluate(model.eval(), kin, val_tasks, dev)
        print('EVAL0', json.dumps(base), flush=True); log.write(json.dumps({'eval': base, 'step': 0}) + '\n')
        best = base['success']
        atomic_save({'model': model.state_dict(), 'cfg': mcfg, 'step': 0, 'fp': kin.fp, 'eval': base}, os.path.join(a.out, 'best.pt'))
    M = C['margin_train_mm']; t0 = time.time()
    saver = Saver(C.get('ckpt_min', 10))             # 2026-10-18：每 10 分钟 + 收到停止请求时存 last.pt（不再只在每 500 次验证时存）
    def save_last(n, ev=None):
        atomic_save({'model': model.state_dict(), 'opt': opt.state_dict(), 'cfg': mcfg, 'step': n, 'fp': kin.fp, 'eval': ev, 'rng': rng_state()}, os.path.join(a.out, 'last.pt')); saver.mark()
    for it in range(it0, C['steps']):
        model.train()
        batch = random.choices(train_tasks, weights=w, k=C['batch'])
        q0, T, R, ex, tr = tasks_to_tensors(batch, dev)
        B = q0.shape[0]
        z = torch.randn(B, H, 6, device=dev)
        q, qp = q0.clone(), q0.clone()
        queue = torch.zeros(B, 0, 6, device=dev)
        subs = torch.zeros(B, 2, 6, device=dev); nsub = torch.zeros(B, dtype=torch.long, device=dev); sub_t = torch.zeros(B, device=dev)
        opt.zero_grad(set_to_none=True)
        stats = {'col': 0.0, 'tgt': 0.0, 'smooth': 0.0}
        # 展开长度随本批任务时长：覆盖本批 exec_s 的 p90 × 1.2（无老师参考题按 6 s），夹在 [plans, plans_max]
        exb = torch.nan_to_num(ex, nan=6.0)
        n_plans = int(min(C['plans_max'], max(C['plans'], math.ceil(1.2 * torch.quantile(exb, 0.9).item() / (3 * 0.02)))))
        nseg = math.ceil(n_plans / C['seg_plans'])
        seg_loss = torch.zeros((), device=dev); prev_d = torch.zeros(B, 6, device=dev)
        for p in range(n_plans):
            tnow = p * 3 * 0.02
            r = queue.shape[1]
            pre = torch.zeros(B, 3, 6, device=dev); pre[:, :r] = queue[:, :r]
            qe = q + queue.sum(1)
            with torch.no_grad():
                if (nsub > 0).any():
                    cur = subs[:, 0]; dd = torch.where(kin.cont, wrap180(cur - qe), cur - qe).abs().max(-1).values
                    adv = (nsub > 0) & ((dd < 5) | (tnow - sub_t > 1.5))
                    subs = torch.where(adv[:, None, None], torch.cat([subs[:, 1:], subs[:, 1:]], 1), subs)
                    nsub = torch.where(adv, nsub - 1, nsub); sub_t = torch.where(adv, torch.full_like(sub_t, tnow), sub_t)
            feat = features(kin, q, qp, T, R, pre, torch.full((B,), r, device=dev), subs[:, 0], (nsub > 0).float())
            with torch.autocast('cuda', dtype=torch.bfloat16, enabled=amp):
                act, val, sub = model.sample_k(feat, 1, z)
            act = act[:, 0].float()
            if p == 0:
                with torch.no_grad():
                    need = sub[:, 12].float() > 0.5
                    subs = q.detach()[:, None] + sub[:, :12].float().view(B, 2, 6) * 180
                    nsub = torch.where(need, torch.full_like(nsub, 2), torch.zeros_like(nsub))
            inc = act.clamp(-1, 1) * step                              # [B,16,6]
            commit = inc[:, r:r + 3]
            if r == 0:   # 冷启动：先提交一块，下一轮才开始执行
                queue = commit; continue
            # 执行队列中的 3 步（上一轮提交的）
            for k in range(3):
                d = queue[:, k]
                qp = q; q = kin.clamp_box(q + d)
                # 碰撞铰链：端点 + 中点（与仿真/运行时的离散检查一致）
                G = kin.gaps(torch.cat([q, 0.5 * (q + qp)], 0), M)
                hinge = torch.relu(M - G['gap'])
                col = torch.zeros(2 * B, device=dev).index_add(0, G['b'], hinge).view(2, B).mean(0)   # 端点与中点取均值：与只查端点时同量纲，w_col 无需重标
                sm = ((d / step) ** 2).sum(-1) + (((d - prev_d) / step) ** 2).sum(-1)
                prev_d = d
                seg_loss = seg_loss + C['w_col'] * col.mean() + C['w_smooth'] * sm.mean() + C['w_len'] * d.abs().sum(-1).mean()
                stats['col'] += col.mean().item(); stats['smooth'] += sm.mean().item()
            queue = commit
            if (p + 1) % C['seg_plans'] == 0 or p == n_plans - 1:
                xf = kin.cum_xf(q)
                pe = kin.fk(q, xf); Rq = kin.orient(q, xf).reshape(B, 9)
                pose = (R.abs().sum(-1) > 0).float()
                lt = ((T - pe) / 100).pow(2).sum(-1) + C['w_rot'] * pose * (Rq - R).pow(2).sum(-1)
                seg_loss = seg_loss + C['w_tgt'] * lt.mean()
                stats['tgt'] += lt.mean().item()
                (seg_loss / nseg).backward()
                seg_loss = torch.zeros((), device=dev)
                q, qp, queue, prev_d = q.detach(), qp.detach(), queue.detach(), prev_d.detach()
        # BC 锚定
        X, Y, V, SG = data.batch(C['batch'] * 4)
        with torch.autocast('cuda', dtype=torch.bfloat16, enabled=amp):
            mem = model.encode(X)
            l_bc = wta_loss(model, mem, Y)[0] if model.head == 'wta' else meanflow_loss(model, mem, Y)[0]
            v, s = model.heads(mem)
        l_bc = l_bc + 0.1 * (v.float() - V).abs().mean() + 0.1 * (s.float() - SG).abs().mean()
        (C['w_bc'] * l_bc).backward()
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), C['grad_clip'])
        opt.step()
        with torch.no_grad():
            e, oe = goal_err(kin, q, T, R)
        rec = {'it': it, 'total': C['steps'], 't': time.time(), 'plans': n_plans, 'err_mm_mean': e.mean().item(), 'reach_frac': ((e < 1) & (oe < 0.5)).float().mean().item(), 'bc': l_bc.item(),
               'gnorm': gn.item(), **stats, 'min': (time.time() - t0) / 60, 'vram': vram.peak_gib()}
        if it % 10 == 0: print(json.dumps(rec), flush=True)
        log.write(json.dumps(rec) + '\n'); log.flush()
        if saver.stop() or (a.stop_at and it + 1 >= a.stop_at):
            save_last(it + 1); log.close()
            print(f'收到停止请求：已存档 last.pt（第 {it + 1} 次），下次运行自动续跑', flush=True); sys.exit(143)
        if saver.due(): save_last(it + 1)
        if (it + 1) % C['eval_every'] == 0 or it == C['steps'] - 1:
            res, _ = evaluate(model.eval(), kin, val_tasks, dev)
            res['it'] = it + 1; res['t'] = time.time()
            print('EVAL', json.dumps(res), flush=True); log.write(json.dumps({'eval': res}) + '\n'); log.flush()
            sd = {'model': model.state_dict(), 'cfg': mcfg, 'step': it + 1, 'fp': kin.fp, 'eval': res}
            save_last(it + 1, res)                                                      # last 带优化器状态供续跑；best 不带（导出/门禁用）
            if res['success'] > best + 1e-4:
                best = res['success']; atomic_save(sd, os.path.join(a.out, 'best.pt'))
    print(f'done, best success {best:.4f}', flush=True)


if __name__ == '__main__':
    main()