| """World v2: free-form captions over a much richer procedural world.
|
|
|
| Text is tokenized as raw UTF-8 bytes, so any string is accepted. What the model
|
| *understands* is whatever this world teaches it: many colours and shapes, single-
|
| object events (moving, falling, jumping, orbiting, fading, changing colour ...)
|
| and two-object interactions (merging with colour mixing, splitting, colliding,
|
| chasing, orbiting, swapping places), each described with many paraphrases.
|
|
|
| Every clip comes with per-frame, per-object ground truth in `meta` for scoring.
|
| """
|
| from __future__ import annotations
|
|
|
| import math
|
| import re
|
| import zlib
|
| from dataclasses import dataclass
|
|
|
| import numpy as np
|
| import torch
|
|
|
|
|
|
|
| PAD, NULL, BOS = 256, 257, 258
|
| BYTE_VOCAB = 259
|
| MAX_BYTES = 256
|
|
|
|
|
| class ByteTokenizer:
|
| """UTF-8 bytes -> ids. No vocabulary: every string is representable."""
|
| vocab_size = BYTE_VOCAB
|
|
|
| def __init__(self, max_len: int = MAX_BYTES):
|
| self.max_len = max_len
|
|
|
| def encode(self, caption: str) -> list[int]:
|
| return [BOS] + list(caption.encode("utf-8"))[: self.max_len - 1]
|
|
|
| def __call__(self, captions: list[str], pad_to_longest: bool = True) -> torch.Tensor:
|
| """[B, L] ids; L = longest caption in the batch (dynamic padding) or max_len."""
|
| encoded = [self.encode(c) for c in captions]
|
| L = max(map(len, encoded)) if pad_to_longest else self.max_len
|
| out = torch.full((len(captions), L), PAD, dtype=torch.long)
|
| for i, ids in enumerate(encoded):
|
| out[i, : len(ids)] = torch.tensor(ids)
|
| return out
|
|
|
| def null(self, n: int, length: int = 1) -> torch.Tensor:
|
| out = torch.full((n, length), PAD, dtype=torch.long)
|
| out[:, 0] = NULL
|
| return out
|
|
|
|
|
| SEP = 259
|
| HASH_OFFSET, HASH_BUCKETS = 260, 16384
|
| HYBRID_VOCAB = HASH_OFFSET + HASH_BUCKETS
|
| MAX_HYBRID = 320
|
| _WORD = re.compile(r"[^\W_]+", re.UNICODE)
|
|
|
|
|
| def word_bucket(word: str) -> int:
|
| """Stable hash of a case-folded word into one of HASH_BUCKETS ids (no vocabulary; any word maps somewhere)."""
|
| return HASH_OFFSET + zlib.crc32(word.casefold().encode("utf-8")) % HASH_BUCKETS
|
|
|
|
|
| class HybridTokenizer(ByteTokenizer):
|
| """[BOS] utf-8 bytes [SEP] hashed words. Bytes keep spelling; hashed words give the encoder word identity
|
| from the first step. Still no vocabulary: every string encodes, nothing is UNK."""
|
| vocab_size = HYBRID_VOCAB
|
|
|
| def __init__(self, max_len: int = MAX_HYBRID, max_bytes: int = MAX_BYTES):
|
| super().__init__(max_len)
|
| self.max_bytes = max_bytes
|
|
|
| def encode(self, caption: str) -> list[int]:
|
| words = [word_bucket(w) for w in _WORD.findall(caption)]
|
| byte_ids = list(caption.encode("utf-8"))[: self.max_bytes - 1]
|
| room = self.max_len - 2 - len(byte_ids)
|
| return [BOS] + byte_ids + [SEP] + words[: max(0, room)]
|
|
|
|
|
|
|
|
|
| COLORS = {
|
| "red": (225, 35, 35), "orange": (250, 140, 20), "yellow": (245, 220, 40), "green": (40, 190, 60),
|
| "cyan": (40, 215, 230), "blue": (45, 85, 240), "purple": (145, 55, 215), "pink": (250, 125, 190),
|
| "brown": (135, 80, 35), "white": (245, 245, 245), "gray": (140, 140, 140), "black": (12, 12, 12),
|
| }
|
| COLOR_WORDS = {
|
| "red": ["red", "red", "red", "crimson", "scarlet"], "orange": ["orange"], "yellow": ["yellow", "yellow", "golden"],
|
| "green": ["green", "green", "lime"], "cyan": ["cyan", "cyan", "turquoise", "light blue"],
|
| "blue": ["blue", "blue", "dark blue"], "purple": ["purple", "purple", "violet"], "pink": ["pink"],
|
| "brown": ["brown"], "white": ["white"], "gray": ["gray", "grey", "silver"], "black": ["black"],
|
| }
|
|
|
| MIX = {
|
| frozenset(["red", "blue"]): "purple", frozenset(["red", "yellow"]): "orange",
|
| frozenset(["blue", "yellow"]): "green", frozenset(["red", "white"]): "pink",
|
| frozenset(["black", "white"]): "gray", frozenset(["blue", "green"]): "cyan",
|
| frozenset(["red", "green"]): "brown", frozenset(["blue", "white"]): "cyan",
|
| frozenset(["purple", "white"]): "pink", frozenset(["orange", "blue"]): "brown",
|
| }
|
|
|
| SHAPES = ["circle", "square", "triangle", "star", "ring", "cross", "heart", "diamond", "hexagon"]
|
| SHAPE_WORDS = {
|
| "circle": ["circle", "circle", "circle", "ball", "dot", "disc"], "square": ["square", "square", "box", "block"],
|
| "triangle": ["triangle"], "star": ["star"], "ring": ["ring", "ring", "donut"], "cross": ["cross", "plus sign"],
|
| "heart": ["heart"], "diamond": ["diamond"], "hexagon": ["hexagon"],
|
| }
|
| ROTATION_INVARIANT = {"circle", "ring"}
|
|
|
| SIZES = {"tiny": 0.085, "small": 0.115, "medium": 0.155, "big": 0.205, "huge": 0.26}
|
| SIZE_WORDS = {"tiny": ["tiny"], "small": ["small", "little"], "medium": [""], "big": ["big", "large"],
|
| "huge": ["huge", "giant"]}
|
|
|
| BACKGROUNDS = {
|
| "black": ((0, 0, 0), (0, 0, 0)), "white": ((250, 250, 250), (250, 250, 250)),
|
| "gray": ((120, 120, 120), (120, 120, 120)), "navy": ((10, 15, 70), (10, 15, 70)),
|
| "sky": ((110, 170, 240), (200, 230, 255)), "sunset": ((90, 40, 130), (250, 150, 60)),
|
| "grass": ((120, 190, 250), (40, 140, 40)),
|
| }
|
| BG_PHRASES = {
|
| "black": ["on a black background", "against a black background", "on black", "in the dark", "in darkness"],
|
| "white": ["on a white background", "against a white background", "on white", "on a blank white page"],
|
| "gray": ["on a gray background", "on a grey background", "against a gray wall"],
|
| "navy": ["at night", "in the night sky", "on a dark blue background", "on a navy background"],
|
| "sky": ["in the sky", "in a blue sky", "against the sky", "in the daytime sky"],
|
| "sunset": ["at sunset", "in a sunset sky", "against a sunset", "at dusk"],
|
| "grass": ["over a field of grass", "above green grass", "over a meadow", "in a grassy field"],
|
| }
|
|
|
| DIRS = {
|
| "left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1),
|
| "up-left": (-1, -1), "up-right": (1, -1), "down-left": (-1, 1), "down-right": (1, 1),
|
| }
|
| DIR_PHRASES = {
|
| "left": ["left", "to the left", "leftward", "toward the left", "from right to left"],
|
| "right": ["right", "to the right", "rightward", "toward the right", "from left to right"],
|
| "up": ["up", "upward", "upwards", "toward the top", "from bottom to top"],
|
| "down": ["down", "downward", "downwards", "toward the bottom", "from top to bottom"],
|
| "up-left": ["up and to the left", "diagonally up left", "toward the top left", "up-left"],
|
| "up-right": ["up and to the right", "diagonally up right", "toward the top right", "up-right"],
|
| "down-left": ["down and to the left", "diagonally down left", "toward the bottom left", "down-left"],
|
| "down-right": ["down and to the right", "diagonally down right", "toward the bottom right", "down-right"],
|
| }
|
| MOVE_VERBS = ["moves", "moves", "slides", "drifts", "travels", "goes", "glides", "floats", "heads"]
|
| MOVE_VERBS_ING = ["moving", "sliding", "drifting", "traveling", "going", "gliding", "floating", "heading"]
|
|
|
| SINGLE_EVENTS = {"move": 4, "fall": 1, "rise": 1, "jump": 1, "bounce": 1, "zigzag": 1, "circle": 1, "grow": 1,
|
| "shrink": 1, "spin": 1, "still": 1, "fade_in": 1, "fade_out": 1, "blink": 1, "recolor": 1.5,
|
| "teleport": 0.7}
|
| PAIR_EVENTS = {"merge": 9, "split": 1.2, "collide": 1, "chase": 1, "orbit": 1, "swap": 1, "independent": 1.8}
|
|
|
|
|
| HELDOUT = {("merge", "star"), ("jump", "heart"), ("orbit", "hexagon"), ("fall", "diamond"), ("spin", "cross"),
|
| ("recolor", "triangle")}
|
|
|
|
|
| @dataclass
|
| class WorldConfig:
|
| frames: int = 16
|
| size: int = 32
|
| supersample: int = 2
|
| p_pair: float = 0.5
|
| p_featured: float = 0.04
|
| merge_weight: float = 9.0
|
| p_mention_bg: float = 0.5
|
| p_default_black: float = 0.7
|
|
|
|
|
| def _choice(rng, seq):
|
| return seq[rng.integers(len(seq))]
|
|
|
|
|
| def _weighted(rng, weights: dict):
|
| keys = list(weights)
|
| p = np.array([weights[k] for k in keys], float)
|
| return keys[rng.choice(len(keys), p=p / p.sum())]
|
|
|
|
|
| def _contrast_ok(rgb, bg: str) -> bool:
|
| top, bottom = (np.array(c, float) for c in BACKGROUNDS[bg])
|
| c = np.array(rgb, float)
|
| return min(np.linalg.norm(c - top), np.linalg.norm(c - bottom)) > 110
|
|
|
|
|
| def nearest_color(rgb) -> str:
|
| c = np.array(rgb, float)
|
| return min(COLORS, key=lambda k: np.linalg.norm(c - np.array(COLORS[k], float)))
|
|
|
|
|
| def mix_colors(a: str, b: str) -> str:
|
| if a == b:
|
| return a
|
| return MIX.get(frozenset([a, b])) or nearest_color((np.array(COLORS[a]) + np.array(COLORS[b])) / 2)
|
|
|
|
|
|
|
|
|
| def noun_phrase(rng, obj: dict, article: str | None = None) -> str:
|
| size = _choice(rng, SIZE_WORDS[obj["size"]])
|
| color = _choice(rng, COLOR_WORDS[obj["color"]])
|
| shape = _choice(rng, SHAPE_WORDS[obj["shape"]])
|
| if article is None:
|
| article = _choice(rng, ["a", "a", "a", "the", "one", ""])
|
| if article.endswith("bigger"):
|
| size = ""
|
| if article == "a" and (size or color)[:1] in "aeiou":
|
| article = "an"
|
| return " ".join(w for w in (article, size, color, shape) if w)
|
|
|
|
|
| def _move_phrase(rng, d: str) -> str:
|
| dp = _choice(rng, DIR_PHRASES[d])
|
| if d == "up" and rng.random() < 0.25:
|
| return _choice(rng, ["rises", "floats up", "goes up"])
|
| return f"{_choice(rng, MOVE_VERBS)} {dp}"
|
|
|
|
|
| def describe_single(rng, obj: dict, ev: dict) -> str:
|
| np_ = noun_phrase(rng, obj)
|
| k = ev["kind"]
|
| if k == "move":
|
| vp = _move_phrase(rng, ev["dir"])
|
| if rng.random() < 0.2:
|
| vp = f"{_choice(rng, MOVE_VERBS_ING)} {_choice(rng, DIR_PHRASES[ev['dir']])}"
|
| return f"{np_} {vp}" if rng.random() < 0.5 else f"{np_}, {vp}"
|
| return f"{np_} {vp}"
|
| if k == "fall":
|
| return f"{np_} {_choice(rng, ['falls', 'falls down', 'drops', 'drops down', 'falls from the top'])}"
|
| if k == "rise":
|
| return f"{np_} {_choice(rng, ['rises from the bottom', 'rises up', 'floats up from the bottom', 'launches upward'])}"
|
| if k == "jump":
|
| d = ev.get("dir")
|
| tail = f" {_choice(rng, ['to the left', 'left'])}" if d == "left" else f" {_choice(rng, ['to the right', 'right'])}" if d == "right" else ""
|
| return f"{np_} {_choice(rng, ['jumps', 'hops', 'leaps'])}{tail}"
|
| if k == "bounce":
|
| return f"{np_} {_choice(rng, ['bounces around', 'bounces', 'bounces off the walls', 'ricochets around'])}"
|
| if k == "zigzag":
|
| return f"{np_} {_choice(rng, ['zigzags', 'moves in a zigzag'])} {_choice(rng, DIR_PHRASES[ev['dir']][:3])}"
|
| if k == "circle":
|
| return f"{np_} {_choice(rng, ['moves in a circle', 'goes around in a circle', 'circles around', 'travels in a loop'])}"
|
| if k == "grow":
|
| return f"{np_} {_choice(rng, ['grows', 'gets bigger', 'expands', 'grows larger', 'inflates'])}"
|
| if k == "shrink":
|
| return f"{np_} {_choice(rng, ['shrinks', 'gets smaller', 'shrinks away', 'deflates'])}"
|
| if k == "spin":
|
| return f"{np_} {_choice(rng, ['spins', 'rotates', 'spins in place', 'twirls'])}"
|
| if k == "still":
|
| return f"{np_} {_choice(rng, ['stays still', 'sits still', 'does not move', 'stays in place', 'is resting'])}"
|
| if k == "fade_in":
|
| return f"{np_} {_choice(rng, ['fades in', 'appears', 'slowly appears', 'materializes'])}"
|
| if k == "fade_out":
|
| return f"{np_} {_choice(rng, ['fades out', 'disappears', 'fades away', 'vanishes'])}"
|
| if k == "blink":
|
| return f"{np_} {_choice(rng, ['blinks', 'flashes', 'flickers', 'blinks on and off'])}"
|
| if k == "recolor":
|
| to = _choice(rng, COLOR_WORDS[ev["to"]])
|
| return f"{np_} {_choice(rng, ['turns', 'changes to', 'becomes', 'changes color to', 'turns into a'])} {to}".replace(
|
| "turns into a " + to, f"turns into a {to} {_choice(rng, SHAPE_WORDS[obj['shape']])}")
|
| if k == "teleport":
|
| return f"{np_} {_choice(rng, ['teleports', 'jumps to another spot', 'teleports across the screen'])}"
|
| raise ValueError(k)
|
|
|
|
|
| def describe_pair(rng, a: dict, b: dict, ev: dict) -> str:
|
| k = ev["kind"]
|
| A, B = noun_phrase(rng, a), noun_phrase(rng, b)
|
| if k == "merge":
|
| res = ev["result"]
|
| verb = _choice(rng, ["merge into", "merge into", "combine into", "fuse into", "blend into", "join to form",
|
| "come together to form", "melt together into"])
|
| tail = f" {verb} {noun_phrase(rng, res, article=_choice(rng, ['a', 'one', 'a single', 'a bigger']))}"
|
| if rng.random() < 0.25:
|
| tail = f" {_choice(rng, ['merge', 'merge together', 'fuse together', 'combine'])}"
|
| return f"{A} and {B}{tail}"
|
| if k == "split":
|
| src = ev["source"]
|
| S = noun_phrase(rng, src)
|
| if rng.random() < 0.25:
|
| return f"{S} {_choice(rng, ['splits in two', 'splits into two', 'divides in half', 'breaks in two'])}"
|
| return f"{S} {_choice(rng, ['splits into', 'breaks into', 'divides into', 'separates into'])} {noun_phrase(rng, a, 'a')} and {noun_phrase(rng, b, 'a')}"
|
| if k == "collide":
|
| return f"{A} and {B} {_choice(rng, ['collide and bounce apart', 'crash into each other', 'bump into each other', 'collide', 'smash together and bounce off'])}"
|
| if k == "chase":
|
| if rng.random() < 0.3:
|
| return f"{B} is chased by {A}"
|
| return f"{A} {_choice(rng, ['chases', 'follows', 'runs after', 'is chasing'])} {B}"
|
| if k == "orbit":
|
| return f"{A} {_choice(rng, ['orbits', 'circles around', 'revolves around', 'goes around'])} {B}"
|
| if k == "swap":
|
| return f"{A} and {B} {_choice(rng, ['swap places', 'trade places', 'switch places', 'swap positions'])}"
|
| if k == "independent":
|
| return f"{describe_single(rng, a, ev['a'])} {_choice(rng, ['and', 'and', 'while', ', and'])} {describe_single(rng, b, ev['b'])}".replace(" , and", ", and")
|
| raise ValueError(k)
|
|
|
|
|
| def _style(rng, caption: str, bg_phrase: str | None) -> str:
|
| if bg_phrase:
|
| caption = f"{bg_phrase}, {caption}" if rng.random() < 0.2 else f"{caption} {bg_phrase}"
|
| if rng.random() < 0.08:
|
| caption = _choice(rng, ["a video of ", "watch as ", "animation: ", "clip: "]) + caption
|
| if rng.random() < 0.3:
|
| caption = caption[0].upper() + caption[1:]
|
| if rng.random() < 0.25:
|
| caption += "."
|
| return caption.strip()
|
|
|
|
|
|
|
|
|
| def _uni(rng, lo, hi, size=None):
|
| """rng.uniform that tolerates an empty range (big objects in a small frame): collapses to the midpoint."""
|
| if hi < lo:
|
| lo = hi = (lo + hi) / 2
|
| return rng.uniform(lo, hi, size=size)
|
|
|
|
|
| def _state(x, y, r, ang, rgb, alpha=1.0):
|
| return [float(x), float(y), float(r), float(ang), tuple(float(c) for c in rgb), float(alpha)]
|
|
|
|
|
| def _single_states(rng, obj: dict, ev: dict, T: int, center=None) -> list:
|
| r = SIZES[obj["size"]] * _uni(rng, 0.92, 1.08)
|
| ang0 = _uni(rng, 0, 2 * math.pi)
|
| rgb = np.array(COLORS[obj["color"]], float)
|
| ts = np.linspace(0, 1, T)
|
| k = ev["kind"]
|
| m = r + 0.03
|
|
|
| def anywhere():
|
| return _uni(rng, m, 1 - m, size=2)
|
|
|
| if k == "move":
|
| dx, dy = DIRS[ev["dir"]]
|
| dist = _uni(rng, 0.4, 0.62) / (math.sqrt(2) if dx and dy else 1)
|
| dist = min(dist, (1 - 2 * m) * 0.98)
|
|
|
| def start(d):
|
| if d > 0:
|
| return _uni(rng, m, 1 - m - dist)
|
| if d < 0:
|
| return _uni(rng, m + dist, 1 - m)
|
| return _uni(rng, m, 1 - m)
|
| x0, y0 = start(dx), start(dy)
|
| return [_state(x0 + dx * dist * t, y0 + dy * dist * t, r, ang0, rgb) for t in ts]
|
| if k in ("fall", "rise"):
|
| x0 = _uni(rng, m, 1 - m)
|
| top, bot = m, 1 - m
|
| prog = ts ** 2 if k == "fall" else 1 - (1 - ts) ** 2
|
| return [_state(x0, top + (bot - top) * p if k == "fall" else bot - (bot - top) * p, r, ang0, rgb) for p in prog]
|
| if k == "jump":
|
| d = {"left": -1, "right": 1}.get(ev.get("dir"), 0)
|
| span = min(_uni(rng, 0.3, 0.5), max(0.0, 1 - 2 * m - 0.02)) * d
|
| x0 = _uni(rng, m + max(0, -span), 1 - m - max(0, span))
|
| ground = 1 - m
|
| height = _uni(rng, 0.35, min(0.6, 1 - 2 * m))
|
| return [_state(x0 + span * t, ground - height * 4 * t * (1 - t), r, ang0, rgb) for t in ts]
|
| if k == "bounce":
|
| lo, hi = m, 1 - m
|
| x, y = anywhere()
|
| th = _uni(rng, 0, 2 * math.pi)
|
| sp = _uni(rng, 1.6, 2.4)
|
| vx, vy = sp * math.cos(th), sp * math.sin(th)
|
| out, dt = [], 1 / (T - 1)
|
| for _ in range(T):
|
| out.append(_state(x, y, r, ang0, rgb))
|
| x, y = x + vx * dt, y + vy * dt
|
| if x < lo or x > hi:
|
| vx, x = -vx, float(np.clip(2 * np.clip(x, lo, hi) - x, lo, hi))
|
| if y < lo or y > hi:
|
| vy, y = -vy, float(np.clip(2 * np.clip(y, lo, hi) - y, lo, hi))
|
| return out
|
| if k == "zigzag":
|
| dx, dy = DIRS[ev["dir"]]
|
| dist = min(_uni(rng, 0.45, 0.6), 1 - 2 * m - 0.1)
|
| amp = _uni(rng, 0.08, 0.14)
|
| along0 = _uni(rng, m, 1 - m - dist) if (dx + dy) > 0 else _uni(rng, m + dist, 1 - m)
|
| across0 = _uni(rng, m + amp, 1 - m - amp)
|
| out = []
|
| for t in ts:
|
| along = along0 + (dx + dy) * dist * t
|
| zig = amp * (2 * abs((t * 3) % 1 - 0.5) * 2 - 1)
|
| x, y = (along, across0 + zig) if dx else (across0 + zig, along)
|
| out.append(_state(x, y, r, ang0, rgb))
|
| return out
|
| if k == "circle":
|
| R = _uni(rng, 0.18, max(0.19, 0.5 - m - 0.02))
|
| cx, cy = center if center is not None else (_uni(rng, m + R, 1 - m - R), _uni(rng, m + R, 1 - m - R))
|
| ph = _uni(rng, 0, 2 * math.pi)
|
| sgn = _choice(rng, [-1, 1])
|
| turns = _uni(rng, 0.8, 1.0)
|
| return [_state(cx + R * math.cos(ph + sgn * 2 * math.pi * turns * t), cy + R * math.sin(ph + sgn * 2 * math.pi * turns * t),
|
| r, ang0, rgb) for t in ts]
|
| if k in ("grow", "shrink"):
|
| mm = r * 1.7 + 0.03
|
| x0, y0 = _uni(rng, mm, 1 - mm, size=2)
|
| f = (lambda t: 0.45 + 1.25 * t) if k == "grow" else (lambda t: 1.7 - 1.25 * t)
|
| return [_state(x0, y0, r * f(t), ang0, rgb) for t in ts]
|
| x0, y0 = anywhere()
|
| if k == "spin":
|
| sp = _choice(rng, [-1, 1]) * _uni(rng, 0.7, 1.1) * math.pi
|
| return [_state(x0, y0, r, ang0 + sp * t, rgb) for t in ts]
|
| if k == "still":
|
| return [_state(x0, y0, r, ang0, rgb) for t in ts]
|
| if k == "fade_in":
|
| return [_state(x0, y0, r, ang0, rgb, min(1, max(0, (t - 0.1) / 0.6))) for t in ts]
|
| if k == "fade_out":
|
| return [_state(x0, y0, r, ang0, rgb, min(1, max(0, (0.9 - t) / 0.6))) for t in ts]
|
| if k == "blink":
|
| period = _choice(rng, [4, 4, 6])
|
| return [_state(x0, y0, r, ang0, rgb, 1.0 if (i // (period // 2)) % 2 == 0 else 0.0) for i in range(T)]
|
| if k == "recolor":
|
| rgb2 = np.array(COLORS[ev["to"]], float)
|
| out = []
|
| for t in ts:
|
| w = min(1, max(0, (t - 0.3) / 0.4))
|
| out.append(_state(x0, y0, r, ang0, (1 - w) * rgb + w * rgb2))
|
| return out
|
| if k == "teleport":
|
| x1, y1 = max((anywhere() for _ in range(8)), key=lambda p: math.hypot(p[0] - x0, p[1] - y0))
|
| cut = rng.integers(T // 3, 2 * T // 3)
|
| return [_state(x0, y0, r, ang0, rgb) if i < cut else _state(x1, y1, r, ang0, rgb) for i in range(T)]
|
| raise ValueError(k)
|
|
|
|
|
| def _pair_states(rng, a: dict, b: dict, ev: dict, T: int):
|
| ra = SIZES[a["size"]] * _uni(rng, 0.92, 1.08)
|
| rb = SIZES[b["size"]] * _uni(rng, 0.92, 1.08)
|
| ca, cb = np.array(COLORS[a["color"]], float), np.array(COLORS[b["color"]], float)
|
| angA, angB = _uni(rng, 0, 2 * math.pi, size=2)
|
| ts = np.linspace(0, 1, T)
|
| k = ev["kind"]
|
| if k in ("merge", "split"):
|
| res = ev["result"] if k == "merge" else ev["source"]
|
| rr = SIZES[res["size"]]
|
| cr = np.array(COLORS[res["color"]], float)
|
| mid = _uni(rng, 0.32, 0.68, size=2)
|
| th = _uni(rng, 0, math.pi)
|
| sep = _uni(rng, 0.62, 0.8)
|
| u = np.array([math.cos(th), math.sin(th)])
|
| pa = np.clip(mid - u * sep / 2, ra + 0.02, 1 - ra - 0.02)
|
| pb = np.clip(mid + u * sep / 2, rb + 0.02, 1 - rb - 0.02)
|
| SA, SB, SR = [], [], []
|
| for t in ts:
|
| tm = t if k == "merge" else 1 - t
|
| p = min(1, tm / 0.6)
|
| f = min(1, max(0, (tm - 0.5) / 0.25))
|
| xa = pa + (mid - pa) * p
|
| xb = pb + (mid - pb) * p
|
| SA.append(_state(*xa, ra * (1 - f) + rr * f, angA, ca * (1 - f) + cr * f, 1 - f))
|
| SB.append(_state(*xb, rb * (1 - f) + rr * f, angB, cb * (1 - f) + cr * f, 1 - f))
|
| SR.append(_state(*mid, rr * (0.6 + 0.4 * f), angA, cr, f))
|
| return [SA, SB, SR]
|
| if k == "collide":
|
| y = _uni(rng, 0.3, 0.7)
|
| vert = rng.random() < 0.3
|
| hit = _uni(rng, 0.4, 0.6)
|
| xa0, xb0 = _uni(rng, 0.05, 0.2) + ra, 1 - _uni(rng, 0.05, 0.2) - rb
|
| contact = (xa0 * rb + xb0 * ra) / (ra + rb)
|
| SA, SB = [], []
|
| for t in ts:
|
| if t <= hit:
|
| s = t / hit
|
| xa = xa0 + (contact - ra - xa0) * s
|
| xb = xb0 + (contact + rb - xb0) * s
|
| else:
|
| s = (t - hit) / (1 - hit)
|
| xa = contact - ra - (contact - ra - xa0) * s * 0.8
|
| xb = contact + rb + (xb0 - contact - rb) * s * 0.8
|
| pa, pb = (y, xa) if vert else (xa, y), (y, xb) if vert else (xb, y)
|
| SA.append(_state(*pa, ra, angA, ca))
|
| SB.append(_state(*pb, rb, angB, cb))
|
| return [SA, SB]
|
| if k == "chase":
|
| lead = _single_states(rng, b, {"kind": _choice(rng, ["move", "circle", "zigzag"]), "dir": _choice(rng, ["left", "right", "up", "down"])}, T)
|
| lag = 4
|
| SA = []
|
| for i in range(T):
|
| j = i - lag
|
| if j >= 0:
|
| x, y = lead[j][0], lead[j][1]
|
| else:
|
| vx, vy = lead[1][0] - lead[0][0], lead[1][1] - lead[0][1]
|
| x, y = lead[0][0] + vx * j, lead[0][1] + vy * j
|
| SA.append(_state(np.clip(x, ra, 1 - ra), np.clip(y, ra, 1 - ra), ra, angA, ca))
|
| SB = [[s[0], s[1], rb, angB, tuple(cb), 1.0] for s in lead]
|
| return [SA, SB]
|
| if k == "orbit":
|
| R = _uni(rng, 0.24, 0.32)
|
| c = (_uni(rng, 0.5 - (0.5 - R - ra - 0.02), 0.5 + (0.5 - R - ra - 0.02)),
|
| _uni(rng, 0.5 - (0.5 - R - ra - 0.02), 0.5 + (0.5 - R - ra - 0.02)))
|
| SA = _single_states(rng, a, {"kind": "circle"}, T, center=c)
|
| SA = [[s[0], s[1], ra, angA, tuple(ca), 1.0] for s in SA]
|
| SB = [_state(*c, rb, angB, cb) for _ in ts]
|
| return [SA, SB]
|
| if k == "swap":
|
| p0 = _uni(rng, 0.25, 0.75, size=2)
|
| th = _uni(rng, 0, 2 * math.pi)
|
| R = _uni(rng, 0.2, 0.28)
|
| c = np.clip(p0, R + 0.12, 1 - R - 0.12)
|
| SA, SB = [], []
|
| for t in ts:
|
| ph = th + math.pi * t
|
| SA.append(_state(c[0] + R * math.cos(ph), c[1] + R * math.sin(ph), ra, angA, ca))
|
| SB.append(_state(c[0] - R * math.cos(ph), c[1] - R * math.sin(ph), rb, angB, cb))
|
| return [SA, SB]
|
| if k == "independent":
|
| return [_single_states(rng, a, ev["a"], T), _single_states(rng, b, ev["b"], T)]
|
| raise ValueError(k)
|
|
|
|
|
|
|
|
|
| def _random_obj(rng, bg: str, shape=None, color=None, size=None) -> dict:
|
| while True:
|
| o = {"shape": shape or _choice(rng, SHAPES), "color": color or _choice(rng, list(COLORS)),
|
| "size": size or _choice(rng, list(SIZES))}
|
| if _contrast_ok(COLORS[o["color"]], bg):
|
| return o
|
| color = None
|
|
|
|
|
| def _random_single_event(rng, obj: dict, bg: str) -> dict:
|
| while True:
|
| k = _weighted(rng, SINGLE_EVENTS)
|
| if k == "spin" and obj["shape"] in ROTATION_INVARIANT:
|
| continue
|
| ev = {"kind": k}
|
| if k in ("move", "zigzag"):
|
| ev["dir"] = _choice(rng, list(DIRS) if k == "move" else ["left", "right", "up", "down"])
|
| if k == "jump":
|
| ev["dir"] = _choice(rng, ["left", "right", None])
|
| if k == "recolor":
|
| to = _choice(rng, [c for c in COLORS if c != obj["color"] and _contrast_ok(COLORS[c], bg)])
|
| ev["to"] = to
|
| return ev
|
|
|
|
|
| def is_heldout(scene: dict) -> bool:
|
| kinds = [scene["event"]["kind"]]
|
| if kinds[0] == "independent":
|
| kinds = [scene["event"]["a"]["kind"], scene["event"]["b"]["kind"]]
|
| shapes = [o["shape"] for o in scene["objects"]]
|
| res = scene["event"].get("result") or scene["event"].get("source")
|
| if res:
|
| shapes.append(res["shape"])
|
| return any((k, sh) in HELDOUT for k in kinds for sh in shapes)
|
|
|
|
|
| def featured_scene(rng, bg: str = "black", mention_bg: bool = False) -> dict:
|
| """The user's example: a red circle and a blue circle merge into a purple circle."""
|
| small = ["tiny", "small", "small", "medium"]
|
| a = {"shape": "circle", "color": "red", "size": _choice(rng, small)}
|
| b = {"shape": "circle", "color": "blue", "size": _choice(rng, small)}
|
| if rng.random() < 0.5:
|
| a, b = b, a
|
| res = {"shape": "circle", "color": "purple", "size": _choice(rng, ["medium", "big"])}
|
| return {"objects": [a, b], "event": {"kind": "merge", "result": res}, "bg": bg, "mention_bg": mention_bg}
|
|
|
|
|
| def random_scene(rng, cfg: WorldConfig, split: str = "train") -> dict:
|
| """split: 'train' (never a HELDOUT combo), 'heldout' (always one), 'any'."""
|
| while True:
|
| scene = _random_scene(rng, cfg)
|
| if split == "any" or (split == "heldout") == is_heldout(scene):
|
| return scene
|
|
|
|
|
| def _random_scene(rng, cfg: WorldConfig) -> dict:
|
| if rng.random() < cfg.p_featured:
|
| mention_bg = rng.random() < cfg.p_mention_bg
|
| bg = _choice(rng, ["black", "white", "gray", "sky", "grass"]) if mention_bg else "black"
|
| return featured_scene(rng, bg, mention_bg)
|
| mention_bg = rng.random() < cfg.p_mention_bg
|
| bg = _choice(rng, list(BACKGROUNDS)) if (mention_bg or rng.random() > cfg.p_default_black) else "black"
|
| if rng.random() >= cfg.p_pair:
|
| o = _random_obj(rng, bg, size=_choice(rng, ["small", "medium", "medium", "big", "big", "huge", "tiny"]))
|
| return {"objects": [o], "event": _random_single_event(rng, o, bg), "bg": bg, "mention_bg": mention_bg}
|
| k = _weighted(rng, {**PAIR_EVENTS, "merge": cfg.merge_weight})
|
| small = ["tiny", "small", "small", "medium"]
|
| if k in ("merge", "split"):
|
| a = _random_obj(rng, bg, size=_choice(rng, small))
|
| partners = [next(iter(pair - {a["color"]})) for pair in MIX if a["color"] in pair]
|
| b_color = _choice(rng, partners) if partners and rng.random() < 0.8 else a["color"]
|
| if not _contrast_ok(COLORS[b_color], bg):
|
| return _random_scene(rng, cfg)
|
| b = _random_obj(rng, bg, color=b_color, size=_choice(rng, small))
|
| if rng.random() < 0.6:
|
| b["shape"] = a["shape"]
|
| mixed = mix_colors(a["color"], b["color"])
|
| if not _contrast_ok(COLORS[mixed], bg):
|
| return _random_scene(rng, cfg)
|
| res = {"shape": a["shape"] if a["shape"] == b["shape"] else "circle", "color": mixed,
|
| "size": _choice(rng, ["medium", "big"])}
|
| ev = {"kind": k, "result" if k == "merge" else "source": res}
|
| return {"objects": [a, b], "event": ev, "bg": bg, "mention_bg": mention_bg}
|
| a = _random_obj(rng, bg, size=_choice(rng, small))
|
| b = _random_obj(rng, bg, size=_choice(rng, small + (["big"] if k == "orbit" else [])))
|
| if a == b:
|
| return _random_scene(rng, cfg)
|
| ev = {"kind": k}
|
| if k == "independent":
|
| ev["a"], ev["b"] = _random_single_event(rng, a, bg), _random_single_event(rng, b, bg)
|
| return {"objects": [a, b], "event": ev, "bg": bg, "mention_bg": mention_bg}
|
|
|
|
|
| def caption_scene(rng, scene: dict) -> str:
|
| objs, ev = scene["objects"], scene["event"]
|
| cap = describe_single(rng, objs[0], ev) if len(objs) == 1 else describe_pair(rng, objs[0], objs[1], ev)
|
| bgp = _choice(rng, BG_PHRASES[scene["bg"]]) if scene["mention_bg"] else None
|
| return _style(rng, cap, bgp)
|
|
|
|
|
|
|
|
|
| def _poly_mask(px, py, verts):
|
| inside = np.zeros(px.shape, bool)
|
| n = len(verts)
|
| for i in range(n):
|
| x1, y1 = verts[i]
|
| x2, y2 = verts[(i + 1) % n]
|
| crosses = (y1 > py) != (y2 > py)
|
| xint = (x2 - x1) * (py - y1) / (y2 - y1 + 1e-12) + x1
|
| inside ^= crosses & (px < xint)
|
| return inside
|
|
|
|
|
| def _regular(n, r, ang, cx, cy, inner=None):
|
| k = 2 * n if inner else n
|
| return [(cx + (inner if inner and i % 2 else r) * math.cos(ang - math.pi / 2 + 2 * math.pi * i / k),
|
| cy + (inner if inner and i % 2 else r) * math.sin(ang - math.pi / 2 + 2 * math.pi * i / k)) for i in range(k)]
|
|
|
|
|
| def shape_mask(shape, px, py, cx, cy, r, ang):
|
| dx, dy = px - cx, py - cy
|
| if shape == "circle":
|
| return dx * dx + dy * dy <= r * r
|
| if shape == "ring":
|
| d2 = dx * dx + dy * dy
|
| return (d2 <= r * r) & (d2 >= (0.55 * r) ** 2)
|
| if shape == "square":
|
| return _poly_mask(px, py, _regular(4, r * 1.15, ang + math.pi / 4, cx, cy))
|
| if shape == "triangle":
|
| return _poly_mask(px, py, _regular(3, r * 1.2, ang, cx, cy))
|
| if shape == "star":
|
| return _poly_mask(px, py, _regular(5, r * 1.25, ang, cx, cy, inner=r * 0.5))
|
| if shape == "hexagon":
|
| return _poly_mask(px, py, _regular(6, r * 1.05, ang, cx, cy))
|
| c, s = math.cos(ang), math.sin(ang)
|
| u, v = (c * dx + s * dy) / r, (-s * dx + c * dy) / r
|
| if shape == "diamond":
|
| return np.abs(u) / 0.75 + np.abs(v) / 1.15 <= 1
|
| if shape == "cross":
|
| arm = 0.33
|
| return ((np.abs(u) <= 1) & (np.abs(v) <= arm)) | ((np.abs(v) <= 1) & (np.abs(u) <= arm))
|
| if shape == "heart":
|
| c, s = math.cos(0.3 * math.sin(ang)), math.sin(0.3 * math.sin(ang))
|
| u, v = (c * dx + s * dy) / r, (-s * dx + c * dy) / r
|
| x, y = u * 1.15, -v * 1.15 + 0.15
|
| return (x * x + y * y - 1) ** 3 - x * x * y ** 3 <= 0
|
| raise ValueError(shape)
|
|
|
|
|
| def _background(bg, S):
|
| top, bot = (np.array(c, np.float32) for c in BACKGROUNDS[bg])
|
| w = np.linspace(0, 1, S, dtype=np.float32)[:, None, None]
|
| return np.broadcast_to(top * (1 - w) + bot * w, (S, S, 3)).copy()
|
|
|
|
|
| def render_scene(scene: dict, rng, cfg: WorldConfig) -> np.ndarray:
|
| T, S, ss = cfg.frames, cfg.size, cfg.supersample
|
| big = S * ss
|
| co = (np.arange(big, dtype=np.float32) + 0.5) / big
|
| px, py = np.meshgrid(co, co)
|
| objs, ev = scene["objects"], scene["event"]
|
| if len(objs) == 1:
|
| tracks = [_single_states(rng, objs[0], ev, T)]
|
| shapes = [objs[0]["shape"]]
|
| else:
|
| tracks = _pair_states(rng, objs[0], objs[1], ev, T)
|
| shapes = [objs[0]["shape"], objs[1]["shape"]]
|
| if ev["kind"] in ("merge", "split"):
|
| shapes.append((ev.get("result") or ev.get("source"))["shape"])
|
| scene["tracks"] = tracks
|
| scene["track_shapes"] = shapes
|
| first = [i for i, tr in enumerate(tracks) if tr[0][5] > 0.5 and tr[0][2] > 0.005]
|
| last = [i for i, tr in enumerate(tracks) if tr[T - 1][5] > 0.5 and tr[T - 1][2] > 0.005]
|
| scene["summary"] = {"initial_count": len(first), "final_count": len(last),
|
| "initial_colors": [nearest_color(tracks[i][0][4]) for i in first],
|
| "final_colors": [nearest_color(tracks[i][T - 1][4]) for i in last],
|
| "final_shapes": [shapes[i] for i in last]}
|
| base = _background(scene["bg"], big)
|
| frames = np.empty((T, S, S, 3), np.uint8)
|
| for f in range(T):
|
| img = base.copy()
|
| for shape, tr in zip(shapes, tracks):
|
| x, y, r, ang, rgb, alpha = tr[f]
|
| if alpha <= 0.01 or r <= 0.005:
|
| continue
|
| m = shape_mask(shape, px, py, x, y, r, ang)
|
| img[m] = img[m] * (1 - alpha) + np.array(rgb, np.float32) * alpha
|
| img = img.reshape(S, ss, S, ss, 3).mean(axis=(1, 3))
|
| frames[f] = np.clip(img + 0.5, 0, 255).astype(np.uint8)
|
| return frames
|
|
|
|
|
| def make_clip(rng, cfg: WorldConfig | None = None, split: str = "train"):
|
| cfg = cfg or WorldConfig()
|
| scene = random_scene(rng, cfg, split)
|
| caption = caption_scene(rng, scene)
|
| video = render_scene(scene, rng, cfg)
|
| scene["caption"] = caption
|
| return video, caption, scene
|
|
|
|
|
| def sample_prompts(n: int, split: str = "train", seed: int = 0, cfg: WorldConfig | None = None) -> list[tuple[str, dict]]:
|
| """Captions + scenes. Call render_scene(scene, rng, cfg) to get ground-truth tracks/summary for a scene."""
|
| rng = np.random.default_rng(seed)
|
| cfg = cfg or WorldConfig()
|
| out = []
|
| for _ in range(n):
|
| scene = random_scene(rng, cfg, split)
|
| out.append((caption_scene(rng, scene), scene))
|
| return out
|
|
|
|
|
| def to_tensor(video: np.ndarray) -> torch.Tensor:
|
| return torch.from_numpy(video).permute(3, 0, 1, 2).float().div(127.5).sub(1)
|
|
|
|
|
| class WorldDataset(torch.utils.data.IterableDataset):
|
| """Infinite stream of (video float [3,T,H,W] in [-1,1], caption)."""
|
|
|
| def __init__(self, cfg: WorldConfig | None = None, seed: int = 0, split: str = "train"):
|
| self.cfg = cfg or WorldConfig()
|
| self.seed = seed
|
| self.split = split
|
|
|
| def __iter__(self):
|
| info = torch.utils.data.get_worker_info()
|
| rng = np.random.default_rng([self.seed, info.id if info else 0])
|
| while True:
|
| video, caption, _ = make_clip(rng, self.cfg, self.split)
|
| yield to_tensor(video), caption
|
|
|
|
|
| def collate(batch):
|
| videos, captions = zip(*batch)
|
| return torch.stack(videos), list(captions)
|
|
|
