RoboTryOn demo

.gitattributes

@@ -0,0 +1,46 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/gr1_stir_pot/episode/video_ego_view.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/gr1_stir_pot/input.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/openarm_box_convey/episode/mask_left.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/openarm_box_convey/episode/object_mask_left.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/openarm_box_convey/episode/video_left.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/openarm_box_convey/input.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/openarm_box_convey/reference.png filter=lfs diff=lfs merge=lfs -text
+examples/rby1_bottle_pour/episode/mask_left.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/rby1_bottle_pour/episode/video_left.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/rby1_bottle_pour/input.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/rby1_bottle_pour/reference.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+prep_assets.py
+assets/

README.md

@@ -0,0 +1,16 @@
+---
+title: RoboTryOn
+emoji: 🤖
+colorFrom: gray
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.50.0
+python_version: "3.10"
+app_file: app.py
+startup_duration_timeout: 1h
+pinned: false
+short_description: Human videos to dexterous robot videos via virtual try-on
+---
+
+Demo of [RoboTryOn](https://github.com/ryan-seungyong-lee/RoboTryOn) ([paper](https://arxiv.org/abs/2611.05163)).
+`robotryon/`, `scripts/infer_video.py` and `configs/` are copies of the code repository.

app.py

@@ -0,0 +1,135 @@
+import os
+import tempfile
+
+import gradio as gr
+import spaces
+
+from core import HERE, Demo
+
+DEMO = Demo(load_models=os.environ.get("ROBOTRYON_UI_ONLY") != "1")
+CLIPS = DEMO.manifest["clips"]
+EMB_KEY = {e["name"]: e["key"] for e in DEMO.manifest["embodiments"]}
+EMB_NAME = {v: k for k, v in EMB_KEY.items()}
+CLIP_LABEL = {c["key"]: f'{c["name"]} ({c["source"]})' for c in CLIPS}
+CLIP_KEY = {v: k for k, v in CLIP_LABEL.items()}
+CODE = "https://github.com/ryan-seungyong-lee/RoboTryOn"
+
+
+def _p(rel):
+    return str(HERE / rel)
+
+
+def _sel(clip_label, emb_name):
+    return CLIP_KEY[clip_label], EMB_KEY[emb_name]
+
+
+def stage1_duration(clip_label, emb_name, *_):
+    return 20 + 25 * DEMO.episode(*_sel(clip_label, emb_name))["num_keyframes"]   # A100: ~34 s per keyframe
+
+
+def stage2_duration(clip_label, emb_name, *_):
+    return 40 + 60 * DEMO.episode(*_sel(clip_label, emb_name))["num_clips"]       # A100: ~60 s per clip
+
+
+@spaces.GPU(duration=stage1_duration)
+def _stage1(clip_label, emb_name, progress=gr.Progress(track_tqdm=True)):
+    clip, emb = _sel(clip_label, emb_name)
+    kf_dir, _ = DEMO.translate(clip, emb, tempfile.mkdtemp(prefix="keyframes_"))
+    return str(kf_dir / "grid.jpg"), (clip, emb, str(kf_dir))
+
+
+@spaces.GPU(duration=stage2_duration, size="xlarge")   # LTX-2.5 needs more than a half H200
+def _stage2(clip_label, emb_name, kf_dir, progress=gr.Progress(track_tqdm=True)):
+    clip, emb = _sel(clip_label, emb_name)
+    return str(DEMO.synthesize(clip, emb, kf_dir, tempfile.mkdtemp(prefix="video_")))
+
+
+def stage1(clip_label, emb_name):
+    return _stage1(clip_label, emb_name)
+
+
+def stage2(clip_label, emb_name, state):
+    if not state or tuple(state[:2]) != _sel(clip_label, emb_name):
+        raise gr.Error("Run stage 1 (keyframe translation) for this video and embodiment first.")
+    return _stage2(clip_label, emb_name, state[2])
+
+
+def pick_clip(evt: gr.SelectData):
+    return CLIP_LABEL[CLIPS[evt.index]["key"]]
+
+
+def show_clip(clip_label, emb_name, state):
+    clip = DEMO.clips[CLIP_KEY[clip_label]]
+    keep = state and tuple(state[:2]) == _sel(clip_label, emb_name)
+    return (_p(clip["video"]), clip["caption"],
+            gr.update(selected_index=[c["key"] for c in CLIPS].index(clip["key"])),
+            *((gr.update(), gr.update()) if keep else (None, None)))
+
+
+def show_embodiment(clip_label, emb_name, state):
+    keep = state and tuple(state[:2]) == _sel(clip_label, emb_name)
+    return (_p(DEMO.embodiments[EMB_KEY[emb_name]]["reference"]),
+            *((gr.update(), gr.update()) if keep else (None, None)))
+
+
+def load_example(clip_label, emb_name):
+    clip, emb = _sel(clip_label, emb_name)
+    ex = next(e for e in DEMO.manifest["examples"] if e["clip"] == clip and e["embodiment"] == emb)
+    return _p(ex["grid"]), _p(ex["video"]), (clip, emb, _p(ex["keyframes"]))
+
+
+CSS = """
+#clips .thumbnail-item { aspect-ratio: 16 / 9 !important; }
+#clips .thumbnail-item img { object-fit: contain !important; }
+#clips .caption-label { font-size: 12px; }
+"""
+
+first = CLIPS[0]
+with gr.Blocks(title="RoboTryOn", css=CSS) as demo:
+    gr.Markdown(
+        "# RoboTryOn\n"
+        "**Synthesizing Dexterous Robot Demonstrations from Human Videos via Virtual Try-On** · "
+        f"[Paper](https://arxiv.org/abs/2611.05163) · "
+        f"[Project page](https://ryan-seungyong-lee.github.io/RoboTryOn) · [Code]({CODE})\n\n"
+        "The input videos are preprocessed in advance (hand reconstruction, retargeting to each "
+        "robot hand, arm masks) due to the time limits of this Space. "
+        f"To run RoboTryOn on your own videos, see the [code]({CODE}).")
+    state = gr.State(None)
+
+    gr.Markdown("### 1. Human video")
+    gallery = gr.Gallery([(_p(c["thumb"]), c["name"]) for c in CLIPS], columns=9,
+                         rows=1, height=120, object_fit="contain", allow_preview=False,
+                         selected_index=0, show_label=False, elem_id="clips")
+    with gr.Row():
+        human = gr.Video(_p(first["video"]), label="Human video", interactive=False, height=220,
+                         scale=1)
+        with gr.Column(scale=2):
+            clip_dd = gr.Dropdown(list(CLIP_LABEL.values()), value=CLIP_LABEL[first["key"]],
+                                  label="Selected video")
+            caption = gr.Textbox(first["caption"], label="Task", interactive=False, lines=3)
+
+    gr.Markdown("### 2. Embodiment")
+    with gr.Row():
+        emb_radio = gr.Radio(list(EMB_KEY), value=next(iter(EMB_KEY)), label="Target embodiment")
+        reference = gr.Image(_p(DEMO.manifest["embodiments"][0]["reference"]),
+                             label="Reference", interactive=False, height=180)
+
+    gr.Markdown("### 3. Generate")
+    with gr.Row():
+        b1 = gr.Button("Stage 1: Keyframe translation", variant="primary")
+        b2 = gr.Button("Stage 2: Video synthesis (uses the stage-1 keyframes)", variant="primary")
+    grid = gr.Image(label="Keyframes: human (top) and robot (bottom)", type="filepath")
+    video = gr.Video(label="Robot video")
+
+    gr.Markdown("### Examples (precomputed)")
+    gr.Examples([[CLIP_LABEL[e["clip"]], EMB_NAME[e["embodiment"]]] for e in DEMO.manifest["examples"]],
+                inputs=[clip_dd, emb_radio], outputs=[grid, video, state], fn=load_example,
+                run_on_click=True, cache_examples=False)
+
+    gallery.select(pick_clip, None, clip_dd)
+    clip_dd.change(show_clip, [clip_dd, emb_radio, state], [human, caption, gallery, grid, video])
+    emb_radio.change(show_embodiment, [clip_dd, emb_radio, state], [reference, grid, video])
+    b1.click(stage1, [clip_dd, emb_radio], [grid, state])
+    b2.click(stage2, [clip_dd, emb_radio, state], video)
+
+demo.queue().launch()

configs/image/infer.yaml

@@ -0,0 +1,32 @@
+# Autoregressive keyframe translation (paper Eq. 3.2). See scripts/infer_keyframes.py.
+
+model:
+  base: black-forest-labs/FLUX.2-klein-base-9B
+checkpoint: null             # trained checkpoint dir (FFT: has transformer/; LoRA: has pytorch_lora_weights.safetensors)
+
+episodes: []                 # preprocessed human episode directories (with keyframes.json)
+reference: {}                # {view: robot-only reference image}; a single path is used for every view
+output_dir: runs/keyframes
+
+views: auto                  # auto (joint stereo when the episode has two views) | stereo | mono
+ref_skeleton: false          # false: reuse the previous keyframe's latent as reference tokens
+                             # true: decode it, paint the current skeleton, re-encode
+steps: 30
+guidance_scale: 3.5          # true CFG against the empty prompt
+seed: 2024                   # keyframe t uses seed + t
+
+image:
+  long_edge: 1024
+  aspect: [16, 9]
+  crop_anchor: bottom
+
+masks:                       # inference dilation = (dilate_min + dilate_max) // 2 + bonus = 58 px
+  dilate_min: 3
+  dilate_max: 71
+  bonus: 21
+  obj_erode: 9               # px, erode the object mask before subtracting it from the hole
+
+caption:                     # per-keyframe hand-state prompt (robotryon.captions)
+  cuts:                      # closure (rad) / dorsumness (deg) edges -> open|half-closed|closed fist, palm|edge-on|back
+    closure: [0.95, 1.15]
+    facing: [49.0, 121.0]

configs/video/infer.yaml

@@ -0,0 +1,35 @@
+# Video synthesis from translated robot keyframes and the human skeleton flow (paper App. B.4).
+model:
+  base: Lightricks/LTX-2.5-Diffusers
+  transformer_subfolder: transformer_full
+  lora: outputs/video/stage2/checkpoint-001000   # directory with pytorch_lora_weights.safetensors
+  attention_backend: native
+
+input:
+  episode: null            # human episode dir: meta.json, keypoints_<view>.npz, keyframes.json
+  keyframes_dir: null      # image-model output for this episode: keyframe_<view>_<frame:06d>.png + keyframes.json
+  views: null              # default: all views in meta.json
+  stereo: true             # generate two views jointly in one sequence; false = one view at a time
+
+output:
+  dir: outputs/video/infer
+
+text:
+  template: "{caption}"
+  embodiment: null         # fills {embodiment} in the template (target robot)
+
+fps: 20
+resolution: [1024, 576]    # output, per view
+skeleton_size: [512, 288]  # as in training: reference downscale 1 at 512x288, 2 at 1024x576
+mask_human: true           # true: skeleton flow on black | false: skeleton drawn on the unmasked human video
+frames_range: [49, 145]    # clip length limits (8k+1)
+seed: 0                    # clip i uses seed + i
+
+two_stage:
+  mode: same_res           # same_res: 8 steps + 3 refine steps at 1024x576
+                           # upsample: 8 steps at 512x288 -> latent upsampler x2 -> 3 steps at 1024x576
+
+skeleton:
+  draw:
+    thickness: 2
+    radius: 3

core.py

@@ -0,0 +1,194 @@
+"""The two demo stages, run with the RoboTryOn code as released (``robotryon/``,
+``scripts/infer_video.py``, ``configs/``).
+
+The human clips are preprocessed in advance for every embodiment (hand reconstruction,
+retargeting to the robot hand, arm masks, keyframe selection) and their prompts are encoded in
+advance, so neither the embodiment models nor the text encoders are needed here.
+"""
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import subprocess
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image, ImageDraw, ImageFont
+
+from robotryon.common.config import load_config
+from robotryon.image.ar_keyframes import KeyframeTranslator
+from robotryon.image.episodes import Episode, FrameReader, load_view_frame
+from robotryon.image.model import load_pipeline as load_image_pipeline
+
+HERE = Path(__file__).resolve().parent
+ASSETS = HERE / "assets"
+_spec = importlib.util.spec_from_file_location("infer_video", HERE / "scripts" / "infer_video.py")
+iv = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(iv)
+
+CKPT_REPO = "RyanL22/RoboTryOn"
+ASSETS_REPO = "RyanL22/RoboTryOn-demo-assets"     # dataset: preprocessed clips, encoded prompts, examples
+FLUX_REPO = "black-forest-labs/FLUX.2-klein-base-9B"
+LTX_REPO = "Lightricks/LTX-2.5-Diffusers"
+# base-model revisions the demo was tested with (later LTX revisions add components)
+FLUX_REVISION = "32773329fbe7e81a90ef971740e8ba4b0364ecf3"
+LTX_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
+FLUX_FILES = ["model_index.json", "scheduler/*", "transformer/*", "vae/*"]
+LTX_FILES = ["model_index.json", "scheduler/*", "transformer_full/*", "vae/*", "audio_vae/*",
+             "vocoder/*", "connectors/*", "diffusion_decoder/*", "duration_head/*",
+             "ltx-2.5-22b-distilled-lora-450-bf16.safetensors"]
+IMAGE_CONFIG = HERE / "configs" / "image" / "infer.yaml"
+VIDEO_CONFIG = HERE / "configs" / "video" / "infer.yaml"
+
+
+def load_manifest() -> dict:
+    if not (ASSETS / "manifest.json").exists():
+        snapshot_download(ASSETS_REPO, repo_type="dataset", local_dir=ASSETS)
+    return json.loads((ASSETS / "manifest.json").read_text())
+
+
+def _dir(env: str, repo: str, files: list[str] | None, revision: str | None = None) -> str:
+    return os.environ.get(env) or snapshot_download(repo, allow_patterns=files, revision=revision)
+
+
+def human_frames(ep: Episode, cfg, frames: list[int], view: str | None = None) -> list[np.ndarray]:
+    """The human frames at ``frames``, cropped and resized as the image model sees them."""
+    c, reader = cfg.image, FrameReader()
+    return [load_view_frame(reader, ep, view or ep.views[0], t, int(c.long_edge),
+                            tuple(c.aspect) if c.get("aspect") else None,
+                            c.get("crop_anchor", "bottom")).image for t in frames]
+
+
+def keyframe_grid(human: list[np.ndarray], robot: list[np.ndarray], frames: list[int],
+                  cols: int = 4, width: int = 384) -> Image.Image:
+    """Human keyframes above their translations, ``cols`` keyframes per row pair."""
+    h = round(width * human[0].shape[0] / human[0].shape[1])
+    rows = (len(frames) + cols - 1) // cols
+    n = min(cols, len(frames))
+    pad = 4
+    grid = Image.new("RGB", (n * width + (n + 1) * pad, rows * (2 * h + pad) + (rows + 1) * pad),
+                     (255, 255, 255))
+    draw = ImageDraw.Draw(grid)
+    font = ImageFont.load_default(size=18)
+    for i, (a, b, t) in enumerate(zip(human, robot, frames)):
+        x = pad + (i % cols) * (width + pad)
+        y = pad + (i // cols) * (2 * h + 2 * pad)
+        grid.paste(Image.fromarray(a).resize((width, h), Image.LANCZOS), (x, y))
+        grid.paste(Image.fromarray(b).resize((width, h), Image.LANCZOS), (x, y + h + pad))
+        draw.text((x + 6, y + 4), f"frame {t}", fill=(255, 255, 255), font=font,
+                  stroke_width=2, stroke_fill=(0, 0, 0))
+    return grid
+
+
+def match_brightness(image: np.ndarray, human: np.ndarray, mask: np.ndarray,
+                     margin: int = 15) -> np.ndarray:
+    """Match the keyframe's lightness (mean and spread of L in Lab) to the human frame.
+
+    Only the background outside the human-arm mask (grown by ``margin`` px) is compared, so the
+    robot itself does not drive the statistics; the same correction is applied to every pixel.
+    """
+    if human.shape[:2] != image.shape[:2]:
+        human = cv2.resize(human, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_AREA)
+    if mask.shape[:2] != image.shape[:2]:
+        mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+    k = np.ones((2 * margin + 1, 2 * margin + 1), np.uint8)
+    bg = cv2.dilate((mask > 0).astype(np.uint8), k) == 0
+    if bg.sum() < 1000:
+        return image
+    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB).astype(np.float32)
+    ref = cv2.cvtColor(human, cv2.COLOR_RGB2LAB).astype(np.float32)
+    src_l, ref_l = lab[..., 0][bg], ref[..., 0][bg]
+    gain = ref_l.std() / max(src_l.std(), 1e-3)
+    lab[..., 0] = (lab[..., 0] - src_l.mean()) * gain + ref_l.mean()
+    return cv2.cvtColor(np.clip(lab, 0, 255).astype(np.uint8), cv2.COLOR_LAB2RGB)
+
+
+def h264(src: Path, dst: Path) -> Path:
+    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(src), "-c:v", "libx264",
+                    "-crf", "18", "-pix_fmt", "yuv420p", str(dst)], check=True)
+    return dst
+
+
+class Demo:
+    """Both models, kept on the CPU between calls and moved to the GPU for one stage."""
+
+    def __init__(self, load_models: bool = True):
+        self.manifest = load_manifest()
+        self.clips = {c["key"]: c for c in self.manifest["clips"]}
+        self.embodiments = {e["key"]: e for e in self.manifest["embodiments"]}
+        self.episodes = self.manifest["episodes"]
+        if not load_models:
+            return
+        self.prompts = torch.load(ASSETS / "prompts.pt")
+        self.negative = torch.load(ASSETS / "negative_prompt.pt")
+        self.ckpt = Path(_dir("ROBOTRYON_CKPT", CKPT_REPO, ["image/*", "video/*"]))
+        self.flux = _dir("FLUX_DIR", FLUX_REPO, FLUX_FILES, FLUX_REVISION)
+        self.ltx = _dir("LTX_DIR", LTX_REPO, LTX_FILES, LTX_REVISION)
+        self.image = load_image_pipeline(self.flux, self.ckpt / "image", device="cpu",
+                                         text_encoder=False)
+        self.video = iv.load_pipeline(self.video_config([]), views=1, device="cpu")
+
+    def episode(self, clip: str, emb: str) -> dict:
+        return self.episodes[f"{emb}/{clip}"]
+
+    def image_config(self, overrides: list[str]):
+        return load_config(IMAGE_CONFIG, [f"model.base={self.flux}",
+                                          f"checkpoint={self.ckpt / 'image'}", *overrides])
+
+    def video_config(self, overrides: list[str]):
+        return load_config(VIDEO_CONFIG, [f"model.base={self.ltx}",
+                                          f"model.lora={self.ckpt / 'video'}", *overrides])
+
+    def translate(self, clip: str, emb: str, out_root: str | Path) -> tuple[Path, Image.Image]:
+        """Stage 1: keyframe translation. Returns the keyframe directory and the grid."""
+        e = self.episode(clip, emb)
+        cfg = self.image_config([f"output_dir={out_root}", *e["image_overrides"]])
+        ep_dir = HERE / e["dir"]
+        ep = Episode.load(ep_dir)
+        ref = np.asarray(Image.open(HERE / self.embodiments[emb]["reference"]).convert("RGB"))
+        self.image.to("cuda")
+        try:
+            tr = KeyframeTranslator(self.image, cfg)
+            tr._prompts.update({k: tuple(x.to("cuda") for x in self.prompts[k])
+                                for k in e["prompts"]})
+            kf_dir = tr.translate_episode(ep_dir, {v: ref for v in ep.views}, out_root)
+        finally:
+            self.image.to("cpu")
+            torch.cuda.empty_cache()
+        frames = json.loads((kf_dir / "keyframes.json").read_text())["indices"]
+        if self.clips[clip]["source"] != "self-recorded":
+            # downloaded human videos (VITRA, EgoDex): the generated keyframes drift in overall
+            # brightness, so match each one to its human frame before stage 2 reads it
+            for v in ep.views:
+                human = human_frames(ep, cfg, frames, v)
+                for t, h in zip(frames, human):
+                    f = kf_dir / f"keyframe_{v}_{t:06d}.png"
+                    mask = np.asarray(Image.open(kf_dir / f"mask_{v}_{t:06d}.png").convert("L"))
+                    img = np.asarray(Image.open(f).convert("RGB"))
+                    Image.fromarray(match_brightness(img, h, mask)).save(f)
+        view = ep.views[0]
+        robot = [np.asarray(Image.open(kf_dir / f"keyframe_{view}_{t:06d}.png").convert("RGB"))
+                 for t in frames]
+        grid = keyframe_grid(human_frames(ep, cfg, frames), robot, frames)
+        grid.save(kf_dir / "grid.jpg", quality=92)
+        return kf_dir, grid
+
+    def synthesize(self, clip: str, emb: str, kf_dir: str | Path, out_dir: str | Path) -> Path:
+        """Stage 2: video synthesis from the stage-1 keyframes. Returns an h264 mp4."""
+        e = self.episode(clip, emb)
+        cfg = self.video_config([f"input.episode={HERE / e['dir']}",
+                                 f"input.keyframes_dir={kf_dir}", f"output.dir={out_dir}"])
+        pe, pm = torch.load(HERE / self.clips[clip]["caption_embeds"])
+        text = (pe, pm, *self.negative)
+        self.video.to("cuda")
+        try:
+            out = iv.synthesize(cfg, self.video, text)
+        finally:
+            self.video.to("cpu")
+            torch.cuda.empty_cache()
+        view = json.loads((out / "meta.json").read_text())["views"][0]
+        return h264(out / f"video_{view}.mp4", out / "robot.mp4")

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,14 @@
+torch==2.8.0
+torchvision==0.23.0
+diffusers @ git+https://github.com/huggingface/diffusers@90b4e34e79a86ec5e7f2437634fe95ecd2108796
+transformers==5.15.0
+peft==0.20.0
+accelerate==1.14.0
+safetensors
+huggingface_hub
+numpy<2
+scipy
+opencv-python-headless
+pillow
+pyyaml
+tqdm

robotryon/captions.py

