worldbench 0.2.0__py3-none-any.whl

worldbench/__init__.py

@@ -0,0 +1,16 @@
+"""WorldBench: robotics world-model evaluation made practical."""
+
+from worldbench.core import Metrics, WorldBench, WorldModelRun, evaluate, load_dataset
+from worldbench.schemas import EvaluationResult, MetricResult
+
+__all__ = [
+    "EvaluationResult",
+    "MetricResult",
+    "Metrics",
+    "WorldBench",
+    "WorldModelRun",
+    "evaluate",
+    "load_dataset",
+]
+
+__version__ = "0.2.0"

worldbench/backends/__init__.py

@@ -0,0 +1,15 @@
+"""WorldBench backends."""
+
+from worldbench.backends.benchmark import BenchmarkBackend
+from worldbench.backends.demo import DemoBackend
+from worldbench.backends.lerobot import create_lerobot_style_demo_source, import_lerobot_repo, import_lerobot_style
+from worldbench.backends.local import LocalBackend
+
+__all__ = [
+    "BenchmarkBackend",
+    "DemoBackend",
+    "LocalBackend",
+    "create_lerobot_style_demo_source",
+    "import_lerobot_repo",
+    "import_lerobot_style",
+]

worldbench/backends/benchmark.py

@@ -0,0 +1,200 @@
+"""Synthetic benchmark scenario generator."""
+
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+from PIL import Image, ImageDraw
+
+from worldbench.utils import ensure_dir, write_json
+
+
+SCENARIO_NAMES = [
+    "push_cube",
+    "action_mismatch",
+    "pre_contact_motion",
+    "object_disappears",
+    "temporal_flicker",
+]
+
+
+@dataclass(frozen=True)
+class ScenarioConfig:
+    name: str
+    failure_mode: str
+
+
+class BenchmarkBackend:
+    """Generate lightweight synthetic benchmark scenarios."""
+
+    width = 128
+    height = 96
+
+    scenarios = [
+        ScenarioConfig("push_cube", "combined control/contact failures"),
+        ScenarioConfig("action_mismatch", "robot moves opposite the commanded action"),
+        ScenarioConfig("pre_contact_motion", "object moves before robot/object contact"),
+        ScenarioConfig("object_disappears", "object disappears during prediction"),
+        ScenarioConfig("temporal_flicker", "prediction flickers or jumps between frames"),
+    ]
+
+    def create(self, output_path: str | Path = "benchmarks", overwrite: bool = True) -> Path:
+        root = Path(output_path)
+        if root.exists() and overwrite:
+            shutil.rmtree(root)
+        ensure_dir(root)
+
+        for scenario in self.scenarios:
+            self._write_scenario(root / scenario.name, scenario)
+        return root
+
+    def _write_scenario(self, root: Path, scenario: ScenarioConfig) -> None:
+        ensure_dir(root)
+        states = _ground_truth_states()
+        actions = _actions()
+        self._write_episode(root, scenario, states, actions)
+        self._write_model(root, "good_model", scenario, states, quality="good")
+        self._write_model(root, "bad_model", scenario, states, quality="bad")
+
+    def _write_episode(
+        self,
+        root: Path,
+        scenario: ScenarioConfig,
+        states: list[dict[str, tuple[int, int]]],
+        actions: list[dict[str, object]],
+    ) -> None:
+        episode = root / "episode_001"
+        frames = ensure_dir(episode / "frames")
+        predictions = ensure_dir(episode / "predictions")
+        for idx, state in enumerate(states):
+            image = self._render_frame(state["robot"], state["object"], f"gt {idx:03d}")
+            image.save(frames / f"{idx:03d}.png")
+            image.save(predictions / f"{idx:03d}.png")
+
+        write_json(episode / "actions.json", actions)
+        write_json(
+            episode / "states.json",
+            [
+                {
+                    "t": idx,
+                    "robot_x": state["robot"][0],
+                    "robot_y": state["robot"][1],
+                    "object_x": state["object"][0],
+                    "object_y": state["object"][1],
+                }
+                for idx, state in enumerate(states)
+            ],
+        )
+        write_json(
+            episode / "metadata.json",
+            {
+                "name": scenario.name,
+                "robot": "synthetic_2d_mobile_manipulator",
+                "task": "push cube",
+                "fps": 5,
+                "description": f"Synthetic benchmark scenario: {scenario.failure_mode}.",
+            },
+        )
+
+    def _write_model(
+        self,
+        root: Path,
+        model_name: str,
+        scenario: ScenarioConfig,
+        states: list[dict[str, tuple[int, int]]],
+        quality: str,
+    ) -> None:
+        model_dir = ensure_dir(root / model_name / "episode_001")
+        generated = states if quality == "good" else _bad_states_for_scenario(scenario.name, states)
+        for idx, state in enumerate(generated):
+            hide_object = quality == "bad" and scenario.name == "object_disappears" and idx in {2, 3, 4, 5, 6, 7, 8, 9}
+            flicker = quality == "bad" and scenario.name == "temporal_flicker" and idx in {4, 6, 8}
+            if quality == "bad" and scenario.name == "push_cube":
+                hide_object = idx == 4
+                flicker = idx == 6
+            label = f"{quality[:4]} {idx:03d}"
+            self._render_frame(
+                state["robot"],
+                state["object"],
+                label,
+                hide_object=hide_object,
+                flicker=flicker,
+            ).save(model_dir / f"{idx:03d}.png")
+
+    def _render_frame(
+        self,
+        robot: tuple[int, int],
+        obj: tuple[int, int],
+        label: str,
+        hide_object: bool = False,
+        flicker: bool = False,
+    ) -> Image.Image:
+        background = (10, 18, 28) if not flicker else (44, 17, 20)
+        image = Image.new("RGB", (self.width, self.height), background)
+        draw = ImageDraw.Draw(image)
+        for x in range(0, self.width, 16):
+            draw.line((x, 0, x, self.height), fill=(25, 38, 50), width=1)
+        for y in range(0, self.height, 16):
+            draw.line((0, y, self.width, y), fill=(25, 38, 50), width=1)
+        draw.rectangle((5, 5, self.width - 6, self.height - 6), outline=(54, 78, 96), width=1)
+        draw.text((8, 7), label, fill=(128, 152, 170))
+
+        rx, ry = robot
+        ox, oy = obj
+        draw.line((rx, ry, ox, oy), fill=(52, 72, 88), width=2)
+        draw.rounded_rectangle((rx - 10, ry - 8, rx + 10, ry + 8), radius=4, fill=(80, 172, 214), outline=(190, 236, 255), width=1)
+        draw.ellipse((rx - 5, ry - 5, rx + 5, ry + 5), fill=(231, 76, 60), outline=(255, 180, 160), width=1)
+        draw.rectangle((rx - 4, ry + 8, rx + 4, ry + 12), fill=(7, 14, 22))
+        if not hide_object:
+            draw.rectangle((ox - 8, oy - 8, ox + 8, oy + 8), fill=(46, 204, 113), outline=(179, 255, 204), width=2)
+        return image
+
+
+def _ground_truth_states() -> list[dict[str, tuple[int, int]]]:
+    robot = (24, 50)
+    obj = (84, 50)
+    states = [{"robot": robot, "object": obj}]
+    for _ in range(9):
+        contact = ((robot[0] - obj[0]) ** 2 + (robot[1] - obj[1]) ** 2) ** 0.5 <= 18.0
+        robot = (robot[0] + 8, robot[1])
+        if contact:
+            obj = (obj[0] + 8, obj[1])
+        states.append({"robot": robot, "object": obj})
+    return states
+
+
+def _actions() -> list[dict[str, object]]:
+    return [
+        {"t": idx, "action": "move_right", "dx": 1.0, "dy": 0.0, "gripper": "open"}
+        for idx in range(9)
+    ]
+
+
+def _bad_states_for_scenario(name: str, states: list[dict[str, tuple[int, int]]]) -> list[dict[str, tuple[int, int]]]:
+    bad: list[dict[str, tuple[int, int]]] = []
+    start_robot = states[0]["robot"]
+    start_object = states[0]["object"]
+    for idx, state in enumerate(states):
+        robot = state["robot"]
+        obj = state["object"]
+        if name == "action_mismatch":
+            robot = (max(14, start_robot[0] - idx * 5), start_robot[1])
+        elif name == "pre_contact_motion":
+            if idx >= 1:
+                obj = (start_object[0] + idx * 7, start_object[1])
+        elif name == "object_disappears":
+            robot = state["robot"]
+        elif name == "temporal_flicker":
+            if idx in {4, 6, 8}:
+                robot = (112, 18)
+                obj = (34 + idx * 3, 76)
+        elif name == "push_cube":
+            robot = (max(14, start_robot[0] - idx * 5), start_robot[1])
+            if idx >= 2:
+                obj = (start_object[0] + idx * 4, start_object[1])
+            if idx == 6:
+                robot = (110, 18)
+        bad.append({"robot": robot, "object": obj})
+    return bad

