Complete shared baseline runtime dependencies

external_baselines/frameskip_adapted/build_required_draws.py

@@ -4,18 +4,12 @@
 from __future__ import annotations
 
 import argparse
-import hashlib
 import json
 from pathlib import Path
 
 from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest
 
 
-def canonical_sha256(payload: object) -> str:
-    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
-    return hashlib.sha256(encoded).hexdigest()
-
-
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--manifest", required=True)
@@ -38,15 +32,10 @@ def main() -> None:
     }
     pruned_draws = 0
     full_draws = 0
-    trace_digest = hashlib.sha256()
-
     for global_draw_id in range(start_draw, end_draw):
         microbatch = (global_draw_id - start_draw) // args.global_microbatch_size
         is_pruned = microbatch % period < args.pruned_batches
         trace = manifest.sample(args.seed, LogicalDraw(0, global_draw_id))
-        trace_digest.update(
-            json.dumps(trace.to_dict(), sort_keys=True, separators=(",", ":")).encode()
-        )
         if is_pruned:
             required[trace.dataset_name].add(trace.trajectory_id)
             pruned_draws += 1
@@ -57,7 +46,6 @@ def main() -> None:
         "schema_version": 1,
         "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
         "manifest_path": str(Path(args.manifest).resolve()),
-        "manifest_sha256": manifest.sha256,
         "sampling_seed": args.seed,
         "start_step": args.start_step,
         "updates": args.updates,
@@ -69,7 +57,6 @@ def main() -> None:
         "full_batches": args.full_batches,
         "pruned_draws": pruned_draws,
         "full_draws": full_draws,
