86 lines
5.2 KiB
Python
86 lines
5.2 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
|
||
|
||
GAP_CHUNK = 1024 # 每次 kin.gaps 的查询点数上限
|
||
|
||
|
||
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)
|
||
# 2026-10-19 修（用户实测）：原来一次喂 64 题 × 64 步 = 4096 点、cap=40 mm 给 kin.gaps，40 mm 内幸存的胶囊对太多，
|
||
# 峰值把显存 / 内存吃满。这里只需判断"≤0（碰撞）"和"> gap_min（起点地板）"，cap 取 gap_min+1 即可，
|
||
# 且按 GAP_CHUNK 个点分块，用完即释放。
|
||
cap = gap_min + 1.0
|
||
flat = tr.reshape(-1, 6)
|
||
with torch.no_grad():
|
||
g = torch.cat([kin.gaps(flat[s:s + GAP_CHUNK], cap)['gm'].min(-1).values for s in range(0, flat.shape[0], GAP_CHUNK)]).view(B, Tn)
|
||
del flat
|
||
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())}})
|
||
del tr, g, per
|
||
if a.device.startswith('cuda'): torch.cuda.empty_cache()
|
||
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()
|