Add files using upload-large-folder tool
50762c1 verified 16 days ago
24.4 kB
"""蒸馏用的 batch 工具：mask / 权重图 / 偏移与文本的延长和开窗。
数据集本身**一行不改**：`actionrope.dataset.AropeLatentDataset` + `arope_collate` 已经给齐
了三个阶段需要的一切（`input_latents [b,48,21,30,52] bf16`、`context [b,21,512,4096] bf16`、
`context_ids [b,21] i64`、`offset_px [b,21,2] f32`、`offset_tok`、`delta_tok`、`action_idx`、
`clip_id`、`kind`）。本文件只做 batch 之上的纯函数，不碰 IO、不碰随机数据增强。
四类工具
--------
1. `to_device` —— 与 `actionrope/train.py:344` 同一份语义（张量搬走，字符串 list 原样）。
2. `make_mask_input` / `make_weight_map` —— stage1/2 的 `mask_input` 与 loss 权重图。
   `mask_input` 直接用 `known_mask(offset_px / 16, 30, 52)`，**不能**写成
   `loss_weight_map(...).eq(1)`：那样一来模型的输入通道就被 loss 配方绑架了——
   `new_weight=1.0`（想关掉 new 区加权时的自然写法）会让权重图恒等于 1，掩码整幅变成
   "全 known"，模型输入被悄悄改掉而不报错；`sigma`/`ramp_latent` 也各自能把这条等价关系
   推到 float32 的舍入边缘（dist 模式下 `d == 0 ⟺ known` 只是巧合，紧贴边界的
   `1 + 5e-10` 在 float32 里就是 1.0）。掩码是**模型输入**，权重图是**loss 配方**，
   两件事必须各自从 offset 现算。
3. `window_offsets` / `window_context` —— stage3 滑窗 KV 的**锚定切片**。绝对 cell k、
   窗口最老保留 cell a 时，窗口内的世界坐标一律减去 `offset_px[a]`，于是整段窗口在 RoPE
   看来就是一条以自身 cell 0 为原点的 21-cell 训练 clip（RoPE 是相对编码，整体平移不变）。
   两者索引口径完全一致：`[:, a:k+1]`。
4. `extend_offsets` / `extend_context` —— rollout 超过 21 cell 时的延长。
   `--rollout_max 26` 意味着 cell 21..26 在数据里根本不存在：`parse_actions` 硬要求 21 个
   cell（`infer.py:127`），`AropeLatentDataset.__getitem__` 硬要求 21 条 prompt。两条流必须
   **同时**延长而且逐下标对齐，否则要么在 `context.shape[1] == 1+N` 的断言处崩，要么更糟：
   几何往前走、文本停在 cell 20，rollout 照跑、loss 照降，但训的是错的东西。
   两条流的延长策略**故意不同**，这不是笔误：
   * 文本可以"保持"最后一条（`_wm/WorldMind/distill/code/stage3_dmd/dmd_gm.py:628-637`
     的 `last = context[:, Fl-1:Fl]; pad = last.expand(...)`）。ARoPE 臂的 prompt 经
     `prompts.strip_action` 剥掉了动作从句，只剩场景描述，而 rollout 不换场景，所以
     "保持最后一条" 在这里是**语义正确**的，不是凑合。
   * 偏移不能"保持"。偏移保持不动 = "玩家站住了"，那是**另一个动作**，不是中性动作
     （V Rising 的 null action 陷阱，GraphGWM/README 记过一次）。所以偏移必须继续积分一个
     速度：`continue`（沿用最后一段速度）或 `resample`（默认，每 5 个 cell 换一次方向，
     方向取 `geometry.ACTION_DIRS` 9 类、速度取 sidecar 实测的 `walk_speed_px_s`）。
速度 → 逐 cell 偏移的换算（不是拍脑袋的 v/4）
------------------------------------------
Wan2.2 因果 VAE 的分组是 cell 0 = 帧 0、cell k = 帧 4k−3…4k（`geometry.cell_frames`），
逐 cell 偏移是 cell 内各帧偏移的**均值**（`geometry.frames_to_cells`），而逐帧偏移按
`frame_off[i] = frame_off[i−1] + v[cell(i)] / FPS` 积分（`infer.py::parse_actions`）。
把这两步合起来，逐 cell 增量是（v 单位 px/s，FPS=16）：
    delta[1]   = 2.5 · v[1] / FPS                      （cell 0 只有帧 0，起步项不同）
    delta[k]   = (1.5 · v[k−1] + 2.5 · v[k]) / FPS     （k ≥ 2）
匀速时退化成 4·v/FPS = v/4 px per cell（一个 cell = 4 帧 = 0.25 s）。换方向的那个 cell 上
1.5/2.5 的加权才是真值，直接写 v/4 会在每个 chunk 边界差 1.5·|Δv|/16（最大 ~13 px ≈ 0.8 个
latent 格）。`test_data.py::test_extend_matches_frame_integration` 用帧级积分逐位核对了这条。
"""
from __future__ import annotations

