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/早停按到达率+开环误差
This commit is contained in:
Hark
@@ -0,0 +1,75 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,61 @@
|
||||
#!/usr/bin/env bash
|
||||
# DAgger 纠偏（2026-10-18 首次正式预训练后新增）：bash ml/dagger/run.sh <头标签> <起点 ckpt>
|
||||
# 每轮：策略闭环采起点（ml.dagger.collect）→ 老师从这些起点重新规划（gen_dataset.cjs --tasks，NPROC 并行）
|
||||
# → 打包成样本（build_dataset.cjs，与主数据同口径）→ 主数据 + 历轮 DAgger 样本一起微调 DAG_STEPS 步。
|
||||
# 环境变量：DAG_ROUNDS（默认 2）、DAG_N（每轮采样题数，默认 8000，每题最多 2 个起点）、DAG_STEPS（每轮微调步数，默认 15000）、
|
||||
# DAG_LR（默认 3e-4）、NPROC、MIN_GAP_MM。
|
||||
# 产物：ml/data/dag_<tag>_r<k>/（起点与老师 shard）、ml/data/samples_dag_<tag>_r<k>/、ml/runs/<tag>_dag<k>/；
|
||||
# 最终模型路径写进 ml/runs/<tag>_dagger.ckpt（m5 读它）。
|
||||
# 断点续跑：采集以 collect.json 为完成标记；老师 shard 先写 .part 再改名、未写完的从断点续写；打包以 .tasks.jsonl 为完成标记；
|
||||
# 微调用该轮 last.pt 续训；整轮完成打 ml/runs/<tag>_dag<k>/done。
|
||||
set -e -o pipefail
|
||||
cd "$(dirname "$0")/../.."
|
||||
PY=${PYTHON:-python}
|
||||
TAG=$1; CK=$2
|
||||
[ -n "$TAG" ] && [ -f "$CK" ] || { echo "usage: run.sh <tag> <ckpt>（$CK 不存在）"; exit 1; }
|
||||
R=${DAG_ROUNDS:-2}; N=${DAG_N:-8000}; STEPS=${DAG_STEPS:-15000}; LR=${DAG_LR:-3e-4}; P=${NPROC:-$(nproc)}
|
||||
CFGF=ml/configs/$TAG.json
|
||||
SAMPLES="ml/data/samples"
|
||||
for ((r=1; r<=R; r++)); do
|
||||
D=ml/data/dag_${TAG}_r$r; SD=ml/data/samples_dag_${TAG}_r$r; OUT=ml/runs/${TAG}_dag$r
|
||||
SAMPLES="$SAMPLES $SD"
|
||||
if [ -f $OUT/done ]; then echo "[dagger] 第 $r 轮已完成，跳过"; CK=$OUT/best.pt; continue; fi
|
||||
if [ -f $D/collect.json ]; then echo "[dagger] 第 $r 轮起点已采集，跳过"
|
||||
else echo "[dagger] 第 $r 轮：用 $CK 闭环采起点（$N 题）"; $PY -m ml.dagger.collect --ckpt $CK --samples ml/data/samples --n $N --seed $r --parts $P --out $D; fi
|
||||
cat $D/collect.json
|
||||
pids=()
|
||||
for f in $D/tasks_*.jsonl; do
|
||||
[ -e "$f" ] || continue
|
||||
k=$(basename $f .jsonl); k=${k#tasks_}; sh=$D/shard_dag_${TAG}_r${r}_$k.jsonl
|
||||
[ -f $sh ] && continue
|
||||
( node ml/gen/gen_dataset.cjs --tasks $f --seed $((9000 + 100 * r + 10#$k)) --out $sh.part 2>> $D/log_$k.txt && mv $sh.part $sh ) &
|
||||
pids+=($!)
|
||||
done
|
||||
if [ ${#pids[@]} -gt 0 ]; then
|
||||
echo "[dagger] 第 $r 轮：老师重新规划（${#pids[@]} 进程）"
|
||||
t0=$(date +%s)
|
||||
while :; do
|
||||
alive=0; for p in "${pids[@]}"; do kill -0 $p 2>/dev/null && alive=$((alive+1)); done
|
||||
[ $alive = 0 ] && break; sleep 30
|
||||
echo "[dagger 节拍] $(( ($(date +%s) - t0) / 60 )) 分钟 运行中 $alive/${#pids[@]} 已写 $(cat $D/*.part $D/shard_*.jsonl 2>/dev/null | wc -l) 条"
|
||||
done
|
||||
for p in "${pids[@]}"; do wait $p || true; done # 判定按产物（同 run_parallel.sh：Git Bash 长跑后 wait 返回码不可靠）
|
||||
for f in $D/tasks_*.jsonl; do
|
||||
k=$(basename $f .jsonl); k=${k#tasks_}; sh=$D/shard_dag_${TAG}_r${r}_$k.jsonl
|
||||
if [ ! -f $sh ]; then
|
||||
n=0; [ -f $sh.part ] && n=$(wc -l < $sh.part); need=$(wc -l < $f)
|
||||
if [ "$n" -ge "$need" ]; then mv $sh.part $sh; else echo "[dagger] $sh 未写完（$n/$need），见 $D/log_$k.txt；重跑本阶段从断点续写"; exit 6; fi
|
||||
fi
|
||||
done
|
||||
fi
|
||||
mkdir -p $SD
|
||||
todo=(); for sh in $D/shard_dag_*.jsonl; do [ -e "$sh" ] || continue; b=$(basename $sh .jsonl); [ -f $SD/$b.tasks.jsonl ] || todo+=("$sh"); done
|
||||
[ ${#todo[@]} = 0 ] || printf '%s\n' "${todo[@]}" | xargs -P "$P" -I{} node ml/prep/build_dataset.cjs {} $SD
|
||||
$PY -m ml.data --index $SD | tail -n 3
|
||||
res=""; [ -f $OUT/last.pt ] && res="--resume $OUT/last.pt"
|
||||
echo "[dagger] 第 $r 轮：微调 $STEPS 步（样本：$SAMPLES）"
|
||||
$PY -m ml.train --samples $SAMPLES --out $OUT --config $CFGF --init $CK --steps $STEPS --lr $LR --warmup 500 $res
|
||||
touch $OUT/done; CK=$OUT/best.pt
|
||||
done
|
||||
echo "$CK" > ml/runs/${TAG}_dagger.ckpt
|
||||
echo "[dagger] 完成：后续阶段使用 $CK"
|
||||
Reference in New Issue
Block a user