Files
Robot_Arm/ml/data.py
T
Hark 2a0a2e55ad ml: 数据生成 + 可微运动学对拍 + MeanFlow/WTA 策略与闭环仿真 + APG/挖错/ONNX 导出与运行手册
- M1 数据生成器:采样/老师/距离下界认证/时间律/统计/一致性门禁;确定性老师(碰撞次数计时 + resetSession)、pos 任务冻结 J6、--consensus 邻域共识过滤、Dr 类。
- 全流程代码与运行手册:数据打包/可微运动学与对拍/MeanFlow+WTA 模型/闭环仿真/APG/挖错循环/ONNX 导出;webapp 策略运行时 ik-policy.js、?policy=1 接入、tests/policy_bench.cjs;ik-solver 增 pairList/clusterExport/criticalPairs/pairGap/leverTable。
2026-10-18 06:40:23 +08:00

56 lines
2.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""样本加载:ml/data/samples/*.{X,Y,V,SG,TID}.npy + *.tasks.jsonl(ml/prep/build_dataset.cjs 产出)
python -m ml.data --index ml/data/samples # 汇总样本/任务数,写 index.json
全部载入内存(200 万样本约 1.5 GB),按任务 split 字段划分 train / val。"""
import argparse, glob, json, os
import numpy as np
import torch
def load_tasks(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
class Samples:
def __init__(self, dirs, device='cpu', split=None):
if isinstance(dirs, str): dirs = [dirs]
xs, ys, vs, ss, ts = [], [], [], [], []
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 self.tasks 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):
i = torch.randint(0, self.n, (bs,), device=self.X.device, generator=gen)
return self.X[i], self.Y[i], self.V[i], self.SG[i]
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
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}
json.dump(info, open(os.path.join(a.index, 'index.json'), 'w'), indent=1)
print(json.dumps(info, indent=1))
if __name__ == '__main__':
main()