import math

import numpy as np
import torch

from actionrope.arope import known_mask, loss_weight_map
from actionrope.geometry import ACTION_DIRS, CELLS, FPS, LAT_H, LAT_W, PX_PER_LATENT

# 一个 cell 覆盖 4 帧（cell 0 除外），见 geometry.cell_frames
CELL_FRAMES = 4
# 帧级积分 + cell 内均值合出来的两个权重（推导见模块 docstring）
_W_PREV, _W_CUR = 1.5, 2.5

# sidecar 实测步速（data/latent/arope_sidecar.pt 的 meta["walk_speed_px_s"]，已乘 |gain|）。
# 横纵不同是透视造成的，不是噪声；`infer.py::parse_actions` 也是按轴分别乘这个数。
# 只作为默认值：真跑的时候优先用 walk_speed_from_sidecar() 从当前 sidecar 现读。
WALK_SPEED_PX_S = (69.24106392992304, 46.012630371093756)

# rollout 延长的默认档位，与 stage3 的 --offset_chunk_cells 5 对齐
DEFAULT_CHUNK_CELLS = 5
DEFAULT_SPEED_MULT = (0.8, 1.2)


# --------------------------------------------------------------------------
# batch 搬运
# --------------------------------------------------------------------------

def to_device(batch: dict, device, non_blocking: bool = True) -> dict:
    """张量搬到 device，其余（clip_id / kind 这种 list[str]）原样带过去。
    与 `actionrope/train.py:344-345` 同一份实现，蒸馏侧不再各写一遍。
    """
    return {
        k: (v.to(device, non_blocking=non_blocking) if isinstance(v, torch.Tensor) else v)
        for k, v in batch.items()
    }


# --------------------------------------------------------------------------
# mask 通道 / loss 权重图
# --------------------------------------------------------------------------

def make_mask_input(offset_px: torch.Tensor, lat_h: int = LAT_H, lat_w: int = LAT_W,
                    angle=None) -> torch.Tensor:
    """known/new 掩码通道 → float32 [b, 1, f, lat_h, lat_w]，喂给 `dit.arope_mask_embedding`。
    offset_px [b, f, 2] px。除以 PX_PER_LATENT=16 换成 latent 格单位，与
    `actionrope/train.py:243-245` 逐位一致（known=1、new=0）。
    stage3 传的是**锚定后**的窗口偏移（见 window_offsets），于是 "known" 的含义变成
    "落在窗口最老那一 cell 的足迹内"——正是模型训练时见过的那种。
    """
    off = _as_bf2(offset_px, "offset_px")
    if angle is not None and bool(torch.any(torch.as_tensor(angle) != 0)):
        known = _known_rotated(off / PX_PER_LATENT, angle, lat_h, lat_w)
    else:
        known = known_mask(off / PX_PER_LATENT, lat_h, lat_w)  # bool [b, f, h, w]
    return known.to(torch.float32).unsqueeze(1)


