Robot_Arm/ml/sim/closed_loop.py

"""闭环仿真（验证 / 挖错 / 门禁共用，PyTorch 批量，无梯度）——按部署节奏运行：
50 Hz 执行；队列 ≤ 3 步时推理一次（已提交前缀 r = 队列长度），提交接下来 3 步；
冷启动用价值头对候选末态打分（短前瞻），之后按与上一计划的一致性选候选；
接近目标 < 5 mm 时用局部 DLS 精修；子目标头在新目标时给出途经点（到达 < 5° 或 1.5 s 超时后切下一个）。
safety=False 时不做任何碰撞拦截（暴露真实错误，挖错用）；safety=True 时近似部署安全层（候选提交步有干涉则换下一条，全不过则停一步）。
  python -m ml.sim.closed_loop --ckpt ml/runs/base/best.pt --samples ml/data/samples --split val --n 1000 [--safety]"""
import argparse, json, math, os
import torch

from ml.common.kin_torch import Kin
from ml.common.features import features, STEP, H, GAP_CAP, D2R, wrap180

DT = 0.02


FIRST_DEG = 5.0   # R7 首次方向判定阈值（度）：小于它的抖动不算选边

def tasks_to_tensors(tasks, device):
    q0 = torch.tensor([t['q0'] for t in tasks], dtype=torch.float32, device=device)
    T = torch.tensor([t['T'] for t in tasks], dtype=torch.float32, device=device)
    R = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in tasks], dtype=torch.float32, device=device)
    # 无老师参考（挖错/门禁题）时 exec_s/travel 为 None → NaN；调用方据此区分口径
    ex = torch.tensor([t['exec_s'] if t.get('exec_s') else float('nan') for t in tasks], dtype=torch.float32, device=device)
    tr = torch.tensor([t['travel'] if t.get('travel') else float('nan') for t in tasks], dtype=torch.float32, device=device)
    return q0, T, R, ex, tr


NOREF_SEC = 6.0      # 无老师参考题的定长时限（s）


def goal_err(kin, q, T, R):
    xf = kin.cum_xf(q)
    p = kin.fk(q, xf); Rq = kin.orient(q, xf).reshape(-1, 9)
    e = (T - p).norm(dim=-1)
    pose = R.abs().sum(-1) > 0
    tr = (Rq * R).sum(-1)
    oe = torch.where(pose, torch.acos(((tr - 1) / 2).clamp(-1, 1)) / D2R, torch.zeros_like(e))
    return e, oe


def dls_steps(kin, q, T, R, n=3, scale=0.3, lam=0.5):
    """局部 DLS 精修：n 步增量 [B,n,6]"""
    step = torch.tensor(STEP, device=q.device)
    pose = (R.abs().sum(-1) > 0).float()
    Rt = R.view(-1, 3, 3)
    def err(x):
        xf = kin.cum_xf(x); p = kin.fk(x, xf); Rq = kin.orient(x, xf)
        M = Rt @ Rq.transpose(1, 2)
        rv = 0.5 * torch.stack([M[:, 2, 1] - M[:, 1, 2], M[:, 0, 2] - M[:, 2, 0], M[:, 1, 0] - M[:, 0, 1]], -1) * 100 * pose[:, None]
        return torch.cat([T - p, rv], -1)
    out = []; x = q.clone()
    for _ in range(n):
        e0 = err(x); J = []
        for j in range(6):
            dx = torch.zeros_like(x); dx[:, j] = 1e-3
            J.append(-(err(x + dx) - e0) / 1e-3)
        J = torch.stack(J, -1)                                    # [B,6(err),6(q)]
        A = J @ J.transpose(1, 2) + lam ** 2 * torch.eye(6, device=q.device)
        y = torch.linalg.solve(A, e0.unsqueeze(-1))
        dq = (J.transpose(1, 2) @ y).squeeze(-1)
        sc = (dq.abs() / (step * scale)).max(-1, keepdim=True).values.clamp(min=1)
        dq = dq / sc
        out.append(dq); x = x + dq
    return torch.stack(out, 1)