-        "logical_trace_sha256": trace_digest.hexdigest(),
         "datasets": {
             name: sorted(ids, key=lambda value: (type(value).__name__, str(value)))
             for name, ids in required.items()
@@ -78,13 +65,11 @@ def main() -> None:
     payload["required_trajectory_count"] = sum(
         len(values) for values in payload["datasets"].values()
     )
-    payload["sha256"] = canonical_sha256(payload)
     destination = Path(args.output)
     destination.parent.mkdir(parents=True, exist_ok=True)
     destination.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
     print(json.dumps({
         "output": str(destination),
-        "sha256": payload["sha256"],
         "pruned_draws": pruned_draws,
         "full_draws": full_draws,
         "required_trajectory_count": payload["required_trajectory_count"],

external_baselines/frameskip_adapted/finalize_cache.py

@@ -1,20 +1,13 @@
 #!/usr/bin/env python3
-"""Validate and merge FrameSkip cache shards into one hash-bound cache."""
+"""Validate and merge FrameSkip cache shards into one portable cache."""
 
 from __future__ import annotations
 
 import argparse
-import hashlib
 import json
 from pathlib import Path
 
 
-def canonical_sha256(payload: object) -> str:
-    return hashlib.sha256(
-        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
-    ).hexdigest()
-
-
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--required-draws", required=True)
@@ -52,7 +45,6 @@ def main() -> None:
         "schema_version": 1,
         "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
         "upstream_commit": "965becf63e9cd5cc339fbf7fde40d1db47efa67a",
-        "required_draws_sha256": required["sha256"],
         "retention_ratio": 0.2,
         "importance": {
             "calculator": "GripperAwareImportanceCalculator",
@@ -79,13 +71,11 @@ def main() -> None:
         ],
         "source_shards": shard_paths,
     }
-    payload["sha256"] = canonical_sha256(payload)
     output = Path(args.output)
     output.parent.mkdir(parents=True, exist_ok=True)
     output.write_text(json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8")
     print(json.dumps({
         "output": str(output),
-        "sha256": payload["sha256"],
         "records": len(payload["records"]),
         "actual_ratio_min": min(r["actual_ratio"] for r in payload["records"]),
         "actual_ratio_max": max(r["actual_ratio"] for r in payload["records"]),

external_baselines/frameskip_adapted/train_frameskip_continuation.py

@@ -68,7 +68,7 @@ def install_adapter(cache_path: str, start_global_draw_id: int) -> None:
             dataset.get_step_data = original_get_step_data
         sample["sampling_trace"]["frameskip"] = {
             "method": cache_payload["method"],
-            "cache_sha256": cache_payload["sha256"],
+            "cache_record_count": len(cache_payload["records"]),
             "is_pruned_minibatch": is_pruned,
             "logical_base_index": logical_index,
             "physical_base_index": actual_index,

external_baselines/frameskip_robotwin_adapted/build_robotwin_cache.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+"""Build a RoboTwin FrameSkip cache with the locked upstream scoring code."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as parquet
+import timm
+from PIL import Image
+
+from starVLA.dataloader.gr00t_lerobot.video import get_frames_by_timestamps
+from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
+
+
+UPSTREAM_COMMIT = "965becf63e9cd5cc339fbf7fde40d1db47efa67a"
+METHOD = "FrameSkip-adapted (Gaussian-TPI fallback)"
+
+
+def load_upstream_module(name: str, path: Path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load upstream module {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def trajectory_paths(root: Path, dataset_name: str, trajectory_id: int) -> tuple[Path, Path]:
+    dataset = root / dataset_name
+    chunk = trajectory_id // 1000
+    parquet_path = dataset / "data" / f"chunk-{chunk:03d}" / f"episode_{trajectory_id:06d}.parquet"
+    video_path = (
+        dataset
+        / "videos"
+        / f"chunk-{chunk:03d}"
+        / "observation.images.cam_high"
+        / f"episode_{trajectory_id:06d}.mp4"
+    )
+    return parquet_path, video_path
+
+
+def load_trajectory(
+    root: Path,
+    dataset_name: str,
+    trajectory_id: int,
+    max_vac_frames: int,
+) -> dict:
+    parquet_path, video_path = trajectory_paths(root, dataset_name, trajectory_id)
+    table = parquet.read_table(parquet_path, columns=["action", "timestamp"])
+    actions = np.asarray(table["action"].to_pylist(), dtype=np.float32)
+    timestamps = np.asarray(table["timestamp"].to_pylist(), dtype=np.float64).reshape(-1)
+    if actions.ndim != 2 or actions.shape[1] != 14:
+        raise RuntimeError(f"unexpected RoboTwin action shape {actions.shape}: {parquet_path}")
+    if len(actions) != len(timestamps):
+        raise RuntimeError(f"action/timestamp length mismatch: {parquet_path}")
+
+    sample_indices = np.linspace(
+        0, len(timestamps) - 1, min(max_vac_frames, len(timestamps)), dtype=int
+    )
+    frames = get_frames_by_timestamps(
+        video_path.as_posix(),
+        timestamps[sample_indices],
+        video_backend="pyav",
+        video_backend_kwargs={},
+    )
+    observations = []
+    for frame in frames:
+        array = np.asarray(frame)
+        if array.ndim == 3 and array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4):
+            array = np.moveaxis(array, 0, -1)
+        observations.append(Image.fromarray(array.astype(np.uint8)))
+    return {"observations": observations, "actions": actions}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", required=True, type=Path)
+    parser.add_argument("--data-root", required=True, type=Path)
+    parser.add_argument("--upstream-root", required=True, type=Path)
+    parser.add_argument("--visual-checkpoint", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--max-trajectories", type=int)
+    args = parser.parse_args()
+
+    candidates = (
+        args.upstream_root / "starVLA" / "starVLA" / "frameskip" / "utils",
+        args.upstream_root / "starVLA" / "frameskip" / "utils",
+    )
+    utils = next(
+        (path for path in candidates if (path / "importance_metrics.py").is_file()),
+        None,
+    )
+    if utils is None:
+        checked = ", ".join(str(path) for path in candidates)
+        raise FileNotFoundError(f"FrameSkip utils not found; checked: {checked}")
+    importance_module = load_upstream_module("frameskip_upstream_importance", utils / "importance_metrics.py")
+    pruner_module = load_upstream_module("frameskip_upstream_pruner", utils / "frame_pruner.py")
+    calculator = importance_module.GripperAwareImportanceCalculator(
+        alpha=0.6,
+        beta=0.2,
+        gamma=0.2,
+        enable_vac=True,
+        visual_encoder_name=None,
+        device=args.device,
+        gripper_boost=2.0,
+    )
+    # Match the verified Bridge cache builder: retain upstream 224px VAC
+    # preprocessing and interpolate DINOv2 positional embeddings dynamically.
+    calculator.visual_encoder = timm.create_model(
+        "vit_small_patch14_dinov2.lvd142m",
+        pretrained=False,
+        num_classes=0,
+        checkpoint_path=str(args.visual_checkpoint),
+        dynamic_img_size=True,
+    ).to(args.device).eval()
+    pruner = pruner_module.FramePruner(
+        compression_ratios=[0.2],
+        preserve_key_stages=True,
+        min_frames_per_trajectory=5,
+    )
+
+    manifest = ValidAnchorManifest.load(args.manifest)
+    records = []
+    for dataset in manifest.datasets:
+        for episode in dataset.episodes:
+            trajectory_id = int(episode.trajectory_id)
+            trajectory = load_trajectory(
+                args.data_root, dataset.name, trajectory_id, max_vac_frames=16
+            )
+            scores = calculator.compute_importance(trajectory)
+            result = pruner.prune_trajectory(trajectory, scores, 0.2)
+            protected = pruner._detect_key_stages(trajectory, scores)
+            retained = set(int(x) for x in result.keep_indices)
+            protected_set = set(int(x) for x in protected)
+            records.append(
+                {
+                    "dataset_name": dataset.name,
+                    "trajectory_id": trajectory_id,
+                    "trajectory_length": int(len(trajectory["actions"])),
+                    "retention_ratio": 0.2,
+                    "actual_ratio": float(result.actual_ratio),
+                    "keep_indices": [int(x) for x in result.keep_indices],
+                    "importance_scores": [float(x) for x in scores],
+                    "protected_indices": sorted(protected_set),
+                    "protected_retained_count": len(protected_set & retained),
+                    "protected_total_count": len(protected_set),
+                    "protected_retention_rate": (
+                        len(protected_set & retained) / len(protected_set)
+                        if protected_set
+                        else 1.0
+                    ),
+                    "vac_frames": len(trajectory["observations"]),
+                }
+            )
+            print(
+                f"cached {dataset.name}/{trajectory_id}: "
+                f"{len(result.keep_indices)}/{len(trajectory['actions'])}",
+                flush=True,
+            )
+            if args.max_trajectories is not None and len(records) >= args.max_trajectories:
+                break
+        if args.max_trajectories is not None and len(records) >= args.max_trajectories:
+            break
+
+    payload = {
+        "schema_version": 1,
+        "method": METHOD,
+        "upstream_commit": UPSTREAM_COMMIT,
+        "required_trajectory_count": sum(
+            len(dataset.episodes) for dataset in manifest.datasets
+        ),
+        "retention_ratio": 0.2,
+        "importance": {
+            "calculator": "GripperAwareImportanceCalculator",
+            "avi_weight": 0.6,
+            "vac_weight": 0.2,
+            "tpi_weight": 0.2,
+            "vac_enabled": True,
+            "max_vac_frames": 16,
+            "visual_encoder": "vit_small_patch14_dinov2.lvd142m",
+            "tpi": "Gaussian sigma_squared=0.2 fallback",
+            "gripper_boost": 2.0,
+        },
+        "pruning": {
+            "first_last": True,
+            "gripper_transition_threshold": 0.1,
+            "top_action_change_percentile": 90,
+            "min_frames_per_trajectory": 5,
+            "mapping": "numpy.searchsorted to the next retained index, clipped at end",
+        },
+        "schedule": {"pruned_minibatches": 5, "full_minibatches": 1},
+        "records": records,
+        "source_shards": [],
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
+    temporary.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    temporary.replace(args.output)
+    print(f"wrote {args.output} records={len(records)}")
+
+
+if __name__ == "__main__":
+    main()

external_baselines/frameskip_robotwin_adapted/train_frameskip_robotwin_continuation.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Run the Florence-RoboTwin FrameSkip continuation from the 10K state."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+from omegaconf import OmegaConf
+
+
+METHOD = "FrameSkip-adapted (Gaussian-TPI fallback)"
+
+
+def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    if payload.get("method") != METHOD:
+        raise RuntimeError(f"unexpected FrameSkip method: {payload.get('method')!r}")
+    if float(payload.get("retention_ratio", -1.0)) != 0.2:
+        raise RuntimeError("FrameSkip cache must use retention_ratio=0.2")
+
+    cache: dict[tuple[str, str], np.ndarray] = {}
+    for record in payload["records"]:
+        key = (record["dataset_name"], str(record["trajectory_id"]))
+        keep = np.asarray(record["keep_indices"], dtype=np.int64)
+        if key in cache:
+            raise RuntimeError(f"duplicate FrameSkip trajectory: {key}")
+        if len(keep) == 0 or np.any(np.diff(keep) <= 0):
+            raise RuntimeError(f"invalid retained indices for {key}")
+        cache[key] = keep
+    return cache, payload
+
+
+def install_adapter(
+    cache_path: str,
+    start_global_draw_id: int,
+    global_microbatch_size: int,
+) -> None:
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+
+    cache, cache_payload = load_cache(cache_path)
+    original_getitem = LeRobotMixtureDataset.__getitem__
+    period = 6  # Five pruned microbatches followed by one full-frame microbatch.
+
+    def frameskip_getitem(self, index):
+        trace = self.sample_trace(index)
+        offset = int(trace.global_draw_id) - start_global_draw_id
+        if offset < 0:
+            raise RuntimeError("FrameSkip draw precedes the 10K continuation boundary")
+
+        microbatch = offset // global_microbatch_size
+        is_pruned = microbatch % period < 5
+        logical_index = int(trace.base_index)
+        actual_index = logical_index
+
+        if is_pruned:
+            key = (trace.dataset_name, str(trace.trajectory_id))
+            try:
+                retained = cache[key]
+            except KeyError as exc:
+                raise RuntimeError(f"FrameSkip cache is missing trajectory {key}") from exc
+            position = int(np.searchsorted(retained, logical_index))
+            actual_index = int(retained[min(position, len(retained) - 1)])
+
+        dataset = self.datasets[trace.dataset_index]
+        original_get_step_data = dataset.get_step_data
+
+        def redirected_get_step_data(trajectory_id, base_index):
+            if str(trajectory_id) != str(trace.trajectory_id):
+                raise RuntimeError("FrameSkip attempted to change trajectory identity")
+            if int(base_index) != logical_index:
+                raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
+            return original_get_step_data(trajectory_id, actual_index)
+
+        dataset.get_step_data = redirected_get_step_data
+        try:
+            sample = original_getitem(self, index)
+        finally:
+            dataset.get_step_data = original_get_step_data
+
+        sample["sampling_trace"]["frameskip"] = {
+            "method": cache_payload["method"],
+            "cache_record_count": len(cache_payload["records"]),
+            "is_pruned_minibatch": is_pruned,
+            "logical_base_index": logical_index,
+            "physical_base_index": actual_index,
+            "retention_ratio": 0.2 if is_pruned else 1.0,
+        }
+        return sample
+
+    LeRobotMixtureDataset.__getitem__ = frameskip_getitem
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, overrides = parser.parse_known_args()
+
+    install_adapter(
+        cache_path=os.environ["FRAMESKIP_CACHE"],
+        start_global_draw_id=int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "640000")),
+        global_microbatch_size=int(os.environ.get("FRAMESKIP_GLOBAL_MICROBATCH_SIZE", "4")),
+    )
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(overrides)))
+    cfg = apply_config_compat(cfg, strict=True)
+    cfg.config_yaml = args.config_yaml
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/frameskip_robotwin_adapted/validate_assets.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Validate transferred RoboTwin assets before a smoke or Slurm launch."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from train_frameskip_robotwin_continuation import load_cache
+
+
+def manifest_support(path: Path) -> set[tuple[str, str]]:
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    support = set()
+    for dataset in payload["datasets"]:
+        name = dataset.get("name", dataset.get("dataset_name"))
+        for episode in dataset.get("episodes", dataset.get("trajectories", [])):
+            trajectory_id = episode.get("trajectory_id", episode.get("id"))
+            support.add((name, str(trajectory_id)))
+    if not support:
+        raise RuntimeError("valid-anchor manifest has no trajectory support")
+    return support
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True, type=Path)
+    parser.add_argument("--manifest", required=True, type=Path)
+    parser.add_argument("--cache", required=True, type=Path)
+    args = parser.parse_args()
+
+    commit_path = args.checkpoint / "checkpoint.json"
+    commit = json.loads(commit_path.read_text(encoding="utf-8"))
+    if commit.get("schema_version") != 1 or int(commit.get("completed_steps", -1)) != 10000:
+        raise RuntimeError("checkpoint commit is not a complete step-10000 state")
+    if int(commit.get("world_size", -1)) != 2:
+        raise RuntimeError("10K checkpoint world size must be 2")
+    inventory = commit.get("files", [])
+    missing = [
+        entry["name"]
+        for entry in inventory
+        if not (args.checkpoint / entry["name"]).is_file()
+        or (args.checkpoint / entry["name"]).stat().st_size != int(entry["bytes"])
+    ]
+    if not inventory or missing:
+        raise RuntimeError(f"incomplete 10K checkpoint inventory: missing={missing}")
+    parent_config = args.checkpoint.parent.parent / "config.full.yaml"
+    if not parent_config.is_file():
+        raise RuntimeError(f"missing parent config required by exact resume: {parent_config}")
+
+    cache, payload = load_cache(str(args.cache))
+    support = manifest_support(args.manifest)
+    cache_support = set(cache)
+    if support != cache_support:
+        raise RuntimeError(
+            "FrameSkip cache support differs from manifest support: "
+            f"missing={len(support - cache_support)}, extra={len(cache_support - support)}"
+        )
+
+    print("FRAMESKIP_ROBOTWIN_ASSET_PREFLIGHT=PASS")
+    print(f"checkpoint_files={len(inventory)}")
+    print(f"trajectory_support={len(support)}")
+    print(f"cache_records={len(payload['records'])}")
+
+
+if __name__ == "__main__":
+    main()

external_baselines/keyframe_focused_adapted/build_ape_cache.py

@@ -0,0 +1,485 @@
+#!/usr/bin/env python3
+"""Build APE weights on the frozen Florence-Bridge valid-anchor support."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from collections import OrderedDict, defaultdict
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as parquet
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest
+
+from external_baselines.keyframe_focused_adapted.keyframe_cache import (
+    ActionCopycat,
+    history_target,
+    q99_normalize,
+)
+
+
+def episode_path(data_root: Path, dataset_name: str, trajectory_id: int) -> Path:
+    return (
+        data_root
+        / dataset_name
+        / "data"
+        / f"chunk-{trajectory_id // 1000:03d}"
+        / f"episode_{trajectory_id:06d}.parquet"
+    )
+
+
+def read_actions(
+    data_root: Path,
+    dataset_name: str,
+    trajectory_id: int,
+    action_dim: int | None,
+) -> np.ndarray:
+    path = episode_path(data_root, dataset_name, trajectory_id)
+    values = parquet.read_table(path, columns=["action"])["action"].to_pylist()
+    actions = np.asarray(values, dtype=np.float32)
+    if actions.ndim != 2 or (action_dim is not None and actions.shape[1] != action_dim):
+        raise RuntimeError(f"unexpected action shape {actions.shape} in {path}")
+    return actions
+
+
+def resolve_statistics_key(
+    dataset_name: str,
+    statistics_payload: dict,
+    explicit_mapping: dict[str, str],
+) -> str:
+    """Resolve Bridge aliases and arbitrary RoboTwin dataset statistics."""
+
+    aliases = {
+        "bridge_orig_lerobot": "oxe_bridge",
+        "fractal20220817_data_lerobot": "oxe_rt1",
+    }
+    candidates = (
+        explicit_mapping.get(dataset_name),
+        aliases.get(dataset_name),
+        dataset_name,
+    )
+    for candidate in candidates:
+        if candidate and candidate in statistics_payload:
+            return candidate
+    if len(statistics_payload) == 1:
+        return next(iter(statistics_payload))
+    raise KeyError(
+        f"no statistics entry for {dataset_name!r}; provide --statistics-key-map"
+    )
+
+
+def requested_support(manifest, args):
+    if args.scope == "full":
+        return {
+            dataset.name: OrderedDict(
+                (int(episode.trajectory_id), tuple(int(x) for x in episode.anchors))
+                for episode in dataset.episodes
+            )
+            for dataset in manifest.datasets
+        }
+    selected = defaultdict(lambda: defaultdict(set))
+    for draw_id in range(args.draw_start, args.draw_start + args.draw_count):
+        trace = manifest.sample(args.draw_seed, LogicalDraw(0, draw_id))
+        selected[trace.dataset_name][int(trace.trajectory_id)].add(int(trace.base_index))
+    return {
+        dataset.name: OrderedDict(
+            (
+                int(episode.trajectory_id),
+                tuple(sorted(selected[dataset.name][int(episode.trajectory_id)])),
+            )
+            for episode in dataset.episodes
+            if selected[dataset.name].get(int(episode.trajectory_id))
+        )
+        for dataset in manifest.datasets
+    }
+
+
+def fitting_support(manifest, max_anchors: int | None):
+    result = {}
+    for dataset in manifest.datasets:
+        episodes = OrderedDict()
+        count = 0
+        for episode in dataset.episodes:
+            anchors = tuple(int(x) for x in episode.anchors if int(x) >= 9)
+            if max_anchors is not None:
+                anchors = anchors[: max(0, max_anchors - count)]
+            if anchors:
+                episodes[int(episode.trajectory_id)] = anchors
+                count += len(anchors)
+            if max_anchors is not None and count >= max_anchors:
+                break
+        if count == 0:
+            raise RuntimeError(f"no copycat fitting anchors for {dataset.name}")
+        result[dataset.name] = episodes
+    return result
+
+
+def examples_for_support(
+    data_root,
+    dataset_name,
+    support,
+    statistics,
+    *,
+    action_dim,
+    binary_action_indices,
+    history=9,
+):
+    inputs, targets = [], []
+    for trajectory_id, anchors in support.items():
+        actions = q99_normalize(
+            read_actions(data_root, dataset_name, trajectory_id, action_dim),
+            statistics,
+            binary_action_indices=binary_action_indices,
+        )
+        for anchor in anchors:
+            x, y = history_target(actions, anchor, history=history)
+            inputs.append(x)
+            targets.append(y)
+    return np.asarray(inputs, dtype=np.float32), np.asarray(targets, dtype=np.float32)
+
+
+def train_copycat(inputs, targets, *, action_dim, epochs, device, seed=0):
+    generator = torch.Generator().manual_seed(seed)
+    permutation = torch.randperm(len(inputs), generator=generator)
+    split = int(len(inputs) * 0.8)
+    train_indices = permutation[:split]
+    val_indices = permutation[split:]
+    x = torch.from_numpy(inputs)
+    y = torch.from_numpy(targets)
+    train = TensorDataset(x[train_indices], y[train_indices])
+    val = TensorDataset(x[val_indices], y[val_indices])
+    train_loader = DataLoader(
+        train, batch_size=512, shuffle=True, generator=generator, num_workers=0
+    )
+    val_loader = DataLoader(val, batch_size=512, shuffle=False, num_workers=0)
+    torch.manual_seed(seed)
+    model = ActionCopycat(action_dim=action_dim).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999))
+    history = []
+    for epoch in range(1, epochs + 1):
+        model.train()
+        losses = []
+        for batch_x, batch_y in train_loader:
+            batch_x = batch_x.to(device)
+            batch_y = batch_y.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            loss = (model(batch_x) - batch_y).square().mean(dim=-1).mean()
+            loss.backward()
+            optimizer.step()
+            losses.append(float(loss.detach().cpu()))
+        model.eval()
+        val_losses = []
+        with torch.no_grad():
+            for batch_x, batch_y in val_loader:
+                batch_x = batch_x.to(device)
+                batch_y = batch_y.to(device)
+                loss = (model(batch_x) - batch_y).square().mean(dim=-1)
+                val_losses.extend(loss.cpu().tolist())
+        history.append(
+            {
+                "epoch": epoch,
+                "train_loss": float(np.mean(losses)),
+                "validation_loss": float(np.mean(val_losses)),
+            }
+        )
+        print(json.dumps(history[-1], sort_keys=True), flush=True)
+        if epoch % max(1, epochs // 3) == 0:
+            for group in optimizer.param_groups:
+                group["lr"] *= 0.1
+    return model, history
+
+
+def prediction_errors(model, inputs, targets, device):
+    errors = []
+    loader = DataLoader(
+        TensorDataset(torch.from_numpy(inputs), torch.from_numpy(targets)),
+        batch_size=4096,
+        shuffle=False,
+        num_workers=0,
+    )
+    model.eval()
+    with torch.no_grad():
+        for batch_x, batch_y in loader:
+            batch_x = batch_x.to(device)
+            batch_y = batch_y.to(device)
+            errors.append(
+                (model(batch_x) - batch_y).square().mean(dim=-1).cpu().numpy()
+            )
+    return np.concatenate(errors)
+
+
+def score_support(
+    model,
+    data_root,
+    dataset_name,
+    support,
+    statistics,
+    device,
+    *,
+    action_dim,
+    binary_action_indices,
+):
+    trajectory_ids, offsets, anchors_out, ape_out = [], [0], [], []
+    model.eval()
+    with torch.no_grad():
+        for trajectory_id, anchors in support.items():
+            actions = q99_normalize(
+                read_actions(data_root, dataset_name, trajectory_id, action_dim),
+                statistics,
+                binary_action_indices=binary_action_indices,
+            )
+            batch_x, batch_y = [], []
+            for anchor in anchors:
+                x, y = history_target(actions, anchor, history=9)
+                batch_x.append(x)
+                batch_y.append(y)
+            x = torch.from_numpy(np.asarray(batch_x, dtype=np.float32)).to(device)
+            y = torch.from_numpy(np.asarray(batch_y, dtype=np.float32)).to(device)
+            ape = (model(x) - y).square().mean(dim=-1).cpu().numpy()
+            trajectory_ids.append(trajectory_id)
+            anchors_out.extend(anchors)
+            ape_out.extend(ape.tolist())
+            offsets.append(len(anchors_out))
+    return {
+        "trajectory_ids": np.asarray(trajectory_ids, dtype=np.int64),
+        "offsets": np.asarray(offsets, dtype=np.int64),
+        "anchors": np.asarray(anchors_out, dtype=np.int32),
+        "ape": np.asarray(ape_out, dtype=np.float32),
+    }
+
+
+def verify_reference_overlap(reference, dataset_index, scored, weights):
+    prefix = f"d{dataset_index}_"
+    reference_trajectory_ids = reference[prefix + "trajectory_ids"]
+    reference_offsets = reference[prefix + "offsets"]
+    reference_anchors = reference[prefix + "anchors"]
+    reference_ape = reference[prefix + "ape"]
+    reference_weights = reference[prefix + "weights"]
+    new_trajectory_positions = {
+        int(trajectory_id): position
+        for position, trajectory_id in enumerate(scored["trajectory_ids"])
+    }
+    checked = 0
+    max_ape_difference = 0.0
+    for reference_position, trajectory_id in enumerate(reference_trajectory_ids):
+        new_position = new_trajectory_positions.get(int(trajectory_id))
+        if new_position is None:
+            raise RuntimeError(
+                f"full cache is missing reference trajectory {int(trajectory_id)}"
+            )
+        reference_slice = slice(
+            int(reference_offsets[reference_position]),
+            int(reference_offsets[reference_position + 1]),
+        )
+        new_slice = slice(
+            int(scored["offsets"][new_position]),
+            int(scored["offsets"][new_position + 1]),
+        )
+        new_anchors = scored["anchors"][new_slice]
+        positions = np.searchsorted(new_anchors, reference_anchors[reference_slice])
+        if np.any(positions >= len(new_anchors)) or not np.array_equal(
+            new_anchors[positions], reference_anchors[reference_slice]
+        ):
+            raise RuntimeError(
+                f"full cache does not contain the reference anchor support for "
+                f"trajectory {int(trajectory_id)}"
+            )
+        new_ape = scored["ape"][new_slice][positions]
+        ape_difference = np.abs(new_ape - reference_ape[reference_slice])
+        if len(ape_difference):
+            max_ape_difference = max(max_ape_difference, float(ape_difference.max()))
+        if not np.array_equal(
+            weights[new_slice][positions], reference_weights[reference_slice]
+        ):
+            raise RuntimeError(
+                f"full cache changed reference weights for trajectory "
+                f"{int(trajectory_id)}"
+            )
+        checked += len(positions)
+    if max_ape_difference > 1e-6:
+        raise RuntimeError(
+            f"full cache changed reference APE scores: max_abs_diff={max_ape_difference}"
+        )
+    return {
+        "checked_anchor_count": checked,
+        "weights_exact_match": True,
+        "max_ape_abs_difference": max_ape_difference,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", type=Path, required=True)
+    parser.add_argument("--data-root", type=Path, required=True)
+    parser.add_argument("--statistics", type=Path, required=True)
+    parser.add_argument(
+        "--statistics-key-map",
+        type=Path,
+        help="Optional JSON mapping from manifest dataset name to statistics key.",
+    )
+    parser.add_argument(
+        "--action-dim",
+        type=int,
+        choices=(7, 14),
+        help="Expected action dimension; inferred from statistics when omitted.",
+    )
+    parser.add_argument(
+        "--binary-action-indices",
+        help=(
+            "Comma-separated zero-based binary action indices. Defaults to index 6 "
+            "for 7D Bridge and none for 14D RoboTwin absolute qpos."
+        ),
+    )
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--scope", choices=("full", "draw-range"), default="full")
+    parser.add_argument("--draw-seed", type=int, default=117465)
+    parser.add_argument("--draw-start", type=int, default=1_280_000)
+    parser.add_argument("--draw-count", type=int, default=128)
+    parser.add_argument("--fit-max-anchors", type=int)
+    parser.add_argument("--epochs", type=int, default=300)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--reference-cache", type=Path)
+    args = parser.parse_args()
+
+    manifest = ValidAnchorManifest.load(args.manifest)
+    statistics_payload = json.loads(args.statistics.read_text(encoding="utf-8"))
+    statistics_key_map = (
+        json.loads(args.statistics_key_map.read_text(encoding="utf-8"))
+        if args.statistics_key_map is not None
+        else {}
+    )
+    reference = None
+    reference_metadata = None
+    if args.reference_cache is not None:
+        if args.scope != "full":
+            raise ValueError("--reference-cache requires --scope full")
+        reference = np.load(args.reference_cache, allow_pickle=False)
+        reference_metadata = json.loads(str(reference["metadata_json"]))
+        if int(reference_metadata["copycat_epochs"]) != args.epochs:
+            raise RuntimeError("reference cache copycat epoch count does not match")
+    scoring = requested_support(manifest, args)
+    fitting = fitting_support(manifest, args.fit_max_anchors)
+    arrays = {}
+    dataset_reports = []
+    for dataset_index, dataset in enumerate(manifest.datasets):
+        dataset_name = dataset.name
+        statistics_key = resolve_statistics_key(
+            dataset_name, statistics_payload, statistics_key_map
+        )
+        action_statistics = statistics_payload[statistics_key]["action"]
+        inferred_action_dim = len(action_statistics["q99"])
+        action_dim = args.action_dim or inferred_action_dim
+        if inferred_action_dim != action_dim:
+            raise RuntimeError(
+                f"statistics action dimension {inferred_action_dim} != {action_dim} "
+                f"for {dataset_name}"
+            )
+        if args.binary_action_indices is None:
+            binary_action_indices = (6,) if action_dim == 7 else ()
+        else:
+            binary_action_indices = tuple(
+                int(value)
+                for value in args.binary_action_indices.split(",")
+                if value.strip()
+            )
+        inputs, targets = examples_for_support(
+            args.data_root,
+            dataset_name,
+            fitting[dataset_name],
+            action_statistics,
+            action_dim=action_dim,
+            binary_action_indices=binary_action_indices,
+        )
+        print(
+            f"training copycat for {dataset_name}: {len(inputs)} anchors",
+            flush=True,
+        )
+        model, train_history = train_copycat(
+            inputs,
+            targets,
+            action_dim=action_dim,
+            epochs=args.epochs,
+            device=torch.device(args.device),
+        )
+        fitted_ape = prediction_errors(
+            model, inputs, targets, torch.device(args.device)
+        )
+        scored = score_support(
+            model,
+            args.data_root,
+            dataset_name,
+            scoring[dataset_name],
+            action_statistics,
+            torch.device(args.device),
+            action_dim=action_dim,
+            binary_action_indices=binary_action_indices,
+        )
+        if reference_metadata is None:
+            threshold_values = scored["ape"] if args.scope == "full" else fitted_ape
+            threshold = float(np.quantile(threshold_values, 0.9))
+        else:
+            reference_report = reference_metadata["datasets"][dataset_index]
+            if reference_report["name"] != dataset_name:
+                raise RuntimeError("reference cache dataset order does not match")
+            if int(reference_report["fit_anchor_count"]) != len(inputs):
+                raise RuntimeError("reference cache fitting support does not match")
+            threshold = float(reference_report["ape_threshold_90th_percentile"])
+        weights = np.where(scored["ape"] >= threshold, 5.0, 1.0).astype(np.float32)
+        overlap_report = None
+        if reference is not None:
+            overlap_report = verify_reference_overlap(
+                reference, dataset_index, scored, weights
+            )
+        for key, value in scored.items():
+            arrays[f"d{dataset_index}_{key}"] = value
+        arrays[f"d{dataset_index}_weights"] = weights
+        dataset_reports.append(
+            {
+                "name": dataset_name,
+                "statistics_key": statistics_key,
+                "action_dim": action_dim,
+                "binary_action_indices": list(binary_action_indices),
+                "fit_anchor_count": len(inputs),
+                "scored_anchor_count": len(scored["anchors"]),
+                "ape_threshold_90th_percentile": threshold,
+                "upweighted_count": int((weights > 1).sum()),
+                "copycat_training": train_history,
+                "reference_overlap": overlap_report,
+            }
+        )
+
+    metadata = {
+        "schema_version": 1,
+        "method": "Keyframe-Focused-adapted",
+        "history_length": 9,
+        "copycat_hidden_units": [300],
+        "copycat_optimizer": {"name": "Adam", "lr": 1e-3, "betas": [0.9, 0.999]},
+        "copycat_epochs": args.epochs,
+        "ape_metric": "mean squared error over normalized action dimensions",
+        "boundary_padding": "repeat earliest physical action",
+        "threshold_scope": "within dataset source",
+        "threshold_percentile": 90,
+        "high_weight": 5.0,
+        "low_weight": 1.0,
+        "cache_scope": args.scope,
+        "threshold_source": (
+            "reference_cache" if args.reference_cache is not None else "computed"
+        ),
+        "dataset_names": [dataset.name for dataset in manifest.datasets],
+        "datasets": dataset_reports,
+    }
+    arrays["metadata_json"] = np.asarray(json.dumps(metadata, sort_keys=True))
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
+    with temporary.open("wb") as destination:
+        np.savez_compressed(destination, **arrays)
+    temporary.replace(args.output)
+    print(f"wrote {args.output}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/keyframe_focused_adapted/configs/keyframe_focused_seed117465_20k_to_20500.yaml

@@ -0,0 +1,118 @@
+run_id: keyframe-focused-adapted-seed117465-20k-to20500
+run_root_dir: ${oc.env:OPTR_RUN_ROOT}
+seed: 117465
+wandb_entity: null
+wandb_project: OPTR-Bridge-baselines
+wandb_group: keyframe-focused
+wandb_job_type: train
+wandb_tags:
+  - baseline
+  - keyframe-focused
+  - bridge
+  - train
+is_debug: false
+version_id: "0.21"
+
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: ${oc.env:FLORENCE_MODEL_PATH}
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  dino:
+    dino_backbone: dinov2_vits14
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    future_action_window_size: 15
+    past_action_window_size: 0
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 2048
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: ${oc.env:BRIDGE_DATA_ROOT}
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: ${oc.env:BRIDGE_STATISTICS_PATH}
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size: [224, 224]
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: ${oc.env:BRIDGE_PROVENANCE_PATH}
+    dataset_content_fingerprint_path: ${oc.env:BRIDGE_FINGERPRINT_PATH}
+    valid_anchor_manifest_path: ${oc.env:BRIDGE_MANIFEST_PATH}
+    temporal_proposal_path: null
+    transform_rng_seed: 298226
+
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps: [20500]
+  evaluation_snapshot_interval: null
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 1.0e-04
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ""
+  trainable_modules: ""
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas: [0.9, 0.95]
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: ${oc.env:COMMON_START_20K}
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  sample_loss_weight_key: keyframe_weight
+  stop_after_steps: 20500
+
+output_dir: ${oc.env:OPTR_RUN_ROOT}/keyframe-focused-adapted-seed117465-20k-to20500

external_baselines/keyframe_focused_adapted/keyframe_cache.py

@@ -0,0 +1,145 @@
+"""APE model and compact cache used by the isolated Keyframe baseline."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+
+
+class ActionCopycat(nn.Module):
+    """Official one-hidden-layer action-only copycat architecture."""
+
+    def __init__(self, action_dim: int = 7, history: int = 9, hidden: int = 300):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(action_dim * history, hidden),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden, action_dim),
+        )
+
+    def forward(self, action_history: torch.Tensor) -> torch.Tensor:
+        return self.net(action_history)
+
+
+def q99_normalize(
+    actions: np.ndarray,
+    statistics: dict,
+    *,
+    binary_action_indices: tuple[int, ...] = (),
+) -> np.ndarray:
+    """Apply q01/q99 normalization with an explicit binary-action contract.
+
+    A zero percentile span is not sufficient evidence that an action dimension is
+    binary.  This matters for RoboTwin absolute qpos, where a joint can be constant
+    in a dataset shard while remaining a continuous joint.  Such continuous,
+    zero-span dimensions are mapped to zero; only explicitly declared indices use
+    the binary threshold.
+    """
+
+    actions = np.asarray(actions, dtype=np.float32)
+    q01 = np.asarray(statistics["q01"], dtype=np.float32)
+    q99 = np.asarray(statistics["q99"], dtype=np.float32)
+    if actions.shape[-1] != len(q01) or q01.shape != q99.shape:
+        raise ValueError(
+            f"action/statistics shape mismatch: actions={actions.shape}, "
+            f"q01={q01.shape}, q99={q99.shape}"
+        )
+    binary = np.zeros(len(q01), dtype=bool)
+    for index in binary_action_indices:
+        if index < 0 or index >= len(q01):
+            raise ValueError(f"binary action index {index} is out of range")
+        binary[index] = True
+    continuous = ~binary
+    scaled = continuous & (q01 != q99)
+    constant = continuous & (q01 == q99)
+    normalized = np.zeros_like(actions, dtype=np.float32)
+    normalized[..., scaled] = (
+        2.0
+        * (actions[..., scaled] - q01[scaled])
+        / (q99[scaled] - q01[scaled])
+        - 1.0
+    )
+    normalized[..., scaled] = np.clip(
+        normalized[..., scaled], -2.2, 2.2
+    )
+    normalized[..., constant] = 0.0
+    normalized[..., binary] = (
+        actions[..., binary] > 0.5
+    ).astype(np.float32)
+    return normalized
+
+
+def history_target(actions: np.ndarray, anchor: int, history: int = 9):
+    """Repeat the earliest physical action at a trajectory boundary."""
+
+    indices = np.arange(anchor - history, anchor, dtype=np.int64)
+    indices = np.clip(indices, 0, len(actions) - 1)
+    return actions[indices].reshape(-1), actions[anchor]
+
+
+@dataclass(frozen=True)
+class CacheRecord:
+    ape: float
+    weight: float
+
+
+class KeyframeWeightCache:
+    """Read a compact NPZ cache without altering sampling."""
+
+    def __init__(self, path: str | Path):
+        self.path = Path(path)
+        payload = np.load(self.path, allow_pickle=False)
+        self.metadata = json.loads(str(payload["metadata_json"].item()))
+        if self.metadata.get("method") != "Keyframe-Focused-adapted":
+            raise RuntimeError("unexpected keyframe cache method")
+        self._datasets = {}
+        for index, name in enumerate(self.metadata["dataset_names"]):
+            trajectory_ids = payload[f"d{index}_trajectory_ids"].astype(np.int64)
+            offsets = payload[f"d{index}_offsets"].astype(np.int64)
+            anchors = payload[f"d{index}_anchors"].astype(np.int64)
+            ape = payload[f"d{index}_ape"].astype(np.float32)
+            weights = payload[f"d{index}_weights"].astype(np.float32)
+            if not (len(anchors) == len(ape) == len(weights)):
+                raise RuntimeError(f"cache array lengths differ for {name}")
+            if len(offsets) != len(trajectory_ids) + 1 or int(offsets[-1]) != len(anchors):
+                raise RuntimeError(f"cache offsets are invalid for {name}")
+            if np.any(~np.isfinite(ape)) or np.any(~np.isfinite(weights)):
+                raise RuntimeError(f"cache contains non-finite values for {name}")
+            if np.any(weights <= 0):
+                raise RuntimeError(f"cache contains non-positive weights for {name}")
+            self._datasets[name] = {
+                "trajectory_to_row": {
+                    int(trajectory_id): row
+                    for row, trajectory_id in enumerate(trajectory_ids)
+                },
+                "offsets": offsets,
+                "anchors": anchors,
+                "ape": ape,
+                "weights": weights,
+            }
+
+    def lookup(self, dataset_name: str, trajectory_id: int, base_index: int) -> CacheRecord:
+        try:
+            dataset = self._datasets[dataset_name]
+            row = dataset["trajectory_to_row"][int(trajectory_id)]
+        except KeyError as exc:
+            raise KeyError(
+                f"keyframe cache has no trajectory {(dataset_name, trajectory_id)}"
+            ) from exc
+        start, end = dataset["offsets"][row : row + 2]
+        local = dataset["anchors"][start:end]
+        position = int(np.searchsorted(local, int(base_index)))
+        if position >= len(local) or int(local[position]) != int(base_index):
+            raise KeyError(
+                f"keyframe cache has no anchor {(dataset_name, trajectory_id, base_index)}"
+            )
+        absolute = int(start + position)
+        return CacheRecord(
+            ape=float(dataset["ape"][absolute]),
+            weight=float(dataset["weights"][absolute]),
+        )

external_baselines/keyframe_focused_adapted/train_keyframe_continuation.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Attach frozen APE loss weights to the unchanged Florence-Bridge draw stream."""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from omegaconf import OmegaConf
+
+from external_baselines.keyframe_focused_adapted.keyframe_cache import (
+    KeyframeWeightCache,
+)
+
+
+def install_adapter(cache_path: str) -> KeyframeWeightCache:
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+    from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
+
+    cache = KeyframeWeightCache(cache_path)
+    if os.environ.get("KEYFRAME_SMOKE_FROZEN_MANIFEST") == "1":
+        original_build_manifest = LeRobotMixtureDataset._build_valid_anchor_manifest
+
+        def load_frozen_manifest(self):
+            path = self.data_cfg.get("valid_anchor_manifest_path")
+            if not path:
+                return original_build_manifest(self)
+            manifest = ValidAnchorManifest.load(path)
+            print(
+                "Smoke-only frozen manifest fast path: "
+                f"{path}",
+                flush=True,
+            )
+            return manifest
+
+        LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest
+
+    original_getitem = LeRobotMixtureDataset.__getitem__
+
+    def keyframe_getitem(self, index):
+        sample = original_getitem(self, index)
+        trace = sample["sampling_trace"]
+        record = cache.lookup(
+            trace["dataset_name"],
+            int(trace["trajectory_id"]),
+            int(trace["base_index"]),
+        )
+        # The logical draw and physical timestep are unchanged; only loss weight is added.
+        trace["keyframe_ape"] = record.ape
+        trace["keyframe_weight"] = record.weight
+        trace["keyframe_cache"] = cache.path.name
+        return sample
+
+    LeRobotMixtureDataset.__getitem__ = keyframe_getitem
+    return cache
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, clipargs = parser.parse_known_args()
+
+    cache = install_adapter(os.environ["KEYFRAME_APE_CACHE"])
+    print(
+        f"Installed Keyframe-Focused-adapted cache {cache.path}",
+        flush=True,
+    )
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
+    cfg = apply_config_compat(cfg, strict=True)
+    if cfg.datasets.vla_data.temporal_proposal_path is not None:
+        raise ValueError("Keyframe loss weighting requires the uniform temporal sampler")
+    if cfg.trainer.sample_loss_weight_key != "keyframe_weight":
+        raise ValueError("trainer.sample_loss_weight_key must be keyframe_weight")
+    cfg.config_yaml = args.config_yaml
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/keyframe_focused_adapted/validate_cache_support.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Verify that a cache covers an unchanged range of uniform logical draws."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest
+
+from external_baselines.keyframe_focused_adapted.keyframe_cache import (
+    KeyframeWeightCache,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", type=Path, required=True)
+    parser.add_argument("--cache", type=Path, required=True)
+    parser.add_argument("--seed", type=int)
+    parser.add_argument("--draw-start", type=int)
+    parser.add_argument("--draw-count", type=int)
+    parser.add_argument("--full-support", action="store_true")
+    args = parser.parse_args()
+
+    manifest = ValidAnchorManifest.load(args.manifest)
+    cache = KeyframeWeightCache(args.cache)
+    if args.full_support:
+        arrays = np.load(args.cache, allow_pickle=False)
+        metadata = json.loads(str(arrays["metadata_json"]))
+        counts: dict[str, int] = {}
+        weights: dict[str, int] = {}
+        for dataset_index, dataset in enumerate(manifest.datasets):
+            prefix = f"d{dataset_index}_"
+            trajectory_ids = arrays[prefix + "trajectory_ids"]
+            offsets = arrays[prefix + "offsets"]
+            anchors = arrays[prefix + "anchors"]
+            cached_weights = arrays[prefix + "weights"]
+            expected_trajectory_ids = np.asarray(
+                [int(episode.trajectory_id) for episode in dataset.episodes],
+                dtype=np.int64,
+            )
+            if not np.array_equal(trajectory_ids, expected_trajectory_ids):
+                raise RuntimeError(
+                    f"cache trajectory support does not match {dataset.name}"
+                )
+            for position, episode in enumerate(dataset.episodes):
+                cached = anchors[int(offsets[position]) : int(offsets[position + 1])]
+                expected = np.asarray(episode.anchors, dtype=np.int32)
+                if not np.array_equal(cached, expected):
+                    raise RuntimeError(
+                        f"cache anchor support does not match "
+                        f"{dataset.name}/{int(episode.trajectory_id)}"
+                    )
+            counts[dataset.name] = int(len(anchors))
+            values, value_counts = np.unique(cached_weights, return_counts=True)
+            for value, count in zip(values, value_counts, strict=True):
+                label = str(float(value))
+                weights[label] = weights.get(label, 0) + int(count)
+        print(
+            json.dumps(
+                {
+                    "status": "PASS",
+                    "mode": "full-support",
+                    "cache": cache.path.name,
+                    "dataset_anchor_counts": counts,
+                    "weight_counts": weights,
+                },
+                sort_keys=True,
+            )
+        )
+        return
+    if args.seed is None or args.draw_start is None or args.draw_count is None:
+        parser.error("draw-range validation requires --seed, --draw-start, and --draw-count")
+    counts: dict[str, int] = {}
+    weights: dict[str, int] = {}
+    for draw_id in range(args.draw_start, args.draw_start + args.draw_count):
+        trace = manifest.sample(args.seed, LogicalDraw(0, draw_id))
+        record = cache.lookup(
+            trace.dataset_name, trace.trajectory_id, trace.base_index
+        )
+        counts[trace.dataset_name] = counts.get(trace.dataset_name, 0) + 1
+        label = str(record.weight)
+        weights[label] = weights.get(label, 0) + 1
+
+    print(
+        json.dumps(
+            {
+                "status": "PASS",
+                "cache": cache.path.name,
+                "draw_start": args.draw_start,
+                "draw_count": args.draw_count,
+                "dataset_counts": counts,
+                "weight_counts": weights,
+            },
+            sort_keys=True,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

external_baselines/loss_is_active/train_active_loss_is_continuation.py

@@ -0,0 +1,330 @@
+#!/usr/bin/env python3
+"""Run online example-level Loss-IS without changing temporal proposals."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import json
+import math
+import os
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from omegaconf import OmegaConf
+
+from external_baselines.loss_is_adapted.loss_is import sample_with_replacement
+
+
+def _stable_seed(*parts: object) -> int:
+    payload = ":".join(str(part) for part in parts).encode("utf-8")
+    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little") % (2**63 - 1)
+
+
+class LossISTrainerMixin:
+    """Always-on distributed adaptation of the official B-to-b sampler."""
+
+    def __init__(self, *args, **kwargs):
+        dataloader = kwargs.get("vla_train_dataloader")
+        if dataloader is None and len(args) >= 3:
+            dataloader = args[2]
+        self._loss_is_dataset = getattr(dataloader, "dataset", None)
+        super().__init__(*args, **kwargs)
+
+        cfg = self.config.trainer.loss_is
+        self._loss_is_factor = int(cfg.candidate_factor)
+        self._loss_is_seed = int(cfg.seed)
+        if self._loss_is_factor != 3:
+            raise ValueError("The locked ICML'18-style baseline uses candidate_factor=3")
+        if self._loss_is_dataset is None or not callable(
+            getattr(self._loss_is_dataset, "__getitem__", None)
+        ):
+            raise TypeError("Loss-IS requires random access to the baseline dataset")
+        self._loss_is_active = True
+        self._loss_is_selection_metrics: list[dict[str, float]] = []
+        self._loss_is_last_metrics: dict[str, float] = {}
+
+    def prepare_training(self) -> None:
+        super().prepare_training()
+        if not bool(getattr(self.config.trainer, "is_resume", False)):
+            return
+
+        # An external common-start checkpoint is the initial fork into Loss-IS.
+        if bool(getattr(self, "_branched_from_external_checkpoint", False)):
+            return
+
+        state_path = Path(self.config.output_dir) / "loss_is_state.json"
+        initial_branch_step = int(
+            getattr(self.config.trainer.loss_is, "initial_branch_step", -1)
+        )
+        if not state_path.is_file() and int(self.completed_steps) == initial_branch_step:
+            print(
+                "Starting always-on Loss-IS from the exact baseline sidecar at "
+                f"step {self.completed_steps}",
+                flush=True,
+            )
+            return
+        if not state_path.is_file():
+            raise FileNotFoundError(f"Loss-IS resume state is missing: {state_path}")
+        payload = json.loads(state_path.read_text(encoding="utf-8"))
+        if int(payload["completed_optimizer_step"]) != int(self.completed_steps):
+            raise RuntimeError("Loss-IS state step does not match the resumed checkpoint")
+        if int(payload["candidate_factor"]) != self._loss_is_factor:
+            raise RuntimeError("Loss-IS candidate factor differs from the resumed state")
+        if payload.get("activation_mode") != "always_on":
+            raise RuntimeError("Loss-IS resume state is not always-on")
+        print(
+            f"Restored always-on Loss-IS state at step {self.completed_steps}",
+            flush=True,
+        )
+
+    def _begin_loss_is_update(self) -> None:
+        if self._microbatch_index != 0:
+            return
+        self._loss_is_active = True
+        self._loss_is_selection_metrics = []
+
+    def _extra_candidate_draw_id(self, base_ids: list[int], slot: int) -> int:
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        low = _stable_seed(
+            self._loss_is_seed,
+            "candidate",
+            self.completed_steps,
+            self._microbatch_index,
+            rank,
+            min(base_ids),
+            slot,
+        ) & ((1 << 60) - 1)
+        return (1 << 61) | low
+
+    def _build_candidate_pool(self, base_batch: list[dict]) -> list[dict]:
+        from starVLA.dataloader.infoflow_sampler import LogicalDraw
+
+        base_ids = [int(item["sampling_trace"]["global_draw_id"]) for item in base_batch]
+        epochs = {int(item["sampling_trace"]["epoch"]) for item in base_batch}
+        if len(epochs) != 1:
+            raise RuntimeError("A Loss-IS microbatch crossed logical epochs")
+        epoch = epochs.pop()
+        candidates = list(base_batch)
+        extra_count = (self._loss_is_factor - 1) * len(base_batch)
+        for slot in range(extra_count):
+            draw_id = self._extra_candidate_draw_id(base_ids, slot)
+            candidates.append(self._loss_is_dataset[LogicalDraw(epoch, draw_id)])
+        candidate_ids = [
+            int(item["sampling_trace"]["global_draw_id"]) for item in candidates
+        ]
+        if len(candidate_ids) != len(set(candidate_ids)):
+            raise RuntimeError("Candidate draw identities must be unique before resampling")
+        return candidates
+
+    @torch.no_grad()
+    def _score_candidates(self, candidates: list[dict], chunk_size: int) -> torch.Tensor:
+        was_training = bool(self.model.training)
+        self.model.eval()
+        scores = []
+        repeat_count = int(self.config.trainer.repeated_diffusion_steps)
+        try:
+            for chunk_index, start in enumerate(range(0, len(candidates), chunk_size)):
+                chunk = candidates[start : start + chunk_size]
+                draw_ids = [
+                    int(item["sampling_trace"]["global_draw_id"]) for item in chunk
+                ]
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    with self._paired_model_rng(
+                        draw_ids,
+                        1_000_000 + self._microbatch_index * self._loss_is_factor + chunk_index,
+                    ) as rng_seeds:
+                        output = self.model.forward(
+                            chunk,
+                            return_per_sample_loss=True,
+                            flow_noise_seed=(
+                                None if rng_seeds is None else rng_seeds["flow_noise"]
+                            ),
+                            return_flow_noise_time=True,
+                        )
+                chunk_scores = output["action_loss"].detach().to(
+                    device="cpu", dtype=torch.float64
+                )
+                if chunk_scores.shape != (len(chunk),):
+                    raise RuntimeError("Candidate scoring did not return one loss per example")
+                noise = output["flow_noise"].detach().cpu()
+                time = output["flow_time"].detach().cpu()
+                expected_noise_shape = (
+                    repeat_count * len(chunk),
+                    *tuple(noise.shape[1:]),
+                )
+                if tuple(noise.shape) != expected_noise_shape:
+                    raise RuntimeError("Candidate scoring returned an invalid flow-noise shape")
+                if tuple(time.shape) != (repeat_count * len(chunk),):
+                    raise RuntimeError("Candidate scoring returned an invalid flow-time shape")
+                noise = noise.reshape(repeat_count, len(chunk), *noise.shape[1:])
+                time = time.reshape(repeat_count, len(chunk))
+                for local_index, item in enumerate(chunk):
+                    # Reuse the scored Flow-Matching realization during backward.
+                    item["loss_is_flow_noise"] = noise[:, local_index].clone()
+                    item["loss_is_flow_time"] = time[:, local_index].clone()
+                scores.append(chunk_scores)
+        finally:
+            if was_training:
+                self.model.train()
+            self._enforce_frozen_module_eval()
+        result = torch.cat(scores)
+        if not torch.isfinite(result).all() or torch.any(result <= 0):
+            raise FloatingPointError("Candidate losses must be finite and strictly positive")
+        return result
+
+    def _selection_generator(self, base_ids: list[int]) -> torch.Generator:
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(
+            _stable_seed(
+                self._loss_is_seed,
+                "selection",
+                self.completed_steps,
+                self._microbatch_index,
+                rank,
+                min(base_ids),
+            )
+        )
+        return generator
+
+    def _get_next_batch(self):
+        base_batch = super()._get_next_batch()
+        self._begin_loss_is_update()
+        base_ids = [int(item["sampling_trace"]["global_draw_id"]) for item in base_batch]
+
+        candidates = self._build_candidate_pool(base_batch)
+        scores = self._score_candidates(candidates, chunk_size=len(base_batch))
+        selection = sample_with_replacement(
+            scores,
+            len(base_batch),
+            generator=self._selection_generator(base_ids),
+        )
+        selected = []
+        for occurrence_id, selected_index, weight, probability in zip(
+            base_ids,
+            selection.indices.tolist(),
+            selection.importance_weights.tolist(),
+            selection.selected_probabilities.tolist(),
+        ):
+            item = copy.deepcopy(candidates[selected_index])
+            trace = item["sampling_trace"]
+            trace["loss_is_candidate_global_draw_id"] = int(trace["global_draw_id"])
+            trace["loss_is_occurrence_id"] = int(occurrence_id)
+            trace["loss_is_importance_weight"] = float(weight)
+            trace["loss_is_probability"] = float(probability)
+            trace["loss_is_candidate_score"] = float(scores[selected_index])
+            trace["loss_is_active"] = True
+            selected.append(item)
+        weights = selection.importance_weights
+        self._loss_is_selection_metrics.append(
+            {
+                "candidate_loss_mean": float(scores.mean()),
+                "candidate_loss_max": float(scores.max()),
+                "weight_mean": float(weights.mean()),
+                "weight_max": float(weights.max()),
+                "ess": float(weights.sum().square() / weights.square().sum()),
+            }
+        )
+        return selected
+
+    def _after_per_sample_training_loss(
+        self,
+        batch_vla,
+        per_sample_loss,
+        *,
+        did_optimizer_step: bool,
+        microbatch_index: int,
+    ) -> None:
+        if per_sample_loss is None:
+            raise RuntimeError("Loss-IS requires per-sample training loss")
+        if not did_optimizer_step:
+            return
+
+        metrics = {
+            "loss_is/active": 1.0,
+            "loss_is/candidate_factor": float(self._loss_is_factor),
+        }
+        if self._loss_is_selection_metrics:
+            for key in self._loss_is_selection_metrics[0]:
+                metrics[f"loss_is/{key}"] = sum(
+                    item[key] for item in self._loss_is_selection_metrics
+                ) / len(self._loss_is_selection_metrics)
+        self._loss_is_last_metrics = metrics
+
+        if self.accelerator.is_main_process:
+            state_path = Path(self.config.output_dir) / "loss_is_state.json"
+            temporary = state_path.with_suffix(".json.tmp")
+            payload = {
+                "schema_version": 1,
+                "completed_optimizer_step": int(self.completed_steps + 1),
+                "activation_mode": "always_on",
+                "active_for_completed_step": True,
+                "candidate_factor": self._loss_is_factor,
+                "selection_scope": "per-rank per-microbatch; globally averaged unbiased estimator",
+                "metrics": metrics,
+            }
+            temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+            os.replace(temporary, state_path)
+
+    def _train_step(self, batch_vla, batch_vlm=None):
+        metrics = super()._train_step(batch_vla, batch_vlm=batch_vlm)
+        metrics.update(self._loss_is_last_metrics)
+        return metrics
+
+
+def install_smoke_manifest_fast_path() -> None:
+    if os.environ.get("LOSS_IS_SMOKE_FROZEN_MANIFEST") != "1":
+        return
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+    from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
+
+    def load_frozen_manifest(self):
+        path = self.data_cfg.get("valid_anchor_manifest_path")
+        if not path:
+            raise ValueError("Smoke fast path requires a frozen valid-anchor manifest")
+        manifest = ValidAnchorManifest.load(path)
+        print(
+            f"Loss-IS smoke frozen manifest: {path} sha256={manifest.sha256}",
+            flush=True,
+        )
+        return manifest
+
+    LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, clipargs = parser.parse_known_args()
+
+    install_smoke_manifest_fast_path()
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
+    cfg = apply_config_compat(cfg, strict=True)
+    if cfg.datasets.vla_data.temporal_proposal_path is not None:
+        raise ValueError("Loss-IS uses the baseline distribution, not a temporal proposal")
+    if cfg.trainer.sample_loss_weight_key != "loss_is_importance_weight":
+        raise ValueError("sample_loss_weight_key must be loss_is_importance_weight")
+    if cfg.trainer.model_rng_trace_key != "loss_is_occurrence_id":
+        raise ValueError("model_rng_trace_key must be loss_is_occurrence_id")
+    cfg.config_yaml = args.config_yaml
+
+    base_trainer = train_starvla.VLATrainer
+
+    class LossISTrainer(LossISTrainerMixin, base_trainer):
+        pass
+
+    train_starvla.VLATrainer = LossISTrainer
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/loss_is_adapted/configs/loss_is_seed117465_20k_to_20500.yaml

@@ -0,0 +1,122 @@
+run_id: loss-is-adapted-seed117465-20k-to20500
+run_root_dir: ${oc.env:OPTR_RUN_ROOT}
+seed: 117465
+wandb_entity: null
+wandb_project: OPTR-Bridge-baselines
+wandb_group: loss-is
+wandb_job_type: train
+wandb_tags:
+  - baseline
+  - loss-is
+  - bridge
+  - train
+is_debug: false
+version_id: "0.21"
+
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: ${oc.env:FLORENCE_MODEL_PATH}
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  dino:
+    dino_backbone: dinov2_vits14
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    future_action_window_size: 15
+    past_action_window_size: 0
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 2048
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: ${oc.env:BRIDGE_DATA_ROOT}
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: ${oc.env:BRIDGE_STATISTICS_PATH}
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size: [224, 224]
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: ${oc.env:BRIDGE_PROVENANCE_PATH}
+    dataset_content_fingerprint_path: ${oc.env:BRIDGE_FINGERPRINT_PATH}
+    valid_anchor_manifest_path: ${oc.env:BRIDGE_MANIFEST_PATH}
+    temporal_proposal_path: null
+    transform_rng_seed: 298226
+
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps: [20500]
+  evaluation_snapshot_interval: null
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 1.0e-04
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ""
+  trainable_modules: ""
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas: [0.9, 0.95]
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: ${oc.env:COMMON_START_20K}
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  model_rng_trace_key: loss_is_occurrence_id
+  sample_loss_weight_key: loss_is_importance_weight
+  stop_after_steps: 20500
+  loss_is:
+    candidate_factor: 3
+    seed: 117465
+
+output_dir: ${oc.env:OPTR_RUN_ROOT}/loss-is-adapted-seed117465-20k-to20500

external_baselines/loss_is_adapted/loss_is.py

@@ -0,0 +1,48 @@
+"""ICML'18-style online loss importance-sampling primitives."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class LossISSelection:
+    indices: torch.Tensor
+    probabilities: torch.Tensor
+    selected_probabilities: torch.Tensor
+    importance_weights: torch.Tensor
+
+
+def sample_with_replacement(
+    scores: torch.Tensor,
+    batch_size: int,
+    *,
+    generator: torch.Generator,
+) -> LossISSelection:
+    """Select ``b`` of ``B`` candidates and return exact ``1/(B g_i)`` weights."""
+
+    scores64 = torch.as_tensor(scores, dtype=torch.float64).flatten().cpu()
+    if scores64.numel() <= batch_size:
+        raise ValueError("candidate pool must be larger than the training batch")
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    if not torch.isfinite(scores64).all() or torch.any(scores64 <= 0):
+        raise ValueError("loss scores must be finite and strictly positive")
+    probability = scores64 / scores64.sum()
+    indices = torch.multinomial(
+        probability,
+        batch_size,
+        replacement=True,
+        generator=generator,
+    )
+    selected_probability = probability[indices]
+    weights = 1.0 / (scores64.numel() * selected_probability)
+    return LossISSelection(
+        indices=indices,
+        probabilities=probability,
+        selected_probabilities=selected_probability,
+        importance_weights=weights,
+    )
+

external_baselines/loss_is_adapted/train_loss_is_continuation.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+"""Run online example-level Loss-IS without changing temporal proposals."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from omegaconf import OmegaConf
+
+from external_baselines.loss_is_adapted.loss_is import sample_with_replacement
+
+
+def _stable_seed(*parts: object) -> int:
+    payload = ":".join(str(part) for part in parts).encode("utf-8")
+    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little") % (2**63 - 1)
+
+
+class LossISTrainerMixin:
+    """Always-on distributed adaptation of the official B-to-b sampler."""
+
+    def __init__(self, *args, **kwargs):
+        dataloader = kwargs.get("vla_train_dataloader")
+        if dataloader is None and len(args) >= 3:
+            dataloader = args[2]
+        self._loss_is_dataset = getattr(dataloader, "dataset", None)
+        super().__init__(*args, **kwargs)
+
+        cfg = self.config.trainer.loss_is
+        self._loss_is_factor = int(cfg.candidate_factor)
+        self._loss_is_seed = int(cfg.seed)
+        if self._loss_is_factor != 3:
+            raise ValueError("The locked ICML'18-style baseline uses candidate_factor=3")
+        if self._loss_is_dataset is None or not callable(
+            getattr(self._loss_is_dataset, "__getitem__", None)
+        ):
+            raise TypeError("Loss-IS requires random access to the baseline dataset")
+        self._loss_is_active = True
+        self._loss_is_selection_metrics: list[dict[str, float]] = []
+        self._loss_is_last_metrics: dict[str, float] = {}
+
+    def _begin_loss_is_update(self) -> None:
+        if self._microbatch_index != 0:
+            return
+        self._loss_is_active = True
+        self._loss_is_selection_metrics = []
+
+    def _extra_candidate_draw_id(self, base_ids: list[int], slot: int) -> int:
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        low = _stable_seed(
+            self._loss_is_seed,
+            "candidate",
+            self.completed_steps,
+            self._microbatch_index,
+            rank,
+            min(base_ids),
+            slot,
+        ) & ((1 << 60) - 1)
+        return (1 << 61) | low
+
+    def _build_candidate_pool(self, base_batch: list[dict]) -> list[dict]:
+        from starVLA.dataloader.infoflow_sampler import LogicalDraw
+
+        base_ids = [int(item["sampling_trace"]["global_draw_id"]) for item in base_batch]
+        epochs = {int(item["sampling_trace"]["epoch"]) for item in base_batch}
+        if len(epochs) != 1:
+            raise RuntimeError("A Loss-IS microbatch crossed logical epochs")
+        epoch = epochs.pop()
+        candidates = list(base_batch)
+        extra_count = (self._loss_is_factor - 1) * len(base_batch)
+        for slot in range(extra_count):
+            draw_id = self._extra_candidate_draw_id(base_ids, slot)
+            candidates.append(self._loss_is_dataset[LogicalDraw(epoch, draw_id)])
+        candidate_ids = [
+            int(item["sampling_trace"]["global_draw_id"]) for item in candidates
+        ]
+        if len(candidate_ids) != len(set(candidate_ids)):
+            raise RuntimeError("Candidate draw identities must be unique before resampling")
+        return candidates
+
+    @torch.no_grad()
+    def _score_candidates(self, candidates: list[dict], chunk_size: int) -> torch.Tensor:
+        was_training = bool(self.model.training)
+        self.model.eval()
+        scores = []
+        repeat_count = int(self.config.trainer.repeated_diffusion_steps)
+        try:
+            for chunk_index, start in enumerate(range(0, len(candidates), chunk_size)):
+                chunk = candidates[start : start + chunk_size]
+                draw_ids = [
+                    int(item["sampling_trace"]["global_draw_id"]) for item in chunk
+                ]
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    with self._paired_model_rng(
+                        draw_ids,
+                        1_000_000 + self._microbatch_index * self._loss_is_factor + chunk_index,
+                    ) as rng_seeds:
+                        output = self.model.forward(
+                            chunk,
+                            return_per_sample_loss=True,
+                            flow_noise_seed=rng_seeds["flow_noise"],
+                            return_flow_noise_time=True,
+                        )
+                chunk_scores = output["action_loss"].detach().to(
+                    device="cpu", dtype=torch.float64
+                )
+                if chunk_scores.shape != (len(chunk),):
+                    raise RuntimeError("Candidate scoring did not return one loss per example")
+                noise = output["flow_noise"].detach().cpu()
+                time = output["flow_time"].detach().cpu()
+                expected_noise_shape = (
+                    repeat_count * len(chunk),
+                    *tuple(noise.shape[1:]),
+                )
+                if tuple(noise.shape) != expected_noise_shape:
+                    raise RuntimeError("Candidate scoring returned an invalid flow-noise shape")
+                if tuple(time.shape) != (repeat_count * len(chunk),):
+                    raise RuntimeError("Candidate scoring returned an invalid flow-time shape")
+                noise = noise.reshape(repeat_count, len(chunk), *noise.shape[1:])
+                time = time.reshape(repeat_count, len(chunk))
+                for local_index, item in enumerate(chunk):
+                    # Reuse the scored Flow-Matching realization during backward.
+                    item["loss_is_flow_noise"] = noise[:, local_index].clone()
+                    item["loss_is_flow_time"] = time[:, local_index].clone()
+                scores.append(chunk_scores)
+        finally:
+            if was_training:
+                self.model.train()
+            self._enforce_frozen_module_eval()
+        result = torch.cat(scores)
+        if not torch.isfinite(result).all() or torch.any(result <= 0):
+            raise FloatingPointError("Candidate losses must be finite and strictly positive")
+        return result
+
+    def _selection_generator(self, base_ids: list[int]) -> torch.Generator:
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(
+            _stable_seed(
+                self._loss_is_seed,
+                "selection",
+                self.completed_steps,
+                self._microbatch_index,
+                rank,
+                min(base_ids),
+            )
+        )
+        return generator
+
+    def _get_next_batch(self):
+        base_batch = super()._get_next_batch()
+        self._begin_loss_is_update()
+        base_ids = [int(item["sampling_trace"]["global_draw_id"]) for item in base_batch]
+
+        candidates = self._build_candidate_pool(base_batch)
+        scores = self._score_candidates(candidates, chunk_size=len(base_batch))
+        selection = sample_with_replacement(
+            scores,
+            len(base_batch),
+            generator=self._selection_generator(base_ids),
+        )
+        selected = []
+        for occurrence_id, selected_index, weight, probability in zip(
+            base_ids,
+            selection.indices.tolist(),
+            selection.importance_weights.tolist(),
+            selection.selected_probabilities.tolist(),
+        ):
+            item = copy.deepcopy(candidates[selected_index])
+            trace = item["sampling_trace"]
+            trace["loss_is_candidate_global_draw_id"] = int(trace["global_draw_id"])
+            trace["loss_is_occurrence_id"] = int(occurrence_id)
+            trace["loss_is_importance_weight"] = float(weight)
+            trace["loss_is_probability"] = float(probability)
+            trace["loss_is_candidate_score"] = float(scores[selected_index])
+            trace["loss_is_active"] = True
+            selected.append(item)
+        weights = selection.importance_weights
+        self._loss_is_selection_metrics.append(
+            {
+                "candidate_loss_mean": float(scores.mean()),
+                "candidate_loss_max": float(scores.max()),
+                "weight_mean": float(weights.mean()),
+                "weight_max": float(weights.max()),
+                "ess": float(weights.sum().square() / weights.square().sum()),
+            }
+        )
+        return selected
+
+    def _after_per_sample_training_loss(
+        self,
+        batch_vla,
+        per_sample_loss,
+        *,
+        did_optimizer_step: bool,
+        microbatch_index: int,
+    ) -> None:
+        if per_sample_loss is None:
+            raise RuntimeError("Loss-IS requires per-sample training loss")
+        if not did_optimizer_step:
+            return
+
+        metrics = {
+            "loss_is/active": 1.0,
+            "loss_is/candidate_factor": float(self._loss_is_factor),
+        }
+        if self._loss_is_selection_metrics:
+            for key in self._loss_is_selection_metrics[0]:
+                metrics[f"loss_is/{key}"] = sum(
+                    item[key] for item in self._loss_is_selection_metrics
+                ) / len(self._loss_is_selection_metrics)
+        self._loss_is_last_metrics = metrics
+
+        if self.accelerator.is_main_process:
+            state_path = Path(self.config.output_dir) / "loss_is_state.json"
+            temporary = state_path.with_suffix(".json.tmp")
+            payload = {
+                "schema_version": 1,
+                "completed_optimizer_step": int(self.completed_steps + 1),
+                "activation_mode": "always_on",
+                "active_for_completed_step": True,
+                "candidate_factor": self._loss_is_factor,
+                "selection_scope": "per-rank per-microbatch; globally averaged unbiased estimator",
+                "metrics": metrics,
+            }
+            temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+            os.replace(temporary, state_path)
+
+    def _train_step(self, batch_vla, batch_vlm=None):
+        metrics = super()._train_step(batch_vla, batch_vlm=batch_vlm)
+        metrics.update(self._loss_is_last_metrics)
+        return metrics
+
+
+def install_smoke_manifest_fast_path() -> None:
+    if os.environ.get("LOSS_IS_SMOKE_FROZEN_MANIFEST") != "1":
+        return
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+    from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
+
+    def load_frozen_manifest(self):
+        path = self.data_cfg.get("valid_anchor_manifest_path")
+        if not path:
+            raise ValueError("Smoke fast path requires a frozen valid-anchor manifest")
+        manifest = ValidAnchorManifest.load(path)
+        print(
+            f"Loss-IS smoke frozen manifest: {path}",
+            flush=True,
+        )
+        return manifest
+
+    LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, clipargs = parser.parse_known_args()
+
+    install_smoke_manifest_fast_path()
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
+    cfg = apply_config_compat(cfg, strict=True)
+    if cfg.datasets.vla_data.temporal_proposal_path is not None:
+        raise ValueError("Loss-IS uses the baseline distribution, not a temporal proposal")
+    if cfg.trainer.sample_loss_weight_key != "loss_is_importance_weight":
+        raise ValueError("sample_loss_weight_key must be loss_is_importance_weight")
+    if cfg.trainer.model_rng_trace_key != "loss_is_occurrence_id":
+        raise ValueError("model_rng_trace_key must be loss_is_occurrence_id")
+    cfg.config_yaml = args.config_yaml
+
+    base_trainer = train_starvla.VLATrainer
+
+    class LossISTrainer(LossISTrainerMixin, base_trainer):
+        pass
+
+    train_starvla.VLATrainer = LossISTrainer
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/qwen_adapted/train_qwen_frameskip_continuation.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""Thin runtime adapter for the locked Qwen-Bridge FrameSkip continuation."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+from omegaconf import OmegaConf
+
+
+def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    if payload.get("method") != "FrameSkip-adapted (Gaussian-TPI fallback)":
+        raise RuntimeError("unexpected FrameSkip cache method")
+    cache = {}
+    for record in payload["records"]:
+        key = (record["dataset_name"], str(record["trajectory_id"]))
+        keep = np.asarray(record["keep_indices"], dtype=np.int64)
+        if len(keep) == 0 or np.any(np.diff(keep) <= 0):
+            raise RuntimeError(f"invalid retained indices for {key}")
+        cache[key] = keep
+    return cache, payload
+
+
+def install_adapter(
+    cache_path: str,
+    start_global_draw_id: int,
+    global_microbatch_size: int,
+) -> None:
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+
+    cache, cache_payload = load_cache(cache_path)
+    if os.environ.get("FRAMESKIP_SMOKE_FROZEN_MANIFEST") == "1":
+        from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
+
+        original_build_manifest = LeRobotMixtureDataset._build_valid_anchor_manifest
+
+        def load_frozen_manifest(self):
+            path = self.data_cfg.get("valid_anchor_manifest_path")
+            if not path:
+                return original_build_manifest(self)
+            manifest = ValidAnchorManifest.load(path)
+            print(
+                f"FrameSkip smoke frozen manifest: {path}",
+                flush=True,
+            )
+            return manifest
+
+        LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest
+
+    original_getitem = LeRobotMixtureDataset.__getitem__
+    period = 6
+
+    def frameskip_getitem(self, index):
+        trace = self.sample_trace(index)
+        offset = trace.global_draw_id - start_global_draw_id
+        if offset < 0:
+            raise RuntimeError("FrameSkip draw precedes the certified continuation start")
+        microbatch = offset // global_microbatch_size
+        is_pruned = microbatch % period < 5
+        logical_index = int(trace.base_index)
+        actual_index = logical_index
+        if is_pruned:
+            key = (trace.dataset_name, str(trace.trajectory_id))
+            if key not in cache:
+                raise RuntimeError(f"FrameSkip cache is missing required trajectory {key}")
+            retained = cache[key]
+            position = int(np.searchsorted(retained, logical_index))
+            position = int(np.clip(position, 0, len(retained) - 1))
+            actual_index = int(retained[position])
+
+        dataset = self.datasets[trace.dataset_index]
+        original_get_step_data = dataset.get_step_data
+
+        def redirected_get_step_data(trajectory_id, base_index):
+            if str(trajectory_id) != str(trace.trajectory_id) or int(base_index) != logical_index:
+                raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
+            return original_get_step_data(trajectory_id, actual_index)
+
+        dataset.get_step_data = redirected_get_step_data
+        try:
+            sample = original_getitem(self, index)
+        finally:
+            dataset.get_step_data = original_get_step_data
+        sample["sampling_trace"]["frameskip"] = {
+            "method": cache_payload["method"],
+            "cache_record_count": len(cache_payload["records"]),
+            "is_pruned_minibatch": is_pruned,
+            "logical_base_index": logical_index,
+            "physical_base_index": actual_index,
+            "retention_ratio": 0.2 if is_pruned else 1.0,
+        }
+        return sample
+
+    LeRobotMixtureDataset.__getitem__ = frameskip_getitem
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, clipargs = parser.parse_known_args()
+    cache_path = os.environ.get("FRAMESKIP_CACHE")
+    if not cache_path:
+        raise RuntimeError("FRAMESKIP_CACHE is required")
+    start_draw = int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "0"))
+    global_microbatch_size = int(
+        os.environ.get("FRAMESKIP_GLOBAL_MICROBATCH_SIZE", "2")
+    )
+    if global_microbatch_size <= 0:
+        raise ValueError("FRAMESKIP_GLOBAL_MICROBATCH_SIZE must be positive")
+    install_adapter(cache_path, start_draw, global_microbatch_size)
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
+    cfg = apply_config_compat(cfg, strict=True)
+    cfg.config_yaml = args.config_yaml
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

external_baselines/qwen_robotwin_adapted/register_qwen_robotwin_smoke.py

@@ -0,0 +1,17 @@
+"""Register the existing one-task RoboTwin smoke mixture."""
+
+from starVLA.dataloader.gr00t_lerobot.registry import (
+    DATASET_NAMED_MIXTURES,
+    ROBOT_TYPE_CONFIG_MAP,
+)
+
+
+MIXTURE_NAME = "qwen_robotwin_smoke_adjust_bottle"
+
+
+def register_qwen_robotwin_smoke() -> None:
+    if "robotwin50" not in ROBOT_TYPE_CONFIG_MAP:
+        raise RuntimeError("The RoboTwin robot configuration is unavailable")
+    DATASET_NAMED_MIXTURES[MIXTURE_NAME] = [
+        ("adjust_bottle", 1.0, "robotwin50"),
+    ]

external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_frameskip.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Run FrameSkip for Qwen-RoboTwin using the shared RoboTwin adapter."""
+
+import os
+
+from external_baselines.qwen_robotwin_adapted.register_qwen_robotwin_smoke import (
+    register_qwen_robotwin_smoke,
+)
+
+
+def main() -> None:
+    if os.environ.get("QWEN_ROBOTWIN_REGISTER_SMOKE") == "1":
+        register_qwen_robotwin_smoke()
+    from external_baselines.frameskip_robotwin_adapted.train_frameskip_robotwin_continuation import (
+        main as run,
+    )
+
+    run()
+
+
+if __name__ == "__main__":
+    main()

external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_keyframe.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Run Keyframe-Focused for Qwen-RoboTwin."""
+
+import os
+
+from external_baselines.qwen_robotwin_adapted.register_qwen_robotwin_smoke import (
+    register_qwen_robotwin_smoke,
+)
+
+
+def main() -> None:
+    if os.environ.get("QWEN_ROBOTWIN_REGISTER_SMOKE") == "1":
+        register_qwen_robotwin_smoke()
+    from external_baselines.keyframe_focused_adapted.train_keyframe_continuation import (
+        main as run,
+    )
+
+    run()
+
+
+if __name__ == "__main__":
+    main()

external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_loss_is.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Run the canonical always-on Loss-IS implementation for Qwen-RoboTwin."""
+
+import os
+
+from external_baselines.qwen_robotwin_adapted.register_qwen_robotwin_smoke import (
+    register_qwen_robotwin_smoke,
+)
+
+
+def main() -> None:
+    if os.environ.get("QWEN_ROBOTWIN_REGISTER_SMOKE") == "1":
+        register_qwen_robotwin_smoke()
+    from external_baselines.loss_is_active.train_active_loss_is_continuation import (
+        main as run,
+    )
+
+    run()
+
+
+if __name__ == "__main__":
+    main()

external_baselines/settings/florence_bridge/frameskip.yaml

@@ -0,0 +1,71 @@
+datasets:
+  vla_data:
+    action_mode: abs
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    data_mix: bridge_rt_infoflow
+    data_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    dataset_content_fingerprint_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    dataset_provenance_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_py: lerobot_datasets
+    expected_world_size: 2
+    num_workers: 4
+    per_device_batch_size: 4
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    precomputed_statistics_path: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    temporal_proposal_path: null
+    traceable_sampling: true
+    valid_anchor_manifest_migration_path: /home/miaochengyu/starVLA-infoflow-codex/external_baselines/frameskip_adapted/preflight/valid_anchor_manifest_migration_20260825.json
+    valid_anchor_manifest_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    video_backend: torchvision_av
+framework:
+  action_model:
+    action_dim: 7
+    action_horizon: 16
+    action_model_type: DiT-B
+    add_pos_embed: true
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    hidden_size: 1024
+    max_seq_len: 1024
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    num_timestep_buckets: 1000
+    state_dim: 7
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-2-large
+output_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/external-baselines/frameskip-adapted-gaussian-tpi-r02-20260825/frameskip-adapted-gaussian-tpi-r02-seed-117465
+run_id: frameskip-adapted-gaussian-tpi-r02-seed-117465
+run_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/external-baselines/frameskip-adapted-gaussian-tpi-r02-20260825
+seed: 117465
+trainer:
+  freeze_modules: ''
+  gradient_accumulation_steps: 8
+  learning_rate:
+    action_model: 0.0001
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+  lr_scheduler_type: cosine_with_min_lr
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  optimizer:
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    name: AdamW
+    weight_decay: 1.0e-08
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07

external_baselines/settings/florence_bridge/keyframe.yaml

@@ -0,0 +1,70 @@
+datasets:
+  vla_data:
+    action_mode: abs
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    data_mix: bridge_rt_infoflow
+    data_root_dir: /home/liuyu/student-workspace/data/runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/runtime/manifests/bridgert/content_fingerprints.json
+    dataset_provenance_path: /home/liuyu/student-workspace/code/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_py: lerobot_datasets
+    expected_world_size: 2
+    num_workers: 4
+    per_device_batch_size: 4
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    precomputed_statistics_path: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    temporal_proposal_path: null
+    traceable_sampling: true
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    video_backend: torchvision_av
+framework:
+  action_model:
+    action_dim: 7
+    action_horizon: 16
+    action_model_type: DiT-B
+    add_pos_embed: true
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    hidden_size: 1024
+    max_seq_len: 1024
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    num_timestep_buckets: 1000
+    state_dim: 7
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-2-large
+output_dir: /home/liuyu/student-workspace/outputs/keyframe_focused_adapted/runs/keyframe-focused-adapted-seed-117465-20k-to-20500
+run_id: keyframe-focused-adapted-seed-117465-20k-to-20500
+run_root_dir: /home/liuyu/student-workspace/outputs/keyframe_focused_adapted/runs
+seed: 117465
+trainer:
+  freeze_modules: ''
+  gradient_accumulation_steps: 8
+  learning_rate:
+    action_model: 0.0001
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+  lr_scheduler_type: cosine_with_min_lr
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  optimizer:
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    name: AdamW
+    weight_decay: 1.0e-08
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07

external_baselines/settings/florence_bridge/loss_is_active.yaml

@@ -0,0 +1,117 @@
+run_id: florence-bridge-active-loss-is-20k-to20500-seed117465
+run_root_dir: /home/liuyu/student-workspace/outputs/active_loss_is_20260907/formal_runs
+seed: 117465
+wandb_entity: null
+wandb_project: Loss-IS-adapted-BridgeRT
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 15
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size:
+    - 224
+    - 224
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: /home/liuyu/student-workspace/code/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/runtime/manifests/bridgert/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    valid_anchor_manifest_migration_path: null
+    temporal_proposal_path: null
+    transform_rng_seed: 298226
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 20500
+  evaluation_snapshot_interval: 1000
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/runtime/runs/florence-scratch-official-100k/florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726/checkpoints/steps_20000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  model_rng_trace_key: loss_is_occurrence_id
+  sample_loss_weight_key: loss_is_importance_weight
+  stop_after_steps: 20500
+  loss_is:
+    candidate_factor: 3
+    seed: 117465

external_baselines/settings/florence_robotwin/frameskip.yaml

@@ -0,0 +1,109 @@
+run_id: florence-robotwin-frameskip-10000-to-10500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/robotwin_external_20260831/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: eager
+    vl_hidden_dim: 1024
+    selected_view_indices:
+    - 0
+    - 1
+    - 2
+    compile_vision_tower: true
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 2
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    temporal_proposal_path: null
+    transform_rng_seed: 289726
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 10500
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 16
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 10500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/florence-groot-robotwin-xvla3view-official150k-seed42-20260804/checkpoints/steps_10000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927
+  evaluation_snapshot_interval: null

external_baselines/settings/florence_robotwin/keyframe.yaml

@@ -0,0 +1,110 @@
+run_id: florence-robotwin-keyframe-10000-to-10500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/robotwin_external_20260831/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: eager
+    vl_hidden_dim: 1024
+    selected_view_indices:
+    - 0
+    - 1
+    - 2
+    compile_vision_tower: true
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 2
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    temporal_proposal_path: null
+    transform_rng_seed: 289726
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 10500
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 16
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 10500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/florence-groot-robotwin-xvla3view-official150k-seed42-20260804/checkpoints/steps_10000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927
+  evaluation_snapshot_interval: null
+  sample_loss_weight_key: keyframe_weight

external_baselines/settings/florence_robotwin/loss_is_active.yaml

@@ -0,0 +1,114 @@
+run_id: florence-robotwin-active-loss-is-10000-to10500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/active_loss_is_20260907/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: eager
+    vl_hidden_dim: 1024
+    selected_view_indices:
+    - 0
+    - 1
+    - 2
+    compile_vision_tower: true
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 2
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    temporal_proposal_path: null
+    transform_rng_seed: 289726
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 10500
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 16
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 10500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/florence-groot-robotwin-xvla3view-official150k-seed42-20260804/checkpoints/steps_10000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927
+  evaluation_snapshot_interval: null
+  sample_loss_weight_key: loss_is_importance_weight
+  model_rng_trace_key: loss_is_occurrence_id
+  loss_is:
+    candidate_factor: 3
+    seed: 484370

external_baselines/settings/qwen_bridge/frameskip.yaml

@@ -0,0 +1,93 @@
+datasets:
+  vla_data:
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    action_mode: abs
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    data_mix: bridge_rt_infoflow
+    data_root_dir: /data/yanfei/qwen_bridge_assets/OXE_LEROBOT_DATASET
+    dataset_content_fingerprint_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    dataset_provenance_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_py: lerobot_datasets
+    expected_world_size: 1
+    num_workers: 1
+    per_device_batch_size: 2
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    precomputed_statistics_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    temporal_proposal_path: null
+    traceable_sampling: true
+    valid_anchor_manifest_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    video_backend: torchvision_av
+framework:
+  action_model:
+    action_dim: 7
+    action_horizon: 16
+    action_model_type: DiT-B
+    add_pos_embed: true
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    hidden_size: 1024
+    max_seq_len: 1024
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    num_timestep_buckets: 1000
+    state_dim: 7
+  name: QwenGR00T
+  qwenvl:
+    attn_implementation: flash_attention_2
+    base_vlm: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+output_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs/qwen3vl-bridge-frameskip-29k-to29500-seed42
+run_id: qwen3vl-bridge-frameskip-29k-to29500-seed42
+run_root_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs
+seed: 42
+trainer:
+  allow_model_only_resume: false
+  eval_interval: null
+  evaluation_snapshot_end_step: null
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: null
+  freeze_modules: ''
+  frozen_modules_eval: false
+  gradient_accumulation_steps: 24
+  gradient_clipping: 1.0
+  is_resume: true
+  learning_rate:
+    action_model: 0.0001
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+  logging_frequency: 10
+  lr_scheduler_type: cosine_with_min_lr
+  max_train_steps: 100000
+  num_warmup_steps: 10000
+  optimizer:
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    name: AdamW8bit
+    weight_decay: 1.0e-08
+  pretrained_checkpoint: null
+  repeated_diffusion_steps: 4
+  save_format: pt
+  save_interval: null
+  save_steps:
+  - 29500
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  stop_after_steps: 29500
+  trainable_modules: ''
+wandb_entity: null
+wandb_project: QwenBridge-CommonStart

external_baselines/settings/qwen_bridge/keyframe.yaml

@@ -0,0 +1,94 @@
+datasets:
+  vla_data:
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    action_mode: abs
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    data_mix: bridge_rt_infoflow
+    data_root_dir: /data/yanfei/qwen_bridge_assets/OXE_LEROBOT_DATASET
+    dataset_content_fingerprint_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    dataset_provenance_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_py: lerobot_datasets
+    expected_world_size: 1
+    num_workers: 1
+    per_device_batch_size: 2
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    precomputed_statistics_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    temporal_proposal_path: null
+    traceable_sampling: true
+    valid_anchor_manifest_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    video_backend: torchvision_av
+framework:
+  action_model:
+    action_dim: 7
+    action_horizon: 16
+    action_model_type: DiT-B
+    add_pos_embed: true
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    hidden_size: 1024
+    max_seq_len: 1024
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    num_timestep_buckets: 1000
+    state_dim: 7
+  name: QwenGR00T
+  qwenvl:
+    attn_implementation: flash_attention_2
+    base_vlm: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+output_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs/qwen3vl-bridge-keyframe-29k-to29500-seed42
+run_id: qwen3vl-bridge-keyframe-29k-to29500-seed42
+run_root_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs
+seed: 42
+trainer:
+  allow_model_only_resume: false
+  eval_interval: null
+  evaluation_snapshot_end_step: null
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: null
+  freeze_modules: ''
+  frozen_modules_eval: false
+  gradient_accumulation_steps: 24
+  gradient_clipping: 1.0
+  is_resume: true
+  learning_rate:
+    action_model: 0.0001
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+  logging_frequency: 10
+  lr_scheduler_type: cosine_with_min_lr
+  max_train_steps: 100000
+  num_warmup_steps: 10000
+  optimizer:
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    name: AdamW8bit
+    weight_decay: 1.0e-08
+  pretrained_checkpoint: null
+  repeated_diffusion_steps: 4
+  sample_loss_weight_key: keyframe_weight
+  save_format: pt
+  save_interval: null
+  save_steps:
+  - 29500
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  stop_after_steps: 29500
+  trainable_modules: ''
+wandb_entity: null
+wandb_project: QwenBridge-CommonStart

external_baselines/settings/qwen_bridge/loss_is_active.yaml

@@ -0,0 +1,99 @@
+datasets:
+  vla_data:
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    action_mode: abs
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    data_mix: bridge_rt_infoflow
+    data_root_dir: /data/yanfei/qwen_bridge_assets/OXE_LEROBOT_DATASET
+    dataset_content_fingerprint_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    dataset_provenance_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_py: lerobot_datasets
+    expected_world_size: 1
+    num_workers: 1
+    per_device_batch_size: 2
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    precomputed_statistics_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    temporal_proposal_path: null
+    traceable_sampling: true
+    valid_anchor_manifest_path: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    video_backend: torchvision_av
+framework:
+  action_model:
+    action_dim: 7
+    action_horizon: 16
+    action_model_type: DiT-B
+    add_pos_embed: true
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    hidden_size: 1024
+    max_seq_len: 1024
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    num_timestep_buckets: 1000
+    state_dim: 7
+  name: QwenGR00T
+  qwenvl:
+    attn_implementation: flash_attention_2
+    base_vlm: /home/yanfei/miaochengyu/robotwin-infoflow/legacy_bridge_runtime/compat/mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+output_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs/qwen3vl-bridge-lossis-29k-to29500-seed42
+run_id: qwen3vl-bridge-lossis-29k-to29500-seed42
+run_root_dir: /data/yanfei/qwen_bridge_external_baselines_29k_20260912/runs
+seed: 42
+trainer:
+  allow_model_only_resume: false
+  eval_interval: null
+  evaluation_snapshot_end_step: null
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: null
+  freeze_modules: ''
+  frozen_modules_eval: false
+  gradient_accumulation_steps: 24
+  gradient_clipping: 1.0
+  is_resume: true
+  learning_rate:
+    action_model: 0.0001
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+  logging_frequency: 10
+  loss_is:
+    candidate_factor: 3
+    initial_branch_step: 29000
+    seed: 42
+  lr_scheduler_type: cosine_with_min_lr
+  max_train_steps: 100000
+  model_rng_trace_key: loss_is_occurrence_id
+  num_warmup_steps: 10000
+  optimizer:
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    name: AdamW8bit
+    weight_decay: 1.0e-08
+  pretrained_checkpoint: null
+  repeated_diffusion_steps: 4
+  sample_loss_weight_key: loss_is_importance_weight
+  save_format: pt
+  save_interval: null
+  save_steps:
+  - 29500
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  stop_after_steps: 29500
+  trainable_modules: ''
+wandb_entity: null
+wandb_project: QwenBridge-CommonStart

external_baselines/settings/qwen_robotwin/frameskip.yaml

@@ -0,0 +1,107 @@
+run_id: qwen-robotwin-frameskip-11000-to-11500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/robotwin_external_20260831/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 2048
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 5
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    transform_rng_seed: 289726
+    temporal_proposal_path: null
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 11500
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: 10000
+  evaluation_snapshot_end_step: 20000
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 4
+  save_format: pt
+  save_final_checkpoint: false
+  save_consolidated_model: false
+  optimizer:
+    name: AdamW8bit
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 11500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/qwen3vl-groot-robotwin-official150k-adamw8bit-mb5-ga4-seed42-20260810/checkpoints/steps_11000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927

external_baselines/settings/qwen_robotwin/keyframe.yaml

@@ -0,0 +1,108 @@
+run_id: qwen-robotwin-keyframe-11000-to-11500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/robotwin_external_20260831/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 2048
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 5
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    transform_rng_seed: 289726
+    temporal_proposal_path: null
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 11500
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: 10000
+  evaluation_snapshot_end_step: 20000
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 4
+  save_format: pt
+  save_final_checkpoint: false
+  save_consolidated_model: false
+  optimizer:
+    name: AdamW8bit
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 11500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/qwen3vl-groot-robotwin-official150k-adamw8bit-mb5-ga4-seed42-20260810/checkpoints/steps_11000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927
+  sample_loss_weight_key: keyframe_weight

external_baselines/settings/qwen_robotwin/loss_is_active.yaml

@@ -0,0 +1,112 @@
+run_id: qwen-robotwin-active-loss-is-11000-to11500-seed484370
+run_root_dir: /home/liuyu/student-workspace/outputs/active_loss_is_20260907/formal_runs
+seed: 484370
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-RoboTwin
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /home/liuyu/student-workspace/data/runtime/playground/Pretrained_models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 2048
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 14
+    state_dim: 0
+    action_horizon: 50
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 2560
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 49
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /home/liuyu/student-workspace/data/robotwin/RoboTwin
+    data_mix: robotwin_all_50
+    action_type: abs_qpos
+    action_mode: abs
+    include_state: false
+    sequential_step_sampling: false
+    per_device_batch_size: 5
+    load_all_data_for_training: true
+    image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    num_workers: 4
+    expected_world_size: 2
+    traceable_sampling: true
+    dataset_provenance_path: /home/liuyu/student-workspace/data/robotwin/manifests/data_provenance.json
+    dataset_content_fingerprint_path: /home/liuyu/student-workspace/data/robotwin/manifests/content_fingerprints.json
+    valid_anchor_manifest_path: /home/liuyu/student-workspace/data/robotwin/manifests/valid_anchor_manifest.json
+    transform_rng_seed: 289726
+    temporal_proposal_path: null
+trainer:
+  use_deepspeed: true
+  max_train_steps: 150000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 11500
+  evaluation_snapshot_interval: null
+  evaluation_snapshot_start_step: 10000
+  evaluation_snapshot_end_step: 20000
+  eval_interval: null
+  repeated_diffusion_steps: 4
+  learning_rate:
+    base: 1.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 4
+  save_format: pt
+  save_final_checkpoint: false
+  save_consolidated_model: false
+  optimizer:
+    name: AdamW8bit
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  stop_after_steps: 11500
+  is_resume: true
+  allow_model_only_resume: false
+  resume_accelerate_checkpoint: /home/liuyu/student-workspace/data/robotwin/common_start/qwen3vl-groot-robotwin-official150k-adamw8bit-mb5-ga4-seed42-20260810/checkpoints/steps_11000_accelerate
+  branch_parent_seed: 42
+  branch_rng_seed: 484370
+  model_rng_seed: 721927
+  sample_loss_weight_key: loss_is_importance_weight
+  model_rng_trace_key: loss_is_occurrence_id
+  loss_is:
+    candidate_factor: 3
+    seed: 484370

starVLA/dataloader/__init__.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 import json
 import os
 from accelerate.logging import get_logger
@@ -100,10 +101,4 @@ def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"):
                 manifest_path = output_dir / "valid_anchor_manifest.json"
                 vla_dataset.valid_anchor_manifest.save(manifest_path)
         return vla_train_dataloader
-    elif dataset_py == "vlm_datasets":
-        from starVLA.dataloader.vlm_datasets import make_vlm_dataloader
-
-        vlm_data_module = make_vlm_dataloader(cfg)
-        vlm_train_dataloader = vlm_data_module["train_dataloader"]
-        
-        return vlm_train_dataloader
+    raise ValueError(f"Unsupported dataset loader: {dataset_py}")

starVLA/dataloader/gr00t_lerobot/data_config.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -854,7 +855,7 @@ class AgilexDataConfig:
     ]
     action_keys = [
         "action.left_joints",
-        "action.right_joints",#@JinhuiYE this order is different from Dataset
+        "action.right_joints",  # This order differs from the dataset schema.
         "action.left_gripper",
         "action.right_gripper",
     ]
@@ -934,7 +935,7 @@ class AgilexData50Config:
     ]
     action_keys = [
         "action.left_joints",
-        "action.right_joints",  # @JinhuiYE this order is different from Dataset
+        "action.right_joints",  # This order differs from the dataset schema.
         "action.left_gripper",
         "action.right_gripper",
     ]

starVLA/dataloader/gr00t_lerobot/datasets.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -55,9 +56,11 @@ from starVLA.dataloader.gr00t_lerobot.transform.state_action import StateActionT
 from starVLA.dataloader.infoflow_sampler import (
     LogicalDraw,
     ValidAnchorManifest,
+    compute_dataset_content_fingerprint,
     validate_dataset_content_inventory,
 )
 from starVLA.dataloader.temporal_proposal import TemporalProposal
+from starVLA.dataloader.relative_pose import relative_values, rotation_type_is_quaternion
 
 from functools import partial
 from typing import Tuple, List
@@ -152,7 +155,7 @@ def calculate_dataset_statistics(parquet_paths: list[Path]) -> dict:
 
 def _normalize_action_mode(mode: str) -> str:
     """Normalize action mode names to {abs, delta, rel}.""" 
-    # @gaoning plz move this, we want dataloader to be independent of the action mode logic, we can move this to transform or a separate utils tool to handle lerobot dataset
+    # Keep action-mode logic isolated from the generic dataset path.
     mode = str(mode).lower()
     if mode in {"absolute", "raw"}:
         mode = "abs"
@@ -390,6 +393,17 @@ def _get_action_col_slices(
     return action_col_slices
 
 
+def _action_slice_is_quaternion(
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_slice: tuple[int, int],
+) -> bool:
+    start, end = action_slice
+    for cfg in lerobot_modality_meta.action.values():
+        if cfg.start == start and cfg.end == end:
+            return rotation_type_is_quaternion(cfg.rotation_type)
+    return False
+
+
 def calculate_delta_action_statistics(
     parquet_paths: list[Path],
     lerobot_modality_meta: "LeRobotModalityMetadata",
@@ -504,7 +518,8 @@ def calculate_rel_action_statistics(
     Calculate action statistics using rel mode.
 
     Rule:
-      - For all t: a_t - s_0
+      - Positions / scalars: a_t - s_0
+      - Quaternion keys: q_rel = inv(q_s0) ⊗ q_a_t (still a unit quaternion)
 
     Mapping rule (only two cases):
       1) Use explicit action_mode_state_map if provided.
@@ -563,7 +578,13 @@ def calculate_rel_action_statistics(
                     if action_part_chunk.shape[1] != state_chunk.shape[1]:
                         raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}")
 
-                    out = action_part_chunk - state_chunk[0]
+                    out = relative_values(
+                        action_part_chunk,
+                        state_chunk[0],
+                        is_quaternion=_action_slice_is_quaternion(
+                            lerobot_modality_meta, a_slice
+                        ),
+                    )
                     action_chunk_full[:, a_slice[0] : a_slice[1]] = out
 
                 accum[action_col].append(action_chunk_full)
@@ -1005,9 +1026,7 @@ class LeRobotSingleDataset(Dataset):
                         video_key = str(col)[len("videos/") : -len("/from_timestamp")]
                         from_timestamps[video_key] = float(value)
 
-                    # TODO auto map key 
                     # Collect video file indices for each video key
-                    #已修改的lerobotv3.0的视频索引（提取视频和文件的索引）
                     video_file_indices = {}
                     for col in timestamp_cols:
                         video_key = str(col)[len("videos/") : -len("/from_timestamp")]
@@ -1303,6 +1322,16 @@ class LeRobotSingleDataset(Dataset):
             self.data_cfg.get("action_mode_state_map", {}) or {}
         )
 
+    def _action_key_is_quaternion(self, action_key: str) -> bool:
+        if not action_key.startswith("action."):
+            return False
+        subkey = action_key.replace("action.", "", 1)
+        try:
+            cfg = self._lerobot_modality_meta.action[subkey]
+        except (AttributeError, KeyError):
+            return False
+        return rotation_type_is_quaternion(cfg.rotation_type)
+
     def _infer_state_key_for_action(self, action_key: str) -> str | None:
         if action_key in self._action_mode_state_map:
             return self._action_mode_state_map[action_key]
@@ -1341,13 +1370,19 @@ class LeRobotSingleDataset(Dataset):
                 )
 
             state0 = state_values[0]
+            is_quaternion = self._action_key_is_quaternion(action_key)
             if self._action_mode == "delta":
                 out = action_values.copy()
+                if is_quaternion:
+                    raise ValueError(
+                        f"action_mode=delta is not defined for quaternion key {action_key}; "
+                        "use rel (SE(3) relative to the current pose) or abs."
+                    )
                 if len(out) > 1:
                     out[1:] = action_values[1:] - action_values[:-1]
                 out[0] = action_values[0] - state0
             elif self._action_mode == "rel":
-                out = action_values - state0
+                out = relative_values(action_values, state0, is_quaternion=is_quaternion)
             else:
                 out = action_values
 
@@ -1523,7 +1558,7 @@ class LeRobotSingleDataset(Dataset):
         data = {}
         # Get the data for all modalities # just for action base data
         self.curr_traj_data = self.get_trajectory_data(trajectory_id)
-        # TODO @JinhuiYE The logic below is poorly implemented. Data reading should be directly based on curr_traj_data.
+        # Read the requested trajectory from the active trajectory table.
         for modality in self.modality_keys:
             # Get the data corresponding to each key in the modality
             for key in self.modality_keys[modality]:
@@ -1660,8 +1695,6 @@ class LeRobotSingleDataset(Dataset):
             episode_meta = self.trajectory_ids_to_metadata[trajectory_id]
 
             video_file_indices = episode_meta.get("videos/file_indices", {})
-            # print(f"{video_file_indices=}")
-            #已修改的lerobotv3.0的视频索引
             if original_key in video_file_indices:
                 video_chunk_index = video_file_indices[original_key]["chunk_index"]
                 video_file_index = video_file_indices[original_key]["file_index"]
@@ -2546,35 +2579,21 @@ class LeRobotMixtureDataset(Dataset):
 
         dataset_provenance = None
         provenance_path = self.data_cfg.get("dataset_provenance_path") if self.data_cfg else None
-        if not provenance_path:
-            raise ValueError(
-                "traceable_sampling requires dataset_provenance_path so the frozen support "
-                "is bound to exact dataset revisions"
-            )
-        provenance_path = Path(provenance_path)
-        with open(provenance_path, "r", encoding="utf-8") as provenance_file:
-            provenance_payload = json.load(provenance_file)
-        provenance_bytes = provenance_path.read_bytes()
-        dataset_provenance = {
-            "sha256": hashlib.sha256(provenance_bytes).hexdigest(),
-            "payload": provenance_payload,
-        }
+        if provenance_path:
+            with open(provenance_path, "r", encoding="utf-8") as provenance_file:
+                dataset_provenance = json.load(provenance_file)
 
         content_fingerprint_path = (
             self.data_cfg.get("dataset_content_fingerprint_path") if self.data_cfg else None
         )
-        if not content_fingerprint_path:
-            raise ValueError(
-                "traceable_sampling requires dataset_content_fingerprint_path; generate the "
-                "full content hashes once before launching distributed training"
-            )
-        content_fingerprint_path = Path(content_fingerprint_path)
-        with open(content_fingerprint_path, "r", encoding="utf-8") as fingerprint_file:
-            content_fingerprints = json.load(fingerprint_file)
-        content_fingerprint_artifact_sha256 = hashlib.sha256(
-            content_fingerprint_path.read_bytes()
-        ).hexdigest()
-
+        if content_fingerprint_path:
+            with open(content_fingerprint_path, "r", encoding="utf-8") as fingerprint_file:
+                content_fingerprints = json.load(fingerprint_file)
+        else:
+            content_fingerprints = {"datasets": {
+                dataset.dataset_name: compute_dataset_content_fingerprint(dataset.dataset_path)
+                for dataset in self.datasets
+            }}
         dataset_identities = []
         for dataset in self.datasets:
             try:
@@ -2592,9 +2611,9 @@ class LeRobotMixtureDataset(Dataset):
                     "modality": dataset._lerobot_modality_meta.model_dump(mode="json"),
                     "transforms": serialize_transforms(dataset),
                     "content": {
-                        "sha256": content_fingerprint["sha256"],
                         "file_count": content_fingerprint["file_count"],
                         "total_bytes": content_fingerprint["total_bytes"],
+                        "files": content_fingerprint["files"],
                     },
                     "delta_indices": {
                         key: value.tolist() for key, value in dataset.delta_indices.items()
@@ -2608,7 +2627,6 @@ class LeRobotMixtureDataset(Dataset):
             "balance_trajectory_weights": self.balance_trajectory_weights,
             "excluded_episodes": excluded_episodes,
             "dataset_provenance": dataset_provenance,
-            "content_fingerprint_artifact_sha256": content_fingerprint_artifact_sha256,
             "datasets": dataset_identities,
         }
         identity = json.loads(json.dumps(identity, sort_keys=True, default=str))
@@ -2630,6 +2648,13 @@ class LeRobotMixtureDataset(Dataset):
             frozen_identity = json.loads(json.dumps(frozen_manifest.identity))
             current_identity = json.loads(json.dumps(current_manifest.identity))
             for identity in (frozen_identity, current_identity):
+                provenance = identity.get("dataset_provenance")
+                if isinstance(provenance, dict) and "payload" in provenance:
+                    identity["dataset_provenance"] = provenance["payload"]
+                identity.pop("content_fingerprint_artifact_sha256", None)
+                for dataset_identity in identity.get("datasets", []):
+                    content = dataset_identity.get("content", {})
+                    content.pop("sha256", None)
                 data_config = identity.get("data_config", {})
                 statistics_path = data_config.get("precomputed_statistics_path")
                 if statistics_path:
@@ -2739,7 +2764,7 @@ class LeRobotMixtureDataset(Dataset):
 
                     dataset, trajectory_id, step = self.sample_step(index)
                     # If dataset has no physical videos (e.g., image frames in parquet
-                    # for VLA-Arena), do not gate sampling on mp4 existence.
+                    # for VLA-Arena), sampling does not require mp4 files.
                     total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0))
                     if total_videos == 0:
                         break

starVLA/dataloader/gr00t_lerobot/transform/state_action.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #

starVLA/dataloader/infoflow_sampler.py

@@ -1,340 +1,6 @@
-"""Deterministic hierarchical sampling primitives for InfoFlow experiments."""
+"""Compatibility alias for the ReTime plugin; preserves legacy imports."""
 
-from __future__ import annotations
+from importlib import import_module
+import sys
 
-from dataclasses import asdict, dataclass
-import hashlib
-import json
-import os
-from pathlib import Path
-from typing import Any, Iterator, Sequence
-
-import numpy as np
-from torch.utils.data import Sampler
-
-
-def _identity_files(dataset_root: Path) -> list[Path]:
-    files = list((dataset_root / "data").rglob("*.parquet"))
-    files.extend((dataset_root / "videos").rglob("*.mp4"))
-    meta_root = dataset_root / "meta"
-    if meta_root.exists():
-        for pattern in ("*.json", "*.jsonl", "*.parquet"):
-            files.extend(meta_root.rglob(pattern))
-    return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix())
-
-
-def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]:
-    """Hash immutable LeRobot data, video, and task/episode metadata files once."""
-    root = Path(dataset_root)
-    digest = hashlib.sha256()
-    entries = []
-    total_bytes = 0
-    for file_path in _identity_files(root):
-        relative_path = file_path.relative_to(root).as_posix()
-        file_size = file_path.stat().st_size
-        entries.append({"path": relative_path, "bytes": file_size})
-        total_bytes += file_size
-        digest.update(relative_path.encode("utf-8"))
-        digest.update(str(file_size).encode("ascii"))
-        with open(file_path, "rb") as data_file:
-            for chunk in iter(lambda: data_file.read(4 * 1024 * 1024), b""):
-                digest.update(chunk)
-    return {
-        "sha256": digest.hexdigest(),
-        "file_count": len(entries),
-        "total_bytes": total_bytes,
-        "files": entries,
-    }
-
-
-def validate_dataset_content_inventory(
-    dataset_root: str | Path, fingerprint: dict[str, Any]
-) -> None:
-    """Cheap per-rank check; full content hashing belongs to fingerprint generation."""
-    root = Path(dataset_root)
-    current_entries = [
-        {
-            "path": path.relative_to(root).as_posix(),
-            "bytes": path.stat().st_size,
-        }
-        for path in _identity_files(root)
-    ]
-    if current_entries != fingerprint.get("files"):
-        raise RuntimeError(
-            f"Dataset file inventory differs from its content fingerprint: {root}"
-        )
-    if len(current_entries) != fingerprint.get("file_count"):
-        raise RuntimeError("Dataset content fingerprint has an inconsistent file count")
-    if sum(entry["bytes"] for entry in current_entries) != fingerprint.get("total_bytes"):
-        raise RuntimeError("Dataset content fingerprint has an inconsistent byte count")
-
-
-def _normalize_probabilities(values: Sequence[float], name: str) -> tuple[float, ...]:
-    probabilities = np.asarray(values, dtype=np.float64)
-    if probabilities.ndim != 1 or probabilities.size == 0:
-        raise ValueError(f"{name} must be a non-empty one-dimensional sequence")
-    if not np.all(np.isfinite(probabilities)) or np.any(probabilities <= 0):
-        raise ValueError(f"{name} must contain finite, strictly positive values")
-    probabilities /= probabilities.sum(dtype=np.float64)
-    return tuple(float(value) for value in probabilities)
-
-
-def _stateless_seed(seed: int, epoch: int, global_draw_id: int) -> int:
-    payload = f"{int(seed)}:{int(epoch)}:{int(global_draw_id)}".encode("utf-8")
-    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little", signed=False)
-
-
-@dataclass(frozen=True)
-class EpisodeSupport:
-    trajectory_index: int
-    trajectory_id: int | str
-    anchors: tuple[int, ...]
-
-    def __post_init__(self) -> None:
-        if not self.anchors:
-            raise ValueError("episode support must contain at least one anchor")
-        if len(set(self.anchors)) != len(self.anchors):
-            raise ValueError("episode anchors must be unique")
-
-
-@dataclass(frozen=True)
-class DatasetSupport:
-    name: str
-    probability: float
-    episodes: tuple[EpisodeSupport, ...]
-    episode_probabilities: tuple[float, ...]
-
-    def __post_init__(self) -> None:
-        if not self.episodes:
-            raise ValueError(f"dataset {self.name!r} has no valid episodes")
-        if len(self.episodes) != len(self.episode_probabilities):
-            raise ValueError("episode support and probability lengths differ")
-        normalized = _normalize_probabilities(self.episode_probabilities, "episode_probabilities")
-        if not np.allclose(normalized, self.episode_probabilities, rtol=0, atol=1e-12):
-            raise ValueError("episode probabilities must already sum to one")
-
-
-@dataclass(frozen=True)
-class LogicalDraw:
-    epoch: int
-    global_draw_id: int
-
-    def __post_init__(self) -> None:
-        if self.epoch < 0 or self.global_draw_id < 0:
-            raise ValueError("epoch and global_draw_id must be non-negative")
-
-
-@dataclass(frozen=True)
-class SampleTrace:
-    manifest_sha256: str
-    epoch: int
-    global_draw_id: int
-    dataset_index: int
-    dataset_name: str
-    trajectory_index: int
-    trajectory_id: int | str
-    anchor_index: int
-    base_index: int
-    baseline_probability: float
-    proposal_sha256: str | None = None
-    proposal_probability: float | None = None
-    importance_weight: float = 1.0
-
-    def to_dict(self) -> dict[str, Any]:
-        return asdict(self)
-
-
-class ValidAnchorManifest:
-    schema_version = 1
-
-    def __init__(self, datasets: Sequence[DatasetSupport], identity: dict[str, Any] | None = None):
-        if not datasets:
-            raise ValueError("manifest must contain at least one dataset")
-        dataset_probabilities = _normalize_probabilities(
-            [dataset.probability for dataset in datasets], "dataset_probabilities"
-        )
-        if not np.allclose(
-            dataset_probabilities,
-            [dataset.probability for dataset in datasets],
-            rtol=0,
-            atol=1e-12,
-        ):
-            raise ValueError("dataset probabilities must already sum to one")
-        self.datasets = tuple(datasets)
-        self.identity = identity or {}
-        self._dataset_probabilities = np.asarray(
-            [dataset.probability for dataset in self.datasets], dtype=np.float64
-        )
-        self._episode_probabilities = tuple(
-            np.asarray(dataset.episode_probabilities, dtype=np.float64)
-            for dataset in self.datasets
-        )
-        canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":"))
-        self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
-
-    @classmethod
-    def from_hierarchy(
-        cls,
-        dataset_names: Sequence[str],
-        dataset_weights: Sequence[float],
-        episodes: Sequence[Sequence[tuple[int | str, Sequence[int], float]]],
-        identity: dict[str, Any] | None = None,
-    ) -> "ValidAnchorManifest":
-        if len(dataset_names) != len(dataset_weights) or len(dataset_names) != len(episodes):
-            raise ValueError("dataset names, weights, and episode supports must have equal lengths")
-        dataset_probabilities = _normalize_probabilities(dataset_weights, "dataset_weights")
-        supports = []
-        for dataset_index, (name, dataset_episodes) in enumerate(zip(dataset_names, episodes)):
-            if not dataset_episodes:
-                raise ValueError(f"dataset {name!r} has no valid episodes")
-            episode_probabilities = _normalize_probabilities(
-                [entry[2] for entry in dataset_episodes], f"episode_weights[{dataset_index}]"
-            )
-            episode_supports = tuple(
-                EpisodeSupport(
-                    trajectory_index=trajectory_index,
-                    trajectory_id=trajectory_id,
-                    anchors=tuple(int(anchor) for anchor in anchors),
-                )
-                for trajectory_index, (trajectory_id, anchors, _) in enumerate(dataset_episodes)
-            )
-            supports.append(
-                DatasetSupport(
-                    name=str(name),
-                    probability=dataset_probabilities[dataset_index],
-                    episodes=episode_supports,
-                    episode_probabilities=episode_probabilities,
-                )
-            )
-        return cls(supports, identity=identity)
-
-    @classmethod
-    def from_dict(cls, payload: dict[str, Any]) -> "ValidAnchorManifest":
-        datasets = []
-        for dataset_payload in payload["datasets"]:
-            episodes = tuple(
-                EpisodeSupport(
-                    trajectory_index=int(episode["trajectory_index"]),
-                    trajectory_id=episode["trajectory_id"],
-                    anchors=tuple(int(anchor) for anchor in episode["anchors"]),
-                )
-                for episode in dataset_payload["episodes"]
-            )
-            datasets.append(
-                DatasetSupport(
-                    name=str(dataset_payload["name"]),
-                    probability=float(dataset_payload["probability"]),
-                    episodes=episodes,
-                    episode_probabilities=tuple(
-                        float(value) for value in dataset_payload["episode_probabilities"]
-                    ),
-                )
-            )
-        manifest = cls(datasets, identity=payload.get("identity", {}))
-        return manifest
-
-    @classmethod
-    def load(cls, path: str | Path) -> "ValidAnchorManifest":
-        with open(path, "r", encoding="utf-8") as manifest_file:
-            return cls.from_dict(json.load(manifest_file))
-
-    def save(self, path: str | Path) -> None:
-        destination = Path(path)
-        destination.parent.mkdir(parents=True, exist_ok=True)
-        temporary = destination.with_suffix(destination.suffix + ".tmp")
-        with open(temporary, "w", encoding="utf-8") as manifest_file:
-            json.dump(self.to_dict(), manifest_file, indent=2)
-        os.replace(temporary, destination)
-
-    def sample(self, seed: int, draw: LogicalDraw) -> SampleTrace:
-        rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id))
-        dataset_index = int(rng.choice(len(self.datasets), p=self._dataset_probabilities))
-        dataset = self.datasets[dataset_index]
-        trajectory_index = int(
-            rng.choice(len(dataset.episodes), p=self._episode_probabilities[dataset_index])
-        )
-        episode = dataset.episodes[trajectory_index]
-        anchor_index = int(rng.integers(len(episode.anchors)))
-        baseline_probability = (
-            dataset.probability
-            * dataset.episode_probabilities[trajectory_index]
-            / len(episode.anchors)
-        )
-        return SampleTrace(
-            manifest_sha256=self.sha256,
-            epoch=int(draw.epoch),
-            global_draw_id=int(draw.global_draw_id),
-            dataset_index=dataset_index,
-            dataset_name=dataset.name,
-            trajectory_index=episode.trajectory_index,
-            trajectory_id=episode.trajectory_id,
-            anchor_index=anchor_index,
-            base_index=episode.anchors[anchor_index],
-            baseline_probability=float(baseline_probability),
-        )
-
-    def to_dict(self, include_hash: bool = True) -> dict[str, Any]:
-        payload = {
-            "schema_version": self.schema_version,
-            "identity": self.identity,
-            "datasets": [asdict(dataset) for dataset in self.datasets],
-        }
-        if include_hash:
-            payload["sha256"] = self.sha256
-        return payload
-
-
-class StatelessDrawSampler(Sampler[LogicalDraw]):
-    """Yield rank-disjoint logical draw ids; sampling itself stays inside the dataset."""
-
-    def __init__(
-        self,
-        num_local_draws: int,
-        rank: int = 0,
-        world_size: int = 1,
-        start_global_draw_id: int = 0,
-    ) -> None:
-        if num_local_draws < 0:
-            raise ValueError("num_local_draws must be non-negative")
-        if world_size <= 0 or rank < 0 or rank >= world_size:
-            raise ValueError("rank must satisfy 0 <= rank < world_size")
-        self.num_local_draws = int(num_local_draws)
-        self.rank = int(rank)
-        self.world_size = int(world_size)
-        self.start_global_draw_id = int(start_global_draw_id)
-        if self.start_global_draw_id < 0:
-            raise ValueError("start_global_draw_id must be non-negative")
-        self.epoch = 0
-        self._resume_epoch = 0
-
-    def set_resume_position(self, epoch: int, start_global_draw_id: int) -> None:
-        """Resume at an inclusive global draw cursor within one epoch."""
-        if epoch < 0 or start_global_draw_id < 0:
-            raise ValueError("resume epoch and draw id must be non-negative")
-        self.epoch = int(epoch)
-        self._resume_epoch = int(epoch)
-        self.start_global_draw_id = int(start_global_draw_id)
-
-    def set_epoch(self, epoch: int) -> None:
-        if epoch < 0:
-            raise ValueError("epoch must be non-negative")
-        self.epoch = int(epoch)
-        if self.epoch != self._resume_epoch:
-            self.start_global_draw_id = 0
-
-    def _first_local_draw_id(self) -> int:
-        return max(
-            0,
-            (self.start_global_draw_id - self.rank + self.world_size - 1)
-            // self.world_size,
-        )
-
-    def __iter__(self) -> Iterator[LogicalDraw]:
-        for local_draw_id in range(self._first_local_draw_id(), self.num_local_draws):
-            yield LogicalDraw(
-                epoch=self.epoch,
-                global_draw_id=local_draw_id * self.world_size + self.rank,
-            )
-
-    def __len__(self) -> int:
-        return max(0, self.num_local_draws - self._first_local_draw_id())
+sys.modules[__name__] = import_module('retime.sampler')

starVLA/dataloader/lerobot_datasets.py

@@ -1,7 +1,6 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved.
-# Modified by [Fangjing Wang/ SUST University] in [2025]. 
 # Modification: [return raw data and suport multi-dataset mixture].
-# Modified by [Jinhui YE/ HKUST University] in [2025]. 
 # Modification: [suport topdowm processing, suport param from config].
 
 import logging
@@ -160,4 +159,4 @@ if __name__ == "__main__":
         if count > 3:
             break
         count += 1
-        pass
+        pass

starVLA/dataloader/relative_pose.py

@@ -0,0 +1,96 @@
+"""SE(3) relative pose helpers for wxyz quaternions.
+
+``action_mode: rel`` must not subtract quaternion components.  Positions are
+``xyz_a - xyz_s0``; rotations are ``q_rel = inv(q_s0) ⊗ q_a``.  The relative
+quaternion is still a unit quaternion, so later rotation-6D conversion stays
+valid with ``absolute: true`` in modality metadata.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+_QUATERNION_DIM = 4
+
+
+def _as_float(values: np.ndarray) -> np.ndarray:
+    return np.asarray(values, dtype=np.float32)
+
+
+def quat_conjugate(quaternion: np.ndarray) -> np.ndarray:
+    """Inverse of a unit wxyz quaternion (conjugate)."""
+    conjugated = np.array(_as_float(quaternion), copy=True)
+    conjugated[..., 1:] *= -1
+    return conjugated
+
+
+def quat_multiply(left: np.ndarray, right: np.ndarray) -> np.ndarray:
+    """Hamilton product of wxyz quaternions, broadcasting on leading axes."""
+    left_q, right_q = np.broadcast_arrays(_as_float(left), _as_float(right))
+    w1, x1, y1, z1 = left_q[..., 0], left_q[..., 1], left_q[..., 2], left_q[..., 3]
+    w2, x2, y2, z2 = right_q[..., 0], right_q[..., 1], right_q[..., 2], right_q[..., 3]
+    product = np.stack(
+        [
+            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        ],
+        axis=-1,
+    )
+    return _normalize_quaternion(product)
+
+
+def _normalize_quaternion(quaternion: np.ndarray) -> np.ndarray:
+    norms = np.linalg.norm(quaternion, axis=-1, keepdims=True)
+    return quaternion / np.clip(norms, 1e-8, None)
+
+
+def relative_quaternion(state_quaternion: np.ndarray, action_quaternion: np.ndarray) -> np.ndarray:
+    """Rotation taking ``state`` onto ``action``: ``inv(q_s) ⊗ q_a``."""
+    return quat_multiply(quat_conjugate(state_quaternion), action_quaternion)
+
+
+def compose_quaternion(state_quaternion: np.ndarray, relative_quaternion_value: np.ndarray) -> np.ndarray:
+    """Apply a relative rotation: ``q_s ⊗ q_rel``."""
+    return quat_multiply(state_quaternion, relative_quaternion_value)
+
+
+def relative_values(
+    action_values: np.ndarray,
+    state0: np.ndarray,
+    *,
+    is_quaternion: bool,
+) -> np.ndarray:
+    """Convert one action block into the ``rel`` representation."""
+    action = _as_float(action_values)
+    origin = _as_float(state0)
+    if is_quaternion:
+        if action.shape[-1] != _QUATERNION_DIM or origin.shape[-1] != _QUATERNION_DIM:
+            raise ValueError(
+                f"quaternion rel requires last dim {_QUATERNION_DIM}, "
+                f"got action {action.shape} state {origin.shape}"
+            )
+        return relative_quaternion(origin, action)
+    return action - origin
+
+
+def compose_values(
+    relative_action: np.ndarray,
+    state0: np.ndarray,
+    *,
+    is_quaternion: bool,
+) -> np.ndarray:
+    """Invert :func:`relative_values`."""
+    relative = _as_float(relative_action)
+    origin = _as_float(state0)
+    if is_quaternion:
+        return compose_quaternion(origin, relative)
+    return relative + origin
+
+
+def rotation_type_is_quaternion(rotation_type: object | None) -> bool:
+    if rotation_type is None:
+        return False
+    value = rotation_type.value if hasattr(rotation_type, "value") else str(rotation_type)
+    return value == "quaternion"

starVLA/dataloader/temporal_proposal.py

@@ -1,229 +1,6 @@
-"""Offline temporal proposals that preserve StarVLA's dataset/episode marginals."""
+"""Compatibility alias for the ReTime plugin; preserves legacy imports."""
 
-from __future__ import annotations
+from importlib import import_module
+import sys
 
-import hashlib
-import json
-import os
-from pathlib import Path
-from typing import Any, Sequence
-
-import numpy as np
-
-from starVLA.dataloader.infoflow_sampler import (
-    LogicalDraw,
-    SampleTrace,
-    ValidAnchorManifest,
-    _stateless_seed,
-)
-
-
-class TemporalProposal:
-    schema_version = 1
-    valid_score_kinds = {
-        "synthetic",
-        "random",
-        "action_change",
-        "residual",
-        "likelihood",
-    }
-
-    def __init__(
-        self,
-        manifest: ValidAnchorManifest,
-        alpha: float,
-        anchor_probabilities: Sequence[Sequence[Sequence[float]]],
-        score_kind: str,
-        scorer_sha256: str,
-        validation_certificate: dict[str, Any] | None = None,
-        legacy_implementation_sha256: str | None = None,
-    ) -> None:
-        if not 0.0 <= alpha < 1.0:
-            raise ValueError("alpha must satisfy 0 <= alpha < 1")
-        if score_kind not in self.valid_score_kinds:
-            raise ValueError(f"unsupported score_kind: {score_kind}")
-        if score_kind == "likelihood" and not validation_certificate:
-            raise ValueError("likelihood proposals require a validation certificate")
-        if not scorer_sha256:
-            raise ValueError("scorer_sha256 is required")
-        if len(anchor_probabilities) != len(manifest.datasets):
-            raise ValueError("proposal dataset count does not match the manifest")
-
-        normalized_probabilities = []
-        importance_weights = []
-        normalized_masses = []
-        for dataset_index, (dataset, dataset_probabilities) in enumerate(
-            zip(manifest.datasets, anchor_probabilities)
-        ):
-            if len(dataset_probabilities) != len(dataset.episodes):
-                raise ValueError(
-                    f"proposal episode count differs for dataset {dataset_index}"
-                )
-            normalized_dataset = []
-            dataset_weights = []
-            normalized_dataset_masses = []
-            for episode_index, (episode, proposal_mass) in enumerate(
-                zip(dataset.episodes, dataset_probabilities)
-            ):
-                mass = np.asarray(proposal_mass, dtype=np.float64)
-                if mass.shape != (len(episode.anchors),):
-                    raise ValueError(
-                        f"proposal anchor count differs for dataset {dataset_index}, "
-                        f"episode {episode_index}"
-                    )
-                if not np.all(np.isfinite(mass)) or np.any(mass < 0) or mass.sum() <= 0:
-                    raise ValueError("proposal masses must be finite, non-negative, and non-zero")
-                residual_probability = mass / mass.sum(dtype=np.float64)
-                baseline_probability = np.full(len(mass), 1.0 / len(mass), dtype=np.float64)
-                safe_probability = (1.0 - alpha) * baseline_probability + alpha * residual_probability
-                weights = baseline_probability / safe_probability
-                if not np.isclose(safe_probability.sum(dtype=np.float64), 1.0, atol=1e-12):
-                    raise ValueError("safe temporal proposal does not sum to one")
-                if np.max(weights) > 1.0 / (1.0 - alpha) + 1e-12:
-                    raise ValueError("importance weight exceeds the analytical safe-mixture bound")
-                normalized_dataset.append(safe_probability)
-                dataset_weights.append(weights)
-                normalized_dataset_masses.append(mass.copy())
-            normalized_probabilities.append(tuple(normalized_dataset))
-            importance_weights.append(tuple(dataset_weights))
-            normalized_masses.append(tuple(normalized_dataset_masses))
-
-        self.manifest_sha256 = manifest.sha256
-        self._manifest_support = manifest.to_dict(include_hash=False)["datasets"]
-        # Schema-v1 proposals included a source-file hash in their content hash.
-        # Preserve that value only when loading an existing artifact; it is not
-        # compared with the current source and new proposals do not emit it.
-        self.legacy_implementation_sha256 = legacy_implementation_sha256
-        self.alpha = float(alpha)
-        self.score_kind = score_kind
-        self.scorer_sha256 = scorer_sha256
-        self.validation_certificate = validation_certificate
-        self.proposal_masses = tuple(normalized_masses)
-        self.anchor_probabilities = tuple(normalized_probabilities)
-        self.importance_weights = tuple(importance_weights)
-        canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":"))
-        self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
-
-    @classmethod
-    def from_masses(
-        cls,
-        manifest: ValidAnchorManifest,
-        alpha: float,
-        masses: Sequence[Sequence[Sequence[float]]],
-        score_kind: str = "synthetic",
-        scorer_sha256: str = "synthetic-fixture",
-        validation_certificate: dict[str, Any] | None = None,
-    ) -> "TemporalProposal":
-        return cls(
-            manifest=manifest,
-            alpha=alpha,
-            anchor_probabilities=masses,
-            score_kind=score_kind,
-            scorer_sha256=scorer_sha256,
-            validation_certificate=validation_certificate,
-        )
-
-    @classmethod
-    def from_dict(
-        cls, manifest: ValidAnchorManifest, payload: dict[str, Any]
-    ) -> "TemporalProposal":
-        if payload.get("schema_version") != cls.schema_version:
-            raise ValueError(
-                f"unsupported temporal proposal schema: {payload.get('schema_version')}"
-            )
-        proposal = cls(
-            manifest=manifest,
-            alpha=float(payload["alpha"]),
-            anchor_probabilities=payload["proposal_masses"],
-            score_kind=str(payload["score_kind"]),
-            scorer_sha256=str(payload["scorer_sha256"]),
-            validation_certificate=payload.get("validation_certificate"),
-            legacy_implementation_sha256=payload.get("implementation_sha256"),
-        )
-        return proposal
-
-    @classmethod
-    def load(
-        cls, path: str | Path, manifest: ValidAnchorManifest
-    ) -> "TemporalProposal":
-        with open(path, "r", encoding="utf-8") as proposal_file:
-            return cls.from_dict(manifest, json.load(proposal_file))
-
-    def save(self, path: str | Path) -> None:
-        destination = Path(path)
-        destination.parent.mkdir(parents=True, exist_ok=True)
-        temporary = destination.with_suffix(destination.suffix + ".tmp")
-        with open(temporary, "w", encoding="utf-8") as proposal_file:
-            json.dump(self.to_dict(), proposal_file, indent=2)
-        os.replace(temporary, destination)
-
-    def sample(
-        self, manifest: ValidAnchorManifest, seed: int, draw: LogicalDraw
-    ) -> SampleTrace:
-        if manifest.to_dict(include_hash=False)["datasets"] != self._manifest_support:
-            raise ValueError("temporal proposal support differs from the active manifest")
-        rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id))
-        dataset_index = int(rng.choice(len(manifest.datasets), p=manifest._dataset_probabilities))
-        dataset = manifest.datasets[dataset_index]
-        trajectory_index = int(
-            rng.choice(
-                len(dataset.episodes), p=manifest._episode_probabilities[dataset_index]
-            )
-        )
-        episode = dataset.episodes[trajectory_index]
-        if self.alpha == 0.0:
-            # Preserve the exact baseline RNG path, not merely its distribution.
-            anchor_index = int(rng.integers(len(episode.anchors)))
-        else:
-            anchor_index = int(
-                rng.choice(
-                    len(episode.anchors),
-                    p=self.anchor_probabilities[dataset_index][trajectory_index],
-                )
-            )
-        baseline_probability = (
-            dataset.probability
-            * dataset.episode_probabilities[trajectory_index]
-            / len(episode.anchors)
-        )
-        proposal_probability = (
-            dataset.probability
-            * dataset.episode_probabilities[trajectory_index]
-            * self.anchor_probabilities[dataset_index][trajectory_index][anchor_index]
-        )
-        return SampleTrace(
-            manifest_sha256=manifest.sha256,
-            epoch=draw.epoch,
-            global_draw_id=draw.global_draw_id,
-            dataset_index=dataset_index,
-            dataset_name=dataset.name,
-            trajectory_index=episode.trajectory_index,
-            trajectory_id=episode.trajectory_id,
-            anchor_index=anchor_index,
-            base_index=episode.anchors[anchor_index],
-            baseline_probability=float(baseline_probability),
-            proposal_sha256=self.sha256,
-            proposal_probability=float(proposal_probability),
-            importance_weight=float(
-                self.importance_weights[dataset_index][trajectory_index][anchor_index]
-            ),
-        )
-
-    def to_dict(self, include_hash: bool = True) -> dict[str, Any]:
-        payload = {
-            "schema_version": self.schema_version,
-            "manifest_sha256": self.manifest_sha256,
-            "alpha": self.alpha,
-            "score_kind": self.score_kind,
-            "scorer_sha256": self.scorer_sha256,
-            "validation_certificate": self.validation_certificate,
-            "proposal_masses": [
-                [probabilities.tolist() for probabilities in dataset]
-                for dataset in self.proposal_masses
-            ],
-        }
-        if self.legacy_implementation_sha256 is not None:
-            payload["implementation_sha256"] = self.legacy_implementation_sha256
-        if include_hash:
-            payload["sha256"] = self.sha256
-        return payload
+sys.modules[__name__] = import_module('retime.proposal')

