fix(kin/diag): kin.gaps 粗筛段也按点分块(COARSE_ENTRIES=3e7,≤446 点/块;实测上一版粗筛 4096 点仍 9.3 GiB);diag_cl 加 D 段(子目标头精度 / 不用子目标 / 老师真实途经点)与目标误差进度比;rollout 加 use_sub、sub_override;本机执行清单更新

This commit is contained in:
Hark
2026-10-18 18:32:45 +00:00
parent 95ba43a1c7
commit 0653cf0d8d
4 changed files with 102 additions and 75 deletions
+23 -59
View File
@@ -3,18 +3,27 @@
本文件是 Hark 对本机计算的唯一要求清单。按顺序执行,每一步把「回传」里的文件或输出原样交给用户,不要跳步,不要自行改代码。 本文件是 Hark 对本机计算的唯一要求清单。按顺序执行,每一步把「回传」里的文件或输出原样交给用户,不要跳步,不要自行改代码。
所有命令在仓库根目录(G:\Robot_v2\机器人手臂)的 Git Bash 里运行。完成的条目由 Hark 在下一次提交里删掉。 所有命令在仓库根目录(G:\Robot_v2\机器人手臂)的 Git Bash 里运行。完成的条目由 Hark 在下一次提交里删掉。
当前状态(2026-10-19 02:00):第二次预训练开环已正常(候选多样性 0.65、8 候选最好 2.9 mm),闭环仍 0% 到达 / 100% 碰撞; ## 已知结论(2026-10-19 02:35 回传)
DAgger 阶段曾把显存吃到约 30 GB,原因是 kin.gaps 一次几千个点、cap=40 mm,已修复(kin.gaps 自动分块),尚未实测。
## 第 0 步:同步代码 - 特征口径 Python 与 JS 一致(最大差 3.6e-6),同噪声动作差 ~1e-5°:排除「闭环输入没见过」。
- 闭环关安全层 0% 到达 / 99.6% 碰撞,首撞中位 1.46 s,主碰撞组对 2-4;全程最小目标误差中位 254 mm:策略根本没朝目标走。开安全层碰撞 5.5%,到达仍 0.4%。
- 开环正常:ol_fk_mm 5.98、ol_best_mm 2.92、ol_div 0.655。老师回放 pass。
- 下一个嫌疑:子目标。训练时特征里的子目标是老师真实途经点,闭环里换成子目标头自己的预测。第 2 步的 D 段专门验证它。
- gaps 显存:上一版只给精算段分块,粗筛段 4096 点仍 9.3 GiB。本版粗筛也按点分块(每块 ≤ 446 点)。
## 第 0 步:同步代码(不要 reset --hard)
本地有未提交的 `ml/monitor.html`、`ml/pstate.py` 改动和一个未推送的合并提交,`reset --hard` 会把它们丢掉。
```bash ```bash
git fetch gitea && git reset --hard gitea/main git add ml/monitor.html ml/pstate.py && git commit -m "monitor: 开环指标两张图;pstate: train 归档加 review_*"
git log --oneline -1 # 应为「fix: 闭环失败定位与显存修复……」或更新 git pull --rebase gitea main
git log --oneline -3 # 最上面应是你的 monitor 提交,其下是 Hark 的「fix(kin/diag): …」
git push gitea main
``` ```
- 训练控制台若开着,先点「停止」,确认没有 python 训练进程在跑。 - 训练控制台若开着先点「停止」,确认没有 python 训练进程在跑。
- 不要启动 DAgger / M5,等 Hark 看完第 1、2 步的结果再定。 - 不要启动 DAgger / M5,等 Hark 看完第 1、2 步结果再定。
## 第 1 步:gaps 显存峰值(不到 1 分钟) ## 第 1 步:gaps 显存峰值(不到 1 分钟)
@@ -32,68 +41,23 @@ EOF
``` ```
- 通过:三行峰值都 ≤ 2 GiB。 - 通过:三行峰值都 ≤ 2 GiB。
- 任务管理器里专用显存接近 24 GB 或共享显存开始上涨:立刻 Ctrl+C,回传已打印的行,不要继续第 2 步。 - 专用显存接近 22 GB 或共享显存开始上涨:立刻 Ctrl+C,回传已打印的行,不要继续第 2 步。
- 回传:终端输出的 3 行。 - 回传:终端 3 行。
## 第 2 步:闭环失败定位(约 2–4 分钟) ## 第 2 步:闭环失败定位(含新加的 D 段,约 5–8 分钟)
```bash ```bash
PYTHONUTF8=1 python -m ml.eval.diag_cl --ckpt ml/runs/base/last.pt --n 256 --out ml/runs/diag_cl.json PYTHONUTF8=1 python -m ml.eval.diag_cl --ckpt ml/runs/base/last.pt --n 256 --out ml/runs/diag_cl.json
``` ```
- 运行期间看一眼任务管理器,专用显存应远低于 24 GB;超过 22 GB 就 Ctrl+C 并回传终端输出。 - 专用显存超过 22 GB 就 Ctrl+C 并回传终端输出。
- 回传:`ml/runs/diag_cl.json`、`ml/runs/replay_teacher.json`、`ml/runs/base/log.jsonl`。 - 新增输出:`C.goal_err_ratio_p50`(0.5/1/2 s 时目标误差 ÷ 起点误差,<1 表示在靠近);`D0_subhead`(冷启动子目标头精度);`D1_no_sub`(不用子目标);`D2_oracle_sub`(用老师真实途经点)。
- 回传:`ml/runs/diag_cl.json`。
## 回传汇总 ## 回传汇总
| 步骤 | 回传内容 | | 步骤 | 回传内容 |
|---|---| |---|---|
| 1 | 终端 3 行峰值 | | 1 | 终端 3 行峰值 |
| 2 | ml/runs/diag_cl.json、ml/runs/replay_teacher.json、ml/runs/base/log.jsonl | | 2 | ml/runs/diag_cl.json |
| 任一步出错 | 终端最后 30 行输出 | | 任一步出错 | 终端最后 30 行输出 |
---
## 执行回传(2026-10-19 02:35,本地 AI)
### 第 0 步:同步代码
分叉:本地 `5a87c06`(含 `ml/eval/diag_collision.py` 等本地提交,晚 02:00 前)与 `gitea/main` `573082d`(03:24 强推)各有 1 个提交;本地未提交改动 `ml/monitor.html`(开环指标两张图)、`ml/pstate.py`(train 归档加 `review_*`)。
已 `git merge 573082d` 合流(工作区内容与 `573082d` + 未提交改动一致),未执行清单里的 `git reset --hard`,无训练进程在跑。
注:本地现比 `gitea/main` 领先 2 个提交(合并提交 + 上述未提交改动尚未提交),未 push。
### 第 1 步:gaps 显存峰值(实测)
```
n= 512 cap= 40 peak 1.20 GiB pairs=10724086
n= 4096 cap= 4 peak 9.30 GiB pairs=14880476
n= 4096 cap= 40 peak 9.30 GiB pairs=87180763
```
**未通过(两条 4096 都是 9.30 GiB,>2 GiB)**。三行是独立新进程各跑一次的结果(排除分配器缓存假象),9.30 GiB 是可复现的确定值。
定位(在 `gaps()` 内插探针,不改行为):
- 峰值全部出在**粗筛段**(`cc/d2/per`,一次材料到 491 组 × 67256 粗对):4096 点第一次 `d2 = (cc[:,cp_a]-cc[:,cp_b]).pow(2)` 的广播中间量就 ~9.3 GiB。`[coarse]` 探针打到 9.30 GiB 后,后续 `[bounds]/[pre-loop]/[post-loop]/[cat-done]` 都不再涨。
- 新加的分块(PAIR_BUDGET=2e6)**只对精算段生效**:本模型几何比注释里假想的大得多——1936 条胶囊 / 491 组 / 67256 粗对 / 104.7 万细对,4096 点粗筛幸存 **2.41 亿对**(cap=4 就已超预算 120 倍),分块确实在切(cap=4 切 121 块、cap=40 切 283 块,每块 10–45 点),每块精算峰值 ≤1.0 GiB。合计峰值还是被粗筛一条打回 9.3 GiB。
- 512 点 cap=40 的 1.20 GiB 同样出自粗筛(512×67256 → 广播中间量 ~1.2 GiB),精算段只占其中一小截。
- 旧版对比(`0b47905`,无任何分块):同一三条独立实测 **4.59 / 16.98 / 36.89 GiB**——旧版 4096×40 的 36.9 GiB 在 24 GB 卡上必 OOM,与"dagger 吃满 24 GB、估计 30 GB"记录吻合。本版把 4096×40 从 36.9 降到 9.30 GiB(约 4.0×),但 4096 两条仍在 9.30,>2 GiB 门槛。
建议修复(未改代码,等你定):粗筛也按 PAIR_BUDGET 同类逻辑分块(复用同一套 bounds),或把粗筛改成按粗对列分块累加 `per`;粗筛分摊到 ~200 万对/块后,预计三条都能回到 ≤2 GiB。
### 第 2 步:闭环失败定位(`ml/runs/diag_cl.json`,完整内容见该文件)
关键数字:
- **A(特征口径):全组通过。** 最大差全部在 fp32 数值噪声量级(最大 3.6e-6,gaps 组;其余 ≤4.8e-7)。Python 现算特征与数据集 X 口径一致,**排除"闭环看到没见过的输入"这条路**。
- **B(噪声下动作差):无害。** 同一噪声两份特征的动作差 mean 5.1e-6°、p90 8.8e-6°,参照动作幅度 14.7°。
- **C(失败形态,n=256):**
- 关安全层:到达 **0%**,碰撞 **99.6%**;首次碰撞 p10/p50/p90 = 0.4/1.46/3.52 s,仅 0.4% 在 0.1 s 内(不是"起步就撞")。
- 碰撞组对:**2-4 共 90 次**(J3-J5),其次 0-4(J1-J5)42、0-3(J1-J4)33、0-5 21、0-2 19、1-6 18。
- 全程最小目标误差 p10/p50/p90 = 93/254/432 mm——**连 5° 都没保持,不是"接近到不了",是根本走不到**。
- 开安全层:到达 0.4%,碰撞 5.5%,刹车中位数 5 次——安全层把碰撞从 99.6% 压到 5.5%,但任务仍然到不了。
### 其余两个回传文件
- `ml/runs/replay_teacher.json`(00:41 开训前自检的产物,非本次新跑):n=512,到达 1.0、碰撞 0.0、min_gap p10 0.5 mm、pass=true。
- `ml/runs/base/log.jsonl`:末行 eval(step 14000)success 0.0、collision_rate 1.0、collision_steps_mean 268、ol_fk_mm 5.98、ol_best_mm 2.92、ol_div 0.655。与 C 段一致。
+36 -15
View File
@@ -118,11 +118,22 @@ class Kin:
return a, b return a, b
# ---------------- 间隙 ---------------- # ---------------- 间隙 ----------------
# 2026-10-19(用户实测 dagger/诊断吃满 24 GB 显存):粗筛后幸存胶囊对数随 点数 × cap 暴涨, # 2026-10-19(用户实测 dagger/诊断吃满 24 GB 显存):本模型 1936 条胶囊 / 491 组 / 67256 粗对 / 104.7 万细对。
# 一次几千个点、cap=40 mm 时中间张量可达十几 GB(训练里 eval 峰值 15.9 GiB 也是它)。 # 显存有两段峰值,都随点数线性涨:
# 现在先按簇级粗筛估出每个点的幸存胶囊对数,按 PAIR_BUDGET 把 batch 切块逐块算,峰值与调用方的 batch 大小无关。 # 粗筛段 cc[:,cp_a]-cc[:,cp_b] 的广播中间量 ≈ 36 B × 点数 × 粗对数(4096 点实测 9.3 GiB,与 cap 无关);
# 精算段 ≈ 每对几百字节 × 幸存细对数(4096 点 cap 40 幸存 2.4 亿对;旧版不分块实测 36.9 GiB)。
# 所以先按 COARSE_ENTRIES 把点切块(粗筛每块 ≤ ~1 GiB),块内再按幸存细对数 PAIR_BUDGET 切(精算每块 ≤ ~1 GiB)。
# 两段不同时存在,峰值与调用方 batch 大小无关。结果与不分块逐位一致(梯度经 cat 保留)。
COARSE_ENTRIES = 30_000_000
PAIR_BUDGET = 2_000_000 PAIR_BUDGET = 2_000_000
@staticmethod
def _cat_gaps(parts, starts):
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, starts)]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
def gaps(self, q, cap=40.0, xf=None, max_pairs=None): def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
"""返回 dict: """返回 dict:
gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap) gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap)
@@ -131,11 +142,25 @@ class Kin:
xf = xf if xf is not None else self.cum_xf(q) xf = xf if xf is not None else self.cum_xf(q)
Rs, Ts = xf Rs, Ts = xf
B = q.shape[0] B = q.shape[0]
step = max(1, self.COARSE_ENTRIES // int(self.cp_a.shape[0]))
if B <= step:
return self._gaps_split(q, cap, xf)
st = list(range(0, B, step))
return self._cat_gaps([self._gaps_split(q[s:s + step], cap, (Rs[s:s + step], Ts[s:s + step])) for s in st], st)
def _coarse(self, q, cap, xf):
Rs, Ts = xf
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
del cc
return d2 <= (self.cp_lim + cap).pow(2)
def _gaps_split(self, q, cap, xf):
"""粗筛已在显存预算内:按幸存细对数再切块做精算"""
Rs, Ts = xf
B = q.shape[0]
with torch.no_grad(): with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g] per = (self._coarse(q, cap, xf).long() * self.cp_cnt).sum(1) # [B] 每点幸存细对数
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1)
per = ((d2 <= (self.cp_lim + cap).pow(2)).long() * self.cp_cnt).sum(1) # [B] 每点幸存胶囊对数
del cc, d2
if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET: if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET:
bounds = [0, B] bounds = [0, B]
else: else:
@@ -144,22 +169,18 @@ class Kin:
if c - base > self.PAIR_BUDGET and k > bounds[-1]: if c - base > self.PAIR_BUDGET and k > bounds[-1]:
bounds.append(k); base = cs[k - 1] bounds.append(k); base = cs[k - 1]
bounds.append(B) bounds.append(B)
del per
if len(bounds) == 2: if len(bounds) == 2:
return self._gaps_one(q, cap, xf) return self._gaps_one(q, cap, xf)
parts = [self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])] return self._cat_gaps([self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])], bounds[:-1])
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, bounds[:-1])]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
def _gaps_one(self, q, cap, xf): def _gaps_one(self, q, cap, xf):
Rs, Ts = xf Rs, Ts = xf
B = q.shape[0] B = q.shape[0]
with torch.no_grad(): with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g] keep = self._coarse(q, cap, xf)
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
keep = d2 <= (self.cp_lim + cap).pow(2)
bb, mm = keep.nonzero(as_tuple=True) bb, mm = keep.nonzero(as_tuple=True)
del keep
cnt = self.cp_cnt[mm] cnt = self.cp_cnt[mm]
rep_b = torch.repeat_interleave(bb, cnt) rep_b = torch.repeat_interleave(bb, cnt)
off = torch.repeat_interleave(self.cp_start[mm] - torch.cumsum(cnt, 0) + cnt, cnt) + torch.arange(int(cnt.sum()), device=q.device) off = torch.repeat_interleave(self.cp_start[mm] - torch.cumsum(cnt, 0) + cnt, cnt) + torch.arange(int(cnt.sum()), device=q.device)
+39
View File
@@ -27,6 +27,9 @@ GROUPS = [('sincos', 0, 12), ('joint_norm', 12, 15), ('dq', 15, 21), ('p', 21, 2
('sub', 68, 74), ('sub_flag', 74, 75), ('gaps', 75, 83)] ('sub', 68, 74), ('sub_flag', 74, 75), ('gaps', 75, 83)]
PROG_T = (0.5, 1.0, 2.0) # 目标误差 / 起点误差 的取样时刻(s)
def pct(v, p): def pct(v, p):
v = sorted(v) v = sorted(v)
return v[min(len(v) - 1, int(len(v) * p))] if v else None return v[min(len(v) - 1, int(len(v) * p))] if v else None
@@ -81,6 +84,7 @@ def main():
# ---------- C:闭环失败形态 ---------- # ---------- C:闭环失败形态 ----------
first_t, col_steps, pairs, min_err, reach = [], [], {}, [], 0 first_t, col_steps, pairs, min_err, reach = [], [], {}, [], 0
prog = {ts: [] for ts in PROG_T}
reach_s, brakes_s, col_s = 0, [], 0 reach_s, brakes_s, col_s = 0, [], 0
for i in range(0, len(tasks), a.bs): for i in range(0, len(tasks), a.bs):
ch = tasks[i:i + a.bs] ch = tasks[i:i + a.bs]
@@ -95,6 +99,9 @@ def main():
Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in ch], dtype=torch.float32, device=dev) Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in ch], dtype=torch.float32, device=dev)
e, _ = goal_err(kin, tr.reshape(-1, 6), Tt.repeat_interleave(Tn, 0), Rt.repeat_interleave(Tn, 0)) e, _ = goal_err(kin, tr.reshape(-1, 6), Tt.repeat_interleave(Tn, 0), Rt.repeat_interleave(Tn, 0))
e = e.view(B, Tn) e = e.view(B, Tn)
for b in range(B):
for ts in PROG_T:
k = min(Tn - 1, int(round(ts / DT))); prog[ts].append((e[b, k] / e[b, 0].clamp(min=1e-3)).item())
# 到达后轨迹被冻结(d=0),之后不再计 # 到达后轨迹被冻结(d=0),之后不再计
for b in range(B): for b in range(B):
if per['done'][b]: reach += 1 if per['done'][b]: reach += 1
@@ -113,6 +120,38 @@ def main():
'first_hit_pair(group a-b)': dict(sorted(pairs.items(), key=lambda kv: -kv[1])[:8]), 'first_hit_pair(group a-b)': dict(sorted(pairs.items(), key=lambda kv: -kv[1])[:8]),
'min_goal_err_mm': {'p10': pct(min_err, .1), 'p50': pct(min_err, .5), 'p90': pct(min_err, .9)}, 'min_goal_err_mm': {'p10': pct(min_err, .1), 'p50': pct(min_err, .5), 'p90': pct(min_err, .9)},
'safety_on': {'reach': reach_s / n, 'collided': col_s / n, 'brakes_p50': pct(brakes_s, .5)}} 'safety_on': {'reach': reach_s / n, 'collided': col_s / n, 'brakes_p50': pct(brakes_s, .5)}}
out['C']['goal_err_ratio_p50'] = {'%.1fs' % ts: pct(v, .5) for ts, v in prog.items()} # <1 = 在靠近目标
# ---------- D:子目标是不是元凶(2026-10-19 加:C 显示全程最小目标误差中位 254 mm,策略根本没朝目标走) ----------
# 训练时特征里的子目标是老师真实途经点(80% 概率给),闭环里却换成子目标头自己的预测;预测偏了,策略就会忠实地朝错的方向走。
sg_by_tid = {}
for i, t in zip(idx, tk):
qrow = torch.atan2(vd.X[i, 0:12:2].float(), vd.X[i, 1:12:2].float()) * 180 / torch.pi
sgr = vd.SG[i].float()
sg_by_tid[t['tid']] = (qrow[None] + sgr[:12].view(2, 6) * 180, bool(sgr[12] > 0.5))
# D0:冷启动子目标头精度(起点特征里抹掉数据集的真子目标,与闭环冷启动同口径)
Xs = X0.clone(); Xs[:, 68:75] = 0
with torch.no_grad(): _, sp = model.heads(model.encode(Xs))
lab = torch.stack([vd.SG[i].float() for i in idx])
need_l = lab[:, 12] > 0.5; need_p = sp[:, 12] > 0.5
via_err = ((sp[:, :6] - lab[:, :6]) * 180).abs().max(-1).values
out['D0_subhead'] = {'n': len(idx), 'need_label_rate': need_l.float().mean().item(), 'need_acc': (need_l == need_p).float().mean().item(),
'via1_maxjoint_err_deg_p50': pct(via_err[need_l].tolist(), .5), 'via1_maxjoint_err_deg_p90': pct(via_err[need_l].tolist(), .9)}
# D1 不用子目标、D2 用老师真实途经点(oracle),同一批题、同一随机种子,对比 C 的基线
def run(tag, **kw):
rc = cc_ = 0; me = []
for i in range(0, len(tasks), a.bs):
ch = tasks[i:i + a.bs]; extra = dict(kw)
if kw.get('sub_override') == 'oracle':
subs = torch.stack([sg_by_tid[t['tid']][0] if t['tid'] in sg_by_tid else torch.zeros(2, 6, device=dev) for t in ch])
need = torch.tensor([sg_by_tid[t['tid']][1] if t['tid'] in sg_by_tid else False for t in ch], device=dev)
extra['sub_override'] = (subs, need)
with torch.no_grad():
_, pp = rollout(model, kin, ch, dev, K=8, safety=False, seed=i, **extra)
rc += int(pp['done'].sum()); cc_ += int((pp['col'] > 0).sum()); me += pp['err'].tolist()
out[tag] = {'reach': rc / n, 'collided': cc_ / n, 'end_goal_err_mm_p50': pct(me, .5)}
run('D1_no_sub', use_sub=False)
run('D2_oracle_sub', sub_override='oracle')
s = json.dumps(out, indent=1, ensure_ascii=False); print(s) s = json.dumps(out, indent=1, ensure_ascii=False); print(s)
os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True) os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
open(a.out, 'w', encoding='utf-8').write(s) open(a.out, 'w', encoding='utf-8').write(s)
+4 -1
View File
@@ -65,7 +65,8 @@ def dls_steps(kin, q, T, R, n=3, scale=0.3, lam=0.5):
@torch.no_grad() @torch.no_grad()
def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None, return_traj=False): def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None, return_traj=False, use_sub=True, sub_override=None):
"""use_sub=False:不用子目标(sf=0);sub_override=(subs [B,2,6] 绝对关节角, need [B] bool):冷启动时用给定途经点代替子目标头(诊断用)"""
g = torch.Generator(device=device); g.manual_seed(seed) g = torch.Generator(device=device); g.manual_seed(seed)
q0, T, R, ex, tr = tasks_to_tensors(tasks, device) q0, T, R, ex, tr = tasks_to_tensors(tasks, device)
B = q0.shape[0] B = q0.shape[0]
@@ -108,6 +109,8 @@ def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None,
sub0 = sub sub0 = sub
need = sub0[:, 12] > 0.5 need = sub0[:, 12] > 0.5
subs = q[:, None] + sub0[:, :12].view(B, 2, 6) * 180 subs = q[:, None] + sub0[:, :12].view(B, 2, 6) * 180
if sub_override is not None: subs, need = sub_override[0].to(device), sub_override[1].to(device)
if not use_sub: need = torch.zeros_like(need)
nsub = torch.where(need, torch.full_like(nsub, 2), torch.zeros_like(nsub)); sub_t.fill_(tnow) nsub = torch.where(need, torch.full_like(nsub, 2), torch.zeros_like(nsub)); sub_t.fill_(tnow)
fut = qe[:, None, None] + torch.cumsum(inc[:, :, r:], 2) # [B,K,H-r,6] fut = qe[:, None, None] + torch.cumsum(inc[:, :, r:], 2) # [B,K,H-r,6]
if last is None: if last is None: