""" Orbis 2 World Model — interactive rollout demo for Hugging Face Spaces. This is the Orbis 2 counterpart to app.py (which drives the single-level Orbis 1 model). Orbis 2 is a hierarchical L1-L2 world model whose rollout script takes an mp4 directly, so the flow differs from app.py: 1. User uploads a short driving clip. 2. We re-encode it to a constant CONTEXT_FPS (= the L1 frame rate) so the rollout logic's exact integer-multiple frame-rate check passes. 3. We call into orbis2_app_engine.RolloutEngine (in-process, not a subprocess — see that module's docstring), which samples the L1 (high-rate) and L2 (low-rate, further back) context windows from the tail of that video. 4. The context is tiled NUM_VIDEOS times so one minibatch rolls out that many independently sampled futures under a single seed. 5. Generated frames (fake_images/sequence_XXXX/*.jpg) are encoded back into mp4s. Checkpoint / config resolution (see resolve_exp_dir / download_checkpoints): * If ORBIS2_EXP_DIR is set (and contains the config + checkpoint), it is used as-is — this is the path that works today against a local training run. * Otherwise, if ORBIS2_HF_REPO is set, the config + checkpoint are pulled from that Hub repo into a local exp dir. Orbis 2 is not on the Hub yet; set these env vars once the checkpoint is uploaded. """ import argparse import hashlib import json import os import random import re import subprocess import tempfile import traceback import uuid from pathlib import Path import cv2 import gradio as gr import numpy as np import spaces from huggingface_hub import snapshot_download from orbis2_app_engine import get_engine import torch print(f"torch version: {torch.__version__}, cuda available: {torch.cuda.is_available()}, cuda version: {torch.version.cuda}") # ---------------------------------------------------------------------------- # Workaround for gradio 5.9.1 / gradio_client bug (gradio-app/gradio#11722): # get_api_info() walks the app's JSON schema and crashes with # "TypeError: argument of type 'bool' is not iterable" # when a schema node is a boolean (e.g. "additionalProperties": true), because # get_type()/_json_schema_to_python_type() assume every schema is a dict. The # main page route calls api_info() on every load, so without this the endpoint # 500s continuously. Make the walker tolerate boolean schemas. # ---------------------------------------------------------------------------- import gradio_client.utils as _gc_utils _orig_get_type = _gc_utils.get_type _orig_json_to_py = _gc_utils._json_schema_to_python_type def _safe_get_type(schema): if isinstance(schema, bool): return "bool" return _orig_get_type(schema) def _safe_json_to_py(schema, defs=None): if isinstance(schema, bool): return "Any" return _orig_json_to_py(schema, defs) _gc_utils.get_type = _safe_get_type _gc_utils._json_schema_to_python_type = _safe_json_to_py # ---------------------------------------------------------------------------- # CLI args # ---------------------------------------------------------------------------- def _parse_args(): parser = argparse.ArgumentParser(description="Orbis 2 hierarchical world model demo") parser.add_argument( "--ui-only", action="store_true", help="Launch only the Gradio UI, without downloading or loading any " "model/checkpoint. For UI/UX testing on a machine without a GPU: " "'Generate rollouts' will show a placeholder error instead of " "actually running inference.", ) return parser.parse_args() ARGS = _parse_args() UI_ONLY = ARGS.ui_only # ---------------------------------------------------------------------------- # Configuration — adjust these to your setup # ---------------------------------------------------------------------------- # Persistent-storage location. Prefer the Space's persistent /data volume so large # downloads/caches survive restarts; fall back to the equivalent path under the home # dir if /data isn't writable (e.g. persistent storage not enabled). def _pick_cache_dir(subpath: str = ".cache/huggingface") -> str: for base in ("/data", os.path.expanduser("~")): candidate = os.path.join(base, subpath) try: probe = Path(candidate) probe.mkdir(parents=True, exist_ok=True) test = probe / ".write_probe" test.touch() test.unlink() return candidate except OSError: continue return os.path.join(os.path.expanduser("~"), subpath) HF_CACHE_DIR = os.path.abspath(os.environ.get("HF_HOME") or _pick_cache_dir()) os.environ["HF_HOME"] = HF_CACHE_DIR HF_HUB_CACHE_DIR = str(Path(HF_CACHE_DIR) / "hub") print(f"[startup] HF cache dir: {HF_CACHE_DIR}" f"{'' if HF_CACHE_DIR.startswith('/data') else ' (ephemeral — /data not writable)'}") # torch.compile artifacts (Inductor/Triton-compiled kernels) are tied to the exact # torch/CUDA/cuDNN build and GPU model they were compiled on -- loading a cache built # on a different combination can silently miscompile or just fail to load. So instead # of one fixed filename, compile_cache_path() below fingerprints the current # environment into the filename, and this dir (persistent across restarts, same as # HF_CACHE_DIR above) is where those per-environment files accumulate. COMPILE_CACHE_DIR = Path(os.path.abspath(_pick_cache_dir(".cache/orbis2_compile_cache"))) print(f"[startup] Compile cache dir: {COMPILE_CACHE_DIR}" f"{'' if str(COMPILE_CACHE_DIR).startswith('/data') else ' (ephemeral — /data not writable)'}") def compile_cache_path() -> Path: """Path to this exact environment's torch.compile cache file. Must be called with CUDA already attached to the process (i.e. from inside a @spaces.GPU-decorated call, not at module import time) -- HF Spaces ZeroGPU only attaches a GPU to that call, not to the main process, so torch.cuda.* would report nothing useful before then.""" parts = [f"torch{torch.__version__}", f"cu{torch.version.cuda or 'none'}"] if torch.cuda.is_available(): parts.append(f"cudnn{torch.backends.cudnn.version()}") parts.append(torch.cuda.get_device_name(0)) parts.append("sm%d%d" % torch.cuda.get_device_capability(0)) fingerprint = re.sub(r"[^A-Za-z0-9]+", "-", "_".join(parts)).strip("-") return COMPILE_CACHE_DIR / f"compile_cache_{fingerprint}.pkl" # --- Orbis 2 models (configs + checkpoints) ---------------------------------- # Three models are needed, laid out as /{L1,L2,tok}/{config,checkpoints/last.ckpt}: # L1 — detail predictor (the world model rollout_demo_v2.py loads directly) # L2 — abstract predictor (frozen; referenced by L1's condition_preprocessor) # tok — tokenizer (referenced by BOTH L1's and L2's tokenizer_config) # # L1's and L2's configs locate their dependencies via $ORBIS2_MODELS_DIR, which we # export below so OmegaConf's expandvars resolves them wherever the tree lands. CONFIG_NAME = "config_distill.yaml" # L1's config, relative to EXP_DIR CKPT_NAME = "checkpoints/last.ckpt" # every model stores its weights here #: (sub-dir, config filename) for each of the required models. MODEL_SPECS = [ ("L1", CONFIG_NAME), ("L2", "config.yaml"), # L1's config_distill.yaml conditions on this frozen, distilled L2 predictor # (a separate Hub folder from "L2") -- must be downloaded too, or L1 fails to # build its condition_preprocessor with a FileNotFoundError on config.yaml. ("L2_distilled", "config.yaml"), ("tok", "config.yaml"), ] # Hub repo holding the three models (root contains L1/, L2/, tok/). HF_CKPT_REPO = os.environ.get("ORBIS2_HF_REPO", "sud0301/orbis2") def _missing_models(root: Path) -> list[str]: """Return the models whose config or checkpoint is absent under `root`.""" return [sub for sub, cfg in MODEL_SPECS if not (root / sub / cfg).exists() or not (root / sub / CKPT_NAME).exists()] def resolve_models_dir() -> Path: """Return a directory containing L1/, L2/ and tok/, downloading it if needed. A local tree wins if it's complete (a dev checkout, or ORBIS2_MODELS_DIR pointing at one). Otherwise we snapshot the Hub repo and use the *cache* path directly rather than copying into the Space — `local_dir=` would duplicate ~20 GB on disk. The cache lives on /data when persistent storage is enabled, so the download survives restarts. """ local = Path(os.environ.get("ORBIS2_MODELS_DIR") or (Path(__file__).parent / "models")).resolve() if not _missing_models(local): print(f"[startup] Using local Orbis 2 models: {local}") return local print(f"[startup] Fetching Orbis 2 models from {HF_CKPT_REPO} (~20 GB, first run only)…") snapshot = Path(snapshot_download( repo_id=HF_CKPT_REPO, allow_patterns=[f"{sub}/**" for sub, _ in MODEL_SPECS], cache_dir=HF_HUB_CACHE_DIR, )) still_missing = _missing_models(snapshot) if still_missing: raise RuntimeError( f"Downloaded {HF_CKPT_REPO} but these models are incomplete: " f"{', '.join(still_missing)}. Each of L1/, L2/, tok/ must contain its " f"config and '{CKPT_NAME}'." ) print(f"[startup] Orbis 2 models ready: {snapshot}") return snapshot FRAME_H, FRAME_W = 288, 512 L1_FRAME_RATE = 10 # Hz to sample L1 context/rollout frames (and output fps) CONTEXT_FPS = L1_FRAME_RATE # we re-encode uploads to this exact fps for the rollout script if UI_ONLY: print("[startup] --ui-only: skipping model/checkpoint download and load (UI/UX test mode).") MODELS_DIR = None EXP_DIR = None ENGINE = None # No model to introspect for the real lookback requirement -- 0 lets the player # auto-seek stay a no-op instead of erroring. MIN_CONTEXT_END_FRAME = 0 MIN_CONTEXT_TIME_S = 0.0 # No model to read model.num_pred_frames from -- 1 is just a placeholder so the # UI-only slider below has *some* seconds-per-step to show; it's never used for # an actual rollout (run_rollout raises before reaching the model in this mode). FRAMES_PER_GEN_STEP = 1 else: MODELS_DIR = resolve_models_dir() os.environ["ORBIS2_MODELS_DIR"] = str(MODELS_DIR) # L1/L2 configs expandvars this EXP_DIR = MODELS_DIR / "L1" # rollout_demo_v2.py --exp_dir # Load the checkpoint once, at Space startup, on CPU (no GPU is attached outside a # @spaces.GPU call). run_rollout() below only moves it to CUDA / compiles it, and # only does that once too -- see orbis2_app_engine.py for why this avoids paying the # multi-minute load+compile cost on every single generation request. ENGINE = get_engine(EXP_DIR, CONFIG_NAME, CKPT_NAME) # SPACES_ZERO_GPU is set by the platform only on ZeroGPU Spaces, where no GPU is # attached to this (the main) process -- only to a @spaces.GPU-decorated call. On # dedicated hardware CUDA is attached for the process's whole lifetime, so move the # model there and compile it once now instead of paying that cost on the first # request (run_rollout's own ENGINE.ensure_ready() call becomes a no-op once this # has already run -- see RolloutEngine.ensure_ready's idempotency). ON_ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU")) if not ON_ZERO_GPU: print("[startup] not running on ZeroGPU (SPACES_ZERO_GPU unset) -- " "moving model to CUDA and compiling now…") ENGINE.ensure_ready(device="cuda", compile=True, compile_artifacts=str(compile_cache_path())) # Earliest point (in a CONTEXT_FPS-encoded upload) that has enough L1+L2 lookback to # serve as the end of a valid context window -- the player auto-seeks here on load # instead of sitting at time 0 (never valid: there's nothing before it). CPU-only, # so safe to compute once here, right after the CPU-loaded ENGINE above. MIN_CONTEXT_END_FRAME = ENGINE.min_context_end_frame(L1_FRAME_RATE) MIN_CONTEXT_TIME_S = MIN_CONTEXT_END_FRAME / CONTEXT_FPS FRAMES_PER_GEN_STEP = ENGINE.frames_per_rollout_step # Seconds of output video one rollout step yields, at L1_FRAME_RATE fps -- lets the # UI expose rollout length in seconds while the model itself still steps in whole # rollout steps (see GEN_SECONDS_STEP / run_rollout's seconds -> steps conversion below). SECONDS_PER_GEN_STEP = FRAMES_PER_GEN_STEP / L1_FRAME_RATE DEFAULT_GEN_STEPS = 10 # rollout steps; each yields model.num_pred_frames frames GEN_STEPS_MIN, GEN_STEPS_MAX, GEN_STEPS_STEP = 4, 30, 2 # bounds on the underlying step count # The slider itself is in seconds (see build_demo()); these are its bounds, derived # from the step-count bounds above via SECONDS_PER_GEN_STEP. GEN_SECONDS_MIN = GEN_STEPS_MIN * SECONDS_PER_GEN_STEP GEN_SECONDS_MAX = GEN_STEPS_MAX * SECONDS_PER_GEN_STEP GEN_SECONDS_STEP = GEN_STEPS_STEP * SECONDS_PER_GEN_STEP DEFAULT_GEN_SECONDS = DEFAULT_GEN_STEPS * SECONDS_PER_GEN_STEP # L1 and L2 are sampled at different step counts. L1 (detail) needs more steps for # image fidelity; L2 (abstract, consistency-distilled) is designed for very few. DEFAULT_L1_STEPS = 4 # L1 sampler steps (NFE) DEFAULT_L2_STEPS = 4 # L2 sampler steps (NFE) — matches the config's l2_pred_NFE OUTPUT_FPS = L1_FRAME_RATE NUM_VIDEOS = 3 # futures rolled out per run, batched into one minibatch DEFAULT_ETA = 0.0 # 0 = deterministic ODE; >0 injects noise each solver step # Very light cornflower-blue background + one darker cornflower-blue accent used # consistently for the title/subtitle and interactive widgets (see build_demo()'s css=). THEME_BG = "#eef3fb" THEME_ACCENT = "#2f56a3" TITLE_HTML = f"""
Orbis 2
A Hierarchical World Model for Driving
""" AUTHOR_HTML = """
Sudhanshu Mittal*, Arian Mousakhan*, Silvio Galesso*, Karim Farid, Johannes Dienert, Rajat Sahay, Thomas Brox — *main contributors
University of Freiburg, Germany
Project page · Code · Orbis 1
""" DESCRIPTION = f""" Upload a short driving clip (a few seconds is enough), or select one of the example clips. The hierarchical model takes the tail of the clip as L1 (high-rate) and L2 (low-rate, further back) context and autoregressively predicts future frames. Each run samples **{NUM_VIDEOS} rollouts** from the same context in a single minibatch, so you see several plausible futures diverge from identical starting conditions. """ PAPER_MD = """ ### Abstract > Current world models typically operate at a single abstraction level, favoring > perceptual fidelity but lacking the spatial and semantic reasoning needed for > downstream driving tasks. We propose a hierarchical driving world model that > separates prediction into a high-level long-horizon scene forecaster and a > low-level detail generator conditioned on it. This design improves both visual > fidelity and spatial-semantic representation quality. We also introduce a > two-stage training strategy: diffusion-forcing pretraining for richer > representations, followed by teacher-forcing fine-tuning for stable > autoregressive rollouts. Our method achieves state-of-the-art results on > standard driving world model benchmarks, including long-horizon fidelity, > counterfactual steering responsiveness, and internal representation quality. ### Method Orbis 2 splits future prediction across two levels of abstraction: - **Abstract predictor (L2)** — operates over a long temporal context to forecast a future state in latent space, capturing abstract scene dynamics over long horizons and enabling steering control. - **Detail predictor (L1)** — conditioned on that abstract prediction, generates fine-grained short-horizon frames, so high-fidelity local prediction stays grounded in long-range temporal context. Training runs in two stages: **diffusion-forcing pretraining** for richer representations, followed by **teacher-forcing fine-tuning** for stable autoregressive rollouts. ### Links & citation ```bibtex @article{orbis2_2026, author = {Mittal, Sudhanshu and Mousakhan, Arian and Galesso, Silvio and Farid, Karim and Dienert, Johannes and Sahay, Rajat and Brox, Thomas}, title = {Orbis 2: A Hierarchical World Model for Driving}, journal = {arXiv preprint arXiv:2607.15898}, year = {2026}, } ``` """ DEMO_MD = f""" **Pipeline.** Your clip is re-encoded to a constant {CONTEXT_FPS} fps. The rollout script samples two context windows from the tail of that video: an **L1** window at {L1_FRAME_RATE} Hz (fine, recent frames) and an **L2** window sampled further back at the frozen L2 predictor's own trained rate (coarse, long-horizon context). Each frame is resized to {FRAME_H}×{FRAME_W}. The **L2 abstract predictor** forecasts a latent future, and the **L1 detail predictor**, conditioned on it, generates the next frames autoregressively. **Rollout length.** *Generated video length (s)* sets how much video to generate; under the hood this is rounded to a whole number of autoregressive steps, each of which emits {FRAMES_PER_GEN_STEP} image frame{"s" if FRAMES_PER_GEN_STEP != 1 else ""} ({SECONDS_PER_GEN_STEP:.2g}s of video per step). **Sampling.** The two levels are sampled independently, with their own step counts. *L1 sampler steps* controls the detail predictor, which integrates the flow-matching ODE and benefits from more steps (sharper frames, slower). *L2 sampler steps* controls the abstract predictor, which is **consistency-distilled** and so is designed to run in very few steps — raising it buys little. **Randomness.** The {NUM_VIDEOS} rollouts share one context and one seed, but the solver draws independent initial noise per batch element, so the futures diverge. Leave *Seed* at -1 for a fresh draw each run, or set it explicitly to reproduce a run exactly. **Steering (optional).** Expand "Draw a steering trajectory" to freehand a path; it's resampled and smoothed, then used to condition the rollout instead of running unconditionally. Leave it untouched, or hit its Clear button, to roll out without steering, as before. **Note.** The clip must be long enough for the L2 look-back window; very short clips are rejected with an explicit error. Frame and step counts are kept modest to fit the ZeroGPU time budget. """ # ---------------------------------------------------------------------------- # Trajectory canvas (freehand steering input) — copied from app_traj.py, whose # header documents these pieces as copy-paste-ready: the pure Python functions # have no Gradio dependency, and CANVAS_HTML_JS / HEAD_SCRIPT are self-contained. # # Canvas coordinate convention: x = horizontal/lateral (X_MIN_M..X_MAX_M), y = # vertical/forward (Y_MIN_M..Y_MAX_M), so points come back as [lateral, forward]. # rollout_demo_v2.py's --trajectory_file expects the opposite, [forward, lateral] # (see trajectory_to_speed_yawrate in orbis2/evaluate/rollout_demo_v2.py) — the # columns are swapped where the file is written, in run_rollout below. # ---------------------------------------------------------------------------- # Both axes span the same number of meters and the canvas is square with the same # px/meter on both axes, so a shape drawn on screen matches its real-world proportions # instead of being stretched. X is centered on 0 (lateral: left/right of start); Y # starts at 0 (forward: trajectory starts at the bottom, drawn upward). TRAJ_EXTENT_M = 8 # meters spanned by each axis TRAJ_X_MIN_M, TRAJ_X_MAX_M = -TRAJ_EXTENT_M / 2, TRAJ_EXTENT_M / 2 TRAJ_Y_MIN_M, TRAJ_Y_MAX_M = 0, TRAJ_EXTENT_M TRAJ_PX_PER_M = 35 # Shrinks the on-screen canvas to 2/3 of its nominal (TRAJ_EXTENT_M * TRAJ_PX_PER_M) # size. Safe to change independently of TRAJ_PX_PER_M: the JS below (metersToPixel / # pixelToMeters) always normalizes by canvas.width/height, never a hardcoded pixel # density, so it scales cleanly with whatever this canvas's actual size ends up being. TRAJ_CANVAS_SCALE = 1.05 #2 / 3 TRAJ_CANVAS_SIZE_PX = int(TRAJ_EXTENT_M * TRAJ_PX_PER_M * TRAJ_CANVAS_SCALE) # Canvas drawing units (comfortable -5..5m / 0..20m) are not calibrated to this # model's dt=1/L1_FRAME_RATE=0.1s; scale positions up before the model sees them # so a drawn stroke implies a realistic speed. Derived empirically: a near-full- # height canvas stroke (~11.5m arc length over 200 samples) implied ~2.9 m/s at # this model's dt, vs. ~24.2 m/s for a known-good reference trajectory (both # figures computed with the same, since-corrected dt assumption — the 8.4x # ratio between them is dt-invariant) — ratio ~8.4x. TRAJ_REAL_WORLD_SCALE = 8.4 TRAJ_N_SAMPLES = 200 # resolution of the authoritative curve sent to the model TRAJ_SMOOTHING_PASSES = 4 # more passes = smoother; 1 pass is a gentle nudge TRAJ_SMOOTHING_WINDOW_FRAC = 0.15 # smoothing window as a fraction of TRAJ_N_SAMPLES # TRAJ_N_SAMPLES points spread evenly over the calibrated dt=1/L1_FRAME_RATE pace # span TRAJ_N_SAMPLES/L1_FRAME_RATE seconds of implied travel time, so point index # i represents i/L1_FRAME_RATE seconds -- used to mark one point per second on the # canvas overlay instead of all TRAJ_N_SAMPLES (too dense to read). TRAJ_SAMPLES_PER_SECOND = L1_FRAME_RATE def resample_by_arclength(points: np.ndarray, n_samples: int) -> np.ndarray: """Resample a raw (possibly unevenly-spaced) path to `n_samples` points evenly spaced by arc length. Freehand strokes have point density that depends on how fast the user moved the mouse, so this normalizes that out before smoothing -- otherwise slow/fast sections would get smoothed by different effective amounts. """ pts = np.asarray(points, dtype=float) if len(pts) < 2: return np.repeat(pts, max(n_samples, 1), axis=0)[:n_samples] deltas = np.diff(pts, axis=0) seg_lengths = np.sqrt((deltas ** 2).sum(axis=1)) cum_len = np.concatenate([[0.0], np.cumsum(seg_lengths)]) total_len = cum_len[-1] if total_len == 0: # User clicked without dragging: no real path, just repeat the point. return np.repeat(pts[:1], n_samples, axis=0) targets = np.linspace(0, total_len, n_samples) x = np.interp(targets, cum_len, pts[:, 0]) y = np.interp(targets, cum_len, pts[:, 1]) return np.stack([x, y], axis=1) def smooth_trajectory( points: np.ndarray, n_samples: int = TRAJ_N_SAMPLES, passes: int = TRAJ_SMOOTHING_PASSES, window_frac: float = TRAJ_SMOOTHING_WINDOW_FRAC, ) -> np.ndarray: """Turn a raw freehand stroke into a smooth model-facing trajectory: resample to even spacing, then repeatedly box-filter both coordinate channels (more passes ~ approximates a Gaussian blur, i.e. "heavy" smoothing). Finally re-anchors the curve so it starts exactly at (0, 0), since smoothing can nudge the very first point slightly, and rotates it about that origin so the initial segment points straight up (+y, zero lateral/x component) regardless of which direction the stroke was actually drawn in. """ pts = resample_by_arclength(points, n_samples) if len(pts) < 2: return pts window = max(int(n_samples * window_frac), 3) if window % 2 == 0: window += 1 # odd window keeps the filter centered kernel = np.ones(window) / window smoothed = pts.copy() for _ in range(passes): # Edge-pad so the curve doesn't shrink/shorten at the endpoints. pad = window // 2 padded_x = np.pad(smoothed[:, 0], pad, mode="edge") padded_y = np.pad(smoothed[:, 1], pad, mode="edge") smoothed = np.stack( [np.convolve(padded_x, kernel, mode="valid"), np.convolve(padded_y, kernel, mode="valid")], axis=1, ) smoothed -= smoothed[0] # force the path to start exactly at (0, 0) # Rotate about the origin so the initial segment points straight up (+y, # zero lateral/x component) -- mirrors the heading-alignment rotation # trajectory_to_model_frame() applies downstream, but at the canvas-space # level so the on-screen overlay/plot already reflects it, not just the # file eventually written for the model. heading0 = np.arctan2(smoothed[1, 1] - smoothed[0, 1], smoothed[1, 0] - smoothed[0, 0]) delta = np.pi / 2 - heading0 cos_d, sin_d = np.cos(delta), np.sin(delta) rotation = np.array([[cos_d, -sin_d], [sin_d, cos_d]]) smoothed = smoothed @ rotation.T return smoothed def parse_points_json(raw_json: str) -> np.ndarray: """Parse the JSON string written by JS (list of [x, y] in real meters, x in roughly [TRAJ_X_MIN_M, TRAJ_X_MAX_M], y in roughly [TRAJ_Y_MIN_M, TRAJ_Y_MAX_M]) into an (N, 2) numpy array. Returns an empty (0, 2) array if the input is empty/invalid, so callers can handle the "nothing drawn yet" state without try/except everywhere. """ try: data = json.loads(raw_json) if raw_json else [] pts = np.array(data, dtype=float) return pts.reshape(-1, 2) if pts.size else np.empty((0, 2)) except (json.JSONDecodeError, ValueError): return np.empty((0, 2)) def points_to_trajectory(raw_json: str, n_samples: int = TRAJ_N_SAMPLES) -> np.ndarray: """End-to-end: raw freehand stroke JSON -> authoritative smoothed (n_samples, 2) trajectory array. This is the function whose output should actually be fed into the model. """ pts = parse_points_json(raw_json) if len(pts) < 2: return np.empty((0, 2)) return smooth_trajectory(pts, n_samples) def trajectory_to_model_frame(trajectory: np.ndarray) -> np.ndarray: """Convert a canvas-drawn trajectory ([x=lateral(right+), y=forward], in canvas meters) into the [forward, lateral] array rollout_demo_v2.py's --trajectory_file expects, applying three corrections the canvas alone doesn't account for: 1. Column swap: canvas is [lateral, forward]; the model wants [forward, lateral]. 2. Sign flip on lateral: the canvas's x-axis is positive-right (screen convention), but rollout_demo_v2.py's own overlay rendering (_panel_coords_fit_trajectory: `px = margin + (lateral_max - lateral) * scale`) shows increasing lateral moving LEFT on screen, i.e. this codebase's convention is positive-lateral = left. Without negating, a rightward-drawn stroke would condition/render as a left turn. 3. Heading alignment: trajectory_to_speed_yawrate/its reconstruction always treats the trajectory's own first segment as pointing exactly along local +forward. A freehand stroke's first smoothed segment is essentially never *exactly* aligned with forward (mouse imprecision, smoothing edge effects) -- even a couple of degrees of initial misalignment, left uncorrected, gets amplified into meters of apparent lateral drift by the far end of a long trajectory once scaled up by TRAJ_REAL_WORLD_SCALE. Rotating so segment 0 already points along +forward before conversion matches what the reconstruction assumes, eliminating that drift. Verified against a real canvas-drawn+smoothed trajectory: this reduces the round-trip (speed/yaw_rate -> reconstructed position) error from several meters down to ~0.16m over a ~97m trajectory. """ traj = np.asarray(trajectory)[:, [1, 0]] * TRAJ_REAL_WORLD_SCALE # [forward, lateral(right+)] traj = traj.copy() traj[:, 1] *= -1 # canvas right+ -> model left+ traj -= traj[0] heading0 = np.arctan2(traj[1, 1] - traj[0, 1], traj[1, 0] - traj[0, 0]) cos_h, sin_h = np.cos(-heading0), np.sin(-heading0) rotation = np.array([[cos_h, -sin_h], [sin_h, cos_h]]) return traj @ rotation.T def render_trajectory_plot(raw_json: str): """Return the dense trajectory array into a gr.State for downstream use, plus that same trajectory as a JSON string -- the JSON string is consumed client-side to draw the smoothed curve back onto the canvas (see the `smoothed_points_json` wiring in build_demo()). """ trajectory = points_to_trajectory(raw_json) smoothed_json = json.dumps(trajectory.tolist()) if len(trajectory) else "[]" return trajectory, smoothed_json # -------------------------------------------------------------------------- # Canvas markup only -- NO """ # ---------------------------------------------------------------------------- # Input-video cursor tracking. Rather than physically cropping the upload, # "Generate rollouts" uses wherever the playhead is left as the end of the # context window (see run_rollout / RolloutEngine.roll_out's context_end_frame). # # Gradio tears down and recreates the