def _known_rotated(off: torch.Tensor, angle, lat_h: int, lat_w: int) -> torch.Tensor:
    """相机在锚点与当前 cell 之间转了 angle 弧度时的 known 掩码（angle = 锚点朝向 − 当前朝向）。
    当前格点 p 先按原口径加上偏移，再绕画幅中心转 −angle 回到**锚点相机系**，落在锚点足迹
    [−0.5, h−0.5] × [−0.5, w−0.5] 内 ⇒ known。angle=0 时逐位退化成 `known_mask`（旋转成恒等）。
    旋转之后 known 区是一块**斜的**多边形 —— 训练里只出现过轴对齐的边界，这一条是已知的分布外。
    """
    from distill.arope_dit import rot2d
    b, f = off.shape[:2]
    dev, dt = off.device, torch.float64
    ang = torch.as_tensor(angle, device=dev, dtype=dt).view(b, f, 1, 1)
    gi = torch.arange(lat_h, device=dev, dtype=dt).view(1, 1, lat_h, 1).expand(b, f, lat_h, lat_w)
    gj = torch.arange(lat_w, device=dev, dtype=dt).view(1, 1, 1, lat_w).expand(b, f, lat_h, lat_w)
    ctr_h, ctr_w = (lat_h - 1) / 2.0, (lat_w - 1) / 2.0
    o = off.to(dt).view(b, f, 1, 1, 2)
    p = torch.stack([gj + o[..., 0] - ctr_w, gi + o[..., 1] - ctr_h], dim=-1)
    q = rot2d(p, -ang)
    wj, wi = q[..., 0] + ctr_w, q[..., 1] + ctr_h
    return (wi >= -0.5) & (wi <= lat_h - 0.5) & (wj >= -0.5) & (wj <= lat_w - 0.5)


def make_weight_map(offset_px: torch.Tensor, lat_h: int = LAT_H, lat_w: int = LAT_W,
                    new_weight: float = 2.0, ramp_latent: float = 0.0, sigma=None) -> torch.Tensor:
    """loss 权重图 → float32 [b, 1, f, lat_h, lat_w]（known=1，new=new_weight）。
    直接转手 `actionrope.arope.loss_weight_map`，把 10k 检查点训练时用的 new 区 ×2 权重
    原样带进 stage1。stage2 的 CD loss **故意不用**这张图（见 PORT_SPEC）。
    """
    off = _as_bf2(offset_px, "offset_px")
    return loss_weight_map(off, lat_h, lat_w, new_weight=new_weight,
                           ramp_latent=ramp_latent, sigma=sigma)


# --------------------------------------------------------------------------
# 滑窗锚定
# --------------------------------------------------------------------------

def window_offsets(offset_px: torch.Tensor, a: int, k: int) -> torch.Tensor:
    """窗口 [a, k] 的锚定偏移 → [b, k−a+1, 2]，第 0 项恒为 0。
    a = max(0, k − kv_window) 是窗口里最老的保留 cell。RoPE 是相对编码，整窗同时平移
    (t, offset) 不改变任何注意力分数，所以 RoPE 重锚定是精确的（`tests/test_world_rope_window.py` a1/a3），
    当前查询 cell 的 known 掩码也精确（d1）。但**缓存里各 cell 的 known 掩码通道是按它们各自提交时的锚点
    a_j = max(0, j−20) < a 固化的**，所以 k>20 的历史与一条锚在 k−20 的训练 clip 并不逐位相同 ——
    这是 stage2→stage3 的输入分布移位，靠 self-forcing 训练穿过（审计 GEO-2）。
    """
    off = _as_bf2(offset_px, "offset_px")
    f = off.shape[1]
    if not 0 <= a <= k < f:
        raise ValueError(f"窗口下标越界: a={a}, k={k}, f={f}")
    return off[:, a:k + 1] - off[:, a:a + 1]


def window_context(context: torch.Tensor, context_ids: torch.Tensor, a: int, k: int):
    """`window_offsets` 的文本孪生：同一段 [a, k] 的 context / context_ids。
    单独写出来只有一个理由——把 "偏移窗口" 和 "文本窗口" 的下标口径钉在同一个函数里，
    stage3 的 scored tail 才不会两边差一格。
    """
    ctx, ids = _check_context(context, context_ids)
    f = ctx.shape[1]
    if not 0 <= a <= k < f:
        raise ValueError(f"窗口下标越界: a={a}, k={k}, f={f}")
    return ctx[:, a:k + 1], ids[:, a:k + 1]