starVLA/model/framework/VLM4A/PI0.py

@@ -0,0 +1,653 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import math
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from starVLA.model.modules.vlm.openpi_transformers.gemma import modeling_gemma
+
+from deployment.model_server.tools.image_tools import resize_with_pad, to_pil_preserve
+from starVLA.model.framework.base_framework import baseframework
+from starVLA.model.framework.share_tools import merge_framework_config
+from starVLA.model.tools import FRAMEWORK_REGISTRY
+from starVLA.model.modules.action_model.OpenPI_ActionHead import (
+    GemmaDims,
+    OpenPI0ActionHead,
+    OpenPIActionHeadBase,
+    make_att_2d_masks,
+    make_att_4d_masks,
+)
+from starVLA.model.modules.vlm.OpenPIPaliGemma import OpenPIPaliGemma
+from starVLA.training.trainer_utils import initialize_overwatch
+
+logger = initialize_overwatch(__name__)
+
+IMAGE_KEYS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+
+
+@dataclass
+class Observation:
+    images: dict[str, torch.Tensor]
+    image_masks: dict[str, torch.Tensor]
+    tokenized_prompt: torch.Tensor
+    tokenized_prompt_mask: torch.Tensor
+    state: torch.Tensor
+
+
+@dataclass
+class PI0DefaultConfig:
+    name: str = "PI0"
+    precision: Literal["bfloat16", "float32"] = "bfloat16"
+
+    action_dim: int = 32
+    action_horizon: int = 50
+    max_state_dim: int = 32
+    discrete_state_input: bool = False
+    max_token_len: int = 200
+    num_inference_steps: int = 10
+
+    image_resolution: list[int] = field(default_factory=lambda: [224, 224])
+    image_keys: list[str] = field(default_factory=lambda: list(IMAGE_KEYS))
+
+    tokenizer: dict = field(default_factory=lambda: {"model_path": "", "pad_state_value": -2.0})
+    paligemma: dict = field(default_factory=lambda: {"model_name_or_path": ""})
+    paligemma_config: dict = field(
+        default_factory=lambda: {
+            "width": 2048,
+            "depth": 18,
+            "mlp_dim": 16384,
+            "num_heads": 8,
+            "num_kv_heads": 1,
+            "head_dim": 256,
+        }
+    )
+    action_expert_config: dict = field(
+        default_factory=lambda: {
+            "width": 1024,
+            "depth": 18,
+            "mlp_dim": 4096,
+            "num_heads": 8,
+            "num_kv_heads": 1,
+            "head_dim": 256,
+        }
+    )
+
+
+def _as_dims(config) -> GemmaDims:
+    return GemmaDims(
+        width=int(config.width),
+        depth=int(config.depth),
+        mlp_dim=int(config.mlp_dim),
+        num_heads=int(config.num_heads),
+        num_kv_heads=int(config.num_kv_heads),
+        head_dim=int(config.head_dim),
+    )
+
+
+class LazyPaliGemmaTokenizer:
+    def __init__(self, max_len: int, model_path: str = ""):
+        self.max_len = int(max_len)
+        self.model_path = str(model_path or "")
+        self._tokenizer = None
+
+    def _load(self):
+        if self._tokenizer is not None:
+            return self._tokenizer
+        if not self.model_path:
+            raise FileNotFoundError(
+                "PI0/PI05 tokenizer.model_path is empty. Set framework.tokenizer.model_path "
+                "to a local paligemma_tokenizer.model path for raw-text examples."
+            )
+        import sentencepiece
+
+        with open(self.model_path, "rb") as f:
+            self._tokenizer = sentencepiece.SentencePieceProcessor(model_proto=f.read())
+        return self._tokenizer
+
+    def tokenize(self, prompt: str, state: np.ndarray | torch.Tensor | None = None):
+        tokenizer = self._load()
+        cleaned_text = str(prompt).strip().replace("_", " ").replace("\n", " ")
+        if state is not None:
+            state_np = np.asarray(state, dtype=np.float32).reshape(-1)
+            discretized_state = np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
+            state_str = " ".join(map(str, discretized_state.tolist()))
+            prompt_text = f"Task: {cleaned_text}, State: {state_str};\nAction: "
+            tokens = tokenizer.encode(prompt_text, add_bos=True)
+        else:
+            tokens = tokenizer.encode(cleaned_text, add_bos=True) + tokenizer.encode("\n")
+
+        token_len = len(tokens)
+        if token_len < self.max_len:
+            padding = [0] * (self.max_len - token_len)
+            mask = [True] * token_len + [False] * len(padding)
+            tokens = tokens + padding
+        else:
+            if token_len > self.max_len:
+                logger.warning("PI tokenizer length %s exceeds max_token_len=%s; truncating", token_len, self.max_len)
+            tokens = tokens[: self.max_len]
+            mask = [True] * self.max_len
+        return np.asarray(tokens, dtype=np.int64), np.asarray(mask, dtype=bool)
+
+
+def _gated_residual(x: torch.Tensor, y: torch.Tensor, gate: torch.Tensor | None):
+    return x + y if gate is None else x + y * gate
+
+
+def _norm(layernorm, hidden_states: torch.Tensor, cond: torch.Tensor | None):
+    try:
+        out = layernorm(hidden_states, cond=cond)
+    except TypeError as exc:
+        if cond is not None:
+            dense = getattr(layernorm, "dense", None)
+            if dense is None:
+                raise RuntimeError(
+                    "PI05 requires Gemma adaRMS support, but this layernorm has no dense module."
+                ) from exc
+            dtype = hidden_states.dtype
+            eps = getattr(layernorm, "variance_epsilon", getattr(layernorm, "eps", 1e-6))
+            variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+            out = hidden_states * torch.rsqrt(variance + eps)
+            scale, shift, gate = dense(cond.to(dense.weight.dtype))[:, None, :].chunk(3, dim=-1)
+            return (out * (1 + scale) + shift).to(dtype), gate.to(dtype)
+        out = layernorm(hidden_states)
+    if isinstance(out, tuple):
+        return out
+    return out, None
+
+
+@FRAMEWORK_REGISTRY.register("PI0")
+@FRAMEWORK_REGISTRY.register("Pi0")
+class PI0(baseframework):
+    default_config_cls = PI0DefaultConfig
+    action_head_cls: type[OpenPIActionHeadBase] = OpenPI0ActionHead
+    model_name = "PI0"
+    discrete_state_input = False
+
+    def __init__(self, config=None, **kwargs):
+        super().__init__()
+        self._init_runtime_config(config)
+        self._build_tokenizer()
+        self._build_model_modules()
+        self.to_precision(self.config.framework.precision)
+        torch.set_float32_matmul_precision("highest")
+
+    def _init_runtime_config(self, config):
+        self.config = merge_framework_config(self.default_config_cls, config)
+        self.config.framework.name = self.model_name
+        self.discrete_state_input = bool(self.config.framework.discrete_state_input)
+
+        self.action_dim = int(self.config.framework.action_dim)
+        self.output_action_dim = int(self.config.framework.get("output_action_dim", self.action_dim))
+        if not 0 < self.output_action_dim <= self.action_dim:
+            raise ValueError("output_action_dim must be between 1 and native action_dim")
+        self.action_horizon = int(self.config.framework.action_horizon)
+        self.max_state_dim = int(self.config.framework.max_state_dim)
+        self.num_inference_steps = int(self.config.framework.num_inference_steps)
+        self.image_resolution = tuple(int(x) for x in self.config.framework.image_resolution)
+        self.image_keys = tuple(self.config.framework.image_keys)
+        self.state_pad_value = float(self.config.framework.tokenizer.pad_state_value)
+
+    def _build_tokenizer(self):
+        self.tokenizer = LazyPaliGemmaTokenizer(
+            max_len=int(self.config.framework.max_token_len),
+            model_path=str(self.config.framework.tokenizer.model_path),
+        )
+
+    def _build_model_modules(self):
+        vlm_config = _as_dims(self.config.framework.paligemma_config)
+        expert_config = _as_dims(self.config.framework.action_expert_config)
+        self.vlm = OpenPIPaliGemma(
+            vlm_config,
+            use_adarms=False,
+            model_name_or_path=str(self.config.framework.paligemma.model_name_or_path),
+        )
+        self.action_head = self.action_head_cls(
+            expert_config,
+            action_dim=self.action_dim,
+            action_horizon=self.action_horizon,
+            max_state_dim=self.max_state_dim,
+        )
+
+    def to_precision(self, precision: Literal["bfloat16", "float32"]):
+        if precision == "float32":
+            self.to(dtype=torch.float32)
+            return
+        if precision != "bfloat16":
+            raise ValueError(f"Unsupported PI precision: {precision}")
+        self.to(dtype=torch.bfloat16)
+        keep_fp32 = (
+            "vision_tower.vision_model.embeddings.patch_embedding.weight",
+            "vision_tower.vision_model.embeddings.patch_embedding.bias",
+            "vision_tower.vision_model.embeddings.position_embedding.weight",
+            "input_layernorm",
+            "post_attention_layernorm",
+            "model.norm",
+        )
+        for name, param in self.named_parameters():
+            if any(selector in name for selector in keep_fp32):
+                param.data = param.data.to(dtype=torch.float32)
+
+    def _pad_array_2d(self, arr: np.ndarray, target_time: int, target_dim: int, pad_value: float) -> np.ndarray:
+        arr = np.asarray(arr, dtype=np.float32)
+        if arr.ndim == 1:
+            arr = arr[None, :]
+        if arr.shape[0] > target_time or arr.shape[1] > target_dim:
+            raise ValueError(f"Input shape {arr.shape} exceeds target {(target_time, target_dim)}")
+        out = np.full((target_time, target_dim), pad_value, dtype=np.float32)
+        out[: arr.shape[0], : arr.shape[1]] = arr
+        return out
+
+    def _pil_to_tensor(self, image) -> torch.Tensor:
+        image_np = np.asarray(image.convert("RGB")).copy()
+        image_np = resize_with_pad(image_np[None], self.image_resolution[0], self.image_resolution[1])[0]
+        return torch.from_numpy(image_np).permute(2, 0, 1).contiguous().float() / 255.0 * 2.0 - 1.0
+
+    def _prepare_examples(self, examples: list[dict], include_actions: bool):
+        images, masks, states, actions, prompts = [], [], [], [], []
+        for example in examples:
+            raw_images = to_pil_preserve(example["image"])
+            if isinstance(raw_images, tuple):
+                raw_images = list(raw_images)
+            elif not isinstance(raw_images, list):
+                raw_images = [raw_images]
+            sample_images, sample_masks = [], []
+            provided_mask = example.get("image_mask")
+            provided_mask = list(provided_mask) if provided_mask is not None else None
+            for idx, _key in enumerate(self.image_keys):
+                has_image = idx < len(raw_images)
+                if has_image:
+                    image = raw_images[idx]
+                else:
+                    image = raw_images[0].copy()
+                    image.paste(0, (0, 0, image.size[0], image.size[1]))
+                active = bool(provided_mask[idx]) if provided_mask is not None and idx < len(provided_mask) else has_image
+                sample_images.append(self._pil_to_tensor(image))
+                sample_masks.append(active)
+            images.append(torch.stack(sample_images))
+            masks.append(sample_masks)
+            prompts.append(str(example.get("lang", example.get("language", ""))))
+
+            state_value = example.get("state")
+            if state_value is None:
+                fill = self.state_pad_value if self.config.framework.discrete_state_input else 0.0
+                state_np = np.full((1, self.max_state_dim), fill, dtype=np.float32)
+            else:
+                state_np = np.asarray(state_value, dtype=np.float32)
+                if state_np.ndim == 2 and state_np.shape[0] > 1:
+                    state_np = state_np[-1:]
+                state_np = self._pad_array_2d(
+                    state_np,
+                    1 if state_np.ndim == 1 else state_np.shape[0],
+                    self.max_state_dim,
+                    self.state_pad_value if self.config.framework.discrete_state_input else 0.0,
+                )
+            states.append(torch.from_numpy(state_np))
+            if include_actions:
+                action_np = self._pad_array_2d(np.asarray(example["action"], dtype=np.float32), self.action_horizon, self.action_dim, 0.0)
+                actions.append(torch.from_numpy(action_np))
+
+        batch = {
+            "image": torch.stack(images),
+            "image_mask": torch.as_tensor(np.asarray(masks), dtype=torch.bool),
+            "state": torch.stack(states),
+            "lang": prompts,
+        }
+        if include_actions:
+            batch["action"] = torch.stack(actions)
+        return batch
+
+    def _build_observation_from_batch(self, batch):
+        device = self.action_head.action_in_proj.weight.device
+        images = {key: batch["image"][:, idx].to(device=device, dtype=torch.float32) for idx, key in enumerate(self.image_keys)}
+        masks = {key: batch["image_mask"][:, idx].to(device=device, dtype=torch.bool) for idx, key in enumerate(self.image_keys)}
+        tokens, token_masks = [], []
+        for idx, prompt in enumerate(batch["lang"]):
+            token_state = batch["state"][idx].squeeze(0) if self.config.framework.discrete_state_input else None
+            token, mask = self.tokenizer.tokenize(prompt, token_state)
+            tokens.append(token)
+            token_masks.append(mask)
+        return Observation(
+            images=images,
+            image_masks=masks,
+            tokenized_prompt=torch.as_tensor(np.asarray(tokens), device=device),
+            tokenized_prompt_mask=torch.as_tensor(np.asarray(token_masks), device=device),
+            state=batch["state"].to(device=device, dtype=torch.float32),
+        )
+
+    def _preprocess_observation(self, observation: Observation):
+        return (
+            [observation.images[key] for key in self.image_keys],
+            [observation.image_masks[key] for key in self.image_keys],
+            observation.tokenized_prompt,
+            observation.tokenized_prompt_mask,
+            observation.state,
+        )
+
+    def embed_prefix(self, images, img_masks, lang_tokens, lang_masks):
+        embs, pad_masks, att_masks = [], [], []
+        for image, image_mask in zip(images, img_masks, strict=True):
+            image_emb = self.vlm.embed_image(image)
+            bsize, num_img_embs = image_emb.shape[:2]
+            embs.append(image_emb)
+            pad_masks.append(image_mask[:, None].expand(bsize, num_img_embs))
+            att_masks += [0] * num_img_embs
+        lang_emb = self.vlm.embed_language_tokens(lang_tokens)
+        embs.append(lang_emb * math.sqrt(lang_emb.shape[-1]))
+        pad_masks.append(lang_masks)
+        att_masks += [0] * lang_emb.shape[1]
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)[None, :]
+        return embs, pad_masks, att_masks.expand(pad_masks.shape[0], -1)
+
+    def forward_with_action_expert(
+        self,
+        *,
+        prefix_embs: torch.Tensor,
+        suffix_embs: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_ids: torch.LongTensor,
+        adarms_cond: torch.Tensor | None,
+    ):
+        streams = [prefix_embs, suffix_embs]
+        conds = [None, adarms_cond]
+        models = [self.vlm.language_model, self.action_head.model]
+        for layer_idx in range(self.vlm.config.text_config.num_hidden_layers):
+            streams = self.forward_shared_gemma_layer(layer_idx, models, streams, conds, attention_mask, position_ids)
+        suffix_out, _ = _norm(models[1].norm, streams[1], conds[1])
+        return suffix_out
+
+    def forward_shared_gemma_layer(self, layer_idx, models, streams, conds, attention_mask, position_ids):
+        query_states, key_states, value_states, gates = [], [], [], []
+        for idx, hidden_states in enumerate(streams):
+            layer = models[idx].layers[layer_idx]
+            hidden_states, gate = _norm(layer.input_layernorm, hidden_states, conds[idx])
+            gates.append(gate)
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+            query_states.append(layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2))
+            key_states.append(layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2))
+            value_states.append(layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2))
+
+        query_states = torch.cat(query_states, dim=2)
+        key_states = torch.cat(key_states, dim=2)
+        value_states = torch.cat(value_states, dim=2)
+        dummy = torch.zeros(
+            query_states.shape[0],
+            query_states.shape[2],
+            query_states.shape[-1],
+            device=query_states.device,
+            dtype=query_states.dtype,
+        )
+        cos, sin = self.vlm.language_model.rotary_emb(dummy, position_ids)
+        query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, unsqueeze_dim=1
+        )
+        att_output, _ = modeling_gemma.eager_attention_forward(
+            self.vlm.language_model.layers[layer_idx].self_attn,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            self.vlm.language_model.layers[layer_idx].self_attn.scaling,
+        )
+        head_dim = self.vlm.language_model.layers[layer_idx].self_attn.head_dim
+        num_heads = self.vlm.language_model.config.num_attention_heads
+        att_output = att_output.reshape(query_states.shape[0], -1, num_heads * head_dim)
+
+        next_streams = []
+        start = 0
+        for idx, hidden_states in enumerate(streams):
+            layer = models[idx].layers[layer_idx]
+            end = start + hidden_states.shape[1]
+            layer_att = att_output[:, start:end]
+            if layer_att.dtype != layer.self_attn.o_proj.weight.dtype:
+                layer_att = layer_att.to(layer.self_attn.o_proj.weight.dtype)
+            out = layer.self_attn.o_proj(layer_att)
+            out = _gated_residual(hidden_states, out, gates[idx])
+            residual = out.clone()
+            out, gate = _norm(layer.post_attention_layernorm, out, conds[idx])
+            if out.dtype != layer.mlp.up_proj.weight.dtype:
+                out = out.to(layer.mlp.up_proj.weight.dtype)
+            out = layer.mlp(out)
+            next_streams.append(_gated_residual(residual, out, gate))
+            start = end
+        return next_streams
+
+    def forward_from_observation(self, observation: Observation, actions: torch.Tensor, noise=None, time=None, return_debug=False, return_per_sample_loss=False, prefix_conditioning=None):
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation)
+        if prefix_conditioning is None:
+            prefix_conditioning = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+        prefix_embs, prefix_pad_masks, prefix_att_masks = prefix_conditioning
+        x_t, u_t, noise, time = self.action_head.prepare_flow_targets(actions, noise=noise, time=time)
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.action_head.embed_action_suffix(state.squeeze(1), x_t, time)
+        if self.vlm.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
+            prefix_embs = prefix_embs.to(torch.bfloat16)
+            suffix_embs = suffix_embs.to(torch.bfloat16)
+        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
+        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
+        attention_mask = make_att_4d_masks(make_att_2d_masks(pad_masks, att_masks))
+        position_ids = torch.cumsum(pad_masks, dim=1) - 1
+        suffix_out = self.forward_with_action_expert(
+            prefix_embs=prefix_embs,
+            suffix_embs=suffix_embs,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            adarms_cond=adarms_cond,
+        )
+        v_t = self.action_head.decode_action_velocity(suffix_out)
+        per_sample_loss = F.mse_loss(u_t, v_t, reduction="none").mean(dim=(1, 2))
+        out = {"v_t": v_t, "action_loss": per_sample_loss if return_per_sample_loss else per_sample_loss.mean()}
+        if return_debug:
+            out.update({"u_t": u_t, "x_t": x_t, "time": time, "noise": noise})
+        return out
+
+    def forward(self, examples: list[dict] = None, **kwargs):
+        '''
+        Training forward for PI0 flow-matching action prediction.
+        '''
+        batch = self._prepare_examples(examples, include_actions=True)
+        observation = self._build_observation_from_batch(batch)
+        actions = batch["action"].to(device=self.action_head.action_in_proj.weight.device, dtype=torch.float32)
+        noise = kwargs.get("noise")
+        time = kwargs.get("time")
+        seed = kwargs.get("flow_noise_seed")
+        if seed is not None and noise is None and time is None:
+            devices = [torch.cuda.current_device()] if torch.cuda.is_available() else []
+            with torch.random.fork_rng(devices=devices):
+                torch.manual_seed(int(seed))
+                noise = self.action_head.sample_action_noise(actions.shape, actions.device)
+                time = self.action_head.sample_flow_time(actions.shape[0], actions.device)
+        return self.forward_from_observation(
+            observation,
+            actions,
+            noise=noise,
+            time=time,
+            return_debug=bool(kwargs.get("return_debug", False)),
+            return_per_sample_loss=bool(kwargs.get("return_per_sample_loss", False)),
+        )
+
+    @torch.inference_mode()
+    def predict_action(self, examples: list[dict] | dict = None, **kwargs):
+        '''
+        Prediction forward for PI0 flow-matching action prediction
+        '''
+        if not isinstance(examples, list):
+            examples = [examples]
+        batch = self._prepare_examples(examples, include_actions=False)
+        observation = self._build_observation_from_batch(batch)
+        actions = self.sample_actions(
+            observation,
+            noise=kwargs.get("noise"),
+            num_steps=int(kwargs.get("num_steps", kwargs.get("num_inference_steps", self.num_inference_steps))),
+        )
+        return {"normalized_actions": actions[..., :self.output_action_dim].float().cpu().numpy()}
+
+    def sample_actions(self, observation: Observation, noise=None, num_steps: int = 10):
+        device = self.action_head.action_in_proj.weight.device
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation)
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+        prefix_att_2d_masks_4d = make_att_4d_masks(make_att_2d_masks(prefix_pad_masks, prefix_att_masks))
+        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
+        self.vlm.language_model.config._attn_implementation = "eager"
+        prefix_output = self.vlm.forward_prefix(
+            inputs_embeds=prefix_embs,
+            attention_mask=prefix_att_2d_masks_4d,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            use_cache=True,
+        )
+        past_key_values = prefix_output.past_key_values
+        bsize = state.shape[0]
+        shape = (bsize, self.action_horizon, self.action_dim)
+        x_t = self.action_head.sample_action_noise(shape, device) if noise is None else torch.as_tensor(noise, device=device, dtype=torch.float32)
+        if x_t.ndim == 2:
+            x_t = x_t[None]
+        dt = torch.tensor(-1.0 / max(num_steps, 1), dtype=torch.float32, device=device)
+        time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        while time >= -dt / 2:
+            v_t = self.denoise_step(state.squeeze(1), prefix_pad_masks, past_key_values, x_t, time.expand(bsize))
+            x_t = x_t + dt * v_t
+            time += dt
+        return x_t
+
+    def denoise_step(self, state, prefix_pad_masks, past_key_values, x_t, timestep):
+        suffix_embs, suffix_pad_masks, suffix_att_masks, _adarms_cond = self.action_head.embed_action_suffix(state, x_t, timestep)
+        suffix_len = suffix_pad_masks.shape[1]
+        bsize, prefix_len = prefix_pad_masks.shape
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(bsize, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        attention_mask = make_att_4d_masks(torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2))
+        position_ids = torch.sum(prefix_pad_masks, dim=-1)[:, None] + torch.cumsum(suffix_pad_masks, dim=1) - 1
+        suffix_embs = suffix_embs.to(torch.bfloat16) if self.vlm.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16 else suffix_embs
+        return self.action_head(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            suffix_embs=suffix_embs,
+            adarms_cond=_adarms_cond,
+        )
+
+    def _is_expected_openpi_missing_key(self, key: str) -> bool:
+        if not getattr(self.action_head, "use_adarms", False):
+            return False
+        prefix = "action_head.gemma_expert.model."
+        if key == prefix + "norm.weight":
+            return True
+        if not key.startswith(prefix + "layers."):
+            return False
+        return key.endswith(".input_layernorm.weight") or key.endswith(".post_attention_layernorm.weight")
+
+    @staticmethod
+    def _resolve_checkpoint_path(checkpoint_path: str | Path) -> Path:
+        path = Path(checkpoint_path)
+        if path.is_dir():
+            for name in ("model.safetensors", "pytorch_model.bin", "model.pt", "checkpoint.pt", "ckpt.pt"):
+                candidate = path / name
+                if candidate.exists():
+                    return candidate
+            raise FileNotFoundError(f"No supported PI checkpoint file found under {path}")
+        return path
+
+    @classmethod
+    def _load_checkpoint_state_dict(cls, checkpoint_path: str | Path):
+        checkpoint_path = cls._resolve_checkpoint_path(checkpoint_path)
+        if str(checkpoint_path).endswith(".safetensors"):
+            from safetensors.torch import load_file
+
+            state_dict = load_file(str(checkpoint_path))
+        else:
+            state_dict = torch.load(checkpoint_path, map_location="cpu")
+            if isinstance(state_dict, dict):
+                for key in ("state_dict", "model_state_dict", "module", "model"):
+                    if key in state_dict and isinstance(state_dict[key], dict):
+                        state_dict = state_dict[key]
+                        break
+        return state_dict
+
+    def _normalize_state_dict(self, state_dict: dict):
+        current = self.state_dict()
+        normalized = {}
+        for key, value in state_dict.items():
+            clean_key = key
+            changed = True
+            while changed:
+                changed = False
+                for prefix in ("model.", "module.", "_orig_mod."):
+                    if clean_key.startswith(prefix):
+                        clean_key = clean_key[len(prefix):]
+                        changed = True
+            if clean_key.startswith("paligemma_with_expert.paligemma.model."):
+                clean_key = clean_key.replace(
+                    "paligemma_with_expert.paligemma.model.",
+                    "vlm.model.model.",
+                    1,
+                )
+            elif clean_key.startswith("paligemma_with_expert.paligemma."):
+                clean_key = clean_key.replace("paligemma_with_expert.paligemma.", "vlm.model.", 1)
+            elif clean_key.startswith("paligemma_with_expert.gemma_expert."):
+                clean_key = clean_key.replace("paligemma_with_expert.gemma_expert.", "action_head.gemma_expert.", 1)
+            elif clean_key.startswith("action_in_proj."):
+                clean_key = clean_key.replace("action_in_proj.", "action_head.action_in_proj.", 1)
+            elif clean_key.startswith("action_out_proj."):
+                clean_key = clean_key.replace("action_out_proj.", "action_head.action_out_proj.", 1)
+            elif clean_key.startswith("state_proj."):
+                clean_key = clean_key.replace("state_proj.", "action_head.state_proj.", 1)
+            elif clean_key.startswith("action_time_mlp_in."):
+                clean_key = clean_key.replace("action_time_mlp_in.", "action_head.action_time_mlp_in.", 1)
+            elif clean_key.startswith("action_time_mlp_out."):
+                clean_key = clean_key.replace("action_time_mlp_out.", "action_head.action_time_mlp_out.", 1)
+            elif clean_key.startswith("time_mlp_in."):
+                clean_key = clean_key.replace("time_mlp_in.", "action_head.time_mlp_in.", 1)
+            elif clean_key.startswith("time_mlp_out."):
+                clean_key = clean_key.replace("time_mlp_out.", "action_head.time_mlp_out.", 1)
+            normalized[clean_key] = value
+        embed_key = "vlm.model.model.language_model.embed_tokens.weight"
+        lm_head_key = "vlm.model.lm_head.weight"
+        if embed_key not in normalized and lm_head_key in normalized and embed_key in current:
+            if normalized[lm_head_key].shape == current[embed_key].shape:
+                normalized[embed_key] = normalized[lm_head_key]
+        return normalized
+
+    def load_pretrained_checkpoint(self, checkpoint_path: str | Path, reload_modules: str | None = None):
+        state_dict = self._normalize_state_dict(self._load_checkpoint_state_dict(checkpoint_path))
+        if reload_modules:
+            for path in [p.strip() for p in reload_modules.split(",") if p.strip()]:
+                module = self
+                for part in path.split("."):
+                    module = getattr(module, part)
+                prefix = path + "."
+                sub_state = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)}
+                module.load_state_dict(sub_state, strict=False)
+            return self
+
+        current = self.state_dict()
+        filtered, shape_mismatch = {}, []
+        for key, value in state_dict.items():
+            if key in current and getattr(value, "shape", None) == current[key].shape:
+                filtered[key] = value
+            elif key in current:
+                shape_mismatch.append((key, tuple(value.shape), tuple(current[key].shape)))
+        missing, unexpected = self.load_state_dict(filtered, strict=False)
+        expected_missing = [key for key in missing if self._is_expected_openpi_missing_key(key)]
+        real_missing = [key for key in missing if key not in expected_missing]
+        logger.info(
+            "%s checkpoint load: checkpoint=%s loaded=%s model=%s missing=%s expected_missing=%s unexpected=%s shape_mismatch=%s",
+            type(self).__name__,
+            len(state_dict),
+            len(filtered),
+            len(current),
+            len(real_missing),
+            len(expected_missing),
+            len(unexpected),
+            len(shape_mismatch),
+        )
+        if real_missing:
+            logger.info("%s first missing keys: %s", type(self).__name__, real_missing[:10])
+        if shape_mismatch:
+            logger.info("%s first shape mismatches: %s", type(self).__name__, shape_mismatch[:10])
+        return self