worldbench/backends/demo.py

@@ -0,0 +1,233 @@
+"""Synthetic demo dataset generator."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+from PIL import Image, ImageDraw
+
+from worldbench.utils import ensure_dir, write_json
+
+
+class DemoBackend:
+    """Generate deterministic synthetic rollouts and model outputs."""
+
+    width = 128
+    height = 96
+
+    def create(self, output_path: str | Path = "examples/demo_dataset", overwrite: bool = True) -> Path:
+        root = Path(output_path)
+        if root.exists() and overwrite:
+            shutil.rmtree(root)
+        ensure_dir(root)
+
+        episodes = [
+            _episode_plan(
+                name="episode_001",
+                task="push cube right",
+                robot_start=(24, 50),
+                object_start=(84, 50),
+                delta=(8, 0),
+                actions=["move_right"] * 7 + ["close_gripper", "move_right"],
+            ),
+            _episode_plan(
+                name="episode_002",
+                task="push cube down",
+                robot_start=(58, 18),
+                object_start=(58, 68),
+                delta=(0, 7),
+                actions=["move_down"] * 7 + ["close_gripper", "move_down"],
+            ),
+        ]
+
+        for plan in episodes:
+            self._write_episode(root, plan)
+
+        self._write_model_outputs(root, episodes, "good_model", quality="good")
+        self._write_model_outputs(root, episodes, "bad_model", quality="bad")
+        self._mirror_example_outputs(root)
+        return root
+
+    def init_structure(self, output_path: str | Path) -> Path:
+        root = Path(output_path)
+        episode = root / "episode_001"
+        ensure_dir(episode / "frames")
+        ensure_dir(episode / "predictions")
+        write_json(
+            episode / "actions.json",
+            [
+                {"t": 0, "action": "move_right", "dx": 1.0, "dy": 0.0, "gripper": "open"},
+                {"t": 1, "action": "close_gripper", "dx": 0.0, "dy": 0.0, "gripper": "closed"},
+            ],
+        )
+        write_json(
+            episode / "states.json",
+            [
+                {"t": 0, "robot_x": 20, "robot_y": 50, "object_x": 80, "object_y": 50},
+                {"t": 1, "robot_x": 30, "robot_y": 50, "object_x": 80, "object_y": 50},
+            ],
+        )
+        write_json(
+            episode / "metadata.json",
+            {
+                "name": "push_cube_template",
+                "robot": "synthetic_2d_arm",
+                "task": "push cube",
+                "fps": 5,
+                "description": "Template WorldBench rollout. Add frames and predictions as numbered PNG files.",
+            },
+        )
+        return root
+
+    def _write_episode(self, root: Path, plan: dict[str, object]) -> None:
+        episode_dir = root / str(plan["name"])
+        frames_dir = episode_dir / "frames"
+        predictions_dir = episode_dir / "predictions"
+        ensure_dir(frames_dir)
+        ensure_dir(predictions_dir)
+        states = plan["states"]
+        actions = plan["actions_json"]
+
+        for idx, state in enumerate(states):
+            image = self._render_frame(state["robot"], state["object"], label=f"gt {idx:03d}")
+            image.save(frames_dir / f"{idx:03d}.png")
+            image.save(predictions_dir / f"{idx:03d}.png")
+
+        write_json(episode_dir / "actions.json", actions)
+        write_json(episode_dir / "states.json", plan["states_json"])
+        write_json(
+            episode_dir / "metadata.json",
+            {
+                "name": str(plan["name"]),
+                "robot": "synthetic_2d_arm",
+                "task": str(plan["task"]),
+                "fps": 5,
+                "description": "Synthetic robot rollout for world-model evaluation.",
+            },
+        )
+
+    def _write_model_outputs(self, root: Path, episodes: list[dict[str, object]], model_name: str, quality: str) -> None:
+        model_root = root / model_name
+        ensure_dir(model_root)
+        for plan in episodes:
+            episode_dir = model_root / str(plan["name"])
+            ensure_dir(episode_dir)
+            states = plan["states"]
+            bad_states = _bad_states(plan) if quality == "bad" else states
+            for idx, state in enumerate(bad_states):
+                hide_object = quality == "bad" and idx in {4}
+                flicker = quality == "bad" and idx in {6}
+                label = f"gt {idx:03d}" if quality == "good" else f"bad {idx:03d}"
+                image = self._render_frame(
+                    state["robot"],
+                    state["object"],
+                    label=label,
+                    hide_object=hide_object,
+                    flicker=flicker,
+                )
+                image.save(episode_dir / f"{idx:03d}.png")
+
+    def _mirror_example_outputs(self, root: Path) -> None:
+        if root.name != "demo_dataset" or root.parent.name != "examples":
+            return
+        for model_name, folder_name in [("good_model", "good_model_outputs"), ("bad_model", "bad_model_outputs")]:
+            target = root.parent / folder_name
+            if target.exists():
+                shutil.rmtree(target)
+            shutil.copytree(root / model_name, target)
+
+    def _render_frame(
+        self,
+        robot: tuple[int, int],
+        obj: tuple[int, int],
+        label: str,
+        hide_object: bool = False,
+        flicker: bool = False,
+    ) -> Image.Image:
+        background = (14, 21, 28) if not flicker else (48, 20, 20)
+        image = Image.new("RGB", (self.width, self.height), background)
+        draw = ImageDraw.Draw(image)
+
+        for x in range(0, self.width, 16):
+            draw.line((x, 0, x, self.height), fill=(24, 34, 43), width=1)
+        for y in range(0, self.height, 16):
+            draw.line((0, y, self.width, y), fill=(24, 34, 43), width=1)
+
+        draw.rectangle((5, 5, self.width - 6, self.height - 6), outline=(58, 75, 88), width=1)
+        draw.text((8, 7), label, fill=(129, 152, 166))
+
+        rx, ry = robot
+        ox, oy = obj
+        draw.line((rx, ry, ox, oy), fill=(51, 68, 80), width=2)
+        draw.ellipse((rx - 8, ry - 8, rx + 8, ry + 8), fill=(231, 76, 60), outline=(255, 180, 160), width=2)
+        draw.rectangle((rx - 4, ry - 14, rx + 4, ry - 8), fill=(255, 202, 120))
+
+        if not hide_object:
+            draw.rectangle((ox - 8, oy - 8, ox + 8, oy + 8), fill=(46, 204, 113), outline=(179, 255, 204), width=2)
+            draw.line((ox - 8, oy, ox + 8, oy), fill=(22, 120, 75), width=1)
+            draw.line((ox, oy - 8, ox, oy + 8), fill=(22, 120, 75), width=1)
+
+        return image
+
+
+def _episode_plan(
+    name: str,
+    task: str,
+    robot_start: tuple[int, int],
+    object_start: tuple[int, int],
+    delta: tuple[int, int],
+    actions: list[str],
+) -> dict[str, object]:
+    states = [{"robot": robot_start, "object": object_start}]
+    states_json = [{"t": 0, "robot_x": robot_start[0], "robot_y": robot_start[1], "object_x": object_start[0], "object_y": object_start[1]}]
+    actions_json = []
+    robot = robot_start
+    obj = object_start
+    for idx, action in enumerate(actions):
+        contact_before_action = ((robot[0] - obj[0]) ** 2 + (robot[1] - obj[1]) ** 2) ** 0.5 <= 18.0
+        if action.startswith("move"):
+            robot = (robot[0] + delta[0], robot[1] + delta[1])
+            if contact_before_action:
+                obj = (obj[0] + delta[0], obj[1] + delta[1])
+        states.append({"robot": robot, "object": obj})
+        states_json.append(
+            {"t": idx + 1, "robot_x": robot[0], "robot_y": robot[1], "object_x": obj[0], "object_y": obj[1]}
+        )
+
+    for idx, action in enumerate(actions):
+        dx = 1.0 if "right" in action else 0.0
+        dy = 1.0 if "down" in action else 0.0
+        actions_json.append(
+            {
+                "t": idx,
+                "action": action,
+                "dx": dx,
+                "dy": dy,
+                "gripper": "closed" if action == "close_gripper" else "open",
+            }
+        )
+    return {"name": name, "task": task, "states": states, "states_json": states_json, "actions_json": actions_json}
+
+
+def _bad_states(plan: dict[str, object]) -> list[dict[str, tuple[int, int]]]:
+    states = plan["states"]
+    bad = []
+    first = states[0]
+    start_robot = first["robot"]
+    for idx, state in enumerate(states):
+        obj = state["object"]
+        if "right" in str(plan["task"]):
+            robot = (max(14, start_robot[0] - idx * 5), start_robot[1])
+            if idx >= 2:
+                obj = (obj[0] + idx * 4, obj[1])
+            if idx == 6:
+                robot = (110, 18)
+        else:
+            robot = (start_robot[0], max(14, start_robot[1] - idx * 4))
+            if idx >= 2:
+                obj = (obj[0], obj[1] + idx * 3)
+            if idx == 6:
+                robot = (104, 20)
+        bad.append({"robot": robot, "object": obj})
+    return bad

worldbench/backends/frame_freeze.py

@@ -0,0 +1,117 @@
+"""Deterministic frame-freeze corruption helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import random
+import shutil
+from pathlib import Path
+
+from worldbench.utils import ensure_dir, list_image_files
+
+
+@dataclass(frozen=True)
+class FrameFreezeResult:
+    """Summary of a frozen prediction sequence."""
+
+    source_frames: int
+    frozen_frames: int
+    output_paths: list[Path]
+
+
+def freeze_frames(
+    source_frames: str | Path,
+    output_frames: str | Path,
+    severity: float = 0.0,
+    seed: int = 42,
+    overwrite: bool = True,
+) -> FrameFreezeResult:
+    """Copy a frame sequence while freezing a deterministic subset of timesteps.
+
+    A frozen frame is copied from the most recent previous output frame, so the
+    visible timeline can repeat short spans without modifying the source data.
+    """
+
+    if not 0.0 <= severity <= 1.0:
+        raise ValueError(f"Severity must be between 0 and 1: {severity}")
+
+    source = Path(source_frames)
+    if not source.is_dir():
+        raise FileNotFoundError(f"Source frame directory does not exist: {source}")
+
+    frames = list_image_files(source)
+    if not frames:
+        raise ValueError(f"No image frames found in {source}")
+
+    destination = Path(output_frames)
+    if destination.exists() and overwrite:
+        shutil.rmtree(destination)
+    elif destination.exists():
+        raise FileExistsError(f"Output frame directory already exists: {destination}")
+    ensure_dir(destination)
+
+    frozen_positions = _frozen_positions(len(frames), severity, seed)
+    frozen_frames = 0
+    output_paths: list[Path] = []
+    previous_output: Path | None = None
+
+    for idx, source_path in enumerate(frames):
+        destination_path = destination / source_path.name
+        if idx == 0 or idx not in frozen_positions or previous_output is None:
+            shutil.copy2(source_path, destination_path)
+        else:
+            shutil.copy2(previous_output, destination_path)
+            frozen_frames += 1
+        previous_output = destination_path
+        output_paths.append(destination_path)
+
+    return FrameFreezeResult(
+        source_frames=len(frames),
+        frozen_frames=frozen_frames,
+        output_paths=output_paths,
+    )
+
+
+def freeze_rollout_predictions(
+    source_dataset: str | Path,
+    output_root: str | Path,
+    severity: float = 0.0,
+    seed: int = 42,
+    overwrite: bool = True,
+) -> list[FrameFreezeResult]:
+    """Create a WorldBench-style prediction tree from a rollout dataset root."""
+
+    source_root = Path(source_dataset)
+    if not source_root.is_dir():
+        raise FileNotFoundError(f"Source dataset directory does not exist: {source_root}")
+
+    destination_root = Path(output_root)
+    if destination_root.exists() and overwrite:
+        shutil.rmtree(destination_root)
+    elif destination_root.exists():
+        raise FileExistsError(f"Output root already exists: {destination_root}")
+    ensure_dir(destination_root)
+
+    results: list[FrameFreezeResult] = []
+    for episode_dir in sorted(
+        path for path in source_root.iterdir() if path.is_dir() and path.name.startswith("episode_")
+    ):
+        frames_dir = episode_dir / "frames"
+        predictions_dir = destination_root / episode_dir.name / "predictions"
+        result = freeze_frames(frames_dir, predictions_dir, severity=severity, seed=seed, overwrite=overwrite)
+        results.append(result)
+    return results
+
+
+def _frozen_positions(frame_count: int, severity: float, seed: int) -> set[int]:
+    if frame_count <= 1 or severity <= 0.0:
+        return set()
+
+    freeze_count = round((frame_count - 1) * severity)
+    freeze_count = max(0, min(frame_count - 1, freeze_count))
+    if freeze_count == 0:
+        return set()
+
+    positions = list(range(1, frame_count))
+    random.Random(seed).shuffle(positions)
+    return set(sorted(positions[:freeze_count]))