File size: 6,501 Bytes
7e86b86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """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 # matches src/data/formats.py motion description
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 # tests set True: t/delta derived from index
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 # need t+1
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))
# angles accumulate unbounded in the data; wrap the *input feature* so the
# expert sees rotation state mod 2pi (the delta target is already wrapped)
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,
}
|