| """Dataset and collate for dual-rate training. |
| |
| Each item: a random (scene, t, staleness delta) triple -> |
| - Qwen input ids for the text context ending at frame index t - delta |
| (header + as many recent frames as fit the token budget + suffix) |
| - normalized current state at t, normalized target delta t -> t+1 |
| - static object features, gravity/dt conditioning vector |
| |
| Frame text rendering replicates src/training/data_loader._format_frame_text |
| byte-for-byte (including the baseline's duplicated "Frame N:" prefix, since |
| frame['description'] itself starts with "Frame N:"). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from .normalization import wrap_angle |
|
|
| VELOCITY_THRESHOLD = 1.0 |
| SUFFIX = "Predict next frame:" |
|
|
|
|
| def render_frame_text(frame_num: int, state: np.ndarray, ids: np.ndarray) -> str: |
| """state: (N, 6). Replicates the baseline text format for one frame.""" |
| moving = int(np.sum((np.abs(state[:, 2]) > VELOCITY_THRESHOLD) |
| | (np.abs(state[:, 3]) > VELOCITY_THRESHOLD))) |
| n = state.shape[0] |
| if moving == 0: |
| desc = "All objects are at rest" |
| elif moving == n: |
| desc = "All objects are in motion" |
| else: |
| desc = f"{moving} of {n} objects are moving" |
| lines = [f"Frame {frame_num}: Frame {frame_num}: {desc}."] |
| for i in range(n): |
| px, py, vx, vy, a, av = (float(x) for x in state[i]) |
| s = f" obj_{int(ids[i])}: pos=({px:.4f}, {py:.4f}), vel=({vx:.4f}, {vy:.4f})" |
| if abs(a) > 0.001 or abs(av) > 0.001: |
| s += f", a={a:.4f}, av={av:.4f}" |
| lines.append(s) |
| lines.append("") |
| return "\n".join(lines) |
|
|
|
|
| def build_context_ids(tokenizer, header_text: str, states: np.ndarray, ids: np.ndarray, |
| ctx_end: int, max_ctx_tokens: int) -> list: |
| """Token ids for header + recent frames [.. ctx_end] + suffix, newest-first fill. |
| |
| Frame index i in `states` is frame number i+1 (frames are 1-indexed in data). |
| """ |
| head = tokenizer(header_text, add_special_tokens=False)["input_ids"] |
| tail = tokenizer("\n" + SUFFIX, add_special_tokens=False)["input_ids"] |
| budget = max_ctx_tokens - len(head) - len(tail) |
| frame_ids = [] |
| for i in range(ctx_end, -1, -1): |
| ft = render_frame_text(i + 1, states[i], ids) |
| fi = tokenizer(ft + "\n", add_special_tokens=False)["input_ids"] |
| if len(fi) > budget: |
| break |
| frame_ids.append(fi) |
| budget -= len(fi) |
| frame_ids.reverse() |
| out = list(head) |
| for fi in frame_ids: |
| out.extend(fi) |
| out.extend(tail) |
| return out |
|
|
|
|
| class DualRateDataset(Dataset): |
| def __init__(self, cache_dir: str, split: str, tokenizer, cfg, normalizer, |
| max_scenes: int | None = None, seed: int = 0): |
| self.files = sorted(Path(cache_dir, split).glob("*.npz")) |
| if max_scenes: |
| rng = np.random.default_rng(seed) |
| idx = rng.permutation(len(self.files))[:max_scenes] |
| self.files = [self.files[i] for i in sorted(idx)] |
| if not self.files: |
| raise FileNotFoundError(f"no cache files in {cache_dir}/{split}") |
| self.tokenizer = tokenizer |
| self.cfg = cfg |
| self.norm = normalizer |
| self.seed = seed |
| self.deterministic = False |
|
|
| def __len__(self): |
| return len(self.files) |
|
|
| def __getitem__(self, idx): |
| z = np.load(self.files[idx]) |
| states, static, ids = z["states"], z["static"], z["ids"] |
| header_text = str(z["header_text"]) |
| gravity, dt = z["gravity"], float(z["dt"]) |
| t_max = states.shape[0] - 2 |
| min_ctx = self.cfg.min_ctx_frames |
|
|
| if self.deterministic: |
| rng = np.random.default_rng(self.seed + idx) |
| else: |
| rng = np.random.default_rng() |
| t = int(rng.integers(min_ctx, t_max + 1)) |
| delta = int(rng.integers(0, self.cfg.tick_k)) |
| ctx_end = max(t - delta, min_ctx - 1) |
|
|
| input_ids = build_context_ids(self.tokenizer, header_text, states, ids, |
| ctx_end, self.cfg.max_ctx_tokens) |
|
|
| cur = states[t].astype(np.float32) |
| target_delta = self.norm.compute_delta(cur, states[t + 1].astype(np.float32)) |
| |
| |
| cur_feat = cur.copy() |
| cur_feat[:, 4] = wrap_angle(cur_feat[:, 4]) |
| gravity_dt = np.array([gravity[0] / 981.0, gravity[1] / 981.0, dt * 60.0], |
| dtype=np.float32) |
| return { |
| "input_ids": input_ids, |
| "cur_state": torch.from_numpy(self.norm.norm_state(cur_feat)), |
| "target_delta": torch.from_numpy(self.norm.norm_delta(target_delta)), |
| "static": torch.from_numpy(static.astype(np.float32)), |
| "gravity_dt": torch.from_numpy(gravity_dt), |
| } |
|
|
|
|
| def collate_dualrate(batch, pad_token_id: int): |
| """Left-pad text, zero-pad objects to max N in batch, build masks.""" |
| b = len(batch) |
| max_l = max(len(x["input_ids"]) for x in batch) |
| max_n = max(x["cur_state"].shape[0] for x in batch) |
| s_dim = batch[0]["cur_state"].shape[-1] |
| f_dim = batch[0]["static"].shape[-1] |
|
|
| input_ids = torch.full((b, max_l), pad_token_id, dtype=torch.long) |
| attention_mask = torch.zeros((b, max_l), dtype=torch.long) |
| cur_state = torch.zeros((b, max_n, s_dim)) |
| target_delta = torch.zeros((b, max_n, s_dim)) |
| static = torch.zeros((b, max_n, f_dim)) |
| obj_mask = torch.zeros((b, max_n), dtype=torch.bool) |
| gravity_dt = torch.stack([x["gravity_dt"] for x in batch]) |
|
|
| for i, x in enumerate(batch): |
| l = len(x["input_ids"]) |
| input_ids[i, max_l - l:] = torch.tensor(x["input_ids"], dtype=torch.long) |
| attention_mask[i, max_l - l:] = 1 |
| n = x["cur_state"].shape[0] |
| cur_state[i, :n] = x["cur_state"] |
| target_delta[i, :n] = x["target_delta"] |
| static[i, :n] = x["static"] |
| obj_mask[i, :n] = True |
|
|
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "cur_state": cur_state, |
| "target_delta": target_delta, |
| "static": static, |
| "obj_mask": obj_mask, |
| "gravity_dt": gravity_dt, |
| } |
|
|