# --------------------------------------------------------------------------
# rollout 延长：偏移
# --------------------------------------------------------------------------

def walk_speed_from_sidecar(sidecar) -> tuple[float, float]:
    """从 sidecar（路径或已加载的 dict）取实测步速 (px/s)。默认常量只是它的冻结副本。"""
    from actionrope.dataset import load_sidecar

    meta = load_sidecar(sidecar)["meta"]
    ws = meta["walk_speed_px_s"]
    if len(ws) != 2:
        raise ValueError(f"sidecar meta 的 walk_speed_px_s 应为 2 个数，得到 {ws!r}")
    return float(ws[0]), float(ws[1])


def append_cell_velocity(offset_px: torch.Tensor, vel: torch.Tensor) -> torch.Tensor:
    """把逐 cell 速度 vel [b, m, 2]（px/s）接到 offset_px [b, F0, 2] 后面 → [b, F0+m, 2]。
    用模块 docstring 的 (1.5·v_prev + 2.5·v_cur)/FPS 增量积分，即帧级积分 + cell 内均值的
    闭式解。第一个新 cell 需要"上一 cell 的速度"，而数据只给了偏移，所以从最后一段已知增量
    反推匀速等效值 v_prev（匀速时精确；变速时那一格的误差 ≤ 1.5·|Δv|/16）。
    """
    off = _as_bf2(offset_px, "offset_px")
    if vel.dim() != 3 or vel.shape[0] != off.shape[0] or vel.shape[2] != 2:
        raise ValueError(f"vel 形状应为 [{off.shape[0]}, m, 2]，得到 {tuple(vel.shape)}")
    b, f0 = off.shape[:2]
    if f0 < 1:
        raise ValueError("offset_px 至少要有 cell 0（形状 [b, ≥1, 2]）")
    m = vel.shape[1]
    if m == 0:
        return off.clone()
    vel = vel.to(device=off.device, dtype=off.dtype)

    v_prev = _implied_last_velocity(off)  # [b, 2] px/s
    out = [off]
    cur = off[:, -1]                       # [b, 2] 最后一个已知 cell 的累计偏移
    k = f0                                 # 即将写入的绝对 cell 下标
    for j in range(m):
        v = vel[:, j]
        if k == 0:
            raise ValueError("cell 0 的偏移恒为 0，不能由速度生成")
        if k == 1:
            # cell 0 只有帧 0，起步项只有 2.5 那一半
            delta = (_W_CUR * v) / FPS
        else:
            delta = (_W_PREV * v_prev + _W_CUR * v) / FPS
        cur = cur + delta
        out.append(cur.unsqueeze(1))
        v_prev = v
        k += 1
    return torch.cat(out, dim=1)