@torch.no_grad()
def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None, return_traj=False):
    g = torch.Generator(device=device); g.manual_seed(seed)
    q0, T, R, ex, tr = tasks_to_tensors(tasks, device)
    B = q0.shape[0]
    step = torch.tensor(STEP, device=device)
    has_ref = torch.isfinite(ex)
    limit = torch.where(has_ref, torch.clamp(2 * ex, min=3.0), torch.full_like(ex, NOREF_SEC)) if max_sec is None else torch.full_like(ex, max_sec)
    nmax = int(math.ceil(limit.max().item() / DT))
    q, qp = q0.clone(), q0.clone()
    queue = torch.zeros(B, 0, 6, device=device)
    last = None
    subs = torch.zeros(B, 2, 6, device=device); nsub = torch.zeros(B, dtype=torch.long, device=device); sub_t = torch.zeros(B, device=device)
    done = torch.zeros(B, dtype=torch.bool, device=device); t_done = torch.full((B,), float('nan'), device=device)
    col = torch.zeros(B, device=device); travel = torch.zeros(B, device=device); brakes = torch.zeros(B, device=device)
    traj = [q.clone()] if return_traj else None
    n_start = int(round(0.5 / DT)); moved_start = None          # R2 起步失败：0.5 s 内关节总行程 < 1°
    j1_acc = torch.zeros(B, device=device); j1_max = torch.zeros(B, device=device); j1_min = torch.zeros(B, device=device)
    j1_first = torch.zeros(B, device=device)   # 2026-10-18 第五轮（审查 N-3）：J1 累计偏离首次超过 FIRST_DEG 时的方向（真正的"首次方向"）
    cold = True
    for it in range(nmax):
        tnow = it * DT
        if queue.shape[1] <= 3:
            r = queue.shape[1]
            pre = torch.zeros(B, 3, 6, device=device); pre[:, :r] = queue[:, :r]
            qe = q + queue.sum(1)
            e, oe = goal_err(kin, qe, T, R)
            near = (e < 5) & ((R.abs().sum(-1) == 0) | (oe < 5))
            # 子目标推进
            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)
            sf = (nsub > 0).float()
            feat = features(kin, q, qp, T, R, pre, torch.full((B,), r, device=device), subs[:, 0], sf)
            z = torch.randn(B * K, H, 6, device=device, generator=g)
            act, val, sub = model.sample_k(feat, K, z)
            inc = act.clamp(-1, 1) * step
            if cold:
                sub0 = sub
                need = sub0[:, 12] > 0.5
                subs = q[:, None] + sub0[:, :12].view(B, 2, 6) * 180
                nsub = torch.where(need, torch.full_like(nsub, 2), torch.zeros_like(nsub)); sub_t.fill_(tnow)
            fut = qe[:, None, None] + torch.cumsum(inc[:, :, r:], 2)          # [B,K,H-r,6]
            if last is None:
                # 短前瞻：候选末态的价值
                qa = fut[:, :, -1].reshape(B * K, 6); qb = fut[:, :, -2].reshape(B * K, 6)
                f2 = features(kin, qa, qb, T.repeat_interleave(K, 0), R.repeat_interleave(K, 0), torch.zeros(B * K, 3, 6, device=device),
                              torch.zeros(B * K, device=device), qa, torch.zeros(B * K, device=device))
                v2, _ = model.heads(model.encode(f2))
                score = -v2.view(B, K)
            else:
                n = min(fut.shape[2], last.shape[1] - 3)
                w = 0.85 ** torch.arange(n, device=device, dtype=torch.float32)
                d2 = (((fut[:, :, :n] - last[:, None, 3:3 + n]) / step) ** 2).sum(-1)
                score = -(d2 * w).sum(-1)
            order = score.argsort(-1, descending=True)
            pick = order[:, 0]
            commit = inc[torch.arange(B), pick, r:r + 3]
            if safety:
                ok_any = torch.zeros(B, dtype=torch.bool, device=device)
                chosen = torch.zeros(B, 3, 6, device=device)
                for kk in range(K):
                    ck = inc[torch.arange(B), order[:, kk], r:r + 3]
                    pts = qe[:, None] + torch.cumsum(ck, 1)
                    mids = torch.cat([qe[:, None], pts[:, :-1]], 1) + 0.5 * ck
                    allq = torch.cat([pts, mids], 1).reshape(-1, 6)
                    okk = (kin.gaps(allq, 0.5)['gm'].min(-1).values > 0).view(B, 6).all(-1)
                    take = okk & ~ok_any
                    chosen = torch.where(take[:, None, None], ck, chosen); ok_any |= okk
                commit = torch.where(ok_any[:, None, None], chosen, torch.zeros_like(chosen))
                brakes += (~ok_any).float()
            if near.any():
                ref = dls_steps(kin, qe, T, R)
                commit = torch.where(near[:, None, None], ref, commit)
            last = torch.cat([qe[:, None], fut[torch.arange(B), pick]], 1)
            queue = torch.cat([queue, commit], 1)
            cold = False
        d = queue[:, 0]; queue = queue[:, 1:]
        d = torch.where(done[:, None], torch.zeros_like(d), d)
        qp = q; q = kin.clamp_box(q + d)
        travel += d.abs().sum(-1)
        j1_acc = j1_acc + d[:, 0]; j1_max = torch.maximum(j1_max, j1_acc); j1_min = torch.minimum(j1_min, j1_acc)
        j1_first = torch.where((j1_first == 0) & (j1_acc.abs() >= FIRST_DEG), torch.sign(j1_acc), j1_first)
        if it + 1 == n_start: moved_start = travel.clone()
        if return_traj: traj.append(q.clone())
        mid = qp + 0.5 * (q - qp)
        gm = kin.gaps(torch.cat([q, mid], 0), 0.5)['gm'].min(-1).values
        hit = ((gm[:B] <= 0) | (gm[B:] <= 0)) & ~done
        col += hit.float()
        e, oe = goal_err(kin, q, T, R)
        newly = (~done) & (e < 1) & (oe < 0.5)
        t_done = torch.where(newly, torch.full_like(t_done, (it + 1) * DT), t_done); done |= newly
        if done.all(): break
    ok = done & (t_done <= limit)
    if moved_start is None: moved_start = travel.clone()
    e_end, _ = goal_err(kin, q, T, R)
    start_fail = (moved_start < 1.0) & ~(done & (t_done <= 0.5))
    okr = ok & has_ref & torch.isfinite(tr)
    res = {'n': B, 'n_ref': int(has_ref.sum()), 'success': ok.float().mean().item(),
           'time_ok_1p3': (ok & has_ref & (t_done <= 1.3 * ex)).float().sum().item() / max(1, int(has_ref.sum())) if has_ref.any() else None,
           'collision_rate': (col > 0).float().mean().item(), 'collision_steps_mean': col.mean().item(),
           'travel_ratio_median': (travel[okr] / tr[okr].clamp(min=1)).median().item() if okr.any() else None,
           'time_mean_s': t_done[ok].mean().item() if ok.any() else None, 'brakes_mean': brakes.mean().item(),
           'start_fail': start_fail.float().mean().item(),
           'time_budget': 'teacher×2(≥3 s)' if bool(has_ref.all()) else ('fixed %.0f s' % (max_sec or NOREF_SEC) if not has_ref.any() else 'mixed')}
    # 绕行侧：J1 累计位移的正向/负向极值哪个更大（+1 / −1），供 R7 左右比例门禁
    # 首次方向；全程没偏离到 FIRST_DEG 的（基本不绕 J1）退回净方向
    side = torch.where(j1_first != 0, j1_first, torch.where(j1_max >= -j1_min, torch.ones_like(j1_max), -torch.ones_like(j1_max)))
    per = {'done': done.cpu(), 't_done': t_done.cpu(), 'col': col.cpu(), 'travel': travel.cpu(), 'err': e_end.cpu(),
           'side': side.cpu(), 'start_fail': start_fail.cpu(), 'brakes': brakes.cpu()}
    if return_traj: per['traj'] = torch.stack(traj, 1).cpu()
    return res, per


