| ''' |
| data.py defines the synthetic benchmark data pipeline for the spectral world model project. |
| |
| Specifies vocabulary and action labels for text-conditioned dynamics |
| Generates moving shape sequences with image frames, text tokens, actions, and rewards |
| Saves train/val/test splits to a compressed .npz benchmark file |
| Provides TransitionDataset for frame-to-frame transition learning |
| Provides SequenceDataset for full-sequence rollout evaluation |
| ''' |
| |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| import json |
| import math |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
|
|
| VOCAB = [ |
| "<pad>", |
| "square", "circle", |
| "stay", "up", "down", "left", "right", |
| "slow", "fast", |
| "x0", "x1", "x2", "x3", |
| "y0", "y1", "y2", "y3", |
| ] |
| TOKEN_TO_ID = {tok: i for i, tok in enumerate(VOCAB)} |
| ACTION_NAMES = ["stay", "up", "down", "left", "right"] |
|
|
|
|
| @dataclass |
| class DatasetConfig: |
| image_size: int = 32 |
| seq_len: int = 8 |
| max_text_len: int = 5 |
| train_sequences: int = 256 |
| val_sequences: int = 64 |
| test_sequences: int = 64 |
| seed: int = 7 |
|
|
|
|
| def _draw_square(canvas: np.ndarray, x: int, y: int, size: int) -> None: |
| h, w = canvas.shape |
| x0 = max(0, x - size // 2); x1 = min(w, x + size // 2 + 1) |
| y0 = max(0, y - size // 2); y1 = min(h, y + size // 2 + 1) |
| canvas[y0:y1, x0:x1] = 1.0 |
|
|
|
|
| def _draw_circle(canvas: np.ndarray, x: int, y: int, radius: int) -> None: |
| h, w = canvas.shape |
| yy, xx = np.ogrid[:h, :w] |
| mask = (xx - x) ** 2 + (yy - y) ** 2 <= radius**2 |
| canvas[mask] = 1.0 |
|
|
|
|
| def _text_tokens(shape: str, action: int, speed: int, x: int, y: int, image_size: int, max_len: int) -> np.ndarray: |
| xb = min(3, max(0, int(4 * x / image_size))) |
| yb = min(3, max(0, int(4 * y / image_size))) |
| words = [shape, ACTION_NAMES[action], "fast" if speed == 2 else "slow", f"x{xb}", f"y{yb}"] |
| ids = [TOKEN_TO_ID[w] for w in words][:max_len] |
| ids += [0] * (max_len - len(ids)) |
| return np.asarray(ids, dtype=np.int64) |
|
|
|
|
| def _apply_action(x: int, y: int, action: int, speed: int, image_size: int, radius: int) -> tuple[int, int]: |
| dx, dy = 0, 0 |
| if action == 1: dy = -speed |
| elif action == 2: dy = speed |
| elif action == 3: dx = -speed |
| elif action == 4: dx = speed |
| x = int(np.clip(x + dx, radius + 1, image_size - radius - 2)) |
| y = int(np.clip(y + dy, radius + 1, image_size - radius - 2)) |
| return x, y |
|
|
|
|
| def generate_sequences(n_sequences: int, cfg: DatasetConfig, rng: np.random.Generator) -> dict[str, np.ndarray]: |
| images = np.zeros((n_sequences, cfg.seq_len, 1, cfg.image_size, cfg.image_size), dtype=np.float32) |
| tokens = np.zeros((n_sequences, cfg.seq_len, cfg.max_text_len), dtype=np.int64) |
| actions = np.zeros((n_sequences, cfg.seq_len - 1), dtype=np.int64) |
| rewards = np.zeros((n_sequences, cfg.seq_len - 1), dtype=np.float32) |
|
|
| for n in range(n_sequences): |
| shape = "square" if rng.random() < 0.5 else "circle" |
| speed = int(rng.choice([1, 2])) |
| radius = 3 if shape == "circle" else 4 |
| x = int(rng.integers(radius + 2, cfg.image_size - radius - 2)) |
| y = int(rng.integers(radius + 2, cfg.image_size - radius - 2)) |
| last_action = 0 |
| for t in range(cfg.seq_len): |
| canvas = np.zeros((cfg.image_size, cfg.image_size), dtype=np.float32) |
| if shape == "square": |
| _draw_square(canvas, x, y, 6) |
| else: |
| _draw_circle(canvas, x, y, radius) |
| |
| canvas += rng.normal(0.0, 0.01, size=canvas.shape).astype(np.float32) |
| canvas = np.clip(canvas, 0.0, 1.0) |
| images[n, t, 0] = canvas |
| tokens[n, t] = _text_tokens(shape, last_action, speed, x, y, cfg.image_size, cfg.max_text_len) |
| if t < cfg.seq_len - 1: |
| action = int(rng.integers(0, len(ACTION_NAMES))) |
| actions[n, t] = action |
| x, y = _apply_action(x, y, action, speed, cfg.image_size, radius) |
| last_action = action |
| |
| rewards[n, t] = float((x > cfg.image_size * 0.65) and (y < cfg.image_size * 0.35)) |
| return {"images": images, "tokens": tokens, "actions": actions, "rewards": rewards} |
|
|
|
|
| def generate_benchmark_npz(path: str | Path, cfg: DatasetConfig = DatasetConfig()) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| rng = np.random.default_rng(cfg.seed) |
| train = generate_sequences(cfg.train_sequences, cfg, rng) |
| val = generate_sequences(cfg.val_sequences, cfg, rng) |
| test = generate_sequences(cfg.test_sequences, cfg, rng) |
| np.savez_compressed( |
| path, |
| train_images=train["images"], train_tokens=train["tokens"], train_actions=train["actions"], train_rewards=train["rewards"], |
| val_images=val["images"], val_tokens=val["tokens"], val_actions=val["actions"], val_rewards=val["rewards"], |
| test_images=test["images"], test_tokens=test["tokens"], test_actions=test["actions"], test_rewards=test["rewards"], |
| vocab=np.asarray(VOCAB), action_names=np.asarray(ACTION_NAMES), |
| config=json.dumps(cfg.__dict__), |
| ) |
|
|
|
|
| class TransitionDataset(Dataset): |
| def __init__(self, npz_path: str | Path, split: str): |
| data = np.load(npz_path, allow_pickle=True) |
| self.images = data[f"{split}_images"] |
| self.tokens = data[f"{split}_tokens"] |
| self.actions = data[f"{split}_actions"] |
| self.seq_len = self.images.shape[1] |
| self.index = [(i, t) for i in range(self.images.shape[0]) for t in range(self.seq_len - 1)] |
|
|
| def __len__(self) -> int: |
| return len(self.index) |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: |
| i, t = self.index[idx] |
| return { |
| "image": torch.from_numpy(self.images[i, t]), |
| "tokens": torch.from_numpy(self.tokens[i, t]).long(), |
| "action": torch.tensor(self.actions[i, t]).long(), |
| "next_image": torch.from_numpy(self.images[i, t + 1]), |
| "next_tokens": torch.from_numpy(self.tokens[i, t + 1]).long(), |
| } |
|
|
|
|
| class SequenceDataset(Dataset): |
| def __init__(self, npz_path: str | Path, split: str): |
| data = np.load(npz_path, allow_pickle=True) |
| self.images = data[f"{split}_images"] |
| self.tokens = data[f"{split}_tokens"] |
| self.actions = data[f"{split}_actions"] |
| self.rewards = data[f"{split}_rewards"] |
|
|
| def __len__(self) -> int: |
| return self.images.shape[0] |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: |
| return { |
| "images": torch.from_numpy(self.images[idx]), |
| "tokens": torch.from_numpy(self.tokens[idx]).long(), |
| "actions": torch.from_numpy(self.actions[idx]).long(), |
| "rewards": torch.from_numpy(self.rewards[idx]), |
| } |
|
|