def extend_offsets(offset_px: torch.Tensor, N: int, mode: str = "resample",
                   chunk: int = DEFAULT_CHUNK_CELLS, rng: torch.Generator | None = None,
                   walk_speed=None, speed_mult: tuple[float, float] = DEFAULT_SPEED_MULT,
                   include_idle: bool = True) -> torch.Tensor:
    """把 [b, F0, 2] 的逐 cell 偏移延长/裁剪到 rollout 需要的 [b, N+1, 2]。
    N = rollout 生成的 cell 数（不含 cell 0），所以长度是 1+N，与
    `dmd_gm.py:628-637` 的 `prompts.shape[1] == 1 + N` 口径一致。
    1+N ≤ F0 时就是切片（原样保留真值），不会引入任何合成运动。
    mode:
      * "resample"（默认）—— 每 `chunk` 个 cell 抽一次动作：方向取 `geometry.ACTION_DIRS`
        的 9 类之一（含 idle），速度按轴取 walk_speed × 倍率（[0.8, 1.2] 均匀），
        与 `infer.py::parse_actions` 的 `ACTION_DIRS[i] * walk_speed * mult` 完全同构，
        对应 ReactiveGWM 的 `dmd_action_extend_mode: preset_chunks`。逐样本独立抽。
      * "continue" —— 沿用最后一段已知增量隐含的匀速（等价于把 delta_tok[-1] 一路重复）。
      * "loop" —— 合成的走出去再原路走回来（见分支内注释）。
      * "retrace" —— 沿真实轨迹倒放回 cell 0 再顺放，回访 clip 自己的 cell（画布记忆训练用）。
    rng: CPU 上的 torch.Generator。抽样一律在 CPU 上做再搬到 offset_px.device，
    这样 4 张卡上的种子行为与设备无关；各 rank 抽到不同方向没关系（数据本来就不同），
    **形状**由 rank-synced 的 N 决定，FSDP 的 all-gather 不会错位。
    """
    off = _as_bf2(offset_px, "offset_px")
    b, f0 = off.shape[:2]
    if N < 0:
        raise ValueError(f"N 应 ≥ 0，得到 {N}")
    want = N + 1
    if want <= f0:
        return off[:, :want].contiguous()
    m = want - f0

    if mode == "continue":
        v = _implied_last_velocity(off).unsqueeze(1).expand(b, m, 2)  # [b, m, 2]
    elif mode == "resample":
        v = _sample_chunk_velocity(b, m, chunk=chunk, rng=rng, walk_speed=walk_speed,
                                   speed_mult=speed_mult, include_idle=include_idle)
        v = v.to(device=off.device, dtype=off.dtype)
    elif mode == "loop":
        # 会**回访**的合成动作（World-Canvas Memory 训练用）：前一半按 resample 抽 chunk 速度走出去，
        # 后一半把这些速度**倒序取反**走回来 —— 回程的每个 cell 都精确落在去程某个 cell 的世界坐标上，
        # 而且间隔 ≥ 窗口长度时正是"被驱逐后再回访"的情形。自然数据（15 s 长 clip）平均只有 4/40 个回访 cell，不够。
        half = (m + 1) // 2
        v_out = _sample_chunk_velocity(b, half, chunk=chunk, rng=rng, walk_speed=walk_speed,
                                       speed_mult=speed_mult, include_idle=False)
        v_back = -torch.flip(v_out, dims=[1])
        v = torch.cat([v_out, v_back], dim=1)[:, :m].to(device=off.device, dtype=off.dtype)
    elif mode == "retrace":
        # 沿**真实**轨迹来回走（World-Canvas Memory 训练用，2026-09-15）：延长段先把 clip 的坐标倒序走回去
        # （cell F0−1+i **精确**等于 cell F0−1−i 的世界坐标，时间距离 2i），走到 cell 0 的位置后再顺着原路
        # 走出去，如此往复。直接抄坐标而不走 append_cell_velocity 的速度积分：那条公式带 (1.5·v_prev+2.5·v_cur)
        # 的跨 cell 平滑，回程会与去程错开几个 px，而回访一致性 loss 与画布选格都要求逐位相同的世界坐标。
        # 与 "loop" 的区别：回访的是 clip 自己的 cell（含真实首帧 cell 0），而且 F0=21 时回程 20 格里有 12 格
        # 与其伙伴的时间距离 > 17（被 recency 窗口驱逐、只能由画布供给），"loop" 要 N≈60 才有同样多。
        assert f0 >= 2, "retrace 至少要 2 个真实 cell"
        back_leg = torch.flip(off[:, :-1], dims=[1])              # [b, F0−1, 2]: cell F0−2, …, 0
        fwd_leg = off[:, 1:]                                      # [b, F0−1, 2]: cell 1, …, F0−1
        legs, n_have, back = [], 0, True
        while n_have < m:
            leg = back_leg if back else fwd_leg
            legs.append(leg); n_have += leg.shape[1]; back = not back
        ext = torch.cat([off, torch.cat(legs, dim=1)[:, :m]], dim=1).contiguous()
        assert ext.shape == (b, want, 2), f"延长失败: {tuple(ext.shape)} != {(b, want, 2)}"
        return ext
    else:
        raise ValueError(f"未知 mode {mode!r}，可选 'resample' | 'continue' | 'loop' | 'retrace'")
    ext = append_cell_velocity(off, v)
    assert ext.shape == (b, want, 2), f"延长失败: {tuple(ext.shape)} != {(b, want, 2)}"
    return ext