starVLA/model/framework/VLM4A/PI05.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+from starVLA.model.framework.VLM4A.PI0 import PI0, PI0DefaultConfig, IMAGE_KEYS
+from starVLA.model.modules.action_model.OpenPI_ActionHead import OpenPI05ActionHead
+from starVLA.model.tools import FRAMEWORK_REGISTRY
+
+
+@dataclass
+class PI05DefaultConfig(PI0DefaultConfig):
+    name: str = "PI05"
+    precision: Literal["bfloat16", "float32"] = "bfloat16"
+
+    action_dim: int = 32
+    action_horizon: int = 50
+    max_state_dim: int = 32
+    discrete_state_input: bool = True
+    max_token_len: int = 200
+    num_inference_steps: int = 10
+
+    image_resolution: list[int] = field(default_factory=lambda: [224, 224])
+    image_keys: list[str] = field(default_factory=lambda: list(IMAGE_KEYS))
+
+    tokenizer: dict = field(default_factory=lambda: {"model_path": "", "pad_state_value": -2.0})
+    paligemma: dict = field(default_factory=lambda: {"model_name_or_path": ""})
+    paligemma_config: dict = field(
+        default_factory=lambda: {
+            "width": 2048,
+            "depth": 18,
+            "mlp_dim": 16384,
+            "num_heads": 8,
+            "num_kv_heads": 1,
+            "head_dim": 256,
+        }
+    )
+    action_expert_config: dict = field(
+        default_factory=lambda: {
+            "width": 1024,
+            "depth": 18,
+            "mlp_dim": 4096,
+            "num_heads": 8,
+            "num_kv_heads": 1,
+            "head_dim": 256,
+        }
+    )
+
+
+@FRAMEWORK_REGISTRY.register("PI05")
+@FRAMEWORK_REGISTRY.register("Pi05")
+class PI05(PI0):
+    """OpenPI PI05 model exposed through the StarVLA framework API."""
+
+    default_config_cls = PI05DefaultConfig
+    action_head_cls = OpenPI05ActionHead
+    model_name = "PI05"
+    discrete_state_input = True

