76 lines
4.7 KiB
Python
76 lines
4.7 KiB
Python
"""DAgger 状态采集（2026-10-18 首次正式预训练后新增）：用当前策略在 train 任务上闭环跑（关安全层），
|
|
从它自己走到的状态里挑起点，交给老师重新规划 → 这些"策略会去、老师数据里没有"的状态拿到正确标签。
|
|
python -m ml.dagger.collect --ckpt ml/runs/base/best.pt --samples ml/data/samples --n 8000 --seed 1 --parts 28 --out ml/data/dag1
|
|
每题最多取 2 个起点：
|
|
- 撞了的题：首次碰撞前 0.2 s 的状态（策略"快要撞上"的地方，最有用）；
|
|
- 另取一个 [0.3 s, 首撞前或结束] 内均匀随机时刻的状态（偏离老师的一般状态）。
|
|
起点必须无碰且间隙 > MIN_GAP_MM（老师从碰撞态出发规划不出来，与 build_dataset 的增广接受同一地板），且离目标 > 5 mm（否则 DLS 接管，不需要老师）。
|
|
输出 tasks_XX.jsonl（按 --parts 均分，供 gen_dataset.cjs --tasks 并行），每行 {q0, T, R, cat, src}。
|
|
为什么需要：纯行为克隆在闭环里会偏出老师分布（首次 21850 步闭环到达 0%），原方案指望 APG/挖错修，但它们没法从 0% 起步。
|
|
老师只用于训练标注，部署仍然无老师。"""
|
|
import argparse, json, os, random
|
|
import torch
|
|
|
|
from ml.common.kin_torch import Kin
|
|
from ml.data import load_tasks
|
|
from ml.model.policy import build
|
|
from ml.sim.closed_loop import rollout, goal_err, tasks_to_tensors, DT
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('--ckpt', required=True); ap.add_argument('--samples', nargs='+', default=['ml/data/samples'])
|
|
ap.add_argument('--n', type=int, default=8000); ap.add_argument('--seed', type=int, default=1)
|
|
ap.add_argument('--parts', type=int, default=8); ap.add_argument('--out', required=True)
|
|
ap.add_argument('--bs', type=int, default=64)
|
|
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
|
|
a = ap.parse_args()
|
|
gap_min = float(os.environ.get('MIN_GAP_MM') or 3)
|
|
os.makedirs(a.out, exist_ok=True)
|
|
kin = Kin(device=a.device)
|
|
ck = torch.load(a.ckpt, map_location=a.device)
|
|
assert ck.get('fp') in (None, kin.fp), '碰撞模型指纹与 checkpoint 不符'
|
|
model = build(ck['cfg']).to(a.device); model.load_state_dict(ck['model']); model.eval()
|
|
pool = [t for t in load_tasks(a.samples[:1]) if t['split'] == 'train' and t.get('cat') not in ('D',)]
|
|
rng = random.Random(a.seed); rng.shuffle(pool); tasks = pool[:a.n]
|
|
out, stat = [], {'tasks': len(tasks), 'collided': 0, 'states': 0, 'rej_gap': 0, 'rej_near': 0}
|
|
for i in range(0, len(tasks), a.bs):
|
|
tb = tasks[i:i + a.bs]
|
|
with torch.no_grad():
|
|
_, per = rollout(model, kin, tb, a.device, K=8, safety=False, seed=a.seed * 100003 + i, return_traj=True)
|
|
tr = per['traj'].to(a.device) # [B,T,6]
|
|
B, Tn, _ = tr.shape
|
|
_, T, R, _, _ = tasks_to_tensors(tb, a.device)
|
|
with torch.no_grad():
|
|
g = torch.cat([kin.gaps(tr[:, s:s + 64].reshape(-1, 6), 40.0)['gm'].min(-1).values.view(B, -1) for s in range(0, Tn, 64)], 1)
|
|
for b, t in enumerate(tb):
|
|
hit = (g[b] <= 0).nonzero()
|
|
c = int(hit[0]) if hit.numel() else Tn
|
|
if hit.numel(): stat['collided'] += 1
|
|
cand = []
|
|
if hit.numel(): cand.append(max(0, c - int(round(0.2 / DT))))
|
|
lo = int(round(0.3 / DT))
|
|
if c - 1 > lo: cand.append(rng.randint(lo, c - 1))
|
|
for k in cand:
|
|
if g[b, k] <= gap_min: stat['rej_gap'] += 1; continue
|
|
q = tr[b, k:k + 1]
|
|
e, oe = goal_err(kin, q, T[b:b + 1], R[b:b + 1])
|
|
if e.item() < 5 and (not t.get('R') or oe.item() < 5): stat['rej_near'] += 1; continue
|
|
out.append({'q0': [round(v, 3) for v in q[0].tolist()], 'T': t['T'], 'R': t.get('R'), 'cat': t.get('cat', 'A'),
|
|
'src': {'tid': t['tid'], 'k': k, 'col': bool(hit.numel())}})
|
|
print(f'[dagger] {min(i + a.bs, len(tasks))}/{len(tasks)} 题，起点 {len(out)}', flush=True)
|
|
stat['states'] = len(out)
|
|
rng.shuffle(out)
|
|
per_part = (len(out) + a.parts - 1) // max(1, a.parts)
|
|
for p in range(a.parts):
|
|
chunk = out[p * per_part:(p + 1) * per_part]
|
|
if not chunk: continue
|
|
fn = os.path.join(a.out, 'tasks_%02d.jsonl' % (p + 1))
|
|
open(fn + '.tmp', 'w', encoding='utf-8').write(''.join(json.dumps(x) + '\n' for x in chunk)); os.replace(fn + '.tmp', fn)
|
|
json.dump(stat, open(os.path.join(a.out, 'collect.json'), 'w', encoding='utf-8'), indent=1) # 最后写：它存在 = 采集完成
|
|
print(json.dumps(stat, ensure_ascii=False), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|