# --------------------------------------------------------------------------
# rollout 延长：文本
# --------------------------------------------------------------------------

def extend_context(context: torch.Tensor, context_ids: torch.Tensor, N: int, mode: str = "hold"):
    """把 [b, F0, L, C] 的逐 cell 文本与 [b, F0] 的 id 延长/裁剪到 1+N 个 cell。
    返回 (context_ext [b, 1+N, L, C], ids_ext [b, 1+N])，与 `extend_offsets` 逐下标对齐。
    mode="hold"（唯一实现）：cell ≥ F0 一律复制**最后一条已写好的** cell 文本及其 id，
    照抄 `_wm/WorldMind/distill/code/stage3_dmd/dmd_gm.py:628-637`。ARoPE 臂的 prompt 已被
    `prompts.strip_action` 剥成"只剩场景"，rollout 又不换场景，所以保持是语义正确的；
    id 必须跟着一起保持，`arope_forward` 靠它去重 text_embedding，id 错了就会拿错一行文本。
    转场语料（burn / hardcut）的 cell 20 已经是烧完之后的目标场景，保持它同样正确。
    注意：文本"保持"而偏移"继续走"，两者策略不同是**有意的**（见模块 docstring），
    不是这里漏了一种模式。
    """
    ctx, ids = _check_context(context, context_ids)
    b, f0 = ctx.shape[:2]
    if N < 0:
        raise ValueError(f"N 应 ≥ 0，得到 {N}")
    if mode != "hold":
        raise ValueError(f"未知 mode {mode!r}：文本只实现 'hold'（保持最后一条已写好的 cell 文本）")
    want = N + 1
    if want <= f0:
        return ctx[:, :want].contiguous(), ids[:, :want].contiguous()
    pad_n = want - f0
    ctx_pad = ctx[:, f0 - 1:f0].expand(b, pad_n, *ctx.shape[2:])
    ids_pad = ids[:, f0 - 1:f0].expand(b, pad_n)
    ctx_ext = torch.cat([ctx, ctx_pad], dim=1)
    ids_ext = torch.cat([ids, ids_pad], dim=1)
    assert ctx_ext.shape[1] == want and ids_ext.shape[1] == want, (
        f"文本延长失败: {ctx_ext.shape[1]}/{ids_ext.shape[1]} != {want}")
    return ctx_ext, ids_ext


def assert_rollout_aligned(offset_ext: torch.Tensor, context_ext: torch.Tensor,
                           ids_ext: torch.Tensor, N: int, latent=None) -> None:
    """rollout 前的对齐闸：三条（可选四条）流的 cell 数必须都是 1+N，且 batch 一致。
    stage3 每次 rollout 之前、每次 teacher/critic forward 之前都调一下——这类错位
    不会报错、只会让 loss 照降而训错东西，所以要用断言把它变响。
    """
    want = N + 1
    b = offset_ext.shape[0]
    if offset_ext.shape != (b, want, 2):
        raise AssertionError(f"offset 形状 {tuple(offset_ext.shape)} != {(b, want, 2)}")
    if context_ext.shape[:2] != (b, want):
        raise AssertionError(f"context 形状 {tuple(context_ext.shape)} 前两维 != {(b, want)}")
    if tuple(ids_ext.shape) != (b, want):
        raise AssertionError(f"context_ids 形状 {tuple(ids_ext.shape)} != {(b, want)}")
    if latent is not None and (latent.shape[0] != b or latent.shape[2] != want):
        raise AssertionError(f"latent 形状 {tuple(latent.shape)} 与 [b={b}, ·, f={want}, ·, ·] 不符")


# --------------------------------------------------------------------------
# 内部
# --------------------------------------------------------------------------