starVLA/model/framework/VLM4A/QwenFast.py

@@ -0,0 +1,174 @@
+"""Qwen2.5-VL + FAST autoregressive action policy for ReTime.
+
+The action tokens and checkpoint follow StarVLA's QwenFast implementation.
+The per-example action-token NLL is exposed for temporal importance sampling.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+from starVLA.model.framework.base_framework import baseframework
+from starVLA.model.tools import FRAMEWORK_REGISTRY
+
+
+class _Qwen25Interface(nn.Module):
+    def __init__(self, model_path: str, *, attn_implementation: str = "sdpa"):
+        super().__init__()
+        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_path, attn_implementation=attn_implementation, dtype=torch.bfloat16
+        )
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.processor.tokenizer.padding_side = "left"
+        self.action_token_min = self.processor.tokenizer.convert_tokens_to_ids("<robot_action_0>")
+        self.action_token_max = self.processor.tokenizer.convert_tokens_to_ids("<robot_action_2047>")
+        if (
+            self.action_token_min < 0
+            or self.action_token_max - self.action_token_min != 2047
+            or self.action_token_max >= self.model.config.text_config.vocab_size
+        ):
+            raise ValueError("QwenFast requires 2048 contiguous <robot_action_i> tokens")
+
+
+class _FastActionEncoder(nn.Module):
+    def __init__(self, path: str, horizon: int, action_dim: int):
+        super().__init__()
+        self.fast_tokenizer = AutoProcessor.from_pretrained(path, trust_remote_code=True)
+        self.fast_tokenizer.time_horizon = horizon
+        self.fast_tokenizer.action_dim = action_dim
+
+    def encode(self, actions: np.ndarray) -> list[list[int]]:
+        tokens = self.fast_tokenizer(actions)
+        if any(not seq or min(seq) < 0 or max(seq) >= 2048 for seq in tokens):
+            raise ValueError("FAST produced an empty or out-of-range action sequence")
+        return tokens
+
+
+@FRAMEWORK_REGISTRY.register("QwenFast")
+class QwenFast(baseframework):
+    """Autoregressive action model compatible with StarVLA's 10K Bridge state."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        framework = config.framework
+        action = framework.action_model
+        self.action_horizon = int(action.action_horizon)
+        self.action_dim = int(action.action_dim)
+        self.qwen_vl_interface = _Qwen25Interface(
+            str(framework.qwenvl.base_vlm),
+            attn_implementation=str(framework.qwenvl.get("attn_implementation", "sdpa")),
+        )
+        self.action_model = _FastActionEncoder(
+            str(action.fast_tokenizer_path), self.action_horizon, self.action_dim
+        )
+        self.cot_prompt = config.datasets.vla_data.get("CoT_prompt", None)
+
+    def _inputs(self, examples: list[dict], *, with_actions: bool):
+        if not examples:
+            raise ValueError("examples must be non-empty")
+        actions = None
+        if with_actions:
+            actions = np.stack(
+                [np.asarray(e["action"], dtype=np.float32)[-self.action_horizon :] for e in examples]
+            )
+            if actions.shape[1:] != (self.action_horizon, self.action_dim):
+                raise ValueError("actions must have the configured horizon and dimension")
+            tokens = self.action_model.encode(actions)
+        messages = []
+        for i, example in enumerate(examples):
+            images = example["image"]
+            if isinstance(images, (Image.Image, np.ndarray)):
+                images = [images]
+            if not images:
+                raise ValueError("each example requires at least one image")
+            prepared_images = []
+            for image in images:
+                if isinstance(image, np.ndarray):
+                    if image.ndim != 3 or image.shape[-1] != 3 or image.dtype != np.uint8:
+                        raise ValueError("NumPy images must be HWC RGB uint8")
+                    image = Image.fromarray(image)
+                if not isinstance(image, Image.Image):
+                    raise ValueError("Images must be PIL or HWC RGB uint8 arrays")
+                prepared_images.append(image)
+            instruction = str(example["lang"])
+            if self.cot_prompt:
+                instruction = self.cot_prompt.replace("{instruction}", instruction)
+            content = [{"type": "image", "image": image} for image in prepared_images]
+            content.append({"type": "text", "text": instruction})
+            message = [{"role": "user", "content": content}]
+            if with_actions:
+                solution = "".join(f"<robot_action_{token}>" for token in tokens[i])
+                message.append({"role": "assistant", "content": [{"type": "text", "text": solution}]})
+            messages.append(message)
+
+        text = [
+            self.qwen_vl_interface.processor.apply_chat_template(
+                message, tokenize=False, add_generation_prompt=True
+            )
+            for message in messages
+        ]
+        image_inputs, video_inputs = process_vision_info(messages)
+        batch = self.qwen_vl_interface.processor(
+            text=text, images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt"
+        )
+        if with_actions:
+            ids = batch["input_ids"]
+            action_mask = (ids >= self.qwen_vl_interface.action_token_min) & (
+                ids <= self.qwen_vl_interface.action_token_max
+            )
+            expected = torch.tensor([len(seq) for seq in tokens])
+            if not torch.equal(action_mask.sum(dim=1).cpu(), expected):
+                raise ValueError("Processor changed the FAST action token sequence")
+            labels = torch.full_like(ids, -100)
+            labels[action_mask] = ids[action_mask]
+            batch["labels"] = labels
+        return batch.to(next(self.qwen_vl_interface.model.parameters()).device)
+
+    def forward(self, examples: list[dict], *, return_per_sample_loss: bool = False, flow_noise_seed=None, **kwargs):
+        if kwargs:
+            raise TypeError(f"Unexpected QwenFast forward arguments: {sorted(kwargs)}")
+        batch = self._inputs(examples, with_actions=True)
+        labels = batch.pop("labels")
+        outputs = self.qwen_vl_interface.model(**batch, use_cache=False, return_dict=True)
+        shifted_labels = labels[:, 1:]
+        mask = shifted_labels != -100
+        if not torch.all(mask.sum(dim=1) > 0):
+            raise ValueError("Every sample must contain an action target")
+        shifted_logits = outputs.logits[:, :-1]
+        token_loss = F.cross_entropy(shifted_logits[mask].float(), shifted_labels[mask], reduction="none")
+        sample_index = torch.arange(len(examples), device=mask.device).unsqueeze(1).expand_as(mask)[mask]
+        sums = torch.zeros(len(examples), dtype=token_loss.dtype, device=token_loss.device)
+        per_sample = sums.index_add(0, sample_index, token_loss) / mask.sum(dim=1)
+        return {"action_loss": per_sample if return_per_sample_loss else per_sample.mean()}
+
+    @torch.inference_mode()
+    def predict_action(self, examples: list[dict], **kwargs):
+        if isinstance(examples, dict):
+            examples = [examples]
+        batch = self._inputs(examples, with_actions=False)
+        prompt_length = batch["input_ids"].shape[1]
+        generated = self.qwen_vl_interface.model.generate(
+            **batch, max_new_tokens=kwargs.pop("max_new_tokens", 256), **kwargs
+        )[:, prompt_length:]
+        low = self.qwen_vl_interface.action_token_min
+        high = self.qwen_vl_interface.action_token_max
+        tokens = [row[(row >= low) & (row <= high)].sub(low).tolist() for row in generated]
+        for sequence in tokens:
+            decoded = self.action_model.fast_tokenizer.bpe_tokenizer.decode(sequence)
+            if len(decoded) != self.action_horizon * self.action_dim:
+                raise ValueError(
+                    "Generated FAST action sequence is incomplete or malformed: "
+                    f"decoded {len(decoded)} coefficients, expected "
+                    f"{self.action_horizon * self.action_dim}"
+                )
+        actions = self.action_model.fast_tokenizer.decode(
+            tokens, time_horizon=self.action_horizon, action_dim=self.action_dim
+        )
+        return {"normalized_actions": actions}

