diff --git a/external_baselines/frameskip_adapted/build_required_draws.py b/external_baselines/frameskip_adapted/build_required_draws.py index d95cf67c56718441640c1f00d8900081ae8ebcf3..cbff1f2bc119113b2b73cd7e60accf2fde82e535 100644 --- a/external_baselines/frameskip_adapted/build_required_draws.py +++ b/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"], diff --git a/external_baselines/frameskip_adapted/finalize_cache.py b/external_baselines/frameskip_adapted/finalize_cache.py index aafb0587cfd9c2b6e66c3e8d92e3d579e234ba67..634b1c9e5c9ac16facf8ef53c1b878ddb4d6fe32 100644 --- a/external_baselines/frameskip_adapted/finalize_cache.py +++ b/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"]), diff --git a/external_baselines/frameskip_adapted/train_frameskip_continuation.py b/external_baselines/frameskip_adapted/train_frameskip_continuation.py index d1293326724fd62d4b682e254d35247f3e160d89..f6fb9816852feb64236bd568bed8111b209f1bd8 100644 --- a/external_baselines/frameskip_adapted/train_frameskip_continuation.py +++ b/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, diff --git a/external_baselines/frameskip_robotwin_adapted/build_robotwin_cache.py b/external_baselines/frameskip_robotwin_adapted/build_robotwin_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..7eef74188372e8200e2331f15fe0d34a0d6e7ca9 --- /dev/null +++ b/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() diff --git a/external_baselines/frameskip_robotwin_adapted/train_frameskip_robotwin_continuation.py b/external_baselines/frameskip_robotwin_adapted/train_frameskip_robotwin_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..95b8b1da1741412b6b3b771d1dfa0f3c7c00d22d --- /dev/null +++ b/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() diff --git a/external_baselines/frameskip_robotwin_adapted/validate_assets.py b/external_baselines/frameskip_robotwin_adapted/validate_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..598c6630e8cc96088dd02775607d4da34be8a070 --- /dev/null +++ b/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() diff --git a/external_baselines/keyframe_focused_adapted/build_ape_cache.py b/external_baselines/keyframe_focused_adapted/build_ape_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..96af93e78e6ee5cbbc8401e8f6e6a0c781fc0187 --- /dev/null +++ b/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() diff --git a/external_baselines/keyframe_focused_adapted/configs/keyframe_focused_seed117465_20k_to_20500.yaml b/external_baselines/keyframe_focused_adapted/configs/keyframe_focused_seed117465_20k_to_20500.yaml new file mode 100644 index 0000000000000000000000000000000000000000..01cdb68c2bd0f16c919c554b17c9522e97458f8e --- /dev/null +++ b/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 diff --git a/external_baselines/keyframe_focused_adapted/keyframe_cache.py b/external_baselines/keyframe_focused_adapted/keyframe_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..4966e0ea0b4fa02231678eb7ae69af82d015952c --- /dev/null +++ b/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]), + ) diff --git a/external_baselines/keyframe_focused_adapted/train_keyframe_continuation.py b/external_baselines/keyframe_focused_adapted/train_keyframe_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2ff32e2f75a4760946d750206ccbb0cdb2f390 --- /dev/null +++ b/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() diff --git a/external_baselines/keyframe_focused_adapted/validate_cache_support.py b/external_baselines/keyframe_focused_adapted/validate_cache_support.py new file mode 100644 index 0000000000000000000000000000000000000000..c90913abc82c0222c6ca43ba97c05633604c381f --- /dev/null +++ b/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() diff --git a/external_baselines/loss_is_active/train_active_loss_is_continuation.py b/external_baselines/loss_is_active/train_active_loss_is_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..13c1845a13167773acb8c17abf09a67275a61b44 --- /dev/null +++ b/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() diff --git a/external_baselines/loss_is_adapted/configs/loss_is_seed117465_20k_to_20500.yaml b/external_baselines/loss_is_adapted/configs/loss_is_seed117465_20k_to_20500.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8efdfa91d3849ffddd6628d3a6ae0fd6932fdd4a --- /dev/null +++ b/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 diff --git a/external_baselines/loss_is_adapted/loss_is.py b/external_baselines/loss_is_adapted/loss_is.py new file mode 100644 index 0000000000000000000000000000000000000000..edc64316f3c7f181926369c2973a5218e453ecc8 --- /dev/null +++ b/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, + ) + diff --git a/external_baselines/loss_is_adapted/train_loss_is_continuation.py b/external_baselines/loss_is_adapted/train_loss_is_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..29e1e8e2761ca3320464cae60bc8f6bf38df4e33 --- /dev/null +++ b/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() diff --git a/external_baselines/qwen_adapted/train_qwen_frameskip_continuation.py b/external_baselines/qwen_adapted/train_qwen_frameskip_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..4e702c01dd0132b68bb0b92189a7604e14074841 --- /dev/null +++ b/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() diff --git a/external_baselines/qwen_robotwin_adapted/register_qwen_robotwin_smoke.py b/external_baselines/qwen_robotwin_adapted/register_qwen_robotwin_smoke.py new file mode 100644 index 0000000000000000000000000000000000000000..5806d15f31fda735839c16b575357882a15f28cd --- /dev/null +++ b/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"), + ] diff --git a/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_frameskip.py b/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_frameskip.py new file mode 100644 index 0000000000000000000000000000000000000000..27fcf88f00c401dc321a7e179b9b4d25e7c4ea96 --- /dev/null +++ b/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() diff --git a/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_keyframe.py b/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_keyframe.py new file mode 100644 index 0000000000000000000000000000000000000000..7d7a1b163bfebaeabb4d31d256e1d0444b68e0f6 --- /dev/null +++ b/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() diff --git a/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_loss_is.py b/external_baselines/qwen_robotwin_adapted/train_qwen_robotwin_loss_is.py new file mode 100644 index 0000000000000000000000000000000000000000..d7f72dba0f1a5a6815486ab90cbc4bf7eb15c746 --- /dev/null +++ b/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() diff --git a/external_baselines/settings/florence_bridge/frameskip.yaml b/external_baselines/settings/florence_bridge/frameskip.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95caee0f012028f166eeb3dd6a1c40971fbaa714 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/florence_bridge/keyframe.yaml b/external_baselines/settings/florence_bridge/keyframe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f0c9855c7f92dd2d4700840a2150cb7a68f5176 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/florence_bridge/loss_is_active.yaml b/external_baselines/settings/florence_bridge/loss_is_active.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b94c61713a40e4b915e961d35f2f120a8ed02387 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/florence_robotwin/frameskip.yaml b/external_baselines/settings/florence_robotwin/frameskip.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49a6ce8c90c8403ad3ef139349c4f86cc8152504 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/florence_robotwin/keyframe.yaml b/external_baselines/settings/florence_robotwin/keyframe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..778789f9734017006c15634ab9306ffabacf6143 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/florence_robotwin/loss_is_active.yaml b/external_baselines/settings/florence_robotwin/loss_is_active.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd1199ad9888d836434ca5b0e197840dc86d42b6 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_bridge/frameskip.yaml b/external_baselines/settings/qwen_bridge/frameskip.yaml new file mode 100644 index 0000000000000000000000000000000000000000..acf8e3f190d817d0af9a5c0f8188610beaf8f631 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_bridge/keyframe.yaml b/external_baselines/settings/qwen_bridge/keyframe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..70873c74138691c9440d7e8f2d9f687d37328903 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_bridge/loss_is_active.yaml b/external_baselines/settings/qwen_bridge/loss_is_active.yaml new file mode 100644 index 0000000000000000000000000000000000000000..32fc3769495992dfa84873564ea5586e4083adb6 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_robotwin/frameskip.yaml b/external_baselines/settings/qwen_robotwin/frameskip.yaml new file mode 100644 index 0000000000000000000000000000000000000000..489bec84cbdc9dafc1a20b798569bcf2fb2bdf8f --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_robotwin/keyframe.yaml b/external_baselines/settings/qwen_robotwin/keyframe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..738b7edc74d299cae048541561fa8b57f3e9eba3 --- /dev/null +++ b/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 diff --git a/external_baselines/settings/qwen_robotwin/loss_is_active.yaml b/external_baselines/settings/qwen_robotwin/loss_is_active.yaml new file mode 100644 index 0000000000000000000000000000000000000000..77122b92722d3bbdc8a1c0d309a19123772c2329 --- /dev/null +++ b/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 diff --git a/starVLA/dataloader/__init__.py b/starVLA/dataloader/__init__.py index 7cb73d66cd1ad0e0329d55bd2f2a8c0f663644ad..58b002c639fa544d4b82b7f5e0d03c3a1b7dd669 100644 --- a/starVLA/dataloader/__init__.py +++ b/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}") diff --git a/starVLA/dataloader/gr00t_lerobot/data_config.py b/starVLA/dataloader/gr00t_lerobot/data_config.py index 618044216ab4af4671418e720df811fea5100445..12f13ce48d87a820eac02cc9c0ee7fe941234818 100644 --- a/starVLA/dataloader/gr00t_lerobot/data_config.py +++ b/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", ] diff --git a/starVLA/dataloader/gr00t_lerobot/datasets.py b/starVLA/dataloader/gr00t_lerobot/datasets.py index 7082123b2ad75e248087f4c5aa51caefb7e79415..39b235b256b3924fc18c7786c5321f2ceb29ea8d 100644 --- a/starVLA/dataloader/gr00t_lerobot/datasets.py +++ b/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 diff --git a/starVLA/dataloader/gr00t_lerobot/transform/state_action.py b/starVLA/dataloader/gr00t_lerobot/transform/state_action.py index 7a740b343187fdedfa1215d7a1d07f4604cf689c..51190a87e5c5092953802999baf986488f875311 100644 --- a/starVLA/dataloader/gr00t_lerobot/transform/state_action.py +++ b/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 # diff --git a/starVLA/dataloader/infoflow_sampler.py b/starVLA/dataloader/infoflow_sampler.py index 11991c847fc7e59984648f13d42f5b9382820ff1..db228ababf79814d5c12e3dc615cffe4b13ea9cf 100644 --- a/starVLA/dataloader/infoflow_sampler.py +++ b/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') diff --git a/starVLA/dataloader/lerobot_datasets.py b/starVLA/dataloader/lerobot_datasets.py index 6e45cbe8c5a9145a2a56ed4a314c7e664d18547f..8c4134eb77d159418979fd5f0f59f073f59c9dd9 100644 --- a/starVLA/dataloader/lerobot_datasets.py +++ b/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 \ No newline at end of file + pass diff --git a/starVLA/dataloader/relative_pose.py b/starVLA/dataloader/relative_pose.py new file mode 100644 index 0000000000000000000000000000000000000000..943c9e15c1ac38fc894586a7b333d0e9c1c2f06a --- /dev/null +++ b/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" diff --git a/starVLA/dataloader/temporal_proposal.py b/starVLA/dataloader/temporal_proposal.py index 4ed3fca3da099ebf45af25d29972d6000356c21f..777256fb7488728579facb8a367d7aed499ef834 100644 --- a/starVLA/dataloader/temporal_proposal.py +++ b/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') diff --git a/starVLA/model/framework/VLM4A/PI0.py b/starVLA/model/framework/VLM4A/PI0.py new file mode 100644 index 0000000000000000000000000000000000000000..33f7e963116bc29bab41dc705395cfe735efecfa --- /dev/null +++ b/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 diff --git a/starVLA/model/framework/VLM4A/PI05.py b/starVLA/model/framework/VLM4A/PI05.py new file mode 100644 index 0000000000000000000000000000000000000000..8a3aa6b3bcd2e47277dd03f9ec298e12a4def4a9 --- /dev/null +++ b/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 diff --git a/starVLA/model/framework/VLM4A/QwenFast.py b/starVLA/model/framework/VLM4A/QwenFast.py new file mode 100644 index 0000000000000000000000000000000000000000..01a91d315b092196ab0e862d4a1c0f9cb770c697 --- /dev/null +++ b/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("") + self.action_token_max = self.processor.tokenizer.convert_tokens_to_ids("") + 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 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"" 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} diff --git a/starVLA/model/framework/VLM4A/QwenGR00T.py b/starVLA/model/framework/VLM4A/QwenGR00T.py index f1ab3786d58307c28020332841b9cdf375db7305..0923e4c87874bae484c1e62f3cb1b8cbe0c5264c 100644 --- a/starVLA/model/framework/VLM4A/QwenGR00T.py +++ b/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( diff --git a/starVLA/model/framework/base_framework.py b/starVLA/model/framework/base_framework.py index 8c4c5fe9c08ebc3cb93801c9ac68731e2b9907a5..47497c0bea94f3dd971d670967d5d80a6ffd7acf 100644 --- a/starVLA/model/framework/base_framework.py +++ b/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): diff --git a/starVLA/model/framework/share_tools.py b/starVLA/model/framework/share_tools.py index 8606daad182d80ba3d59dee136c8e982c07b1ee1..a8c8c4291065acd84096dbf1703292f37dc73a70 100644 --- a/starVLA/model/framework/share_tools.py +++ b/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" diff --git a/starVLA/model/modules/action_model/GR00T_ActionHeader.py b/starVLA/model/modules/action_model/GR00T_ActionHeader.py index 67e102de57894a2548ce4f8df9ab3eb09032df4e..9ca51428c8aa387a5ba8b85461a690c35dc8c832 100644 --- a/starVLA/model/modules/action_model/GR00T_ActionHeader.py +++ b/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( diff --git a/starVLA/model/modules/action_model/OpenPI_ActionHead.py b/starVLA/model/modules/action_model/OpenPI_ActionHead.py new file mode 100644 index 0000000000000000000000000000000000000000..1b0648e973e55ac0a96f62643b4d3b22f23dac81 --- /dev/null +++ b/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 diff --git a/starVLA/model/modules/action_model/__init__.py b/starVLA/model/modules/action_model/__init__.py index 6c16f8f21133c45b778d3ccda0f5e84580c5c235..b09bb11c9b2bf0ed931612e81f6f729cddd0c162 100644 --- a/starVLA/model/modules/action_model/__init__.py +++ b/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.""" diff --git a/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py b/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py index 4394c6444e08d27053ce5d139326bd4e7a3fc7e8..412dd8a7f8a6fe08d799c4b1f8e38176a8eb2ed9 100644 --- a/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py +++ b/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 diff --git a/starVLA/model/modules/vlm/Florence2.py b/starVLA/model/modules/vlm/Florence2.py index a2e000b9e5eb285efe8fa5e8fcc1049f7ad4adc2..a937144ab85a43ecd300baeac3ca37dd5dac6447 100644 --- a/starVLA/model/modules/vlm/Florence2.py +++ b/starVLA/model/modules/vlm/Florence2.py @@ -1,6 +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 [Jinhui YE / HKUST University] in [2025]. from typing import Optional @@ -227,7 +227,27 @@ class _Florence_Interface(nn.Module): dim=1, ) enc_out.last_hidden_state = action_conditioning - enc_out.hidden_states = [action_conditioning] + # `hidden_states` MUST be written through the mapping API, not as a plain + # attribute. `enc_out` is a transformers `ModelOutput` (an OrderedDict + # subclass): plain attribute assignment for a key that is not already in + # the dict stores it *outside* the mapping, so `enc_out.keys()` would only + # ever report `last_hidden_state`. + # + # DeepSpeed ZeRO-3 walks the module output with `apply_to_tensors_only`, + # which iterates a dict output by `keys()`, and replaces every tensor it + # finds with `PreBackwardFunctionForModule.apply(tensor)`. That wrapper is + # what triggers `pre_sub_module_backward_function` and therefore the + # backward-time all-gather of this module's own parameters. + # + # `auxiliary_view_embeddings` (see __init__) is a parameter held directly + # on this wrapper and consumed by the free function + # `_partition_multiview_features`, so it is gathered only by that hook. + # When `hidden_states` lived outside the mapping, the framework read the + # unwrapped tensor (QwenGR00T.forward uses `hidden_states[-1]`), the hook + # never fired, the parameter stayed partitioned to a 0-numel shard, and + # backward died in AccumulateGrad with + # "The size of tensor a (0) must match the size of tensor b (1024)". + enc_out["hidden_states"] = (action_conditioning,) enc_out.auxiliary_visual_tokens = auxiliary_features.shape[1] # last_hidden = qwenvl_outputs.hidden_states[-1] # [B, L, H] return enc_out diff --git a/starVLA/model/modules/vlm/OpenPIPaliGemma.py b/starVLA/model/modules/vlm/OpenPIPaliGemma.py new file mode 100644 index 0000000000000000000000000000000000000000..4cf67189f2b70712e83a3e9b0b3c09721918e6af --- /dev/null +++ b/starVLA/model/modules/vlm/OpenPIPaliGemma.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import inspect + +import torch +from torch import nn +from transformers.models.auto import CONFIG_MAPPING + +from starVLA.model.modules.vlm.openpi_transformers.gemma.configuration_gemma import GemmaConfig +from starVLA.model.modules.vlm.openpi_transformers.gemma.modeling_gemma import GemmaModel +from starVLA.model.modules.vlm.openpi_transformers.siglip.modeling_siglip import SiglipVisionModel + + +class OpenPIMultiModalProjector(nn.Module): + def __init__(self, config): + super().__init__() + self.linear = nn.Linear(config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True) + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + return self.linear(image_features) + + +class OpenPIPaliGemmaCore(nn.Module): + def __init__(self, hf_config, text_config: GemmaConfig): + super().__init__() + self.config = hf_config + self.vision_tower = SiglipVisionModel(hf_config.vision_config) + self.multi_modal_projector = OpenPIMultiModalProjector(hf_config) + self.language_model = GemmaModel(text_config) + self.vocab_size = text_config.vocab_size + self.pad_token_id = hf_config.pad_token_id if hf_config.pad_token_id is not None else -1 + + def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + image_outputs = self.vision_tower(pixel_values) + return self.multi_modal_projector(image_outputs.last_hidden_state) + + +class OpenPIPaliGemmaContainer(nn.Module): + def __init__(self, hf_config, text_config: GemmaConfig): + super().__init__() + self.config = hf_config + self.model = OpenPIPaliGemmaCore(hf_config, text_config) + self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) + self.model.language_model.embed_tokens.weight = self.lm_head.weight + + @property + def language_model(self): + return self.model.language_model + + @property + def vision_tower(self): + return self.model.vision_tower + + @property + def multi_modal_projector(self): + return self.model.multi_modal_projector + + def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + return self.model.get_image_features(pixel_values) + + +class OpenPIPaliGemma(nn.Module): + """PaliGemma VLM component used by OpenPI PI0/PI05.""" + + def __init__(self, config, *, use_adarms: bool = False, model_name_or_path: str = ""): + super().__init__() + hf_config = CONFIG_MAPPING["paligemma"]() + hf_config._vocab_size = 257152 + hf_config.image_token_index = 257152 + hf_config.text_config.hidden_size = config.width + hf_config.text_config.intermediate_size = config.mlp_dim + hf_config.text_config.num_attention_heads = config.num_heads + hf_config.text_config.head_dim = config.head_dim + hf_config.text_config.num_hidden_layers = config.depth + hf_config.text_config.num_key_value_heads = config.num_kv_heads + hf_config.text_config.hidden_activation = "gelu_pytorch_tanh" + hf_config.text_config.torch_dtype = "float32" + hf_config.text_config.vocab_size = 257152 + hf_config.text_config.use_adarms = bool(use_adarms) + hf_config.text_config.adarms_cond_dim = config.width if use_adarms else None + hf_config.vision_config.intermediate_size = 4304 + hf_config.vision_config.projection_dim = config.width + hf_config.vision_config.projector_hidden_act = "gelu_fast" + hf_config.vision_config.torch_dtype = "float32" + + if model_name_or_path: + raise ValueError( + "OpenPIPaliGemma uses StarVLA vendored OpenPI modules. " + "Load weights through PI0/PI05.load_pretrained_checkpoint instead of model_name_or_path." + ) + self.model = OpenPIPaliGemmaContainer( + hf_config, + self._build_openpi_text_config(config, use_adarms=use_adarms), + ) + + @staticmethod + def _build_openpi_text_config(config, *, use_adarms: bool): + return 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, + ) + + @property + def config(self): + return self.model.config + + @property + def language_model(self): + return self.model.language_model + + @property + def vision_tower(self): + return self.model.vision_tower + + def embed_image(self, image: torch.Tensor): + vision_model = self.vision_tower.vision_model + if hasattr(vision_model, "config"): + vision_model.config._attn_implementation = "eager" + try: + return self.model.get_image_features(image) + except RuntimeError as exc: + if image.device.type != "cuda" or "expected scalar type Float but found BFloat16" not in str(exc): + raise + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + return self.model.get_image_features(image) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.language_model.embed_tokens(tokens) + + def forward_prefix( + self, + *, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.LongTensor | None, + past_key_values, + use_cache: bool | 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.language_model.forward).parameters: + kwargs["adarms_cond"] = None + return self.language_model.forward(**kwargs) diff --git a/starVLA/model/modules/vlm/QWen3.py b/starVLA/model/modules/vlm/QWen3.py new file mode 100644 index 0000000000000000000000000000000000000000..9675be5a0694d44bde95b8e3674c06096046cc1b --- /dev/null +++ b/starVLA/model/modules/vlm/QWen3.py @@ -0,0 +1,192 @@ +# 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"); + +from typing import Optional + +import torch +from starVLA.training.trainer_utils import initialize_overwatch +from transformers import AutoProcessor, Qwen3VLForConditionalGeneration +from transformers.modeling_outputs import CausalLMOutputWithPast + +logger = initialize_overwatch(__name__) + +IGNORE_INDEX = -100 +IMAGE_TOKEN_INDEX = 151655 +VIDEO_TOKEN_INDEX = 151656 +DEFAULT_IMAGE_TOKEN = "" +DEFAULT_VIDEO_TOKEN = "