def _as_bf2(off: torch.Tensor, name: str) -> torch.Tensor:
    if not isinstance(off, torch.Tensor):
        raise TypeError(f"{name} 应为 Tensor，得到 {type(off)}")
    if off.dim() != 3 or off.shape[-1] != 2:
        raise ValueError(f"{name} 形状应为 [b, f, 2]，得到 {tuple(off.shape)}")
    return off


def _check_context(context: torch.Tensor, context_ids: torch.Tensor):
    if context.dim() != 4:
        raise ValueError(f"context 形状应为 [b, f, L, C]，得到 {tuple(context.shape)}")
    if context_ids.dim() != 2:
        raise ValueError(f"context_ids 形状应为 [b, f]，得到 {tuple(context_ids.shape)}")
    if context.shape[:2] != context_ids.shape:
        raise ValueError(f"context {tuple(context.shape)[:2]} 与 context_ids "
                         f"{tuple(context_ids.shape)} 的 (b, f) 不一致")
    return context, context_ids


def _implied_last_velocity(off: torch.Tensor) -> torch.Tensor:
    """从最后一段已知增量反推匀速等效速度 [b, 2]（px/s）。
    匀速时 delta[k] = 4·v/FPS（k ≥ 2），delta[1] = 2.5·v/FPS；F0 = 1（只有 cell 0）时
    没有任何增量可用，只能取 0（此时 'continue' 退化成"站住不动"，调用方应当用 'resample'）。
    """
    b, f0 = off.shape[:2]
    if f0 < 2:
        return off.new_zeros(b, 2)
    delta = off[:, -1] - off[:, -2]
    w = _W_CUR if f0 == 2 else (_W_PREV + _W_CUR)  # 2.5（cell 1） / 4.0（cell ≥ 2）
    return delta * (FPS / w)


def _sample_chunk_velocity(b: int, m: int, chunk: int, rng: torch.Generator | None,
                           walk_speed, speed_mult: tuple[float, float],
                           include_idle: bool) -> torch.Tensor:
    """逐样本、每 chunk 个 cell 抽一次动作 → 逐 cell 速度 [b, m, 2]（px/s，float32 CPU）。"""
    if chunk < 1:
        raise ValueError(f"chunk 应 ≥ 1，得到 {chunk}")
    if rng is not None and rng.device.type != "cpu":
        raise ValueError("rng 必须是 CPU 上的 torch.Generator（抽样固定在 CPU 做，保证跨设备一致）")
    lo, hi = float(speed_mult[0]), float(speed_mult[1])
    if not 0.0 <= lo <= hi:
        raise ValueError(f"speed_mult 应为 0 ≤ lo ≤ hi，得到 {speed_mult!r}")
    ws = WALK_SPEED_PX_S if walk_speed is None else walk_speed
    ws = torch.tensor([float(ws[0]), float(ws[1])], dtype=torch.float32)
    dirs = torch.from_numpy(np.asarray(ACTION_DIRS, dtype=np.float32))  # [9, 2] 单位方向，0 = idle
    n_dir = dirs.shape[0]
    lo_idx = 0 if include_idle else 1

    n_chunk = math.ceil(m / chunk)
    idx = torch.randint(lo_idx, n_dir, (b, n_chunk), generator=rng)          # [b, n_chunk]
    mult = torch.rand((b, n_chunk, 1), generator=rng) * (hi - lo) + lo       # [b, n_chunk, 1]
    v_chunk = dirs[idx] * ws.view(1, 1, 2) * mult                            # [b, n_chunk, 2]
    # 每个 chunk 的速度铺满它覆盖的 cell（最后一个 chunk 可能不满，截断）
    return v_chunk.repeat_interleave(chunk, dim=1)[:, :m].contiguous()


__all__ = [
    "ACTION_DIRS", "CELLS", "CELL_FRAMES", "DEFAULT_CHUNK_CELLS", "DEFAULT_SPEED_MULT",
    "WALK_SPEED_PX_S", "append_cell_velocity", "assert_rollout_aligned", "extend_context",
    "extend_offsets", "make_mask_input", "make_weight_map", "to_device",
    "walk_speed_from_sidecar", "window_context", "window_offsets",
]