starVLA/model/framework/VLM4A/QwenGR00T.py

@@ -1,7 +1,6 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # Copyright 2025 starVLA community. All rights reserved.
 # Licensed under the MIT License, Version 1.0 (the "License");
-# Implemented by [Junqiu YU / Fudan University] in [2025].
-# Design and Merged by [Jinhui YE / HKUST University] in [2025].
 """
 Qwen-GR00T Framework
 A lightweight implementation that Qwen-VL + Flow-matching head to directly predict continuous actions
@@ -178,6 +177,7 @@ class Qwen_GR00T(baseframework):
         """ """
         return_per_sample_loss = bool(kwargs.pop("return_per_sample_loss", False))
         flow_noise_seed = kwargs.pop("flow_noise_seed", None)
+        return_flow_noise_time = bool(kwargs.pop("return_flow_noise_time", False))
         if kwargs:
             raise TypeError(f"Unexpected forward arguments: {sorted(kwargs)}")
         batch_images = [example["image"] for example in examples]  #  [B，[PLT]]
@@ -215,11 +215,42 @@ class Qwen_GR00T(baseframework):
                 state = torch.tensor(np.array(state), device=last_hidden.device, dtype=last_hidden.dtype)
                 state_repeated = state.repeat(repeated_diffusion_steps, 1, 1)
 
+            explicit_flow = [
+                "loss_is_flow_noise" in example and "loss_is_flow_time" in example
+                for example in examples
+            ]
+            if any(explicit_flow) and not all(explicit_flow):
+                raise ValueError("A batch cannot mix explicit and sampled flow realizations")
+            flow_noise = None
+            flow_time = None
+            if all(explicit_flow):
+                flow_noise = torch.stack(
+                    [
+                        torch.as_tensor(example["loss_is_flow_noise"])
+                        for example in examples
+                    ],
+                    dim=1,
+                ).reshape(actions_target_repeated.shape)
+                flow_time = torch.stack(
+                    [
+                        torch.as_tensor(example["loss_is_flow_time"])
+                        for example in examples
+                    ],
+                    dim=1,
+                ).reshape(actions_target_repeated.shape[0])
+
             action_loss = self.action_model(
                 last_hidden_repeated, actions_target_repeated, state_repeated,
                 reduction="none" if return_per_sample_loss else "mean",
                 flow_noise_seed=flow_noise_seed,
+                flow_noise=flow_noise,
+                flow_time=flow_time,
+                return_flow_noise_time=return_flow_noise_time,
             )  # (B, chunk_len, action_dim)
+            sampled_noise = None
+            sampled_time = None
+            if return_flow_noise_time:
+                action_loss, sampled_noise, sampled_time = action_loss
             if return_per_sample_loss:
                 action_loss = aggregate_repeated_losses(
                     action_loss,
@@ -227,7 +258,11 @@ class Qwen_GR00T(baseframework):
                     batch_size=len(examples),
                 )
 
-        return {"action_loss": action_loss}
+        output = {"action_loss": action_loss}
+        if return_flow_noise_time:
+            output["flow_noise"] = sampled_noise
+            output["flow_time"] = sampled_time
+        return output
 
     @torch.inference_mode()
     def predict_action(

starVLA/model/framework/base_framework.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 """
 Base framework abstraction providing:
 - Pretrained loading (config + normalization stats + weights)
