115 lines
7.3 KiB
Python
115 lines
7.3 KiB
Python
"""样本加载：ml/data/samples/*.{X,Y,V,SG,TID}.npy + *.tasks.jsonl + *.meta.json（ml/prep/build_dataset.cjs 产出）
|
||
python -m ml.data --index ml/data/samples # 汇总样本/任务数，写 index.json
|
||
全部载入内存（200 万样本约 1.5 GB），按任务 split 字段划分 train / val。
|
||
动作尺度（2026-10-19）：Y = 每步增量 / (act_scale·STEP)。act_scale 取自各 shard 的 meta.json（DATA_VER 2 起写；
|
||
旧数据没有 meta → act_scale 1.0、data_ver 1，即原口径）。同一次训练的所有目录（主数据 + DAgger）必须同一个 act_scale。"""
|
||
import argparse, glob, json, os
|
||
import numpy as np
|
||
import torch
|
||
|
||
PREFIX_SL, MASK_SL, DQ_SL, SUB_SL = slice(47, 65), slice(65, 68), slice(15, 21), slice(68, 75) # 特征分组（features.py 拼接顺序）
|
||
|
||
|
||
def load_tasks(dirs):
|
||
if isinstance(dirs, str): dirs = [dirs]
|
||
tasks = []
|
||
for d in dirs:
|
||
for f in sorted(glob.glob(os.path.join(d, '*.tasks.jsonl'))):
|
||
for line in open(f, encoding='utf-8'):
|
||
if line.strip():
|
||
tasks.append(json.loads(line))
|
||
return tasks
|
||
|
||
|
||
def read_meta(d):
|
||
"""目录内各 shard 的打包参数（data_ver、act_scale、增广参数、sig）。没有 meta 的旧打包 → {'data_ver': 1, 'act_scale': 1.0}。
|
||
一部分有、一部分没有，或 act_scale 不一致 → 拒绝（新旧标签尺度混在一起会让动作幅度错 25%）。"""
|
||
xs = sorted(glob.glob(os.path.join(d, '*.X.npy')))
|
||
ms = [json.load(open(f[:-6] + '.meta.json', encoding='utf-8')) for f in xs if os.path.exists(f[:-6] + '.meta.json')]
|
||
if not ms: return {'data_ver': 1, 'act_scale': 1.0}
|
||
if len(ms) != len(xs): raise SystemExit(f'{d}：{len(xs) - len(ms)} 个 shard 缺 meta.json（新旧打包混在一起），重跑 prep')
|
||
sc = {float(m.get('act_scale', 1.0)) for m in ms}
|
||
if len(sc) != 1: raise SystemExit(f'{d}：各 shard 的 act_scale 不一致 {sorted(sc)}，重跑 prep')
|
||
return ms[0]
|
||
|
||
|
||
class Samples:
|
||
def __init__(self, dirs, device='cpu', split=None, val_dirs=None):
|
||
"""val_dirs：只用这些目录的 val 任务划 val（默认 = 全部 dirs）。2026-10-19（审查 review_closedloop F11）：
|
||
DAgger 目录的任务起点带速度、来自策略状态，不该混进预训练的 val；train.py / apg 传主数据目录，DAgger 样本全部算 train。"""
|
||
if isinstance(dirs, str): dirs = [dirs]
|
||
xs, ys, vs, ss, ts = [], [], [], [], []
|
||
metas = [read_meta(d) for d in dirs]
|
||
sc = {float(m['act_scale']) for m in metas}
|
||
if len(sc) > 1: raise SystemExit(f'样本目录的 act_scale 不一致 {sorted(sc)}（{dirs}）：DAgger 样本须与主数据同一版打包')
|
||
self.act_scale = sc.pop(); self.data_ver = int(metas[0].get('data_ver', 1)) if metas else 1
|
||
for d in dirs:
|
||
for fx in sorted(glob.glob(os.path.join(d, '*.X.npy'))):
|
||
b = fx[:-6]
|
||
xs.append(np.load(fx)); ys.append(np.load(b + '.Y.npy')); vs.append(np.load(b + '.V.npy'))
|
||
ss.append(np.load(b + '.SG.npy')); ts.append(np.load(b + '.TID.npy'))
|
||
if not xs: raise SystemExit(f'没有样本：{dirs}（先跑 ml/prep/run_prep.sh）')
|
||
self.tasks = load_tasks(dirs)
|
||
val_tid = {t['tid'] for t in load_tasks(val_dirs if val_dirs is not None else dirs) if t['split'] == 'val'}
|
||
X, Y, V, SG, TID = map(np.concatenate, (xs, ys, vs, ss, ts))
|
||
is_val = np.isin(TID, np.fromiter(val_tid, dtype=np.int64)) if val_tid else np.zeros(len(TID), bool)
|
||
sel = {'train': ~is_val, 'val': is_val, None: np.ones(len(TID), bool)}[split]
|
||
t = lambda a: torch.from_numpy(np.ascontiguousarray(a[sel])).to(device)
|
||
self.X, self.Y, self.V, self.SG, self.TID = t(X), t(Y), t(V), t(SG), t(TID)
|
||
self.n = self.X.shape[0]
|
||
|
||
def batch(self, bs, gen=None, aug=None):
|
||
"""aug：训练配置（train 段），按 aug_pre_drop / aug_pre_sigma / aug_dq_sigma 做条件输入增广（见 augment）；None = 原样。"""
|
||
i = torch.randint(0, self.n, (bs,), device=self.X.device, generator=gen)
|
||
X, Y = self.X[i], self.Y[i]
|
||
if aug: X, Y = self.augment(X.clone(), Y.clone(), aug, gen)
|
||
return X, Y, self.V[i], self.SG[i]
|
||
|
||
def augment(self, X, Y, cfg, gen=None):
|
||
"""条件输入增广（2026-10-19，审查 review_data_actions C3 / review_model_train P1-4：前缀与 dq 的 copycat 捷径）。
|
||
训练集里 r=3 时要执行的 Y[3:6] 几乎就是前缀第 3 步的延续，dq 又永远是老师的干净速度，模型容易学成「照抄上一块」，
|
||
闭环里自己的误差经前缀 / dq 原样带回。三种扰动，标签都保持自洽：
|
||
- 前缀丢弃（aug_pre_drop）：r>0 的行以该概率清空前缀与掩码 → 变成 r=0 样本。Y 本来就从当前 q 起算，r=0 时标签不变；
|
||
- 前缀噪声（aug_pre_sigma，STEP 为单位）：r>0 的行一半加噪声 ε。前缀 = 已提交、一定会执行的步，所以标签前 r 步同样 +ε，
|
||
第 r 步减去 Σε 把位置拉回老师轨迹（亏空在下一步补齐，与增广收回同理）；
|
||
- dq 噪声（aug_dq_sigma，STEP 为单位）：一半的行给速度特征加噪声，状态与标签不变。"""
|
||
B, dev = X.shape[0], X.device
|
||
rnd = lambda *s: torch.rand(*s, device=dev, generator=gen)
|
||
rn = lambda *s: torch.randn(*s, device=dev, generator=gen)
|
||
pl = X[:, MASK_SL].sum(-1).round().long()
|
||
pd = float(cfg.get('aug_pre_drop', 0.0))
|
||
if pd > 0:
|
||
m = (rnd(B) < pd) & (pl > 0)
|
||
X[:, 47:68] = torch.where(m[:, None], torch.zeros_like(X[:, 47:68]), X[:, 47:68]); pl = torch.where(m, torch.zeros_like(pl), pl)
|
||
ps = float(cfg.get('aug_pre_sigma', 0.0))
|
||
if ps > 0:
|
||
m = (rnd(B) < 0.5) & (pl > 0)
|
||
km = (torch.arange(3, device=dev)[None] < pl[:, None]) & m[:, None] # [B,3] 前缀有效步
|
||
eps = rn(B, 3, 6) * ps * km[..., None].to(X.dtype) # STEP 为单位
|
||
X[:, PREFIX_SL] = X[:, PREFIX_SL] + eps.reshape(B, 18)
|
||
Y[:, :3] = Y[:, :3] + eps / self.act_scale # Y 以 act_scale·STEP 为单位
|
||
Y[torch.arange(B, device=dev), pl.clamp(max=Y.shape[1] - 1)] -= eps.sum(1) / self.act_scale
|
||
Y.clamp_(-1, 1)
|
||
ds = float(cfg.get('aug_dq_sigma', 0.0))
|
||
if ds > 0:
|
||
m = (rnd(B) < 0.5).to(X.dtype)
|
||
X[:, DQ_SL] = X[:, DQ_SL] + rn(B, 6) * ds * m[:, None]
|
||
return X, Y
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(); ap.add_argument('--index', default='ml/data/samples'); a = ap.parse_args()
|
||
tr, va = Samples(a.index, split='train'), Samples(a.index, split='val')
|
||
tasks = load_tasks([a.index])
|
||
cats = {}
|
||
for t in tasks: cats[t['cat'] + '/' + t['mode']] = cats.get(t['cat'] + '/' + t['mode'], 0) + 1
|
||
meta = read_meta(a.index)
|
||
info = {'train_samples': tr.n, 'val_samples': va.n, 'tasks': len(tasks), 'val_tasks': sum(t['split'] == 'val' for t in tasks), 'by_cat_mode': cats,
|
||
'data_ver': int(meta.get('data_ver', 1)), 'act_scale': float(meta.get('act_scale', 1.0)), 'sig': meta.get('sig')}
|
||
json.dump(info, open(os.path.join(a.index, 'index.json'), 'w'), indent=1)
|
||
print(json.dumps(info, indent=1))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|