@@ -0,0 +1,146 @@
+"""Text prompts: per-hand state words and the embodiment phrase.
+
+Image model. A keyframe's prompt states two facts per hand that a 2D skeleton conveys weakly at
+the sizes hands occupy in the frame -- how closed the hand is and which side faces the camera::
+
+    Match the hand pose to the drawn skeleton.
+    Left hand: open, palm toward camera. Right hand: closed fist, back of hand toward camera.
+
+``closure`` is the MCP + PIP flexion averaged over the four fingers (radians);
+``dorsumness`` is 180 minus the angle between the knuckle-plane normal and the ray from the hand
+to the camera (degrees, high = back of the hand toward the lens). Both are read from the
+3D hand in the camera frame -- the reconstructed human hand, or the robot hand from forward
+kinematics, with the same formula -- and cut into words (``HandStateCuts``). A hand that is not
+tracked is ``not visible``.
+
+Each episode stores both values per view and frame in ``hand_state_<view>.npz`` (``closure``,
+``dorsumness``, ``(T, 2)`` float32, NaN where a hand is not tracked), written by the preprocessing
+for human clips (reconstructed hands) and robot episodes (forward kinematics), so training and
+inference read the same quantity.
+
+Video model. The prompt is the task caption (``text.template``).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+INSTRUCTION = "Match the hand pose to the drawn skeleton."
+CLOSURE = ("open", "half-closed", "closed fist")
+FACING = ("palm toward camera", "edge-on", "back of hand toward camera")
+NOT_VISIBLE = "not visible"
+
+# MediaPipe / MANO 21-joint layout: (MCP, PIP, DIP, tip) per finger.
+_CHAINS = {"index": (5, 6, 7, 8), "middle": (9, 10, 11, 12), "ring": (13, 14, 15, 16),
+           "little": (17, 18, 19, 20)}
+
+
+def _angle(a, b, c) -> float:
+    """Flexion at b: 0 when a-b-c is straight."""
+    u, v = a - b, c - b
+    u = u / max(np.linalg.norm(u), 1e-9)
+    v = v / max(np.linalg.norm(v), 1e-9)
+    return float(np.pi - np.arccos(np.clip(np.dot(u, v), -1.0, 1.0)))
+
+
+def hand_fields(j21: np.ndarray, hand: str) -> tuple[float, float]:
+    """``(21, 3)`` joints of one hand in the camera frame (camera at the origin, looking down +z)
+    -> ``(closure, dorsumness)``; NaN when the hand is not finite."""
+    j = np.asarray(j21, np.float64)
+    if not np.isfinite(j).all():
+        return float("nan"), float("nan")
+    w = j[0]
+    closure = np.mean([_angle(w, j[a], j[b]) + _angle(j[a], j[b], j[c])
+                       for a, b, c, _ in _CHAINS.values()])
+    n = np.cross(j[5] - w, j[17] - w)
+    ray = -j.mean(0)                                   # hand -> camera
+    nn, nr = np.linalg.norm(n), np.linalg.norm(ray)
+    if nn < 1e-9 or nr < 1e-9:
+        return float(closure), float("nan")
+    d = 180.0 - float(np.degrees(np.arccos(np.clip(np.dot(n, ray) / (nn * nr), -1.0, 1.0))))
+    # the normal leaves through the palm on the right hand and the back on the left
+    return float(closure), (d if hand == "left" else 180.0 - d)
+
+
+@dataclass
+class HandStateCuts:
+    """Edges that turn ``closure`` / ``dorsumness`` into words (lower, upper)."""
+    closure: tuple[float, float] = (0.95, 1.15)
+    facing: tuple[float, float] = (49.0, 121.0)
+
+    @classmethod
+    def from_config(cls, cfg) -> "HandStateCuts":
+        cfg = cfg or {}
+        return cls(tuple(float(x) for x in cfg.get("closure", cls.closure)),
+                   tuple(float(x) for x in cfg.get("facing", cls.facing)))
+
+    def words(self, closure: float, dorsumness: float) -> str:
+        if not (np.isfinite(closure) and np.isfinite(dorsumness)):
+            return NOT_VISIBLE
+        ci = int(closure >= self.closure[0]) + int(closure >= self.closure[1])
+        fi = int(dorsumness >= self.facing[0]) + int(dorsumness >= self.facing[1])
+        return f"{CLOSURE[ci]}, {FACING[fi]}"
+
+
+def hand_state_prompt(left: str, right: str) -> str:
+    """The image prompt from the two hands' words."""
+    return f"{INSTRUCTION} Left hand: {left}. Right hand: {right}."
+
+
+def all_prompts() -> list[str]:
+    """Every prompt training can produce: each pair of hand descriptions, the hand-dropout prompt
+    (the instruction alone) and the empty (unconditional) prompt."""
+    cells = [f"{c}, {f}" for c in CLOSURE for f in FACING] + [NOT_VISIBLE]
+    return [hand_state_prompt(a, b) for a in cells for b in cells] + [INSTRUCTION, ""]
+
+
+def sample_prompt(words: tuple[str, str], rng: np.random.Generator, p_uncond: float,
+                  p_drop: float) -> str:
+    """Training prompt: empty with ``p_uncond``, hand words dropped with ``p_drop``."""
+    u = float(rng.random())
+    if u < p_uncond:
+        return ""
+    if u < p_uncond + p_drop:
+        return INSTRUCTION
+    return hand_state_prompt(*words)
+
+
+def hand_state_arrays(joints3d: np.ndarray, valid: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """``(T, 2, 21|17, 3)`` camera-frame joints and ``(T, 2)`` validity -> ``closure`` and
+    ``dorsumness`` ``(T, 2)`` float32 (NaN where invalid)."""
+    from .embodiment.retarget import to21
+    j = to21(np.asarray(joints3d, np.float64))
+    T = j.shape[0]
+    clo = np.full((T, 2), np.nan, np.float32)
+    dor = np.full((T, 2), np.nan, np.float32)
+    for t in range(T):
+        for h, name in enumerate(("left", "right")):
+            if valid[t, h]:
+                clo[t, h], dor[t, h] = hand_fields(j[t, h], name)
+    return clo, dor
+
+
+def save_hand_state(path, closure: np.ndarray, dorsumness: np.ndarray) -> None:
+    np.savez_compressed(path, closure=closure.astype(np.float32), dorsumness=dorsumness.astype(np.float32))
+
+
+def load_hand_state(path) -> tuple[np.ndarray, np.ndarray]:
+    z = np.load(path)
+    return z["closure"], z["dorsumness"]
+
+
+def words_at(states: list[tuple[np.ndarray, np.ndarray]], t: int,
+             cuts: HandStateCuts) -> tuple[str, str]:
+    """Words of both hands at frame ``t`` from one or more views' hand states (averaged)."""
+    out = []
+    for h in range(2):
+        vals = [(c[t, h], d[t, h]) for c, d in states
+                if t < len(c) and np.isfinite(c[t, h]) and np.isfinite(d[t, h])]
+        if not vals:
+            out.append(NOT_VISIBLE)
+            continue
+        cl, dn = np.mean(np.asarray(vals, np.float64), axis=0)
+        out.append(cuts.words(float(cl), float(dn)))
+    return out[0], out[1]

robotryon/common/config.py

@@ -0,0 +1,115 @@
+"""YAML configs with dotted command-line overrides.
+
+Every script takes ``--config path.yaml`` followed by any number of ``key.sub=value``
+overrides, e.g.::
+
+    python scripts/train_image.py --config configs/image/klein9b_lora.yaml \
+        train.max_steps=20 data.root=/path/to/data
+
+A config may name a parent with ``_base_: other.yaml`` (relative to itself); the child is merged
+on top of it. Values are parsed with YAML rules, so ``1e-5``, ``true``, ``[1, 2]`` and ``null``
+work as expected on the command line.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+from pathlib import Path
+from typing import Any, Sequence
+
+import yaml
+
+
+class Config(dict):
+    """A dict whose keys are also attributes (recursively)."""
+
+    def __getattr__(self, key: str) -> Any:
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(key) from e
+
+    def __setattr__(self, key: str, value: Any) -> None:
+        self[key] = value
+
+    def get_path(self, dotted: str, default: Any = None) -> Any:
+        node: Any = self
+        for k in dotted.split("."):
+            if not isinstance(node, dict) or k not in node:
+                return default
+            node = node[k]
+        return node
+
+    def to_dict(self) -> dict:
+        return _plain(self)
+
+
+def _wrap(x: Any) -> Any:
+    if isinstance(x, dict):
+        return Config({k: _wrap(v) for k, v in x.items()})
+    if isinstance(x, list):
+        return [_wrap(v) for v in x]
+    return x
+
+
+def _plain(x: Any) -> Any:
+    if isinstance(x, dict):
+        return {k: _plain(v) for k, v in x.items()}
+    if isinstance(x, list):
+        return [_plain(v) for v in x]
+    return x
+
+
+def merge(base: dict, over: dict) -> dict:
+    out = copy.deepcopy(base)
+    for k, v in over.items():
+        if isinstance(v, dict) and isinstance(out.get(k), dict):
+            out[k] = merge(out[k], v)
+        else:
+            out[k] = copy.deepcopy(v)
+    return out
+
+
+def _set_dotted(d: dict, dotted: str, value: Any) -> None:
+    keys = dotted.split(".")
+    for k in keys[:-1]:
+        d = d.setdefault(k, {})
+    d[keys[-1]] = value
+
+
+def _parse_value(raw: str) -> Any:
+    v = yaml.safe_load(raw)
+    if isinstance(v, str):          # YAML 1.1 reads "1e-5" as a string; accept it as a float
+        try:
+            return float(v)
+        except ValueError:
+            pass
+    return v
+
+
+def load_config(path: str | Path, overrides: Sequence[str] = ()) -> Config:
+    path = Path(path)
+    cfg = yaml.safe_load(path.read_text()) or {}
+    base = cfg.pop("_base_", None)
+    if base is not None:
+        cfg = merge(load_config(path.parent / base).to_dict(), cfg)
+    for item in overrides:
+        if "=" not in item:
+            raise ValueError(f"override must be key=value, got {item!r}")
+        key, raw = item.split("=", 1)
+        _set_dotted(cfg, key, _parse_value(raw))
+    return _wrap(cfg)
+
+
+def parse_args(description: str = "", argv: Sequence[str] | None = None) -> Config:
+    """``--config file.yaml [key=value ...]`` -> config."""
+    p = argparse.ArgumentParser(description=description)
+    p.add_argument("--config", required=True, help="YAML config file")
+    p.add_argument("overrides", nargs="*", help="dotted overrides, e.g. train.lr=1e-5")
+    a = p.parse_args(argv)
+    return load_config(a.config, a.overrides)
+
+
+def save_config(cfg: dict, path: str | Path) -> None:
+    Path(path).write_text(yaml.safe_dump(_plain(cfg), sort_keys=False))

robotryon/embodiment/__init__.py

@@ -0,0 +1,25 @@
+"""Robot embodiments: kinematics, cameras, silhouettes and hand retargeting from a YAML spec.
+
+``load_embodiment("my_robot.yaml")`` returns a :class:`MujocoEmbodiment`; see
+``configs/embodiments/gr1.yaml`` for every field.
+"""
+
+from .base import OFFSCREEN_TOLERANCE, Embodiment
+from .camera import CameraModel
+from .spec import EmbodimentSpec, load_spec, parse_spec
+
+
+def load_embodiment(path_or_name):
+    """Embodiment from a YAML path, or from ``configs/embodiments/<name>.yaml`` (relative to
+    the working directory) when given a bare name."""
+    from pathlib import Path
+
+    from .mujoco_embodiment import MujocoEmbodiment   # mujoco is an optional dependency
+    p = Path(path_or_name)
+    if not p.exists() and p.suffix == "":
+        p = Path("configs/embodiments") / f"{path_or_name}.yaml"
+    return MujocoEmbodiment(load_spec(p))
+
+
+__all__ = ["Embodiment", "CameraModel", "EmbodimentSpec", "OFFSCREEN_TOLERANCE",
+           "load_embodiment", "load_spec", "parse_spec"]

robotryon/embodiment/base.py

@@ -0,0 +1,144 @@
+"""The interface every robot is used through.
+
+The generative models never see joint angles, only images, a mask of the arm and hand, and a
+17-keypoint hand skeleton. What a robot must provide is therefore small:
+
+* ``fk(q)``: wrist pose and the 17 hand keypoints of each hand in the robot base frame;
+* ``camera_pose(q, view)`` and ``project(points, view)``: the head camera(s), to draw the
+  skeleton of a robot frame and to place a human hand recorded from the same viewpoint;
+* ``silhouette(q, view)``: the rendered robot, which together with SAM3 gives the robot mask;
+* ``retarget_hand(human_kp3d)``: the robot hand posed like a reconstructed human hand.
+
+Arrays over hands always use the order ``(left, right)``; a hand the robot does not have is NaN.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from robotryon.skeleton.topology import NUM_HANDS, NUM_KEYPOINTS
+
+from .spec import HANDS, EmbodimentSpec
+
+# A keypoint off screen is still drawn (OpenCV clips the bone) unless it lies further than
+# this many frame sizes outside the image, where the projection is degenerate.
+OFFSCREEN_TOLERANCE = 3.0
+
+
+class Embodiment(ABC):
+    spec: EmbodimentSpec
+
+    @property
+    def name(self) -> str:
+        return self.spec.name
+
+    @property
+    def views(self) -> list[str]:
+        return self.spec.views
+
+    @property
+    def joint_names(self) -> list[str | None]:
+        """Layout of the recorded joint vector ``q`` (``None`` = slot not in the model)."""
+        return self.spec.joints
+
+    # --- kinematics ----------------------------------------------------------------------
+    @abstractmethod
+    def fk(self, q) -> dict[str, dict[str, np.ndarray]]:
+        """Forward kinematics in the base frame, metres.
+
+        Args:
+            q: joint vector in ``joint_names`` order, or a ``{joint: value}`` dict.
+        Returns:
+            ``{hand: {"wrist_pos": (3,), "wrist_rot": (3, 3), "keypoints": (17, 3)}}`` for each
+            hand the robot has (``"left"``/``"right"``); keypoint 0 is the wrist and
+            ``wrist_rot`` columns are the wrist frame's axes in the base frame.
+        """
+
+    @property
+    def arm_joint_slices(self) -> dict[str, slice]:
+        """Slice of the joint/action vector holding each arm's joints (YAML ``arm_joints``)."""
+        return {h: slice(*r) for h, r in self.spec.arm_joints.items()}
+
+    def hand_keypoints(self, q) -> np.ndarray:
+        """``(2, 17, 3)`` keypoints of ``fk(q)`` in (left, right) order, NaN for a missing hand."""
+        out, f = empty_hands(), self.fk(q)
+        for hi, h in enumerate(HANDS):
+            if h in f:
+                out[hi] = f[h]["keypoints"]
+        return out
+
+    @abstractmethod
+    def camera_pose(self, q, view: str) -> np.ndarray:
+        """4x4 camera-to-base pose of ``view`` (optical frame) at joints ``q``.
+
+        ``q=None`` uses the default pose (useful for fixed or stereo-relative quantities).
+        """
+
+    @abstractmethod
+    def silhouette(self, q, view: str, image_size=None) -> np.ndarray:
+        """Rendered robot arm+hand mask, bool ``(H, W)``."""
+
+    @abstractmethod
+    def retarget_hand(self, human_kp3d: np.ndarray, valid=None, prev: dict | None = None):
+        """Robot hand posed like a human hand (vector retargeting, see ``retarget``).
+
+        Args:
+            human_kp3d: ``(2, 21|17, 3)`` human joints (MediaPipe order) in any frame.
+            valid: ``(2,)`` bool, hands to retarget (default: finite ones).
+            prev: the joints returned for the previous frame of the clip; they warm-start the
+                optimization and its smoothing term (None for the first frame).
+        Returns:
+            ``(2, 17, 3)`` robot keypoints in the same frame (NaN for skipped hands) and a
+            ``{joint: value}`` dict of the retargeted robot hand joints.
+        """
+
+    # --- cameras -------------------------------------------------------------------------
+    def camera(self, view: str):
+        return self.spec.cameras[view].intrinsics
+
+    def stereo_extrinsic(self, view_a: str, view_b: str) -> np.ndarray:
+        """4x4 transform from ``view_b``'s camera frame to ``view_a``'s."""
+        Ta, Tb = self.camera_pose(None, view_a), self.camera_pose(None, view_b)
+        return np.linalg.inv(Ta) @ Tb
+
+    def project(self, points, view: str, q=None, frame: str = "base", image_size=None):
+        """``(..., 3)`` points -> ``(..., 2)`` pixels and ``(...)`` depth-ok mask.
+
+        ``frame="base"`` needs ``q`` (the camera may move with the head);
+        ``frame="camera"`` takes points already in ``view``'s optical frame.
+        """
+        pts = np.asarray(points, np.float64)
+        if frame == "base":
+            T = np.linalg.inv(self.camera_pose(q, view))
+            pts = pts @ T[:3, :3].T + T[:3, 3]
+        elif frame != "camera":
+            raise ValueError(f"frame must be 'base' or 'camera', got {frame!r}")
+        return self.camera(view).project(pts, image_size)
+
+    def keypoints_2d(self, points, view: str, image_size, q=None, frame: str = "base"):
+        """Project ``(2, 17, 3)`` keypoints into the skeleton format.
+
+        Returns ``kp`` float32 ``(2, 17, 2)`` (0 where invisible) and ``vis`` bool ``(2, 17)``:
+        in front of the camera and within ``OFFSCREEN_TOLERANCE`` frames of the image.
+        """
+        px, ok = self.project(points, view, q=q, frame=frame, image_size=image_size)
+        W, H = image_size
+        m = OFFSCREEN_TOLERANCE
+        u, v = px[..., 0], px[..., 1]
+        with np.errstate(invalid="ignore"):
+            vis = ok & (u > -m * W) & (u < (1 + m) * W) & (v > -m * H) & (v < (1 + m) * H)
+        kp = np.where(vis[..., None], px, 0.0).astype(np.float32)
+        return kp, vis
+
+    def keypoints_in_view(self, q, view: str, image_size):
+        """Projected skeleton ``(kp, vis)`` of the robot at joints ``q`` (see keypoints_2d)."""
+        return self.keypoints_2d(self.hand_keypoints(q), view, image_size, q=q)
+
+
+def empty_hands(d: int = 3) -> np.ndarray:
+    return np.full((NUM_HANDS, NUM_KEYPOINTS, d), np.nan)
+
+
+__all__ = ["Embodiment", "HANDS", "OFFSCREEN_TOLERANCE", "empty_hands"]

robotryon/embodiment/camera.py

@@ -0,0 +1,103 @@
+"""Camera intrinsics: projection of camera-frame points to pixels and back to rays.
+
+A camera is calibrated at one resolution (``width`` x ``height``); every method takes the
+image size it should work at and rescales the intrinsics, so the same calibration serves a
+video stored at any resolution with the same aspect. Points are in the OpenCV optical frame
+(x right, y down, z forward). Two lens models are supported:
+
+* ``pinhole``: OpenCV plumb-bob distortion ``[k1, k2, p1, p2, k3]`` (empty = rectified).
+* ``fisheye``: OpenCV equidistant model ``[k1, k2, k3, k4]`` (zeros = ideal equidistant).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+
+
+@dataclass
+class CameraModel:
+    width: int
+    height: int
+    fx: float
+    fy: float
+    cx: float
+    cy: float
+    model: str = "pinhole"
+    dist: tuple = ()
+    # Additive pixel offset applied after projection (a fitted principal-point residual),
+    # given at the calibration resolution.
+    pp_offset: tuple = (0.0, 0.0)
+    # Points closer than this to the lens are inside the robot's own head and are not drawn.
+    min_depth: float = 0.15
+    extra: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        if self.model not in ("pinhole", "fisheye"):
+            raise ValueError(f"unknown camera model {self.model!r}")
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "CameraModel":
+        known = {"width", "height", "fx", "fy", "cx", "cy", "model", "dist", "pp_offset",
+                 "min_depth"}
+        kw = {k: d[k] for k in known if k in d}
+        kw["dist"] = tuple(float(x) for x in d.get("dist", ()) or ())
+        kw["pp_offset"] = tuple(float(x) for x in d.get("pp_offset", (0.0, 0.0)))
+        return cls(**kw, extra={k: v for k, v in d.items() if k not in known})
+
+    # --- intrinsics at a given resolution -------------------------------------------------
+    def _scale(self, image_size):
+        if image_size is None:
+            return 1.0, 1.0
+        return image_size[0] / self.width, image_size[1] / self.height
+
+    def K(self, image_size=None) -> np.ndarray:
+        sx, sy = self._scale(image_size)
+        return np.array([[self.fx * sx, 0.0, self.cx * sx],
+                         [0.0, self.fy * sy, self.cy * sy],
+                         [0.0, 0.0, 1.0]])
+
+    def D(self) -> np.ndarray:
+        n = 4 if self.model == "fisheye" else 5
+        d = np.zeros(n)
+        d[:len(self.dist)] = self.dist
+        return d
+
+    def offset(self, image_size=None) -> np.ndarray:
+        sx, sy = self._scale(image_size)
+        return np.array([self.pp_offset[0] * sx, self.pp_offset[1] * sy])
+
+    # --- projection ------------------------------------------------------------------------
+    def project(self, pts_cam: np.ndarray, image_size=None):
+        """``(..., 3)`` camera-frame points -> ``(..., 2)`` pixels and ``(...)`` depth-ok mask.
+
+        Points behind ``min_depth`` get NaN pixels.
+        """
+        pts = np.asarray(pts_cam, np.float64)
+        shape = pts.shape[:-1]
+        flat = pts.reshape(-1, 3)
+        ok = np.isfinite(flat).all(1) & (flat[:, 2] > self.min_depth)
+        px = np.full((flat.shape[0], 2), np.nan)
+        if ok.any():
+            obj = flat[ok].reshape(-1, 1, 3)
+            K, D = self.K(image_size), self.D()
+            if self.model == "fisheye":
+                uv, _ = cv2.fisheye.projectPoints(obj, np.zeros(3), np.zeros(3), K, D)
+            else:
+                uv, _ = cv2.projectPoints(obj, np.zeros(3), np.zeros(3), K, D)
+            px[ok] = uv.reshape(-1, 2) + self.offset(image_size)
+        return px.reshape(shape + (2,)), ok.reshape(shape)
+
+    def normalized(self, pixels: np.ndarray, image_size=None) -> np.ndarray:
+        """``(..., 2)`` pixels -> ``(..., 2)`` undistorted normalized coordinates (x/z, y/z)."""
+        uv = np.asarray(pixels, np.float64)
+        shape = uv.shape[:-1]
+        uv = (uv.reshape(-1, 2) - self.offset(image_size)).reshape(-1, 1, 2)
+        K, D = self.K(image_size), self.D()
+        if self.model == "fisheye":
+            n = cv2.fisheye.undistortPoints(uv, K, D)
+        else:
+            n = cv2.undistortPoints(uv, K, D)
+        return n.reshape(shape + (2,))

robotryon/embodiment/mujoco_embodiment.py

@@ -0,0 +1,320 @@
+"""MuJoCo implementation of :class:`Embodiment`, generic over the YAML description.
+
+The robot model (MJCF, or URDF which MuJoCo also parses) is supplied by the user; nothing
+robot-specific is coded here. Forward kinematics is ``mj_kinematics`` on a pose built from
+the recorded joint vector: unlisted joints stay at ``qpos0``, coupled passive joints are
+filled from their sources and ``fixed_joints`` are applied last. A head camera mounted on a
+model frame therefore follows the recorded head joints.
+
+Pin ``mujoco==3.11.0``: with 3.12 site-based camera poses silently collapse (every point
+lands at camera depth ~0), so a different version raises unless
+``ROBOTRYON_ALLOW_ANY_MUJOCO=1``.
+
+Rendering (``silhouette``) needs an OpenGL backend; set ``MUJOCO_GL=egl`` (GPU) or
+``MUJOCO_GL=osmesa`` (CPU) before the first import of mujoco.
+"""
+
+from __future__ import annotations
+
+import os
+
+import cv2
+import numpy as np
+
+from . import retarget as RT
+from .base import Embodiment, empty_hands
+from .spec import HANDS, EmbodimentSpec, load_spec
+
+MUJOCO_PIN = "3.11.0"
+_SIL_HIDDEN_GROUP = 5
+
+
+def _import_mujoco():
+    import mujoco
+    if mujoco.__version__ != MUJOCO_PIN and os.environ.get("ROBOTRYON_ALLOW_ANY_MUJOCO") != "1":
+        raise RuntimeError(
+            f"mujoco {mujoco.__version__} != {MUJOCO_PIN}: newer versions break site-based "
+            f"camera poses without an error. Install mujoco=={MUJOCO_PIN} or set "
+            f"ROBOTRYON_ALLOW_ANY_MUJOCO=1.")
+    return mujoco
+
+
+class MujocoEmbodiment(Embodiment):
+    def __init__(self, spec: EmbodimentSpec | str):
+        self.spec = load_spec(spec) if isinstance(spec, (str, os.PathLike)) else spec
+        mj = self._mj = _import_mujoco()
+        if not os.path.exists(self.spec.model_path):
+            raise FileNotFoundError(f"robot model not found: {self.spec.model_path}")
+        self.model = mj.MjModel.from_xml_path(self.spec.model_path)
+        self.data = mj.MjData(self.model)
+
+        names = set(n for n in self.spec.joints if n) | set(self.spec.fixed_joints)
+        names |= {c[0] for c in self.spec.coupling} | {c[2] for c in self.spec.coupling}
+        names |= {j for js in self.spec.retarget.joints.values() for j in js}
+        self._adr = {n: self._joint_adr(n) for n in names}
+        self._rt_bounds = {h: self._joint_limits(js) for h, js in self.spec.retarget.joints.items()}
+        self._rt_palm: dict[str, float] = {}
+        self._kp = {h: [(*self._frame(n), off) for n, off in
+                        zip(self.spec.keypoints[h], self.spec.keypoint_offsets[h])]
+                    for h in self.spec.hands}
+        # the wrist pose: the explicit wrist_frame, else keypoint 0 (with its offset)
+        self._wrist = {h: (*self._frame(self.spec.wrist_frame[h]),
+                           self.spec.keypoint_offsets[h][0]
+                           if self.spec.wrist_frame[h] == self.spec.keypoints[h][0] else None)
+                       for h in self.spec.hands}
+        self._mount = {v: (self._frame(c.mount.frame) if c.mount.frame else None)
+                       for v, c in self.spec.cameras.items()}
+        self._renderers: dict = {}
+
+    # --- name resolution -----------------------------------------------------------------
+    def _joint_adr(self, name: str) -> int:
+        jid = self._mj.mj_name2id(self.model, self._mj.mjtObj.mjOBJ_JOINT, name)
+        if jid < 0:
+            raise ValueError(f"joint {name!r} not in {self.spec.model_path}")
+        return int(self.model.jnt_qposadr[jid])
+
+    def _joint_limits(self, names) -> tuple[np.ndarray, np.ndarray]:
+        """``(lo, hi)`` of each joint: the model's limits ([-pi, pi] when unlimited), narrowed
+        by ``retarget.limits``."""
+        lo, hi = [], []
+        for n in names:
+            jid = self._mj.mj_name2id(self.model, self._mj.mjtObj.mjOBJ_JOINT, n)
+            r = self.model.jnt_range[jid]
+            ok = bool(self.model.jnt_limited[jid]) and r[1] > r[0]
+            a, b = (float(r[0]), float(r[1])) if ok else (-np.pi, np.pi)
+            if n in self.spec.retarget.limits:
+                ua, ub = self.spec.retarget.limits[n]
+                a, b = max(a, ua), min(b, ub)
+                if b <= a:
+                    raise ValueError(f"retarget.limits for {n} [{ua}, {ub}] leaves no range "
+                                     f"inside the model limits")
+            lo.append(a)
+            hi.append(b)
+        return np.asarray(lo), np.asarray(hi)
+
+    def _frame(self, name: str) -> tuple[str, int]:
+        """``site:x`` / ``body:x`` / ``x`` (site first, then body) -> (kind, id)."""
+        kind, _, bare = name.rpartition(":")
+        T = self._mj.mjtObj
+        for k in ([kind] if kind else ["site", "body"]):
+            oid = self._mj.mj_name2id(self.model, T.mjOBJ_SITE if k == "site" else T.mjOBJ_BODY,
+                                      bare)
+            if oid >= 0:
+                return k, oid
+        raise ValueError(f"frame {name!r} not in {self.spec.model_path}")
+
+    def _pos(self, f):
+        p = (self.data.site_xpos if f[0] == "site" else self.data.xpos)[f[1]].copy()
+        if len(f) > 2 and f[2] is not None:                  # point fixed in the frame
+            p = p + self._rot(f) @ f[2]
+        return p
+
+    def _rot(self, f):
+        m = self.data.site_xmat if f[0] == "site" else self.data.xmat
+        return m[f[1]].reshape(3, 3).copy()
+
+    # --- kinematics ----------------------------------------------------------------------
+    def joint_dict(self, q) -> dict[str, float]:
+        """Joint vector (``spec.joints`` order) or dict -> complete ``{joint: value}``."""
+        if q is None:
+            d = {}
+        elif isinstance(q, dict):
+            d = {k: float(v) for k, v in q.items()}
+        else:
+            q = np.asarray(q, np.float64).reshape(-1)
+            if len(q) != len(self.spec.joints):
+                raise ValueError(f"expected {len(self.spec.joints)} joints, got {len(q)}")
+            d = {n: float(v) * self.spec.joint_scale.get(n, 1.0) for n, v in zip(self.spec.joints, q) if n}
+        for joint, ratio, src in self.spec.coupling:
+            if src in d:
+                d[joint] = ratio * d[src]
+        d.update(self.spec.fixed_joints)
+        return d
+
+    def set_state(self, q) -> None:
+        self.data.qpos[:] = self.model.qpos0
+        for n, v in self.joint_dict(q).items():
+            adr = self._adr.get(n)
+            if adr is None:
+                adr = self._adr[n] = self._joint_adr(n)
+            self.data.qpos[adr] = v
+        self._mj.mj_kinematics(self.model, self.data)
+
+    def _hand_points(self) -> np.ndarray:
+        """Keypoints of the current state, base frame, with the per-finger length scale.
+
+        The scale pulls each finger's distal points toward its knuckle
+        (``p <- p_knuckle + s * (p - p_knuckle)``), a fitted correction for model fingers that
+        render longer than the real hand.
+        """
+        out = empty_hands()
+        for hi, h in enumerate(HANDS):
+            if h not in self._kp:
+                continue
+            p = np.stack([self._pos(f) for f in self._kp[h]])
+            for finger, (k, distal) in RT.CHAINS17.items():
+                s = self.spec.finger_scale.get(finger)
+                if s is not None:
+                    p[list(distal)] = p[k] + s * (p[list(distal)] - p[k])
+            out[hi] = p
+        return out
+
+    def fk(self, q) -> dict[str, dict[str, np.ndarray]]:
+        self.set_state(q)
+        kp = self._hand_points()
+        return {h: {"wrist_pos": self._pos(self._wrist[h]), "wrist_rot": self._rot(self._wrist[h]),
+                    "keypoints": kp[HANDS.index(h)]}
+                for h in HANDS if h in self._wrist}
+
+    def hand_keypoints(self, q) -> np.ndarray:
+        self.set_state(q)
+        return self._hand_points()
+
+    def _camera_pose_now(self, view: str) -> np.ndarray:
+        m = self.spec.cameras[view].mount
+        if m.extrinsic is not None:
+            return m.extrinsic.copy()
+        f = self._mount[view]
+        R_frame = self._rot(f)
+        T = np.eye(4)
+        T[:3, :3] = R_frame @ m.rotation.T
+        T[:3, 3] = self._pos(f) + R_frame @ m.translation
+        return T
+
+    def camera_pose(self, q, view: str) -> np.ndarray:
+        self.set_state(q)
+        return self._camera_pose_now(view)
+
+    def keypoints_in_view(self, q, view: str, image_size):
+        """One FK pass -> projected skeleton ``kp (2,17,2)``, ``vis (2,17)`` of a robot frame."""
+        self.set_state(q)
+        T = np.linalg.inv(self._camera_pose_now(view))
+        pts = self._hand_points() @ T[:3, :3].T + T[:3, 3]
+        return self.keypoints_2d(pts, view, image_size, frame="camera")
+
+    # --- rendering -----------------------------------------------------------------------
+    def _renderer(self, view: str, image_size, scale: float):
+        key = (view, tuple(image_size), scale)
+        if key in self._renderers:
+            return self._renderers[key]
+        mj, model = self._mj, self.model
+        inc, exc = tuple(self.spec.silhouette_bodies), tuple(self.spec.silhouette_exclude)
+        for g in range(model.ngeom):
+            body = mj.mj_id2name(model, mj.mjtObj.mjOBJ_BODY, int(model.geom_bodyid[g])) or ""
+            show = body.startswith(inc) and not (exc and body.startswith(exc))
+            model.geom_group[g] = 0 if show else _SIL_HIDDEN_GROUP
+        # Render a symmetric wide pinhole, then remap each output pixel through the real
+        # camera's (distorted, off-centre) ray so the silhouette matches the projection.
+        W, H = image_size
+        w, h = max(8, round(W * scale)), max(8, round(H * scale))
+        cam = self.camera(view)
+        u, v = np.meshgrid((np.arange(w) + 0.5) * W / w - 0.5, (np.arange(h) + 0.5) * H / h - 0.5)
+        n = cam.normalized(np.stack([u, v], -1), image_size)
+        f = cam.K(image_size)[1, 1] * h / H
+        ext = np.abs(n.reshape(-1, 2)).max(0) * 1.03
+        Wp, Hp = int(2 * np.ceil(ext[0] * f)) + 2, int(2 * np.ceil(ext[1] * f)) + 2
+        mx = (n[..., 0] * f + Wp / 2.0 - 0.5).astype(np.float32)
+        my = (n[..., 1] * f + Hp / 2.0 - 0.5).astype(np.float32)
+        model.vis.global_.offwidth = max(int(model.vis.global_.offwidth), Wp)
+        model.vis.global_.offheight = max(int(model.vis.global_.offheight), Hp)
+        model.vis.map.znear = 0.001
+        model.vis.quality.offsamples = 0          # MSAA would blend segmentation ids
+        r = mj.Renderer(model, height=Hp, width=Wp)
+        r.enable_segmentation_rendering()
+        opt = mj.MjvOption()
+        opt.geomgroup[:] = 0
+        opt.geomgroup[0] = 1
+        opt.sitegroup[:] = 0
+        mcam = mj.MjvCamera()
+        mj.mjv_defaultCamera(mcam)
+        self._renderers[key] = (r, opt, mcam, mx, my, f, Hp)
+        return self._renderers[key]
+
+    def silhouette(self, q, view: str, image_size=None, scale: float = 0.5) -> np.ndarray:
+        cam = self.camera(view)
+        image_size = tuple(image_size or (cam.width, cam.height))
+        r, opt, mcam, mx, my, f, Hp = self._renderer(view, image_size, scale)
+        self.set_state(q)
+        T = self._camera_pose_now(view)
+        r.update_scene(self.data, camera=mcam, scene_option=opt)
+        near, far = 0.02, 6.0
+        half_h = near * (Hp / 2.0) / f
+        for gc in r.scene.camera:
+            gc.pos[:] = T[:3, 3]
+            gc.forward[:] = T[:3, 2]
+            gc.up[:] = -T[:3, 1]
+            gc.frustum_center = 0.0
+            gc.frustum_bottom, gc.frustum_top = -half_h, half_h
+            gc.frustum_near, gc.frustum_far = near, far
+            gc.orthographic = 0
+        seg = r.render()[:, :, 0]
+        sil = cv2.remap((seg >= 0).astype(np.uint8), mx, my, cv2.INTER_NEAREST,
+                        borderMode=cv2.BORDER_CONSTANT, borderValue=0)
+        return cv2.resize(sil, image_size, interpolation=cv2.INTER_NEAREST) > 0
+
+    # --- retargeting ---------------------------------------------------------------------
+    def _hand_vectors(self, hand: str, x: np.ndarray):
+        """Retargeting vectors (wrist frame) of one robot hand at hand-joint values ``x``."""
+        self.set_state(dict(zip(self.spec.retarget.joints[hand], x)))
+        return RT.local_vectors(self._hand_points()[HANDS.index(hand)])
+
+    def _palm(self, hand: str) -> float:
+        if hand not in self._rt_palm:
+            lo, hi = self._rt_bounds[hand]
+            self.set_state(dict(zip(self.spec.retarget.joints[hand], np.clip(0.0, lo, hi))))
+            self._rt_palm[hand] = RT.palm_size(self._hand_points()[HANDS.index(hand)])
+        return self._rt_palm[hand]
+
+    def _solve_hand(self, hand: str, h17: np.ndarray, x_prev=None):
+        """AnyTeleop vector retargeting of one hand: hand joints (model order) or None."""
+        from scipy.optimize import least_squares
+
+        rt = self.spec.retarget
+        target = RT.local_vectors(h17)
+        if target is None:
+            return None
+        scale = rt.scaling if rt.scaling is not None else self._palm(hand) / max(RT.palm_size(h17), 1e-6)
+        target = target * scale
+        lo, hi = self._rt_bounds[hand]
+        x0 = np.clip(np.zeros(len(lo)) if x_prev is None else np.asarray(x_prev, np.float64), lo, hi)
+        x0 = np.clip(x0, lo + 1e-9, hi - 1e-9)
+        w = np.sqrt(rt.smoothing)
+
+        def residual(x):
+            v = self._hand_vectors(hand, x)
+            r = np.zeros(target.size) if v is None else (v - target).ravel()
+            return r if x_prev is None else np.concatenate([r, w * (x - x0)])
+
+        sol = least_squares(residual, x0, bounds=(lo, hi), loss="huber", f_scale=rt.huber_delta,
+                            diff_step=1e-4, max_nfev=200)
+        return sol.x
+
+    def retarget_hand(self, human_kp3d, valid=None, prev=None):
+        j21 = RT.to21(human_kp3d)
+        if valid is None:
+            valid = np.isfinite(j21).all(axis=(-1, -2))
+        rt = self.spec.retarget
+        q = {}
+        for hi, h in enumerate(HANDS):
+            joints = rt.joints.get(h)
+            if not valid[hi] or not joints:
+                continue
+            x_prev = None
+            if prev and all(n in prev for n in joints):
+                x_prev = [prev[n] for n in joints]
+            x = self._solve_hand(h, RT.human17(j21[hi]), x_prev)
+            if x is not None:
+                q.update({n: float(v) for n, v in zip(joints, x)})
+        out = empty_hands()
+        if not q:
+            return out, q
+        robot = self.hand_keypoints(q)
+        for hi, h in enumerate(HANDS):
+            if valid[hi] and all(n in q for n in rt.joints.get(h) or [None]):
+                placed = RT.place_on_human(robot[hi], j21[hi], rt.bone_match)
+                if placed is not None:
+                    out[hi] = placed
+        return out, q
+
+
+__all__ = ["MujocoEmbodiment", "MUJOCO_PIN"]

robotryon/embodiment/retarget.py

@@ -0,0 +1,122 @@
+"""Retarget a reconstructed human hand onto a robot hand (AnyTeleop vector retargeting).
+
+The skeleton that conditions the generative models must look like the target robot's hand,
+so a human hand is not drawn directly. Instead (Appendix B.2.2):
+
+1. Both hands are expressed in a wrist frame (origin at the wrist, x toward the knuckle
+   centroid, z along the knuckle-plane normal) built identically from the human joints and
+   from the robot's 17 keypoints.
+2. Following the vector retargeting of AnyTeleop (Qin et al., 2023), the robot hand joints are
+   optimized so that the robot's wrist-to-fingertip and wrist-to-middle-joint vectors match the
+   human's, scaled by the ratio of the palm sizes, within the joint limits of the robot model,
+   with a small penalty on the change from the previous frame (``VECTORS17``,
+   ``MujocoEmbodiment.retarget_hand``).
+3. Robot forward kinematics at those joints gives the robot hand's 17 keypoints. Bone
+   matching stretches each robot finger to the human's knuckle distance and finger length
+   while keeping the robot's directions (the robot's own curl), so the fingertips reach what
+   the human touches.
+4. The result is placed at the human wrist frame, in the human's camera frame.
+
+The frame helpers and the vector set live here; the optimization needs FK and is done by the
+embodiment.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from robotryon.skeleton.topology import MEDIAPIPE21_TO_17
+
+# MediaPipe / MANO 21-joint layout: wrist, then four joints per finger, proximal to distal.
+CHAINS21 = {"thumb": (1, 2, 3, 4), "index": (5, 6, 7, 8), "middle": (9, 10, 11, 12),
+            "ring": (13, 14, 15, 16), "little": (17, 18, 19, 20)}
+# 17-keypoint layout (robotryon.skeleton.topology): (knuckle, distal points) per finger.
+CHAINS17 = {"thumb": (1, (2, 3, 4)), "index": (5, (6, 7)), "middle": (8, (9, 10)),
+            "ring": (11, (12, 13)), "little": (14, (15, 16))}
+FINGERS4 = ("index", "middle", "ring", "little")
+
+# Robot wrist frame from the 17-keypoint layout: wrist, index knuckle, little knuckle.
+WRIST17, INDEX17, LITTLE17 = 0, 5, 14
+
+# Retargeting vectors as (origin, target) pairs of the 17-keypoint layout: wrist to every
+# fingertip and wrist to the joint before it (thumb distal, index/middle/ring/little middle).
+TIPS17 = tuple(distal[-1] for _, distal in CHAINS17.values())
+MIDS17 = (3, 6, 9, 12, 15)
+VECTORS17 = tuple((WRIST17, t) for t in TIPS17) + tuple((WRIST17, m) for m in MIDS17)
+
+
+def to21(kp: np.ndarray) -> np.ndarray:
+    """``(..., 17|21, 3)`` -> 21 joints. Missing DIP joints are put midway between PIP and
+    tip, which leaves every finger length and direction used below unchanged."""
+    kp = np.asarray(kp, np.float64)
+    if kp.shape[-2] == 21:
+        return kp
+    out = np.zeros(kp.shape[:-2] + (21, 3))
+    out[..., MEDIAPIPE21_TO_17, :] = kp
+    for pip, dip, tip in ((6, 7, 8), (10, 11, 12), (14, 15, 16), (18, 19, 20)):
+        out[..., dip, :] = 0.5 * (out[..., pip, :] + out[..., tip, :])
+    return out
+
+
+def wrist_frame(w, imcp, lmcp):
+    """Right-handed frame (columns = axes) at the wrist, or None when degenerate."""
+    x = (imcp + lmcp) * 0.5 - w
+    n = np.cross(imcp - w, lmcp - w)
+    nx, nn = np.linalg.norm(x), np.linalg.norm(n)
+    if nx < 1e-9 or nn < 1e-9:
+        return None
+    x, n = x / nx, n / nn
+    return np.stack([x, np.cross(n, x), n], axis=1)
+
+
+def palm_size(p17: np.ndarray) -> float:
+    """Mean wrist-to-knuckle distance of the four fingers (17-keypoint layout)."""
+    return float(np.mean([np.linalg.norm(p17[CHAINS17[f][0]] - p17[WRIST17]) for f in FINGERS4]))
+
+
+def local_vectors(p17: np.ndarray):
+    """``(17, 3)`` keypoints -> ``(len(VECTORS17), 3)`` vectors in the hand's wrist frame,
+    or None when the frame is degenerate."""
+    R = wrist_frame(p17[WRIST17], p17[INDEX17], p17[LITTLE17])
+    if R is None:
+        return None
+    o, t = zip(*VECTORS17)
+    return (p17[list(t)] - p17[list(o)]) @ R
+
+
+def human17(j21: np.ndarray) -> np.ndarray:
+    """``(21, 3)`` human joints -> the 17-keypoint layout."""
+    return np.asarray(j21, np.float64)[list(MEDIAPIPE21_TO_17)]
+
+
+def bone_match(local17: np.ndarray, j21: np.ndarray) -> np.ndarray:
+    """Stretch robot finger segments (wrist-frame 17 keypoints) to the human's lengths.
+
+    Each knuckle is moved along its own direction to the human wrist-to-knuckle distance, then
+    the rest of the finger is scaled about the knuckle so knuckle-to-tip matches the human's.
+    Directions stay the robot's.
+    """
+    out = local17.copy()
+    for f, (k, distal) in CHAINS17.items():
+        hm, ht = CHAINS21[f][0], CHAINS21[f][-1]
+        r_palm, h_palm = np.linalg.norm(local17[k]), np.linalg.norm(j21[hm] - j21[0])
+        if r_palm > 1e-6 and h_palm > 1e-6:
+            out[k] = local17[k] * (h_palm / r_palm)
+        r_fin = np.linalg.norm(local17[distal[-1]] - local17[k])
+        h_fin = np.linalg.norm(j21[ht] - j21[hm])
+        s = h_fin / r_fin if (r_fin > 1e-6 and h_fin > 1e-6) else 1.0
+        for i in distal:
+            out[i] = out[k] + s * (local17[i] - local17[k])
+    return out
+
+
+def place_on_human(robot17: np.ndarray, j21: np.ndarray, match_bones: bool = True):
+    """Move robot keypoints (any frame) onto the human wrist frame; None if degenerate."""
+    Rr = wrist_frame(robot17[WRIST17], robot17[INDEX17], robot17[LITTLE17])
+    Rh = wrist_frame(j21[0], j21[5], j21[17])
+    if Rr is None or Rh is None:
+        return None
+    local = (robot17 - robot17[WRIST17]) @ Rr
+    if match_bones:
+        local = bone_match(local, j21)
+    return j21[0] + local @ Rh.T