@@ -94,7 +95,7 @@ def build_framework(cfg): # The single entry point for building different model
 
 
 # PreTrainedModel, AutoModel, PretrainedConfig,  are so good, find sometime to study them
-# TODO @JinhuiYE find sometime to merge yaml config with transformer config
+# Framework construction merges the YAML and transformer configurations.
 
 
 class baseframework(PreTrainedModel):

starVLA/model/framework/share_tools.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 """
 Shared configuration / utility helpers for framework components:
 - NamespaceWithGet: lightweight namespace behaving like a dict
@@ -349,7 +350,7 @@ def read_model_config(pretrained_checkpoint):
         with open(dataset_statistics_json, "r") as f:
             norm_stats = json.load(f)
     else:
-        overwatch.error(f"❌ Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
+        overwatch.error(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
         raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
     return global_cfg, norm_stats
 
@@ -387,30 +388,20 @@ def read_mode_config(pretrained_checkpoint):
             apply_config_compat(ocfg)
             global_cfg = OmegaConf.to_container(ocfg, resolve=True)
         except Exception as e:
-            overwatch.error(f"❌ Failed to load YAML config `{config_yaml}`: {e}")
+            overwatch.error(f"Failed to load YAML config `{config_yaml}`: {e}")
             raise
 
         # Load Dataset Statistics for Action Denormalization
         with open(dataset_statistics_json, "r") as f:
             norm_stats = json.load(f)
     else:
-        overwatch.error(f"❌ Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
+        overwatch.error(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
         raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.")
     return global_cfg, norm_stats
 
 
-# =============================================================================
-# Config compatibility / "tightening" layer (introduced in version_id "0.21").
-#
-# Goal: keep user-facing YAMLs short and unambiguous while preserving full
-# back-compat for old checkpoints' config.yaml. See bar/config_收紧.md for
-# the design rationale.
-#
-# This function is *idempotent* — calling it multiple times yields the same
-# result. It does NOT touch framework class signatures; instead it normalises
-# the OmegaConf tree so that downstream framework __init__ code (which still
-# reads e.g. `future_action_window_size`) keeps working unchanged.
-# =============================================================================
+# Normalize legacy configuration keys before model construction. The
+# transformation is idempotent and leaves framework class signatures intact.
 
 CONFIG_VERSION = "0.21"
 

starVLA/model/modules/action_model/GR00T_ActionHeader.py

@@ -1,5 +1,5 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved.
-# Modified by [Junqiu YU/ Fudan University] in [2025].
 # Modification: [rm and add some connect adapter to match with starVLA, e.g., "rm "].
 # Action repeat is inspired by CogACT
 
@@ -311,7 +311,21 @@ class FlowmatchingActionHead(nn.Module):
         self,
         actions: torch.Tensor,
         flow_noise_seed: int | None,
+        flow_noise: torch.Tensor | None = None,
+        flow_time: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (flow_noise is None) != (flow_time is None):
+            raise ValueError("flow_noise and flow_time must be supplied together")
+        if flow_noise is not None:
+            if tuple(flow_noise.shape) != tuple(actions.shape):
+                raise ValueError("Explicit flow noise does not match the action shape")
+            if tuple(flow_time.shape) != (actions.shape[0],):
+                raise ValueError("Explicit flow time does not match the batch shape")
+            return (
+                flow_noise.to(device=actions.device, dtype=actions.dtype),
+                flow_time.to(device=actions.device, dtype=actions.dtype),
+            )
+
         def sample() -> tuple[torch.Tensor, torch.Tensor]:
             noise = torch.randn(
                 actions.shape,
@@ -391,6 +405,9 @@ class FlowmatchingActionHead(nn.Module):
         encoder_attention_mask=None,
         reduction: str = "mean",
         flow_noise_seed: int | None = None,
+        flow_noise: torch.Tensor | None = None,
+        flow_time: torch.Tensor | None = None,
+        return_flow_noise_time: bool = False,
     ):
         """
         vl_embs: shape (B, seq_length, feature_dim)