def evaluate(model, kin, tasks, device, K=8, safety=False, bs=256, seed=0):
    if not tasks:     # 小样本摸底时 val 可能为空
        return {'n': 0, 'success': float('nan'), 'collision_rate': float('nan'), 'start_fail': float('nan'), 'brakes_mean': float('nan')}, []
    agg, pers = {}, []
    for i in range(0, len(tasks), bs):
        r, p = rollout(model, kin, tasks[i:i + bs], device, K=K, safety=safety, seed=seed + i)
        w = r['n']; pers.append((r, p))
        for k, v in r.items():
            if k == 'n' or v is None or isinstance(v, str): continue
            agg[k] = agg.get(k, 0.0) + v * w
    n = sum(r['n'] for r, _ in pers)
    out = {k: v / n for k, v in agg.items()}; out['n'] = n
    out['n_ref'] = sum(r['n_ref'] for r, _ in pers)
    out['time_budget'] = pers[0][0]['time_budget'] if len({r['time_budget'] for r, _ in pers}) == 1 else 'mixed'
    return out, pers


def main():
    from ml.model.policy import build
    from ml.data import load_tasks
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True); ap.add_argument('--samples', default='ml/data/samples')
    ap.add_argument('--tasks', default=None, help='任务 jsonl（覆盖 --samples/--split）')
    ap.add_argument('--split', default='val'); ap.add_argument('--n', type=int, default=1000)
    ap.add_argument('--safety', action='store_true'); ap.add_argument('--K', type=int, default=8)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--out', default=None)
    a = ap.parse_args()
    ck = torch.load(a.ckpt, map_location=a.device)
    model = build(ck['cfg']).to(a.device); model.load_state_dict(ck['model']); model.eval()
    kin = Kin(device=a.device)
    assert ck.get('fp') in (None, kin.fp), '碰撞模型指纹与 checkpoint 不符'
    if a.tasks: tasks = [json.loads(l) for l in open(a.tasks, encoding='utf-8') if l.strip()]
    else: tasks = [t for t in load_tasks([a.samples]) if t['split'] == a.split]
    tasks = tasks[:a.n]
    res, pers = evaluate(model, kin, tasks, a.device, K=a.K, safety=a.safety)
    by = {}
    flat_done = torch.cat([p['done'] for _, p in pers]); flat_col = torch.cat([p['col'] for _, p in pers])
    for t, d, c in zip(tasks, flat_done.tolist(), flat_col.tolist()):
        k = t.get('cat', '?') + '/' + t.get('mode', '?'); b = by.setdefault(k, [0, 0, 0]); b[0] += 1; b[1] += int(d); b[2] += int(c > 0)
    res['by_cat_mode'] = {k: {'n': v[0], 'success': v[1] / v[0], 'collision': v[2] / v[0]} for k, v in by.items()}
    s = json.dumps(res, indent=1); print(s)
    if a.out:                                    # 先写 .tmp 再改名：中断不会留下半截结果被当成已完成（M5 续跑按文件存在跳过）
        open(a.out + '.tmp', 'w').write(s); os.replace(a.out + '.tmp', a.out)


if __name__ == '__main__':
    main()