robotryon/embodiment/spec.py

@@ -0,0 +1,201 @@
+"""Embodiment description parsed from YAML (see ``configs/embodiments/gr1.yaml``).
+
+Everything that differs between robots lives here as data: the kinematic model file, the
+layout of the recorded joint vector, passive joints coupled to actuated ones, the link that
+stands for each of the 17 hand keypoints, the cameras, the links that make up the robot's
+silhouette, and which hand joints the human-hand retargeting moves. Parsing is pure Python
+so configs can be validated without a physics engine.
+
+Strings may contain ``{s}``, which is replaced by each hand's side token (``hands.sides``), so
+a mirror-symmetric robot is described once for both hands.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import yaml
+
+from robotryon.skeleton.topology import NUM_KEYPOINTS
+
+from .camera import CameraModel
+
+HANDS = ("left", "right")
+
+
+@dataclass
+class CameraMount:
+    """Where a camera sits: on a model frame that moves with the robot, or fixed in the base.
+
+    ``frame`` is ``site:<name>``/``body:<name>`` (or a bare name, resolved site first);
+    ``rotation`` maps the frame's axes to the optical frame (``p_opt = R @ p_frame``);
+    ``translation`` offsets the lens origin in the mount frame (e.g. a stereo baseline).
+    ``extrinsic`` instead gives a fixed 4x4 camera-to-base pose.
+    """
+    frame: str | None = None
+    rotation: np.ndarray = field(default_factory=lambda: np.eye(3))
+    translation: np.ndarray = field(default_factory=lambda: np.zeros(3))
+    extrinsic: np.ndarray | None = None
+
+
+@dataclass
+class CameraSpec:
+    intrinsics: CameraModel
+    mount: CameraMount
+    video_key: str | None = None
+
+
+@dataclass
+class RetargetSpec:
+    """Vector retargeting of a human hand onto the robot hand (``robotryon.embodiment.retarget``).
+
+    ``joints``: per hand, the actuated hand joints the optimization moves; coupled joints follow
+    their source through ``coupling`` and the limits come from the robot model.
+    ``scaling``: factor on the human vectors, or None for the ratio of the palm sizes.
+    ``smoothing``: weight of the squared change from the previous frame's joints.
+    ``huber_delta``: Huber threshold (m) of the vector residuals.
+    ``limits``: ``{joint: (lo, hi)}`` narrowing the model's joint limits for the retargeting, e.g.
+    to keep finger joints out of hyperextension the model allows but the real hand never uses.
+    """
+    joints: dict[str, list[str]] = field(default_factory=dict)
+    limits: dict[str, tuple[float, float]] = field(default_factory=dict)
+    scaling: float | None = None
+    smoothing: float = 4e-3
+    huber_delta: float = 0.02
+    bone_match: bool = True
+
+
+@dataclass
+class EmbodimentSpec:
+    name: str
+    model_path: str
+    joints: list[str | None]
+    joint_scale: dict[str, float]                     # joint -> model value / recorded value
+    arm_joints: dict[str, tuple[int, int]]            # hand -> [start, stop) in the joint vector
+    coupling: list[tuple[str, float, str]]            # (joint, ratio, source joint)
+    fixed_joints: dict[str, float]
+    keypoints: dict[str, list[str]]                   # hand -> 17 frame names
+    keypoint_offsets: dict[str, list]                 # hand -> 17 offsets (m, in the frame) or None
+    wrist_frame: dict[str, str]                       # hand -> frame giving the wrist pose
+    finger_scale: dict[str, float]
+    cameras: dict[str, CameraSpec]
+    silhouette_bodies: list[str]
+    silhouette_exclude: list[str]
+    retarget: RetargetSpec = field(default_factory=RetargetSpec)
+    raw: dict = field(default_factory=dict)
+
+    @property
+    def hands(self) -> list[str]:
+        return [h for h in HANDS if h in self.keypoints]
+
+    @property
+    def views(self) -> list[str]:
+        return list(self.cameras)
+
+
+def _fmt(x: Any, s: str) -> Any:
+    return x.replace("{s}", s) if isinstance(x, str) else x
+
+
+def _per_hand(value, hand: str):
+    """``value`` or ``value[hand]`` when it is given per hand."""
+    if isinstance(value, dict) and hand in value:
+        return value[hand]
+    return value
+
+
+def parse_spec(cfg: dict, base_dir: str | Path | None = None) -> EmbodimentSpec:
+    hands_cfg = cfg.get("hands", {})
+    sides = hands_cfg.get("sides", {"left": "left", "right": "right"})
+
+    keypoints, keypoint_offsets, wrist_frame = {}, {}, {}
+    for hand, tok in sides.items():
+        if hand not in HANDS:
+            raise ValueError(f"hands.sides keys must be left/right, got {hand!r}")
+        entries = (hands_cfg.get(hand) or {}).get("keypoints") or hands_cfg.get("keypoints")
+        if entries is None:
+            raise ValueError(f"no keypoints given for the {hand} hand")
+        if len(entries) != NUM_KEYPOINTS:
+            raise ValueError(f"{hand} hand needs {NUM_KEYPOINTS} keypoints, got {len(entries)}")
+        # an entry is a frame name, or {frame: name, offset: [x, y, z]} for a point fixed in it
+        names = [_fmt(e["frame"] if isinstance(e, dict) else e, tok) for e in entries]
+        keypoint_offsets[hand] = [np.asarray(_per_hand(e["offset"], hand), np.float64).reshape(3)
+                                  if isinstance(e, dict) and e.get("offset") is not None else None
+                                  for e in entries]
+        keypoints[hand] = names
+        wf = (hands_cfg.get(hand) or {}).get("wrist_frame") or hands_cfg.get("wrist_frame")
+        wrist_frame[hand] = _fmt(wf, tok) if wf else names[0]
+
+    coupling = []
+    for c in cfg.get("coupling", []) or []:
+        templated = "{s}" in c["joint"] or "{s}" in c["source"]
+        for tok in (sides.values() if templated else [""]):
+            coupling.append((_fmt(c["joint"], tok), float(c["ratio"]), _fmt(c["source"], tok)))
+
+    cameras = {}
+    for view, c in (cfg.get("cameras") or {}).items():
+        m = c.get("mount", {}) or {}
+        ext = c.get("extrinsic")
+        mount = CameraMount(
+            frame=m.get("frame"),
+            rotation=np.asarray(m.get("rotation", np.eye(3)), np.float64).reshape(3, 3),
+            translation=np.asarray(m.get("translation", [0, 0, 0]), np.float64).reshape(3),
+            extrinsic=None if ext is None else np.asarray(ext, np.float64).reshape(4, 4))
+        if mount.frame is None and mount.extrinsic is None:
+            raise ValueError(f"camera {view!r} needs mount.frame or extrinsic")
+        cameras[view] = CameraSpec(CameraModel.from_dict(c["intrinsics"]), mount,
+                                   c.get("video_key"))
+
+    rt = cfg.get("retarget", {}) or {}
+    if "drivers" in rt:
+        raise ValueError("retarget.drivers (percentile retargeting) was replaced by vector "
+                         "retargeting: list the hand joints under retarget.joints instead")
+    rt_joints = {hand: [_fmt(j, tok) for j in (_per_hand(rt.get("joints") or [], hand) or [])]
+                 for hand, tok in sides.items()}
+    limits = {}
+    for j, r in (rt.get("limits") or {}).items():
+        for hand, tok in sides.items():
+            rr = _per_hand(r, hand)
+            limits[_fmt(j, tok)] = (float(rr[0]), float(rr[1]))
+    scaling = rt.get("scaling", "auto")
+    retarget = RetargetSpec(
+        joints=rt_joints,
+        limits=limits,
+        scaling=None if scaling in (None, "auto") else float(scaling),
+        smoothing=float(rt.get("smoothing", 4e-3)),
+        huber_delta=float(rt.get("huber_delta", 0.02)),
+        bone_match=bool(rt.get("bone_match", True)))
+
+    model_path = str(cfg["model_path"])
+    if base_dir is not None and not Path(model_path).is_absolute():
+        model_path = str(Path(base_dir) / model_path)
+    sil = cfg.get("silhouette", {}) or {}
+    toks = list(sides.values())
+    expand = lambda xs: sorted({_fmt(x, t) for x in xs or [] for t in toks})  # noqa: E731
+    return EmbodimentSpec(
+        name=str(cfg.get("name", "robot")),
+        model_path=model_path,
+        joints=[j["name"] if isinstance(j, dict) else j for j in cfg.get("joints", [])],
+        joint_scale={j["name"]: float(j["scale"]) for j in cfg.get("joints", [])
+                     if isinstance(j, dict) and "scale" in j},
+        arm_joints={h: (int(r[0]), int(r[1])) for h, r in (cfg.get("arm_joints") or {}).items()},
+        coupling=coupling,
+        fixed_joints={k: float(v) for k, v in (cfg.get("fixed_joints") or {}).items()},
+        keypoints=keypoints,
+        keypoint_offsets=keypoint_offsets,
+        wrist_frame=wrist_frame,
+        finger_scale={k: float(v) for k, v in (hands_cfg.get("finger_scale") or {}).items()},
+        cameras=cameras,
+        silhouette_bodies=expand(sil.get("bodies")),
+        silhouette_exclude=expand(sil.get("exclude")),
+        retarget=retarget,
+        raw=cfg,
+    )
+
+
+def load_spec(path: str | Path) -> EmbodimentSpec:
+    path = Path(path)
+    return parse_spec(yaml.safe_load(path.read_text()), base_dir=path.parent)

robotryon/idm/__init__.py

@@ -0,0 +1,5 @@
+"""Inverse dynamics model: training data, model, action labelling and wrist correction.
+
+Torch-dependent modules (``model``, ``data``, ``labeling``) are imported explicitly by the
+scripts; ``wrist_correction`` only needs numpy/scipy.
+"""

robotryon/idm/data.py

@@ -0,0 +1,319 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted from GR00T-Dreams (Apache-2.0): gr00t/data/dataset.py (LeRobot v2.1 reader),
+# gr00t/data/transform/{video,state_action,mirror}.py, gr00t/model/transforms_idm.py and the
+# multi-embodiment mixture. Rewritten as one compact pipeline; the eval-mode image path and
+# the q99 normalisation are numerically identical to the original.
+"""IDM training data: LeRobot v2.1 datasets of robot videos with joint actions.
+
+Each embodiment is declared in the config (no robot list in code)::
+
+    name: openarm_rh56f1           # also the key of its statistics in the checkpoint metadata
+    embodiment_id: 29              # row of the per-embodiment action encoder/decoder
+    root: /path/to/lerobot_dataset
+    views: [camera_ego_left, camera_ego_right]     # LeRobot video keys; one key = mono
+    view_ids: [0, 1, 2, 3]         # view-embedding id of [t0 v0, t0 v1, t1 v0, t1 v1]
+    action_groups: [neck_joints, left_arm_joints, ...]   # modality.json groups, raw order
+    hand_groups: [left_hand_joints, right_hand_joints]   # smoothed with the short SG window
+    norm: q99                      # or q99_const (near-constant joints -> 0)
+    video_size: [512, 288]         # dataset frame size; labelled videos are resized to it
+    mirror: {prob, view_swap, swap, sign}      # optional left/right reflection augmentation
+
+A sample is frames (t, t+K) of every view and the chunk a_{t:t+K}, normalised to [-1, 1] with
+the dataset's q01/q99 and zero-padded to the model's action_dim.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+IMAGE_SIZE = 224
+_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+CONST_EPS = 1e-3
+
+
+# ----------------------------------------------------------------------------- normalisation
+def normalize(x: torch.Tensor, q01: torch.Tensor, q99: torch.Tensor, mode: str) -> torch.Tensor:
+    q01, q99 = q01.to(x.dtype), q99.to(x.dtype)
+    if mode == "q99":                       # constant dims (q01 == q99) pass through
+        live = q01 != q99
+        out = x.clone()
+    elif mode == "q99_const":               # near-constant dims are mapped to 0
+        live = (q99 - q01) > CONST_EPS
+        out = torch.zeros_like(x)
+    else:
+        raise ValueError(f"unknown norm mode {mode!r}")
+    out[..., live] = 2 * (x[..., live] - q01[live]) / (q99[live] - q01[live]) - 1
+    return out.clamp(-1, 1)
+
+
+def unnormalize(x: torch.Tensor, q01: torch.Tensor, q99: torch.Tensor, mode: str) -> torch.Tensor:
+    q01, q99 = q01.to(x.dtype), q99.to(x.dtype)
+    out = (x + 1) / 2 * (q99 - q01) + q01
+    if mode == "q99_const":
+        out = torch.where((q99 - q01) > CONST_EPS, out, ((q01 + q99) / 2).expand_as(x))
+    return out
+
+
+def group_stats(stats: dict, groups: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
+    """Concatenate the per-group q01/q99 (metadata["statistics"]["action"]) in group order."""
+    q01 = np.concatenate([stats[g]["q01"] for g in groups])
+    q99 = np.concatenate([stats[g]["q99"] for g in groups])
+    return torch.tensor(q01.tolist()), torch.tensor(q99.tolist())
+
+
+# ----------------------------------------------------------------------------- images
+def preprocess_frames(frames: np.ndarray, train: bool = False, fov_range=None,
+                      color_jitter=None) -> np.ndarray:
+    """(N, H, W, 3) uint8 -> (N, 3, 224, 224) float32, ImageNet-normalised.
+
+    All N frames (every view and both timesteps of one sample) share one random draw, so the
+    stereo geometry and the temporal pair stay consistent. Train-time augmentation: field of
+    view (random zoom-in crop or zoom-out black padding at a random offset) and colour jitter.
+    """
+    import torchvision.transforms.v2 as T
+    import torchvision.transforms.v2.functional as TF
+
+    x = torch.from_numpy(np.ascontiguousarray(frames)).to(torch.float32).div(255.0).permute(0, 3, 1, 2)
+    if train and fov_range is not None:
+        x = _random_fov(x, *fov_range)
+    x = TF.resize(x, [IMAGE_SIZE, IMAGE_SIZE], interpolation=T.InterpolationMode.BILINEAR,
+                  antialias=True)
+    if train and color_jitter is not None:
+        x = T.ColorJitter(*color_jitter)(x)
+    u8 = (x.permute(0, 2, 3, 1) * 255).to(torch.uint8).numpy()
+    # Channels-last memory (a transposed view, no copy): cuDNN picks its convolution kernel by
+    # layout, and this is the layout the IDM was trained and evaluated with.
+    return ((u8.astype(np.float32) / 255.0 - _MEAN) / _STD).transpose(0, 3, 1, 2)
+
+
+def _random_fov(x: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
+    import torchvision.transforms.v2.functional as TF
+
+    H, W = x.shape[-2:]
+    s = float(torch.empty(1).uniform_(lo, hi))
+    if s <= 1.0:
+        h, w = max(1, round(H * s)), max(1, round(W * s))
+        top, left = int(torch.randint(0, H - h + 1, (1,))), int(torch.randint(0, W - w + 1, (1,)))
+        return x[..., top:top + h, left:left + w]
+    h, w = round(H * s), round(W * s)
+    top, left = int(torch.randint(0, h - H + 1, (1,))), int(torch.randint(0, w - W + 1, (1,)))
+    return TF.pad(x, [left, top, w - W - left, h - H - top], fill=0.0)
+
+
+def model_inputs(images: np.ndarray, spec: dict, horizon: int, action_dim: int) -> dict:
+    """The per-sample fields every batch needs besides the action target."""
+    n_native = native_dim(spec)
+    mask = np.zeros((horizon, action_dim), dtype=bool)
+    mask[:, :n_native] = True
+    view_ids = list(spec["view_ids"])
+    assert len(view_ids) == len(images), f"view_ids {view_ids} for {len(images)} images"
+    return dict(images=images, view_ids=np.asarray(view_ids, np.int64), n_images=len(images),
+                embodiment_id=int(spec["embodiment_id"]), actions_mask=mask)
+
+
+def collate(samples: list[dict]) -> dict:
+    out = {}
+    for k in samples[0]:
+        vals = [s[k] for s in samples]
+        cat = k in ("images", "view_ids")
+        out[k] = torch.from_numpy(np.concatenate(vals) if cat else np.stack(vals)) \
+            if isinstance(vals[0], np.ndarray) else torch.tensor(vals)
+    return out
+
+
+# ----------------------------------------------------------------------------- LeRobot reader
+class LeRobotMeta:
+    """meta/{info,modality,episodes,stats} of a LeRobot v2.1 dataset, reduced to what the IDM
+    needs: per-group column slices of the action, video paths and episode lengths."""
+
+    def __init__(self, root: str | Path):
+        self.root = Path(root)
+        m = self.root / "meta"
+        self.info = json.loads((m / "info.json").read_text())
+        self.modality = json.loads((m / "modality.json").read_text())
+        self.episodes = [json.loads(l) for l in (m / "episodes.jsonl").read_text().splitlines() if l.strip()]
+        self.stats = json.loads((m / "stats.json").read_text())
+
+    def action_slice(self, group: str) -> tuple[str, int, int]:
+        g = self.modality["action"][group]
+        return g.get("original_key") or "action", g["start"], g["end"]
+
+    def action_statistics(self, groups: list[str]) -> dict:
+        out = {}
+        for g in groups:
+            col, a, b = self.action_slice(g)
+            out[g] = {k: list(np.asarray(v)[a:b]) for k, v in self.stats[col].items()}
+        return out
+
+    def parquet(self, ep: int) -> Path:
+        return self.root / self.info["data_path"].format(
+            episode_chunk=ep // self.info["chunks_size"], episode_index=ep)
+
+    def video(self, ep: int, view: str) -> Path:
+        key = self.modality["video"][view].get("original_key") or view
+        return self.root / self.info["video_path"].format(
+            episode_chunk=ep // self.info["chunks_size"], episode_index=ep, video_key=key)
+
+
+def native_dim(spec: dict) -> int:
+    return int(spec["action_dim"])
+
+
+def frames_at(path: Path, timestamps: np.ndarray) -> np.ndarray:
+    """Frames nearest to the given timestamps (seconds), (N, H, W, 3) uint8 RGB."""
+    import decord
+
+    vr = decord.VideoReader(str(path))
+    ts = vr.get_frame_timestamp(range(len(vr)))[:, :1]
+    return vr.get_batch(np.abs(ts - timestamps).argmin(axis=0)).asnumpy()
+
+
+class Mirror:
+    """Left/right reflection: flip the frames, exchange the two eyes, swap left/right joint
+    groups (with a per-joint sign) and negate the joints listed in ``sign``. Declared per
+    embodiment because the sign map depends on the kinematics."""
+
+    def __init__(self, cfg: dict, groups: list[str], views: list[str]):
+        self.prob = float(cfg.get("prob", 0.0))
+        self.view_perm = list(range(len(views)))
+        for a, b in cfg.get("view_swap", []):
+            ia, ib = views.index(a), views.index(b)
+            self.view_perm[ia], self.view_perm[ib] = ib, ia
+        self.swaps = [(groups.index(a), groups.index(b), s) for a, b, s in cfg.get("swap", [])]
+        self.signs = {groups.index(g): np.asarray(s) for g, s in (cfg.get("sign") or {}).items()}
+
+    def __call__(self, frames: list[np.ndarray], parts: list[np.ndarray]):
+        if self.prob <= 0 or float(torch.rand(())) >= self.prob:
+            return frames, parts
+        frames = [np.ascontiguousarray(frames[i][..., ::-1, :]) for i in self.view_perm]
+        parts = list(parts)
+        for ia, ib, s in self.swaps:
+            s = 1.0 if s is None else np.asarray(s)
+            parts[ia], parts[ib] = parts[ib] * s, parts[ia] * s
+        for i, s in self.signs.items():
+            parts[i] = parts[i] * s
+        return frames, parts
+
+
+class LeRobotIDMDataset(Dataset):
+    """Every frame t of every episode is a sample: frames (t, t+K) of all views (clamped to the
+    episode), and actions a_t .. a_{t+H-1} (the last action repeats past the end)."""
+
+    def __init__(self, spec: dict, model_cfg: dict, data_cfg: dict, train: bool = True):
+        self.spec, self.train = spec, train
+        self.meta = LeRobotMeta(spec["root"])
+        self.groups = list(spec["action_groups"])
+        self.H, self.A = int(model_cfg["action_horizon"]), int(model_cfg["action_dim"])
+        self.K = int(data_cfg["frame_gap"])
+        eps = self.meta.episodes[: data_cfg.get("max_episodes") or None]
+        self.steps = [(e["episode_index"], t, e["length"]) for e in eps for t in range(e["length"])]
+        self.statistics = self.meta.action_statistics(self.groups)
+        self.q01, self.q99 = group_stats(self.statistics, self.groups)
+        widths = [self.meta.action_slice(g)[2] - self.meta.action_slice(g)[1] for g in self.groups]
+        assert sum(widths) == native_dim(spec) <= self.A, (spec["name"], widths, self.A)
+        self.mirror = Mirror(spec.get("mirror") or {}, self.groups, list(spec["views"]))
+        self.fov = spec.get("fov_range", data_cfg.get("fov_range"))
+        self.jitter = data_cfg.get("color_jitter")
+        self._cache = (None, None)
+
+    def __len__(self):
+        return len(self.steps)
+
+    def metadata(self) -> dict:
+        return {"statistics": {"action": self.statistics}, "spec": dict(self.spec)}
+
+    def _episode(self, ep: int):
+        if self._cache[0] != ep:
+            import pandas as pd
+            self._cache = (ep, pd.read_parquet(self.meta.parquet(ep)))
+        return self._cache[1]
+
+    def __getitem__(self, i: int) -> dict:
+        ep, t, L = self.steps[i]
+        df = self._episode(ep)
+        obs = np.clip(np.array([t, t + self.K]), 0, L - 1)
+        ts = df["timestamp"].to_numpy()[obs]
+        frames = [frames_at(self.meta.video(ep, v), ts) for v in self.spec["views"]]
+        rows = np.minimum(np.arange(t, t + self.H), L - 1)
+        parts = []
+        for g in self.groups:
+            col, a, b = self.meta.action_slice(g)
+            parts.append(np.stack(df[col].to_numpy()[rows])[:, a:b].astype(np.float64))
+        if self.train:
+            frames, parts = self.mirror(frames, parts)
+        # (v, t) -> images ordered [t0 v0, t0 v1, t1 v0, t1 v1]
+        stacked = np.stack(frames, 1).reshape(-1, *frames[0].shape[1:])
+        images = preprocess_frames(stacked, self.train, self.fov, self.jitter)
+        act = normalize(torch.from_numpy(np.concatenate(parts, 1)), self.q01, self.q99,
+                        self.spec.get("norm", "q99")).numpy()
+        sample = model_inputs(images, self.spec, self.H, self.A)
+        sample["actions"] = np.pad(act, ((0, 0), (0, self.A - act.shape[1]))).astype(np.float32)
+        return sample
+
+
+class MixtureDataset(Dataset):
+    """Weighted mixture of embodiments. Each epoch, embodiment k fills round(w_k * N) of the
+    N = sum(len) slots with a fresh random permutation of its samples (repeated if up-weighted)."""
+
+    def __init__(self, datasets: list[LeRobotIDMDataset], weights: list[float], seed: int = 42):
+        self.datasets = datasets
+        w = np.asarray(weights, np.float64)
+        self.weights = w / w.sum()
+        self.seed, self.N = seed, sum(len(d) for d in datasets)
+        self.epoch, self.table = None, None
+        self.set_epoch(0)
+
+    def set_epoch(self, epoch: int) -> None:
+        if epoch == self.epoch:
+            return
+        self.epoch = epoch
+        rng = np.random.default_rng(self.seed + epoch)
+        raw = self.weights * self.N
+        slots = np.floor(raw).astype(int)
+        for i in np.argsort(-(raw - slots))[: self.N - slots.sum()]:
+            slots[i] += 1
+        rows = []
+        for k, (ds, n_slots) in enumerate(zip(self.datasets, slots)):
+            reps, rem = divmod(int(n_slots), len(ds))
+            local = np.concatenate([rng.permutation(len(ds)) for _ in range(reps)]
+                                   + [rng.permutation(len(ds))[:rem]])
+            rows.append(np.stack([np.full(n_slots, k), local], 1))
+        self.table = np.concatenate(rows).astype(np.int64)
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, i):
+        k, j = self.table[i]
+        return self.datasets[k][j]
+
+
+def embodiment_specs(data_cfg: dict) -> dict[str, dict]:
+    """Embodiments declared in the config, keyed by name, with ``data.roots`` overrides applied."""
+    roots = data_cfg.get("roots") or {}
+    specs = {}
+    for e in data_cfg["embodiments"]:
+        spec = dict(e)
+        spec["root"] = roots.get(spec["name"], spec.get("root"))
+        specs[spec["name"]] = spec
+    return specs

robotryon/idm/dit.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted from GR00T-Dreams (Apache-2.0), gr00t/model/action_head/cross_attention_dit.py.
+# Changes: plain nn.Module instead of diffusers ModelMixin/ConfigMixin, unused options removed.
+# Parameter names are unchanged so GR00T-Dreams IDM checkpoints load as-is.
+"""Diffusion transformer (DiT) of the flow-matching action head, and the self-attention
+transformer that mixes the visual tokens before the DiT cross-attends to them."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention import Attention, FeedForward
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from torch import nn
+
+
+class TimestepEncoder(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timesteps):
+        dtype = next(self.parameters()).dtype
+        return self.timestep_embedder(self.time_proj(timesteps).to(dtype))
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, embedding_dim: int, norm_eps: float = 1e-5):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, embedding_dim * 2)
+        self.norm = nn.LayerNorm(embedding_dim, norm_eps, False)
+
+    def forward(self, x, temb):
+        scale, shift = self.linear(self.silu(temb)).chunk(2, dim=1)
+        return self.norm(x) * (1 + scale[:, None]) + shift[:, None]
+
+
+class BasicTransformerBlock(nn.Module):
+    """Attention (self, or cross when encoder states are given) + feed-forward, pre-norm."""
+
+    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int,
+                 dropout: float = 0.0, activation_fn: str = "geglu", attention_bias: bool = False,
+                 upcast_attention: bool = False, norm_elementwise_affine: bool = True,
+                 norm_type: str = "layer_norm", norm_eps: float = 1e-5, final_dropout: bool = False):
+        super().__init__()
+        self.norm_type = norm_type
+        if norm_type == "ada_norm":
+            self.norm1 = AdaLayerNorm(dim)
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+        self.attn1 = Attention(query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim,
+                               dropout=dropout, bias=attention_bias, cross_attention_dim=None,
+                               upcast_attention=upcast_attention, out_bias=True)
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn,
+                              final_dropout=final_dropout, inner_dim=None, bias=True)
+        self.final_dropout = nn.Dropout(dropout) if final_dropout else None
+
+    def forward(self, hidden_states, encoder_hidden_states=None, temb=None):
+        if self.norm_type == "ada_norm":
+            h = self.norm1(hidden_states, temb)
+        else:
+            h = self.norm1(hidden_states)
+        attn = self.attn1(h, encoder_hidden_states=encoder_hidden_states, attention_mask=None)
+        if self.final_dropout:
+            attn = self.final_dropout(attn)
+        hidden_states = attn + hidden_states
+        return self.ff(self.norm3(hidden_states)) + hidden_states
+
+
+class DiT(nn.Module):
+    """Action-token transformer conditioned on the flow time (AdaLN) and on the visual tokens.
+    With ``interleave_self_attention`` the odd blocks are self-attention only."""
+
+    def __init__(self, num_attention_heads: int = 16, attention_head_dim: int = 64,
+                 output_dim: int = 1024, num_layers: int = 8, dropout: float = 0.2,
+                 attention_bias: bool = True, activation_fn: str = "gelu-approximate",
+                 upcast_attention: bool = False, norm_type: str = "ada_norm",
+                 norm_elementwise_affine: bool = False, norm_eps: float = 1e-5,
+                 final_dropout: bool = True, interleave_self_attention: bool = True):
+        super().__init__()
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.interleave_self_attention = interleave_self_attention
+        self.timestep_encoder = TimestepEncoder(self.inner_dim)
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlock(self.inner_dim, num_attention_heads, attention_head_dim,
+                                  dropout=dropout, activation_fn=activation_fn,
+                                  attention_bias=attention_bias, upcast_attention=upcast_attention,
+                                  norm_type=norm_type, norm_elementwise_affine=norm_elementwise_affine,
+                                  norm_eps=norm_eps, final_dropout=final_dropout)
+            for _ in range(num_layers)])
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+        self.proj_out_2 = nn.Linear(self.inner_dim, output_dim)
+
+    def forward(self, hidden_states, encoder_hidden_states, timestep):
+        temb = self.timestep_encoder(timestep)
+        hidden_states = hidden_states.contiguous()
+        encoder_hidden_states = encoder_hidden_states.contiguous()
+        for idx, block in enumerate(self.transformer_blocks):
+            cross = not (idx % 2 == 1 and self.interleave_self_attention)
+            hidden_states = block(hidden_states, encoder_hidden_states if cross else None, temb)
+        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        return self.proj_out_2(hidden_states)
+
+
+class SelfAttentionTransformer(nn.Module):
+    def __init__(self, num_attention_heads: int = 16, attention_head_dim: int = 64,
+                 num_layers: int = 4, dropout: float = 0.2, attention_bias: bool = True,
+                 activation_fn: str = "gelu-approximate", upcast_attention: bool = False,
+                 final_dropout: bool = True):
+        super().__init__()
+        inner = num_attention_heads * attention_head_dim
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlock(inner, num_attention_heads, attention_head_dim, dropout=dropout,
+                                  activation_fn=activation_fn, attention_bias=attention_bias,
+                                  upcast_attention=upcast_attention, final_dropout=final_dropout)
+            for _ in range(num_layers)])
+
+    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None):
+        hidden_states = hidden_states.contiguous()
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states)
+        return hidden_states

robotryon/idm/labeling.py

@@ -0,0 +1,99 @@
+"""Action labels for a generated robot video (paper Sec. 3.4, App. B.2.6).
+
+1. Sliding window, stride 1: at every frame t the IDM reads frames (t, t+K) (clamped to the
+   clip) and predicts the chunk a_t .. a_{t+H-1}.
+2. Chunk averaging: frame u appears in the H chunks that start at u-H+1 .. u; its label is the
+   mean of those predictions (fewer at the clip ends).
+3. Savitzky-Golay smoothing (polynomial order 2) per joint group: a long window on the arm-like
+   groups, a short one on the hand groups so that grasp onsets stay sharp.
+4. Optional wrist correction (``wrist_correction.py``).
+5. Proprioceptive state from the actions: q_t = a_{t-d} (the robot reaches a commanded pose
+   d frames later); the first d states repeat a_0.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from scipy.signal import savgol_filter
+
+from .data import collate, model_inputs, preprocess_frames, unnormalize
+
+
+class ChunkAccumulator:
+    """Average every chunk prediction that covers a frame."""
+
+    def __init__(self, length: int, action_dim: int, horizon: int):
+        self.L, self.H = int(length), int(horizon)
+        self._sum = np.zeros((self.L, action_dim), dtype=np.float64)
+        self._cnt = np.zeros((self.L, 1), dtype=np.float64)
+
+    def add(self, start: int, chunk: np.ndarray) -> None:
+        n = min(self.H, len(chunk), self.L - start)
+        if n > 0:
+            self._sum[start:start + n] += chunk[:n]
+            self._cnt[start:start + n] += 1.0
+
+    def result(self) -> np.ndarray:
+        if (self._cnt == 0).any():
+            raise RuntimeError("some frames are covered by no chunk")
+        return (self._sum / self._cnt).astype(np.float32)
+
+
+def smooth_actions(actions: np.ndarray, groups: dict[str, slice], hand_groups, window: int,
+                   hand_window: int, order: int = 2) -> np.ndarray:
+    """SG filter per group along time; hand groups use ``hand_window``. A window longer than the
+    clip (or not above the order) leaves the group unchanged."""
+    out = np.array(actions, dtype=np.float32, copy=True)
+    for name, sl in groups.items():
+        w = int(hand_window if name in hand_groups else window) | 1
+        if order < w <= len(out):
+            out[:, sl] = savgol_filter(out[:, sl], w, order, axis=0, mode="nearest")
+    return out
+
+
+def state_from_action(actions: np.ndarray, delay: int) -> np.ndarray:
+    """q_t = a_{t-delay}; q_t = a_0 for t < delay."""
+    a = np.asarray(actions)
+    d = min(int(delay), len(a))
+    return np.concatenate([np.repeat(a[:1], d, axis=0), a[: len(a) - d]], axis=0)
+
+
+def group_slices(spec: dict, widths: dict[str, int]) -> dict[str, slice]:
+    out, s = {}, 0
+    for g in spec["action_groups"]:
+        out[g] = slice(s, s + widths[g])
+        s += widths[g]
+    return out
+
+
+@torch.no_grad()
+def predict_actions(model, views: list[np.ndarray], spec: dict, stats: dict, frame_gap: int,
+                    batch_size: int, generator: torch.Generator | None = None) -> np.ndarray:
+    """Chunk-averaged raw actions (T, native_dim) for one clip.
+
+    ``views``: one (T, H, W, 3) uint8 RGB array per view in ``spec["views"]`` order, already at
+    the embodiment's ``video_size``. ``stats``: per-group action statistics of this embodiment.
+    """
+    H, A = model.cfg["action_horizon"], model.cfg["action_dim"]
+    groups = list(spec["action_groups"])
+    q01 = torch.tensor(np.concatenate([stats[g]["q01"] for g in groups]).tolist())
+    q99 = torch.tensor(np.concatenate([stats[g]["q99"] for g in groups]).tolist())
+    n_native = len(q01)
+    T = min(len(v) for v in views)
+    device = next(model.parameters()).device
+    acc = ChunkAccumulator(T, n_native, H)
+
+    def sample(t):
+        obs = np.clip(np.array([t, t + frame_gap]), 0, T - 1)
+        frames = np.stack([v[obs] for v in views], 1).reshape(-1, *views[0].shape[1:])
+        return model_inputs(preprocess_frames(frames), spec, H, A)
+
+    for s in range(0, T, batch_size):
+        batch = collate([sample(t) for t in range(s, min(s + batch_size, T))])
+        batch = {k: v.to(device) for k, v in batch.items()}
+        pred = model.get_action(batch, generator=generator).cpu()[..., :n_native]
+        raw = unnormalize(pred, q01, q99, spec.get("norm", "q99")).numpy()
+        for b, chunk in enumerate(raw):
+            acc.add(s + b, chunk)
+    return acc.result()

robotryon/idm/model.py

@@ -0,0 +1,327 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted from GR00T-Dreams (Apache-2.0): gr00t/model/idm.py,
+# gr00t/model/action_head/flow_matching_action_head_idm.py, multimodal_projector.py.
+# Changes: Depth Anything V2 vision tower instead of SigLIP, no hydra, plain safetensors I/O,
+# image tokens placed from a per-sample image count instead of token-id arrays.
+"""Inverse dynamics model h_psi (paper Sec. 3.4, App. B.2.6).
+
+Given two frames o_t and o_{t+K} (K = 20, 1 s at 20 fps; one or two camera views each), the IDM
+predicts the action chunk a_{t:t+K} by flow matching:
+
+    images --DINOv2 trunk of Depth Anything V2-L--> 16x16 patch tokens (CLS dropped)
+           --linear + 2x(2x2 pool MLP)--> 16 tokens per image, + view embedding
+    tokens --self-attention transformer--> visual context (left zero-padded to 112 tokens)
+    noisy chunk a^tau --per-embodiment action encoder--> DiT (AdaLN on tau, cross-attention to
+           the context) --per-embodiment decoder--> velocity v;  loss ||v - (a - eps)||^2
+
+One model is shared by all embodiments. Actions are zero-padded to ``action_dim`` (the widest
+robot) and the padded dimensions are masked out of the input and the loss; the action encoder
+and decoder have separate weights per embodiment id. At inference the chunk is integrated
+from Gaussian noise with ``num_inference_timesteps`` Euler steps.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.distributions import Beta
+
+from .dit import DiT, SelfAttentionTransformer
+
+DEFAULT_MODEL_CFG = dict(
+    vision_tower="depth-anything/Depth-Anything-V2-Large-hf",
+    vision_hidden_size=1024, hidden_size=1024,
+    action_dim=66, action_horizon=20,
+    max_num_embodiments=32, max_num_views=6, max_seq_len=1024,
+    max_vl_tokens=112, tokens_per_image=16,
+    frozen_vision_layers=[11],
+    num_inference_timesteps=16, num_timestep_buckets=1000,
+    noise_beta_alpha=1.5, noise_beta_beta=1.0, noise_s=0.999,
+    dit=dict(num_layers=8, num_attention_heads=16, attention_head_dim=64, dropout=0.2,
+             output_dim=1024, interleave_self_attention=True),
+    vl_self_attention=dict(num_layers=4, num_attention_heads=16, attention_head_dim=64,
+                           dropout=0.2),
+)
+
+
+# ----------------------------------------------------------------------------- vision tower
+class DepthAnythingTower(nn.Module):
+    """The DINOv2 trunk of Depth Anything V2 (the DPT neck and depth head are discarded).
+    Input: ImageNet-normalised 224x224 images. Output: (N, 256, D) post-layernorm patch tokens
+    of the last layer, CLS token removed."""
+
+    def __init__(self, name: str, pretrained: bool = True):
+        super().__init__()
+        from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation
+
+        if pretrained:
+            model = DepthAnythingForDepthEstimation.from_pretrained(name)
+        else:
+            model = DepthAnythingForDepthEstimation._from_config(DepthAnythingConfig.from_pretrained(name))
+        self.backbone = model.backbone
+        # Only read with bool_masked_pos, which the backbone never passes: no gradient (DDP).
+        self.backbone.embeddings.mask_token.requires_grad_(False)
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        return self.backbone(pixel_values).feature_maps[-1][:, 1:]
+
+
+class _DownSample(nn.Module):
+    """Concatenate 2x2 neighbouring tokens of a square grid: (N, h*w, C) -> (N, h*w/4, 4C)."""
+
+    def forward(self, x):
+        n, L, c = x.shape
+        h = w = int(L ** 0.5)
+        x = x.reshape(n, h, w, c)
+        if w % 2 or h % 2:
+            x = F.pad(x, (0, 0, 0, h % 2, 0, w % 2))
+            n, w, h, c = x.shape
+        x = x.view(n, w, h // 2, c * 2).permute(0, 2, 1, 3).contiguous()
+        return x.view(n, h // 2, w // 2, c * 4).reshape(n, -1, c * 4)
+
+
+class MultimodalProjector(nn.Module):
+    """``mlp_doubledownsample``: 256 patch tokens -> 64 -> 16 tokens."""
+
+    def __init__(self, mm_hidden: int, hidden: int):
+        super().__init__()
+        self.layers = nn.Sequential(
+            _DownSample(), nn.LayerNorm(mm_hidden * 4), nn.Linear(mm_hidden * 4, mm_hidden),
+            nn.GELU(), nn.Linear(mm_hidden, mm_hidden),
+            _DownSample(), nn.LayerNorm(mm_hidden * 4), nn.Linear(mm_hidden * 4, hidden),
+            nn.GELU(), nn.Linear(hidden, hidden))
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+# ----------------------------------------------------------------------------- action head
+class CategorySpecificLinear(nn.Module):
+    def __init__(self, num_categories, input_dim, hidden_dim):
+        super().__init__()
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x, cat_ids):
+        return torch.bmm(x, self.W[cat_ids]) + self.b[cat_ids].unsqueeze(1)
+
+
+class CategorySpecificMLP(nn.Module):
+    def __init__(self, num_categories, input_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x, cat_ids):
+        return self.layer2(F.relu(self.layer1(x, cat_ids)), cat_ids)
+
+
+def sinusoidal(timesteps: torch.Tensor, dim: int) -> torch.Tensor:
+    """(B, T) -> (B, T, dim) sin/cos encoding."""
+    timesteps = timesteps.float()
+    half = dim // 2
+    exponent = -torch.arange(half, dtype=torch.float, device=timesteps.device) * (
+        torch.log(torch.tensor(10000.0)) / half)
+    freqs = timesteps.unsqueeze(-1) * exponent.exp()
+    return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    def __init__(self, action_dim, hidden_size, num_embodiments):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
+
+    def forward(self, actions, timesteps, cat_ids):
+        B, T, _ = actions.shape
+        a_emb = self.W1(actions, cat_ids)
+        tau = sinusoidal(timesteps.unsqueeze(1).expand(-1, T), self.hidden_size).to(a_emb.dtype)
+        x = self.W2(torch.cat([a_emb, tau], dim=-1), cat_ids)
+        return self.W3(x * torch.sigmoid(x), cat_ids)
+
+
+class IDM(nn.Module):
+    """Flow-matching inverse dynamics model. See the module docstring.
+
+    Batch (see ``data.collate``): ``images`` (N, 3, 224, 224) for all samples concatenated in
+    order [t0 view0, t0 view1, t1 view0, ...]; ``view_ids`` (N,); ``n_images`` (B,);
+    ``embodiment_id`` (B,); ``actions_mask`` (B, H, A) bool, True on the robot's native
+    dimensions; ``actions`` (B, H, A) normalised targets (training only).
+    """
+
+    def __init__(self, cfg: dict | None = None, pretrained_vision: bool = True):
+        super().__init__()
+        c = {**DEFAULT_MODEL_CFG, **(cfg or {})}
+        self.cfg = c
+        D = c["hidden_size"]
+        self.vision = DepthAnythingTower(c["vision_tower"], pretrained=pretrained_vision)
+        self.vision_projector = nn.Linear(c["vision_hidden_size"], D)
+        self.mm_projector = MultimodalProjector(c["vision_hidden_size"], D)
+        self.view_embedding = nn.Embedding(c["max_num_views"], D)
+        self.position_embedding = nn.Embedding(c["max_seq_len"], D)
+        nn.init.normal_(self.view_embedding.weight, std=0.02)
+        nn.init.normal_(self.position_embedding.weight, std=0.02)
+        self.vl_self_attention = SelfAttentionTransformer(**c["vl_self_attention"])
+        self.dit = DiT(**c["dit"])
+        self.action_encoder = MultiEmbodimentActionEncoder(c["action_dim"], D, c["max_num_embodiments"])
+        self.action_decoder = CategorySpecificMLP(c["max_num_embodiments"], D, D, c["action_dim"])
+        self.beta_dist = Beta(c["noise_beta_alpha"], c["noise_beta_beta"])
+        # One ViT layer is kept frozen, as in the trained checkpoint.
+        for i in c["frozen_vision_layers"]:
+            self.vision.backbone.encoder.layer[i].requires_grad_(False)
+
+    # -- encoding -------------------------------------------------------------------------
+    def encode(self, batch) -> torch.Tensor:
+        """Visual context (B, max_vl_tokens, D): each sample's image tokens right-aligned,
+        zeros in front (the zeros are attended to, exactly as in training)."""
+        x = self.mm_projector(self.vision_projector(self.vision(batch["images"])))
+        x = x + self.view_embedding(batch["view_ids"]).unsqueeze(1)
+        n_tok = batch["n_images"] * self.cfg["tokens_per_image"]
+        L = self.cfg["max_vl_tokens"]
+        assert int(n_tok.max()) <= L, f"{int(n_tok.max())} image tokens > max_vl_tokens={L}"
+        pos = torch.arange(L, device=x.device)
+        mask = (pos[None] >= (L - n_tok)[:, None]).unsqueeze(-1).expand(-1, -1, x.shape[-1])
+        vl = torch.zeros(len(n_tok), L, x.shape[-1], dtype=x.dtype, device=x.device)
+        return self.vl_self_attention(vl.masked_scatter(mask, x))
+
+    def _velocity(self, context, actions, t_disc, t_dit, emb_id):
+        sa = self.action_encoder(actions, t_disc, emb_id)
+        T = sa.shape[1]
+        sa = sa + self.position_embedding(torch.arange(T, device=sa.device)).unsqueeze(0)
+        out = self.dit(sa, context, t_dit)
+        return self.action_decoder(out, emb_id)[:, -T:]
+
+    # -- training -------------------------------------------------------------------------
+    def forward(self, batch) -> torch.Tensor:
+        """Flow-matching loss over the native (unpadded) action dimensions."""
+        context = self.encode(batch)
+        actions, mask = batch["actions"], batch["actions_mask"].to(batch["actions"].dtype)
+        noise = torch.randn_like(actions)
+        s = self.cfg["noise_s"]
+        t = (s - self.beta_dist.sample([actions.shape[0]]).to(actions.device, actions.dtype)) / s
+        t = t[:, None, None]
+        noisy = ((1 - t) * noise + t * actions) * mask
+        t_disc = (t[:, 0, 0] * self.cfg["num_timestep_buckets"]).long()
+        pred = self._velocity(context, noisy, t_disc, t_disc, batch["embodiment_id"])
+        loss = F.mse_loss(pred, actions - noise, reduction="none") * mask
+        return loss.sum() / mask.sum()
+
+    # -- inference ------------------------------------------------------------------------
+    @torch.inference_mode()
+    def get_action(self, batch, generator: torch.Generator | None = None) -> torch.Tensor:
+        """Integrate the flow from noise; returns normalised, padded actions (B, H, A)."""
+        emb_id = batch["embodiment_id"]
+        images = batch["images"]
+        B, H, A = len(emb_id), self.cfg["action_horizon"], self.cfg["action_dim"]
+        x = torch.randn((B, H, A), dtype=images.dtype, device=images.device, generator=generator)
+        mask = batch["actions_mask"].to(x.dtype)
+        x = x * mask
+        context = self.encode(batch)
+        n = self.cfg["num_inference_timesteps"]
+        for i in range(n):
+            t = int(i / float(n) * self.cfg["num_timestep_buckets"])
+            v = self._velocity(context, x, (torch.ones(B) * t).to(x.device),
+                               torch.tensor([t], device=x.device, dtype=torch.long), emb_id)
+            x = (x + (1.0 / n) * v) * mask
+        return x
+
+    # -- I/O ------------------------------------------------------------------------------
+    def save_pretrained(self, out_dir: str | Path, metadata: dict | None = None) -> None:
+        from safetensors.torch import save_file
+
+        out = Path(out_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        sd = {k: v.detach().contiguous().cpu() for k, v in self.state_dict().items()}
+        save_file(sd, str(out / "model.safetensors"), metadata={"format": "pt"})
+        (out / "idm_config.json").write_text(json.dumps(self.cfg, indent=1))
+        if metadata is not None:
+            (out / "metadata.json").write_text(json.dumps(metadata, indent=1))
+
+    @classmethod
+    def from_pretrained(cls, ckpt_dir: str | Path, device="cpu") -> "IDM":
+        """Load a checkpoint written by ``save_pretrained`` or by the GR00T-Dreams IDM trainer."""
+        from safetensors.torch import load_file
+
+        d = Path(ckpt_dir)
+        if (d / "idm_config.json").exists():
+            cfg = json.loads((d / "idm_config.json").read_text())
+        else:
+            cfg = _legacy_config(json.loads((d / "config.json").read_text()))
+        model = cls(cfg, pretrained_vision=False)
+        sd = load_file(str(d / "model.safetensors"))
+        if any(k.startswith("action_head.") for k in sd):
+            sd = _legacy_state_dict(sd)
+        model.load_state_dict(sd, strict=True)
+        return model.to(device).eval()
+
+
+def load_metadata(ckpt_dir: str | Path) -> dict:
+    """Per-embodiment normalisation statistics saved with the checkpoint, keyed by name."""
+    d = Path(ckpt_dir)
+    for p in (d / "metadata.json", d / "experiment_cfg" / "metadata.json"):
+        if p.exists():
+            return json.loads(p.read_text())
+    raise FileNotFoundError(f"no metadata.json in {d}")
+
+
+def _legacy_config(c: dict) -> dict:
+    h = c["action_head_cfg"]["config"]
+    keep = ("num_layers", "num_attention_heads", "attention_head_dim", "dropout")
+    return dict(
+        vision_tower=h["siglip_model_cfg"]["pretrained_model_name_or_path"],
+        vision_hidden_size=h["siglip_hidden_size"], hidden_size=h["hidden_size"],
+        action_dim=h["action_dim"], action_horizon=h["action_horizon"],
+        max_num_embodiments=h.get("max_num_embodiments", 32), max_num_views=h["max_num_views"],
+        max_seq_len=h.get("max_seq_len", 1024),
+        num_inference_timesteps=h["num_inference_timesteps"],
+        num_timestep_buckets=h["num_timestep_buckets"],
+        noise_beta_alpha=h["noise_beta_alpha"], noise_beta_beta=h["noise_beta_beta"],
+        noise_s=h["noise_s"],
+        dit={**{k: h["diffusion_model_cfg"][k] for k in keep},
+             "output_dim": h["diffusion_model_cfg"]["output_dim"],
+             "interleave_self_attention": h["diffusion_model_cfg"]["interleave_self_attention"]},
+        vl_self_attention={k: h["vl_self_attention_cfg"][k] for k in keep},
+    )
+
+
+_LEGACY_RENAMES = (
+    (r"^action_head\.siglip_model\.vision_model\.backbone\.", "vision.backbone."),
+    (r"^action_head\.model\.", "dit."),
+    (r"^action_head\.vl_self_attention_model\.", "vl_self_attention."),
+    (r"^action_head\.", ""),
+)
+
+
+def _legacy_state_dict(sd: dict) -> dict:
+    out = {}
+    for k, v in sd.items():
+        if k.endswith(("siglip_model.logit_scale", "siglip_model.logit_bias")):
+            continue
+        for pat, rep in _LEGACY_RENAMES:
+            k2 = re.sub(pat, rep, k)
+            if k2 != k:
+                k = k2
+                break
+        out[k] = v
+    return out

robotryon/idm/wrist_correction.py

@@ -0,0 +1,229 @@
+"""Wrist correction of IDM action labels (paper Sec. 3.4, App. B.2.6).
+
+The IDM reads joint angles from pixels and places the wrist less accurately than the 3D hand
+reconstruction of the human source video. For each hand we therefore move the wrist, and only
+the wrist, toward the reconstruction:
+
+1. The reconstructed human wrist (camera frame) is mapped to the robot base frame with the
+   head-camera extrinsics T_base_cam.
+2. The robot hand implied by the IDM joints (FK) is scaled by ``hand_scale``
+   about its wrist; the target wrist is the one at which the scaled hand's end-effector keypoint
+   coincides with the reconstructed wrist.
+3. Targets are kept only on high-confidence frames. The offset d = target - IDM wrist is
+   cleaned: excursion rejection (hysteresis around a running median of accepted offsets),
+   linear fill of short gaps, zero-phase 2nd-order Butterworth low-pass per usable segment.
+4. Damped least squares IK on the arm joints of that side: wrist position to the IDM wrist + d,
+   wrist orientation held at the FK orientation of the IDM arm joints (weighted rotation
+   residual), damping toward the IDM arm joints, numerical Jacobian of FK, warm-started from
+   the previous frame of the segment.
+5. The solution is blended into the IDM joints with weights ramping 0 -> 1 over ``ramp`` frames
+   at segment edges (not at the clip ends), then the arm tracks get SG smoothing. Hand joints
+   and all other joints are left untouched.
+
+Kinematics interface required (minimal protocol; ``robotryon.embodiment.Embodiment`` is
+expected to satisfy it)::
+
+    class Kinematics(Protocol):
+        arm_joint_slices: dict[str, slice]   # {"left": slice, "right": slice} into the action
+        def fk(self, q: np.ndarray) -> dict[str, dict[str, np.ndarray]]:
+            # {"left": {"wrist_pos": (3,), "wrist_rot": (3, 3), "keypoints": (17, 3)}, "right": ...}
+            # robot base frame, metres; keypoints use robotryon.skeleton order (0 = wrist)
+
+Hands are indexed 0 = left, 1 = right, as in ``robotryon.skeleton``.
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+from scipy.signal import butter, filtfilt, savgol_filter
+
+SIDES = ("left", "right")
+
+DEFAULTS = dict(
+    hand_scale=1.0, ee_keypoint=0, min_confidence=0.5,
+    reject_enter_mm=150.0, reject_exit_mm=80.0, reject_hard_mm=500.0, reject_ref_n=15,
+    gap_fill=6, cutoff_hz=2.0, fps=20.0, ramp=8, smooth_window=9, smooth_order=2,
+    ik_iters=30, w_rot=2000.0, damping=0.5, step_cap=0.08, fd_eps=1e-3,
+)
+
+
+# ----------------------------------------------------------------------------- track cleaning
+def segments(mask: np.ndarray) -> list[tuple[int, int]]:
+    """Half-open [a, b) runs of True."""
+    out, a = [], None
+    for i, m in enumerate(mask):
+        if m and a is None:
+            a = i
+        if not m and a is not None:
+            out.append((a, i))
+            a = None
+    if a is not None:
+        out.append((a, len(mask)))
+    return out
+
+
+def reject_excursions(d, ok, enter, exit_, hard, ref_n=15):
+    """Refuse frames whose offset leaves the level it has been holding (running median of the
+    last ``ref_n`` accepted offsets); a refusal lasts until the offset is back within ``exit_``.
+    ``hard`` caps |d| before a level is established."""
+    good, buf, ref = ok.copy(), [], None
+    for t in range(len(d)):
+        if not ok[t]:
+            continue
+        if np.linalg.norm(d[t]) > hard:
+            good[t] = False
+            continue
+        if ref is None:
+            buf.append(d[t])
+            if len(buf) >= 5:
+                ref = np.median(np.array(buf[-ref_n:]), axis=0)
+            continue
+        dev = float(np.linalg.norm(d[t] - ref))
+        if dev > enter or (dev > exit_ and not good[max(t - 1, 0)]):
+            good[t] = False
+        else:
+            buf.append(d[t])
+            ref = np.median(np.array(buf[-ref_n:]), axis=0)
+    return good
+
+
+def fill_short_gaps(track, ok, max_gap):
+    """Linearly interpolate refused runs of <= max_gap frames between usable frames (in place);
+    returns the new usable mask."""
+    ok = ok.copy()
+    for a, b in segments(~ok):
+        if a == 0 or b == len(ok) or b - a > max_gap:
+            continue
+        for k in range(track.shape[1]):
+            track[a:b, k] = np.interp(np.arange(a, b), [a - 1, b], [track[a - 1, k], track[b, k]])
+        ok[a:b] = True
+    return ok
+
+
+def lowpass(track, ok, cutoff_hz, fps):
+    """Zero-phase Butterworth per usable segment (segments too short for filtfilt are kept)."""
+    b, a = butter(2, cutoff_hz / (fps / 2.0), "low")
+    out = track.copy()
+    for lo, hi in segments(ok):
+        if hi - lo > 3 * max(len(a), len(b)) + 2:
+            out[lo:hi] = np.stack([filtfilt(b, a, track[lo:hi, c]) for c in range(track.shape[1])], 1)
+    return out
+
+
+def ramp_weights(ok, ramp):
+    """1 inside usable segments, ramping linearly at segment edges that are not clip ends."""
+    T, w = len(ok), np.zeros(len(ok))
+    for a, b in segments(ok):
+        for t in range(a, b):
+            up = 1.0 if a == 0 else (t - a + 1) / ramp
+            dn = 1.0 if b == T else (b - t) / ramp
+            w[t] = min(1.0, up, dn)
+    return w
+
+
+# ----------------------------------------------------------------------------- IK
+def rotvec(R_cur, R_ref):
+    """Axis-angle of R_cur R_ref^T."""
+    from scipy.spatial.transform import Rotation
+
+    return Rotation.from_matrix(R_cur @ R_ref.T).as_rotvec()
+
+
+def ik_wrist(fk: Callable, q, side, sl, target_mm, R_ref, q_ref, iters=30, w_rot=2000.0,
+             damping=0.5, step_cap=0.08, fd_eps=1e-3):
+    """Damped least squares on the joints ``sl``: wrist position (mm) to ``target_mm``,
+    orientation pinned to ``R_ref`` (``w_rot`` mm per rad), damping toward ``q_ref``.
+    Returns (q, position residual in mm before the last step)."""
+    s = np.array(q, dtype=np.float64)
+    n = sl.stop - sl.start
+
+    def res(x):
+        h = fk(x)[side]
+        return np.concatenate([np.asarray(h["wrist_pos"]) * 1000.0 - target_mm,
+                               w_rot * rotvec(np.asarray(h["wrist_rot"]), R_ref)])
+
+    lam, last = damping * 1e4, np.inf
+    for _ in range(iters):
+        r = res(s)
+        if not np.all(np.isfinite(r)):
+            return np.array(q, dtype=np.float64), np.inf
+        last = float(np.linalg.norm(r[:3]))
+        J = np.zeros((6, n))
+        for j in range(n):
+            s2 = s.copy()
+            s2[sl.start + j] += fd_eps
+            J[:, j] = (res(s2) - r) / fd_eps
+        if not np.all(np.isfinite(J)):
+            return np.array(q, dtype=np.float64), np.inf
+        dq = np.linalg.solve(J.T @ J + (1.0 + lam) * np.eye(n), -J.T @ r - lam * (s[sl] - q_ref[sl]))
+        norm = np.linalg.norm(dq)
+        if norm > step_cap:
+            dq *= step_cap / norm
+        s[sl] += dq
+        if norm < 1e-4:
+            break
+    return s, last
+
+
+# ----------------------------------------------------------------------------- main entry
+def camera_to_base(points_cam: np.ndarray, T_base_cam: np.ndarray) -> np.ndarray:
+    """(T, 2, 3) camera-frame points -> base frame; T_base_cam is (4, 4) or (T, 4, 4)."""
+    Tm = np.broadcast_to(np.asarray(T_base_cam, np.float64), (len(points_cam), 4, 4))
+    return np.einsum("tij,thj->thi", Tm[:, :3, :3], points_cam) + Tm[:, None, :3, 3]
+
+
+def correct_wrists(actions: np.ndarray, kin, wrist_base: np.ndarray, confidence: np.ndarray,
+                   cfg: dict | None = None) -> tuple[np.ndarray, dict]:
+    """Correct the arm joints of ``actions`` (T, D) so the wrists follow ``wrist_base``.
+
+    wrist_base: (T, 2, 3) reconstructed human wrists in the robot base frame (metres, NaN when
+    missing); confidence: (T, 2). Returns the corrected actions and per-side statistics.
+    """
+    c = {**DEFAULTS, **(cfg or {})}
+    act = np.asarray(actions, np.float64)
+    T = len(act)
+    out = act.copy()
+    fk_idm = [kin.fk(act[t]) for t in range(T)]
+    info = {}
+    for hi, side in enumerate(SIDES):
+        sl = kin.arm_joint_slices[side]
+        sl = sl if isinstance(sl, slice) else slice(*sl)
+        p_idm = np.stack([np.asarray(f[side]["wrist_pos"]) for f in fk_idm]) * 1000.0
+        R_idm = [np.asarray(f[side]["wrist_rot"]) for f in fk_idm]
+        # end effector of the scaled IDM hand, relative to its wrist
+        ee = np.stack([np.asarray(f[side]["keypoints"])[c["ee_keypoint"]] for f in fk_idm]) * 1000.0
+        target = wrist_base[:, hi] * 1000.0 - c["hand_scale"] * (ee - p_idm)
+        d = target - p_idm
+        ok = (confidence[:, hi] >= c["min_confidence"]) & np.isfinite(d).all(1)
+        n_conf = int(ok.sum())
+        d = np.where(ok[:, None], d, 0.0)
+        ok = reject_excursions(d, ok, c["reject_enter_mm"], c["reject_exit_mm"],
+                               c["reject_hard_mm"], c["reject_ref_n"])
+        n_rej = n_conf - int(ok.sum())
+        ok = fill_short_gaps(d, ok, c["gap_fill"])
+        d = lowpass(d, ok, c["cutoff_hz"], c["fps"])
+        q_ik, prev, resid = act[:, sl].copy(), None, []
+        for t in range(T):
+            if not ok[t]:
+                prev = None
+                continue
+            cur = act[t].copy()
+            if prev is not None:
+                cur[sl] = prev
+            new, r = ik_wrist(kin.fk, cur, side, sl, p_idm[t] + d[t], R_idm[t], act[t],
+                              c["ik_iters"], c["w_rot"], c["damping"], c["step_cap"], c["fd_eps"])
+            if not np.isfinite(r):
+                ok[t], prev = False, None
+                continue
+            q_ik[t] = prev = new[sl]
+            resid.append(r)
+        w = ramp_weights(ok, c["ramp"])
+        blended = act[:, sl] + w[:, None] * (q_ik - act[:, sl])
+        if T > c["smooth_window"]:
+            blended = savgol_filter(blended, c["smooth_window"], c["smooth_order"], axis=0, mode="nearest")
+        out[:, sl] = blended
+        info[side] = dict(confident_frames=n_conf, rejected=n_rej, corrected_frames=int((w > 0).sum()),
+                          ik_residual_mm_median=float(np.median(resid)) if resid else None)
+    return out.astype(np.asarray(actions).dtype), info

robotryon/image/__init__.py

@@ -0,0 +1,14 @@
+"""Keyframe translation f_theta (paper Sec. 3.2): FLUX.2 [klein] 9B Base as a reference-guided
+inpainting model.
+
+Modules:
+  fill_channels  widened patch embedding, channel-concatenated [noisy | conditioned frame | mask]
+  stereo_ids     RoPE ids with the camera view on the 4th axis
+  episodes       data layout access and model-resolution geometry
+  dataset        training samples (reference window, robot-only exemplar, masks, skeleton)
+  text           prompts, caption dropout and the text-embedding cache
+  training       flow-matching loss
+  model          FFT / LoRA build, save and load
+  sampler        true-CFG Euler sampler returning latent tokens
+  ar_keyframes   autoregressive translation with latent reuse
+"""

robotryon/image/ar_keyframes.py

@@ -0,0 +1,161 @@
+"""Autoregressive keyframe translation (paper Eq. 3.2, App. B.2.3).
+
+Human keyframes are translated in temporal order. Keyframe t_j is conditioned on
+
+* the conditioned human frame Phi(o^h; M^h, S^h): the human arm mask (convex hull per arm,
+  fixed dilation, eroded object removed) blacked out and the retargeted robot-hand skeleton drawn
+  on top, plus the mask itself as an input channel;
+* the reference: for the first keyframe, the encoded robot-only reference image supplied by the
+  user (it selects the embodiment); for every later keyframe, the DENOISED LATENT TOKENS of the
+  previous keyframe, reused directly as reference tokens (no decode / re-encode).
+
+``ref_skeleton: true`` replaces latent reuse by a pixel-space reference: the previous keyframe is
+decoded, the current keyframe's skeleton is painted on it, and the result is re-encoded (the first
+keyframe's reference image is painted the same way).
+
+Stereo episodes are translated jointly (both views in one token sequence) by default; ``views:
+mono`` translates each view as its own chain.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ..masks import MaskAug, inference_mask
+from ..skeleton import condition_frame, draw_skeleton, paint
+from .episodes import Episode, FrameReader, load_view_frame
+from .fill_channels import pack
+from .sampler import LatentCodec, ViewCondition, denoise
+from ..captions import HandStateCuts, hand_state_arrays, hand_state_prompt, load_hand_state, words_at
+
+
+@dataclass
+class TranslatedKeyframe:
+    image: np.ndarray        # uint8 (H, W, 3) robot keyframe
+    cond: np.ndarray         # uint8 (H, W, 3) conditioned human frame
+    mask: np.ndarray         # uint8 (H, W) {0, 1} inference mask
+    tokens: torch.Tensor     # (1, h*w, 128) denoised latent tokens (next keyframe's reference)
+
+
+class KeyframeTranslator:
+    def __init__(self, pipe, cfg, device="cuda"):
+        self.pipe, self.cfg, self.device = pipe, cfg, device
+        self.codec = LatentCodec(pipe.vae, device)
+        self.reader = FrameReader()
+        self.dilate = MaskAug.from_config(cfg.masks).infer_dilate()
+        self.ref_skeleton = bool(cfg.get("ref_skeleton", False))
+        self._prompts: dict[str, tuple] = {}
+
+    def prompt_fn(self, ep: Episode, episode_dir: Path):
+        """``prompt(t)``: the instruction and each hand's closure / facing words at keyframe t
+        (``robotryon.captions``), read from the episode's ``hand_state_<view>.npz`` (or computed
+        from ``hands.npz`` when absent)."""
+        cuts = HandStateCuts.from_config(self.cfg.caption.get("cuts"))
+        states = []
+        for i, v in enumerate(ep.views):
+            f = episode_dir / f"hand_state_{v}.npz"
+            if f.exists():
+                states.append(load_hand_state(f))
+        if not states:
+            z = np.load(episode_dir / "hands.npz")
+            states = [hand_state_arrays(z["joints3d"][i], z["valid"][i].astype(bool))
+                      for i in range(z["joints3d"].shape[0])]
+        return lambda t: hand_state_prompt(*words_at(states, int(t), cuts))
+
+    def encode_prompt(self, text: str):
+        if text not in self._prompts:
+            self._prompts[text] = self.pipe.encode_prompt(prompt=text, device=self.device)
+        return self._prompts[text]
+
+    def condition(self, ep: Episode, view: str, t: int):
+        """(frame, conditioned frame, mask, skeleton image) of one human keyframe at model resolution."""
+        c = self.cfg
+        f = load_view_frame(self.reader, ep, view, t, int(c.image.long_edge),
+                            tuple(c.image.aspect) if c.image.get("aspect") else None,
+                            c.image.get("crop_anchor", "bottom"))
+        H, W = f.image.shape[:2]
+        mask = inference_mask(f.hulls, W, H, self.dilate, f.object_mask, int(c.masks.obj_erode))
+        skel = draw_skeleton(f.kp, f.vis, W, H)
+        cond = condition_frame(f.image, mask, skel)
+        return f.image, cond, mask, skel
+
+    def reference(self, view: str, prev: dict, references: dict, skel: np.ndarray) -> torch.Tensor:
+        """Reference tokens of one view for the current keyframe."""
+        if view in prev and not self.ref_skeleton:
+            return prev[view].tokens                       # latent reuse
+        if view in prev:
+            img = prev[view].image
+        else:
+            H, W = skel.shape[:2]
+            img = np.asarray(Image.fromarray(references[view]).resize((W, H), Image.LANCZOS))
+        return pack(self.codec.encode(paint(img, skel) if self.ref_skeleton else img))
+
+    @torch.no_grad()
+    def translate_chain(self, views: list[str], keyframes: list[int], references: dict[str, np.ndarray],
+                        prompt, conditions) -> dict:
+        """One chain over ``keyframes``; ``views`` are denoised jointly.
+
+        ``prompt`` is a string or ``prompt(t)`` giving each keyframe its own prompt.
+        ``conditions(view, t)`` returns ``(frame, conditioned frame, mask, skeleton image)``.
+        Returns ``{view: {t: TranslatedKeyframe}}``.
+        """
+        c = self.cfg
+        prompt_at = prompt if callable(prompt) else (lambda t: prompt)
+        neg, neg_ids = self.encode_prompt("") if float(c.guidance_scale) > 1.0 else (None, None)
+        prev: dict[str, TranslatedKeyframe] = {}
+        out: dict[str, dict] = {v: {} for v in views}
+        for t in keyframes:
+            conds, frames = [], []
+            for v in views:
+                _, cond, mask, skel = conditions(v, t)
+                mask_t = torch.from_numpy(mask.astype(np.float32))[None, None].to(self.device)
+                conds.append(ViewCondition(self.codec.encode(cond), mask_t,
+                                           self.reference(v, prev, references, skel)))
+                frames.append((cond, mask))
+            pos, pos_ids = self.encode_prompt(prompt_at(t))
+            gen = torch.Generator(device=self.device).manual_seed(int(c.seed) + int(t))
+            tokens = denoise(self.pipe.transformer, self.pipe.scheduler, conds, pos, pos_ids, neg, neg_ids,
+                             steps=int(c.steps), guidance_scale=float(c.guidance_scale), generator=gen)
+            for v, tok, (cond, mask) in zip(views, tokens, frames):
+                H, W = mask.shape
+                img = self.codec.decode_tokens(tok, H // 16, W // 16)
+                out[v][t] = prev[v] = TranslatedKeyframe(img, cond, mask, tok)
+        return out
+
+    def translate_episode(self, episode_dir: str | Path, references: dict[str, np.ndarray],
+                          out_dir: str | Path) -> Path:
+        """Translate every keyframe of a preprocessed human episode and write the results."""
+        c = self.cfg
+        ep = Episode.load(episode_dir)
+        kf = json.loads((Path(episode_dir) / "keyframes.json").read_text())["indices"]
+        keyframes = sorted({int(i) for i in kf} | {0})
+        stereo = len(ep.views) >= 2 and c.views in ("auto", "stereo")
+        if c.views == "stereo" and not stereo:
+            raise ValueError(f"{episode_dir}: views=stereo needs two camera views, found {ep.views}")
+        chains = [ep.views[:2]] if stereo else [[v] for v in ep.views]
+        missing = [v for v in ep.views if v not in references]
+        if missing:
+            raise ValueError(f"no reference image for view(s) {missing}")
+        prompt_at = self.prompt_fn(ep, Path(episode_dir))
+        prompts = {int(t): prompt_at(t) for t in keyframes}
+        out = Path(out_dir) / ep.path.name
+        out.mkdir(parents=True, exist_ok=True)
+        for views in chains:
+            res = self.translate_chain(views, keyframes, references, prompt_at,
+                                       lambda v, t: self.condition(ep, v, t))
+            for v, per_t in res.items():
+                for t, k in per_t.items():
+                    Image.fromarray(k.image).save(out / f"keyframe_{v}_{t:06d}.png")
+                    Image.fromarray(k.cond).save(out / f"cond_{v}_{t:06d}.png")
+                    Image.fromarray(k.mask * 255).save(out / f"mask_{v}_{t:06d}.png")
+        (out / "keyframes.json").write_text(json.dumps({
+            "indices": keyframes, "views": ep.views, "joint_stereo": stereo, "prompts": prompts,
+            "ref_skeleton": self.ref_skeleton, "steps": int(c.steps),
+            "guidance_scale": float(c.guidance_scale),
+            "seed": int(c.seed)}, indent=1))
+        return out

robotryon/image/dataset.py

@@ -0,0 +1,227 @@
+"""Training samples for the keyframe-translation model (paper Sec. 3.2, App. B.2.3).
+
+A sample is a robot frame o_t of one episode, seen by one camera (mono) or both (stereo):
+
+* target   o_t
+* cond     Phi(o_t; M_t, S_t): the randomised shape-agnostic mask blacked out, the augmented
+           hand skeleton (rendered on the fly from keypoints) drawn on top
+* mask     M_t
+* reference o_t' with t' drawn uniformly from [t - 5 s, t - 2 s] of the same episode (frame 0 if
+           the target is too early). With probability ``p_robot_only`` the reference is instead a
+           "robot-only" exemplar: a frame of ANOTHER episode of the same embodiment with
+           everything but the robot (its tight mask, eroded) filled with white. This is
+           the reference the first keyframe gets at inference.
+* prompt   the hand-state prompt of the target frame (each hand's closure / facing words from the
+           episode's ``hand_state_<view>.npz``; see ``robotryon.captions``), as a cached
+           embedding.
+
+With probability ``ref_skeleton_p`` the target's (augmented) skeleton is also painted on the
+reference, sequential or robot-only alike; otherwise the reference is left clean.
+
+The two views of a stereo sample share every random draw (mask dilations, skeleton corruption,
+reference frame, field-of-view crop): they are two views of one instant.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ..masks import MaskAug, fov_augment, robot_only
+from ..skeleton import SkeletonAugment, condition_frame, draw_skeleton, paint
+from .episodes import Episode, FrameReader, crop_resize, load_view_frame, output_geometry, scan_root
+from ..captions import HandStateCuts, all_prompts, load_hand_state, sample_prompt, words_at
+from .text import TextEmbedCache
+
+
+def _to_tensor(img: Image.Image) -> torch.Tensor:
+    a = np.asarray(img, dtype=np.float32)
+    if a.ndim == 2:
+        return torch.from_numpy(a / 255.0)[None]
+    return torch.from_numpy(a / 127.5 - 1.0).permute(2, 0, 1).contiguous()
+
+
+class KeyframeDataset(torch.utils.data.Dataset):
+    """Indexed by ``(episode_index, frame, view_indices, seed)`` tuples from :class:`FillBatchSampler`."""
+
+    def __init__(self, cfg, text_cache: TextEmbedCache | None):
+        d = cfg.data
+        roots = d.roots if isinstance(d.roots, (list, tuple)) else [d.roots]
+        self.root_weights = []
+        self.episodes: list[Episode] = []
+        for i, r in enumerate(roots):
+            path, w = (r.path, float(r.get("weight", 1.0))) if isinstance(r, dict) else (r, 1.0)
+            eps = scan_root(path, i)
+            if not eps:
+                raise FileNotFoundError(f"no episodes (*/meta.json) under {path}")
+            self.episodes += eps
+            self.root_weights.append(w)
+        self.long_edge = int(cfg.image.long_edge)
+        self.aspect = tuple(cfg.image.aspect) if cfg.image.get("aspect") else None
+        self.anchor = cfg.image.get("crop_anchor", "bottom")
+        self.ref_window = tuple(d.ref_window_sec)
+        self.p_robot_only = float(d.p_robot_only)
+        self.ref_skeleton_p = float(d.ref_skeleton_p)
+        self.robot_only_erode = int(d.robot_only.erode)
+        self.robot_only_bg = int(d.robot_only.background)
+        self.robot_only_min_area = float(d.robot_only.min_area)
+        self.bonus_tags = set(cfg.masks.get("bonus_embodiments", []) or [])
+        self.mask_aug = MaskAug.from_config(cfg.masks)
+        self.fov = cfg.masks.fov
+        self.skel_aug = SkeletonAugment.from_config(cfg.skeleton.augment)
+        self.thickness = tuple(cfg.skeleton.thickness)
+        self.cap = cfg.caption
+        self.text_cache = text_cache
+        self.reader: FrameReader | None = None
+        self.by_embodiment: dict[str, list[int]] = {}
+        for i, ep in enumerate(self.episodes):
+            self.by_embodiment.setdefault(ep.embodiment, []).append(i)
+
+    def __len__(self) -> int:
+        return len(self.episodes)
+
+    def hand_states(self, ep: Episode, views: list[str]) -> list:
+        return [load_hand_state(ep.path / f"hand_state_{v}.npz") for v in views
+                if (ep.path / f"hand_state_{v}.npz").exists()]
+
+    def all_prompts(self) -> set[str]:
+        return set(all_prompts())
+
+    # -- reference ------------------------------------------------------------------------------
+    def reference_index(self, ep: Episode, t: int, rng: np.random.Generator) -> int:
+        lo, hi = self.ref_window
+        newest = t - int(round(lo * ep.fps))
+        oldest = max(0, t - int(round(hi * ep.fps)))
+        if newest < oldest:
+            return 0
+        return int(rng.integers(oldest, newest + 1))
+
+    def draw_exemplar(self, ep_idx: int, views: list[str], rng) -> tuple[int, int] | None:
+        """(episode, frame) of another episode of the same embodiment with a visible robot."""
+        ep = self.episodes[ep_idx]
+        pool = [i for i in self.by_embodiment.get(ep.embodiment, [])
+                if i != ep_idx and all(v in self.episodes[i].views for v in views)]
+        if not pool:
+            return None
+        ex = int(pool[int(rng.integers(0, len(pool)))])
+        ex_ep = self.episodes[ex]
+        for _ in range(8):
+            t = int(rng.integers(0, ex_ep.num_frames))
+            m = self.reader.binary(ex_ep.mask(views[0]), t)
+            if m.mean() >= self.robot_only_min_area:
+                return ex, t
+        return None
+
+    def robot_only(self, ex: int, t: int, view: str, size) -> np.ndarray:
+        ep = self.episodes[ex]
+        rgb = self.reader.rgb(ep.video(view), t)
+        src = (rgb.shape[1], rgb.shape[0])
+        box, _ = output_geometry(src[0], src[1], self.long_edge, self.aspect, self.anchor)
+        img = crop_resize(rgb, box, size, src)
+        m = crop_resize(self.reader.binary(ep.mask(view), t).astype(np.uint8) * 255, box, size, src)
+        return robot_only(img, m > 127, self.robot_only_erode, self.robot_only_bg)
+
+    # -- one view -------------------------------------------------------------------------------
+    def view_streams(self, ep: Episode, view: str, t: int, ref_t: int, exemplar, ref_skeleton: bool,
+                     rng) -> dict:
+        f = load_view_frame(self.reader, ep, view, t, self.long_edge, self.aspect, self.anchor)
+        H, W = f.image.shape[:2]
+        mask, _ = self.mask_aug.train_mask(f.hulls, W, H, rng, bonus=ep.embodiment in self.bonus_tags,
+                                           object_mask=f.object_mask)
+        kp, vis = self.skel_aug(f.kp, f.vis, rng)
+        th = int(rng.integers(self.thickness[0], self.thickness[1] + 1))
+        skel = draw_skeleton(kp, vis, W, H, thickness=th, radius=max(1, int(round(th * 1.5))))
+        cond = condition_frame(f.image, mask, skel)
+        if exemplar is not None:
+            ref = self.robot_only(exemplar[0], exemplar[1], view, (W, H))
+        else:
+            rgb = self.reader.rgb(ep.video(view), ref_t)
+            src = (rgb.shape[1], rgb.shape[0])
+            box, _ = output_geometry(src[0], src[1], self.long_edge, self.aspect, self.anchor)
+            ref = crop_resize(rgb, box, (W, H), src)
+        if ref_skeleton:
+            ref = paint(ref, skel)
+        return {"target": f.image, "cond": cond, "mask": mask * 255, "reference": ref,
+                "has_hull": bool(f.hulls)}
+
+    def __getitem__(self, item) -> dict:
+        ep_idx, t, view_idx, seed = item
+        if self.reader is None:
+            self.reader = FrameReader()
+        ep = self.episodes[ep_idx]
+        views = [ep.views[i] for i in view_idx]
+        rng = np.random.default_rng([int(seed), 0])
+        exemplar = None
+        if float(rng.random()) < self.p_robot_only:
+            exemplar = self.draw_exemplar(ep_idx, views, rng)
+        ref_skeleton = float(rng.random()) < self.ref_skeleton_p
+        for attempt in range(8):           # prefer frames where the arm is visible
+            ref_t = self.reference_index(ep, t, rng)
+            out = [self.view_streams(ep, v, t, ref_t, exemplar, ref_skeleton,
+                                     np.random.default_rng([int(seed), 1, attempt])) for v in views]
+            if out[0]["has_hull"]:
+                break
+            t = int(rng.integers(0, ep.num_frames))
+        keys = ("target", "cond", "mask", "reference")
+        imgs = [Image.fromarray(np.asarray(o[k], np.uint8)) for o in out for k in keys]
+        imgs, _ = fov_augment(imgs, rng, float(self.fov.crop_min), float(self.fov.scale_max),
+                              float(self.fov.get("flip_prob", 0.0)))
+        n = len(keys)
+        sample = {k: torch.stack([_to_tensor(imgs[v * n + j]) for v in range(len(views))])
+                  for j, k in enumerate(keys)}
+        words = words_at(self.hand_states(ep, views), t, HandStateCuts.from_config(self.cap.get("cuts")))
+        prompt = sample_prompt(words, rng, float(self.cap.p_uncond), float(self.cap.p_drop_hands))
+        sample["prompt"] = prompt
+        if self.text_cache is not None:
+            sample["prompt_embeds"] = self.text_cache.load(prompt)
+        sample["robot_only"] = exemplar is not None
+        sample["ref_skeleton"] = ref_skeleton
+        return sample
+
+
+def collate(batch: list[dict]) -> dict:
+    out = {k: torch.stack([b[k] for b in batch]) for k in ("target", "cond", "mask", "reference")}
+    if "prompt_embeds" in batch[0]:
+        out["prompt_embeds"] = torch.stack([b["prompt_embeds"] for b in batch])
+    out["prompt"] = [b["prompt"] for b in batch]
+    out["robot_only"] = [b["robot_only"] for b in batch]
+    return out
+
+
+class FillBatchSampler(torch.utils.data.Sampler):
+    """Endless stream of homogeneous batches for one rank.
+
+    A batch is all-stereo or all-mono (the token sequence length depends on it): stereo with
+    probability ``stereo_p`` (paper: stereo : mono = 3 : 1) when stereo episodes exist. Items pick
+    a data root by weight, an episode uniformly within it, a frame uniformly, and for mono a view
+    uniformly. Each rank draws its own stream from ``(seed, rank)``.
+    """
+
+    def __init__(self, dataset: KeyframeDataset, batch_size: int, stereo_p: float, seed: int,
+                 rank: int = 0):
+        self.ds, self.bs, self.stereo_p = dataset, int(batch_size), float(stereo_p)
+        self.rng = np.random.default_rng([int(seed), int(rank)])
+        eps = dataset.episodes
+        n_roots = len(dataset.root_weights)
+        self.pools = {s: [[i for i, e in enumerate(eps) if e.root == r and (len(e.views) >= 2 or not s)]
+                          for r in range(n_roots)] for s in (True, False)}
+        self.has_stereo = any(self.pools[True])
+
+    def _root_probs(self, stereo: bool) -> np.ndarray:
+        w = np.array([wt if self.pools[stereo][r] else 0.0
+                      for r, wt in enumerate(self.ds.root_weights)], np.float64)
+        return w / w.sum()
+
+    def __iter__(self):
+        while True:
+            stereo = self.has_stereo and float(self.rng.random()) < self.stereo_p
+            p = self._root_probs(stereo)
+            batch = []
+            for r in self.rng.choice(len(p), size=self.bs, p=p):
+                pool = self.pools[stereo][int(r)]
+                e = pool[int(self.rng.integers(0, len(pool)))]
+                ep = self.ds.episodes[e]
+                t = int(self.rng.integers(0, ep.num_frames))
+                views = (0, 1) if stereo else (int(self.rng.integers(0, len(ep.views))),)
+                batch.append((e, t, views, int(self.rng.integers(0, 2 ** 31))))
+            yield batch

robotryon/image/episodes.py

@@ -0,0 +1,173 @@
+"""Episode directories in the shared RoboTryOn data layout, and random frame access.
+
+::
+
+    <root>/<episode_id>/
+      meta.json              {"fps", "views", "num_frames", "caption", "embodiment"}
+      video_<view>.mp4       RGB
+      keypoints_<view>.npz   kp (T, 2, 17, 2) px at video resolution, vis (T, 2, 17)
+      mask_<view>.mp4        tight arm + hand mask (gray, > 127 = masked)
+      object_mask_<view>.mp4 optional manipulated-object mask
+"""
+from __future__ import annotations
+
+import json
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+@dataclass
+class Episode:
+    path: Path
+    views: list[str]
+    num_frames: int
+    fps: float
+    caption: str = ""
+    embodiment: str = ""
+    root: int = 0
+    extra: dict = field(default_factory=dict)
+
+    @classmethod
+    def load(cls, path: str | Path, root: int = 0) -> "Episode":
+        path = Path(path)
+        m = json.loads((path / "meta.json").read_text())
+        return cls(path=path, views=list(m["views"]), num_frames=int(m["num_frames"]),
+                   fps=float(m["fps"]), caption=str(m.get("caption", "")),
+                   embodiment=str(m.get("embodiment", "")), root=root, extra=m)
+
+    def video(self, view: str) -> Path:
+        return self.path / f"video_{view}.mp4"
+
+    def mask(self, view: str) -> Path:
+        return self.path / f"mask_{view}.mp4"
+
+    def object_mask(self, view: str) -> Path | None:
+        p = self.path / f"object_mask_{view}.mp4"
+        return p if p.exists() else None
+
+    def keypoints(self, view: str) -> Path:
+        return self.path / f"keypoints_{view}.npz"
+
+
+def scan_root(root: str | Path, root_index: int = 0) -> list[Episode]:
+    return [Episode.load(p.parent, root_index) for p in sorted(Path(root).glob("*/meta.json"))]
+
+
+class FrameReader:
+    """Small LRU of open ``cv2.VideoCapture`` objects and keypoint files (one per worker)."""
+
+    def __init__(self, max_open: int = 16):
+        self._caps: OrderedDict[str, cv2.VideoCapture] = OrderedDict()
+        self._kps: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict()
+        self.max_open = max_open
+
+    def _cap(self, path: Path) -> cv2.VideoCapture:
+        key = str(path)
+        if key in self._caps:
+            self._caps.move_to_end(key)
+            return self._caps[key]
+        cap = cv2.VideoCapture(key)
+        if not cap.isOpened():
+            raise FileNotFoundError(f"cannot open video {key}")
+        self._caps[key] = cap
+        while len(self._caps) > self.max_open:
+            self._caps.popitem(last=False)[1].release()
+        return cap
+
+    def rgb(self, path: Path, t: int) -> np.ndarray:
+        cap = self._cap(path)
+        cap.set(cv2.CAP_PROP_POS_FRAMES, int(t))
+        ok, bgr = cap.read()
+        if not ok:
+            raise RuntimeError(f"cannot decode frame {t} of {path}")
+        return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+
+    def binary(self, path: Path, t: int) -> np.ndarray:
+        """Gray mask frame -> bool ``(H, W)`` (> 127)."""
+        return self.rgb(path, t).max(axis=-1) > 127
+
+    def keypoints(self, path: Path) -> tuple[np.ndarray, np.ndarray]:
+        key = str(path)
+        if key not in self._kps:
+            z = np.load(key)
+            self._kps[key] = (z["kp"].astype(np.float32), z["vis"].astype(bool))
+            while len(self._kps) > 4 * self.max_open:
+                self._kps.popitem(last=False)
+        return self._kps[key]
+
+
+def output_geometry(src_w: int, src_h: int, long_edge: int, aspect: tuple[int, int] | None,
+                    anchor: str = "bottom") -> tuple[tuple[int, int, int, int], tuple[int, int]]:
+    """Crop box ``(x0, y0, w, h)`` in source pixels and output size ``(W, H)``.
+
+    A source with a different aspect ratio than ``aspect`` is cropped to it first ("bottom" keeps
+    the bottom edge, where the arms enter; "center" crops symmetrically). The long edge is then
+    resized to ``long_edge`` and both sides snapped to multiples of 16 (VAE /8 and 2x2 packing).
+    """
+    x0, y0, w, h = 0, 0, src_w, src_h
+    if aspect is not None and abs(src_w / src_h - aspect[0] / aspect[1]) > 1e-3:
+        if src_w / src_h < aspect[0] / aspect[1]:          # too tall: crop height
+            h = int(round(src_w * aspect[1] / aspect[0]))
+            y0 = src_h - h if anchor == "bottom" else (src_h - h) // 2
+        else:                                                # too wide: crop width, centred
+            w = int(round(src_h * aspect[0] / aspect[1]))
+            x0 = (src_w - w) // 2
+    s = long_edge / max(w, h)
+    snap = lambda v: max(16, int(round(v / 16) * 16))      # noqa: E731
+    return (x0, y0, w, h), (snap(w * s), snap(h * s))
+
+
+def crop_resize(img: np.ndarray, box, size, src_size, nearest: bool = False) -> np.ndarray:
+    """Crop ``box`` (source-frame pixels) out of ``img`` and resize to ``size``.
+
+    ``img`` may have a different resolution than the source frame ``src_size`` (e.g. a mask video
+    stored smaller); the box is rescaled accordingly. RGB uses PIL Lanczos (as the original
+    pipeline did), masks use nearest neighbour.
+    """
+    fx, fy = img.shape[1] / src_size[0], img.shape[0] / src_size[1]
+    x0, y0 = int(round(box[0] * fx)), int(round(box[1] * fy))
+    x1, y1 = int(round((box[0] + box[2]) * fx)), int(round((box[1] + box[3]) * fy))
+    sub = img[y0:y1, x0:x1]
+    if nearest:
+        return cv2.resize(sub.astype(np.uint8), size, interpolation=cv2.INTER_NEAREST)
+    pil = Image.fromarray(sub)
+    return np.asarray(pil if pil.size == tuple(size) else pil.resize(size, Image.LANCZOS))
+
+
+@dataclass
+class ViewFrame:
+    """One camera frame mapped to the model resolution.
+
+    image: uint8 ``(H, W, 3)``; hulls: per-arm convex hulls in output pixels; kp/vis: hand
+    keypoints in output pixels; object_mask: uint8 ``(H, W)`` or None.
+    """
+    image: np.ndarray
+    hulls: list
+    kp: np.ndarray
+    vis: np.ndarray
+    object_mask: np.ndarray | None
+
+
+def load_view_frame(reader: FrameReader, ep: Episode, view: str, t: int, long_edge: int,
+                    aspect=None, anchor: str = "bottom", with_object: bool = True) -> ViewFrame:
+    from ..masks import arm_hulls, scale_hulls
+    rgb = reader.rgb(ep.video(view), t)
+    src = (rgb.shape[1], rgb.shape[0])
+    box, size = output_geometry(src[0], src[1], long_edge, aspect, anchor)
+    sx, sy = size[0] / box[2], size[1] / box[3]
+    image = crop_resize(rgb, box, size, src)
+    m = reader.binary(ep.mask(view), t)
+    mx, my = src[0] / m.shape[1], src[1] / m.shape[0]          # mask px -> source px
+    hulls = scale_hulls(arm_hulls(m), mx * sx, my * sy, (box[0] / mx, box[1] / my))
+    kp, vis = reader.keypoints(ep.keypoints(view))
+    kp_t = (kp[t] - np.array(box[:2], np.float32)) * np.array([sx, sy], np.float32)
+    om = None
+    if with_object and ep.object_mask(view) is not None:
+        om = crop_resize(reader.binary(ep.object_mask(view), t).astype(np.uint8), box, size, src,
+                         nearest=True)
+    return ViewFrame(image, hulls, kp_t, vis[t].copy(), om)

robotryon/image/fill_channels.py

@@ -0,0 +1,92 @@
+"""Channel-concatenated inpainting conditioning for the FLUX.2 [klein] transformer.
+
+The conditioned frame and the mask are pixel-aligned with the target, so they ride the CHANNEL
+axis of the target tokens (FLUX-Fill style)::
+
+    target token = [ noisy latent (128) | conditioned-frame latent (128) | folded mask (256) ]
+
+The patch embedding ``x_embedder`` is ``Linear(128, inner_dim)`` in the base model and is applied
+once to the whole token sequence, so it is widened to 512 input features. The new columns are
+zero-initialised: before training the widened model reproduces the base model exactly, and a
+conditioning stream it cannot use degrades to "ignored" rather than to noise. The widened
+embedding is always trained in full (FFT and LoRA).
+
+Reference tokens are spatially unaligned with the target and ride the SEQUENCE axis instead; they
+carry only the 128 real latent channels and are zero-padded to 512, which contributes exactly
+what the base model would because the new columns start at zero.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+BASE_IN = 128          # 32 VAE latent channels x 2x2 patchify
+MASK_IN = 256          # 8x8 pixel fold -> 64 channels at H/8, then 2x2 patchify
+FILL_IN = BASE_IN + BASE_IN + MASK_IN      # 512
+
+
+def widen_x_embedder(transformer, new_in: int = FILL_IN) -> nn.Linear:
+    """Grow ``transformer.x_embedder`` in place (old columns kept, new ones zero). Idempotent.
+
+    Also records the new width in the model config (``in_channels=512, out_channels=128``) so a
+    full checkpoint written with ``save_pretrained`` reloads with ``from_pretrained`` directly.
+    """
+    old: nn.Linear = transformer.x_embedder
+    if old.in_features == new_in:
+        return old
+    if old.in_features != BASE_IN:
+        raise ValueError(f"x_embedder has {old.in_features} inputs, expected {BASE_IN} or {new_in}")
+    new = nn.Linear(new_in, old.out_features, bias=old.bias is not None,
+                    device=old.weight.device, dtype=old.weight.dtype)
+    with torch.no_grad():
+        new.weight.zero_()
+        new.weight[:, :BASE_IN].copy_(old.weight)
+        if old.bias is not None:
+            new.bias.copy_(old.bias)
+    transformer.x_embedder = new
+    out_ch = transformer.config.out_channels or transformer.config.in_channels
+    transformer.register_to_config(in_channels=new_in, out_channels=out_ch)
+    return new
+
+
+def fold_mask_to_latent(mask: torch.Tensor, vae_scale: int = 8) -> torch.Tensor:
+    """``(B, 1, H, W)`` mask in [0, 1] -> ``(B, 256, H/16, W/16)``.
+
+    The mask is not interpolated: each 8x8 pixel block is folded into 64 channels (as in FLUX
+    Fill) so sub-latent edge precision survives, then the same 2x2 patchify as the latents.
+    """
+    b, c, h, w = mask.shape
+    if c != 1:
+        raise ValueError(f"expected a single-channel mask, got {c}")
+    if h % (vae_scale * 2) or w % (vae_scale * 2):
+        raise ValueError(f"mask {h}x{w} must be a multiple of {vae_scale * 2}")
+    lh, lw = h // vae_scale, w // vae_scale
+    m = mask.view(b, lh, vae_scale, lw, vae_scale)
+    m = m.permute(0, 2, 4, 1, 3).reshape(b, vae_scale * vae_scale, lh, lw)
+    m = m.view(b, vae_scale * vae_scale, lh // 2, 2, lw // 2, 2)
+    return m.permute(0, 1, 3, 5, 2, 4).reshape(b, vae_scale * vae_scale * 4, lh // 2, lw // 2)
+
+
+def pad_condition_tokens(packed: torch.Tensor, new_in: int = FILL_IN) -> torch.Tensor:
+    """Zero-pad ``(B, N, 128)`` reference tokens to the widened embedding width."""
+    if packed.shape[-1] == new_in:
+        return packed
+    return F.pad(packed, (0, new_in - packed.shape[-1]))
+
+
+def pack(latents: torch.Tensor) -> torch.Tensor:
+    """``(B, C, h, w)`` -> ``(B, h*w, C)`` tokens (same as ``Flux2KleinPipeline._pack_latents``)."""
+    b, c, h, w = latents.shape
+    return latents.reshape(b, c, h * w).permute(0, 2, 1)
+
+
+def unpack(tokens: torch.Tensor, h: int, w: int) -> torch.Tensor:
+    """Inverse of :func:`pack`."""
+    b, n, c = tokens.shape
+    return tokens.permute(0, 2, 1).reshape(b, c, h, w)
+
+
+def fill_tokens(noisy: torch.Tensor, cond: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Target tokens ``[noisy | cond | mask]`` from latents ``(B,128,h,w)`` and pixel mask ``(B,1,H,W)``."""
+    return torch.cat([pack(noisy), pack(cond), pack(fold_mask_to_latent(mask.to(cond.dtype)))], dim=-1)

robotryon/image/model.py

@@ -0,0 +1,138 @@
+"""Build, save and load the keyframe-translation model f_theta (FLUX.2 [klein] 9B Base).
+
+Two training modes share everything but the trainable set:
+
+* ``fft``: the whole transformer is trained (FSDP, paper Table 4). Checkpoint:
+  ``<ckpt>/transformer/`` written by ``save_pretrained`` (sharded safetensors, config records the
+  widened ``in_channels=512``).
+* ``lora``: LoRA adapters on the attention / fused QKV-MLP projections of every block plus the
+  widened ``x_embedder`` trained in full. Checkpoint: ``<ckpt>/pytorch_lora_weights.safetensors``
+  (diffusers LoRA format) and ``<ckpt>/x_embedder.safetensors``.
+
+``load_pipeline`` recognises either checkpoint type.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+from .fill_channels import FILL_IN, widen_x_embedder
+
+LORA_FILE = "pytorch_lora_weights.safetensors"
+X_EMBEDDER_FILE = "x_embedder.safetensors"
+NUM_SINGLE_BLOCKS = 24
+
+
+def lora_targets() -> list[str]:
+    """Attention projections of the double-stream blocks, the fused QKV+MLP input projection of
+    the single-stream blocks, and the single-stream attention output projections."""
+    return (["to_k", "to_q", "to_v", "to_out.0", "to_qkv_mlp_proj"]
+            + [f"single_transformer_blocks.{i}.attn.to_out" for i in range(NUM_SINGLE_BLOCKS)])
+
+
+def load_transformer(base: str, dtype: torch.dtype):
+    from diffusers import Flux2Transformer2DModel
+    return Flux2Transformer2DModel.from_pretrained(base, subfolder="transformer", torch_dtype=dtype)
+
+
+def build_trainable(base: str, mode: str, lora_cfg=None):
+    """Transformer ready for training (x_embedder widened, trainable set marked).
+
+    fft: loaded in fp32 (FSDP keeps fp32 master shards and computes in bf16), all parameters
+    trainable. lora: base weights frozen in bf16, adapters and x_embedder trainable in fp32.
+    """
+    if mode == "fft":
+        tr = load_transformer(base, torch.float32)
+        widen_x_embedder(tr, FILL_IN)
+        tr.requires_grad_(True)
+    elif mode == "lora":
+        from peft import LoraConfig
+        tr = load_transformer(base, torch.bfloat16)
+        tr.requires_grad_(False)
+        tr.add_adapter(LoraConfig(r=int(lora_cfg.rank), lora_alpha=int(lora_cfg.alpha),
+                                  lora_dropout=float(lora_cfg.get("dropout", 0.0)),
+                                  init_lora_weights="gaussian", target_modules=lora_targets()))
+        widen_x_embedder(tr, FILL_IN).requires_grad_(True)
+        for p in tr.parameters():
+            if p.requires_grad:
+                p.data = p.data.float()
+    else:
+        raise ValueError(f"train.mode must be 'fft' or 'lora', got {mode!r}")
+    return tr
+
+
+def save_lora(transformer, out_dir: str | Path) -> None:
+    """LoRA tensors in diffusers format plus the full widened x_embedder."""
+    from diffusers import Flux2KleinPipeline
+    from peft.utils import get_peft_model_state_dict
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    lora = get_peft_model_state_dict(transformer)
+    lora = {k: v.detach().to(torch.bfloat16).cpu().contiguous() for k, v in lora.items()}
+    Flux2KleinPipeline.save_lora_weights(out_dir, transformer_lora_layers=lora)
+    save_file({"x_embedder.weight": transformer.x_embedder.weight.detach().float().cpu().contiguous()},
+              str(out_dir / X_EMBEDDER_FILE))
+
+
+def save_full(transformer, out_dir: str | Path, state_dict: dict | None = None) -> None:
+    """Full transformer as a diffusers model directory ``<out_dir>/transformer`` (bf16 shards).
+
+    ``state_dict`` is the gathered full state dict when training under FSDP (only the main
+    process calls this); the unwrapped module supplies the config.
+    """
+    import json
+    from huggingface_hub import split_torch_state_dict_into_shards
+    d = Path(out_dir) / "transformer"
+    d.mkdir(parents=True, exist_ok=True)
+    sd = state_dict if state_dict is not None else transformer.state_dict()
+    sd = {k: v.detach().to(torch.bfloat16).cpu().contiguous() for k, v in sd.items()}
+    split = split_torch_state_dict_into_shards(
+        sd, filename_pattern="diffusion_pytorch_model{suffix}.safetensors", max_shard_size="5GB")
+    for fname, keys in split.filename_to_tensors.items():
+        save_file({k: sd[k] for k in keys}, str(d / fname), metadata={"format": "pt"})
+    if split.is_sharded:
+        index = {"metadata": split.metadata, "weight_map": split.tensor_to_filename}
+        (d / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps(index, indent=2))
+    transformer.save_config(d)
+
+
+def checkpoint_kind(ckpt: str | Path) -> str:
+    ckpt = Path(ckpt)
+    if (ckpt / LORA_FILE).exists():
+        return "lora"
+    if (ckpt / "transformer" / "config.json").exists():
+        return "fft"
+    raise FileNotFoundError(f"{ckpt} holds neither {LORA_FILE} nor transformer/config.json")
+
+
+def load_pipeline(base: str, ckpt: str | Path, device: str = "cuda", text_encoder: bool = True):
+    """``Flux2KleinPipeline`` in bf16 with the trained keyframe-translation weights.
+
+    ``text_encoder=False`` skips the text encoder (prompts encoded beforehand).
+
+    LoRA: widen ``x_embedder`` BEFORE loading the adapter (it was fitted to the 512-wide
+    embedding), copy the trained embedding, then load the adapter unfused.
+    """
+    from diffusers import Flux2KleinPipeline
+    ckpt = Path(ckpt)
+    kind = checkpoint_kind(ckpt)
+    kw = {} if text_encoder else {"text_encoder": None, "tokenizer": None}
+    if kind == "fft":
+        from diffusers import Flux2Transformer2DModel
+        tr = Flux2Transformer2DModel.from_pretrained(ckpt / "transformer", torch_dtype=torch.bfloat16)
+        pipe = Flux2KleinPipeline.from_pretrained(base, transformer=tr, torch_dtype=torch.bfloat16, **kw)
+        pipe.to(device)
+    else:
+        pipe = Flux2KleinPipeline.from_pretrained(base, torch_dtype=torch.bfloat16, **kw).to(device)
+        widen_x_embedder(pipe.transformer, FILL_IN)
+        w = load_file(str(ckpt / X_EMBEDDER_FILE))["x_embedder.weight"]
+        with torch.no_grad():
+            pipe.transformer.x_embedder.weight.copy_(w.to(pipe.transformer.x_embedder.weight.dtype))
+        pipe.load_lora_weights(str(ckpt), weight_name=LORA_FILE)
+    if pipe.transformer.x_embedder.in_features != FILL_IN:
+        raise ValueError(f"{ckpt}: x_embedder has {pipe.transformer.x_embedder.in_features} inputs")
+    pipe.transformer.requires_grad_(False)
+    return pipe
+

robotryon/image/sampler.py

@@ -0,0 +1,121 @@
+"""Inpainting sampler for keyframe translation (true classifier-free guidance, Euler flow matching).
+
+``Flux2KleinPipeline.__call__`` cannot express channel-concatenated conditioning, so the denoising
+loop is written out here. The token layout must match training exactly (see ``fill_channels`` and
+``stereo_ids``)::
+
+    [ target_v (T=0, L=v): noisy | conditioned frame | mask ]  for every view v
+    [ reference_v (T=10, L=v): 128 latent channels, zero-padded ]  for every view v
+
+The schedule is the pipeline's: ``sigmas = linspace(1, 1/steps, steps)`` shifted by the empirical
+``mu`` of the total target token count (both views for stereo). FLUX.2 [klein] Base has no guidance
+embedding, so guidance is true CFG against the empty prompt:
+``v = v_uncond + s * (v_cond - v_uncond)``.
+
+The sampler returns the final latent TOKENS. In autoregressive inference these tokens are fed
+directly as the next keyframe's reference tokens, without decoding and re-encoding (App. B.2.3).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from .fill_channels import BASE_IN, FILL_IN, fold_mask_to_latent, pack, pad_condition_tokens, unpack
+from .stereo_ids import fill_ids
+
+
+class LatentCodec:
+    """Pixels <-> normalised, 2x2-patchified FLUX.2 VAE latents ``(1, 128, H/16, W/16)``."""
+
+    def __init__(self, vae, device):
+        self.vae, self.device = vae, device
+        self.mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device)
+        self.std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(device)
+
+    @staticmethod
+    def to_tensor(img: np.ndarray) -> torch.Tensor:
+        """uint8 ``(H, W, 3)`` -> float ``(1, 3, H, W)`` in [-1, 1]."""
+        return torch.from_numpy(np.asarray(img).astype(np.float32) / 127.5 - 1).permute(2, 0, 1)[None]
+
+    @staticmethod
+    def patchify(lat: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = lat.shape
+        lat = lat.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4)
+        return lat.reshape(b, c * 4, h // 2, w // 2)
+
+    @staticmethod
+    def unpatchify(lat: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = lat.shape
+        lat = lat.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3)
+        return lat.reshape(b, c // 4, h * 2, w * 2)
+
+    def encode_tensor(self, px: torch.Tensor) -> torch.Tensor:
+        lat = self.vae.encode(px.to(self.device, self.vae.dtype)).latent_dist.mode()
+        return (self.patchify(lat) - self.mean) / self.std
+
+    def encode(self, img: np.ndarray) -> torch.Tensor:
+        return self.encode_tensor(self.to_tensor(img))
+
+    def decode_tokens(self, tokens: torch.Tensor, h: int, w: int) -> np.ndarray:
+        """``(1, h*w, 128)`` tokens -> uint8 ``(16h, 16w, 3)`` image."""
+        lat = unpack(tokens, h, w) * self.std + self.mean
+        img = self.vae.decode(self.unpatchify(lat).to(self.vae.dtype), return_dict=False)[0]
+        return ((img[0].float().permute(1, 2, 0).clamp(-1, 1).cpu().numpy() + 1) * 127.5).astype(np.uint8)
+
+
+@dataclass
+class ViewCondition:
+    """Conditioning of one camera view.
+
+    cond_latent: ``(1, 128, h, w)`` latent of the conditioned frame Phi(o; M, S).
+    mask: ``(1, 1, 16h, 16w)`` pixel mask in [0, 1].
+    reference_tokens: ``(1, h*w, 128)`` reference tokens (an encoded image, or the previous
+        keyframe's denoised tokens).
+    """
+    cond_latent: torch.Tensor
+    mask: torch.Tensor
+    reference_tokens: torch.Tensor
+
+
+@torch.no_grad()
+def denoise(transformer, scheduler, views: list[ViewCondition], prompt_embeds, text_ids,
+            neg_embeds=None, neg_text_ids=None, steps: int = 30, guidance_scale: float = 3.5,
+            generator: torch.Generator | None = None) -> list[torch.Tensor]:
+    """Jointly denoise one or more views. Returns one ``(1, h*w, 128)`` token tensor per view."""
+    from diffusers.pipelines.flux2.pipeline_flux2_klein import compute_empirical_mu
+    device, dtype = prompt_embeds.device, transformer.dtype
+    n = len(views)
+    _, _, h, w = views[0].cond_latent.shape
+    S = h * w
+    cond = [torch.cat([pack(v.cond_latent.float()),
+                       pack(fold_mask_to_latent(v.mask.to(device, torch.float32)))], dim=-1).to(dtype)
+            for v in views]
+    refs = [pad_condition_tokens(v.reference_tokens.to(dtype), FILL_IN) for v in views]
+    ids = fill_ids(h, w, n).to(device)
+    lats = [torch.randn((1, S, BASE_IN), generator=generator, device=device, dtype=dtype)
+            for _ in range(n)]
+
+    use_cfg = guidance_scale is not None and guidance_scale > 1.0 and neg_embeds is not None
+    sigmas = np.linspace(1.0, 1 / steps, steps)
+    scheduler.set_timesteps(sigmas=sigmas, device=device, mu=compute_empirical_mu(S * n, steps))
+    scheduler.set_begin_index(0)
+    for t in scheduler.timesteps:
+        hidden = torch.cat([torch.cat([lats[i], cond[i]], dim=-1) for i in range(n)] + refs, dim=1)
+        tt = t.expand(1).to(dtype) / 1000
+        pred = transformer(hidden_states=hidden, timestep=tt, guidance=None,
+                           encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=ids,
+                           return_dict=False)[0]
+        if use_cfg:
+            neg = transformer(hidden_states=hidden, timestep=tt, guidance=None,
+                              encoder_hidden_states=neg_embeds, txt_ids=neg_text_ids, img_ids=ids,
+                              return_dict=False)[0]
+            pred = neg + guidance_scale * (pred - neg)
+        # One Euler step per view against the same t: rewind the scheduler's step counter.
+        idx = scheduler.step_index
+        for i in range(n):
+            scheduler._step_index = idx
+            lats[i] = scheduler.step(pred[:, i * S:(i + 1) * S], t, lats[i],
+                                     return_dict=False)[0].to(dtype)
+    return lats

robotryon/image/stereo_ids.py

@@ -0,0 +1,34 @@
+"""RoPE position ids for one or two camera views in one token sequence.
+
+FLUX.2 gives every image token a 4D position ``(T, H, W, L)``. The pipeline sets ``L = 0`` for all
+image tokens (only text tokens use it), so for images ``L`` is a free axis; stereo generation
+uses it as the camera-view index. ``T`` keeps its learned meaning: targets at ``T = 0`` and the
+reference at ``T = 10``. The token order is::
+
+    [target_v0, target_v1, ..., reference_v0, reference_v1, ...]
+
+With one view this is byte-identical to the stock ``_prepare_latent_ids`` + ``_prepare_image_ids``.
+"""
+from __future__ import annotations
+
+import torch
+
+TARGET_T = 0
+REFERENCE_T = 10
+
+
+def view_image_ids(sizes: list[tuple[int, int]], view_ids: list[int],
+                   t_coords: list[int]) -> torch.Tensor:
+    """``sizes[i] = (h, w)`` token grid of image i -> ``(1, sum h*w, 4)`` int64 position ids."""
+    if not (len(sizes) == len(view_ids) == len(t_coords)):
+        raise ValueError("one size, view id and t coordinate per image")
+    out = [torch.cartesian_prod(torch.tensor([int(t)]), torch.arange(h), torch.arange(w),
+                                torch.tensor([int(v)]))
+           for (h, w), v, t in zip(sizes, view_ids, t_coords)]
+    return torch.cat(out, dim=0).unsqueeze(0)
+
+
+def fill_ids(h: int, w: int, n_views: int) -> torch.Tensor:
+    """Ids for ``n_views`` targets followed by ``n_views`` references, all on an ``h x w`` grid."""
+    return view_image_ids([(h, w)] * (2 * n_views), list(range(n_views)) * 2,
+                          [TARGET_T] * n_views + [REFERENCE_T] * n_views)

robotryon/image/text.py

@@ -0,0 +1,56 @@
+"""The on-disk text-embedding cache of the image model's prompts (``robotryon.captions``).
+
+Qwen3 prompt embeddings are ``(512, 12288)`` per prompt, too large to keep for every prompt in
+memory, so they are computed once at trainer start-up and stored as one safetensors file per
+prompt, named by a hash of the text.
+"""
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+MAX_SEQUENCE_LENGTH = 512
+
+
+class TextEmbedCache:
+    """One ``<sha1>.safetensors`` file per prompt holding ``prompt_embeds`` ``(L, D)``."""
+
+    def __init__(self, root: str | Path):
+        self.root = Path(root)
+
+    def path(self, prompt: str) -> Path:
+        return self.root / f"{hashlib.sha1(prompt.encode('utf-8')).hexdigest()}.safetensors"
+
+    def missing(self, prompts) -> list[str]:
+        return sorted({p for p in prompts if not self.path(p).exists()})
+
+    def load(self, prompt: str) -> torch.Tensor:
+        p = self.path(prompt)
+        if not p.exists():
+            raise FileNotFoundError(f"no cached embedding for prompt {prompt!r}")
+        return load_file(str(p))["prompt_embeds"]
+
+    @torch.no_grad()
+    def fill(self, text_pipe, prompts, device, batch_size: int = 16) -> int:
+        """Encode and store the given prompts with a ``Flux2KleinPipeline`` holding the text encoder."""
+        self.root.mkdir(parents=True, exist_ok=True)
+        prompts = list(prompts)
+        for i in range(0, len(prompts), batch_size):
+            chunk = prompts[i:i + batch_size]
+            emb, _ = text_pipe.encode_prompt(prompt=chunk, device=device,
+                                             max_sequence_length=MAX_SEQUENCE_LENGTH)
+            for p, e in zip(chunk, emb):
+                tmp = self.path(p).with_suffix(".tmp")
+                save_file({"prompt_embeds": e.to(torch.bfloat16).contiguous().cpu()}, str(tmp))
+                tmp.replace(self.path(p))
+        return len(prompts)
+
+
+def text_ids_for(prompt_embeds: torch.Tensor) -> torch.Tensor:
+    """``(B, L, D)`` -> ``(B, L, 4)`` text position ids (same as ``Flux2KleinPipeline._prepare_text_ids``)."""
+    b, n, _ = prompt_embeds.shape
+    ids = torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(n))
+    return ids.unsqueeze(0).expand(b, -1, -1).to(prompt_embeds.device)

robotryon/image/training.py

@@ -0,0 +1,48 @@
+"""Flow-matching objective of the keyframe-translation model (paper Eq. 3.1).
+
+With ``z = E_img(o_t)`` and ``z_tau = (1 - tau) z + tau eps``, the transformer predicts
+``eps - z`` from the target tokens ``[z_tau | E_img(Phi(o_t; M_t, S_t)) | fold(M_t)]`` and the
+reference tokens ``E_img(o_t')``. ``tau`` is drawn uniformly from the base model's training
+sigma grid, one value per sample shared by its views. For stereo samples the loss is the mean
+over the two views.
+"""
+from __future__ import annotations
+
+import torch
+
+from .fill_channels import FILL_IN, fill_tokens, pack, pad_condition_tokens, unpack
+from .sampler import LatentCodec
+from .stereo_ids import fill_ids
+from .text import text_ids_for
+
+
+def flow_matching_loss(transformer, codec: LatentCodec, sched_sigmas: torch.Tensor,
+                       sched_timesteps: torch.Tensor, batch: dict, device) -> torch.Tensor:
+    B, V = batch["target"].shape[:2]
+    with torch.no_grad():
+        enc = [{k: codec.encode_tensor(batch[k][:, v].to(device))
+                for k in ("target", "cond", "reference")} for v in range(V)]
+    z0 = enc[0]["target"]
+    _, _, h, w = z0.shape
+    n = sched_timesteps.numel()
+    idx = (torch.rand(B) * n).long().clamp(max=n - 1)
+    sigma = sched_sigmas[idx].to(device=device, dtype=z0.dtype).view(B, 1, 1, 1)
+    timestep = sched_timesteps[idx].to(device)
+
+    targets, tokens = [], []
+    for v, e in enumerate(enc):
+        noise = torch.randn_like(e["target"])
+        noisy = (1.0 - sigma) * e["target"] + sigma * noise
+        tokens.append(fill_tokens(noisy, e["cond"], batch["mask"][:, v].to(device)))
+        targets.append(noise - e["target"])
+    tokens += [pad_condition_tokens(pack(e["reference"]), FILL_IN) for e in enc]
+    hidden = torch.cat(tokens, dim=1).to(torch.bfloat16)
+    ids = fill_ids(h, w, V).to(device).expand(B, -1, -1)
+    prompt_embeds = batch["prompt_embeds"].to(device)
+    pred = transformer(hidden_states=hidden, timestep=timestep / 1000, guidance=None,
+                       encoder_hidden_states=prompt_embeds, txt_ids=text_ids_for(prompt_embeds),
+                       img_ids=ids, return_dict=False)[0]
+    S = h * w
+    losses = [((unpack(pred[:, v * S:(v + 1) * S], h, w).float() - targets[v].float()) ** 2)
+              .reshape(B, -1).mean(1).mean() for v in range(V)]
+    return sum(losses) / V

robotryon/masks/__init__.py

@@ -0,0 +1,14 @@
+"""Shape-agnostic arm/hand masks for keyframe translation (paper Sec. 3.2, App. B.2.3).
+
+The mask must hide the silhouette of whatever embodiment it covers, so it is never the raw
+segmentation: every arm component becomes its convex hull, which is then dilated. During training
+the dilation is random per arm, the manipulated object is cut out of the hole with some
+probability, and the whole sample is jointly rescaled (field-of-view crop / pad). At inference the
+dilation is fixed and the eroded object mask is always subtracted.
+"""
+from .hull import arm_hulls, fill_hulls, scale_hulls
+from .augment import MaskAug, fov_augment, inference_mask
+from .reference import robot_only
+
+__all__ = ["arm_hulls", "fill_hulls", "scale_hulls", "MaskAug", "fov_augment", "inference_mask",
+           "robot_only"]

robotryon/masks/augment.py

@@ -0,0 +1,131 @@
+"""Train-time mask randomisation, the fixed inference mask, and the joint field-of-view augmentation.
+
+Training (``MaskAug.train_mask``):
+  * one random dilation per arm hull, so two arms in one frame already carry different
+    hole-to-arm ratios and the model cannot read the hand size off the hole;
+  * episodes flagged as "bonus" (thin human arms) get a raised dilation floor, because a hull
+    hugging a thin arm leaves no room for a thicker robot arm;
+  * with probability ``p_object_cut``, where an object mask exists, the eroded object is cut out
+    of the hole (a concave hole), which is what inference always does.
+
+Inference (``inference_mask``): every hull dilated by a fixed kernel (the middle of the training
+range plus the bonus), then the eroded object mask is subtracted so the object is preserved.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from .hull import fill_hulls
+
+
+def _ellipse(k: int) -> np.ndarray:
+    k = int(k) | 1
+    return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+
+
+@dataclass
+class MaskAug:
+    dilate_min: int = 3
+    dilate_max: int = 71
+    bonus: int = 21                 # added to the dilation floor for "bonus" (human) episodes
+    p_object_cut: float = 0.5
+    obj_erode_min: int = 3
+    obj_erode_max: int = 21
+    min_keep_frac: float = 0.4      # reject a cut that leaves less than this fraction of the hole
+
+    @classmethod
+    def from_config(cls, cfg) -> "MaskAug":
+        cfg = cfg or {}
+        return cls(**{k: cfg[k] for k in cls.__dataclass_fields__ if k in cfg})
+
+    def sample_dilate(self, rng: np.random.Generator, bonus: bool = False) -> int:
+        lo = self.dilate_min + (self.bonus if bonus else 0)
+        return int(rng.integers(lo, max(lo + 1, self.dilate_max) + 1)) | 1
+
+    def sample_obj_erode(self, rng: np.random.Generator) -> int:
+        return int(rng.integers(self.obj_erode_min, self.obj_erode_max + 1)) | 1
+
+    def infer_dilate(self) -> int:
+        """Fixed inference dilation: middle of the training range plus the bonus."""
+        return (self.dilate_min + self.dilate_max) // 2 + self.bonus
+
+    def train_mask(self, polys: list[np.ndarray], width: int, height: int,
+                   rng: np.random.Generator, bonus: bool = False,
+                   object_mask: np.ndarray | None = None) -> tuple[np.ndarray, int]:
+        """Randomised training hole. Returns ``(mask uint8 {0,1}, object erode kernel or 0)``."""
+        mask = np.zeros((height, width), np.uint8)
+        for one in fill_hulls(polys, width, height):
+            mask |= cv2.dilate(one, _ellipse(self.sample_dilate(rng, bonus)))
+        if object_mask is None:
+            return mask, 0
+        if float(rng.random()) >= self.p_object_cut:
+            return mask, 0
+        om = (np.asarray(object_mask) > 0).astype(np.uint8)
+        if om.sum() == 0:
+            return mask, 0
+        e = self.sample_obj_erode(rng)
+        om = cv2.erode(om, _ellipse(e))
+        cut = (mask.astype(bool) & ~om.astype(bool)).astype(np.uint8)
+        hull_px = int(mask.sum())
+        if cut.sum() > 0 and (hull_px == 0 or cut.sum() / hull_px > self.min_keep_frac):
+            return cut, e
+        return mask, 0
+
+
+def inference_mask(polys: list[np.ndarray], width: int, height: int, dilate: int,
+                   object_mask: np.ndarray | None = None, obj_erode: int = 9) -> np.ndarray:
+    """Fixed-dilation hull mask minus the eroded object -> ``(H, W)`` uint8 {0, 1}."""
+    mask = np.zeros((height, width), np.uint8)
+    for one in fill_hulls(polys, width, height):
+        mask |= one
+    if mask.any():
+        mask = cv2.dilate(mask, _ellipse(dilate))
+    if object_mask is not None:
+        om = (np.asarray(object_mask) > 0).astype(np.uint8)
+        if obj_erode > 0:
+            om = cv2.erode(om, _ellipse(obj_erode))
+        mask = (mask.astype(bool) & ~om.astype(bool)).astype(np.uint8)
+    return mask
+
+
+def fov_augment(imgs: list[Image.Image], rng: np.random.Generator, crop_min: float,
+                scale_max: float = 1.0, flip_prob: float = 0.0) -> tuple[list[Image.Image], float]:
+    """One shared field-of-view draw for every stream of a sample (all views included).
+
+    ``scale ~ U[crop_min, max(1, scale_max)]``. Below 1 a ``scale``-sized window is cropped and
+    resized back; above 1 the image is pasted at a random offset onto a black canvas ``scale``
+    times larger and resized back, so the scene shrinks and black borders appear (the mask is
+    black there too: "not a hole"). The model thus learns that the camera's field of view is not
+    fixed.
+    """
+    if not imgs:
+        return imgs, 1.0
+    w, h = imgs[0].size
+    if any(im.size != (w, h) for im in imgs):
+        raise ValueError("fov_augment needs equally sized images")
+    scale = float(rng.uniform(crop_min, max(1.0, float(scale_max))))
+    if scale > 1.0 + 1e-6:
+        cw, ch = int(round(w * scale)), int(round(h * scale))
+        left = int(rng.integers(0, cw - w + 1))
+        top = int(rng.integers(0, ch - h + 1))
+        flip = float(rng.random()) < flip_prob
+        out = []
+        for im in imgs:
+            canvas = Image.new(im.mode, (cw, ch), 0)
+            canvas.paste(im, (left, top))
+            im2 = canvas.resize((w, h), Image.LANCZOS)
+            out.append(im2.transpose(Image.FLIP_LEFT_RIGHT) if flip else im2)
+        return out, scale
+    cw, ch = max(8, int(round(w * scale))), max(8, int(round(h * scale)))
+    left = int(rng.integers(0, w - cw + 1)) if w > cw else 0
+    top = int(rng.integers(0, h - ch + 1)) if h > ch else 0
+    flip = float(rng.random()) < flip_prob
+    out = []
+    for im in imgs:
+        im2 = im.crop((left, top, left + cw, top + ch)).resize((w, h), Image.LANCZOS)
+        out.append(im2.transpose(Image.FLIP_LEFT_RIGHT) if flip else im2)
+    return out, scale

robotryon/masks/hull.py

@@ -0,0 +1,53 @@
+"""Convex hull per arm.
+
+A segmentation silhouette gives the embodiment away (a robot arm has cables and a different
+outline than a human arm). Each connected component of the mask is therefore replaced by its
+convex hull. Small gaps are closed before the component split, so a limb and a thin cable that
+pinches to a one-pixel thread stay one component instead of producing a separate cable-shaped
+polygon.
+"""
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+PRE_DILATE = 9            # closing kernel before the component split (px at mask resolution)
+MIN_COMPONENT_PX = 256    # smaller components are segmentation noise, not an arm
+
+
+def arm_hulls(mask: np.ndarray, pre_dilate: int = PRE_DILATE,
+              min_component_px: int = MIN_COMPONENT_PX) -> list[np.ndarray]:
+    """Binary mask ``(H, W)`` -> list of ``(K, 2)`` float32 hull polygons in mask pixels."""
+    m = (np.asarray(mask) > 0).astype(np.uint8)
+    if m.sum() == 0:
+        return []
+    if pre_dilate > 0:
+        k = pre_dilate | 1
+        m = cv2.dilate(m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)))
+    n, lab = cv2.connectedComponents(m)
+    polys = []
+    for i in range(1, n):
+        ys, xs = np.where(lab == i)
+        if ys.size < min_component_px:
+            continue
+        pts = np.stack([xs, ys], axis=1).astype(np.int32)
+        polys.append(cv2.convexHull(pts).reshape(-1, 2).astype(np.float32))
+    return polys
+
+
+def scale_hulls(polys: list[np.ndarray], sx: float, sy: float,
+                offset: tuple[float, float] = (0.0, 0.0)) -> list[np.ndarray]:
+    """Map polygons from mask pixels to output pixels: ``(p - offset) * (sx, sy)``."""
+    off = np.asarray(offset, np.float32)
+    s = np.array([sx, sy], np.float32)
+    return [(p - off) * s for p in polys]
+
+
+def fill_hulls(polys: list[np.ndarray], width: int, height: int) -> list[np.ndarray]:
+    """Rasterise each polygon separately -> list of ``(H, W)`` uint8 {0, 1} masks."""
+    out = []
+    for p in polys:
+        one = np.zeros((height, width), np.uint8)
+        cv2.fillConvexPoly(one, p.astype(np.int32), 1)
+        out.append(one)
+    return out

robotryon/masks/reference.py

@@ -0,0 +1,29 @@
+"""Robot-only reference image: the robot's pixels on a plain white background.
+
+The first keyframe of a demonstration has no translated predecessor, so its reference is a real
+frame of the target robot with everything but the robot removed; this selects the embodiment.
+Training uses the same construction for its robot-only exemplars.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+WHITE = 255
+
+
+def robot_only(image: np.ndarray, mask: np.ndarray, erode: int = 7,
+               background: int = WHITE) -> np.ndarray:
+    """Keep the robot pixels of ``image`` (H, W, 3) where ``mask`` (H, W) is set.
+
+    The mask is eroded by ``erode`` px so no background leaks in at the silhouette; everything
+    else is filled with ``background``.
+    """
+    m = np.asarray(mask).astype(bool)
+    if erode > 0:
+        k = int(erode) | 1
+        m = cv2.erode(m.astype(np.uint8), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))) > 0
+    out = np.full_like(image, background)
+    out[m] = image[m]
+    return out

robotryon/skeleton/__init__.py

@@ -0,0 +1,8 @@
+from .augment import SkeletonAugment
+from .render import condition_frame, draw_skeleton, draw_skeleton_video, paint
+from .topology import EDGES, KEYPOINT_NAMES, MEDIAPIPE21_TO_17, NUM_HANDS, NUM_KEYPOINTS
+
+__all__ = [
+    "SkeletonAugment", "condition_frame", "draw_skeleton", "draw_skeleton_video", "paint",
+    "EDGES", "KEYPOINT_NAMES", "MEDIAPIPE21_TO_17", "NUM_HANDS", "NUM_KEYPOINTS",
+]

robotryon/skeleton/augment.py

@@ -0,0 +1,75 @@
+"""Training-time skeleton augmentation.
+
+Human keypoints at inference come from a reconstructed and retargeted hand, which is imprecise.
+To make the generative models treat the skeleton as a coarse guide rather than an exact pose,
+robot keypoints are corrupted during training:
+
+* each keypoint is displaced by up to half the length of the finger segment it belongs to,
+* the whole skeleton is dropped with probability ``p_drop_all``,
+* each keypoint is dropped independently with probability ``p_drop_keypoint``.
+
+For video, the displacement is sampled once per clip and per keypoint and then held fixed in
+time, so the skeleton flow stays temporally coherent; keypoint dropout is also per clip.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from .topology import segment_length_per_keypoint
+
+
+@dataclass
+class SkeletonAugment:
+    max_shift_frac: float = 0.5
+    p_drop_all: float = 0.2
+    p_drop_keypoint: float = 0.2
+    enabled: bool = True
+
+    @classmethod
+    def from_config(cls, cfg) -> "SkeletonAugment":
+        if cfg is None:
+            return cls(enabled=False)
+        return cls(**{k: cfg[k] for k in ("max_shift_frac", "p_drop_all", "p_drop_keypoint",
+                                          "enabled") if k in cfg})
+
+    def __call__(self, kp: np.ndarray, vis: np.ndarray, rng: np.random.Generator,
+                 temporal: bool | None = None) -> tuple[np.ndarray, np.ndarray]:
+        """Corrupt keypoints.
+
+        Args:
+            kp: ``(H, 17, 2)`` for one frame or ``(T, H, 17, 2)`` for a clip, in pixels.
+            vis: matching visibility, ``(H, 17)`` or ``(T, H, 17)``.
+            rng: random generator.
+            temporal: treat the first axis as time. Inferred from ``kp.ndim == 4`` if omitted.
+
+        Returns:
+            ``(kp, vis)``: new arrays, the inputs are not modified.
+        """
+        kp = np.array(kp, dtype=np.float32, copy=True)
+        vis = np.array(vis, dtype=bool, copy=True)
+        if not self.enabled:
+            return kp, vis
+        if temporal is None:
+            temporal = kp.ndim == 4
+        shape = kp.shape[1:-1] if temporal else kp.shape[:-1]          # (H, 17)
+
+        if rng.random() < self.p_drop_all:
+            vis[...] = False
+            return kp, vis
+
+        # Segment length from the clip median, so a single bad frame cannot blow up the shift.
+        seg = segment_length_per_keypoint(kp, vis)
+        if temporal:
+            seg = np.median(np.where(vis, seg, np.nan), axis=0)
+            seg = np.nan_to_num(seg, nan=0.0)
+        radius = self.max_shift_frac * seg * np.sqrt(rng.random(shape))   # uniform in a disc
+        angle = rng.uniform(0.0, 2 * np.pi, shape)
+        shift = np.stack([radius * np.cos(angle), radius * np.sin(angle)], -1).astype(np.float32)
+        kp = kp + (shift[None] if temporal else shift)
+
+        keep = rng.random(shape) >= self.p_drop_keypoint
+        vis &= keep[None] if temporal else keep
+        return kp, vis

robotryon/skeleton/render.py

@@ -0,0 +1,75 @@
+"""Draw hand skeletons as 2D images.
+
+The same renderer is used for every embodiment and for humans, so the skeleton carries pose
+only and never reveals what kind of hand produced it.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+from .topology import COLORS, EDGES, GROUP_OF
+
+
+def draw_skeleton(kp: np.ndarray, vis: np.ndarray, width: int, height: int,
+                  thickness: int = 2, radius: int = 3, supersample: int = 2) -> np.ndarray:
+    """Render one frame of hand skeletons on a black canvas.
+
+    Args:
+        kp: ``(H, 17, 2)`` pixel coordinates.
+        vis: ``(H, 17)`` visibility. Invisible or non-finite keypoints are skipped together with
+            their edges.
+        width, height: canvas size in pixels.
+
+    Returns:
+        ``(height, width, 3)`` uint8 RGB image.
+    """
+    s = max(1, int(supersample))
+    canvas = np.zeros((height * s, width * s, 3), np.uint8)
+    # Same stroke convention as the original renderer: sizes are given on the supersampled
+    # canvas at half scale, so thickness=2, radius=3 is the production look.
+    lw, dr = max(1, thickness * s // 2), max(1, radius * s // 2)
+    kp = np.asarray(kp, np.float64)
+    vis = np.asarray(vis, bool) & np.isfinite(kp).all(-1)
+    for h in range(kp.shape[0]):
+        pts = np.round(kp[h] * s).astype(np.int64) if vis[h].any() else None
+        if pts is None:
+            continue
+        for a, b, g in EDGES:
+            if vis[h, a] and vis[h, b]:
+                cv2.line(canvas, tuple(pts[a]), tuple(pts[b]), COLORS[g], lw, cv2.LINE_AA)
+        for j in range(kp.shape[1]):
+            if vis[h, j]:
+                cv2.circle(canvas, tuple(pts[j]), dr, COLORS[GROUP_OF[j]], -1, cv2.LINE_AA)
+    if s > 1:
+        canvas = cv2.resize(canvas, (width, height), interpolation=cv2.INTER_AREA)
+    return canvas
+
+
+def draw_skeleton_video(kp: np.ndarray, vis: np.ndarray, width: int, height: int,
+                        **kw) -> np.ndarray:
+    """``(T, H, 17, 2)`` keypoints -> ``(T, height, width, 3)`` uint8 skeleton video."""
+    return np.stack([draw_skeleton(kp[t], vis[t], width, height, **kw)
+                     for t in range(kp.shape[0])])
+
+
+def paint(frame: np.ndarray, skeleton: np.ndarray) -> np.ndarray:
+    """Overlay a rendered skeleton on a frame; skeleton pixels win wherever drawn."""
+    out = frame.copy()
+    lit = skeleton.max(axis=-1) > 0
+    out[lit] = skeleton[lit]
+    return out
+
+
+def condition_frame(frame: np.ndarray, mask: np.ndarray | None,
+                    skeleton: np.ndarray) -> np.ndarray:
+    """Phi(o; M, S) = ((1 - M) * o) (+) S: black out the mask region, then draw the skeleton.
+
+    The skeleton is drawn even slightly outside the mask so fingertips are never clipped.
+    ``mask=None`` leaves the frame unmasked (mask-free inference).
+    """
+    out = frame.copy()
+    if mask is not None:
+        out[np.asarray(mask).astype(bool)] = 0
+    return paint(out, skeleton)

robotryon/skeleton/topology.py

@@ -0,0 +1,90 @@
+"""Hand keypoint topology shared by every embodiment.
+
+Each hand has 17 keypoints in MediaPipe order with the four distal (DIP) joints removed:
+
+    0           wrist
+    1, 2, 3, 4  thumb: base (CMC), MCP, IP, tip
+    5, 6, 7     index: knuckle (MCP), middle (PIP), tip
+    8, 9, 10    middle
+    11, 12, 13  ring
+    14, 15, 16  little
+
+Robot hands map their links onto the same 17 points, so hands with different degrees of
+freedom share one skeleton. Arrays of keypoints use shape ``(..., 2, 17, D)``: hand 0 is the
+left hand, hand 1 is the right hand.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+NUM_KEYPOINTS = 17
+NUM_HANDS = 2
+
+KEYPOINT_NAMES = (
+    "wrist",
+    "thumb_cmc", "thumb_mcp", "thumb_ip", "thumb_tip",
+    "index_mcp", "index_pip", "index_tip",
+    "middle_mcp", "middle_pip", "middle_tip",
+    "ring_mcp", "ring_pip", "ring_tip",
+    "little_mcp", "little_pip", "little_tip",
+)
+
+# MediaPipe-21 index of each of the 17 keypoints (the DIP joints 7/11/15/19 are dropped).
+MEDIAPIPE21_TO_17 = (0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20)
+
+GROUP_OF = (
+    "palm",
+    "thumb", "thumb", "thumb", "thumb",
+    "index", "index", "index",
+    "middle", "middle", "middle",
+    "ring", "ring", "ring",
+    "little", "little", "little",
+)
+
+# (parent, child, group). The palm connects the wrist to the thumb base and the little knuckle
+# and links the knuckles to one another; each finger is a chain from its knuckle to its tip.
+EDGES: tuple[tuple[int, int, str], ...] = (
+    (0, 1, "palm"), (0, 5, "palm"), (5, 8, "palm"), (8, 11, "palm"), (11, 14, "palm"),
+    (0, 14, "palm"),
+    (1, 2, "thumb"), (2, 3, "thumb"), (3, 4, "thumb"),
+    (5, 6, "index"), (6, 7, "index"),
+    (8, 9, "middle"), (9, 10, "middle"),
+    (11, 12, "ring"), (12, 13, "ring"),
+    (14, 15, "little"), (15, 16, "little"),
+)
+
+COLORS = {
+    "palm": (255, 170, 60), "thumb": (255, 232, 70), "index": (110, 235, 120),
+    "middle": (95, 205, 255), "ring": (185, 150, 255), "little": (255, 125, 205),
+}
+
+# Finger segments a keypoint belongs to, used to scale the training-time perturbation.
+# A keypoint is displaced by at most half the length of the shortest finger segment it touches.
+_FINGER_EDGES = tuple((a, b) for a, b, g in EDGES if g != "palm")
+
+
+def segment_length_per_keypoint(kp: np.ndarray, vis: np.ndarray | None = None) -> np.ndarray:
+    """Length of the finger segment each keypoint belongs to.
+
+    Args:
+        kp: ``(..., 17, D)`` keypoints of one or more hands.
+        vis: optional ``(..., 17)`` visibility; invisible endpoints do not contribute.
+
+    Returns:
+        ``(..., 17)`` lengths. The wrist uses its distance to the nearest knuckle. A keypoint
+        without any valid segment gets 0, so it is never displaced.
+    """
+    kp = np.asarray(kp, dtype=np.float64)
+    if vis is None:
+        vis = np.isfinite(kp).all(-1)
+    out = np.full(kp.shape[:-1], np.inf)
+    edges = _FINGER_EDGES + ((0, 5), (0, 1))
+    for a, b in edges:
+        ok = vis[..., a] & vis[..., b]
+        d = np.linalg.norm(np.nan_to_num(kp[..., a, :] - kp[..., b, :]), axis=-1)
+        d = np.where(ok, d, np.inf)
+        out[..., a] = np.minimum(out[..., a], d)
+        out[..., b] = np.minimum(out[..., b], d)
+    out[~np.isfinite(out)] = 0.0
+    return out

robotryon/video/__init__.py

@@ -0,0 +1,15 @@
+"""Video model g_phi: an in-context LoRA on LTX-2.5 that synthesizes the robot video between
+translated keyframes under the guidance of the hand-skeleton flow (paper Sec. 3.3).
+
+Modules:
+    conditioning  training sequence layout, keyframe sampling, flow-matching loss
+    dataset       random training clips with on-the-fly skeleton rendering and augmentation
+    rope4d        fourth RoPE axis distinguishing the two stereo views
+    pipeline      diffusers pipeline subclass and the two-stage generation of one clip
+    chunks        splitting long videos at keyframes and concatenating the clips
+    text          caption embeddings (cached once; the text encoder is never co-resident)
+    training      gradient checkpointing and LoRA weight I/O
+    io            mp4 / png helpers
+
+Heavy dependencies (torch, diffusers) are imported by the submodules that need them.
+"""

robotryon/video/chunks.py

@@ -0,0 +1,164 @@
+"""Splitting a long demonstration into clips the video model can generate at once.
+
+The video VAE takes clips of ``8k + 1`` frames, and the model is trained on clips of at most
+145 frames (about 7 s at 20 fps). A longer demonstration is therefore cut at its keyframes:
+each clip runs from one keyframe to a later one, holds every keyframe in between, and the next
+clip starts at the keyframe the previous one ended on. Adjacent clips thus **share their
+boundary keyframe**, so the arm pose cannot jump at the seam, and the concatenation drops the
+duplicated boundary frame of every clip after the first.
+
+A clip is ``8k + 1`` frames long, which usually reaches past its last keyframe; those trailing
+frames are anchored to nothing and are discarded at concatenation. Two corner cases follow the
+original generation code:
+
+* two consecutive keyframes further apart than the cap: the clip is stretched past the cap
+  rather than leaving a hole in the video (logged as off-distribution);
+* the end of the video leaves less than the minimum clip length: a short tail clip of at least
+  9 frames is generated instead of dropping the end of the demonstration, pulling the last
+  keyframe onto the clip's final frame if the ``8k + 1`` constraint requires it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+
+import numpy as np
+
+VAE_TEMPORAL = 8
+
+
+def allowed_frame_counts(lo: int, hi: int) -> list[int]:
+    """Clip lengths in ``[lo, hi]`` the VAE accepts (``F = 8k + 1``)."""
+    return [f for f in range(lo, hi + 1) if (f - 1) % VAE_TEMPORAL == 0]
+
+
+def next_8k1(n: int) -> int:
+    """Smallest ``8k + 1`` that is >= n."""
+    return n if (n - 1) % VAE_TEMPORAL == 0 else n + (VAE_TEMPORAL - (n - 1) % VAE_TEMPORAL)
+
+
+def prev_8k1(n: int) -> int:
+    """Largest ``8k + 1`` that is <= n (0 if none)."""
+    return n - (n - 1) % VAE_TEMPORAL if n >= 1 else 0
+
+
+def latent_index(pixel: int, num_frames: int) -> int:
+    """Latent frame holding a pixel frame: latent 0 is pixel 0, latent j>0 is [8j-7, 8j]."""
+    lat_n = (num_frames - 1) // VAE_TEMPORAL + 1
+    return 0 if pixel == 0 else min((pixel + 7) // VAE_TEMPORAL, lat_n - 1)
+
+
+@dataclass
+class Chunk:
+    """One generated clip: ``frames`` frames starting at ``start`` of the episode.
+
+    ``kf_slots`` index into the episode's keyframe list, ``kf_local`` are the same keyframes
+    as frame indices inside the clip (first is always 0).
+    """
+    start: int
+    frames: int
+    kf_slots: list[int]
+    kf_local: list[int]
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def _tail_chunk(kf: list[int], i: int, start: int, avail: int, log) -> Chunk | None:
+    """A short last clip covering what is left after keyframe ``i``."""
+    F = prev_8k1(avail)
+    if F < 9:
+        return None
+    ids = [t for t in range(i, len(kf)) if kf[t] <= start + F - 1] or [i]
+    if ids[-1] + 1 < len(kf):
+        ids.append(ids[-1] + 1)                 # include the last keyframe, pulled to F-1
+    local = [min(kf[t] - start, F - 1) for t in ids]
+    if len(set(local)) < 2:
+        return None
+    log(f"  tail clip: frames {start}..{start + F - 1} ({F}f, below the training minimum) "
+        f"keyframes {ids} @ {local}")
+    return Chunk(start, F, ids, local)
+
+
+def plan_chunks(kf: list[int], total: int, frames_range=(49, 145), log=print) -> list[Chunk]:
+    """Cut a ``total``-frame video with keyframes ``kf`` (sorted frame indices) into clips.
+
+    Each clip starts at a keyframe and ends at a later one, holding as many keyframes as fit
+    in ``frames_range[1]`` frames; its length is the smallest allowed ``8k + 1`` that covers
+    them (at least ``frames_range[0]``).
+    """
+    choices = allowed_frame_counts(*frames_range)
+    lo, hi = choices[0], choices[-1]
+    out: list[Chunk] = []
+    i = 0
+    while i < len(kf) - 1:
+        j = i + 1
+        while j + 1 < len(kf) and (kf[j + 1] - kf[i] + 1) <= hi:
+            j += 1
+        span = kf[j] - kf[i] + 1
+        start = kf[i]
+        if span > hi:
+            # Two consecutive keyframes are further apart than the cap: stretch the clip.
+            F = next_8k1(span)
+            if start + F > total:
+                F = prev_8k1(total - start)
+                if F < lo:
+                    tc = _tail_chunk(kf, i, start, total - start, log)
+                    if tc is not None:
+                        out.append(tc)
+                    break
+            log(f"  keyframes {i}->{j} are {span} frames apart, above the {hi}-frame cap: "
+                f"generating {F} frames")
+            out.append(Chunk(start, F, [i, j], [0, min(kf[j] - start, F - 1)]))
+            i = j
+            continue
+        avail = total - start
+        F = next((f for f in choices if f >= max(span, lo) and f <= avail), None)
+        if F is None:
+            # The video ends before a long-enough clip: keep the start (frame 0 of a clip is
+            # the substituted keyframe) and pull the last keyframe onto the clip's last frame.
+            F = max((f for f in choices if f <= avail), default=None)
+            if F is None or F < lo:
+                tc = _tail_chunk(kf, i, start, avail, log)
+                if tc is not None:
+                    out.append(tc)
+                break
+            log(f"  last clip: {span} -> {F} frames, last keyframe pulled to {F - 1}")
+        local = [min(k - start, F - 1) for k in kf[i:j + 1]]
+        out.append(Chunk(start, F, list(range(i, j + 1)), local))
+        i = j
+    return out
+
+
+def drop_colliding_keyframes(kf_local: list[int], num_frames: int) -> list[int]:
+    """Positions (into ``kf_local``) to keep so no two keyframes share a latent frame.
+
+    Short clips can put two keyframes into the same latent slab (17 frames, keyframes at
+    [0, 11, 16] -> latents [0, 2, 2]), which the pipeline cannot represent. The first and last
+    keyframes are kept; colliding interior ones are dropped.
+    """
+    if len(kf_local) <= 2:
+        return list(range(len(kf_local)))
+    last = latent_index(kf_local[-1], num_frames)
+    keep, seen = [], set()
+    for n, loc in enumerate(kf_local):
+        li = latent_index(loc, num_frames)
+        if n < len(kf_local) - 1 and (li in seen or li == last):
+            continue
+        seen.add(li)
+        keep.append(n)
+    return keep
+
+
+def concat_chunks(chunks: list[Chunk], clips: list[np.ndarray]) -> np.ndarray:
+    """Join generated clips ``(F, H, W, 3)`` into one video.
+
+    Clip 0 contributes frames ``0..kf_local[-1]``; every later clip starts one frame in, since
+    its frame 0 is the boundary keyframe the previous clip already ended on. Frames after a
+    clip's last keyframe are dropped.
+    """
+    parts = []
+    for n, (ch, clip) in enumerate(zip(chunks, clips)):
+        lo = ch.kf_local[0] + (1 if n else 0)
+        parts.append(clip[lo:ch.kf_local[-1] + 1])
+    return np.concatenate(parts, axis=0)

robotryon/video/conditioning.py

@@ -0,0 +1,150 @@
+"""Assembling one training step of the in-context video model g_phi (paper Eq. 3.3).
+
+A training sequence holds three groups of tokens, in the order of diffusers'
+``LTX2InContextPipeline``::
+
+    [ target tokens | keyframe tokens | skeleton-flow reference tokens ]
+
+* **target**: the clip ``w = E_vid(o_{1:T})``, noised by flow matching to ``w_tau``. The model
+  predicts the velocity ``eps - w`` on these tokens.
+* **keyframes** ``o_T``: each keyframe is one pixel frame, VAE-encoded on its own (exactly what
+  inference can build from a PNG), clean (timestep 0), and placed at its true time by RoPE
+  coordinates ``[f, f+1) / fps``. The keyframe at frame 0 instead *replaces* latent frame 0 of
+  the target, which encodes exactly that one frame; it is excluded from the loss.
+* **reference**: the rendered skeleton flow ``S_{1:T}``, VAE-encoded like the clip, clean, at
+  ``1 / ref_downscale`` resolution with its spatial coordinates scaled back onto the target
+  grid.
+
+Stereo clips hold both views in every group, each view encoded separately and laid end to end
+inside the group; a fourth RoPE axis (rope4d.py) tells the views apart. Mono clips use plain
+three-axis coordinates. The transformer is bidirectional, so a keyframe appended at the end
+of the sequence conditions every frame.
+"""
+
+from __future__ import annotations
+
+import torch
+
+AUDIO_TOKENS = 1        # LTX-2 is audio-video; training feeds a one-token silent audio stream
+
+
+def sample_keyframe_indices(num_frames: int, count_weights: dict[int, float],
+                            include_last_p: float, generator: torch.Generator | None = None
+                            ) -> list[int]:
+    """Keyframes of one training clip (paper App. B.2.4): 1 to 4, the first frame always.
+
+    The count is drawn from ``count_weights``; the remaining keyframes sit at random frames.
+    With probability ``include_last_p`` one of them is the clip's last frame, which is where
+    inference always places the shared boundary keyframe of a chunk.
+    """
+    counts = sorted(count_weights)
+    probs = torch.tensor([float(count_weights[c]) for c in counts])
+    n = min(counts[int(torch.multinomial(probs, 1, generator=generator))], num_frames)
+    idx = {0}
+    if n >= 2 and torch.rand((), generator=generator) < include_last_p:
+        idx.add(num_frames - 1)
+    pool = torch.randperm(num_frames - 1, generator=generator) + 1     # frames 1..F-1
+    for f in pool.tolist():
+        if len(idx) >= n:
+            break
+        idx.add(f)
+    return sorted(idx)
+
+
+def _per_view(coords: torch.Tensor, num_views: int) -> list[torch.Tensor]:
+    if num_views == 1:
+        return [coords]                       # mono: no view axis, pretrained RoPE untouched
+    from .rope4d import add_view_coord
+    return [add_view_coord(coords, v) for v in range(num_views)]
+
+
+def build_video_coords(rope, num_latent_frames: int, height: int, width: int,
+                       ref_downscale: int, device, fps: float, keyframe_indices: list[int],
+                       num_views: int = 1) -> torch.Tensor:
+    """RoPE coordinates for ``[target | keyframes | reference]``, matching ``build_step``.
+
+    ``height``/``width`` are one view's latent grid. Keyframe tokens are one pixel frame wide
+    at ``[f, f+1)`` (already divided by fps, like ``prepare_video_coords``).
+    """
+    tgt = rope.prepare_video_coords(1, num_latent_frames, height, width, device=device, fps=fps)
+    parts = _per_view(tgt, num_views)
+    if keyframe_indices:
+        one = rope.prepare_video_coords(1, 1, height, width, device=device, fps=fps)
+        for f in keyframe_indices:
+            c = one.clone()
+            c[:, 0, :, 0] = float(f) / fps
+            c[:, 0, :, 1] = float(f + 1) / fps
+            parts += _per_view(c, num_views)
+    if ref_downscale == 1:
+        ref = tgt.clone()
+    else:
+        ref = rope.prepare_video_coords(1, num_latent_frames, height // ref_downscale,
+                                        width // ref_downscale, device=device, fps=fps)
+        ref[:, 1:3] *= ref_downscale
+    parts += _per_view(ref, num_views)
+    return torch.cat(parts, dim=2)
+
+
+def build_step(target: torch.Tensor, reference: torch.Tensor,
+               keyframes: dict[int, torch.Tensor], sigma: torch.Tensor,
+               tokens_per_frame: int, num_views: int = 1):
+    """Noise the target and lay out one sequence.
+
+    Args:
+        target: ``(1, V * seq, C)`` packed clean target tokens, view-major.
+        reference: ``(1, V * ref_seq, C)`` packed skeleton-flow tokens.
+        keyframes: ``{pixel frame: (1, V * tokens_per_frame, C)}``, one-frame encodes.
+        sigma: ``(1,)`` flow-matching noise level in [0, 1].
+        tokens_per_frame: spatial tokens of one view's latent frame.
+
+    Returns:
+        ``(hidden, timestep, velocity, pinned, appended)``: the input tokens, per-token
+        timesteps (``sigma * 1000`` on noised tokens, 0 on clean ones), the velocity target
+        for the first ``V * seq`` outputs, a mask of target tokens replaced by the frame-0
+        keyframe, and the appended keyframe indices in sequence order.
+    """
+    b, seq, _ = target.shape
+    per_view = seq // num_views
+    noise = torch.randn_like(target)
+    s = sigma.view(-1, 1, 1).to(target.dtype)
+    noisy = (1.0 - s) * target + s * noise
+    pinned = torch.zeros(b, seq, dtype=torch.bool, device=target.device)
+    if 0 in keyframes:                         # frame 0 replaces latent frame 0 of each view
+        noisy = noisy.clone()
+        for v in range(num_views):
+            dst, src = v * per_view, v * tokens_per_frame
+            noisy[:, dst:dst + tokens_per_frame] = keyframes[0][:, src:src + tokens_per_frame]
+            pinned[:, dst:dst + tokens_per_frame] = True
+    appended = sorted(f for f in keyframes if f != 0)
+    hidden = torch.cat([noisy] + [keyframes[f] for f in appended] + [reference], dim=1)
+    ts = (sigma.view(-1, 1) * 1000.0).to(target.dtype).expand(b, hidden.shape[1]).clone()
+    ts[:, :seq] = torch.where(pinned, torch.zeros_like(ts[:, :seq]), ts[:, :seq])
+    ts[:, seq:] = 0.0
+    return hidden, ts, noise - target, pinned, appended
+
+
+def flow_matching_loss(pred: torch.Tensor, velocity: torch.Tensor,
+                       pinned: torch.Tensor) -> torch.Tensor:
+    """MSE between predicted and true velocity over the noised target tokens."""
+    keep = (~pinned).unsqueeze(-1)
+    diff = (pred.float() - velocity.float()) ** 2 * keep
+    return diff.sum() / keep.expand_as(diff).sum().clamp(min=1)
+
+
+def shifted_logit_normal_sigma(seq_len: int, device=None, std: float = 1.0, eps: float = 1e-3,
+                               uniform_prob: float = 0.1) -> torch.Tensor:
+    """Noise level for one sample, as LTX-2's trainer draws it.
+
+    Adapted from LTX-2 ``ltx_trainer.timestep_samplers.ShiftedLogitNormalTimestepSampler``
+    (Apache-2.0): a logit-normal whose mean shifts with the token count, stretched between its
+    0.5th and 99.9th percentiles, mixed with 10% uniform samples.
+    """
+    m = (2.05 - 0.95) / (4096 - 1024)
+    mu = m * seq_len + (0.95 - m * 1024)
+    x = torch.sigmoid(torch.randn((1,), device=device) * std + mu)
+    p999 = torch.sigmoid(torch.tensor(mu + 3.0902 * std, device=device))
+    p005 = torch.sigmoid(torch.tensor(mu - 2.5758 * std, device=device))
+    z = (x - p005) / (p999 - p005)
+    z = torch.clamp(torch.where(z >= eps, z, 2 * eps - z), 0, 1)
+    uniform = (1 - eps) * torch.rand((1,), device=device) + eps
+    return torch.where(torch.rand((1,), device=device) > uniform_prob, z, uniform)

robotryon/video/dataset.py

@@ -0,0 +1,167 @@
+"""Training clips for the video model: robot video, its skeleton flow, and its caption.
+
+Episodes follow the shared format (``<root>/<episode>/meta.json``, ``video_<view>.mp4``,
+``keypoints_<view>.npz``). One item is a random window of one episode:
+
+* length ``F`` drawn uniformly from the ``8k + 1`` values in ``frames_range`` that fit,
+  start drawn uniformly;
+* stereo episodes give both views (``V = 2``), except with probability ``mono_p`` where one
+  random view is used (``V = 1``), so one model learns stereo and mono generation;
+* one crop rectangle (bottom-anchored, scale in ``[min_scale, 1]``) and one horizontal flip
+  for every view and for the skeleton. A mirrored stereo pair also swaps its views, since the
+  reflected right image is what a left camera would see;
+* the skeleton flow ``S_{1:T}`` is rendered on the fly from the keypoints, after the paper's
+  augmentation (``robotryon.skeleton.SkeletonAugment``: shift every keypoint by up to half its
+  finger segment, drop the whole skeleton with p=0.2, each keypoint with p=0.2), held fixed
+  over the clip. The views of a stereo clip are augmented jointly (the hands of both views
+  share one drop-all decision).
+
+The skeleton is drawn directly at ``skeleton_size`` (the reference resolution), the same
+resolution inference renders at.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from ..skeleton import SkeletonAugment, draw_skeleton_video
+from .chunks import allowed_frame_counts
+from .io import read_frames
+
+
+@dataclass
+class Episode:
+    path: Path
+    views: tuple[str, ...]
+    num_frames: int
+    caption: str
+    root_index: int
+
+
+def caption_text(template: str, meta: dict, embodiment: str | None = None) -> str:
+    """The text prompt: ``template`` filled with the episode's caption (and embodiment)."""
+    return template.format(caption=meta.get("caption", ""),
+                           embodiment=embodiment or meta.get("embodiment", ""))
+
+
+def load_episodes(roots: list, min_frames: int, fps: float,
+                  caption_template: str = "{caption}") -> tuple[list[Episode], list[float]]:
+    """All usable episodes under ``roots`` and their sampling weights.
+
+    ``roots`` items are paths or ``{path, weight}``. A root's weight is split over its
+    episodes in proportion to their length.
+    """
+    episodes, weights = [], []
+    for ri, r in enumerate(roots):
+        path, w = (r, 1.0) if isinstance(r, str) else (r["path"], float(r.get("weight", 1.0)))
+        found = []
+        for m in sorted(Path(path).glob("*/meta.json")):
+            meta = json.loads(m.read_text())
+            if abs(float(meta["fps"]) - fps) > 1e-6:
+                raise ValueError(f"{m}: fps {meta['fps']} != {fps}; resample the episode first")
+            if int(meta["num_frames"]) < min_frames:
+                continue
+            found.append(Episode(m.parent, tuple(meta["views"]), int(meta["num_frames"]),
+                                 caption_text(caption_template, meta), ri))
+        if not found:
+            raise FileNotFoundError(f"no episode with >= {min_frames} frames under {path}")
+        total = sum(e.num_frames for e in found)
+        episodes += found
+        weights += [w * e.num_frames / total for e in found]
+    return episodes, weights
+
+
+def crop_rect(W: int, H: int, aspect: float, scale: float, rng: random.Random):
+    """A crop of the given aspect, bottom-anchored (arms enter from the bottom edge)."""
+    bw, bh = (int(round(H * aspect)), H) if W / H > aspect else (W, int(round(W / aspect)))
+    w = max(8, int(round(bw * scale)))
+    h = min(bh, max(8, int(round(w / aspect))))
+    return rng.randint(0, max(0, W - w)), H - h, w, h
+
+
+def crop_frames(frames: np.ndarray, rect, flip: bool, size) -> np.ndarray:
+    x, y, w, h = rect
+    out = np.empty((len(frames), size[1], size[0], 3), np.uint8)
+    for i, f in enumerate(frames):
+        c = f[y:y + h, x:x + w]
+        out[i] = cv2.resize(c[:, ::-1] if flip else c, size, interpolation=cv2.INTER_AREA)
+    return out
+
+
+def crop_keypoints(kp: np.ndarray, rect, flip: bool, size) -> np.ndarray:
+    """Map pixel keypoints through the same crop, flip and resize as the frames."""
+    x, y, w, h = rect
+    out = np.array(kp, np.float32, copy=True)
+    out[..., 0] = (out[..., 0] - x + 0.5) * size[0] / w - 0.5
+    out[..., 1] = (out[..., 1] - y + 0.5) * size[1] / h - 0.5
+    if flip:
+        out[..., 0] = size[0] - 1 - out[..., 0]
+    return out
+
+
+def to_tensor(a: np.ndarray) -> torch.Tensor:
+    """``(V, F, H, W, 3)`` uint8 RGB -> ``(V, 3, F, H, W)`` float in [-1, 1]."""
+    return torch.from_numpy(np.ascontiguousarray(a)).permute(0, 4, 1, 2, 3).float() / 127.5 - 1
+
+
+class VideoClipDataset(Dataset):
+    """Random training windows. Items are unbatched dicts (clip lengths differ)."""
+
+    def __init__(self, episodes: list[Episode], out_size, skeleton_size, frames_range,
+                 min_scale: float, flip_p: float, mono_p: float, augment: SkeletonAugment,
+                 draw_kw: dict | None = None, fps: float = 20.0, seed: int = 0):
+        self.episodes = episodes
+        self.out_size, self.skeleton_size = tuple(out_size), tuple(skeleton_size)
+        self.choices = allowed_frame_counts(*frames_range)
+        self.min_scale, self.flip_p, self.mono_p = min_scale, flip_p, mono_p
+        self.augment, self.draw_kw = augment, dict(draw_kw or {})
+        self.fps, self.seed, self.epoch = fps, seed, 0
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+    def __len__(self) -> int:
+        return len(self.episodes)
+
+    def __getitem__(self, idx: int) -> dict:
+        ep = self.episodes[idx]
+        rng = random.Random(hash((self.seed, self.epoch, idx)) & 0xFFFFFFFF)
+        nrng = np.random.default_rng(rng.getrandbits(64))
+        F = rng.choice([f for f in self.choices if f <= ep.num_frames])
+        start = rng.randint(0, ep.num_frames - F)
+        views = list(ep.views)
+        if len(views) > 1 and rng.random() < self.mono_p:
+            views = [rng.choice(views)]
+
+        frames = [read_frames(ep.path / f"video_{v}.mp4", start, F) for v in views]
+        kps, vis = [], []
+        for v in views:
+            z = np.load(ep.path / f"keypoints_{v}.npz")
+            kps.append(z["kp"][start:start + F])
+            vis.append(z["vis"][start:start + F].astype(bool))
+
+        H, W = frames[0].shape[1:3]
+        rect = crop_rect(W, H, self.out_size[0] / self.out_size[1],
+                         rng.uniform(self.min_scale, 1.0), rng)
+        flip = rng.random() < self.flip_p
+        video = np.stack([crop_frames(f, rect, flip, self.out_size) for f in frames])
+        kp = np.concatenate([crop_keypoints(k, rect, flip, self.skeleton_size) for k in kps], 1)
+        kp, vs = self.augment(kp, np.concatenate(vis, 1), nrng, temporal=True)
+        sw, sh = self.skeleton_size
+        skel = np.stack([draw_skeleton_video(kp[:, 2 * i:2 * i + 2], vs[:, 2 * i:2 * i + 2],
+                                             sw, sh, **self.draw_kw)
+                         for i in range(len(views))])
+        if flip and len(views) == 2:                  # a mirrored stereo pair swaps cameras
+            video, skel = video[::-1], skel[::-1]
+        return {"video": to_tensor(video), "condition": to_tensor(skel),
+                "caption": ep.caption, "num_frames": F, "fps": self.fps,
+                "meta": {"episode": ep.path.name, "start": start, "views": len(views),
+                         "flip": flip}}

robotryon/video/io.py

@@ -0,0 +1,52 @@
+"""Minimal mp4 / png I/O for the video model (RGB uint8 arrays, OpenCV underneath)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+
+def read_frames(path: Path, start: int, count: int, size=None) -> np.ndarray:
+    """Frames ``[start, start + count)`` as ``(count, H, W, 3)`` RGB, optionally resized
+    (INTER_AREA) to ``size = (W, H)``."""
+    cap = cv2.VideoCapture(str(path))
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start)
+    frames = []
+    for _ in range(count):
+        ok, f = cap.read()
+        if not ok:
+            break
+        if size is not None:
+            f = cv2.resize(f, tuple(size), interpolation=cv2.INTER_AREA)
+        frames.append(f[..., ::-1])
+    cap.release()
+    if len(frames) < count:
+        raise RuntimeError(f"{path}: wanted {count} frames from {start}, got {len(frames)}")
+    return np.ascontiguousarray(np.stack(frames))
+
+
+def num_frames(path: Path) -> int:
+    cap = cv2.VideoCapture(str(path))
+    n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+    return n
+
+
+def read_image(path: Path) -> np.ndarray:
+    """PNG -> ``(H, W, 3)`` RGB uint8 at its own resolution."""
+    img = cv2.imread(str(path))
+    if img is None:
+        raise FileNotFoundError(path)
+    return np.ascontiguousarray(img[..., ::-1])
+
+
+def write_video(frames: np.ndarray, path: Path, fps: float) -> None:
+    """``(F, H, W, 3)`` RGB uint8 -> mp4 (mp4v)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    h, w = frames.shape[1:3]
+    wr = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+    for f in frames:
+        wr.write(np.ascontiguousarray(f[..., ::-1]))
+    wr.release()

robotryon/video/pipeline.py

@@ -0,0 +1,271 @@
+"""Video synthesis with the in-context LoRA (paper Eq. 3.4 and App. B.4).
+
+``StereoIcLoraPipeline`` extends diffusers' ``LTX2InContextPipeline`` with two things:
+
+* **stereo**: ``conditions`` and ``reference_conditions`` are per-view lists. Each view is
+  prepared by the parent and the token groups are interleaved by role,
+  ``[base(L) base(R) | kf(L) kf(R) | ref(L) ref(R)]``, so the reference tokens stay at the end
+  of the sequence where the parent's attention mask expects them. A fourth RoPE axis
+  (rope4d.py) tells the views apart; mono uses the pipeline's own coordinates.
+* **exact keyframe times**: diffusers places a keyframe at a latent index, which snaps it to
+  pixel frame ``8k - 7``. Training places keyframes at their exact pixel frame ``[f, f+1)``,
+  so the appended keyframe coordinates are rewritten to the requested pixel frames. The
+  keyframe at frame 0 replaces latent frame 0 and needs no rewrite.
+
+``generate`` runs the two-stage recipe on one clip. Both stages use the full LTX-2.5
+transformer with the IC-LoRA and Lightricks' distilled LoRA, the published distilled sigma
+schedules (8 steps, then 3 refinement steps from sigma 0.909) and neutral guidance. Both stages
+are conditioned on the same keyframes and skeleton flow. ``two_stage.mode``:
+
+* ``same_res``: stage 1 already at the output resolution; stage 2 refines at that resolution.
+* ``upsample``: stage 1 at half the output resolution, the stage-1 latents are upsampled x2 by
+  the LTX-2.5 spatial latent upsampler, and stage 2 refines at the output resolution.
+
+The skeleton flow is always given at one ``skeleton_size`` (512x288); each stage encodes it
+with ``reference_downscale = stage width / skeleton width`` (1 at 512x288, 2 at 1024x576),
+exactly as the two training stages did.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import numpy as np
+import torch
+from diffusers.pipelines.ltx2 import LTX2InContextPipeline
+
+from .chunks import latent_index
+from .rope4d import add_view_coord, install_view_axis_rope, set_view_coord
+
+MAX_SEQ = 256                       # text tokens; training used the same
+DISTILLED_LORA = "ltx-2.5-22b-distilled-lora-450-bf16.safetensors"
+
+
+def snapshot_dir(repo_or_path: str) -> str:
+    """A local directory for a Hugging Face repo id (newest cached snapshot) or a path."""
+    if os.path.isdir(repo_or_path):
+        return repo_or_path
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.constants import HF_HUB_CACHE
+    slug = "models--" + repo_or_path.replace("/", "--")
+    hits = sorted(glob.glob(os.path.join(HF_HUB_CACHE, slug, "snapshots", "*")))
+    return hits[-1] if hits else snapshot_download(repo_or_path)
+
+
+class StereoIcLoraPipeline(LTX2InContextPipeline):
+    """``LTX2InContextPipeline`` with a view axis and exact-frame keyframes."""
+
+    num_views = 1
+    exact_kf_px: list[int] | None = None       # pixel frames of the appended keyframes
+
+    def _unpack_latents(self, latents, num_frames, height, width, patch_size=1,
+                        patch_size_t=1):
+        """Unpack the views' base tokens into a batch ``(V, C, F, H, W)`` for decoding."""
+        per_view = num_frames * height * width
+        views = latents.shape[1] // per_view
+        if views < 2:
+            return super()._unpack_latents(latents, num_frames, height, width, patch_size,
+                                           patch_size_t)
+        return torch.cat([super(StereoIcLoraPipeline, self)._unpack_latents(
+            latents[:, v * per_view:(v + 1) * per_view], num_frames, height, width,
+            patch_size, patch_size_t) for v in range(views)], dim=0)
+
+    def _retime_keyframes(self, kf_coords, n_kf, kw):
+        """Move the appended keyframe blocks to their exact pixel frames ``[f, f+1) / fps``."""
+        if not self.exact_kf_px or not n_kf:
+            return kf_coords
+        fps = float(kw["frame_rate"])
+        tpf = (int(kw["height"]) // 32) * (int(kw["width"]) // 32)
+        if n_kf != len(self.exact_kf_px) * tpf:
+            raise ValueError(f"{n_kf} appended keyframe tokens do not match "
+                             f"{len(self.exact_kf_px)} keyframes of {tpf} tokens")
+        out = kf_coords.clone()
+        for j, f in enumerate(self.exact_kf_px):
+            sl = slice(j * tpf, (j + 1) * tpf)
+            out[:, 0, sl, 0] = float(f) / fps
+            out[:, 0, sl, 1] = float(f + 1) / fps
+        return out
+
+    def prepare_latents(self, conditions=None, reference_conditions=None, **kw):
+        views = self.num_views
+        conditions = conditions or [None] * views
+        reference_conditions = reference_conditions or [None] * views
+        lat_in = kw.pop("latents", None)       # stage 2: stage-1 (upsampled) latents per view
+
+        def view_latents(v):
+            if lat_in is None:
+                return None
+            if lat_in.dim() == 5 and lat_in.shape[0] == views:
+                return lat_in[v:v + 1]
+            raise ValueError(f"latents {tuple(lat_in.shape)} do not split into {views} views")
+
+        self._single_view_coords = True        # reference coords from the parent stay 3-axis
+        try:
+            per = [super(StereoIcLoraPipeline, self).prepare_latents(
+                       conditions=conditions[v], reference_conditions=reference_conditions[v],
+                       latents=view_latents(v), **kw) for v in range(views)]
+        finally:
+            self._single_view_coords = False
+
+        groups = {k: [[] for _ in range(3)] for k in ("lat", "mask", "clean")}
+        kf_co, ref_co, ref_masks, ref_counts = [], [], [], []
+        for v, (lat, mask, clean, acoords, nref, rmask, rcounts) in enumerate(per):
+            n_app = 0 if acoords is None else acoords.shape[2]
+            n_kf, n_base = n_app - nref, lat.shape[1] - n_app
+            bounds = (slice(0, n_base), slice(n_base, n_base + n_kf),
+                      slice(n_base + n_kf, lat.shape[1]))
+            for name, src in (("lat", lat), ("mask", mask), ("clean", clean)):
+                for g, sl in enumerate(bounds):
+                    groups[name][g].append(src[:, sl])
+            if acoords is not None:
+                a4 = set_view_coord(acoords, v) if views > 1 else acoords
+                kf_co.append(self._retime_keyframes(a4[:, :, :n_kf], n_kf, kw))
+                ref_co.append(a4[:, :, n_kf:])
+            if rmask is not None:
+                ref_masks.append(rmask)
+            ref_counts.extend(rcounts or [])
+
+        def flat(g):
+            return torch.cat(g[0] + g[1] + g[2], dim=1)
+        coords = torch.cat(kf_co + ref_co, dim=2) if (kf_co or ref_co) else None
+        return (flat(groups["lat"]), flat(groups["mask"]), flat(groups["clean"]), coords,
+                sum(x.shape[1] for x in groups["lat"][2]),
+                torch.cat(ref_masks, dim=1) if ref_masks else None, ref_counts)
+
+
+def _install_stereo(pipe: StereoIcLoraPipeline) -> None:
+    """View-axis RoPE, and a base grid that covers every view (view axis stamped per view)."""
+    install_view_axis_rope(pipe.transformer)
+    pipe._single_view_coords = False
+    rope = pipe.transformer.rope
+    orig = rope.prepare_video_coords
+
+    def widened(*args, **kwargs):
+        plain = orig(*args, **kwargs)
+        if pipe._single_view_coords:
+            return plain
+        return torch.cat([add_view_coord(plain, v) for v in range(pipe.num_views)], dim=2)
+
+    rope.prepare_video_coords = widened
+
+
+def build_pipeline(lora: str, *, base: str = "Lightricks/LTX-2.5-Diffusers",
+                   transformer_subfolder: str = "transformer_full", views: int = 1,
+                   attention_backend: str | None = "native", upsample: bool = False,
+                   device: str = "cuda") -> StereoIcLoraPipeline:
+    """Load LTX-2.5 + IC-LoRA + distilled LoRA, without the text encoder.
+
+    Args:
+        lora: checkpoint directory holding ``pytorch_lora_weights.safetensors`` (or the file).
+        views: views generated jointly in one sequence (2 = stereo, 1 = mono).
+        upsample: also load the spatial latent upsampler (``two_stage.mode: upsample``).
+    """
+    from diffusers import LTX2VideoTransformer3DModel
+
+    root = snapshot_dir(base)
+    pipe = StereoIcLoraPipeline.from_pretrained(
+        root, transformer=LTX2VideoTransformer3DModel.from_pretrained(
+            root, subfolder=transformer_subfolder, dtype=torch.bfloat16),
+        text_encoder=None, tokenizer=None, dtype=torch.bfloat16)
+    pipe.num_views = int(views)
+    if views > 1:
+        _install_stereo(pipe)
+    lora_dir, name = ((os.path.dirname(lora), os.path.basename(lora)) if os.path.isfile(lora)
+                      else (lora, "pytorch_lora_weights.safetensors"))
+    pipe.load_lora_weights(lora_dir, weight_name=name, adapter_name="iclora")
+    if not any("lora" in k for k, _ in pipe.transformer.named_parameters()):
+        raise RuntimeError(f"no LoRA weights attached from {lora}")
+    pipe.load_lora_weights(root, weight_name=DISTILLED_LORA, adapter_name="distilled")
+    pipe.set_adapters(["iclora", "distilled"], [1.0, 1.0])
+    if attention_backend:
+        pipe.transformer.set_attention_backend(attention_backend)
+    pipe.vae.enable_tiling()
+    pipe.vae.enable_slicing()
+    pipe.to(device)
+    pipe.upsampler = None
+    if upsample:
+        from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline
+        from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
+        model = LTX2LatentUpsamplerModel.from_pretrained(
+            root, subfolder="latent_upsampler", dtype=torch.bfloat16).to(device)
+        pipe.upsampler = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=model)
+    return pipe
+
+
+def _u8(img: torch.Tensor) -> np.ndarray:
+    """``(3, H, W)`` in [-1, 1] -> ``(1, H, W, 3)`` uint8 (keyframe convention)."""
+    return ((img.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0)[None].cpu().numpy()
+
+
+def _f01(video: torch.Tensor) -> np.ndarray:
+    """``(3, F, H, W)`` in [-1, 1] -> ``(F, H, W, 3)`` float in [0, 1] (reference convention)."""
+    return ((video.clamp(-1, 1) + 1) / 2).permute(1, 2, 3, 0).float().cpu().numpy()
+
+
+def _to_uint8(out: torch.Tensor, views: int) -> torch.Tensor:
+    """Pipeline ``pt`` output in [0, 1], ``(V, F, 3, H, W)`` -> ``(V, F, H, W, 3)`` uint8."""
+    v = out if out.shape[0] == views else out[0]
+    frames = [x if x.shape[0] == 3 else x.permute(1, 0, 2, 3) for x in v]   # -> (3, F, H, W)
+    return torch.stack([(x.float().clamp(0, 1) * 255).byte().permute(1, 2, 3, 0).cpu()
+                        for x in frames])
+
+
+@torch.no_grad()
+def generate(pipe: StereoIcLoraPipeline, *, skeleton: torch.Tensor,
+             keyframes: list[tuple[int, torch.Tensor]], text: tuple, num_frames: int,
+             size: tuple[int, int], fps: float, seed: int = 0,
+             mode: str = "same_res") -> torch.Tensor:
+    """Synthesize one clip.
+
+    Args:
+        skeleton: ``(V, 3, F, h, w)`` skeleton flow in [-1, 1] at the reference resolution.
+        keyframes: ``[(pixel frame, (V, 3, H, W) in [-1, 1])]`` sorted, first at frame 0. Any
+            resolution: the pipeline resizes them for each stage.
+        text: ``(prompt_embeds, mask, negative_embeds, negative_mask)`` (see text.py).
+        size: output ``(width, height)`` of one view.
+        mode: ``same_res`` or ``upsample``.
+
+    Returns:
+        ``(V, F, height, width, 3)`` uint8 RGB.
+    """
+    from diffusers.pipelines.ltx2 import LTX2ReferenceCondition, LTX2VideoCondition
+    from diffusers.pipelines.ltx2.utils import (DISTILLED_SIGMA_VALUES,
+                                                STAGE_2_DISTILLED_SIGMA_VALUES)
+    if mode not in ("same_res", "upsample"):
+        raise ValueError(f"two_stage.mode must be same_res or upsample, got {mode!r}")
+    views = int(skeleton.shape[0])
+    if views != pipe.num_views:
+        raise ValueError(f"{views} views given to a {pipe.num_views}-view pipeline")
+    px = [p for p, _ in keyframes]
+    lat = [latent_index(p, num_frames) for p in px]
+    if px != sorted(set(px)) or len(set(lat)) != len(lat):
+        raise ValueError(f"keyframes must be sorted and in distinct latent frames: {px}")
+    pipe.exact_kf_px = [p for p, li in zip(px, lat) if li != 0]
+
+    W, H = size
+    skel_w = skeleton.shape[-1]
+    W1, H1 = (W // 2, H // 2) if mode == "upsample" else (W, H)
+    pe, pm, npe, npm = (x.cuda() for x in text)
+    kw = dict(prompt_embeds=pe, prompt_attention_mask=pm, negative_prompt_embeds=npe,
+              negative_prompt_attention_mask=npm,
+              conditions=[[LTX2VideoCondition(frames=_u8(img[v]), index=li, crf=0)
+                           for li, (_, img) in zip(lat, keyframes)] for v in range(views)],
+              reference_conditions=[[LTX2ReferenceCondition(frames=_f01(skeleton[v]))]
+                                    for v in range(views)],
+              reference_downscale_factor=W1 // skel_w, max_sequence_length=MAX_SEQ,
+              height=H1, width=W1, num_frames=num_frames, frame_rate=fps,
+              generator=torch.Generator("cuda").manual_seed(seed), return_dict=False,
+              audio_latents=torch.zeros(1, pipe.audio_latent_channels, 1,
+                                        pipe.audio_mel_bins // pipe.audio_vae_mel_compression_ratio,
+                                        device="cuda", dtype=torch.float32),
+              guidance_scale=1.0, stg_scale=0.0, modality_scale=1.0, audio_guidance_scale=1.0,
+              audio_stg_scale=0.0, audio_modality_scale=1.0)
+
+    lat1 = pipe(sigmas=list(DISTILLED_SIGMA_VALUES), output_type="latent", **kw)[0]
+    if mode == "upsample":
+        lat1 = pipe.upsampler(latents=lat1, output_type="latent", return_dict=False)[0]
+    s2 = list(STAGE_2_DISTILLED_SIGMA_VALUES)
+    kw.update(height=H, width=W, reference_downscale_factor=W // skel_w)
+    out = pipe(sigmas=s2, latents=lat1, noise_scale=s2[0], output_type="pt", **kw)[0]
+    return _to_uint8(out, views)

robotryon/video/rope4d.py

@@ -0,0 +1,129 @@
+"""A fourth RoPE axis for the camera view.
+
+Both stereo views are generated in one sequence: each view is encoded by the VAE on its own
+and the token runs are concatenated. To let attention relate the same (t, h, w) position in the
+two cameras, every token also carries a view coordinate.
+
+Passing four axes to the stock LTX-2 rotary embedding would re-cut its channel-to-frequency
+map (``dim // (2 * num_axes)`` frequencies per axis) and perturb every pretrained head at once.
+Instead the t/h/w frequencies stay exactly where they were, and the view rotation is written
+into the last ``VIEW_CHANNELS_PER_HEAD`` channels of every head's slice (64 of 2048 channels),
+using the lowest frequency of the table. With two views the coordinates sit a quarter turn
+apart. Three-axis coordinates fall through to the original implementation unchanged, so mono
+clips see the pretrained embedding exactly.
+"""
+
+from __future__ import annotations
+
+import types
+
+import torch
+
+VIEW_CHANNELS_PER_HEAD = 2
+BASE_VIEWS = 2
+
+
+def view_channel_indices(dim: int, num_heads: int, per_head: int, device=None) -> torch.Tensor:
+    """Channels taken by the view axis: the last ``per_head`` of each head's slice, head-major."""
+    head_half = (dim // 2) // num_heads
+    if per_head >= head_half:
+        raise ValueError(f"per_head={per_head} would take a head's whole {head_half} channels")
+    return torch.tensor([h * head_half + head_half - per_head + c
+                         for h in range(num_heads) for c in range(per_head)],
+                        dtype=torch.long, device=device)
+
+
+def _forward_with_view_axis(self, coords: torch.Tensor, device=None):
+    """Drop-in for ``LTX2AudioVideoRotaryPosEmbed.forward`` accepting a 4th (view) axis."""
+    device = device or coords.device
+    if coords.shape[1] != 4:
+        return self._forward_3axis(coords, device)
+    if self.rope_type != "split":
+        raise NotImplementedError(f"view axis needs rope_type='split', got {self.rope_type!r}")
+    if coords.ndim == 4:                                      # [start, end) -> midpoints
+        start, end = coords.chunk(2, dim=-1)
+        coords = ((start + end) / 2.0).squeeze(-1)
+    spatial, view = coords[:, :3], coords[:, 3]
+
+    # t/h/w exactly as the original: three axes, dim // 6 frequencies each
+    max_positions = (self.base_num_frames, self.base_height, self.base_width)
+    grid = torch.stack([spatial[:, i] / max_positions[i] for i in range(3)], dim=-1).to(device)
+    freqs_dtype = torch.float64 if self.double_precision else torch.float32
+    pow_indices = torch.pow(self.theta, torch.linspace(0.0, 1.0, steps=self.dim // 6,
+                                                       dtype=freqs_dtype, device=device))
+    freqs_base = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)
+    angles = ((grid.unsqueeze(-1) * 2 - 1) * freqs_base).transpose(-1, -2).flatten(2)
+
+    # pad the angle (cos 1 / sin 0 in the original) to dim // 2
+    pad = self.dim // 2 - angles.shape[-1]
+    if pad < 0:
+        raise RuntimeError(f"angle width {angles.shape[-1]} exceeds dim//2")
+    if pad > 0:
+        angles = torch.cat([torch.zeros_like(angles[:, :, :pad]), angles], dim=-1)
+
+    # the view rotation in the tail channels of every head
+    idx = view_channel_indices(self.dim, self.num_attention_heads, VIEW_CHANNELS_PER_HEAD,
+                               device=device)
+    view_angle = ((view.to(device).unsqueeze(-1) / BASE_VIEWS * 2 - 1)
+                  * freqs_base[:VIEW_CHANNELS_PER_HEAD])
+    angles = angles.clone()
+    angles[:, :, idx] = view_angle.repeat(1, 1, self.num_attention_heads).to(angles.dtype)
+
+    b, n = angles.shape[:2]
+    cos = angles.cos().reshape(b, n, self.num_attention_heads, -1).swapaxes(1, 2)
+    sin = angles.sin().reshape(b, n, self.num_attention_heads, -1).swapaxes(1, 2)
+    return cos, sin
+
+
+def install_view_axis_rope(transformer_or_rope) -> None:
+    """Patch the video rope (only) to accept view coordinates. Idempotent."""
+    rope = getattr(transformer_or_rope, "rope", transformer_or_rope)
+    if getattr(rope, "_view_axis_installed", False):
+        return
+    rope._forward_3axis = types.MethodType(type(rope).forward, rope)
+    rope.forward = types.MethodType(_forward_with_view_axis, rope)
+    rope._view_axis_installed = True
+
+
+def add_view_coord(coords: torch.Tensor, view: int) -> torch.Tensor:
+    """``(B, 3, N, 2)`` -> ``(B, 4, N, 2)`` with the view axis constant at ``[view, view+1)``."""
+    b, _, n, _ = coords.shape
+    v = torch.empty(b, 1, n, 2, dtype=coords.dtype, device=coords.device)
+    v[..., 0] = float(view)
+    v[..., 1] = float(view + 1)
+    return torch.cat([coords, v], dim=1)
+
+
+def set_view_coord(coords: torch.Tensor, view: int) -> torch.Tensor:
+    """Set the view axis, appending it if the coordinates have only three axes."""
+    if coords.shape[1] != 4:
+        return add_view_coord(coords, view)
+    out = coords.clone()
+    out[:, 3, :, 0] = float(view)
+    out[:, 3, :, 1] = float(view + 1)
+    return out
+
+
+@torch.no_grad()
+def check_bands_preserved(rope, device=None, atol: float = 1e-3) -> dict:
+    """Verify the patched rope changes only the view channels (run before training).
+
+    Channels count as changed when they differ by more than ``atol``: the 3-axis and 4-axis
+    paths are computed separately and can differ by float rounding on some CPUs."""
+    device = device or torch.device("cpu")
+    base = rope.prepare_video_coords(1, 4, 9, 16, device=device, fps=20.0)
+    cos3, sin3 = rope._forward_3axis(base, device)
+    cos4, sin4 = rope(add_view_coord(base, 0))
+    idx = view_channel_indices(rope.dim, rope.num_attention_heads, VIEW_CHANNELS_PER_HEAD,
+                               device=cos3.device)
+    keep = torch.ones(rope.dim // 2, dtype=torch.bool, device=cos3.device)
+    keep[idx] = False
+    keep = keep.view(rope.num_attention_heads, -1)[None, :, None, :]
+    d_cos, d_sin = (cos4 - cos3).abs(), (sin4 - sin3).abs()
+    changed = ((d_cos > atol) | (d_sin > atol)).any(dim=2).squeeze(0)
+    return {"n_changed": int(changed.sum()),
+            "n_expected": rope.num_attention_heads * VIEW_CHANNELS_PER_HEAD,
+            "heads_touched": int(changed.any(dim=-1).sum()),
+            "num_heads": rope.num_attention_heads,
+            "max_abs_diff_off_view": float(torch.max((d_cos * keep).max(),
+                                                     (d_sin * keep).max()))}