Upload folder using huggingface_hub

README.md

@@ -1,13 +1,28 @@
 ---
-title: Latent Dynamics Reasoning
-emoji: 👀
-colorFrom: yellow
-colorTo: purple
+title: Latent Dynamics Reasoning (LDR)
+emoji: 🎬
+colorFrom: red
+colorTo: pink
 sdk: gradio
-sdk_version: 6.22.0
-python_version: '3.12'
+sdk_version: 6.15.1
 app_file: app.py
-pinned: false
+short_description: Extrapolate physical video from 3 conditioning frames
+python_version: "3.10"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Latent Dynamics Reasoning (LDR)
+
+Extrapolative video world model that rolls out future video frames from three
+initial conditioning images, reasoning about latent physical dynamics.
+
+Upload three frames of a physical scene, pick the task type (uniform,
+parabola, collision, looming, bouncing), and the model generates ~29 future
+frames of predicted motion.
+
+Based on: [Learning How the World Evolves: Extrapolative Video World Models
+via Latent Dynamics Reasoning](https://lat-dyn-reason.github.io/)
+
+- [Model](https://huggingface.co/haodongli/LDR) (Apache-2.0)
+- [Code](https://github.com/Lat-Dyn-Reason/Lat-Dyn-Reason)
+- [Project Page](https://lat-dyn-reason.github.io/)

app.py

@@ -0,0 +1,219 @@
+"""LDR — Extrapolative Video World Models via Latent Dynamics Reasoning.
+
+Upload three conditioning frames, pick a task type, and the model rolls out
+~29 future frames of physical motion (uniform, parabola, collision, looming,
+or bouncing) and returns them as a video.
+"""
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST be before torch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import gradio as gr
+import tempfile
+import imageio
+import imageio.v3 as iio
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+from ldr import build_ldr
+
+# ---------------------------------------------------------------------------
+# Model loading (module scope — ZeroGPU intercepts .to("cuda"))
+# ---------------------------------------------------------------------------
+
+MODEL_ID = "haodongli/LDR"
+
+# The joint-5task checkpoint handles all five tasks in one model.
+# Single-task checkpoints are smaller but only work for one task.
+# We load both 256x256 single-task and joint so users can pick.
+CKPT_CONFIGS = {
+    "uniform":   ("256x256/single_task/uniform.pt",   256),
+    "parabola":  ("256x256/single_task/parabola.pt",  256),
+    "collision": ("256x256/single_task/collision.pt", 256),
+    "looming":   ("256x256/single_task/looming.pt",   256),
+    "bouncing":  ("256x256/single_task/bouncing.pt",  256),
+    "joint":     ("256x256/joint_task/joint_5task.pt", 256),
+}
+
+OURS_ARCH = dict(num_pred=29, width=256, accel_scale=0.5, n_kp=16,
+                 warp_flow_res=64, kappa_init=0.15)
+
+
+def _load_checkpoint(ckpt_path, img_size):
+    """Load a .pt checkpoint into an LDR model (mirrors eval.build_model)."""
+    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    arch = dict(OURS_ARCH)
+    cargs = ck.get("args") if isinstance(ck, dict) else None
+    if cargs:
+        arch.update(width=cargs.get("accel_width", 256),
+                    n_kp=cargs.get("n_kp", 16),
+                    accel_scale=cargs.get("accel_scale", 0.5),
+                    kappa_init=cargs.get("kappa_init", cargs.get("accel_init_damp", 0.15)))
+    model = build_ldr(**arch).to("cuda").eval()
+    sd = ck["model"] if (isinstance(ck, dict) and "model" in ck) else ck
+    sd = {k[7:] if k.startswith("module.") else k: v for k, v in sd.items()}
+    model.load_state_dict(sd)
+    return model
+
+
+# Pre-download all checkpoints at module scope so they are on disk for the
+# worker to stream.
+_ckpt_cache = {}  # task_name -> (model, img_size)
+for _task, (_rel, _img) in CKPT_CONFIGS.items():
+    _path = hf_hub_download(MODEL_ID, _rel, repo_type="model")
+    _model = _load_checkpoint(_path, _img)
+    _ckpt_cache[_task] = (_model, _img)
+
+
+def _get_model(task):
+    return _ckpt_cache[task]
+
+
+# ---------------------------------------------------------------------------
+# Frame utilities (ported from infer.py)
+# ---------------------------------------------------------------------------
+
+def _resize(ft, fr, img_size):
+    """Resize a frame tensor (N, 3, H, W) and its numpy companion to img_size."""
+    if ft.shape[-1] != img_size:
+        ft = nn.functional.interpolate(ft, size=(img_size, img_size),
+                                       mode="bilinear", align_corners=False)
+        fr = (((ft + 1) * 127.5).clamp(0, 255).byte().permute(0, 2, 3, 1).numpy())
+    return ft, fr
+
+
+def _load_frames(images, img_size):
+    """Load a list of PIL images → (float_tensor [-1,1], numpy uint8 frames)."""
+    fr = np.stack([np.asarray(img.convert("RGB")) for img in images], 0)
+    ft = (torch.from_numpy(fr.astype(np.float32)).permute(0, 3, 1, 2) / 127.5 - 1.0)
+    return _resize(ft, fr, img_size)
+
+
+def _gen(model, frames_t, frames_np, nc):
+    """Run the model and concatenate conditioning frames with predictions."""
+    cond_t = frames_t[:nc].to("cuda")
+    with torch.no_grad():
+        pred = model(cond_t.unsqueeze(0), nc, full=False,
+                      cond_img=cond_t[nc - 1:nc])
+        pf = ((pred.squeeze(0).clamp(-1, 1) + 1) * 127.5).byte().cpu() \
+              .permute(0, 2, 3, 1).numpy()
+    return np.concatenate([frames_np[:nc], pf], axis=0)
+
+
+# ---------------------------------------------------------------------------
+# Inference function
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=120)
+def generate(frame1, frame2, frame3, task="uniform",
+             fps=8, progress=gr.Progress(track_tqdm=True)):
+    """Roll out future video frames from three conditioning images.
+
+    Args:
+        frame1: First conditioning frame (earliest in the sequence).
+        frame2: Second conditioning frame.
+        frame3: Third conditioning frame (latest in the sequence).
+        task: Physics task type — uniform, parabola, collision, looming, or bouncing.
+             Use "joint" for the model trained on all five tasks.
+        fps: Output video frames per second.
+    Returns:
+        Path to the generated MP4 video file.
+    """
+    model, img_size = _get_model(task)
+    images = [frame1, frame2, frame3]
+    ft, fr = _load_frames(images, img_size)
+    all_frames = _gen(model, ft, fr, num_cond=3)
+
+    # Write to a unique temp file (concurrency-safe)
+    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False, dir="/tmp")
+    tmp.close()
+    imageio.mimwrite(tmp.name, list(all_frames), fps=fps, codec="libx264",
+                     quality=8, output_params=["-pix_fmt", "yuv420p"])
+    return tmp.name
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown("""
+    # Latent Dynamics Reasoning (LDR)
+    Upload three conditioning frames of a physical scene and the model
+    extrapolates ~29 future frames as video. Pick the physics task type that
+    matches your input (uniform motion, parabola, collision, looming, or
+    bouncing), or use the **joint** model trained on all five.
+
+    [Model](https://huggingface.co/haodongli/LDR) ·
+    [Code](https://github.com/Lat-Dyn-Reason/Lat-Dyn-Reason) ·
+    [Project Page](https://lat-dyn-reason.github.io/)
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            frame1 = gr.Image(label="Frame 1 (t=0)", type="pil",
+                              height=200)
+            frame2 = gr.Image(label="Frame 2 (t=1)", type="pil",
+                              height=200)
+            frame3 = gr.Image(label="Frame 3 (t=2)", type="pil",
+                              height=200)
+        with gr.Column(scale=2):
+            task = gr.Dropdown(
+                choices=["uniform", "parabola", "collision", "looming",
+                         "bouncing", "joint"],
+                value="uniform",
+                label="Task / Checkpoint",
+                info="Match the physics type of your scene. "
+                     "'joint' handles all five but slightly lower quality.")
+            fps = gr.Slider(4, 16, value=8, step=1, label="Output FPS")
+            run = gr.Button("Generate Video", variant="primary")
+
+    output = gr.Video(label="Generated Video (conditioning + prediction)")
+
+    with gr.Accordion("Examples", open=True):
+        gr.Examples(
+            examples=[
+                # [frame1, frame2, frame3, task, fps]
+                ["examples/uniform/00.png", "examples/uniform/01.png",
+                 "examples/uniform/02.png", "uniform", 8],
+                ["examples/parabola/00.png", "examples/parabola/01.png",
+                 "examples/parabola/02.png", "parabola", 8],
+                ["examples/collision/00.png", "examples/collision/01.png",
+                 "examples/collision/02.png", "collision", 8],
+                ["examples/looming/00.png", "examples/looming/01.png",
+                 "examples/looming/02.png", "looming", 8],
+                ["examples/bouncing/00.png", "examples/bouncing/01.png",
+                 "examples/bouncing/02.png", "bouncing", 8],
+                ["examples/uniform_pikachu/00.png",
+                 "examples/uniform_pikachu/01.png",
+                 "examples/uniform_pikachu/02.png", "uniform", 8],
+                ["examples/uniform_soccer/00.png",
+                 "examples/uniform_soccer/01.png",
+                 "examples/uniform_soccer/02.png", "uniform", 8],
+            ],
+            inputs=[frame1, frame2, frame3, task, fps],
+            outputs=output,
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run.click(
+        fn=generate,
+        inputs=[frame1, frame2, frame3, task, fps],
+        outputs=output,
+        api_name="generate",
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

ldr/__init__.py

@@ -0,0 +1,4 @@
+"""Latent Dynamics Reasoning (LDR). See ldr/rollout.md for the rollout derivation."""
+from .model import build_ldr, LDR, PerceptualPyramidLoss
+
+__all__ = ["build_ldr", "LDR", "PerceptualPyramidLoss"]

ldr/metrics.py

@@ -0,0 +1,326 @@
+"""PhyWorld white-box parser: extract each ball's center/radius from predicted frames and score vs GT."""
+import numpy as np
+import pandas as pd
+
+left_color = (255, 0, 0)
+right_color = (0, 0, 255)
+COLORS = [left_color, right_color]
+WORLD_SCALE = 10.0
+
+
+def parse_state_from_image(image_rgb, default_y, default_r1, thres=0.15, color=0):
+    image_copy = image_rgb.copy()
+    all_scaled_circles = []
+    if color == 0:
+        circle_mask = (image_copy[:, :, 0] > 127) & (image_copy[:, :, 1] < 127) & (image_copy[:, :, 2] < 127)
+    elif color == 1:
+        circle_mask = (image_copy[:, :, 0] < 127) & (image_copy[:, :, 1] < 127) & (image_copy[:, :, 2] > 127)
+    else:
+        raise ValueError("Invalid color")
+
+    area = np.sum(circle_mask)
+    radius = np.sqrt(area / np.pi)
+    scaled_radius = radius / image_copy.shape[1] * WORLD_SCALE
+    if area > thres:
+        center = (
+            np.mean(np.nonzero(circle_mask)[1]),
+            np.mean(np.nonzero(circle_mask)[0]),
+        )
+        scaled_center = (
+            center[0] / image_copy.shape[1] * WORLD_SCALE,
+            center[1] / image_copy.shape[1] * WORLD_SCALE,
+        )
+    else:
+        scaled_center = (WORLD_SCALE, default_y)
+        scaled_radius = default_r1
+
+    all_scaled_circles.append([*scaled_center, scaled_radius])
+    return np.array(all_scaled_circles)
+
+
+def parse_state_from_image_collision(image_rgb, default_y, default_r1, default_r2, thres=0.15):
+    image_copy = image_rgb.copy()
+    all_scaled_circles = []
+    for ball_id, color in enumerate(COLORS):
+        if color == (255, 0, 0):
+            circle_mask = (image_copy[:, :, 0] > 127) & (image_copy[:, :, 1] < 127) & (image_copy[:, :, 2] < 127)
+        elif color == (0, 0, 255):
+            circle_mask = (image_copy[:, :, 0] < 127) & (image_copy[:, :, 1] < 127) & (image_copy[:, :, 2] > 127)
+        else:
+            raise ValueError("Invalid color")
+
+        area = np.sum(circle_mask)
+        radius = np.sqrt(area / np.pi)
+        scaled_radius = radius / image_copy.shape[1] * WORLD_SCALE
+        if area > thres:
+            center = (
+                np.mean(np.nonzero(circle_mask)[1]),
+                np.mean(np.nonzero(circle_mask)[0]),
+            )
+            scaled_center = (
+                center[0] / image_copy.shape[1] * WORLD_SCALE,
+                center[1] / image_copy.shape[1] * WORLD_SCALE,
+            )
+        else:
+            if ball_id == 0:
+                scaled_center = (0, default_y)
+                scaled_radius = default_r1
+            else:
+                scaled_center = (WORLD_SCALE, default_y)
+                scaled_radius = default_r2
+
+        all_scaled_circles.append([*scaled_center, scaled_radius])
+    return np.array(all_scaled_circles)
+
+
+def get_last_ema(values_list, span):
+    series = pd.Series(values_list)
+    ema = series.ewm(span=span, adjust=False).mean()
+    return ema.iloc[-1]
+
+
+def xy_metrics(list_a, list_b):
+    distances_x, distances_y = [], []
+    assert len(list_a) == len(list_b) == 1
+    for elem_a, elem_b in zip(list_a, list_b):
+        if not np.any(np.isnan(elem_a)) and not np.any(np.isnan(elem_b)):
+            distances_x.append(np.abs(elem_a[0] - elem_b[0]))
+            distances_y.append(np.abs(elem_a[1] - elem_b[1]))
+    x_error_avg = np.mean(distances_x) if distances_x else np.nan
+    y_error_avg = np.mean(distances_y) if distances_y else np.nan
+    return x_error_avg, y_error_avg
+
+
+def evaluate_xy(rollout_frames, gt_features, init, mode, gamma=0.98, sample_freq=1, pred_states=None):
+    assert sample_freq == 1, 'there may be some bugs if it is greater than 1'
+
+    # keep only frames where the ball is fully in view
+    left_ball_r = init[0]
+    left_ball_init_v = init[1]
+    left_ball_m = init[0]**2
+    index = []
+    CONDITION_FRAMES = 4
+    for i, state in enumerate(gt_features):
+        if i < CONDITION_FRAMES-1:
+            continue
+        left_ball_x, left_ball_y = state[0], state[1]
+        if left_ball_x - left_ball_r >= 0 and left_ball_x + left_ball_r <= WORLD_SCALE \
+            and left_ball_y - left_ball_r >= 0 and left_ball_y + left_ball_r <= WORLD_SCALE:
+            index.append(i)
+
+    if rollout_frames is not None:
+        assert len(rollout_frames) == len(gt_features), f'{len(rollout_frames)}, {len(gt_features)}'
+        rollout_frames = rollout_frames[index]
+        gt_features = gt_features[index]
+
+        x_error_list, y_error_list, r_list = [], [], []
+        x_pos_list, y_pos_list = [], []
+        default_y, default_r1 = np.nan, np.nan
+        for rollout_frame, gt_feature in zip(rollout_frames[sample_freq-1::sample_freq], gt_features[sample_freq-1::sample_freq]):
+            parsed_state = parse_state_from_image(rollout_frame, default_y, default_r1, color=init[2] if len(init) >= 3 else 0)
+            default_y, default_r1 = parsed_state[0][1], parsed_state[0][2]
+            x_pos_list.append(parsed_state[:, 0])
+            y_pos_list.append(parsed_state[:, 1])
+            r_list.append(parsed_state[:, 2])
+            parsed_state = parsed_state[:, :2]
+            x_error, y_error = xy_metrics([gt_feature], parsed_state)
+            x_error_list.append(x_error)
+            y_error_list.append(y_error)
+        span = len(rollout_frames[sample_freq-1::sample_freq])
+        ema_x_error = get_last_ema(x_error_list, span=span)
+        ema_y_error = get_last_ema(y_error_list, span=span)
+
+    else:
+        assert pred_states is not None and len(pred_states) == len(gt_features), f'{len(pred_states)}, {len(gt_features)}'
+        pred_states = pred_states[index]
+        gt_features = gt_features[index]
+        x_pos_list = pred_states[:, 0]
+        y_pos_list = pred_states[:, 1]
+        r_list = pred_states[:, 2]
+
+    r1_list = r_list
+    r1_min, r1_max = min(r1_list), max(r1_list)
+    delta_r = r1_max - r1_min
+
+    gt_x_pos = gt_features[:, 0]
+    gt_x_vel = np.diff(gt_x_pos, axis=0) / (0.1 * sample_freq)
+    gt_y_pos = gt_features[:, 1]
+    gt_y_vel = np.diff(gt_y_pos, axis=0) / (0.1 * sample_freq)
+
+    x_pos = np.array(x_pos_list)
+    y_pos = np.array(y_pos_list)
+    if x_pos.shape[-1] == 1:
+        x_pos = np.squeeze(x_pos, -1)
+        y_pos = np.squeeze(y_pos, -1)
+
+    x_vel = np.diff(x_pos, axis=0) / (0.1 * sample_freq)
+    y_vel = np.diff(y_pos, axis=0) / (0.1 * sample_freq)
+
+    x_err_avg = np.mean(np.abs(x_pos - gt_x_pos))
+    y_err_avg = np.mean(np.abs(y_pos - gt_y_pos))
+
+    x_vel_err_avg = np.mean(np.abs(x_vel - gt_x_vel))
+    y_vel_err_avg = np.mean(np.abs(y_vel - gt_y_vel))
+
+    if mode == 'all':
+        return {
+            'init': init,
+            'x_pos': x_pos,
+            'y_pos': y_pos,
+            'x_vel': x_vel,
+            'y_vel': y_vel,
+            'r': r_list,
+            'gt_x_pos': gt_x_pos,
+            'gt_y_pos': gt_y_pos,
+            'abs_x_err_avg': x_err_avg,
+            'abs_y_err_avg': y_err_avg,
+            'delta_r': delta_r,
+            'abs_x_vel_err_avg': x_vel_err_avg,
+            'abs_y_vel_err_avg': y_vel_err_avg,
+        }
+    else:
+        return ema_x_error, ema_y_error
+
+
+def evaluate_xy_collision(rollout_frames, gt_features, init, mode, gamma=0.98, sample_freq=1, pred_states=None):
+    assert sample_freq == 1, 'there may be some bugs if it is greater than 1'
+
+    left_ball_r, right_ball_r = init[:2]
+    left_ball_init_v, right_ball_init_v = init[2:4]
+    left_ball_m, right_ball_m = init[:2]**2
+    index = []
+    CONDITION_FRAMES = 4
+    for i, state in enumerate(gt_features):
+        if i < CONDITION_FRAMES-1:
+            continue
+        left_ball_x = state[0, 0]
+        right_ball_x = state[1, 0]
+        if left_ball_x - left_ball_r >= 0 and right_ball_x + right_ball_r <= WORLD_SCALE:
+            index.append(i)
+
+    if rollout_frames is not None:
+        assert len(rollout_frames) == len(gt_features), f'{len(rollout_frames)}, {len(gt_features)}'
+        rollout_frames = rollout_frames[index]
+        gt_features = gt_features[index]
+
+        r_list = []
+        x_pos_list = []
+        y_pos_list = []
+        default_y, default_r1, default_r2 = np.nan, np.nan, np.nan
+        for rollout_frame, gt_feature in zip(rollout_frames[sample_freq-1::sample_freq], gt_features[sample_freq-1::sample_freq]):
+            parsed_state = parse_state_from_image_collision(rollout_frame, default_y, default_r1, default_r2)
+            default_y, default_r1, default_r2 = parsed_state[0][1], parsed_state[0][2], parsed_state[1][2]
+            x_pos_list.append(parsed_state[:, 0])
+            y_pos_list.append(parsed_state[:, 1])
+            r_list.append(parsed_state[:, 2])
+            parsed_state = parsed_state[:, :2]
+
+    else:
+        assert pred_states is not None and len(pred_states) == len(gt_features), f'{len(pred_states)}, {len(gt_features)}'
+        pred_states = pred_states[index]
+        gt_features = gt_features[index]
+        x_pos_list = pred_states[:, :2]
+        y_pos_list = pred_states[:, 2:4]
+        r_list = pred_states[:, 4:]
+
+    # only use frames before collision to avoid out-of-vision region
+    MIN_POST_FRAMES = 8
+    r_list = r_list[:(MIN_POST_FRAMES-CONDITION_FRAMES+1)//sample_freq]
+    r1_list = [x[0] for x in r_list]
+    r2_list = [x[1] for x in r_list]
+    r1_min, r1_max = min(r1_list), max(r1_list)
+    r2_min, r2_max = min(r2_list), max(r2_list)
+    delta_r = max(r1_max - r1_min, r2_max - r2_min)
+
+    gt_x_pos = gt_features[:, :, 0]
+    gt_y_pos = gt_features[:, :, 1]
+    gt_x_vel = np.diff(gt_x_pos, axis=0) / (0.1 * sample_freq)
+
+    # collision frame = first large change in GT x-velocity
+    collision_index = len(gt_x_vel) - 1
+    for i in range(1, len(gt_x_vel)):
+        if np.abs(gt_x_vel[i, 0] - gt_x_vel[i-1, 0]) > 0.1 or np.abs(gt_x_vel[i, 1] - gt_x_vel[i-1, 1]) > 0.1:
+            collision_index = i
+            break
+
+    x_pos = np.array(x_pos_list)
+    x_vel = np.diff(x_pos, axis=0) / (0.1 * sample_freq)
+
+    pre_x_err_avg = np.mean(np.abs(x_pos[:collision_index+1] - gt_x_pos[:collision_index+1]))
+    post_x_err_avg = np.mean(np.abs(x_pos[collision_index+1:] - gt_x_pos[collision_index+1:]))
+
+    y_pos = np.array(y_pos_list)
+    pre_y_err_avg = np.mean(np.abs(y_pos[:collision_index+1] - gt_y_pos[:collision_index+1]))
+    post_y_err_avg = np.mean(np.abs(y_pos[collision_index+1:] - gt_y_pos[collision_index+1:]))
+
+    pre_vel_err_avg = np.mean(np.abs(x_vel[:collision_index] - gt_x_vel[:collision_index]))
+    post_vel_err_avg = np.mean(np.abs(x_vel[collision_index+1:] - gt_x_vel[collision_index+1:]))
+
+    momentum = left_ball_m * x_vel[:, 0] + right_ball_m * x_vel[:, 1]
+    gt_momentum = left_ball_m * left_ball_init_v - right_ball_m * right_ball_init_v
+    pre_momentum_error_avg = np.mean(np.abs(momentum[:collision_index] - gt_momentum))
+    post_momentum_error_avg = np.mean(np.abs(momentum[collision_index+1:] - gt_momentum))
+
+    energy = left_ball_m * x_vel[:, 0]**2 / 2 + right_ball_m * x_vel[:, 1]**2 / 2
+    gt_energy = left_ball_m * left_ball_init_v**2 / 2 + right_ball_m * right_ball_init_v**2 / 2
+    pre_energy_error_avg = np.mean(np.abs(energy[:collision_index] - gt_energy))
+    post_energy_error_avg = np.mean(np.abs(energy[collision_index+1:] - gt_energy))
+
+    return {
+        'init': init,
+        'collision_index': int(collision_index),
+        'pre_x_vel': x_vel[:collision_index],
+        'gt_pre_x_vel': gt_x_vel[:collision_index],
+        'post_x_vel': x_vel[collision_index+1:],
+        'gt_post_x_vel': gt_x_vel[collision_index+1:],
+        'x_pos': x_pos,
+        'gt_x_pos': gt_x_pos,
+        'x_vel': x_vel,
+        'abs_x_err_avg': float(np.mean(np.abs(x_pos - gt_x_pos))),
+        'abs_y_err_avg': float(np.mean(np.abs(y_pos - gt_y_pos))),
+        'y_pos': y_pos,
+        'gt_y_pos': gt_y_pos,
+        'delta_r': delta_r,
+        'pre_x_err_avg': pre_x_err_avg,
+        'post_x_err_avg': post_x_err_avg,
+        'pre_y_err_avg': pre_y_err_avg,
+        'post_y_err_avg': post_y_err_avg,
+        'pre_vel_err_avg': pre_vel_err_avg,
+        'post_vel_err_avg': post_vel_err_avg,
+        'pre_momentum_error_avg': pre_momentum_error_avg,
+        'post_momentum_error_avg': post_momentum_error_avg,
+        'pre_energy_error_avg': pre_energy_error_avg,
+        'post_energy_error_avg': post_energy_error_avg,
+    }
+
+
+def evaluate_xy_looming(rollout_frames, gt_pos, gt_radius, init, mode='all', sample_freq=1):
+    """Looming (scale dynamics): parsed ball radius r(t) vs GT radius, plus x/y position; primary axis = radius."""
+    CONDITION_FRAMES = 4
+    m = min(len(rollout_frames), len(gt_pos), len(gt_radius))
+    idx = []
+    for i in range(m):
+        if i < CONDITION_FRAMES - 1:
+            continue
+        x, y, r = float(gt_pos[i, 0]), float(gt_pos[i, 1]), float(gt_radius[i])
+        if x - r >= 0 and x + r <= WORLD_SCALE and y - r >= 0 and y + r <= WORLD_SCALE:
+            idx.append(i)
+    if len(idx) == 0:
+        return {'abs_r_err_avg': np.nan, 'abs_x_err_avg': np.nan, 'abs_y_err_avg': np.nan, 'abs_r_vel_err_avg': np.nan}
+    xs, ys, rs = [], [], []
+    default_y, default_r = np.nan, np.nan
+    for i in idx:
+        ps = parse_state_from_image(rollout_frames[i], default_y, default_r, color=0)
+        default_y, default_r = ps[0][1], ps[0][2]
+        xs.append(float(ps[0][0])); ys.append(float(ps[0][1])); rs.append(float(ps[0][2]))
+    xs, ys, rs = np.array(xs), np.array(ys), np.array(rs)
+    gx = np.array([float(gt_pos[i, 0]) for i in idx]); gy = np.array([float(gt_pos[i, 1]) for i in idx])
+    gr = np.array([float(gt_radius[i]) for i in idx])
+    r_vel_err = float(np.mean(np.abs(np.diff(rs) - np.diff(gr)) / (0.1 * sample_freq))) if len(rs) > 1 else np.nan
+    return {
+        'init': init, 'r': rs, 'gt_r': gr, 'x_pos': xs, 'y_pos': ys,
+        'abs_r_err_avg': float(np.mean(np.abs(rs - gr))),
+        'abs_x_err_avg': float(np.mean(np.abs(xs - gx))),
+        'abs_y_err_avg': float(np.mean(np.abs(ys - gy))),
+        'abs_r_vel_err_avg': r_vel_err,
+    }

ldr/model.py

@@ -0,0 +1,455 @@
+"""LDR model: structured-latent encoder, kinematic-integration rollout, warp-render decoder, perceptual loss."""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class _KeypointEnc(nn.Module):
+    """Frame (N,3,H,W) -> K heatmaps at H/8 -> marginal soft-argmax -> structured latent (N,K,3)=(mu_x,mu_y,sigma)."""
+    def __init__(self, n_kp, gn=8):
+        super().__init__()
+        self.n_kp = n_kp
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 32, 7, 1, 3), nn.GroupNorm(gn, 32), nn.SiLU(),
+            nn.Conv2d(32, 64, 3, 2, 1), nn.GroupNorm(gn, 64), nn.SiLU(),
+            nn.Conv2d(64, 64, 3, 1, 1), nn.GroupNorm(gn, 64), nn.SiLU(),
+            nn.Conv2d(64, 96, 3, 2, 1), nn.GroupNorm(gn, 96), nn.SiLU(),
+            nn.Conv2d(96, 96, 3, 1, 1), nn.GroupNorm(gn, 96), nn.SiLU(),
+            nn.Conv2d(96, 128, 3, 2, 1), nn.GroupNorm(gn, 128), nn.SiLU(),
+            nn.Conv2d(128, 128, 3, 1, 1), nn.GroupNorm(gn, 128), nn.SiLU(),
+            nn.Conv2d(128, n_kp, 1))
+
+    def forward(self, x):
+        hm = self.net(x)
+        H, W = hm.shape[-2:]
+        gx = torch.linspace(-1, 1, W, device=x.device).view(1, 1, W)
+        gy = torch.linspace(-1, 1, H, device=x.device).view(1, 1, H)
+        px = torch.softmax(hm.mean(2), dim=-1)
+        py = torch.softmax(hm.mean(3), dim=-1)
+        cx = (px * gx).sum(-1); cy = (py * gy).sum(-1)             # centroid mu
+        vx = (px * (gx - cx.unsqueeze(-1)) ** 2).sum(-1)
+        vy = (py * (gy - cy.unsqueeze(-1)) ** 2).sum(-1)
+        s = torch.sqrt(0.5 * (vx + vy) + 1e-6)                     # extent sigma
+        return torch.stack([cx, cy, s], dim=-1)
+
+
+def _make_coord_grid(h, w, device, dtype=torch.float32):
+    y = torch.linspace(-1, 1, h, device=device, dtype=dtype)
+    x = torch.linspace(-1, 1, w, device=device, dtype=dtype)
+    xx = x.view(1, w).expand(h, w)
+    yy = y.view(h, 1).expand(h, w)
+    return torch.stack([xx, yy], dim=-1)
+
+
+def _region_gaussian(centers, stds, res, device):
+    grid = _make_coord_grid(res, res, device).view(1, 1, res, res, 2)
+    c = centers.view(*centers.shape[:-1], 1, 1, 2)
+    var = (stds ** 2).view(*stds.shape, 1, 1).clamp(min=1e-6)
+    d2 = ((grid - c) ** 2).sum(-1)
+    return torch.exp(-0.5 * d2 / var)
+
+
+class _GNSameBlock(nn.Module):
+    def __init__(self, cin, cout, kernel_size=3, padding=1, gn=8):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cout, kernel_size, padding=padding)
+        self.norm = nn.GroupNorm(gn, cout)
+
+    def forward(self, x):
+        return F.relu(self.norm(self.conv(x)))
+
+
+class _GNDownBlock(nn.Module):
+    def __init__(self, cin, cout, gn=8):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm = nn.GroupNorm(gn, cout)
+        self.pool = nn.AvgPool2d(2)
+
+    def forward(self, x):
+        return self.pool(F.relu(self.norm(self.conv(x))))
+
+
+class _GNUpBlock(nn.Module):
+    def __init__(self, cin, cout, gn=8):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm = nn.GroupNorm(gn, cout)
+
+    def forward(self, x):
+        return F.relu(self.norm(self.conv(F.interpolate(x, scale_factor=2))))
+
+
+class _GNResBlock(nn.Module):
+    def __init__(self, c, gn=8):
+        super().__init__()
+        self.n1 = nn.GroupNorm(gn, c); self.c1 = nn.Conv2d(c, c, 3, padding=1)
+        self.n2 = nn.GroupNorm(gn, c); self.c2 = nn.Conv2d(c, c, 3, padding=1)
+
+    def forward(self, x):
+        out = self.c1(F.relu(self.n1(x)))
+        out = self.c2(F.relu(self.n2(out)))
+        return out + x
+
+
+class _HGEncoder(nn.Module):
+    def __init__(self, block_expansion, in_features, num_blocks, max_features, gn=8):
+        super().__init__()
+        blocks = []
+        for i in range(num_blocks):
+            cin = in_features if i == 0 else min(max_features, block_expansion * (2 ** i))
+            cout = min(max_features, block_expansion * (2 ** (i + 1)))
+            blocks.append(_GNDownBlock(cin, cout, gn=gn))
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, x):
+        outs = [x]
+        for b in self.blocks:
+            outs.append(b(outs[-1]))
+        return outs
+
+
+class _HGDecoder(nn.Module):
+    def __init__(self, block_expansion, in_features, num_blocks, max_features, gn=8):
+        super().__init__()
+        ups = []
+        for i in range(num_blocks)[::-1]:
+            cin = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))
+            cout = min(max_features, block_expansion * (2 ** i))
+            ups.append(_GNUpBlock(cin, cout, gn=gn))
+        self.ups = nn.ModuleList(ups)
+        self.out_filters = block_expansion + in_features
+
+    def forward(self, x):
+        out = x.pop()
+        for up in self.ups:
+            out = up(out)
+            out = torch.cat([out, x.pop()], dim=1)
+        return out
+
+
+class _Hourglass(nn.Module):
+    """Small GroupNorm U-Net used by the dense-motion mask predictor."""
+    def __init__(self, block_expansion, in_features, num_blocks, max_features, gn=8):
+        super().__init__()
+        self.enc = _HGEncoder(block_expansion, in_features, num_blocks, max_features, gn=gn)
+        self.dec = _HGDecoder(block_expansion, in_features, num_blocks, max_features, gn=gn)
+        self.out_filters = self.dec.out_filters
+
+    def forward(self, x):
+        return self.dec(self.enc(x))
+
+
+class _MeasuredWarpRenderer(nn.Module):
+    """Warp the fixed cond frame to the predicted pose (FOMM/MRAA similarity flow + Gao splat occlusion)."""
+
+    def __init__(self, num_kp, block_expansion=32, max_features=256, num_down_blocks=2,
+                 num_bottleneck_blocks=3, flow_res=64, mask_block_expansion=32,
+                 mask_num_blocks=4, mask_max_features=256,
+                 occ_lo=1e-3, occ_hi=2.0, scale_min=0.04, bg_resp=0.1, gn=8):
+        super().__init__()
+        self.num_kp = num_kp; self.flow_res = flow_res
+        self.occ_lo = occ_lo; self.occ_hi = occ_hi
+        self.scale_min = scale_min; self.bg_resp = bg_resp
+        in_feat = (num_kp + 1) * (3 + 1)
+        self.hourglass = _Hourglass(mask_block_expansion, in_feat, mask_num_blocks, mask_max_features, gn=gn)
+        self.mask = nn.Conv2d(self.hourglass.out_filters, num_kp + 1, 7, padding=3)
+        self.first = _GNSameBlock(3, block_expansion, kernel_size=7, padding=3, gn=gn)
+        down, up = [], []
+        for i in range(num_down_blocks):
+            down.append(_GNDownBlock(min(max_features, block_expansion * (2 ** i)),
+                                     min(max_features, block_expansion * (2 ** (i + 1))), gn=gn))
+        for i in range(num_down_blocks):
+            up.append(_GNUpBlock(min(max_features, block_expansion * (2 ** (num_down_blocks - i))),
+                                 min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1))), gn=gn))
+        self.down_blocks = nn.ModuleList(down)
+        self.up_blocks = nn.ModuleList(up)
+        self.bottleneck = nn.Sequential()
+        bc = min(max_features, block_expansion * (2 ** num_down_blocks))
+        for i in range(num_bottleneck_blocks):
+            self.bottleneck.add_module('r' + str(i), _GNResBlock(bc, gn=gn))
+        self.final = nn.Conv2d(block_expansion, 3, 7, padding=3)
+
+    # (1) per-region similarity backward flows: output(driving) grid -> source(cond) grid
+    def _sparse_backward(self, coords_t, coords_cond, res, device):
+        N = coords_t.shape[0]
+        grid = _make_coord_grid(res, res, device).view(1, 1, res, res, 2)
+        kp_t = coords_t[..., :2].reshape(N, self.num_kp, 1, 1, 2)
+        kp_c = coords_cond[..., :2].reshape(N, self.num_kp, 1, 1, 2)
+        s_t = coords_t[..., 2].clamp(min=self.scale_min).reshape(N, self.num_kp, 1, 1, 1)
+        s_c = coords_cond[..., 2].clamp(min=self.scale_min).reshape(N, self.num_kp, 1, 1, 1)
+        region = kp_c + (s_c / s_t) * (grid - kp_t)
+        bg = grid.expand(N, 1, res, res, 2)
+        return torch.cat([bg, region], dim=1)
+
+    def _heatmaps(self, coords_t, coords_cond, res, device):
+        gt = _region_gaussian(coords_t[..., :2], coords_t[..., 2].clamp(min=self.scale_min), res, device)
+        gc = _region_gaussian(coords_cond[..., :2], coords_cond[..., 2].clamp(min=self.scale_min), res, device)
+        hm = gt - gc
+        bg = torch.zeros(hm.shape[0], 1, res, res, device=device, dtype=hm.dtype)
+        return torch.cat([bg, hm], dim=1).unsqueeze(2)
+
+    # (2) dense flow = softmax-mask blend of the K+1 sparse flows (MRAA)
+    def _dense_flow(self, cond_img, coords_t, coords_cond):
+        N = coords_t.shape[0]; r = self.flow_res; device = cond_img.device
+        src = F.interpolate(cond_img, size=(r, r), mode='bilinear', align_corners=False)
+        sparse = self._sparse_backward(coords_t, coords_cond, r, device)
+        src_rep = src.unsqueeze(1).expand(N, self.num_kp + 1, 3, r, r).reshape(N * (self.num_kp + 1), 3, r, r)
+        deformed = F.grid_sample(src_rep, sparse.reshape(N * (self.num_kp + 1), r, r, 2),
+                                 align_corners=True, padding_mode='border')
+        deformed = deformed.view(N, self.num_kp + 1, 3, r, r)
+        hm = self._heatmaps(coords_t, coords_cond, r, device)
+        inp = torch.cat([hm, deformed], dim=2).reshape(N, (self.num_kp + 1) * 4, r, r)
+        mask = F.softmax(self.mask(self.hourglass(inp)), dim=1)
+        flow = (sparse.permute(0, 1, 4, 2, 3) * mask.unsqueeze(2)).sum(1)
+        return flow.permute(0, 2, 3, 1)
+
+    # (3) deterministic occlusion via forward splat (Gao), a measurement (no grad)
+    @torch.no_grad()
+    def _occlusion(self, coords_t, coords_cond):
+        N = coords_t.shape[0]; r = self.flow_res; device = coords_t.device
+        P = _make_coord_grid(r, r, device).view(1, 1, r, r, 2)
+        kp_t = coords_t[..., :2].reshape(N, self.num_kp, 1, 1, 2)
+        kp_c = coords_cond[..., :2].reshape(N, self.num_kp, 1, 1, 2)
+        s_t = coords_t[..., 2].clamp(min=self.scale_min).reshape(N, self.num_kp, 1, 1, 1)
+        s_c = coords_cond[..., 2].clamp(min=self.scale_min).reshape(N, self.num_kp, 1, 1, 1)
+        fwd = kp_t + (s_t / s_c) * (P - kp_c)
+        d2 = ((P - kp_c) ** 2).sum(-1)
+        w = torch.exp(-0.5 * d2 / (s_c.squeeze(-1) ** 2))
+        denom = self.bg_resp + w.sum(1, keepdim=True)
+        a_k = (w / denom).unsqueeze(-1)
+        a_bg = (self.bg_resp / denom).unsqueeze(-1)
+        Pxy = P.expand(N, 1, r, r, 2)
+        F_flow = (a_bg * Pxy + (a_k * fwd).sum(1, keepdim=True)).squeeze(1)
+        fx = (F_flow[..., 0] * 0.5 + 0.5) * (r - 1)
+        fy = (F_flow[..., 1] * 0.5 + 0.5) * (r - 1)
+        E = self._bilinear_splat(fx, fy, r, N)
+        m = ((E > self.occ_lo) & (E < self.occ_hi)).float().unsqueeze(1)
+        return m, E
+
+    @staticmethod
+    def _bilinear_splat(fx, fy, r, N):
+        device = fx.device
+        x0 = torch.floor(fx); y0 = torch.floor(fy)
+        wx = fx - x0; wy = fy - y0
+        x0 = x0.long(); y0 = y0.long(); x1 = x0 + 1; y1 = y0 + 1
+        E = torch.zeros(N, r * r, device=device)
+
+        def scat(xi, yi, wgt):
+            xi = xi.clamp(0, r - 1); yi = yi.clamp(0, r - 1)
+            E.scatter_add_(1, (yi * r + xi).reshape(N, -1), wgt.reshape(N, -1))
+        scat(x0, y0, (1 - wx) * (1 - wy)); scat(x1, y0, wx * (1 - wy))
+        scat(x0, y1, (1 - wx) * wy);       scat(x1, y1, wx * wy)
+        return E.view(N, r, r)
+
+    @staticmethod
+    def _deform(inp, flow):
+        _, ho, wo, _ = flow.shape
+        _, _, h, w = inp.shape
+        if ho != h or wo != w:
+            flow = F.interpolate(flow.permute(0, 3, 1, 2), size=(h, w), mode='bilinear',
+                                 align_corners=False).permute(0, 2, 3, 1)
+        return F.grid_sample(inp, flow, align_corners=True, padding_mode='border')
+
+    @staticmethod
+    def _gate(warped, prev, occ):
+        if occ.shape[2:] != warped.shape[2:]:
+            occ = F.interpolate(occ, size=warped.shape[2:], mode='bilinear', align_corners=False)
+        if prev is None:
+            return warped * occ
+        return warped * occ + prev * (1 - occ)
+
+    # (4) generator: MRAA skips + occlusion-gated warp + final source-pixel blend
+    def forward(self, cond_img, coords_t, coords_cond):
+        flow = self._dense_flow(cond_img, coords_t, coords_cond)
+        occ, _ = self._occlusion(coords_t, coords_cond)
+        out = self.first(cond_img)
+        skips = [out]
+        for db in self.down_blocks:
+            out = db(out); skips.append(out)
+        out = self._gate(self._deform(out, flow), None, occ)
+        out = self.bottleneck(out)
+        for i, ub in enumerate(self.up_blocks):
+            out = self._gate(self._deform(skips[-(i + 1)], flow), out, occ)
+            out = ub(out)
+        out = self._gate(self._deform(skips[0], flow), out, occ)
+        out = torch.sigmoid(self.final(out))
+        return self._gate(self._deform(cond_img, flow), out, occ)
+
+
+class _AntiAliasInterpolation2d(nn.Module):
+    def __init__(self, channels, scale):
+        super().__init__()
+        sigma = (1 / scale - 1) / 2
+        kernel_size = 2 * round(sigma * 4) + 1
+        self.ka = kernel_size // 2
+        self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka
+        kernel = 1
+        grids = torch.meshgrid([torch.arange(kernel_size, dtype=torch.float32)] * 2, indexing='ij')
+        for size, mgrid in zip([kernel_size, kernel_size], grids):
+            mean = (size - 1) / 2
+            kernel = kernel * torch.exp(-(mgrid - mean) ** 2 / (2 * sigma ** 2))
+        kernel = kernel / kernel.sum()
+        kernel = kernel.view(1, 1, *kernel.shape).repeat(channels, 1, 1, 1)
+        self.register_buffer('weight', kernel)
+        self.groups = channels; self.scale = scale
+        self.int_inv_scale = int(1 / scale)
+
+    def forward(self, x):
+        if self.scale == 1.0:
+            return x
+        out = F.pad(x, (self.ka, self.kb, self.ka, self.kb))
+        out = F.conv2d(out, weight=self.weight, groups=self.groups)
+        return out[:, :, ::self.int_inv_scale, ::self.int_inv_scale]
+
+
+class _ImagePyramide(nn.Module):
+    def __init__(self, scales, num_channels=3):
+        super().__init__()
+        self.scales = list(scales)
+        self.downs = nn.ModuleDict({str(s).replace('.', '-'): _AntiAliasInterpolation2d(num_channels, s)
+                                    for s in scales})
+
+    def forward(self, x):
+        return {'prediction_' + str(s): self.downs[str(s).replace('.', '-')](x) for s in self.scales}
+
+
+def _load_vgg19_features():
+    """Load ImageNet VGG-19 features from PHYWORLD_VGG19_PATH, else timm 'vgg19.tv_in1k' (== torchvision weights)."""
+    import os
+    from torchvision import models
+    net = models.vgg19(weights=None)
+    tried = []
+    path = os.environ.get('PHYWORLD_VGG19_PATH', '')
+    if path and os.path.exists(path):
+        try:
+            sd = torch.load(path, map_location='cpu', weights_only=False)
+            sd = {(k[9:] if k.startswith('features.') else k): v for k, v in sd.items()}
+            sd = {k: v for k, v in sd.items() if k in net.features.state_dict()}
+            net.features.load_state_dict(sd, strict=True)
+            return net.features
+        except Exception as e:
+            tried.append(f'path={e}')
+    try:
+        import timm
+        m = timm.create_model('vgg19.tv_in1k', pretrained=True)
+        sd = {k[9:]: v for k, v in m.state_dict().items() if k.startswith('features.')}
+        net.features.load_state_dict(sd, strict=True)
+        return net.features
+    except Exception as e:
+        tried.append(f'timm={e}')
+    raise RuntimeError('VGG19 features unavailable (' + ' | '.join(tried) + ')')
+
+
+class _Vgg19(nn.Module):
+    def __init__(self):
+        super().__init__()
+        f = _load_vgg19_features()
+        self.slices = nn.ModuleList()
+        for lo, hi in [(0, 2), (2, 7), (7, 12), (12, 21), (21, 30)]:
+            s = nn.Sequential()
+            for x in range(lo, hi):
+                s.add_module(str(x), f[x])
+            self.slices.append(s)
+        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+        for p in self.parameters():
+            p.requires_grad_(False)
+
+    def forward(self, x):
+        x = (x - self.mean) / self.std
+        outs = []
+        for s in self.slices:
+            x = s(x); outs.append(x)
+        return outs
+
+
+class PerceptualPyramidLoss(nn.Module):
+    """Multi-scale VGG-19 perceptual loss (FOMM); falls back to multi-scale pixel-L1 if VGG is unavailable."""
+    def __init__(self, scales=(1, 0.5, 0.25, 0.125), slice_weights=(1., 1., 1., 1., 1.)):
+        super().__init__()
+        self.scales = list(scales); self.slice_weights = slice_weights
+        self.pyramid = _ImagePyramide(self.scales, 3)
+        try:
+            self.vgg = _Vgg19(); self.use_vgg = True; self.note = 'vgg19-perceptual'
+        except Exception as e:
+            self.vgg = None; self.use_vgg = False; self.note = 'MS-L1-fallback:' + str(e)[:100]
+
+    def forward(self, pred, target):
+        pred = (pred.clamp(-1, 1) + 1) * 0.5
+        target = (target.clamp(-1, 1) + 1) * 0.5
+        pp = self.pyramid(pred); pt = self.pyramid(target)
+        total = pred.sum() * 0.0
+        for s in self.scales:
+            a = pp['prediction_' + str(s)]; b = pt['prediction_' + str(s)]
+            if self.use_vgg:
+                xv = self.vgg(a); yv = self.vgg(b)
+                for i, wgt in enumerate(self.slice_weights):
+                    total = total + wgt * (xv[i] - yv[i].detach()).abs().mean()
+            else:
+                total = total + (a - b.detach()).abs().mean()
+        return total
+
+
+class LDR(nn.Module):
+    """Encode each frame to a structured latent, roll it forward by kinematic integration, decode by warping the cond frame."""
+
+    def __init__(self, n_kp=16, num_pred=29, width=256, accel_scale=0.5, warp_flow_res=64, kappa_init=0.15):
+        super().__init__()
+        self.n_kp = n_kp; self.num_pred = num_pred; self.accel_scale = accel_scale
+        self.cdim = 3
+        # kappa (softplus of log_kappa): uncertainty gate for the measured second-order init
+        self.log_kappa = nn.Parameter(torch.log(torch.expm1(torch.tensor(max(float(kappa_init), 1e-3)))))
+        self.enc = _KeypointEnc(n_kp)
+        d = n_kp * self.cdim
+        self.g = nn.Sequential(nn.Linear(2 * d, width), nn.SiLU(),
+                               nn.Linear(width, width), nn.SiLU(),
+                               nn.Linear(width, d))
+        nn.init.zeros_(self.g[-1].weight); nn.init.zeros_(self.g[-1].bias)  # zero-init: rollout starts as pure inertia
+        self.warp = _MeasuredWarpRenderer(num_kp=n_kp, flow_res=warp_flow_res)
+
+    def _residual(self, s, v):
+        return self.accel_scale * torch.tanh(self.g(torch.cat([s, v], dim=1)))
+
+    def rollout(self, c_cond, n=None):
+        n = self.num_pred if n is None else n
+        B = c_cond.shape[0]
+        assert c_cond.shape[1] >= 3, 'kinematic initialization needs 3 conditioning latents'
+        s = c_cond[:, -1].reshape(B, -1)
+        v = (c_cond[:, -1] - c_cond[:, -2]).reshape(B, -1)
+        a0_raw = (c_cond[:, -1] - 2 * c_cond[:, -2] + c_cond[:, -3]).reshape(B, -1)
+        kappa = F.softplus(self.log_kappa)
+        a0 = (a0_raw * a0_raw) / (a0_raw * a0_raw + kappa * kappa) * a0_raw
+        out = []
+        for _ in range(n):
+            v = v + a0 + self._residual(s, v)
+            s = s + v
+            out.append(s)
+        return torch.stack(out, 1).view(B, n, self.n_kp, self.cdim)
+
+    def _enc_seq(self, frames):
+        B, L = frames.shape[:2]
+        return self.enc(frames.reshape(B * L, *frames.shape[2:])).view(B, L, self.n_kp, self.cdim)
+
+    def _decode_seq(self, cond_img, coords_seq, coords_cond):
+        B, T = coords_seq.shape[:2]; H, W = cond_img.shape[-2:]
+        src01 = ((cond_img.clamp(-1, 1) + 1) * 0.5).unsqueeze(1).expand(B, T, 3, H, W).reshape(B * T, 3, H, W)
+        ct = coords_seq.reshape(B * T, self.n_kp, self.cdim)
+        cc = coords_cond.unsqueeze(1).expand(B, T, self.n_kp, self.cdim).reshape(B * T, self.n_kp, self.cdim)
+        out01 = self.warp(src01, ct, cc)
+        return (out01 * 2 - 1).view(B, T, 3, H, W)
+
+    def forward(self, frames, nc, full=False, horizon=None, cond_img=None):
+        coords = self._enc_seq(frames)
+        coords_cond = coords[:, nc - 1]
+        if not full:
+            return self._decode_seq(cond_img, self.rollout(coords[:, :nc]), coords_cond)
+        dec_ae = self._decode_seq(cond_img, coords, coords_cond)
+        roll = self.rollout(coords[:, :nc], horizon)
+        dec_roll = self._decode_seq(cond_img, roll, coords_cond)
+        return dec_roll, dec_ae, roll, coords
+
+
+def build_ldr(n_kp=16, num_pred=29, width=256, accel_scale=0.5, warp_flow_res=64, kappa_init=0.15):
+    return LDR(n_kp=n_kp, num_pred=num_pred, width=width, accel_scale=accel_scale,
+               warp_flow_res=warp_flow_res, kappa_init=kappa_init)

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+numpy
+imageio
+imageio-ffmpeg
+pillow