@@ -400,6 +417,8 @@ class FlowmatchingActionHead(nn.Module):
         noise, t = self._sample_training_noise_and_time(
             actions,
             flow_noise_seed,
+            flow_noise=flow_noise,
+            flow_time=flow_time,
         )
         t = t[:, None, None]  # shape (B,1,1) for broadcast
 
@@ -416,7 +435,10 @@ class FlowmatchingActionHead(nn.Module):
 
         # Slice out only the action portion of pred and target.
         elementwise_loss = (pred_actions - velocity) ** 2
-        return reduce_flow_matching_loss(elementwise_loss, reduction=reduction)
+        loss = reduce_flow_matching_loss(elementwise_loss, reduction=reduction)
+        if return_flow_noise_time:
+            return loss, noise, t[:, 0, 0]
+        return loss
 
     @torch.no_grad()
     def predict_action(

starVLA/model/modules/action_model/OpenPI_ActionHead.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import inspect
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+from starVLA.model.modules.vlm.openpi_transformers.gemma.configuration_gemma import GemmaConfig
+from starVLA.model.modules.vlm.openpi_transformers.gemma.modeling_gemma import GemmaForCausalLM
+
+
+@dataclass
+class GemmaDims:
+    width: int
+    depth: int
+    mlp_dim: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device: torch.device,
+) -> Tensor:
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    if time.ndim != 1:
+        raise ValueError("time must have shape [batch_size]")
+
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float32, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * torch.pi
+    sin_input = scaling_factor[None, :] * time.to(torch.float32)[:, None]
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def sample_beta(alpha: float, beta: float, bsize: int, device: torch.device) -> torch.Tensor:
+    alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device)
+    beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device)
+    return torch.distributions.Beta(alpha_t, beta_t).sample((bsize,))
+
+
+def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
+    if att_masks.ndim != 2 or pad_masks.ndim != 2:
+        raise ValueError(f"Expected 2D masks, got {pad_masks.ndim=} {att_masks.ndim=}")
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+def make_att_4d_masks(att_2d_masks: torch.Tensor) -> torch.Tensor:
+    return torch.where(att_2d_masks[:, None, :, :], 0.0, -2.3819763e38)
+
+
+
+class OpenPIGemma(nn.Module):
+    """Gemma transformer wrapper used by OpenPI action heads."""
+
+    def __init__(self, config: GemmaDims, *, use_adarms: bool = False):
+        super().__init__()
+        hf_config = GemmaConfig(
+            head_dim=config.head_dim,
+            hidden_size=config.width,
+            intermediate_size=config.mlp_dim,
+            num_attention_heads=config.num_heads,
+            num_hidden_layers=config.depth,
+            num_key_value_heads=config.num_kv_heads,
+            vocab_size=257152,
+            hidden_activation="gelu_pytorch_tanh",
+            torch_dtype="float32",
+            use_adarms=bool(use_adarms),
+            adarms_cond_dim=config.width if use_adarms else None,
+        )
+        gemma = GemmaForCausalLM(config=hf_config)
+        gemma.model.embed_tokens = None
+        self.model = gemma.model
+        self.lm_head = gemma.lm_head
+        self.validate_transformers_api(use_adarms=use_adarms)
+        if use_adarms:
+            self._ensure_adarms_norms(config.width)
+
+    def _ensure_adarms_norms(self, cond_dim: int) -> None:
+        for layer in self.model.layers:
+            self._ensure_adarms_norm(layer.input_layernorm, cond_dim)
+            self._ensure_adarms_norm(layer.post_attention_layernorm, cond_dim)
+        self._ensure_adarms_norm(self.model.norm, cond_dim)
+
+    @staticmethod
+    def _ensure_adarms_norm(norm: nn.Module, cond_dim: int) -> None:
+        if getattr(norm, "dense", None) is None:
+            width = int(norm.weight.shape[0])
+            norm.cond_dim = int(cond_dim)
+            norm.dense = nn.Linear(cond_dim, width * 3)
+            nn.init.zeros_(norm.dense.weight)
+        if hasattr(norm, "weight"):
+            norm.weight.data.zero_()
+            norm.weight.requires_grad_(False)
+
+    def validate_transformers_api(self, *, use_adarms: bool) -> None:
+        if not use_adarms:
+            return
+        norm_sig = inspect.signature(self.model.layers[0].input_layernorm.forward)
+        model_sig = inspect.signature(self.model.forward)
+        if "cond" not in norm_sig.parameters or "adarms_cond" not in model_sig.parameters:
+            raise RuntimeError(
+                "PI05 requires StarVLA vendored OpenPI-compatible Gemma adaRMS support. "
+                "The imported Gemma does not expose GemmaRMSNorm.forward(cond=...) and "
+                "GemmaModel.forward(adarms_cond=...). "
+                f"Current signatures: GemmaRMSNorm.forward{norm_sig}, GemmaModel.forward{model_sig}"
+            )
+
+
+class OpenPIActionHeadBase(nn.Module):
+    """OpenPI action expert + action head.
+
+    This module owns the action-side Gemma expert, suffix embedding, action
+    projections, and the suffix-only denoise forward used during inference.
+    """
+
+    use_adarms: bool = False
+
+    def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int):
+        super().__init__()
+        self.action_dim = int(action_dim)
+        self.action_horizon = int(action_horizon)
+        self.max_state_dim = int(max_state_dim)
+        self.width = int(config.width)
+        self.gemma_expert = OpenPIGemma(config, use_adarms=self.use_adarms)
+        self.action_in_proj = nn.Linear(self.action_dim, self.width)
+        self.action_out_proj = nn.Linear(self.width, self.action_dim)
+
+    @property
+    def model(self):
+        return self.gemma_expert.model
+
+    @property
+    def layers(self):
+        return self.gemma_expert.model.layers
+
+    @property
+    def norm(self):
+        return self.gemma_expert.model.norm
+
+    def set_gemma_expert_eager_attention(self):
+        self.gemma_expert.model.config._attn_implementation = "eager"
+
+    def supports_cached_denoise(self) -> bool:
+        return True
+
+    def forward_gemma_expert(
+        self,
+        *,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.LongTensor | None,
+        past_key_values,
+        use_cache: bool | None,
+        adarms_cond: torch.Tensor | None,
+    ):
+        kwargs = {
+            "inputs_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": use_cache,
+        }
+        if "adarms_cond" in inspect.signature(self.gemma_expert.model.forward).parameters:
+            kwargs["adarms_cond"] = adarms_cond
+        elif adarms_cond is not None:
+            raise RuntimeError(
+                "PI05 requires StarVLA vendored OpenPI-compatible Gemma adaRMS support, "
+                "but the imported action expert does not expose it."
+            )
+        return self.gemma_expert.model.forward(**kwargs)
+
+    def sample_action_noise(self, shape, device: torch.device):
+        return torch.normal(0.0, 1.0, size=shape, dtype=torch.float32, device=device)
+
+    def sample_flow_time(self, bsize: int, device: torch.device):
+        return (sample_beta(1.5, 1.0, bsize, device) * 0.999 + 0.001).to(torch.float32)
+
+    def prepare_flow_targets(self, actions: torch.Tensor, noise=None, time=None):
+        actions = actions.to(device=self.action_in_proj.weight.device, dtype=torch.float32)
+        noise = self.sample_action_noise(actions.shape, actions.device) if noise is None else noise.to(actions.device, torch.float32)
+        time = self.sample_flow_time(actions.shape[0], actions.device) if time is None else time.to(actions.device, torch.float32)
+        x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions
+        u_t = noise - actions
+        return x_t, u_t, noise, time
+
+    def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor):
+        raise NotImplementedError
+
+    def decode_action_velocity(self, action_hidden_states: torch.Tensor) -> torch.Tensor:
+        action_hidden_states = action_hidden_states[:, -self.action_horizon :].to(self.action_out_proj.weight.dtype)
+        return self.action_out_proj(action_hidden_states).to(torch.float32)
+
+    def forward(
+        self,
+        *,
+        attention_mask: torch.Tensor,
+        position_ids: torch.LongTensor,
+        past_key_values,
+        suffix_embs: torch.Tensor,
+        adarms_cond: torch.Tensor | None,
+        force_bfloat16: bool = False,
+    ):
+        if force_bfloat16:
+            suffix_embs = suffix_embs.to(torch.bfloat16)
+        self.set_gemma_expert_eager_attention()
+        suffix_output = self.forward_gemma_expert(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=suffix_embs,
+            use_cache=False,
+            adarms_cond=adarms_cond,
+        )
+        suffix_out = suffix_output.last_hidden_state
+        return self.decode_action_velocity(suffix_out)
+
+
+
+class OpenPI0ActionHead(OpenPIActionHeadBase):
+    """PI0 suffix path: state token plus action-time MLP, no adaRMS conditioning."""
+
+    use_adarms = False
+
+    def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int):
+        super().__init__(config, action_dim=action_dim, action_horizon=action_horizon, max_state_dim=max_state_dim)
+        self.state_proj = nn.Linear(self.max_state_dim, self.width)
+        self.action_time_mlp_in = nn.Linear(2 * self.width, self.width)
+        self.action_time_mlp_out = nn.Linear(self.width, self.width)
+
+    def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor):
+        state_emb = self.state_proj(state.to(self.state_proj.weight.dtype))
+        embs = [state_emb[:, None, :]]
+        pad_masks = [torch.ones(state_emb.shape[0], 1, dtype=torch.bool, device=state_emb.device)]
+        att_masks = [1]
+
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.width,
+            min_period=4e-3,
+            max_period=4.0,
+            device=timestep.device,
+        ).to(self.action_in_proj.weight.dtype)
+        action_emb = self.action_in_proj(noisy_actions.to(self.action_in_proj.weight.dtype))
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = self.action_time_mlp_out(
+            F.silu(self.action_time_mlp_in(torch.cat([action_emb, time_emb], dim=2)))
+        )
+        embs.append(action_time_emb)
+        pad_masks.append(torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=action_time_emb.device))
+        att_masks += [1] + ([0] * (self.action_horizon - 1))
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device)[None, :]
+        return embs, pad_masks, att_masks.expand(embs.shape[0], -1), None
+
+
+class OpenPI05ActionHead(OpenPIActionHeadBase):
+    """PI05 suffix path: action tokens conditioned by adaRMS time embedding."""
+
+    use_adarms = True
+
+    def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int):
+        super().__init__(config, action_dim=action_dim, action_horizon=action_horizon, max_state_dim=max_state_dim)
+        self.time_mlp_in = nn.Linear(self.width, self.width)
+        self.time_mlp_out = nn.Linear(self.width, self.width)
+
+    def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor):
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.width,
+            min_period=4e-3,
+            max_period=4.0,
+            device=timestep.device,
+        ).to(self.time_mlp_in.weight.dtype)
+        action_time_emb = self.action_in_proj(noisy_actions.to(self.action_in_proj.weight.dtype))
+        adarms_cond = F.silu(self.time_mlp_out(F.silu(self.time_mlp_in(time_emb))))
+        pad_masks = torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=action_time_emb.device)
+        att_masks = torch.tensor([1] + ([0] * (self.action_horizon - 1)), dtype=torch.bool, device=action_time_emb.device)
+        return action_time_emb, pad_masks, att_masks[None, :].expand(action_time_emb.shape[0], -1), adarms_cond

starVLA/model/modules/action_model/__init__.py

@@ -1 +1,2 @@
+# Modified in the ReTime research artifact; see NOTICE.
 """Action models included in the Florence-Bridge OPTR release."""

starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py

@@ -1,3 +1,4 @@
+# Modified in the ReTime research artifact; see NOTICE.
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -166,7 +167,7 @@ class BasicTransformerBlock(nn.Module):
         attn_output = self.attn1(
             norm_hidden_states,
             encoder_hidden_states=encoder_hidden_states,
-            attention_mask=encoder_attention_mask,  # @JinhuiYE original attention_mask=attention_mask
+            attention_mask=encoder_attention_mask,
         )
         if self.final_dropout:
             attn_output = self.final_dropout(attn_output)
@@ -276,7 +277,7 @@ class DiT(ModelMixin, ConfigMixin):
         hidden_states = hidden_states.contiguous()
         encoder_hidden_states = encoder_hidden_states.contiguous()
 
-        all_hidden_states = [hidden_states]
+        all_hidden_states = [hidden_states] if return_all_hidden_states else None
 
         # Process through transformer blocks
         for idx, block in enumerate(self.transformer_blocks):
@@ -296,7 +297,8 @@ class DiT(ModelMixin, ConfigMixin):
                     encoder_attention_mask=encoder_attention_mask,
                     temb=temb,
                 )
-            all_hidden_states.append(hidden_states)
+            if all_hidden_states is not None:
+                all_hidden_states.append(hidden_states)
 
         # Output processing
         conditioning = temb
@@ -364,12 +366,13 @@ class SelfAttentionTransformer(ModelMixin, ConfigMixin):
     ):
         # Process through transformer blocks - single pass through the blocks
         hidden_states = hidden_states.contiguous()
-        all_hidden_states = [hidden_states]
+        all_hidden_states = [hidden_states] if return_all_hidden_states else None
 
         # Process through transformer blocks
         for idx, block in enumerate(self.transformer_blocks):
             hidden_states = block(hidden_states)
-            all_hidden_states.append(hidden_states)
+            if all_hidden_states is not None:
+                all_hidden_states.append(hidden_states)
 
         if return_all_hidden_states:
             return hidden_states, all_hidden_states