mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 03:06:01 +00:00
Compare commits
12 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Martino Russi
|
ddcb61f7dd | ||
|
Steven Palma
|
a9879e69ed | ||
|
Steven Palma
|
9d82bb9871 | ||
|
Steven Palma
|
c5371d0691 | ||
|
Steven Palma
|
b2c062c0f4 | ||
|
Maxime Ellerbach
|
051b13573e | ||
|
Pepijn
|
7de2e4c1ef | ||
|
Nikodem Bartnik
|
8db50611c2 | ||
|
Maxime Ellerbach
|
92f96f33b3 | ||
|
Steven Palma
|
d4b3ca569c | ||
|
Steven Palma
|
3f2179f3b6 | ||
|
Nikodem Bartnik
|
867b58cfb2 |
@@ -1,155 +0,0 @@
|
||||
"""Classify each sub-dataset: is it SO-100/101, and what joint encoding is it in?
|
||||
|
||||
Detection = robot_type string (recording-time signal) cross-checked against the
|
||||
per-episode stats min/max (magnitude + exact-boundary saturation). Mismatches are
|
||||
flagged as `ambiguous` for manual review rather than silently converted.
|
||||
"""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
# so100/so101 (+ _follower/_bimanual), so_follower, and bimanual bi_so* (bi_so_follower,
|
||||
# bi_so100_follower, ...; 12-dim).
|
||||
SO_PREFIXES = ("so100", "so101", "so_", "bi_so")
|
||||
SO_EXACT: set[str] = set()
|
||||
# Robots that superficially look SO-like but are NOT in scope for the joint fix:
|
||||
NEVER_FIX = {"koch", "koch_follower", "koch_bimanual", "moss", "moss_follower"}
|
||||
|
||||
RAD_MAX = 3.5 # |val| below this => radians
|
||||
DEG_MIN = 105.0 # |val| above this => old-convention degrees
|
||||
SAT_ATOL = 0.5 # closeness to +/-100 / 0 / 100 counted as normalization saturation
|
||||
|
||||
|
||||
def is_so_robot_type(rt: str) -> bool:
|
||||
"""True if the recorded ``robot_type`` denotes an in-scope SO-100/101 arm."""
|
||||
return bool(rt) and (rt.startswith(SO_PREFIXES) or rt in SO_EXACT) and rt not in NEVER_FIX
|
||||
|
||||
|
||||
SO_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper")
|
||||
|
||||
|
||||
def is_end_effector(info: dict) -> bool:
|
||||
"""True if action/observation.state are task-space end-effector features (e.g. ``ee_x``,
|
||||
``ee_roll``) rather than joint angles. Such datasets are out of scope for the joint fix."""
|
||||
feats = info.get("features", {})
|
||||
for key in ("action", "observation.state"):
|
||||
names = [str(n).lower() for n in (feats.get(key, {}).get("names") or [])]
|
||||
if any(n.startswith("ee_") or "end_effector" in n or "eef" in n for n in names):
|
||||
return True
|
||||
if {"x", "y", "z"} <= set(names):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _leading_so_joints(names: list[str], dim: int) -> int:
|
||||
"""Number of LEADING joints (a multiple of 6) that match the SO joint order in blocks of 6.
|
||||
When names are absent, fall back to the full dim if it's already a multiple of 6, else 0."""
|
||||
if not names:
|
||||
return dim if dim and dim % 6 == 0 else 0
|
||||
k = 0
|
||||
while (k + 1) * 6 <= len(names) and all(SO_JOINTS[i] in names[k * 6 + i] for i in range(6)):
|
||||
k += 1
|
||||
return k * 6
|
||||
|
||||
|
||||
def so_joint_count(info: dict, key: str) -> int:
|
||||
"""Leading SO-arm joint count for one feature (``action`` / ``observation.state``). Trailing
|
||||
non-SO columns (bbox, appended EE pose, ...) are excluded so only the genuine SO joints are
|
||||
ever degrees-converted."""
|
||||
feat = info.get("features", {}).get(key, {})
|
||||
dim = (feat.get("shape") or [0])[0]
|
||||
names = [str(n).lower() for n in (feat.get("names") or [])]
|
||||
return _leading_so_joints(names, dim)
|
||||
|
||||
|
||||
def is_mislabeled_so(info: dict) -> bool:
|
||||
"""True when ``robot_type`` claims SO but no leading 6-DOF SO joint block can be substantiated
|
||||
from action/observation.state (wrong dim, or names that don't match the SO set). When the first
|
||||
6 joint names DO match, the SO block is honored (and processed) even if extra columns follow."""
|
||||
return max(so_joint_count(info, "action"), so_joint_count(info, "observation.state")) == 0
|
||||
|
||||
|
||||
def load_info(root: Path) -> dict:
|
||||
return json.loads((Path(root) / "meta" / "info.json").read_text())
|
||||
|
||||
|
||||
def _global_bounds(root: Path):
|
||||
"""Per-joint global min/max over action (fallback observation.state), across episodes."""
|
||||
lo = hi = None
|
||||
key_used = None
|
||||
with open(Path(root) / "meta" / "episodes_stats.jsonl") as f:
|
||||
for line in f:
|
||||
s = json.loads(line)["stats"]
|
||||
key = "action" if "action" in s else ("observation.state" if "observation.state" in s else None)
|
||||
if key is None:
|
||||
continue
|
||||
key_used = key
|
||||
mn = np.asarray(s[key]["min"], dtype=float)
|
||||
mx = np.asarray(s[key]["max"], dtype=float)
|
||||
lo = mn if lo is None else np.minimum(lo, mn)
|
||||
hi = mx if hi is None else np.maximum(hi, mx)
|
||||
return lo, hi, key_used
|
||||
|
||||
|
||||
def encoding_from_bounds(lo, hi, rt: str) -> dict:
|
||||
"""Detect the SO-arm joint encoding from per-joint global min/max and the robot_type name.
|
||||
|
||||
Layout-agnostic (v2.1 episodes_stats or v3.0 stats.json both reduce to lo/hi here), so it is
|
||||
the single source of truth for the degrees_old / degrees_new / normalized / radians decision.
|
||||
"""
|
||||
lo = np.asarray(lo, dtype=float)
|
||||
hi = np.asarray(hi, dtype=float)
|
||||
maxabs = float(np.nanmax(np.abs(np.concatenate([lo, hi]))))
|
||||
# saturation on any arm joint (index != gripper) at +/-100, or gripper at 0/100
|
||||
n = 6
|
||||
sat = False
|
||||
for a in range(len(hi) // n):
|
||||
arm_hi, arm_lo = hi[a * n:a * n + n], lo[a * n:a * n + n]
|
||||
joints_hi, joints_lo = arm_hi[:5], arm_lo[:5]
|
||||
grip_hi, grip_lo = arm_hi[5], arm_lo[5]
|
||||
sat |= bool(np.any(np.isclose(joints_hi, 100, atol=SAT_ATOL)) or
|
||||
np.any(np.isclose(joints_lo, -100, atol=SAT_ATOL)) or
|
||||
np.isclose(grip_hi, 100, atol=SAT_ATOL) or np.isclose(grip_lo, 0, atol=SAT_ATOL))
|
||||
|
||||
if maxabs <= RAD_MAX:
|
||||
enc = "radians"
|
||||
elif maxabs > DEG_MIN:
|
||||
enc = "degrees_old"
|
||||
elif sat:
|
||||
enc = "normalized"
|
||||
else:
|
||||
enc = "degrees_new"
|
||||
|
||||
name_says_new = rt.endswith(("_follower", "_bimanual"))
|
||||
ambiguous = (enc == "degrees_old" and name_says_new) or (enc in ("normalized", "degrees_new") and not name_says_new)
|
||||
return {"encoding": enc, "maxabs": round(maxabs, 2), "saturates": sat, "ambiguous": ambiguous}
|
||||
|
||||
|
||||
def classify(root) -> dict:
|
||||
root = Path(root)
|
||||
info = load_info(root)
|
||||
rt = info.get("robot_type", "") or ""
|
||||
dim = (info.get("features", {}).get("action", {}).get("shape") or [None])[0]
|
||||
out = {"root": str(root), "robot_type": rt, "action_dim": dim,
|
||||
"codebase_version": info.get("codebase_version"), "ambiguous": False}
|
||||
|
||||
if is_end_effector(info):
|
||||
return {**out, "is_so": False, "encoding": "end_effector",
|
||||
"note": "task-space end-effector features"}
|
||||
|
||||
if is_so_robot_type(rt) and is_mislabeled_so(info):
|
||||
return {**out, "is_so": False, "encoding": "non_so", "mislabeled_so": True,
|
||||
"note": "robot_type claims SO but joint dim/names don't match a 6-DOF SO arm"}
|
||||
|
||||
is_so = is_so_robot_type(rt)
|
||||
if not is_so:
|
||||
return {**out, "is_so": False, "encoding": "non_so"}
|
||||
|
||||
lo, hi, key_used = _global_bounds(root)
|
||||
if lo is None:
|
||||
return {**out, "is_so": True, "encoding": "unknown", "ambiguous": True,
|
||||
"note": "no action/state stats found"}
|
||||
|
||||
n = so_joint_count(info, key_used) or len(hi) # ignore trailing non-joint columns
|
||||
return {**out, "is_so": True, "stats_key": key_used, "so_dim": n,
|
||||
**encoding_from_bounds(lo[:n], hi[:n], rt)}
|
||||
@@ -1,55 +0,0 @@
|
||||
"""Extract one sub-dataset from a LeRobotDataset monorepo into a standalone HF dataset repo.
|
||||
|
||||
Downloads only ``{src-repo}/{folder}/...`` (scoped listing, no whole-repo walk) and re-uploads
|
||||
its contents at the ROOT of a NEW dataset repo, so the result is a self-contained LeRobotDataset.
|
||||
|
||||
python extract_dataset.py --folder wannrrr/etnai --dst-repo CarolinePascal/etnai
|
||||
"""
|
||||
import argparse
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
from run_migration import download_subfolder # noqa: E402
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
ap.add_argument("--folder", required=True, metavar="USER/DATASET",
|
||||
help="Sub-dataset path within the source monorepo, e.g. 'wannrrr/etnai'.")
|
||||
ap.add_argument("--dst-repo", required=True, metavar="ORG/NAME",
|
||||
help="New standalone destination dataset repo (must differ from the source).")
|
||||
ap.add_argument("--src-repo", default="lerobot/community_dataset_v3", metavar="ORG/NAME",
|
||||
help="Source monorepo to pull the sub-dataset from.")
|
||||
ap.add_argument("--work-dir", default="./extract_work", help="Local scratch directory.")
|
||||
ap.add_argument("--private", action="store_true", help="Create the destination repo as private.")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.dst_repo in (args.src_repo, args.folder):
|
||||
ap.error("--dst-repo must be a new repo name, distinct from the source repo/folder.")
|
||||
|
||||
local = Path(args.work_dir) / args.folder
|
||||
if local.parent.exists():
|
||||
shutil.rmtree(local.parent, ignore_errors=True)
|
||||
|
||||
print(f"downloading {args.src_repo}/{args.folder} ...", file=sys.stderr)
|
||||
download_subfolder(args.folder, args.work_dir, repo=args.src_repo)
|
||||
if not (local / "meta" / "info.json").exists():
|
||||
ap.error(f"'{args.folder}' is not a LeRobotDataset (no meta/info.json) in {args.src_repo}")
|
||||
|
||||
api = HfApi()
|
||||
api.create_repo(args.dst_repo, repo_type="dataset", private=args.private, exist_ok=True)
|
||||
print(f"uploading -> {args.dst_repo} ...", file=sys.stderr)
|
||||
api.upload_folder(repo_id=args.dst_repo, repo_type="dataset", folder_path=str(local),
|
||||
commit_message=f"Standalone copy of {args.folder} from {args.src_repo}")
|
||||
shutil.rmtree(Path(args.work_dir) / args.folder.split("/")[0], ignore_errors=True)
|
||||
print(f"done: https://huggingface.co/datasets/{args.dst_repo}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,260 +0,0 @@
|
||||
"""Rewrite observation.state / action to degrees in a LOCAL v2.1 SO-arm dataset, then
|
||||
regenerate meta/episodes_stats.jsonl (action & state only; other features preserved).
|
||||
Run this BEFORE the stock v2.1->v3.0 converter so its stats aggregation stays correct.
|
||||
"""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
import so_arm_frame
|
||||
from classify import classify, load_info, so_joint_count
|
||||
|
||||
VALUE_COLS = ("observation.state", "action")
|
||||
|
||||
|
||||
def _stack(col_values) -> np.ndarray:
|
||||
return np.stack([np.asarray(v, dtype=np.float64) for v in col_values]) # (N, D)
|
||||
|
||||
|
||||
def _set_robot_type(root: Path, robot_type: str) -> None:
|
||||
info_path = root / "meta" / "info.json"
|
||||
info = json.loads(info_path.read_text())
|
||||
info["robot_type"] = robot_type
|
||||
info_path.write_text(json.dumps(info, indent=4))
|
||||
|
||||
|
||||
def _rewrite_parquet(root: Path, encoding: str, so_dims: dict) -> None:
|
||||
for pq in sorted((root / "data").glob("*/*.parquet")):
|
||||
df = pd.read_parquet(pq)
|
||||
changed = False
|
||||
for col in VALUE_COLS:
|
||||
n = so_dims.get(col, 0)
|
||||
if col in df.columns and n:
|
||||
full = _stack(df[col].values) # (N, D)
|
||||
full[:, :n] = so_arm_frame.to_degrees(full[:, :n], encoding, n_joints_per_arm=6)
|
||||
df[col] = list(full.astype(np.float32))
|
||||
changed = True
|
||||
if changed:
|
||||
df.to_parquet(pq, index=False)
|
||||
|
||||
|
||||
def _regen_episode_stats(root: Path) -> None:
|
||||
stats_path = root / "meta" / "episodes_stats.jsonl"
|
||||
orig = {}
|
||||
with open(stats_path) as f:
|
||||
for line in f:
|
||||
e = json.loads(line)
|
||||
orig[e["episode_index"]] = e
|
||||
for pq in sorted((root / "data").glob("*/*.parquet")):
|
||||
df = pd.read_parquet(pq)
|
||||
for ep in np.unique(df["episode_index"].values):
|
||||
ep = int(ep)
|
||||
sub = df[df["episode_index"] == ep]
|
||||
entry = orig.get(ep)
|
||||
if entry is None:
|
||||
continue
|
||||
for col in VALUE_COLS:
|
||||
if col in sub.columns:
|
||||
a = _stack(sub[col].values) # (n, D)
|
||||
entry["stats"][col] = {
|
||||
"min": a.min(0).tolist(), "max": a.max(0).tolist(),
|
||||
"mean": a.mean(0).tolist(), "std": a.std(0).tolist(),
|
||||
"count": [int(a.shape[0])],
|
||||
}
|
||||
with open(stats_path, "w") as f:
|
||||
for ep in sorted(orig):
|
||||
f.write(json.dumps(orig[ep]) + "\n")
|
||||
|
||||
|
||||
def _read_jsonl(path: Path) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return [json.loads(line) for line in f if line.strip()]
|
||||
|
||||
|
||||
def _write_jsonl(path: Path, rows: list[dict]) -> None:
|
||||
with open(path, "w") as f:
|
||||
for r in rows:
|
||||
f.write(json.dumps(r) + "\n")
|
||||
|
||||
|
||||
def data_video_episode_mismatch(root) -> str | None:
|
||||
"""Return a description when the data files and any camera's video files disagree on the
|
||||
episode count (dataset can't be migrated, e.g. 'All cams dont have same number of episodes'),
|
||||
else None. Datasets without videos never mismatch here."""
|
||||
root = Path(root)
|
||||
info = json.loads((root / "meta" / "info.json").read_text())
|
||||
counts = {"data": len(list((root / "data").glob("*/episode_*.parquet")))}
|
||||
for k, f in info.get("features", {}).items():
|
||||
if f.get("dtype") == "video":
|
||||
counts[k] = len(list((root / "videos").glob(f"*/{k}/episode_*.mp4")))
|
||||
if len(counts) > 1 and len(set(counts.values())) > 1:
|
||||
return f"data/video episode counts disagree: {counts}"
|
||||
return None
|
||||
|
||||
|
||||
def _file_ep_indices(root: Path, pattern: str) -> list[int]:
|
||||
return sorted(int(p.stem.split("_")[-1]) for p in root.glob(pattern))
|
||||
|
||||
|
||||
def reindex_episodes(root) -> str | None:
|
||||
"""Compact non-contiguous episode indices to 0..N-1 when every source agrees on the set.
|
||||
|
||||
Some datasets (e.g. '*_clean' variants) had episodes deleted, leaving gaps in the episode
|
||||
numbering (data, videos, and metadata all skip the same indices, e.g. {20, 37, 38, 39}). The
|
||||
stock v2.1->v3.0 converter renumbers data/videos by sorted file order (0..N-1) but reads the
|
||||
original gapped indices from episodes.jsonl, so the two disagree and it raises
|
||||
"Number of episodes is not the same". When the data files, every camera's videos, and both
|
||||
metadata files list the *exact same* episode index set, remap it to 0..N-1 everywhere so the
|
||||
converter's positional alignment holds. Returns a note if remapped, else None (already
|
||||
contiguous, or the sources disagree -> unsafe to touch)."""
|
||||
root = Path(root)
|
||||
info = json.loads((root / "meta" / "info.json").read_text())
|
||||
|
||||
ref = _file_ep_indices(root, "data/*/episode_*.parquet")
|
||||
if not ref:
|
||||
return None
|
||||
sources = {"data": ref}
|
||||
vkeys = [k for k, f in info.get("features", {}).items() if f.get("dtype") == "video"]
|
||||
for k in vkeys:
|
||||
sources[k] = _file_ep_indices(root, f"videos/*/{k}/episode_*.mp4")
|
||||
eps = _read_jsonl(root / "meta" / "episodes.jsonl")
|
||||
stats = _read_jsonl(root / "meta" / "episodes_stats.jsonl")
|
||||
sources["episodes"] = sorted(e["episode_index"] for e in eps)
|
||||
sources["episodes_stats"] = sorted(s["episode_index"] for s in stats)
|
||||
|
||||
if any(v != ref for v in sources.values()):
|
||||
return None # sources disagree on the episode set -> not safe to reindex here
|
||||
n = len(ref)
|
||||
if ref == list(range(n)):
|
||||
return None # already contiguous
|
||||
|
||||
remap = {old: new for new, old in enumerate(ref)}
|
||||
|
||||
# Data: rewrite episode_index (and rebuild the global 'index'), then rename the file. Ascending
|
||||
# order is collision-free because new <= old for every episode.
|
||||
running = 0
|
||||
for old in ref:
|
||||
matches = list((root / "data").glob(f"*/episode_{old:06d}.parquet"))
|
||||
if not matches:
|
||||
return None
|
||||
pq = matches[0]
|
||||
df = pd.read_parquet(pq)
|
||||
if "episode_index" in df.columns:
|
||||
df["episode_index"] = remap[old]
|
||||
if "index" in df.columns:
|
||||
df["index"] = np.arange(running, running + len(df), dtype=df["index"].dtype)
|
||||
running += len(df)
|
||||
df.to_parquet(pq, index=False)
|
||||
dst = pq.with_name(f"episode_{remap[old]:06d}.parquet")
|
||||
if dst != pq:
|
||||
pq.rename(dst)
|
||||
|
||||
# Videos: rename per camera (ascending -> collision-free).
|
||||
for k in vkeys:
|
||||
for old in ref:
|
||||
for mp4 in (root / "videos").glob(f"*/{k}/episode_{old:06d}.mp4"):
|
||||
dst = mp4.with_name(f"episode_{remap[old]:06d}.mp4")
|
||||
if dst != mp4:
|
||||
mp4.rename(dst)
|
||||
|
||||
for e in eps:
|
||||
e["episode_index"] = remap[e["episode_index"]]
|
||||
for s in stats:
|
||||
s["episode_index"] = remap[s["episode_index"]]
|
||||
_write_jsonl(root / "meta" / "episodes.jsonl", sorted(eps, key=lambda e: e["episode_index"]))
|
||||
_write_jsonl(root / "meta" / "episodes_stats.jsonl", sorted(stats, key=lambda s: s["episode_index"]))
|
||||
|
||||
info["total_episodes"] = n
|
||||
info["total_frames"] = int(running)
|
||||
if "total_videos" in info:
|
||||
info["total_videos"] = n * len(vkeys)
|
||||
info["splits"] = {"train": f"0:{n}"}
|
||||
(root / "meta" / "info.json").write_text(json.dumps(info, indent=4))
|
||||
return f"episode indices compacted to 0..{n - 1} (dropped gaps {sorted(set(range(ref[-1] + 1)) - set(ref))})"
|
||||
|
||||
|
||||
def reconcile_episode_count(root) -> str | None:
|
||||
"""When the data files and video files agree on an episode count N but the metadata lists a
|
||||
different count, rewrite the metadata (episodes.jsonl, episodes_stats.jsonl, info.json) to N.
|
||||
|
||||
Only the safe direction is handled: trimming metadata that lists MORE episodes than actually
|
||||
exist. If the data itself is non-contiguous, the videos disagree with the data, or the metadata
|
||||
lists FEWER episodes than the data (which would require fabricating per-episode stats), nothing
|
||||
is changed and the stock converter's mismatch error is left to surface. Returns a note on fix."""
|
||||
root = Path(root)
|
||||
info = json.loads((root / "meta" / "info.json").read_text())
|
||||
|
||||
data_idx = sorted(int(p.stem.split("_")[-1]) for p in (root / "data").glob("*/episode_*.parquet"))
|
||||
n = len(data_idx)
|
||||
if n == 0 or data_idx != list(range(n)):
|
||||
return None
|
||||
|
||||
vkeys = [k for k, f in info.get("features", {}).items() if f.get("dtype") == "video"]
|
||||
for k in vkeys:
|
||||
if len(list((root / "videos").glob(f"*/{k}/episode_*.mp4"))) != n:
|
||||
return None # data and videos disagree -> out of scope for this fix
|
||||
|
||||
eps_path = root / "meta" / "episodes.jsonl"
|
||||
stats_path = root / "meta" / "episodes_stats.jsonl"
|
||||
eps, stats = _read_jsonl(eps_path), _read_jsonl(stats_path)
|
||||
if len(eps) == n and len(stats) == n:
|
||||
return None
|
||||
|
||||
eps_keep = [e for e in eps if e.get("episode_index", -1) < n]
|
||||
stats_keep = [s for s in stats if s.get("episode_index", -1) < n]
|
||||
if len(eps_keep) != n or len(stats_keep) != n:
|
||||
return None # metadata is missing episodes present in the data -> can't safely fabricate
|
||||
|
||||
dropped = max(len(eps), len(stats)) - n
|
||||
_write_jsonl(eps_path, eps_keep)
|
||||
_write_jsonl(stats_path, stats_keep)
|
||||
info["total_episodes"] = n
|
||||
info["total_frames"] = int(sum(e.get("length", 0) for e in eps_keep))
|
||||
if "total_videos" in info:
|
||||
info["total_videos"] = n * len(vkeys)
|
||||
info["splits"] = {"train": f"0:{n}"}
|
||||
(root / "meta" / "info.json").write_text(json.dumps(info, indent=4))
|
||||
return f"metadata episode count reconciled to {n} (data & videos agree; dropped {dropped} stale meta entries)"
|
||||
|
||||
|
||||
def fix_dataset_in_place(root) -> dict:
|
||||
"""Returns the classification dict augmented with the action taken."""
|
||||
root = Path(root)
|
||||
cls = classify(root)
|
||||
if cls.get("mislabeled_so"):
|
||||
# robot_type claims SO but the joints prove otherwise (wrong dim or non-SO names).
|
||||
# Relabel to 'unknown' and migrate structurally rather than degrees-converting on a
|
||||
# false assumption; the joint values are left exactly as recorded.
|
||||
_set_robot_type(root, "unknown")
|
||||
return {**cls, "robot_type": "unknown", "converted": False,
|
||||
"action": f"structural v2.1->v3.0 only; robot_type relabeled '{cls.get('robot_type')}'"
|
||||
"->'unknown' (joints don't match a 6-DOF SO arm), joint values left unchanged"}
|
||||
enc = cls.get("encoding")
|
||||
if not cls.get("is_so") or enc in ("radians", "unknown", "non_so"):
|
||||
reason = {
|
||||
"non_so": "not an SO-100/101 dataset",
|
||||
"radians": "SO-arm joints already in radians",
|
||||
"unknown": "SO-arm but joint encoding could not be determined",
|
||||
}.get(enc, "no joint conversion applicable")
|
||||
return {**cls, "converted": False,
|
||||
"action": f"structural v2.1->v3.0 only ({reason}); joint values left unchanged"}
|
||||
if enc == "normalized" and not so_arm_frame.CANON_IS_CALIBRATED:
|
||||
# Without per-robot calibration the un-normalization is an identity (placeholder
|
||||
# spans == 100), so rewriting is pointless. Keep the normalized values as-is and let
|
||||
# the dataset card flag them APPROXIMATE instead.
|
||||
return {**cls, "converted": False,
|
||||
"action": "structural v2.1->v3.0 only; joint values kept in normalized units "
|
||||
"(-100..100 / 0..100), NOT converted to degrees (uncalibrated -> APPROXIMATE)"}
|
||||
# drop stray files that would otherwise be uploaded
|
||||
for junk in (root / "meta").glob("info.json.bak"):
|
||||
junk.unlink()
|
||||
info = load_info(root)
|
||||
so_dims = {c: so_joint_count(info, c) for c in VALUE_COLS}
|
||||
_rewrite_parquet(root, enc, so_dims)
|
||||
_regen_episode_stats(root)
|
||||
full_dims = {c: (info.get("features", {}).get(c, {}).get("shape") or [0])[0] for c in VALUE_COLS}
|
||||
partial = any(0 < so_dims[c] < full_dims[c] for c in VALUE_COLS)
|
||||
tail = " (leading SO joints only; trailing non-joint columns left unchanged)" if partial else ""
|
||||
return {**cls, "converted": True,
|
||||
"action": f"structural v2.1->v3.0 + joint values converted ({enc} -> degrees){tail}"}
|
||||
@@ -1,133 +0,0 @@
|
||||
"""Migrate the SO-100/101 datasets referenced by ``allenai/MolmoAct2-SO100_101-Dataset``.
|
||||
|
||||
That repo does NOT store the datasets themselves; it lists them. Each
|
||||
``language_annotations/{user}/{dataset}/...`` folder name is the HF repo id of a *standalone*
|
||||
LeRobotDataset. This script derives those repo ids, drops any already present in
|
||||
``lerobot/community_dataset_v3`` (already migrated) and in the destination (resume), then runs
|
||||
the exact same per-dataset pipeline as ``run_migration.py`` on each remaining standalone repo
|
||||
(download whole repo -> SO-arm joint fix -> v2.1->v3.0 convert -> card -> upload -> cleanup).
|
||||
|
||||
python migrate_molmoact.py --dst-repo lerobot/community_dataset_v3 --work-dir ./molmo_work
|
||||
Flags mirror run_migration.py: --only-classify, --no-push, --folder-name USER/DATASET [...],
|
||||
--limit N, --reference-repo (the "already migrated" set to skip against).
|
||||
"""
|
||||
import argparse
|
||||
import csv
|
||||
import shutil
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
from classify import classify # noqa: E402
|
||||
from run_migration import already_done, list_datasets, migrate_one # noqa: E402
|
||||
|
||||
LIST_REPO = "allenai/MolmoAct2-SO100_101-Dataset"
|
||||
REFERENCE_REPO = "lerobot/community_dataset_v3" # the "already migrated" set to skip against
|
||||
ANNOTATIONS_PREFIX = "language_annotations/"
|
||||
|
||||
|
||||
def list_molmoact_datasets(api: HfApi, repo: str = LIST_REPO) -> list[str]:
|
||||
"""Standalone dataset repo ids (``{user}/{dataset}``) derived from the folder names under
|
||||
``language_annotations/`` in the MolmoAct listing repo."""
|
||||
files = api.list_repo_files(repo, repo_type="dataset")
|
||||
return sorted({"/".join(f.split("/")[1:3]) for f in files
|
||||
if f.startswith(ANNOTATIONS_PREFIX) and len(f.split("/")) >= 3})
|
||||
|
||||
|
||||
def pending_datasets(api: HfApi, subs: list[str], dst_repo: str | None,
|
||||
reference_repo: str, no_upload: bool, only_classify: bool) -> list[str]:
|
||||
"""Drop ids already in the reference repo (already migrated) and, unless classify/no-push,
|
||||
ids already in the destination repo (resume)."""
|
||||
skip = set(list_datasets(api, reference_repo))
|
||||
if not only_classify and not no_upload and dst_repo and dst_repo != reference_repo:
|
||||
skip |= {p[: -len("/meta/info.json")] for p in api.list_repo_files(dst_repo, repo_type="dataset")
|
||||
if p.endswith("/meta/info.json")}
|
||||
return [s for s in subs if s not in skip]
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Migrate the standalone SO-100/101 datasets listed by "
|
||||
f"{LIST_REPO} to LeRobotDataset v3.0 (degrees), skipping any already present "
|
||||
f"in --reference-repo. One dataset at a time (download -> fix -> convert -> "
|
||||
"upload -> cleanup); resumable.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
ap.add_argument("--dst-repo", default=REFERENCE_REPO, metavar="ORG/NAME",
|
||||
help="Destination HF dataset repo to push the converted v3.0 datasets to "
|
||||
"(created if missing).")
|
||||
ap.add_argument("--reference-repo", default=REFERENCE_REPO, metavar="ORG/NAME",
|
||||
help="Repo whose datasets are considered already migrated and skipped.")
|
||||
ap.add_argument("--work-dir", default="./molmo_work", metavar="DIR",
|
||||
help="Local scratch directory (one dataset lives here at a time on a push run).")
|
||||
ap.add_argument("--manifest", default="manifest_molmoact.csv", metavar="CSV",
|
||||
help="CSV log appended to as datasets are processed. Reused across resumed runs.")
|
||||
ap.add_argument("--limit", type=int, default=None, metavar="N",
|
||||
help="Process only the first N pending datasets (alphabetical). Ignored with "
|
||||
"--folder-name.")
|
||||
ap.add_argument("--folder-name", nargs="+", default=None, metavar="USER/DATASET",
|
||||
help="One or more specific standalone repo ids to process (must appear in the "
|
||||
f"{LIST_REPO} listing).")
|
||||
ap.add_argument("--only-classify", action="store_true",
|
||||
help="Detect robot type + joint encoding and write the manifest only; no "
|
||||
"download of data, convert, or push.")
|
||||
ap.add_argument("--no-push", action="store_true",
|
||||
help="Fix + convert locally but do NOT upload; output kept under --work-dir.")
|
||||
args = ap.parse_args()
|
||||
no_upload = args.no_push
|
||||
|
||||
api = HfApi()
|
||||
all_ids = list_molmoact_datasets(api)
|
||||
if args.folder_name:
|
||||
wanted = {n.strip("/") for n in args.folder_name}
|
||||
subs = [s for s in all_ids if s in wanted]
|
||||
missing = wanted - set(subs)
|
||||
if missing:
|
||||
print(f"warning: not in {LIST_REPO} listing: {', '.join(sorted(missing))}", file=sys.stderr)
|
||||
else:
|
||||
subs = pending_datasets(api, all_ids, args.dst_repo, args.reference_repo, no_upload, args.only_classify)
|
||||
if args.limit:
|
||||
subs = subs[: args.limit]
|
||||
print(f"{len(subs)} dataset(s) to process (of {len(all_ids)} listed)", file=sys.stderr)
|
||||
if not subs:
|
||||
return
|
||||
|
||||
if not args.only_classify and not no_upload:
|
||||
api.create_repo(args.dst_repo, repo_type="dataset", exist_ok=True)
|
||||
dst_files = set() if (args.only_classify or no_upload) else set(
|
||||
api.list_repo_files(args.dst_repo, repo_type="dataset"))
|
||||
|
||||
first = not Path(args.manifest).exists()
|
||||
with open(args.manifest, "a", newline="") as mf:
|
||||
w = None
|
||||
for i, sub in enumerate(subs):
|
||||
try:
|
||||
if args.only_classify:
|
||||
from huggingface_hub import snapshot_download
|
||||
local = Path(args.work_dir) / sub
|
||||
snapshot_download(repo_id=sub, repo_type="dataset", local_dir=str(local),
|
||||
allow_patterns=["meta/*"])
|
||||
row = {"root": sub, **classify(local)}
|
||||
shutil.rmtree(Path(args.work_dir) / sub.split("/")[0], ignore_errors=True)
|
||||
elif not no_upload and already_done(api, args.dst_repo, sub, dst_files):
|
||||
row = {"root": sub, "action": "skipped: already present in destination repo"}
|
||||
else:
|
||||
row = migrate_one(api, args.dst_repo, sub, args.work_dir, no_upload, standalone=True)
|
||||
except Exception as e:
|
||||
row = {"root": sub, "action": f"ERROR: {e}"}
|
||||
traceback.print_exc()
|
||||
if w is None:
|
||||
w = csv.DictWriter(mf, fieldnames=sorted(
|
||||
{"root", "robot_type", "is_so", "encoding", "action_dim",
|
||||
"maxabs", "ambiguous", "action", "codebase_version", "note"}))
|
||||
if first:
|
||||
w.writeheader()
|
||||
w.writerow({k: row.get(k) for k in w.fieldnames})
|
||||
mf.flush()
|
||||
print(f"[{i+1}/{len(subs)}] {sub}: {row.get('action')}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,137 +0,0 @@
|
||||
"""Prune datasets from the migrated destination repo based on the run manifest(s).
|
||||
|
||||
Selects, from the manifest CSV(s) written by run_migration.py / slurm_migrate.py, the datasets to
|
||||
remove and deletes their folders from the destination repo. Two independent filters:
|
||||
|
||||
--errored rows whose migration action starts with "ERROR:" (failed to convert; usually
|
||||
absent from the repo, but any partial upload left behind is cleaned up).
|
||||
--mislabeled-so rows whose robot_type claims SO but the dataset isn't a real 6-DOF SO arm
|
||||
(action_dim not a multiple of 6, or a classification note saying so).
|
||||
|
||||
Only folders actually present in the destination repo are touched. Dry-run by default; pass --yes
|
||||
to perform the deletions.
|
||||
|
||||
python prune_destination.py --manifest /fsx/$USER/cdv3_manifests/manifest_*.csv \
|
||||
--dst-repo lerobot/community_dataset_v3 --errored --mislabeled-so # dry-run
|
||||
python prune_destination.py --manifest manifest_*.csv --errored --mislabeled-so --yes
|
||||
"""
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pandas as pd
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
from classify import is_so_robot_type
|
||||
|
||||
DST_REPO = "lerobot/community_dataset_v3"
|
||||
|
||||
|
||||
def _is_errored(row) -> bool:
|
||||
return str(row.get("action") or "").strip().upper().startswith("ERROR")
|
||||
|
||||
|
||||
def _is_mislabeled_so(row) -> bool:
|
||||
if "claims SO but" in str(row.get("note") or ""):
|
||||
return True
|
||||
if not is_so_robot_type(str(row.get("robot_type") or "")):
|
||||
return False
|
||||
try:
|
||||
return int(float(row["action_dim"])) % 6 != 0
|
||||
except (TypeError, ValueError, KeyError):
|
||||
return False
|
||||
|
||||
|
||||
def select(df: pd.DataFrame, present: set[str], errored: bool, mislabeled: bool) -> dict[str, str]:
|
||||
"""Map each dataset root that is present in the repo AND matches an enabled filter to a reason.
|
||||
'errored' takes precedence over 'mislabeled-so' when a row matches both."""
|
||||
out: dict[str, str] = {}
|
||||
for _, row in df.iterrows():
|
||||
root = str(row.get("root") or "").strip()
|
||||
if not root or root not in present or root in out:
|
||||
continue
|
||||
if errored and _is_errored(row):
|
||||
out[root] = "errored"
|
||||
elif mislabeled and _is_mislabeled_so(row):
|
||||
out[root] = "mislabeled-so"
|
||||
return out
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--manifest", nargs="+", required=True, metavar="CSV",
|
||||
help="One or more manifest CSVs from the migration run (per-rank files ok).")
|
||||
ap.add_argument("--dst-repo", default=DST_REPO, metavar="ORG/NAME",
|
||||
help="Destination HF dataset repo to prune.")
|
||||
ap.add_argument("--errored", action="store_true",
|
||||
help="Delete datasets that errored during migration.")
|
||||
ap.add_argument("--mislabeled-so", action="store_true",
|
||||
help="Delete datasets labelled SO but not a real 6-DOF SO arm.")
|
||||
ap.add_argument("--yes", action="store_true",
|
||||
help="Actually delete. Without it, only prints what would be deleted (dry-run).")
|
||||
ap.add_argument("--list-missing", action="store_true",
|
||||
help="Report datasets in the manifest that are NOT present in the destination "
|
||||
"repo (grouped by their migration action), then exit without deleting.")
|
||||
ap.add_argument("--report", action="store_true",
|
||||
help="List all three categories (errored, mislabeled-so, missing) without "
|
||||
"deleting anything, then exit. '*' marks datasets absent from the repo.")
|
||||
args = ap.parse_args()
|
||||
if not (args.errored or args.mislabeled_so or args.list_missing or args.report):
|
||||
ap.error("enable at least one of: --errored, --mislabeled-so, --list-missing, --report")
|
||||
|
||||
df = pd.concat([pd.read_csv(p) for p in args.manifest], ignore_index=True)
|
||||
if "root" not in df.columns:
|
||||
ap.error("manifest has no 'root' column; is this a run_migration.py manifest?")
|
||||
df = df.drop_duplicates(subset="root", keep="last")
|
||||
|
||||
api = HfApi()
|
||||
present = {p[: -len("/meta/info.json")]
|
||||
for p in api.list_repo_files(args.dst_repo, repo_type="dataset")
|
||||
if p.endswith("/meta/info.json")}
|
||||
|
||||
if args.report or args.list_missing:
|
||||
def _dump(title, sub):
|
||||
print(f"== {title} ({len(sub)}) ==")
|
||||
for root in sorted(sub):
|
||||
print(f" {' ' if root in present else '*'} {root}")
|
||||
print()
|
||||
missing = set(df.loc[~df["root"].isin(present), "root"])
|
||||
if args.report:
|
||||
_dump("errored", set(df.loc[df.apply(_is_errored, axis=1), "root"]))
|
||||
_dump("mislabeled-so", set(df.loc[df.apply(_is_mislabeled_so, axis=1), "root"]))
|
||||
_dump("missing from repo", missing)
|
||||
print(f"{df['root'].nunique()} unique manifest rows, {len(present)} datasets in "
|
||||
f"{args.dst_repo}. '*' = absent from repo.", file=sys.stderr)
|
||||
return
|
||||
|
||||
to_delete = select(df, present, args.errored, args.mislabeled_so)
|
||||
for root, reason in sorted(to_delete.items()):
|
||||
print(f"{reason:14s} {root}")
|
||||
print(f"\n{len(to_delete)} dataset(s) present in {args.dst_repo} match "
|
||||
f"({df['root'].nunique()} rows in manifest, {len(present)} datasets in repo).", file=sys.stderr)
|
||||
|
||||
if not args.yes:
|
||||
print("dry-run: nothing deleted. re-run with --yes to delete.", file=sys.stderr)
|
||||
return
|
||||
failed = []
|
||||
for root, reason in sorted(to_delete.items()):
|
||||
for attempt in range(1, 4): # transient Hub ReadTimeouts are common; retry with backoff
|
||||
try:
|
||||
api.delete_folder(path_in_repo=root, repo_id=args.dst_repo, repo_type="dataset",
|
||||
commit_message=f"Prune {root} ({reason})")
|
||||
print(f"deleted {root} ({reason})", file=sys.stderr)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == 3:
|
||||
failed.append(root)
|
||||
print(f"FAILED {root}: {e}", file=sys.stderr)
|
||||
else:
|
||||
time.sleep(2 ** attempt)
|
||||
if failed:
|
||||
print(f"\n{len(failed)} deletion(s) failed (likely transient); safe to re-run to retry: "
|
||||
f"{failed}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,300 +0,0 @@
|
||||
"""End-to-end migration of the community_dataset_v3 monorepo to v3.0 + SO-arm degrees.
|
||||
|
||||
For each `{user}/{dataset}` sub-dataset: stream-download it, fix SO-arm joint values
|
||||
(if applicable), run the stock v2.1->v3.0 structural converter locally, upload the v3.0
|
||||
result under the same path into a NEW repo, then delete the local copy. Resumable.
|
||||
|
||||
uv run python run_migration.py --dst-repo HuggingFaceVLA/community_dataset_v3_degrees \
|
||||
--work-dir /big/disk/cdv3_work --manifest manifest.csv
|
||||
Flags: --only-classify (just write manifest), --no-push (fix+convert locally, keep output,
|
||||
no upload), --folder-name A [B ...] (target specific dataset folders), --limit N.
|
||||
Uncalibrated `normalized` datasets keep their normalized joint units (flagged APPROXIMATE on
|
||||
the card); paste fitted CANON ranges in so_arm_frame.py to convert them to degrees instead.
|
||||
"""
|
||||
import argparse, csv, json, shutil, sys, traceback
|
||||
from pathlib import Path
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
import so_arm_frame
|
||||
from classify import classify, is_end_effector, load_info
|
||||
from fix_dataset import (
|
||||
data_video_episode_mismatch,
|
||||
fix_dataset_in_place,
|
||||
reconcile_episode_count,
|
||||
reindex_episodes,
|
||||
)
|
||||
|
||||
SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
|
||||
NORMALIZED_TAG = "normalized" # card tag for datasets whose SO joints stay in normalized units
|
||||
|
||||
|
||||
def download_subfolder(sub: str, work_dir: str, patterns: list[str] | None = None, repo: str = SRC_REPO) -> None:
|
||||
"""Download only ``{repo}/{sub}/...`` into ``work_dir``.
|
||||
|
||||
``snapshot_download`` walks the entire repo tree (``list_repo_tree(recursive=True)``
|
||||
with no path scope) before applying ``allow_patterns``. On this 791-dataset monorepo
|
||||
that whole-repo enumeration is pathologically slow and looks like a hang. Listing the
|
||||
scoped ``path_in_repo=sub`` subtree and fetching its files directly avoids it.
|
||||
"""
|
||||
from fnmatch import fnmatch
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
from huggingface_hub.hf_api import RepoFile
|
||||
|
||||
api = HfApi()
|
||||
for entry in api.list_repo_tree(repo, path_in_repo=sub, repo_type="dataset", recursive=True):
|
||||
if not isinstance(entry, RepoFile):
|
||||
continue
|
||||
if patterns and not any(fnmatch(entry.path, pat) for pat in patterns):
|
||||
continue
|
||||
hf_hub_download(repo, filename=entry.path, repo_type="dataset", local_dir=work_dir)
|
||||
|
||||
|
||||
def list_datasets(api: HfApi, repo: str) -> list[str]:
|
||||
files = api.list_repo_files(repo, repo_type="dataset")
|
||||
roots = {p[: -len("/meta/info.json")] for p in files if p.endswith("/meta/info.json")}
|
||||
return sorted(roots)
|
||||
|
||||
|
||||
def resolve_folders(api: HfApi, repo: str, names: list[str]) -> list[str]:
|
||||
"""Expand each --folder-name into concrete dataset roots (folders that contain
|
||||
meta/info.json). A name may be a full dataset path (returned as-is) or a namespace/prefix
|
||||
like 'Beegbrain' (expanded to every dataset beneath it). Unknown names pass through so
|
||||
they surface as a clear per-item error instead of a confusing FileNotFoundError."""
|
||||
out: list[str] = []
|
||||
for name in names:
|
||||
name = name.strip("/")
|
||||
try:
|
||||
paths = [e.path for e in api.list_repo_tree(
|
||||
repo, path_in_repo=name, recursive=True, repo_type="dataset")]
|
||||
except Exception:
|
||||
out.append(name) # let it fail loudly downstream
|
||||
continue
|
||||
roots = sorted({p[: -len("/meta/info.json")] for p in paths if p.endswith("/meta/info.json")})
|
||||
out.extend(roots or [name])
|
||||
seen: set[str] = set()
|
||||
return [r for r in out if not (r in seen or seen.add(r))]
|
||||
|
||||
|
||||
def already_done(api: HfApi, dst: str, sub: str, dst_files: set[str]) -> bool:
|
||||
return f"{sub}/meta/info.json" in dst_files # present in target => skip (resume)
|
||||
|
||||
|
||||
def _write_dataset_card(local: Path, sub: str, result: dict, standalone: bool = False) -> None:
|
||||
"""Regenerate the sub-dataset's card the way LeRobot does (create_lerobot_dataset_card
|
||||
from meta/info.json), then append a migration section documenting provenance and the
|
||||
joint-encoding fix. When ``standalone`` is set, ``sub`` is itself the source dataset's HF
|
||||
repo id (rather than a folder inside the ``SRC_REPO`` monorepo)."""
|
||||
enc = result.get("encoding")
|
||||
converted_degrees = bool(result.get("converted"))
|
||||
approx = enc == "normalized" and not so_arm_frame.CANON_IS_CALIBRATED
|
||||
|
||||
enc_labels = {
|
||||
"degrees_old": "legacy degrees (old community frame, pre-#777 convention)",
|
||||
"degrees_new": "degrees (recorded with `use_degrees=True`)",
|
||||
"normalized": "normalized units (joints -100..100, gripper 0..100)",
|
||||
"radians": "radians",
|
||||
"unknown": "undetermined",
|
||||
}
|
||||
joint_actions = {
|
||||
"degrees_old": "per-joint offsets and axis directions corrected to the post-#777 frame (values stay in degrees)",
|
||||
"degrees_new": "already in the post-#777 degrees frame; values unchanged",
|
||||
"normalized": ("un-normalized to physical degrees using calibrated joint ranges"
|
||||
if converted_degrees else
|
||||
"left in normalized units (-100..100 joints, 0..100 gripper); NOT converted to degrees"),
|
||||
"radians": "left unchanged (already in radians)",
|
||||
"unknown": "left unchanged (encoding could not be determined)",
|
||||
}
|
||||
|
||||
lines = [
|
||||
"## Migration to LeRobotDataset v3.0",
|
||||
"",
|
||||
"Migrated to LeRobotDataset **v3.0**"
|
||||
+ (" with SO-100/101 joint state/action mapped to the post-#777 physical frame (in degrees)."
|
||||
if converted_degrees else "."),
|
||||
"",
|
||||
(f"- Source: [`{sub}`](https://huggingface.co/datasets/{sub})" if standalone else
|
||||
f"- Source: [`{SRC_REPO}`](https://huggingface.co/datasets/{SRC_REPO}/tree/main/{sub}) (`{sub}`)"),
|
||||
"- Codebase version: v2.1 -> v3.0",
|
||||
]
|
||||
if result.get("is_so"):
|
||||
lines += [
|
||||
f"- Original joint encoding: {enc_labels.get(enc, enc)}",
|
||||
f"- Joint values: {joint_actions.get(enc, 'left unchanged')}",
|
||||
f"- Robot type: `{result.get('robot_type')}`",
|
||||
f"- Action dimension: {result.get('action_dim')}",
|
||||
]
|
||||
else:
|
||||
lines += ["- Joint values: not applicable (not an SO-100/101 dataset)"]
|
||||
if approx:
|
||||
lines += ["", "> **Note:** per-robot calibration was unavailable, so joint state/action were "
|
||||
"left in their original *normalized* units (-100..100 joints, 0..100 gripper) rather "
|
||||
"than converted to physical degrees. Treat these joint values as APPROXIMATE."]
|
||||
if result.get("ambiguous"):
|
||||
lines += ["", "> **Note:** joint-encoding detection was flagged ambiguous; conversion used the "
|
||||
"best-guess encoding above and may warrant manual review."]
|
||||
|
||||
section = "\n".join(lines) + "\n"
|
||||
|
||||
readme = local / "README.md"
|
||||
try:
|
||||
try:
|
||||
from lerobot.datasets.utils import create_lerobot_dataset_card
|
||||
except ImportError:
|
||||
from lerobot.common.datasets.utils import create_lerobot_dataset_card
|
||||
|
||||
class _Info(dict): # satisfies both the dict and .to_dict() card variants
|
||||
def to_dict(self):
|
||||
return dict(self)
|
||||
|
||||
rt = result.get("robot_type") or None
|
||||
tags = [rt] if rt else []
|
||||
if enc == "normalized" and not converted_degrees:
|
||||
tags.append(NORMALIZED_TAG)
|
||||
card = create_lerobot_dataset_card(
|
||||
tags=tags or None,
|
||||
dataset_info=_Info(load_info(local)),
|
||||
license="apache-2.0",
|
||||
repo_id=sub,
|
||||
)
|
||||
card.text = card.text.rstrip() + "\n\n" + section
|
||||
card.save(str(readme))
|
||||
except Exception:
|
||||
# LeRobot card generator unavailable at runtime: keep the standalone migration note.
|
||||
if readme.exists():
|
||||
readme.write_text(readme.read_text().rstrip() + "\n\n" + section)
|
||||
else:
|
||||
readme.write_text(f"# {sub}\n\n" + section)
|
||||
|
||||
|
||||
def migrate_one(api, dst_repo, sub, work_dir, no_upload, src_repo: str = SRC_REPO,
|
||||
standalone: bool = False) -> dict:
|
||||
local = Path(work_dir) / sub
|
||||
if local.parent.exists():
|
||||
shutil.rmtree(local.parent, ignore_errors=True) # clean any partial
|
||||
if standalone:
|
||||
# ``sub`` is a self-contained HF dataset repo (not a monorepo folder): pull it whole.
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id=sub, repo_type="dataset", local_dir=str(local))
|
||||
else:
|
||||
download_subfolder(sub, work_dir, repo=src_repo)
|
||||
|
||||
info = load_info(local)
|
||||
if info.get("codebase_version") != "v2.1":
|
||||
return {"root": sub, "action": f"skipped: source codebase is {info.get('codebase_version')} (expected v2.1)"}
|
||||
if is_end_effector(info):
|
||||
return {"root": sub, "robot_type": info.get("robot_type"),
|
||||
"action": "skipped: end-effector (task-space) dataset, out of scope"}
|
||||
mismatch = data_video_episode_mismatch(local)
|
||||
if mismatch:
|
||||
return {"root": sub, "robot_type": info.get("robot_type"),
|
||||
"action": f"skipped: {mismatch}"}
|
||||
|
||||
result = fix_dataset_in_place(local) # SO-arm value fix (or structural_only)
|
||||
|
||||
reconciled = reconcile_episode_count(local) # align stale meta counts to data+video files
|
||||
if reconciled:
|
||||
result["action"] = f"{result['action']}; {reconciled}"
|
||||
|
||||
reindexed = reindex_episodes(local) # compact non-contiguous episode indices to 0..N-1
|
||||
if reindexed:
|
||||
result["action"] = f"{result['action']}; {reindexed}"
|
||||
|
||||
from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset
|
||||
convert_dataset(repo_id=sub, root=str(local), push_to_hub=False) # v2.1 -> v3.0, in place
|
||||
|
||||
_write_dataset_card(local, sub, result, standalone=standalone) # document conversion in the card
|
||||
|
||||
base = {k: result.get(k) for k in
|
||||
("robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action")}
|
||||
if no_upload:
|
||||
# keep the converted output on disk for inspection; do NOT delete or push
|
||||
base["action"] = f"{base['action']}; not pushed (kept locally at {local})"
|
||||
return {"root": sub, **base}
|
||||
api.upload_folder(repo_id=dst_repo, repo_type="dataset", folder_path=str(local),
|
||||
path_in_repo=sub, commit_message=f"Add {sub} (v3.0, {result['action']})")
|
||||
shutil.rmtree(Path(work_dir) / sub.split("/")[0], ignore_errors=True) # drop after successful push
|
||||
return {"root": sub, **base}
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Migrate the HuggingFaceVLA/community_dataset_v3 monorepo to LeRobotDataset "
|
||||
"v3.0, converting SO-100/101 joint state/action to physical degrees along "
|
||||
"the way. Processes one sub-dataset at a time (download -> fix -> convert -> "
|
||||
"upload -> cleanup) and is resumable.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
ap.add_argument("--dst-repo", default=None, metavar="ORG/NAME",
|
||||
help="Destination HF dataset repo to push the converted v3.0 datasets to "
|
||||
"(created if missing). Required unless --no-push or --only-classify.")
|
||||
ap.add_argument("--work-dir", default="./cdv3_work", metavar="DIR",
|
||||
help="Local scratch directory used to download, convert, and (unless pushing) "
|
||||
"retain each sub-dataset. Only one dataset lives here at a time on a push run.")
|
||||
ap.add_argument("--manifest", default="manifest.csv", metavar="CSV",
|
||||
help="CSV log appended to as datasets are processed (robot_type, detected "
|
||||
"encoding, action taken, errors). Reused across resumed runs.")
|
||||
ap.add_argument("--limit", type=int, default=None, metavar="N",
|
||||
help="Process only the first N sub-datasets (alphabetical). Ignored when "
|
||||
"--folder-name is given. Useful for a quick end-to-end smoke test.")
|
||||
ap.add_argument("--folder-name", nargs="+", default=None, metavar="USER/DATASET",
|
||||
help=f"One or more folders WITHIN the {SRC_REPO} monorepo to process. Either a "
|
||||
"full dataset path ('Beegbrain/draw_pixel_art') or a whole namespace "
|
||||
"('Beegbrain'), which expands to every dataset under it. Skips the full "
|
||||
"791-dataset listing.")
|
||||
ap.add_argument("--only-classify", action="store_true",
|
||||
help="Detect each dataset's robot type and joint encoding and write the "
|
||||
"manifest, without downloading data, converting, or pushing. Run this "
|
||||
"first to review scope (especially rows flagged ambiguous=True).")
|
||||
ap.add_argument("--no-push", action="store_true",
|
||||
help="Fix + convert locally but do NOT upload; the converted v3.0 output is "
|
||||
"kept under --work-dir for inspection instead of being deleted.")
|
||||
args = ap.parse_args()
|
||||
no_upload = args.no_push
|
||||
if not no_upload and not args.only_classify and not args.dst_repo:
|
||||
ap.error("--dst-repo is required unless --no-push or --only-classify is set.")
|
||||
|
||||
api = HfApi()
|
||||
if args.folder_name:
|
||||
subs = resolve_folders(api, SRC_REPO, args.folder_name)
|
||||
print(f"targeting {len(subs)} sub-dataset(s): {', '.join(subs)}", file=sys.stderr)
|
||||
else:
|
||||
subs = list_datasets(api, SRC_REPO)
|
||||
if args.limit:
|
||||
subs = subs[: args.limit]
|
||||
print(f"{len(subs)} sub-datasets found", file=sys.stderr)
|
||||
|
||||
if not args.only_classify and not no_upload:
|
||||
api.create_repo(args.dst_repo, repo_type="dataset", exist_ok=True)
|
||||
dst_files = set() if (args.only_classify or no_upload) else set(
|
||||
api.list_repo_files(args.dst_repo, repo_type="dataset"))
|
||||
|
||||
first = not Path(args.manifest).exists()
|
||||
with open(args.manifest, "a", newline="") as mf:
|
||||
w = None
|
||||
for i, sub in enumerate(subs):
|
||||
try:
|
||||
if args.only_classify:
|
||||
# classify without full download: fetch just the meta/ of this sub
|
||||
download_subfolder(sub, args.work_dir, patterns=[f"{sub}/meta/*"])
|
||||
row = {"root": sub, **classify(Path(args.work_dir) / sub)}
|
||||
shutil.rmtree(Path(args.work_dir) / sub.split("/")[0], ignore_errors=True)
|
||||
elif not no_upload and already_done(api, args.dst_repo, sub, dst_files):
|
||||
row = {"root": sub, "action": "skipped: already present in destination repo"}
|
||||
else:
|
||||
row = migrate_one(api, args.dst_repo, sub, args.work_dir, no_upload)
|
||||
except Exception as e:
|
||||
row = {"root": sub, "action": f"ERROR: {e}"}
|
||||
traceback.print_exc()
|
||||
if w is None:
|
||||
w = csv.DictWriter(mf, fieldnames=sorted(
|
||||
{"root", "robot_type", "is_so", "encoding", "action_dim",
|
||||
"maxabs", "ambiguous", "action", "codebase_version", "note"}))
|
||||
if first:
|
||||
w.writeheader()
|
||||
w.writerow({k: row.get(k) for k in w.fieldnames})
|
||||
mf.flush()
|
||||
print(f"[{i+1}/{len(subs)}] {sub}: {row.get('action')}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,285 +0,0 @@
|
||||
"""SLURM-distributed driver for run_migration.py.
|
||||
|
||||
Fans the ``community_dataset_v3`` -> v3.0 migration out across SLURM workers, mirroring the
|
||||
datatrove pattern used by ``examples/dataset/slurm_recompute_stats.py``. This is a *map-only*
|
||||
job: each worker owns a stride ``subs[rank::world_size]`` of the sub-datasets and runs the
|
||||
exact same per-dataset pipeline as ``run_migration.py`` (download -> fix -> v2.1->v3.0 convert
|
||||
-> upload -> cleanup). There is no aggregate step.
|
||||
|
||||
Resume is twofold and free: (1) each worker skips any sub-dataset already present in the
|
||||
destination repo (``already_done``), and (2) datatrove skips ranks whose completion marker
|
||||
exists. Re-run the identical command to mop up failures.
|
||||
|
||||
Example (numeric smoke test on one namespace, no SLURM):
|
||||
python slurm_migrate.py --slurm 0 --workers 1 \
|
||||
--dst-repo HuggingFaceVLA/community_dataset_v3_degrees \
|
||||
--work-dir ./cdv3_work --manifest-dir ./cdv3_manifests \
|
||||
--folder-name Beegbrain
|
||||
|
||||
Full run on the cluster:
|
||||
python slurm_migrate.py \
|
||||
--dst-repo HuggingFaceVLA/community_dataset_v3_degrees \
|
||||
--work-dir /fsx/$USER/cdv3_work \
|
||||
--manifest-dir /fsx/$USER/cdv3_manifests \
|
||||
--logs-dir /fsx/$USER/logs/cdv3_migrate \
|
||||
--workers 64 --partition hopper-cpu --qos normal \
|
||||
--cpus-per-task 4 --mem-per-cpu 4G \
|
||||
--env-command "source /fsx/$USER/venvs/lerobot/bin/activate; export HF_TOKEN=<token>"
|
||||
|
||||
IMPORTANT: workers must reach the internet (HF download + upload) and have a write-scoped
|
||||
HF token (HF_TOKEN) in --env-command. Keep --workers modest (many concurrent commits to one
|
||||
destination repo contend); rely on resume passes to clear transient upload failures.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from datatrove.executor import LocalPipelineExecutor
|
||||
from datatrove.executor.slurm import SlurmPipelineExecutor
|
||||
from datatrove.pipeline.base import PipelineStep
|
||||
|
||||
MIGRATION_DIR = str(Path(__file__).resolve().parent)
|
||||
|
||||
|
||||
class MigrateShard(PipelineStep):
|
||||
"""Each worker migrates its ``subs[rank::world_size]`` slice of sub-datasets."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
subs,
|
||||
dst_repo,
|
||||
work_dir,
|
||||
manifest_dir,
|
||||
migration_dir,
|
||||
no_push=False,
|
||||
only_classify=False,
|
||||
standalone=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.subs = subs
|
||||
self.dst_repo = dst_repo
|
||||
self.work_dir = work_dir
|
||||
self.manifest_dir = manifest_dir
|
||||
self.migration_dir = migration_dir
|
||||
self.no_push = no_push
|
||||
self.only_classify = only_classify
|
||||
self.standalone = standalone
|
||||
|
||||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
||||
# Pickled onto the worker: keep self-contained. The migration package dir must be on
|
||||
# sys.path so ``run_migration`` and its siblings (classify/fix_dataset/so_arm_frame)
|
||||
# import.
|
||||
import csv
|
||||
import logging
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
if self.migration_dir not in sys.path:
|
||||
sys.path.insert(0, self.migration_dir)
|
||||
|
||||
from classify import classify
|
||||
from huggingface_hub import HfApi
|
||||
from run_migration import already_done, download_subfolder, migrate_one
|
||||
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
init_logging()
|
||||
|
||||
my_subs = self.subs[rank::world_size]
|
||||
if not my_subs:
|
||||
logging.info(f"Rank {rank}: no sub-datasets assigned")
|
||||
return
|
||||
logging.info(f"Rank {rank}: {len(my_subs)} / {len(self.subs)} sub-datasets")
|
||||
|
||||
# Per-rank scratch and manifest so workers never collide (migrate_one wipes
|
||||
# work_dir/<namespace> around each dataset).
|
||||
work_dir = str(Path(self.work_dir) / f"rank_{rank:05d}")
|
||||
Path(work_dir).mkdir(parents=True, exist_ok=True)
|
||||
Path(self.manifest_dir).mkdir(parents=True, exist_ok=True)
|
||||
manifest = Path(self.manifest_dir) / f"manifest_{rank:05d}.csv"
|
||||
|
||||
api = HfApi()
|
||||
# Resume prefetch. A single transient Hub read timeout here must NOT kill the whole
|
||||
# rank (and skip its entire dataset slice), so retry with backoff and, as a last
|
||||
# resort, fall back to an empty set (already-present datasets are re-checked per item
|
||||
# and, for --source molmoact, were already filtered out on the submit node).
|
||||
dst_files: set = set()
|
||||
if not self.only_classify and not self.no_push:
|
||||
for attempt in range(5):
|
||||
try:
|
||||
dst_files = set(api.list_repo_files(self.dst_repo, repo_type="dataset"))
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == 4:
|
||||
logging.warning(f"Rank {rank}: could not list {self.dst_repo} after 5 "
|
||||
f"tries ({e}); proceeding without a resume set.")
|
||||
else:
|
||||
time.sleep(5 * (attempt + 1))
|
||||
|
||||
fieldnames = sorted(
|
||||
{"root", "robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action"}
|
||||
)
|
||||
write_header = not manifest.exists()
|
||||
with open(manifest, "a", newline="") as mf:
|
||||
w = csv.DictWriter(mf, fieldnames=fieldnames)
|
||||
if write_header:
|
||||
w.writeheader()
|
||||
for i, sub in enumerate(my_subs):
|
||||
try:
|
||||
if self.only_classify:
|
||||
if self.standalone:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
snapshot_download(repo_id=sub, repo_type="dataset",
|
||||
local_dir=str(Path(work_dir) / sub),
|
||||
allow_patterns=["meta/*"])
|
||||
else:
|
||||
download_subfolder(sub, work_dir, patterns=[f"{sub}/meta/*"])
|
||||
row = {"root": sub, **classify(Path(work_dir) / sub)}
|
||||
shutil.rmtree(Path(work_dir) / sub.split("/")[0], ignore_errors=True)
|
||||
elif not self.no_push and already_done(api, self.dst_repo, sub, dst_files):
|
||||
row = {"root": sub, "action": "skipped: already present in destination repo"}
|
||||
else:
|
||||
row = migrate_one(api, self.dst_repo, sub, work_dir, self.no_push,
|
||||
standalone=self.standalone)
|
||||
except Exception as e:
|
||||
row = {"root": sub, "action": f"ERROR: {e}"}
|
||||
traceback.print_exc()
|
||||
w.writerow({k: row.get(k) for k in fieldnames})
|
||||
mf.flush()
|
||||
logging.info(f"Rank {rank} [{i + 1}/{len(my_subs)}] {sub}: {row.get('action')}")
|
||||
|
||||
|
||||
def _mem_gb(mem: str) -> int:
|
||||
s = str(mem).strip().lower().rstrip("b").rstrip("g")
|
||||
return int(float(s))
|
||||
|
||||
|
||||
def _make_executor(pipeline, logs_dir, job_name, slurm, workers, time, partition, cpus, mem, qos, env_command, venv_path):
|
||||
kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)}
|
||||
if slurm:
|
||||
kwargs.update(
|
||||
{
|
||||
"job_name": job_name,
|
||||
"tasks": workers,
|
||||
"workers": workers,
|
||||
"time": time,
|
||||
"partition": partition,
|
||||
"cpus_per_task": cpus,
|
||||
"mem_per_cpu_gb": _mem_gb(mem),
|
||||
"sbatch_args": {},
|
||||
}
|
||||
)
|
||||
if qos:
|
||||
kwargs["qos"] = qos
|
||||
if venv_path:
|
||||
kwargs["venv_path"] = venv_path
|
||||
if env_command:
|
||||
kwargs["env_command"] = env_command
|
||||
return SlurmPipelineExecutor(**kwargs)
|
||||
kwargs.update({"tasks": workers, "workers": 1})
|
||||
return LocalPipelineExecutor(**kwargs)
|
||||
|
||||
|
||||
def main():
|
||||
import sys
|
||||
|
||||
if MIGRATION_DIR not in sys.path:
|
||||
sys.path.insert(0, MIGRATION_DIR)
|
||||
from huggingface_hub import HfApi
|
||||
from migrate_molmoact import REFERENCE_REPO, list_molmoact_datasets, pending_datasets
|
||||
from run_migration import SRC_REPO, list_datasets, resolve_folders
|
||||
|
||||
p = argparse.ArgumentParser(
|
||||
description="SLURM-distributed migration to LeRobotDataset v3.0 (map-only). Source is "
|
||||
"either the community_dataset_v3 monorepo or the standalone datasets listed "
|
||||
"by allenai/MolmoAct2-SO100_101-Dataset.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
p.add_argument("--source", choices=("monorepo", "molmoact"), default="monorepo",
|
||||
help="'monorepo': HuggingFaceVLA/community_dataset_v3 subfolders. "
|
||||
"'molmoact': standalone datasets listed by allenai/MolmoAct2-SO100_101-Dataset.")
|
||||
p.add_argument("--reference-repo", default=REFERENCE_REPO, metavar="ORG/NAME",
|
||||
help="(--source molmoact) Repo whose datasets are already migrated and skipped.")
|
||||
p.add_argument("--dst-repo", default=None, metavar="ORG/NAME", help="Destination HF dataset repo.")
|
||||
p.add_argument("--work-dir", default="./cdv3_work", help="Scratch root; each rank gets a subdir.")
|
||||
p.add_argument("--manifest-dir", default="./cdv3_manifests", help="Per-rank manifest CSVs land here.")
|
||||
p.add_argument("--logs-dir", type=Path, default=Path("logs"), help="datatrove logs dir.")
|
||||
p.add_argument("--job-name", default="cdv3_migrate", help="SLURM job name.")
|
||||
p.add_argument("--workers", type=int, default=64, help="Number of parallel SLURM tasks.")
|
||||
p.add_argument("--slurm", type=int, default=1, help="1 = submit via SLURM; 0 = run locally.")
|
||||
p.add_argument("--partition", default=None, help="SLURM partition, e.g. 'hopper-cpu'.")
|
||||
p.add_argument("--qos", default=None, help="SLURM QoS, e.g. 'normal'.")
|
||||
p.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per SLURM task.")
|
||||
p.add_argument("--mem-per-cpu", default="4G", help="Memory per CPU, e.g. '4G'.")
|
||||
p.add_argument("--time", default="24:00:00", help="Wall-clock limit per task.")
|
||||
p.add_argument("--venv-path", default=None, help="venv activate script sourced on each worker.")
|
||||
p.add_argument("--env-command", default=None, help="Raw shell snippet run before python (export HF_TOKEN, etc.).")
|
||||
p.add_argument("--folder-name", nargs="+", default=None, help="Target specific folders/namespaces instead of all.")
|
||||
p.add_argument("--limit", type=int, default=None, help="Only the first N sub-datasets (ignored with --folder-name).")
|
||||
p.add_argument("--only-classify", action="store_true", help="Only classify + write manifest; no convert/upload.")
|
||||
p.add_argument("--no-push", action="store_true", help="Fix + convert locally, keep output, do not upload.")
|
||||
args = p.parse_args()
|
||||
|
||||
if not args.no_push and not args.only_classify and not args.dst_repo:
|
||||
p.error("--dst-repo is required unless --no-push or --only-classify is set.")
|
||||
|
||||
api = HfApi()
|
||||
standalone = args.source == "molmoact"
|
||||
if standalone:
|
||||
all_ids = list_molmoact_datasets(api)
|
||||
if args.folder_name:
|
||||
wanted = {n.strip("/") for n in args.folder_name}
|
||||
subs = [s for s in all_ids if s in wanted]
|
||||
else:
|
||||
subs = pending_datasets(api, all_ids, args.dst_repo, args.reference_repo,
|
||||
args.no_push, args.only_classify)
|
||||
if args.limit:
|
||||
subs = subs[: args.limit]
|
||||
elif args.folder_name:
|
||||
subs = resolve_folders(api, SRC_REPO, args.folder_name)
|
||||
else:
|
||||
subs = list_datasets(api, SRC_REPO)
|
||||
if args.limit:
|
||||
subs = subs[: args.limit]
|
||||
print(f"{len(subs)} sub-datasets targeted", file=sys.stderr)
|
||||
if not subs:
|
||||
p.error("no sub-datasets resolved")
|
||||
|
||||
# Create the destination repo once on the submit node so workers don't race on it.
|
||||
if not args.only_classify and not args.no_push:
|
||||
api.create_repo(args.dst_repo, repo_type="dataset", exist_ok=True)
|
||||
|
||||
executor = _make_executor(
|
||||
pipeline=[
|
||||
MigrateShard(
|
||||
subs,
|
||||
args.dst_repo,
|
||||
args.work_dir,
|
||||
args.manifest_dir,
|
||||
MIGRATION_DIR,
|
||||
no_push=args.no_push,
|
||||
only_classify=args.only_classify,
|
||||
standalone=standalone,
|
||||
)
|
||||
],
|
||||
logs_dir=args.logs_dir,
|
||||
job_name=args.job_name,
|
||||
slurm=args.slurm == 1,
|
||||
workers=args.workers,
|
||||
time=args.time,
|
||||
partition=args.partition,
|
||||
cpus=args.cpus_per_task,
|
||||
mem=args.mem_per_cpu,
|
||||
qos=args.qos,
|
||||
env_command=args.env_command,
|
||||
venv_path=args.venv_path,
|
||||
)
|
||||
executor.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,69 +0,0 @@
|
||||
"""SO-100/101 joint-frame conversion to physical degrees (post-#777 convention).
|
||||
|
||||
Two calibration-free branches + one that needs an assumed canonical range:
|
||||
|
||||
* degrees_old (bare robot_type `so100`/`so101`, |vals|>~180): PR #3879 old->new
|
||||
convention (sign flip shoulder_lift, +90 deg shoulder_lift/elbow_flex).
|
||||
EXACT.
|
||||
* degrees_new (`*_follower` recorded with use_degrees=True, not saturated): already
|
||||
degrees. EXACT.
|
||||
* normalized (`*_follower`, -100..100 joints / 0..100 gripper, saturates at bounds):
|
||||
5 arm joints are mid-range-zero, only the SCALE is missing (per-robot
|
||||
range_min/max not stored) -> use assumed canonical spans below. APPROXIMATE.
|
||||
The gripper (0..100) is kept in its native frame, matching degrees_new.
|
||||
* radians -> untouched.
|
||||
|
||||
Joint order per arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper.
|
||||
Bimanual (12-dim) tiles the 6-joint block twice.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
JOINT_ORDER = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
||||
|
||||
# --- PR #3879 (degrees). old(community frame) <-> new(v3.0 / post-#777) frame. ---
|
||||
SIGNS = np.array([1.0, -1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float64)
|
||||
OFFSETS_DEG = np.array([0.0, 90.0, 90.0, 0.0, 0.0, 0.0], dtype=np.float64)
|
||||
|
||||
# --- Canonical per-joint spans (DEGREES) used ONLY to invert the -100..100 normalization of
|
||||
# the 5 arm joints (RANGE_M100_100) when per-robot calibration is unavailable: normalized
|
||||
# +/-100 -> +/-HALF_RANGE. The gripper (RANGE_0_100) is left in its native 0..100 frame in
|
||||
# every SO dataset, so it needs no canonical span. THESE ARE PLACEHOLDERS — run
|
||||
# calibrate_canonical_ranges.py and paste the fitted values here before a production run. ---
|
||||
CANON_HALF_RANGE_DEG = np.array([100.0, 100.0, 100.0, 100.0, 100.0], dtype=np.float64) # 5 arm joints
|
||||
CANON_IS_CALIBRATED = False # flipped to True once you paste fitted values
|
||||
|
||||
|
||||
def _convert_arm(x: np.ndarray, encoding: str) -> np.ndarray:
|
||||
"""x: (..., 6) for a single SO arm -> degrees (..., 6)."""
|
||||
x = np.asarray(x, dtype=np.float64)
|
||||
if encoding == "radians":
|
||||
return x
|
||||
if encoding == "degrees_old":
|
||||
return SIGNS * (x - OFFSETS_DEG)
|
||||
if encoding == "degrees_new":
|
||||
return x
|
||||
if encoding == "normalized":
|
||||
new_deg = np.array(x, dtype=np.float64)
|
||||
new_deg[..., :5] = (x[..., :5] / 100.0) * CANON_HALF_RANGE_DEG
|
||||
# gripper is RANGE_0_100 in every SO dataset (including use_degrees=True / degrees_new),
|
||||
# so it is already frame-consistent and must be left untouched, not remapped to +/-deg.
|
||||
return new_deg
|
||||
raise ValueError(f"unknown encoding: {encoding!r}")
|
||||
|
||||
|
||||
def to_degrees(arr, encoding: str, n_joints_per_arm: int = 6) -> np.ndarray:
|
||||
"""arr: (..., D) with D a multiple of 6. Returns float32 degrees, same shape."""
|
||||
arr = np.asarray(arr, dtype=np.float64)
|
||||
d = arr.shape[-1]
|
||||
if d % n_joints_per_arm != 0:
|
||||
raise ValueError(f"action/state dim {d} is not a multiple of {n_joints_per_arm}")
|
||||
if encoding == "normalized" and not CANON_IS_CALIBRATED:
|
||||
raise RuntimeError(
|
||||
"CANON ranges are placeholders. Run calibrate_canonical_ranges.py and set "
|
||||
"CANON_* + CANON_IS_CALIBRATED=True before converting 'normalized' datasets to degrees."
|
||||
)
|
||||
out = np.empty_like(arr)
|
||||
for a in range(d // n_joints_per_arm):
|
||||
sl = slice(a * n_joints_per_arm, (a + 1) * n_joints_per_arm)
|
||||
out[..., sl] = _convert_arm(arr[..., sl], encoding)
|
||||
return out.astype(np.float32)
|
||||
@@ -1,117 +0,0 @@
|
||||
"""Tag the migrated (v3.0) datasets whose SO-arm joints are still in normalized units.
|
||||
|
||||
Walks every `{user}/{dataset}` sub-dataset in the destination repo, re-classifies it from its
|
||||
v3.0 metadata (meta/info.json + meta/stats.json), and adds the `normalized` card tag to any whose
|
||||
joint state/action are in normalized units (-100..100 / 0..100) rather than physical degrees. These
|
||||
are the datasets run_migration.py left un-converted (uncalibrated -> APPROXIMATE). Idempotent:
|
||||
already-tagged datasets are skipped. Dry-run by default; pass --yes to actually push the edited card.
|
||||
|
||||
python tag_normalized.py --dst-repo lerobot/community_dataset_v3 # dry-run
|
||||
python tag_normalized.py --dst-repo lerobot/community_dataset_v3 --yes
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from huggingface_hub import DatasetCard, HfApi, hf_hub_download
|
||||
|
||||
from classify import (
|
||||
encoding_from_bounds,
|
||||
is_end_effector,
|
||||
is_mislabeled_so,
|
||||
is_so_robot_type,
|
||||
load_info,
|
||||
so_joint_count,
|
||||
)
|
||||
from run_migration import NORMALIZED_TAG
|
||||
|
||||
DST_REPO = "lerobot/community_dataset_v3"
|
||||
|
||||
|
||||
def is_normalized(root: Path) -> bool:
|
||||
"""True when the SO-arm joints of a v3.0 sub-dataset at ``root`` are still normalized."""
|
||||
info = load_info(root)
|
||||
rt = info.get("robot_type", "") or ""
|
||||
if is_end_effector(info) or not is_so_robot_type(rt) or is_mislabeled_so(info):
|
||||
return False
|
||||
stats = json.loads((root / "meta" / "stats.json").read_text())
|
||||
key = next((k for k in ("action", "observation.state") if k in stats), None)
|
||||
if key is None:
|
||||
return False
|
||||
lo = np.asarray(stats[key]["min"], dtype=float)
|
||||
hi = np.asarray(stats[key]["max"], dtype=float)
|
||||
n = so_joint_count(info, key) or len(hi)
|
||||
return encoding_from_bounds(lo[:n], hi[:n], rt)["encoding"] == "normalized"
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--dst-repo", default=DST_REPO, metavar="ORG/NAME",
|
||||
help="Destination HF dataset repo whose sub-datasets are inspected and tagged.")
|
||||
ap.add_argument("--limit", type=int, default=None, metavar="N",
|
||||
help="Inspect only the first N sub-datasets (alphabetical). Useful for a smoke test.")
|
||||
ap.add_argument("--yes", action="store_true",
|
||||
help="Actually push the tag. Without it, only prints what would be tagged (dry-run).")
|
||||
args = ap.parse_args()
|
||||
|
||||
api = HfApi()
|
||||
subs = sorted(p[: -len("/meta/info.json")]
|
||||
for p in api.list_repo_files(args.dst_repo, repo_type="dataset")
|
||||
if p.endswith("/meta/info.json"))
|
||||
if args.limit:
|
||||
subs = subs[: args.limit]
|
||||
print(f"{len(subs)} sub-datasets in {args.dst_repo}", file=sys.stderr)
|
||||
|
||||
tagged, already, failed = [], [], []
|
||||
for i, sub in enumerate(subs):
|
||||
try:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
for f in ("meta/info.json", "meta/stats.json"):
|
||||
hf_hub_download(args.dst_repo, f"{sub}/{f}", repo_type="dataset", local_dir=tmp)
|
||||
if not is_normalized(Path(tmp) / sub):
|
||||
continue
|
||||
readme = hf_hub_download(args.dst_repo, f"{sub}/README.md", repo_type="dataset", local_dir=tmp)
|
||||
card = DatasetCard.load(readme)
|
||||
card_tags = list(card.data.tags or [])
|
||||
if NORMALIZED_TAG in card_tags:
|
||||
already.append(sub)
|
||||
continue
|
||||
if not args.yes:
|
||||
tagged.append(sub)
|
||||
continue
|
||||
card.data.tags = card_tags + [NORMALIZED_TAG]
|
||||
card.save(readme)
|
||||
for attempt in range(1, 4): # transient Hub ReadTimeouts are common; retry w/ backoff
|
||||
try:
|
||||
api.upload_file(path_or_fileobj=readme, path_in_repo=f"{sub}/README.md",
|
||||
repo_id=args.dst_repo, repo_type="dataset",
|
||||
commit_message=f"Tag {sub} '{NORMALIZED_TAG}'")
|
||||
tagged.append(sub)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == 3:
|
||||
failed.append(sub)
|
||||
print(f"FAILED {sub}: {e}", file=sys.stderr)
|
||||
else:
|
||||
time.sleep(2 ** attempt)
|
||||
except Exception as e:
|
||||
failed.append(sub)
|
||||
print(f"ERROR {sub}: {e}", file=sys.stderr)
|
||||
print(f"[{i + 1}/{len(subs)}] {sub}", file=sys.stderr)
|
||||
|
||||
verb = "tagged" if args.yes else "would tag"
|
||||
for sub in tagged:
|
||||
print(f"{verb}: {sub}")
|
||||
print(f"\n{len(tagged)} {verb} '{NORMALIZED_TAG}', {len(already)} already tagged, "
|
||||
f"{len(failed)} failed.", file=sys.stderr)
|
||||
if not args.yes and tagged:
|
||||
print("dry-run: nothing pushed. re-run with --yes to apply.", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -150,14 +150,14 @@ class MyPolicy(PreTrainedPolicy):
|
||||
|
||||
The methods called by the train/eval loops:
|
||||
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
|
||||
Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
|
||||
|
||||
@@ -295,12 +295,10 @@ The file names are load-bearing: the factory does lazy imports by name, and the
|
||||
|
||||
### Wiring
|
||||
|
||||
Four places need to know about your policy. All by name.
|
||||
Two places need to know about your policy. All by name.
|
||||
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
|
||||
3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
|
||||
4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
|
||||
Mirror an existing policy that's structurally similar to yours; the diff is small.
|
||||
|
||||
@@ -332,6 +330,10 @@ This way:
|
||||
|
||||
Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
|
||||
|
||||
### Avoid copying a modeling file — subclass it
|
||||
|
||||
If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
|
||||
|
||||
### Benchmarks and a published checkpoint
|
||||
|
||||
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
|
||||
@@ -367,7 +369,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
|
||||
The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
|
||||
|
||||
- [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
|
||||
- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
|
||||
- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
|
||||
- [ ] `make_my_policy_pre_post_processors` follows the naming convention.
|
||||
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
|
||||
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
|
||||
|
||||
@@ -46,8 +46,11 @@ CMD = (
|
||||
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
|
||||
"pip install --no-deps "
|
||||
"'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
|
||||
# Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
|
||||
# draccus 0.11, which break lerobot at import time.
|
||||
"pip install --upgrade-strategy only-if-needed "
|
||||
"datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
|
||||
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"openai && "
|
||||
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
|
||||
"export VLLM_VIDEO_BACKEND=pyav && "
|
||||
|
||||
@@ -413,8 +413,6 @@ ignore = [
|
||||
"__init__.py" = ["F401", "F403", "E402"]
|
||||
# E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
|
||||
"src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
|
||||
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
|
||||
|
||||
[tool.ruff.lint.isort]
|
||||
combine-as-imports = true
|
||||
known-first-party = ["lerobot"]
|
||||
|
||||
@@ -205,24 +205,30 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
|
||||
f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
|
||||
# HACK: Parse the original config to get the config subclass, so that we can
|
||||
# apply cli overrides.
|
||||
# This is very ugly, ideally we'd like to be able to do that natively with draccus
|
||||
# something like --policy.path (in addition to --policy.type)
|
||||
with draccus.config_type("json"):
|
||||
orig_config = draccus.parse(cls, config_file, args=[])
|
||||
|
||||
if config_file is None:
|
||||
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
|
||||
|
||||
with open(config_file) as f:
|
||||
config = json.load(f)
|
||||
|
||||
config.pop("type")
|
||||
# Resolve the concrete config subclass from the serialized "type" tag, then parse
|
||||
# the config (with CLI overrides) directly for that class. The "type" key is
|
||||
# stripped because draccus only consumes it when parsing the registry base class.
|
||||
policy_type = config.pop("type", None)
|
||||
if policy_type is None:
|
||||
raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
|
||||
try:
|
||||
config_cls = cls.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
|
||||
f"Available policy types: {cls.get_known_choices()}"
|
||||
) from e
|
||||
|
||||
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
|
||||
json.dump(config, f)
|
||||
config_file = f.name
|
||||
|
||||
cli_overrides = policy_kwargs.pop("cli_overrides", [])
|
||||
with draccus.config_type("json"):
|
||||
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
|
||||
return draccus.parse(config_cls, config_file, args=cli_overrides)
|
||||
|
||||
@@ -729,7 +729,6 @@ def concatenate_video_files(
|
||||
stream_map[input_stream.index].time_base = input_stream.time_base
|
||||
|
||||
# Demux + remux packets (no re-encode)
|
||||
last_dts: dict[int, int] = {}
|
||||
for packet in input_container.demux():
|
||||
# Skip packets from un-mapped streams
|
||||
if packet.stream.index not in stream_map:
|
||||
@@ -739,19 +738,6 @@ def concatenate_video_files(
|
||||
if packet.dts is None:
|
||||
continue
|
||||
|
||||
# Enforce strictly increasing DTS. Clips encoded with B-frames start at a negative DTS, so the
|
||||
# concat demuxer can emit the first packet of a clip with a DTS equal to the last of the previous
|
||||
# one; the MP4 muxer rejects duplicate/decreasing DTS with [Errno 22]. Nudge such packets forward.
|
||||
prev_dts = last_dts.get(packet.stream.index)
|
||||
if prev_dts is not None and packet.dts <= prev_dts:
|
||||
shift = prev_dts + 1 - packet.dts
|
||||
packet.dts += shift
|
||||
if packet.pts is not None:
|
||||
packet.pts += shift
|
||||
if packet.pts is not None and packet.pts < packet.dts:
|
||||
packet.pts = packet.dts
|
||||
last_dts[packet.stream.index] = packet.dts
|
||||
|
||||
output_stream = stream_map[packet.stream.index]
|
||||
packet.stream = output_stream
|
||||
output_container.mux(packet)
|
||||
|
||||
@@ -32,6 +32,7 @@ from .pretrained import PreTrainedPolicy as PreTrainedPolicy
|
||||
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
|
||||
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
|
||||
from .utils import make_robot_action, prepare_observation_for_inference
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
|
||||
@@ -57,6 +58,7 @@ __all__ = [
|
||||
"PI05Config",
|
||||
"SmolVLAConfig",
|
||||
"TDMPCConfig",
|
||||
"VLAJEPAConfig",
|
||||
"VQBeTConfig",
|
||||
"WallXConfig",
|
||||
"XVLAConfig",
|
||||
|
||||
@@ -18,17 +18,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_act import ACTConfig
|
||||
|
||||
@@ -54,34 +47,4 @@ def make_act_pre_post_processors(
|
||||
tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
|
||||
pre-processor pipeline and the post-processor pipeline.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Flow-matching sampling primitives shared across policies.
|
||||
|
||||
Canonical versions of the beta-distributed timestep sampler and the forward-Euler
|
||||
denoising loop (with its real-time-chunking hook) that the openpi-derived policies
|
||||
(pi0, pi05, smolvla, eo1) historically each carried a copy of. All functions are
|
||||
stateless; adopting them does not affect checkpoints.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
|
||||
|
||||
|
||||
def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor: # see openpi (exact copy)
|
||||
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
|
||||
alpha_t = torch.tensor(alpha, dtype=torch.float32)
|
||||
beta_t = torch.tensor(beta, dtype=torch.float32)
|
||||
dist = torch.distributions.Beta(alpha_t, beta_t)
|
||||
return dist.sample((bsize,)).to(device)
|
||||
|
||||
|
||||
def sample_noise(shape, device) -> Tensor:
|
||||
"""Standard-normal float32 noise, the flow-matching x_1 sample."""
|
||||
return torch.normal(
|
||||
mean=0.0,
|
||||
std=1.0,
|
||||
size=shape,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
def sample_time_beta(bsize: int, device, *, alpha: float, beta: float, scale: float, offset: float) -> Tensor:
|
||||
"""Beta-distributed flow-matching timesteps: ``Beta(alpha, beta) * scale + offset`` (openpi convention)."""
|
||||
time_beta = sample_beta(alpha, beta, bsize, device)
|
||||
time = time_beta * scale + offset
|
||||
return time.to(dtype=torch.float32, device=device)
|
||||
|
||||
|
||||
def euler_integrate(
|
||||
denoise_fn: Callable[[Tensor, Tensor], Tensor],
|
||||
noise: Tensor,
|
||||
num_steps: int,
|
||||
*,
|
||||
forward_euler: bool = False,
|
||||
rtc_processor: "RTCProcessor | None" = None,
|
||||
rtc_enabled: bool = False,
|
||||
inference_delay: int | None = None,
|
||||
prev_chunk_left_over: Tensor | None = None,
|
||||
execution_horizon: int | None = None,
|
||||
) -> Tensor:
|
||||
"""Euler integration of a velocity field between the noise and action endpoints.
|
||||
|
||||
Two integration conventions are supported via ``forward_euler``:
|
||||
|
||||
* Backward (default, openpi: pi0, pi05, eo1, smolvla): integrates from t=1 (noise) to
|
||||
t=0 (actions) with ``dt = -1/num_steps`` and ``time = 1.0 + step*dt``.
|
||||
* Forward (groot, evo1, wall_x): integrates from t=0 (noise) to t=1 (actions) with
|
||||
``dt = +1/num_steps`` and ``time = step*dt``.
|
||||
|
||||
In both cases the update is ``x_t <- x_t + dt * v_t``, with the optional
|
||||
real-time-chunking (RTC) guidance hook wrapping the velocity computation and debug
|
||||
tracking after each step.
|
||||
|
||||
Args:
|
||||
denoise_fn: Computes the velocity ``v_t`` from ``(x_t, time_tensor)`` where
|
||||
``time_tensor`` is a float32 tensor of shape ``(batch_size,)``. The returned
|
||||
velocity must have the same shape and dtype as ``x_t``.
|
||||
noise: Initial sample of shape ``(batch_size, ...)``. This is ``x_1`` for the
|
||||
backward convention and ``x_0`` for the forward convention.
|
||||
num_steps: Number of Euler steps.
|
||||
forward_euler: If ``True`` use the forward convention (start at t=0); otherwise
|
||||
use the backward openpi convention (start at t=1).
|
||||
rtc_processor: Optional RTC processor. Debug tracking fires whenever it is set and
|
||||
has debugging enabled, even if RTC guidance itself is disabled (this mirrors
|
||||
the historical per-policy loops).
|
||||
rtc_enabled: Whether to route the velocity computation through
|
||||
``rtc_processor.denoise_step`` (requires ``rtc_processor``).
|
||||
inference_delay: RTC guidance parameter, forwarded verbatim.
|
||||
prev_chunk_left_over: RTC guidance parameter, forwarded verbatim.
|
||||
execution_horizon: RTC guidance parameter, forwarded verbatim.
|
||||
"""
|
||||
bsize = noise.shape[0]
|
||||
device = noise.device
|
||||
|
||||
dt = 1.0 / num_steps if forward_euler else -1.0 / num_steps
|
||||
t_start = 0.0 if forward_euler else 1.0
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = t_start + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return denoise_fn(input_x_t, current_timestep)
|
||||
|
||||
if rtc_enabled:
|
||||
v_t = rtc_processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk_left_over,
|
||||
inference_delay=inference_delay,
|
||||
time=time,
|
||||
original_denoise_step_partial=denoise_step_partial_call,
|
||||
execution_horizon=execution_horizon,
|
||||
)
|
||||
else:
|
||||
v_t = denoise_step_partial_call(x_t)
|
||||
|
||||
x_t = x_t + dt * v_t
|
||||
|
||||
if rtc_processor is not None and rtc_processor.is_debug_enabled():
|
||||
rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
|
||||
|
||||
return x_t
|
||||
@@ -0,0 +1,243 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Helpers shared by the openpi-derived VLA policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla).
|
||||
|
||||
These are the canonical versions of functions that historically were copy-pasted per
|
||||
policy. They are pure (no parameters, no module state), so importing them from here
|
||||
instead of a policy-local copy has no effect on checkpoints.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: N812
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
|
||||
from lerobot.utils.device_utils import get_safe_dtype
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers import DynamicCache
|
||||
else:
|
||||
DynamicCache = None
|
||||
|
||||
|
||||
def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy)
|
||||
time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
|
||||
) -> Tensor:
|
||||
"""Computes sine-cosine positional embedding vectors for scalar positions."""
|
||||
if dimension % 2 != 0:
|
||||
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
|
||||
|
||||
if time.ndim != 1:
|
||||
raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
|
||||
|
||||
dtype = get_safe_dtype(torch.float64, device.type)
|
||||
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
|
||||
period = min_period * (max_period / min_period) ** fraction
|
||||
|
||||
# Compute the outer product
|
||||
scaling_factor = 1.0 / period * 2 * math.pi
|
||||
sin_input = scaling_factor[None, :] * time[:, None]
|
||||
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
|
||||
|
||||
|
||||
def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor: # see openpi (exact copy)
|
||||
"""Copied from big_vision.
|
||||
|
||||
Tokens can attend to valid inputs tokens which have a cumulative mask_ar
|
||||
smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
|
||||
setup several types of attention, for example:
|
||||
|
||||
[[1 1 1 1 1 1]]: pure causal attention.
|
||||
|
||||
[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
|
||||
themselves and the last 3 tokens have a causal attention. The first
|
||||
entry could also be a 1 without changing behaviour.
|
||||
|
||||
[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
|
||||
block can attend all previous blocks and all tokens on the same block.
|
||||
|
||||
Args:
|
||||
input_mask: bool[B, N] true if its part of the input, false if padding.
|
||||
mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
|
||||
it and 0 where it shares the same attention mask as the previous token.
|
||||
"""
|
||||
if att_masks.ndim != 2:
|
||||
raise ValueError(att_masks.ndim)
|
||||
if pad_masks.ndim != 2:
|
||||
raise ValueError(pad_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 prepare_attention_masks_4d(att_2d_masks: Tensor, dtype: torch.dtype | None = None) -> Tensor:
|
||||
"""Expand boolean 2D attention masks to the additive 4D layout expected by transformers.
|
||||
|
||||
Valid positions become 0.0 and masked positions the large negative openpi constant.
|
||||
"""
|
||||
att_2d_masks_4d = att_2d_masks[:, None, :, :]
|
||||
result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
|
||||
def clone_past_key_values(past_key_values):
|
||||
"""Clone the DynamicCache returned by prefix prefill for compiled denoising."""
|
||||
if DynamicCache is None:
|
||||
require_package("transformers", extra="transformers-dep")
|
||||
|
||||
return DynamicCache(
|
||||
tuple(
|
||||
(keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def pad_vector(vector: Tensor, new_dim: int, *, truncate: bool = False) -> Tensor:
|
||||
"""Pad the last dimension of a vector to new_dim with zeros.
|
||||
|
||||
Can be (batch_size x sequence_length x features_dimension)
|
||||
or (batch_size x features_dimension)
|
||||
|
||||
With ``truncate=False`` (openpi behavior), vectors whose last dimension is already
|
||||
>= new_dim are returned unchanged. With ``truncate=True`` (xVLA behavior), the last
|
||||
dimension is truncated to exactly ``new_dim`` (which may be 0).
|
||||
"""
|
||||
if vector.shape[-1] == new_dim:
|
||||
return vector
|
||||
if not truncate:
|
||||
if vector.shape[-1] >= new_dim:
|
||||
return vector
|
||||
return F.pad(vector, (0, new_dim - vector.shape[-1]))
|
||||
shape = list(vector.shape)
|
||||
current_dim = shape[-1]
|
||||
shape[-1] = new_dim
|
||||
new_vector = vector.new_zeros(*shape)
|
||||
length = min(current_dim, new_dim)
|
||||
new_vector[..., :length] = vector[..., :length]
|
||||
return new_vector
|
||||
|
||||
|
||||
def resize_with_pad_torch( # see openpi `resize_with_pad_torch` (exact copy)
|
||||
images: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
mode: str = "bilinear",
|
||||
) -> torch.Tensor:
|
||||
"""PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
|
||||
by padding with black. If the image is float32, it must be in the range [-1, 1].
|
||||
|
||||
Padding is centered (openpi convention). For the top-left-padding variant used by
|
||||
smolvla/xvla, see :func:`resize_with_pad`.
|
||||
|
||||
Args:
|
||||
images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
|
||||
height: Target height
|
||||
width: Target width
|
||||
mode: Interpolation mode ('bilinear', 'nearest', etc.)
|
||||
|
||||
Returns:
|
||||
Resized and padded tensor with same shape format as input
|
||||
"""
|
||||
# Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
|
||||
if images.shape[-1] <= 4: # Assume channels-last format
|
||||
channels_last = True
|
||||
if images.dim() == 3:
|
||||
images = images.unsqueeze(0) # Add batch dimension
|
||||
images = images.permute(0, 3, 1, 2) # [b, h, w, c] -> [b, c, h, w]
|
||||
else:
|
||||
channels_last = False
|
||||
if images.dim() == 3:
|
||||
images = images.unsqueeze(0) # Add batch dimension
|
||||
|
||||
batch_size, channels, cur_height, cur_width = images.shape
|
||||
|
||||
# Calculate resize ratio
|
||||
ratio = max(cur_width / width, cur_height / height)
|
||||
resized_height = int(cur_height / ratio)
|
||||
resized_width = int(cur_width / ratio)
|
||||
|
||||
# Resize
|
||||
resized_images = F.interpolate(
|
||||
images,
|
||||
size=(resized_height, resized_width),
|
||||
mode=mode,
|
||||
align_corners=False if mode == "bilinear" else None,
|
||||
)
|
||||
|
||||
# Handle dtype-specific clipping
|
||||
if images.dtype == torch.uint8:
|
||||
resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
|
||||
elif images.dtype == torch.float32:
|
||||
resized_images = resized_images.clamp(0.0, 1.0)
|
||||
else:
|
||||
raise ValueError(f"Unsupported image dtype: {images.dtype}")
|
||||
|
||||
# Calculate padding
|
||||
pad_h0, remainder_h = divmod(height - resized_height, 2)
|
||||
pad_h1 = pad_h0 + remainder_h
|
||||
pad_w0, remainder_w = divmod(width - resized_width, 2)
|
||||
pad_w1 = pad_w0 + remainder_w
|
||||
|
||||
# Pad
|
||||
constant_value = 0 if images.dtype == torch.uint8 else 0.0
|
||||
padded_images = F.pad(
|
||||
resized_images,
|
||||
(pad_w0, pad_w1, pad_h0, pad_h1), # left, right, top, bottom
|
||||
mode="constant",
|
||||
value=constant_value,
|
||||
)
|
||||
|
||||
# Convert back to original format if needed
|
||||
if channels_last:
|
||||
padded_images = padded_images.permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c]
|
||||
|
||||
return padded_images
|
||||
|
||||
|
||||
def resize_with_pad(img: torch.Tensor, height: int, width: int, *, pad_value: float) -> torch.Tensor:
|
||||
"""Resize a (b, c, h, w) image without distortion, padding on the LEFT and TOP.
|
||||
|
||||
This is the smolvla/xvla convention. For the centered-padding openpi variant, see
|
||||
:func:`resize_with_pad_torch`. ``pad_value`` is keyword-only on purpose: callers
|
||||
historically used different values (0, -1) and must state their choice explicitly.
|
||||
"""
|
||||
if img.ndim != 4:
|
||||
raise ValueError(f"(b,c,h,w) expected, but got {img.shape}")
|
||||
|
||||
current_height, current_width = img.shape[2:]
|
||||
if current_height == height and current_width == width:
|
||||
return img
|
||||
|
||||
ratio = max(current_width / width, current_height / height)
|
||||
resized_height = int(current_height / ratio)
|
||||
resized_width = int(current_width / ratio)
|
||||
resized_img = F.interpolate(
|
||||
img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
|
||||
)
|
||||
|
||||
pad_height = max(0, height - resized_height)
|
||||
pad_width = max(0, width - resized_width)
|
||||
padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
|
||||
return padded_img
|
||||
@@ -19,17 +19,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_diffusion import DiffusionConfig
|
||||
|
||||
@@ -63,32 +56,4 @@ def make_diffusion_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -23,24 +23,16 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.types import TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
from .configuration_eo1 import EO1Config
|
||||
@@ -242,14 +234,12 @@ def make_eo1_pre_post_processors(
|
||||
]:
|
||||
"""Build pre/post processor pipelines for EO1."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
|
||||
EO1QwenProcessorStep(
|
||||
processor_name=config.vlm_base,
|
||||
@@ -257,27 +247,12 @@ def make_eo1_pre_post_processors(
|
||||
image_max_pixels=config.image_max_pixels,
|
||||
use_fast_processor=config.use_fast_processor,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -27,9 +27,11 @@ from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
else:
|
||||
AutoModel = None
|
||||
AutoTokenizer = None
|
||||
is_flash_attn_2_available = None
|
||||
|
||||
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD = (0.229, 0.224, 0.225)
|
||||
@@ -135,9 +137,13 @@ class InternVL3Embedder(nn.Module):
|
||||
raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
|
||||
self.model_dtype = model_dtype
|
||||
|
||||
attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
|
||||
attn_implementation = (
|
||||
"flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager"
|
||||
)
|
||||
if use_flash_attn and attn_implementation == "eager":
|
||||
logger.warning("flash_attn is not installed. Falling back to eager attention.")
|
||||
logger.warning(
|
||||
"Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
|
||||
)
|
||||
|
||||
self.model = AutoModel.from_pretrained(
|
||||
model_name,
|
||||
@@ -359,11 +365,3 @@ class InternVL3Embedder(nn.Module):
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.model.parameters()).device
|
||||
|
||||
|
||||
def _flash_attn_available() -> bool:
|
||||
try:
|
||||
import flash_attn # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
return False
|
||||
return True
|
||||
|
||||
+66
-318
@@ -17,6 +17,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, TypedDict, Unpack
|
||||
|
||||
@@ -44,26 +45,10 @@ from lerobot.utils.constants import (
|
||||
)
|
||||
from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||
|
||||
from .act.configuration_act import ACTConfig
|
||||
from .diffusion.configuration_diffusion import DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config
|
||||
from .evo1.configuration_evo1 import Evo1Config
|
||||
from .fastwam.configuration_fastwam import FastWAMConfig
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig
|
||||
from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
from .pi0.configuration_pi0 import PI0Config
|
||||
from .pi05.configuration_pi05 import PI05Config
|
||||
from .pretrained import PreTrainedPolicy
|
||||
from .smolvla.configuration_smolvla import SmolVLAConfig
|
||||
from .tdmpc.configuration_tdmpc import TDMPCConfig
|
||||
from .utils import validate_visual_features_consistency
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig
|
||||
|
||||
|
||||
def _reconnect_relative_absolute_steps(
|
||||
@@ -88,100 +73,23 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
"""
|
||||
Retrieves a policy class by its registered name.
|
||||
|
||||
This function uses dynamic imports to avoid loading all policy classes into memory
|
||||
at once, improving startup time and reducing dependencies.
|
||||
Resolution is convention-based: the draccus-registered config class of ``name`` is
|
||||
looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
|
||||
the ``<X>Policy`` class is imported from there. The modeling module is only imported
|
||||
at call time, keeping heavy optional dependencies lazy. This works for both built-in
|
||||
policies and third-party lerobot plugins (anything registered via
|
||||
``@PreTrainedConfig.register_subclass``).
|
||||
|
||||
Args:
|
||||
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
||||
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
|
||||
"molmoact2", "eo1", "evo1".
|
||||
name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
|
||||
Returns:
|
||||
The policy class corresponding to the given name.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the policy name is not recognized.
|
||||
ValueError: If the policy name is not registered.
|
||||
ImportError: If the policy's optional dependencies are not installed.
|
||||
"""
|
||||
if name == "tdmpc":
|
||||
from .tdmpc.modeling_tdmpc import TDMPCPolicy
|
||||
|
||||
return TDMPCPolicy
|
||||
elif name == "diffusion":
|
||||
from .diffusion.modeling_diffusion import DiffusionPolicy
|
||||
|
||||
return DiffusionPolicy
|
||||
elif name == "act":
|
||||
from .act.modeling_act import ACTPolicy
|
||||
|
||||
return ACTPolicy
|
||||
elif name == "multi_task_dit":
|
||||
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
|
||||
|
||||
return MultiTaskDiTPolicy
|
||||
elif name == "vqbet":
|
||||
from .vqbet.modeling_vqbet import VQBeTPolicy
|
||||
|
||||
return VQBeTPolicy
|
||||
elif name == "pi0":
|
||||
from .pi0.modeling_pi0 import PI0Policy
|
||||
|
||||
return PI0Policy
|
||||
elif name == "pi0_fast":
|
||||
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
|
||||
|
||||
return PI0FastPolicy
|
||||
elif name == "pi05":
|
||||
from .pi05.modeling_pi05 import PI05Policy
|
||||
|
||||
return PI05Policy
|
||||
elif name == "gaussian_actor":
|
||||
from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
|
||||
|
||||
return GaussianActorPolicy
|
||||
elif name == "smolvla":
|
||||
from .smolvla.modeling_smolvla import SmolVLAPolicy
|
||||
|
||||
return SmolVLAPolicy
|
||||
elif name == "groot":
|
||||
from .groot.modeling_groot import GrootPolicy
|
||||
|
||||
return GrootPolicy
|
||||
elif name == "xvla":
|
||||
from .xvla.modeling_xvla import XVLAPolicy
|
||||
|
||||
return XVLAPolicy
|
||||
elif name == "wall_x":
|
||||
from .wall_x.modeling_wall_x import WallXPolicy
|
||||
|
||||
return WallXPolicy
|
||||
elif name == "eo1":
|
||||
from .eo1.modeling_eo1 import EO1Policy
|
||||
|
||||
return EO1Policy
|
||||
elif name == "molmoact2":
|
||||
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
|
||||
|
||||
return MolmoAct2Policy
|
||||
elif name == "vla_jepa":
|
||||
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
|
||||
|
||||
return VLAJEPAPolicy
|
||||
elif name == "lingbot_va":
|
||||
from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
|
||||
|
||||
return LingBotVAPolicy
|
||||
elif name == "fastwam":
|
||||
from .fastwam.modeling_fastwam import FastWAMPolicy
|
||||
|
||||
return FastWAMPolicy
|
||||
elif name == "evo1":
|
||||
from .evo1.modeling_evo1 import Evo1Policy
|
||||
|
||||
return Evo1Policy
|
||||
else:
|
||||
try:
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{name}' is not available.") from e
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
|
||||
|
||||
def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
@@ -192,9 +100,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
mapping a string identifier to the corresponding config class.
|
||||
|
||||
Args:
|
||||
policy_type: The type of the policy. Supported types include "tdmpc",
|
||||
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
|
||||
"smolvla", "wall_x", "molmoact2", "eo1", "evo1".
|
||||
policy_type: The registered type of the policy (any name registered via
|
||||
``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
|
||||
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
||||
|
||||
Returns:
|
||||
@@ -203,48 +110,11 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
Raises:
|
||||
ValueError: If the `policy_type` is not recognized.
|
||||
"""
|
||||
if policy_type == "tdmpc":
|
||||
return TDMPCConfig(**kwargs)
|
||||
elif policy_type == "diffusion":
|
||||
return DiffusionConfig(**kwargs)
|
||||
elif policy_type == "act":
|
||||
return ACTConfig(**kwargs)
|
||||
elif policy_type == "multi_task_dit":
|
||||
return MultiTaskDiTConfig(**kwargs)
|
||||
elif policy_type == "vqbet":
|
||||
return VQBeTConfig(**kwargs)
|
||||
elif policy_type == "pi0":
|
||||
return PI0Config(**kwargs)
|
||||
elif policy_type == "pi05":
|
||||
return PI05Config(**kwargs)
|
||||
elif policy_type == "gaussian_actor":
|
||||
return GaussianActorConfig(**kwargs)
|
||||
elif policy_type == "smolvla":
|
||||
return SmolVLAConfig(**kwargs)
|
||||
elif policy_type == "groot":
|
||||
return GrootConfig(**kwargs)
|
||||
elif policy_type == "xvla":
|
||||
return XVLAConfig(**kwargs)
|
||||
elif policy_type == "wall_x":
|
||||
return WallXConfig(**kwargs)
|
||||
elif policy_type == "eo1":
|
||||
return EO1Config(**kwargs)
|
||||
elif policy_type == "molmoact2":
|
||||
return MolmoAct2Config(**kwargs)
|
||||
elif policy_type == "vla_jepa":
|
||||
return VLAJEPAConfig(**kwargs)
|
||||
elif policy_type == "lingbot_va":
|
||||
return LingBotVAConfig(**kwargs)
|
||||
elif policy_type == "fastwam":
|
||||
return FastWAMConfig(**kwargs)
|
||||
elif policy_type == "evo1":
|
||||
return Evo1Config(**kwargs)
|
||||
else:
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
return config_cls(**kwargs)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
return config_cls(**kwargs)
|
||||
|
||||
|
||||
class ProcessorConfigKwargs(TypedDict, total=False):
|
||||
@@ -298,8 +168,7 @@ def make_pre_post_processors(
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If a processor factory is not implemented for the given
|
||||
policy configuration type.
|
||||
ValueError: If no processor factory exists for the given policy configuration type.
|
||||
"""
|
||||
if pretrained_path:
|
||||
if isinstance(policy_cfg, GrootConfig):
|
||||
@@ -351,166 +220,13 @@ def make_pre_post_processors(
|
||||
)
|
||||
return preprocessor, postprocessor
|
||||
|
||||
# Create a new processor based on policy type
|
||||
if isinstance(policy_cfg, TDMPCConfig):
|
||||
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
|
||||
|
||||
processors = make_tdmpc_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, DiffusionConfig):
|
||||
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
|
||||
|
||||
processors = make_diffusion_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, ACTConfig):
|
||||
from .act.processor_act import make_act_pre_post_processors
|
||||
|
||||
processors = make_act_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MultiTaskDiTConfig):
|
||||
from .multi_task_dit.processor_multi_task_dit import (
|
||||
make_multi_task_dit_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_multi_task_dit_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VQBeTConfig):
|
||||
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
|
||||
|
||||
processors = make_vqbet_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI0Config):
|
||||
from .pi0.processor_pi0 import make_pi0_pre_post_processors
|
||||
|
||||
processors = make_pi0_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI05Config):
|
||||
from .pi05.processor_pi05 import make_pi05_pre_post_processors
|
||||
|
||||
processors = make_pi05_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GaussianActorConfig):
|
||||
from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
|
||||
|
||||
processors = make_gaussian_actor_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, SmolVLAConfig):
|
||||
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
|
||||
|
||||
processors = make_smolvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GrootConfig):
|
||||
from .groot.processor_groot import make_groot_pre_post_processors
|
||||
|
||||
processors = make_groot_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, XVLAConfig):
|
||||
from .xvla.processor_xvla import (
|
||||
make_xvla_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_xvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, WallXConfig):
|
||||
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
|
||||
|
||||
processors = make_wall_x_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, EO1Config):
|
||||
from .eo1.processor_eo1 import make_eo1_pre_post_processors
|
||||
|
||||
processors = make_eo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
elif isinstance(policy_cfg, Evo1Config):
|
||||
from .evo1.processor_evo1 import make_evo1_pre_post_processors
|
||||
|
||||
processors = make_evo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MolmoAct2Config):
|
||||
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
|
||||
|
||||
processors = make_molmoact2_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VLAJEPAConfig):
|
||||
from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
|
||||
|
||||
processors = make_vla_jepa_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, LingBotVAConfig):
|
||||
from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
|
||||
|
||||
processors = make_lingbot_va_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, FastWAMConfig):
|
||||
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
|
||||
|
||||
processors = make_fastwam_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
else:
|
||||
try:
|
||||
processors = _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
|
||||
|
||||
return processors
|
||||
# Create new processors from the policy config, resolving the per-policy factory
|
||||
# function by naming convention (lazy import keeps optional dependencies optional).
|
||||
return _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
|
||||
def make_policy(
|
||||
@@ -654,10 +370,12 @@ def make_policy(
|
||||
return policy
|
||||
|
||||
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
|
||||
"""Get policy class from its registered name using dynamic imports.
|
||||
|
||||
This is used as a helper function to import policies from 3rd party lerobot plugins.
|
||||
Works for built-in policies and 3rd party lerobot plugins alike: the config class
|
||||
registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
|
||||
class is imported from the sibling ``modeling_*`` module by naming convention.
|
||||
|
||||
Args:
|
||||
name: The name of the policy.
|
||||
@@ -683,22 +401,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
|
||||
"configuration_", "modeling_"
|
||||
) # e.g., configuration_diffusion -> modeling_diffusion
|
||||
|
||||
module = importlib.import_module(module_path)
|
||||
policy_cls = getattr(module, cls_name)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The modeling_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Policy class for '{name}' is not implemented.") from e
|
||||
raise
|
||||
policy_cls = getattr(module, cls_name, None)
|
||||
if policy_cls is None:
|
||||
raise ValueError(
|
||||
f"Policy class '{cls_name}' not found in '{module_path}'. "
|
||||
f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
|
||||
)
|
||||
return policy_cls
|
||||
|
||||
|
||||
def _make_processors_from_policy_config(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
dataset_meta: Any | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Create pre- and post-processors from a policy configuration using dynamic imports.
|
||||
|
||||
This is used as a helper function to import processor factories from 3rd party lerobot plugins.
|
||||
Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
|
||||
by naming convention. Works for built-in policies and 3rd party lerobot plugins.
|
||||
|
||||
Args:
|
||||
config: The policy configuration object.
|
||||
dataset_stats: Dataset statistics for normalization.
|
||||
dataset_meta: Dataset metadata, forwarded only to factories that declare a
|
||||
``dataset_meta`` parameter (e.g. groot, molmoact2).
|
||||
Returns:
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
"""
|
||||
@@ -711,6 +446,19 @@ def _make_processors_from_policy_config(
|
||||
logging.debug(
|
||||
f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
|
||||
)
|
||||
module = importlib.import_module(module_path)
|
||||
function = getattr(module, function_name)
|
||||
return function(config, dataset_stats=dataset_stats)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The processor_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
|
||||
raise
|
||||
function = getattr(module, function_name, None)
|
||||
if function is None:
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
|
||||
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
|
||||
if "dataset_meta" in inspect.signature(function).parameters:
|
||||
call_kwargs["dataset_meta"] = dataset_meta
|
||||
return function(config, **call_kwargs)
|
||||
|
||||
@@ -22,20 +22,11 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
ActionProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
@@ -105,38 +96,20 @@ def make_fastwam_pre_post_processors(
|
||||
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
|
||||
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
|
||||
|
||||
steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
device=config.device,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
),
|
||||
steps.unnormalize,
|
||||
]
|
||||
if config.toggle_action_dimensions:
|
||||
output_steps.append(
|
||||
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
|
||||
)
|
||||
output_steps.append(DeviceProcessorStep(device="cpu"))
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -20,17 +20,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_gaussian_actor import GaussianActorConfig
|
||||
|
||||
@@ -62,33 +55,4 @@ def make_gaussian_actor_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
# Add remaining processors
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -25,19 +25,12 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
|
||||
from .configuration_lingbot_va import LingBotVAConfig
|
||||
@@ -52,15 +45,13 @@ def make_lingbot_va_pre_post_processors(
|
||||
]:
|
||||
"""Build the pre/post processor pipelines for LingBot-VA."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
# Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
|
||||
@@ -70,18 +61,7 @@ def make_lingbot_va_pre_post_processors(
|
||||
norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -19,18 +19,12 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
|
||||
@@ -66,9 +60,11 @@ def make_multi_task_dit_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_encoder_name,
|
||||
padding=config.tokenizer_padding,
|
||||
@@ -76,32 +72,12 @@ def make_multi_task_dit_pre_post_processors(
|
||||
max_length=config.tokenizer_max_length,
|
||||
truncation=config.tokenizer_truncation,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -21,22 +21,16 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_pi0 import PI0Config
|
||||
|
||||
@@ -136,10 +130,12 @@ def make_pi0_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -147,32 +143,15 @@ def make_pi0_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -24,26 +24,17 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi05 import PI05Config
|
||||
|
||||
@@ -135,18 +126,16 @@ def make_pi05_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
relative_step,
|
||||
# NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
|
||||
# because the tokenizer step expects normalized state in [-1, 1] range for discretization
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -154,26 +143,13 @@ def make_pi05_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -25,26 +25,17 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
ActionTokenizerProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi0_fast import PI0FastConfig
|
||||
|
||||
@@ -135,6 +126,8 @@ def make_pi0_fast_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
|
||||
# This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
|
||||
# caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
|
||||
@@ -144,14 +137,10 @@ def make_pi0_fast_pre_post_processors(
|
||||
# before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
|
||||
# continues to receive normalized state in [-1, 1] as expected.
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_tokenizer_name,
|
||||
@@ -165,26 +154,13 @@ def make_pi0_fast_pre_post_processors(
|
||||
fast_skip_tokens=config.fast_skip_tokens,
|
||||
paligemma_tokenizer_name=config.text_tokenizer_name,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -23,8 +23,6 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -34,6 +32,7 @@ from torch import Tensor, nn
|
||||
from lerobot.__version__ import __version__
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
from .utils import log_model_loading_keys
|
||||
@@ -221,26 +220,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
log_model_loading_keys(missing_keys, unexpected_keys)
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
@abc.abstractmethod
|
||||
|
||||
@@ -19,19 +19,13 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NewLineTaskProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_smolvla import SmolVLAConfig
|
||||
|
||||
@@ -66,9 +60,11 @@ def make_smolvla_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
NewLineTaskProcessorStep(),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.vlm_model_name,
|
||||
@@ -76,28 +72,11 @@ def make_smolvla_pre_post_processors(
|
||||
padding_side="right",
|
||||
max_length=config.tokenizer_max_length,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -19,17 +19,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_tdmpc import TDMPCConfig
|
||||
|
||||
@@ -61,32 +54,4 @@ def make_tdmpc_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -20,20 +20,16 @@ import torch
|
||||
|
||||
from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
EnvTransition,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TransitionKey,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
|
||||
@@ -112,15 +108,12 @@ def make_vla_jepa_pre_post_processors(
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
features = {**config.input_features, **config.output_features}
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps: list[ProcessorStep] = []
|
||||
if config.clip_normalized_actions:
|
||||
@@ -129,6 +122,8 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
# NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA
|
||||
# unnormalizes over BOTH input and output features.
|
||||
output_steps.append(
|
||||
UnnormalizerProcessorStep(
|
||||
features=features,
|
||||
@@ -140,16 +135,5 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
output_steps.append(DeviceProcessorStep(device="cpu"))
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -20,17 +20,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_vqbet import VQBeTConfig
|
||||
|
||||
@@ -62,32 +55,4 @@ def make_vqbet_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -58,10 +58,14 @@ class WallXConfig(PreTrainedConfig):
|
||||
# Action prediction mode: "diffusion" or "fast"
|
||||
prediction_mode: str = "diffusion"
|
||||
|
||||
# Attention Implementation, options: "eager", "flash_attention_2", "sdpa"
|
||||
# NOTE: flash-attn==2.7.4.post1 is required for flash_attention_2 implementation
|
||||
# Wall-X's bidirectional action-token islands currently require eager attention.
|
||||
attn_implementation: str = "eager"
|
||||
|
||||
# Vision attention is independent from the text action-token mask. ``auto`` uses
|
||||
# PyTorch's packed variable-length attention when the runtime supports it and
|
||||
# otherwise falls back to the native per-chunk SDPA implementation.
|
||||
vision_attn_implementation: str = "auto"
|
||||
|
||||
# ==================== Optimizer Presets ====================
|
||||
optimizer_lr: float = 2e-5
|
||||
optimizer_betas: tuple[float, float] = (0.9, 0.95)
|
||||
@@ -86,6 +90,18 @@ class WallXConfig(PreTrainedConfig):
|
||||
if self.prediction_mode not in ["diffusion", "fast"]:
|
||||
raise ValueError(f"prediction_mode must be 'diffusion' or 'fast', got {self.prediction_mode}")
|
||||
|
||||
if self.attn_implementation != "eager":
|
||||
raise ValueError(
|
||||
"Wall-X currently supports only attn_implementation='eager' because its "
|
||||
"bidirectional action-token islands require an explicit attention mask."
|
||||
)
|
||||
|
||||
if self.vision_attn_implementation not in {"auto", "sdpa", "varlen"}:
|
||||
raise ValueError(
|
||||
"vision_attn_implementation must be one of 'auto', 'sdpa', or 'varlen', got "
|
||||
f"{self.vision_attn_implementation!r}"
|
||||
)
|
||||
|
||||
# Assign use_fast_tokenizer based on prediction_mode
|
||||
if self.prediction_mode == "fast":
|
||||
self.use_fast_tokenizer = True
|
||||
|
||||
@@ -43,11 +43,14 @@ from typing import TYPE_CHECKING, Any
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
import torch.nn.functional as functional
|
||||
from safetensors import SafetensorError
|
||||
from safetensors.torch import load_file
|
||||
from torch import Tensor
|
||||
from torch.distributions import Beta
|
||||
from torch.nn import CrossEntropyLoss
|
||||
from torchvision.transforms import InterpolationMode
|
||||
from torchvision.transforms.v2 import functional as tv_functional
|
||||
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
from lerobot.utils.import_utils import (
|
||||
@@ -74,17 +77,17 @@ if TYPE_CHECKING or _wallx_deps_available:
|
||||
from qwen_vl_utils.vision_process import smart_resize
|
||||
from torchdiffeq import odeint
|
||||
from transformers import AutoProcessor, BatchFeature
|
||||
from transformers.cache_utils import StaticCache
|
||||
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
||||
Qwen2_5_VisionTransformerPretrainedModel,
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
)
|
||||
from transformers.utils import is_torchdynamo_compiling
|
||||
from transformers.utils import cached_file, is_torchdynamo_compiling
|
||||
|
||||
from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
|
||||
from .qwen_model.qwen2_5_vl_moe import (
|
||||
Qwen2_5_VisionTransformerPretrainedModel,
|
||||
from .qwen_model import (
|
||||
Qwen2_5_VLACausalLMOutputWithPast,
|
||||
Qwen2_5_VLConfig,
|
||||
Qwen2_5_VLMoEModel,
|
||||
configure_wall_x_vision_attention,
|
||||
)
|
||||
else:
|
||||
LoraConfig = None
|
||||
@@ -93,13 +96,14 @@ else:
|
||||
odeint = None
|
||||
AutoProcessor = None
|
||||
BatchFeature = None
|
||||
StaticCache = None
|
||||
Qwen2_5_VLForConditionalGeneration = None
|
||||
cached_file = None
|
||||
is_torchdynamo_compiling = None
|
||||
Qwen2_5_VLConfig = None
|
||||
Qwen2_5_VisionTransformerPretrainedModel = None
|
||||
Qwen2_5_VLACausalLMOutputWithPast = None
|
||||
Qwen2_5_VLMoEModel = None
|
||||
configure_wall_x_vision_attention = None
|
||||
|
||||
from .utils import (
|
||||
get_wallx_normal_text,
|
||||
@@ -111,6 +115,75 @@ from .utils import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _wall_x_resize_dimensions(height: int, width: int) -> tuple[int, int, int, int]:
|
||||
"""Return the intermediate and final Wall-X resize dimensions as ``(H, W, H, W)``."""
|
||||
if RESOLUTION == -1:
|
||||
intermediate_height, intermediate_width = height, width
|
||||
elif width > height:
|
||||
intermediate_width = RESOLUTION
|
||||
intermediate_height = int(RESOLUTION * height / width)
|
||||
else:
|
||||
intermediate_height = RESOLUTION
|
||||
intermediate_width = int(RESOLUTION * width / height)
|
||||
|
||||
resized_height, resized_width = smart_resize(
|
||||
intermediate_height,
|
||||
intermediate_width,
|
||||
factor=IMAGE_FACTOR,
|
||||
min_pixels=MIN_PIXELS,
|
||||
max_pixels=MAX_PIXELS,
|
||||
)
|
||||
return intermediate_height, intermediate_width, resized_height, resized_width
|
||||
|
||||
|
||||
def _resize_wall_x_image_batch(images: Tensor) -> tuple[Tensor, tuple[int, int, int, int]]:
|
||||
"""Quantize and resize a BCHW camera batch without leaving its current device."""
|
||||
if images.ndim != 4:
|
||||
raise ValueError(f"Wall-X images must be BCHW tensors, got shape {tuple(images.shape)}")
|
||||
|
||||
original_height, original_width = images.shape[-2:]
|
||||
intermediate_height, intermediate_width, resized_height, resized_width = _wall_x_resize_dimensions(
|
||||
original_height, original_width
|
||||
)
|
||||
|
||||
if images.is_floating_point():
|
||||
# Match the previous PIL path, which quantized via `(image * 255).to(torch.uint8)`.
|
||||
images = (images * 255).to(torch.uint8)
|
||||
elif images.dtype != torch.uint8:
|
||||
raise TypeError(f"Wall-X images must be floating point or uint8, got {images.dtype}")
|
||||
|
||||
if images.shape[-2:] != (intermediate_height, intermediate_width):
|
||||
images = tv_functional.resize(
|
||||
images,
|
||||
[intermediate_height, intermediate_width],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True,
|
||||
)
|
||||
if images.shape[-2:] != (resized_height, resized_width):
|
||||
images = tv_functional.resize(
|
||||
images,
|
||||
[resized_height, resized_width],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True,
|
||||
)
|
||||
|
||||
return images, (original_height, original_width, resized_height, resized_width)
|
||||
|
||||
|
||||
def _prepare_wall_x_image_inputs(
|
||||
batch: dict[str, Any], img_keys: list[str]
|
||||
) -> tuple[list[list[Tensor]], dict[str, tuple[int, int, int, int]]]:
|
||||
"""Resize each camera as a batch, then restore sample-major/camera-minor ordering."""
|
||||
resized_by_key: dict[str, Tensor] = {}
|
||||
dimensions_by_key: dict[str, tuple[int, int, int, int]] = {}
|
||||
for key in img_keys:
|
||||
resized_by_key[key], dimensions_by_key[key] = _resize_wall_x_image_batch(batch[key])
|
||||
|
||||
batch_size = batch[img_keys[0]].shape[0]
|
||||
image_inputs = [[resized_by_key[key][i] for key in img_keys] for i in range(batch_size)]
|
||||
return image_inputs, dimensions_by_key
|
||||
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
"""Sinusoidal positional embedding for diffusion timesteps."""
|
||||
|
||||
@@ -246,7 +319,7 @@ class ActionHead(nn.Module):
|
||||
flow = flow.to(torch.float32)
|
||||
|
||||
action_pred = self.action_proj_back(action_hidden_states)
|
||||
loss = F.mse_loss(action_pred, flow, reduction="none")
|
||||
loss = functional.mse_loss(action_pred, flow, reduction="none")
|
||||
|
||||
if dof_mask is not None:
|
||||
dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]).to(torch.float32)
|
||||
@@ -254,7 +327,7 @@ class ActionHead(nn.Module):
|
||||
|
||||
return loss
|
||||
|
||||
def proprioception_proj(self, proprioception, dof_mask=None, use_history=False):
|
||||
def proprioception_proj(self, proprioception, dof_mask=None):
|
||||
"""Project proprioceptive data to hidden space."""
|
||||
# Ensure proper device and dtype alignment
|
||||
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
||||
@@ -264,10 +337,7 @@ class ActionHead(nn.Module):
|
||||
if dof_mask is not None:
|
||||
# Concatenate proprioception with DOF mask
|
||||
# TODO: Use variable-based dimension checking for better flexibility
|
||||
if use_history:
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
else:
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
|
||||
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
||||
dtype=self.propri_proj.weight.dtype
|
||||
@@ -281,7 +351,7 @@ class ActionHead(nn.Module):
|
||||
_Qwen2_5_VLForAction_Base = Qwen2_5_VLForConditionalGeneration if _wallx_deps_available else nn.Module
|
||||
|
||||
|
||||
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base): # noqa: N801
|
||||
"""
|
||||
Qwen2.5 Vision-Language Mixture of Experts model for action processing.
|
||||
|
||||
@@ -305,6 +375,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config=None,
|
||||
action_tokenizer_path=None,
|
||||
attn_implementation: str = "eager",
|
||||
vision_attn_implementation: str = "auto",
|
||||
cache_dir: str | PathLike | None = None,
|
||||
force_download: bool = False,
|
||||
local_files_only: bool = False,
|
||||
@@ -321,11 +392,14 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path
|
||||
action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config
|
||||
attn_implementation (str, optional): Attention implementation, if None will load from default config
|
||||
vision_attn_implementation (str, optional): Vision attention backend. ``auto`` uses packed
|
||||
variable-length attention when supported and otherwise falls back to SDPA.
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Qwen2_5_VLMoEForAction: Loaded model instance
|
||||
"""
|
||||
Qwen2_5_VLMoEModel._require_eager_attention(attn_implementation)
|
||||
if config is None:
|
||||
config = cls.config_class.from_pretrained(
|
||||
pretrained_name_or_path,
|
||||
@@ -339,7 +413,15 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
)
|
||||
if attn_implementation is not None:
|
||||
config._attn_implementation = attn_implementation
|
||||
processor = AutoProcessor.from_pretrained(pretrained_name_or_path, use_fast=True)
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
pretrained_name_or_path,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
local_files_only=local_files_only,
|
||||
token=token,
|
||||
revision=revision,
|
||||
use_fast=True,
|
||||
)
|
||||
if action_tokenizer_path is not None:
|
||||
action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
|
||||
processor.action_processor = action_tokenizer
|
||||
@@ -351,41 +433,41 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config.text_config.pad_token_id = processor.tokenizer.pad_token_id
|
||||
|
||||
# Initialize model with configuration and processor
|
||||
model = cls(config, processor=processor, action_tokenizer=action_tokenizer, **kwargs)
|
||||
model = cls(
|
||||
config,
|
||||
processor=processor,
|
||||
action_tokenizer=action_tokenizer,
|
||||
vision_attn_implementation=vision_attn_implementation,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Resize token embeddings to match processor tokenizer vocabulary size
|
||||
model.resize_token_embeddings(len(processor.tokenizer))
|
||||
|
||||
# Try to load the model.safetensors file
|
||||
print(f"Loading model from: {pretrained_name_or_path}")
|
||||
logger.info("Loading Wall-X model from %s", pretrained_name_or_path)
|
||||
try:
|
||||
from transformers.utils import cached_file
|
||||
|
||||
# Try safetensors first
|
||||
resolved_file = cached_file(
|
||||
pretrained_name_or_path,
|
||||
"model.safetensors",
|
||||
cache_dir=kwargs.get("cache_dir"),
|
||||
force_download=kwargs.get("force_download", False),
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
resume_download=kwargs.get("resume_download"),
|
||||
proxies=kwargs.get("proxies"),
|
||||
token=kwargs.get("token"),
|
||||
revision=kwargs.get("revision"),
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
from safetensors.torch import load_file
|
||||
|
||||
sd = load_file(resolved_file)
|
||||
print("✓ Loaded state dict from model.safetensors")
|
||||
except Exception as e:
|
||||
print(f"Could not load state dict from remote files: {e}")
|
||||
print("Returning model without loading pretrained weights")
|
||||
return model
|
||||
except (OSError, SafetensorError) as error:
|
||||
raise OSError(
|
||||
f"Failed to load pretrained Wall-X weights from {pretrained_name_or_path!r}"
|
||||
) from error
|
||||
logger.info("Loaded Wall-X state dict from model.safetensors")
|
||||
|
||||
state_dict = {}
|
||||
# filter normalizer statistic params
|
||||
del_keys = []
|
||||
for key in sd.keys():
|
||||
for key in sd:
|
||||
if "action_preprocessor.normalizer" in key:
|
||||
del_keys.append(key)
|
||||
for key in del_keys:
|
||||
@@ -404,6 +486,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
action_tokenizer=None,
|
||||
action_mapper=None,
|
||||
flow_loss_weight=1.0,
|
||||
vision_attn_implementation: str = "auto",
|
||||
):
|
||||
"""
|
||||
Initialize the Qwen2.5 VLMoE model for action processing.
|
||||
@@ -416,10 +499,16 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
action_mapper: Action mapping utility
|
||||
flow_loss_weight (float): Weight for flow loss computation
|
||||
"""
|
||||
Qwen2_5_VLMoEModel._require_eager_attention(config._attn_implementation)
|
||||
config._attn_implementation = "eager"
|
||||
# Text needs eager attention for action-token islands. Vision has no such
|
||||
# constraint, so keep its portable native fallback on SDPA.
|
||||
config.vision_config._attn_implementation = "sdpa"
|
||||
super().__init__(config)
|
||||
|
||||
# Initialize vision transformer and language model components
|
||||
self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
|
||||
configure_wall_x_vision_attention(self.visual, vision_attn_implementation)
|
||||
self.model = Qwen2_5_VLMoEModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
@@ -457,7 +546,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
|
||||
params_to_keep_float32 = []
|
||||
|
||||
for name, param in self.named_parameters():
|
||||
for name, _param in self.named_parameters():
|
||||
if "input_layernorm" in name or "post_attention_layernorm" in name or "model.norm" in name:
|
||||
params_to_keep_float32.append(name)
|
||||
if "action_preprocessor" in name:
|
||||
@@ -491,7 +580,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
"action_token_id": action_token_id,
|
||||
}
|
||||
|
||||
def add_lora(self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1):
|
||||
def add_lora(self, r=8, lora_alpha=32, target_modules=None, lora_dropout=0.1):
|
||||
"""
|
||||
Add LoRA (Low-Rank Adaptation) adapters to the model.
|
||||
|
||||
@@ -501,6 +590,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
target_modules (list): List of module names to apply LoRA to
|
||||
lora_dropout (float): Dropout probability for LoRA layers
|
||||
"""
|
||||
if target_modules is None:
|
||||
target_modules = ["q_proj", "v_proj"]
|
||||
|
||||
config = LoraConfig(
|
||||
r=r,
|
||||
lora_alpha=lora_alpha,
|
||||
@@ -795,6 +887,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if rope_deltas is not None:
|
||||
self.rope_deltas = rope_deltas
|
||||
|
||||
# Calculate RoPE position IDs if not provided
|
||||
# Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation
|
||||
if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
|
||||
@@ -833,7 +928,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process image embeddings
|
||||
if pixel_values is not None:
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
|
||||
mask = input_ids == self.config.image_token_id
|
||||
mask_unsqueezed = mask.unsqueeze(-1)
|
||||
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
|
||||
@@ -845,7 +940,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process video embeddings
|
||||
if pixel_values_videos is not None:
|
||||
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
|
||||
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
|
||||
n_video_features = video_embeds.shape[0]
|
||||
|
||||
@@ -869,7 +964,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
proprioception = self.action_preprocessor.proprioception_proj(
|
||||
proprioception,
|
||||
agent_pos_mask,
|
||||
use_history=proprioception.shape[1] > 1,
|
||||
)
|
||||
mask = input_ids == self.action_token_id_set["propri_token_id"]
|
||||
mask_unsqueezed = mask.unsqueeze(-1)
|
||||
@@ -919,6 +1013,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
@@ -1107,7 +1202,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process image embeddings
|
||||
if pixel_values is not None:
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
|
||||
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
|
||||
n_image_features = image_embeds.shape[0]
|
||||
|
||||
@@ -1128,7 +1223,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process video embeddings
|
||||
if pixel_values_videos is not None:
|
||||
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
|
||||
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
|
||||
n_video_features = video_embeds.shape[0]
|
||||
|
||||
@@ -1153,7 +1248,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
proprio_embed = self.action_preprocessor.proprioception_proj(
|
||||
proprioception,
|
||||
agent_pos_mask,
|
||||
use_history=proprioception.shape[1] > 1,
|
||||
)
|
||||
proprioception_mask = input_ids == self.action_token_id_set["propri_token_id"]
|
||||
proprio_embed = proprio_embed.to(torch.bfloat16)
|
||||
@@ -1202,25 +1296,37 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
|
||||
# Split input sequence for text and fast modes (not needed for diffusion)
|
||||
if predict_mode == "text" or predict_mode == "fast":
|
||||
# Look for generation prompt tokens: <|im_start|>assistant
|
||||
generation_prompt = "<|im_start|>assistant\n"
|
||||
generation_prompt_ids = torch.tensor(
|
||||
[151644, 77091], device=input_ids.device, dtype=input_ids.dtype
|
||||
)
|
||||
matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
|
||||
input_ids[0, 1:] == generation_prompt_ids[1]
|
||||
self.processor.tokenizer.encode(generation_prompt, add_special_tokens=False),
|
||||
device=input_ids.device,
|
||||
dtype=input_ids.dtype,
|
||||
)
|
||||
prompt_length = generation_prompt_ids.numel()
|
||||
if prompt_length == 0:
|
||||
raise ValueError(f"Tokenizer produced no tokens for generation prompt {generation_prompt!r}")
|
||||
if input_ids.shape[1] < prompt_length:
|
||||
matches = torch.empty(0, device=input_ids.device, dtype=torch.bool)
|
||||
else:
|
||||
matches = (
|
||||
input_ids[0]
|
||||
.unfold(dimension=0, size=prompt_length, step=1)
|
||||
.eq(generation_prompt_ids)
|
||||
.all(dim=-1)
|
||||
)
|
||||
|
||||
if matches.any():
|
||||
split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
|
||||
prompt_end = split_pos + prompt_length
|
||||
# Extract ground truth output tokens (including newline)
|
||||
gt_output_ids = input_ids[:, split_pos + 3 :]
|
||||
gt_output_ids = input_ids[:, prompt_end:]
|
||||
# Remove output part from input, keeping prompt
|
||||
input_ids = input_ids[:, : split_pos + 3]
|
||||
inputs_embeds = inputs_embeds[:, : split_pos + 3, :]
|
||||
input_ids = input_ids[:, :prompt_end]
|
||||
inputs_embeds = inputs_embeds[:, :prompt_end, :]
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask[:, : split_pos + 3]
|
||||
attention_mask = attention_mask[:, :prompt_end]
|
||||
if labels is not None:
|
||||
labels = labels[:, split_pos + 3 :]
|
||||
labels = labels[:, prompt_end:]
|
||||
else:
|
||||
raise ValueError(
|
||||
"input_ids does not contain the generation prompt tokens <|im_start|>assistant"
|
||||
@@ -1255,7 +1361,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
use_cache=True,
|
||||
pad_token_id=self.processor.tokenizer.pad_token_id,
|
||||
temperature=(1.0 if not re_generate else 0.7), # Higher temperature for regeneration
|
||||
do_sample=(False if not re_generate else True), # Enable sampling for regeneration
|
||||
do_sample=re_generate, # Enable sampling for regeneration
|
||||
)
|
||||
|
||||
# Decode generated and ground truth text
|
||||
@@ -1524,27 +1630,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
else:
|
||||
model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
|
||||
|
||||
# Prepare 4D causal attention mask for static cache
|
||||
if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
|
||||
if model_inputs["inputs_embeds"] is not None:
|
||||
batch_size, sequence_length, _ = inputs_embeds.shape
|
||||
device = inputs_embeds.device
|
||||
else:
|
||||
batch_size, sequence_length = input_ids.shape
|
||||
device = input_ids.device
|
||||
|
||||
attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
|
||||
attention_mask,
|
||||
sequence_length=sequence_length,
|
||||
target_length=past_key_values.get_max_cache_shape(),
|
||||
dtype=self.lm_head.weight.dtype,
|
||||
device=device,
|
||||
cache_position=cache_position,
|
||||
batch_size=batch_size,
|
||||
config=self.config,
|
||||
past_key_values=past_key_values,
|
||||
)
|
||||
|
||||
# Assemble all model inputs for generation
|
||||
model_inputs.update(
|
||||
{
|
||||
@@ -1749,6 +1834,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
pretrained_name_or_path=config.pretrained_name_or_path,
|
||||
action_tokenizer_path=config.action_tokenizer_path,
|
||||
attn_implementation=config.attn_implementation,
|
||||
vision_attn_implementation=config.vision_attn_implementation,
|
||||
)
|
||||
self.model.to(config.device)
|
||||
self.model.to_bfloat16_for_selected_params()
|
||||
@@ -1768,6 +1854,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
def preprocess_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
*,
|
||||
compute_position_ids: bool = False,
|
||||
) -> BatchFeature:
|
||||
"""
|
||||
Convert a batch of LeRobot dataset items to Wall-X model input format.
|
||||
@@ -1789,50 +1877,21 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
# Get batch size from state tensor
|
||||
batch_size = batch[OBS_STATE].shape[0]
|
||||
|
||||
# ==================== PROCESS ALL SAMPLES ====================
|
||||
all_image_inputs = []
|
||||
all_texts = []
|
||||
|
||||
# Find image keys in batch
|
||||
img_keys = [key for key in self.config.image_features if key in batch]
|
||||
if not img_keys:
|
||||
raise ValueError("Wall-X requires at least one image feature in each batch")
|
||||
|
||||
# Resize one camera batch at a time on the tensors' current device. Reassembling
|
||||
# sample-major keeps image_grid_thw aligned with each sample's image placeholders.
|
||||
all_image_inputs, dimensions_by_key = _prepare_wall_x_image_inputs(batch, img_keys)
|
||||
all_texts = []
|
||||
|
||||
# Preserve the existing grounding behavior for multi-camera inputs: the old camera
|
||||
# loop left these values set to the final configured camera's dimensions.
|
||||
orig_height, orig_width, resized_height, resized_width = dimensions_by_key[img_keys[-1]]
|
||||
|
||||
for i in range(batch_size):
|
||||
# Vision preprocessing per sample
|
||||
processed_frames = []
|
||||
orig_height, orig_width = None, None
|
||||
resized_height, resized_width = None, None
|
||||
|
||||
for key in img_keys:
|
||||
current_obs = batch[key][i].clone() # (C, H, W)
|
||||
if current_obs.dim() == 3:
|
||||
current_obs = current_obs.permute(1, 2, 0) # (H, W, C)
|
||||
|
||||
img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
|
||||
orig_width, orig_height = img_pil.size
|
||||
|
||||
target_size = RESOLUTION
|
||||
if target_size != -1:
|
||||
if orig_width > orig_height:
|
||||
new_width = target_size
|
||||
new_height = int(target_size * orig_height / orig_width)
|
||||
else:
|
||||
new_height = target_size
|
||||
new_width = int(target_size * orig_width / orig_height)
|
||||
img_pil = img_pil.resize((new_width, new_height))
|
||||
|
||||
current_width, current_height = img_pil.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
current_height,
|
||||
current_width,
|
||||
factor=IMAGE_FACTOR,
|
||||
min_pixels=MIN_PIXELS,
|
||||
max_pixels=MAX_PIXELS,
|
||||
)
|
||||
resized_img = img_pil.resize((resized_width, resized_height))
|
||||
processed_frames.append(resized_img)
|
||||
|
||||
all_image_inputs.append(processed_frames)
|
||||
|
||||
# Text preprocessing
|
||||
task_text = batch["task"][i] if isinstance(batch["task"], list) else batch["task"]
|
||||
instruction_info = {"instruction": task_text}
|
||||
@@ -1859,8 +1918,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
agent_pos_mask = (~torch.isnan(agent_pos)).float()
|
||||
agent_pos = agent_pos.nan_to_num(nan=0.0)
|
||||
|
||||
if agent_pos.shape[-1] != 20:
|
||||
pad_size = 20 - agent_pos.shape[-1]
|
||||
if agent_pos.shape[-1] < self.config.max_state_dim:
|
||||
pad_size = self.config.max_state_dim - agent_pos.shape[-1]
|
||||
agent_pos = torch.cat(
|
||||
[
|
||||
agent_pos,
|
||||
@@ -1880,6 +1939,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
elif agent_pos.shape[-1] > self.config.max_state_dim:
|
||||
raise ValueError(
|
||||
f"State dimension {agent_pos.shape[-1]} exceeds max_state_dim {self.config.max_state_dim}"
|
||||
)
|
||||
|
||||
# ==================== PROCESS ACTIONS ====================
|
||||
action = batch.get(ACTION) # (batch_size, chunk_size, action_dim)
|
||||
@@ -1889,8 +1952,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
dof_mask = (~torch.isnan(action)).float()
|
||||
action = action.nan_to_num(nan=0.0)
|
||||
|
||||
if action.shape[-1] != 20:
|
||||
pad_size = 20 - action.shape[-1]
|
||||
if action.shape[-1] < self.config.max_action_dim:
|
||||
pad_size = self.config.max_action_dim - action.shape[-1]
|
||||
action = torch.cat(
|
||||
[action, torch.zeros(action.shape[0], action.shape[1], pad_size, device=action.device)],
|
||||
dim=-1,
|
||||
@@ -1902,6 +1965,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
elif action.shape[-1] > self.config.max_action_dim:
|
||||
raise ValueError(
|
||||
f"Action dimension {action.shape[-1]} exceeds max_action_dim {self.config.max_action_dim}"
|
||||
)
|
||||
else:
|
||||
action_dim = self.config.output_features[ACTION].shape[0]
|
||||
dof_mask = torch.cat(
|
||||
@@ -1910,7 +1977,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
batch_size, self.config.chunk_size, action_dim, device=batch[OBS_STATE].device
|
||||
),
|
||||
torch.zeros(
|
||||
batch_size, self.config.chunk_size, 20 - action_dim, device=batch[OBS_STATE].device
|
||||
batch_size,
|
||||
self.config.chunk_size,
|
||||
self.config.max_action_dim - action_dim,
|
||||
device=batch[OBS_STATE].device,
|
||||
),
|
||||
],
|
||||
dim=-1,
|
||||
@@ -1930,12 +2000,26 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
text=all_texts,
|
||||
images=all_image_inputs,
|
||||
videos=None,
|
||||
device=batch[OBS_STATE].device,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
max_length=TOKENIZER_MAX_LENGTH,
|
||||
)
|
||||
|
||||
if compute_position_ids:
|
||||
# Qwen's RoPE indexing uses Python list/scalar conversions. Run it while the
|
||||
# tokenizer and grid metadata are still on CPU, then move the compact result.
|
||||
position_ids, rope_deltas = self.model.get_rope_index(
|
||||
inputs.input_ids,
|
||||
inputs.get("image_grid_thw"),
|
||||
inputs.get("video_grid_thw"),
|
||||
inputs.get("second_per_grid_ts"),
|
||||
inputs.attention_mask,
|
||||
)
|
||||
inputs["position_ids"] = position_ids
|
||||
inputs["rope_deltas"] = rope_deltas
|
||||
|
||||
# ==================== ADDITIONAL INPUTS ====================
|
||||
action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids("<|action|>")
|
||||
moe_token_types = inputs.input_ids == action_token_id
|
||||
@@ -1952,7 +2036,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
)
|
||||
|
||||
# Move all tensors to the correct device
|
||||
device = self.config.device
|
||||
device = batch[OBS_STATE].device
|
||||
for key, value in inputs.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
inputs[key] = value.to(device)
|
||||
@@ -1972,9 +2056,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
Returns:
|
||||
tuple: (loss, loss_dict)
|
||||
"""
|
||||
batch = self.preprocess_inputs(
|
||||
batch,
|
||||
)
|
||||
batch = self.preprocess_inputs(batch, compute_position_ids=True)
|
||||
|
||||
# Call the underlying model's forward with mode="train"
|
||||
outputs = self.model(**batch, mode="train")
|
||||
@@ -1982,19 +2064,19 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
# Extract losses from output
|
||||
loss = outputs.loss
|
||||
loss_dict = {
|
||||
"loss": loss.item() if loss is not None else 0.0,
|
||||
"loss": loss.detach() if loss is not None else 0.0,
|
||||
}
|
||||
|
||||
if outputs.flow_loss is not None:
|
||||
loss_dict["flow_loss"] = outputs.flow_loss.item()
|
||||
loss_dict["flow_loss"] = outputs.flow_loss.detach()
|
||||
if outputs.cross_entropy_loss is not None:
|
||||
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.item()
|
||||
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.detach()
|
||||
|
||||
# Add channel losses if available
|
||||
if outputs.channel_loss_dict is not None:
|
||||
for key, value in outputs.channel_loss_dict.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
loss_dict[f"channel_{key}"] = value.item()
|
||||
loss_dict[f"channel_{key}"] = value.detach()
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
|
||||
@@ -20,19 +20,13 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_wall_x import WallXConfig
|
||||
|
||||
@@ -65,37 +59,22 @@ def make_wall_x_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
WallXTaskProcessor(), # Process task description
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="wall_x_task_processor")
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .configuration_qwen2_5_vl import (
|
||||
Qwen2_5_VLConfig,
|
||||
Qwen2_5_VLTextConfig,
|
||||
Qwen2_5_VLVisionConfig,
|
||||
)
|
||||
from .qwen2_5_vl_moe import (
|
||||
BlockSparseMLP,
|
||||
Qwen2_5_VLACausalLMOutputWithPast,
|
||||
Qwen2_5_VLDecoderLayer_with_MoE,
|
||||
Qwen2_5_VLMoEModel,
|
||||
SparseMoeBlock,
|
||||
)
|
||||
from .vision_attention import (
|
||||
WallXVisionAttention,
|
||||
configure_wall_x_vision_attention,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BlockSparseMLP",
|
||||
"Qwen2_5_VLACausalLMOutputWithPast",
|
||||
"Qwen2_5_VLConfig",
|
||||
"Qwen2_5_VLDecoderLayer_with_MoE",
|
||||
"Qwen2_5_VLMoEModel",
|
||||
"Qwen2_5_VLTextConfig",
|
||||
"Qwen2_5_VLVisionConfig",
|
||||
"SparseMoeBlock",
|
||||
"WallXVisionAttention",
|
||||
"configure_wall_x_vision_attention",
|
||||
]
|
||||
@@ -1,250 +1,114 @@
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_rope_utils import rope_config_validation
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Wall-X configuration extensions for the native Transformers Qwen2.5-VL config."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub.dataclasses import strict
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
|
||||
Qwen2_5_VLConfig as TransformersQwen2_5_VLConfig,
|
||||
Qwen2_5_VLTextConfig as TransformersQwen2_5_VLTextConfig,
|
||||
Qwen2_5_VLVisionConfig,
|
||||
)
|
||||
else:
|
||||
|
||||
@dataclass
|
||||
class _TransformersConfigFallback:
|
||||
"""Import-safe stand-in used only when Transformers is unavailable."""
|
||||
|
||||
TransformersQwen2_5_VLConfig = _TransformersConfigFallback
|
||||
TransformersQwen2_5_VLTextConfig = _TransformersConfigFallback
|
||||
Qwen2_5_VLVisionConfig = None
|
||||
|
||||
# Wall-X checkpoints pre0.6.0 use the legacy, flat Qwen2.5-VL config layout. The native
|
||||
# ``Qwen2_5_VLConfig`` accepts that layout and moves text-model fields into its
|
||||
# ``text_config`` sub-config, so only the Wall-X-specific MoE fields need to be
|
||||
# declared here.
|
||||
_LEGACY_TEXT_ATTRIBUTES = {
|
||||
"attention_dropout",
|
||||
"attention_moe",
|
||||
"dim_inputs",
|
||||
"dof_config",
|
||||
"experts",
|
||||
"hidden_act",
|
||||
"hidden_size",
|
||||
"initializer_range",
|
||||
"intermediate_size",
|
||||
"layer_types",
|
||||
"max_position_embeddings",
|
||||
"max_window_layers",
|
||||
"mlp_moe",
|
||||
"noise_scheduler",
|
||||
"num_attention_heads",
|
||||
"num_experts",
|
||||
"num_hidden_layers",
|
||||
"num_key_value_heads",
|
||||
"pad_token_id",
|
||||
"rms_norm_eps",
|
||||
"sliding_window",
|
||||
"use_cache",
|
||||
"use_sliding_window",
|
||||
"vocab_size",
|
||||
}
|
||||
|
||||
|
||||
class Qwen2_5_VLVisionConfig(PretrainedConfig):
|
||||
model_type = "qwen2_5_vl"
|
||||
base_config_key = "vision_config"
|
||||
@strict
|
||||
class Qwen2_5_VLTextConfig(TransformersQwen2_5_VLTextConfig): # noqa: N801
|
||||
"""Native Qwen2.5-VL text config plus Wall-X's hard-routed MoE settings."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
depth=32,
|
||||
hidden_size=3584,
|
||||
hidden_act="silu",
|
||||
intermediate_size=3420,
|
||||
num_heads=16,
|
||||
in_channels=3,
|
||||
patch_size=14,
|
||||
spatial_merge_size=2,
|
||||
temporal_patch_size=2,
|
||||
tokens_per_second=4,
|
||||
window_size=112,
|
||||
out_hidden_size=3584,
|
||||
fullatt_block_indexes=[7, 15, 23, 31],
|
||||
initializer_range=0.02,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
num_experts: int = 4
|
||||
experts: list[dict] | None = None
|
||||
dof_config: dict | None = None
|
||||
noise_scheduler: dict | None = None
|
||||
dim_inputs: tuple[int, ...] | list[int] = (1536, 1536)
|
||||
attention_moe: bool = False
|
||||
mlp_moe: bool = False
|
||||
|
||||
self.depth = depth
|
||||
self.hidden_size = hidden_size
|
||||
self.hidden_act = hidden_act
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_heads = num_heads
|
||||
self.in_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.spatial_merge_size = spatial_merge_size
|
||||
self.temporal_patch_size = temporal_patch_size
|
||||
self.tokens_per_second = tokens_per_second
|
||||
self.window_size = window_size
|
||||
self.fullatt_block_indexes = fullatt_block_indexes
|
||||
self.out_hidden_size = out_hidden_size
|
||||
self.initializer_range = initializer_range
|
||||
def __post_init__(self, **kwargs):
|
||||
self.dim_inputs = tuple(self.dim_inputs)
|
||||
super().__post_init__(**kwargs)
|
||||
|
||||
|
||||
class Qwen2_5_VLConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
|
||||
Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||
with the defaults will yield a similar configuration to that of
|
||||
Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
|
||||
@strict
|
||||
class Qwen2_5_VLConfig(TransformersQwen2_5_VLConfig): # noqa: N801
|
||||
"""Native composite Qwen2.5-VL config with a Wall-X text sub-config.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
The native composite loader supports both current nested configs and the
|
||||
flat layout used by existing ``wall-oss-flow`` checkpoints.
|
||||
"""
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 152064):
|
||||
Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`Qwen2_5_VLModel`]
|
||||
hidden_size (`int`, *optional*, defaults to 8192):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 29568):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 80):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 64):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
num_key_value_heads (`int`, *optional*, defaults to 8):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details checkout [this
|
||||
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||
The maximum sequence length that this model might ever be used with.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-05):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether the model's input and output word embeddings should be tied.
|
||||
rope_theta (`float`, *optional*, defaults to 1000000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use sliding window attention.
|
||||
sliding_window (`int`, *optional*, defaults to 4096):
|
||||
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
|
||||
max_window_layers (`int`, *optional*, defaults to 80):
|
||||
The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
vision_config (`Dict`, *optional*):
|
||||
The config for the visual encoder initialization.
|
||||
rope_scaling (`Dict`, *optional*):
|
||||
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||
accordingly.
|
||||
Expected contents:
|
||||
`rope_type` (`str`):
|
||||
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||
'llama3'], with 'default' being the original RoPE implementation.
|
||||
`factor` (`float`, *optional*):
|
||||
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||
original maximum pre-trained length.
|
||||
`original_max_position_embeddings` (`int`, *optional*):
|
||||
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||
pretraining.
|
||||
`attention_factor` (`float`, *optional*):
|
||||
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||
`factor` field to infer the suggested value.
|
||||
`beta_fast` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 32.
|
||||
`beta_slow` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 1.
|
||||
`short_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`long_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`low_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||
`high_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
|
||||
|
||||
>>> # Initializing a Qwen2_5_VL style configuration
|
||||
>>> configuration = Qwen2_5_VLConfig()
|
||||
|
||||
>>> # Initializing a model from the Qwen2-VL-7B style configuration
|
||||
>>> model = Qwen2_5_VLForConditionalGeneration(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "qwen2_5_vl"
|
||||
sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
# Default tensor parallel plan for base model `Qwen2_5_VL`
|
||||
base_model_tp_plan = {
|
||||
"layers.*.self_attn.q_proj": "colwise",
|
||||
"layers.*.self_attn.k_proj": "colwise",
|
||||
"layers.*.self_attn.v_proj": "colwise",
|
||||
"layers.*.self_attn.o_proj": "rowwise",
|
||||
"layers.*.mlp.gate_proj": "colwise",
|
||||
"layers.*.mlp.up_proj": "colwise",
|
||||
"layers.*.mlp.down_proj": "rowwise",
|
||||
}
|
||||
base_model_pp_plan = {
|
||||
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
||||
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
||||
"norm": (["hidden_states"], ["hidden_states"]),
|
||||
sub_configs = {
|
||||
"vision_config": Qwen2_5_VLVisionConfig,
|
||||
"text_config": Qwen2_5_VLTextConfig,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=152064,
|
||||
hidden_size=8192,
|
||||
intermediate_size=29568,
|
||||
num_hidden_layers=80,
|
||||
num_attention_heads=64,
|
||||
num_key_value_heads=8,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=32768,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-05,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=1000000.0,
|
||||
use_sliding_window=False,
|
||||
sliding_window=4096,
|
||||
max_window_layers=80,
|
||||
attention_dropout=0.0,
|
||||
vision_config=None,
|
||||
rope_scaling=None,
|
||||
num_experts=4,
|
||||
experts=None,
|
||||
dof_config=None,
|
||||
noise_scheduler=None,
|
||||
dim_inputs=(1536, 1536),
|
||||
attention_moe=False,
|
||||
mlp_moe=False,
|
||||
**kwargs,
|
||||
):
|
||||
if isinstance(vision_config, dict):
|
||||
self.vision_config = self.sub_configs["vision_config"](**vision_config)
|
||||
elif vision_config is None:
|
||||
self.vision_config = self.sub_configs["vision_config"]()
|
||||
def __getattr__(self, name):
|
||||
"""Keep legacy direct access to fields now owned by ``text_config``.
|
||||
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.use_sliding_window = use_sliding_window
|
||||
self.sliding_window = sliding_window
|
||||
self.max_window_layers = max_window_layers
|
||||
self.layer_types = ["dense"] * num_hidden_layers
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.attention_dropout = attention_dropout
|
||||
self.rope_scaling = rope_scaling
|
||||
|
||||
self.num_experts = num_experts
|
||||
self.experts = experts
|
||||
self.dof_config = dof_config
|
||||
self.noise_scheduler = noise_scheduler
|
||||
self.dim_inputs = tuple(dim_inputs)
|
||||
self.attention_moe = attention_moe
|
||||
self.mlp_moe = mlp_moe
|
||||
|
||||
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||
if self.rope_scaling["type"] == "mrope":
|
||||
self.rope_scaling["type"] = "default"
|
||||
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||
rope_config_validation(self, ignore_keys={"mrope_section"})
|
||||
|
||||
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
|
||||
|
||||
@property
|
||||
def text_config(self):
|
||||
return self
|
||||
|
||||
|
||||
__all__ = ["Qwen2_5_VLConfig"]
|
||||
Wall-X historically used a flat config and accesses fields such as
|
||||
``hidden_size`` and ``num_experts`` directly. Forwarding unknown
|
||||
attributes preserves that API without duplicating the native config.
|
||||
"""
|
||||
text_config = self.__dict__.get("text_config")
|
||||
if name in _LEGACY_TEXT_ATTRIBUTES and text_config is not None and hasattr(text_config, name):
|
||||
return getattr(text_config, name)
|
||||
raise AttributeError(f"{type(self).__name__!s} has no attribute {name!r}")
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,208 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Wall-X vision attention backends.
|
||||
|
||||
Qwen2.5-VL's native non-Flash vision path splits a packed image sequence into
|
||||
Python-level chunks before calling attention. Wall-X batches many camera frames,
|
||||
so that path launches thousands of tiny attention operations per training step.
|
||||
This module keeps the native SDPA path as a portable fallback and adds a packed
|
||||
``torch.nn.attention.varlen`` path that consumes Qwen's existing ``cu_seqlens``
|
||||
metadata directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
||||
Qwen2_5_VLVisionAttention,
|
||||
apply_rotary_pos_emb_vision,
|
||||
)
|
||||
else:
|
||||
Qwen2_5_VLVisionAttention = nn.Module
|
||||
apply_rotary_pos_emb_vision = None
|
||||
|
||||
try:
|
||||
from torch.nn.attention.varlen import varlen_attn as _varlen_attn
|
||||
except ImportError: # torch<2.10
|
||||
_varlen_attn = None
|
||||
|
||||
_VARLEN_USES_WINDOW_SIZE = (
|
||||
_varlen_attn is not None and "window_size" in inspect.signature(_varlen_attn).parameters
|
||||
)
|
||||
|
||||
|
||||
VisionAttentionBackend = Literal["auto", "sdpa", "varlen"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def _log_resolved_backend(requested: str, resolved: str) -> None:
|
||||
logger.info("Wall-X vision attention backend: %s (requested: %s)", resolved, requested)
|
||||
|
||||
|
||||
def _varlen_unavailable_reason(
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> str | None:
|
||||
if _varlen_attn is None:
|
||||
return "torch.nn.attention.varlen is unavailable (PyTorch 2.10 or newer is required)"
|
||||
if position_embeddings is None:
|
||||
return "precomputed vision position embeddings were not provided"
|
||||
if hidden_states.device.type != "cuda" or torch.version.cuda is None:
|
||||
return "packed varlen attention requires an NVIDIA CUDA device"
|
||||
if hidden_states.dtype not in {torch.float16, torch.bfloat16}:
|
||||
return f"packed varlen attention requires float16 or bfloat16 inputs, got {hidden_states.dtype}"
|
||||
major, _minor = torch.cuda.get_device_capability(hidden_states.device)
|
||||
if major < 8:
|
||||
return "packed varlen attention requires an NVIDIA Ampere GPU or newer"
|
||||
return None
|
||||
|
||||
|
||||
def _supports_varlen_attention(
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> bool:
|
||||
return _varlen_unavailable_reason(hidden_states, position_embeddings) is None
|
||||
|
||||
|
||||
class WallXVisionAttention(Qwen2_5_VLVisionAttention):
|
||||
"""Qwen2.5-VL vision attention with packed varlen and native SDPA fallback."""
|
||||
|
||||
def __init__(self, config, backend: VisionAttentionBackend):
|
||||
super().__init__(config)
|
||||
self.wallx_backend = backend
|
||||
self._resolved_backend_key = None
|
||||
self._resolved_backend = None
|
||||
|
||||
def _resolve_backend(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> str:
|
||||
key = (
|
||||
hidden_states.device.type,
|
||||
hidden_states.device.index,
|
||||
hidden_states.dtype,
|
||||
position_embeddings is not None,
|
||||
)
|
||||
if self._resolved_backend_key == key:
|
||||
return self._resolved_backend
|
||||
|
||||
use_varlen = self.wallx_backend != "sdpa" and _supports_varlen_attention(
|
||||
hidden_states, position_embeddings
|
||||
)
|
||||
if self.wallx_backend == "varlen" and not use_varlen:
|
||||
reason = _varlen_unavailable_reason(hidden_states, position_embeddings)
|
||||
raise RuntimeError(f"Wall-X vision_attn_implementation='varlen' cannot be used: {reason}")
|
||||
|
||||
resolved_backend = "varlen" if use_varlen else "sdpa"
|
||||
self._resolved_backend_key = key
|
||||
self._resolved_backend = resolved_backend
|
||||
_log_resolved_backend(self.wallx_backend, resolved_backend)
|
||||
return resolved_backend
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor,
|
||||
rotary_pos_emb: torch.Tensor | None = None,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
del rotary_pos_emb
|
||||
|
||||
if self._resolve_backend(hidden_states, position_embeddings) == "sdpa":
|
||||
return super().forward(
|
||||
hidden_states=hidden_states,
|
||||
cu_seqlens=cu_seqlens,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
seq_length = hidden_states.shape[0]
|
||||
query_states, key_states, value_states = (
|
||||
self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
|
||||
)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb_vision(
|
||||
query_states,
|
||||
key_states,
|
||||
cos,
|
||||
sin,
|
||||
)
|
||||
|
||||
if cu_seqlens.dtype != torch.int32:
|
||||
cu_seqlens = cu_seqlens.to(dtype=torch.int32)
|
||||
max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item())
|
||||
varlen_kwargs = {"scale": self.scaling}
|
||||
if _VARLEN_USES_WINDOW_SIZE:
|
||||
varlen_kwargs["window_size"] = (-1, -1)
|
||||
else: # Stable PyTorch 2.10 API; pre-release variants used window_size.
|
||||
varlen_kwargs["is_causal"] = False
|
||||
attn_output = _varlen_attn(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
cu_seqlens,
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
max_seqlen,
|
||||
**varlen_kwargs,
|
||||
)
|
||||
attn_output = attn_output.reshape(seq_length, -1).contiguous()
|
||||
return self.proj(attn_output)
|
||||
|
||||
|
||||
def configure_wall_x_vision_attention(
|
||||
vision_model: nn.Module,
|
||||
backend: VisionAttentionBackend,
|
||||
) -> None:
|
||||
"""Install Wall-X's scoped packed attention without changing checkpoint keys."""
|
||||
if backend == "sdpa":
|
||||
_log_resolved_backend(backend, "sdpa")
|
||||
return
|
||||
if backend == "varlen" and _varlen_attn is None:
|
||||
raise RuntimeError(
|
||||
"Wall-X vision_attn_implementation='varlen' requires torch.nn.attention.varlen "
|
||||
"from PyTorch 2.10 or newer"
|
||||
)
|
||||
if backend == "auto" and _varlen_attn is None:
|
||||
_log_resolved_backend(backend, "sdpa")
|
||||
return
|
||||
|
||||
for block in vision_model.blocks:
|
||||
previous_attention = block.attn
|
||||
replacement = WallXVisionAttention(previous_attention.config, backend=backend)
|
||||
replacement.to(
|
||||
device=previous_attention.qkv.weight.device,
|
||||
dtype=previous_attention.qkv.weight.dtype,
|
||||
)
|
||||
replacement.load_state_dict(previous_attention.state_dict(), strict=True)
|
||||
replacement.train(previous_attention.training)
|
||||
block.attn = replacement
|
||||
@@ -116,6 +116,7 @@ def preprocesser_call(
|
||||
images: list | Any | None = None,
|
||||
text: str | list[str] | None = None,
|
||||
videos: list | Any | None = None,
|
||||
device: torch.device | str | None = None,
|
||||
padding: bool | str = False,
|
||||
truncation: bool | None = None,
|
||||
max_length: int | None = None,
|
||||
@@ -134,6 +135,7 @@ def preprocesser_call(
|
||||
images: Input images (PIL, numpy arrays, or torch tensors)
|
||||
text: Text or list of texts to tokenize
|
||||
videos: Input videos (numpy arrays or torch tensors)
|
||||
device: Device on which image/video preprocessing should run
|
||||
padding: Whether to pad sequences to same length
|
||||
truncation: Whether to truncate sequences longer than max_length
|
||||
max_length: Maximum length for truncation/padding
|
||||
@@ -151,7 +153,11 @@ def preprocesser_call(
|
||||
"""
|
||||
# Process image inputs
|
||||
if images is not None and len(images) > 0:
|
||||
image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
|
||||
image_inputs = processor.image_processor(
|
||||
images=images,
|
||||
return_tensors=return_tensors,
|
||||
device=device,
|
||||
)
|
||||
image_grid_thw = image_inputs["image_grid_thw"]
|
||||
else:
|
||||
image_inputs = {}
|
||||
@@ -159,7 +165,11 @@ def preprocesser_call(
|
||||
|
||||
# Process video inputs
|
||||
if videos is not None:
|
||||
videos_inputs = processor.image_processor(videos=videos, return_tensors=return_tensors)
|
||||
videos_inputs = processor.image_processor(
|
||||
videos=videos,
|
||||
return_tensors=return_tensors,
|
||||
device=device,
|
||||
)
|
||||
video_grid_thw = videos_inputs["video_grid_thw"]
|
||||
else:
|
||||
videos_inputs = {}
|
||||
@@ -413,10 +423,7 @@ def get_task_instruction(
|
||||
}
|
||||
)
|
||||
|
||||
if priority_order is not None:
|
||||
priority_order = OrderedDict(priority_order)
|
||||
else:
|
||||
priority_order = default_priority_order
|
||||
priority_order = OrderedDict(priority_order) if priority_order is not None else default_priority_order
|
||||
|
||||
got_instruction = False
|
||||
task_instruction = ""
|
||||
@@ -424,9 +431,8 @@ def get_task_instruction(
|
||||
# Sample instruction components based on priority probabilities
|
||||
for key, prob in priority_order.items():
|
||||
if key in frame_instruction_info and frame_instruction_info[key] != "":
|
||||
if got_instruction:
|
||||
if random.random() >= prob:
|
||||
continue
|
||||
if got_instruction and random.random() >= prob:
|
||||
continue
|
||||
|
||||
task_instruction += f"\n{frame_instruction_info[key]}"
|
||||
got_instruction = True
|
||||
@@ -538,10 +544,7 @@ def img_key_mapping(img_keys: list[str]) -> list[str]:
|
||||
if key in CAMERA_NAME_MAPPING:
|
||||
key = CAMERA_NAME_MAPPING[key]
|
||||
else:
|
||||
if "view" in key:
|
||||
key = key.replace("_", " ")
|
||||
else:
|
||||
key = key + " view"
|
||||
key = key.replace("_", " ") if "view" in key else key + " view"
|
||||
processed_img_keys.append(key)
|
||||
return processed_img_keys
|
||||
|
||||
|
||||
@@ -22,19 +22,14 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
ObservationProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
@@ -42,8 +37,6 @@ from lerobot.utils.constants import (
|
||||
OBS_IMAGES,
|
||||
OBS_PREFIX,
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_xvla import XVLAConfig
|
||||
@@ -61,10 +54,11 @@ def make_xvla_pre_post_processors(
|
||||
Build the LeRobot processor pipelines for XVLA.
|
||||
"""
|
||||
|
||||
features = {**config.input_features, **config.output_features}
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.tokenizer_name,
|
||||
max_length=config.tokenizer_max_length,
|
||||
@@ -74,32 +68,15 @@ def make_xvla_pre_post_processors(
|
||||
XVLAImageToFloatProcessorStep(),
|
||||
XVLAImageNetNormalizeProcessorStep(),
|
||||
XVLAAddDomainIdProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
|
||||
# Custom XVLA processor steps
|
||||
|
||||
@@ -42,10 +42,14 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
|
||||
from .factory import (
|
||||
DefaultPolicyProcessorSteps,
|
||||
make_default_policy_processor_steps,
|
||||
make_default_pre_post_processors,
|
||||
make_default_processors,
|
||||
make_default_robot_action_processor,
|
||||
make_default_robot_observation_processor,
|
||||
make_default_teleop_action_processor,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from .gym_action_processor import (
|
||||
Numpy2TorchActionProcessorStep,
|
||||
@@ -129,10 +133,14 @@ __all__ = [
|
||||
"ImageCropResizeProcessorStep",
|
||||
"InfoProcessorStep",
|
||||
"InterventionActionProcessorStep",
|
||||
"DefaultPolicyProcessorSteps",
|
||||
"make_default_policy_processor_steps",
|
||||
"make_default_pre_post_processors",
|
||||
"make_default_processors",
|
||||
"make_default_teleop_action_processor",
|
||||
"make_default_robot_action_processor",
|
||||
"make_default_robot_observation_processor",
|
||||
"make_policy_processor_pipelines",
|
||||
"AbsoluteActionsProcessorStep",
|
||||
"RelativeActionsProcessorStep",
|
||||
"MapDeltaActionToRobotActionStep",
|
||||
|
||||
@@ -14,15 +14,33 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .batch_processor import AddBatchDimensionProcessorStep
|
||||
from .converters import (
|
||||
observation_to_transition,
|
||||
policy_action_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_policy_action,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
|
||||
from .pipeline import (
|
||||
IdentityProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RobotProcessorPipeline,
|
||||
)
|
||||
from .rename_processor import RenameObservationsProcessorStep
|
||||
|
||||
|
||||
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
|
||||
@@ -61,3 +79,97 @@ def make_default_processors():
|
||||
robot_action_processor = make_default_robot_action_processor()
|
||||
robot_observation_processor = make_default_robot_observation_processor()
|
||||
return (teleop_action_processor, robot_action_processor, robot_observation_processor)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DefaultPolicyProcessorSteps:
|
||||
"""The canonical processor steps shared by most policies' pre/post pipelines.
|
||||
|
||||
Policies compose these in their own order (step ORDER is a Hub-serialized contract
|
||||
and intentionally stays explicit per policy) and interleave their custom steps.
|
||||
"""
|
||||
|
||||
rename_observations: RenameObservationsProcessorStep
|
||||
add_batch_dim: AddBatchDimensionProcessorStep
|
||||
to_device: DeviceProcessorStep
|
||||
normalize: NormalizerProcessorStep
|
||||
unnormalize: UnnormalizerProcessorStep
|
||||
to_cpu: DeviceProcessorStep
|
||||
|
||||
|
||||
def make_default_policy_processor_steps(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> DefaultPolicyProcessorSteps:
|
||||
"""Construct the canonical policy processor steps from a policy config.
|
||||
|
||||
Args:
|
||||
config: A `PreTrainedConfig` providing `device`, `input_features`,
|
||||
`output_features` and `normalization_mapping`.
|
||||
dataset_stats: Dataset statistics used for (un)normalization.
|
||||
normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
|
||||
their normalization stats to the policy device; most leave it unset).
|
||||
"""
|
||||
return DefaultPolicyProcessorSteps(
|
||||
rename_observations=RenameObservationsProcessorStep(rename_map={}),
|
||||
add_batch_dim=AddBatchDimensionProcessorStep(),
|
||||
to_device=DeviceProcessorStep(device=config.device),
|
||||
normalize=NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=normalizer_device,
|
||||
),
|
||||
unnormalize=UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
to_cpu=DeviceProcessorStep(device="cpu"),
|
||||
)
|
||||
|
||||
|
||||
def make_policy_processor_pipelines(
|
||||
input_steps: list[ProcessorStep],
|
||||
output_steps: list[ProcessorStep],
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""Wrap pre/post step lists into the canonical policy pipeline pair.
|
||||
|
||||
Uses the standard pipeline names (which determine the serialized JSON filenames on
|
||||
the Hub) and the standard policy-action converters on the postprocessor.
|
||||
"""
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_default_pre_post_processors(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
|
||||
and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
|
||||
compose `make_default_policy_processor_steps` themselves instead.
|
||||
"""
|
||||
s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
|
||||
return make_policy_processor_pipelines(
|
||||
input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
|
||||
output_steps=[s.unnormalize, s.to_cpu],
|
||||
)
|
||||
|
||||
@@ -21,8 +21,6 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -30,6 +28,7 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.configs.rewards import RewardModelConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -129,29 +128,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
if missing_keys:
|
||||
logging.warning(f"Missing key(s) when loading model: {missing_keys}")
|
||||
if unexpected_keys:
|
||||
logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
def get_optim_params(self):
|
||||
|
||||
@@ -28,7 +28,12 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
|
||||
from lerobot.annotations.steerable_pipeline.executor import Executor
|
||||
@@ -42,6 +47,12 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
|
||||
from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
|
||||
from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
|
||||
from lerobot.configs import parser
|
||||
from lerobot.utils.import_utils import _datasets_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _datasets_available:
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
from lerobot.datasets.utils import create_lerobot_dataset_card
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -50,8 +61,6 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
|
||||
if cfg.root is not None:
|
||||
return Path(cfg.root)
|
||||
if cfg.repo_id is not None:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
|
||||
raise ValueError("Either --root or --repo_id must be provided.")
|
||||
|
||||
@@ -125,7 +134,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
|
||||
Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
|
||||
"""
|
||||
from huggingface_hub import HfApi # noqa: PLC0415
|
||||
require_package("datasets", "dataset")
|
||||
|
||||
repo_id = cfg.new_repo_id or cfg.repo_id
|
||||
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
||||
@@ -143,33 +152,26 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
repo_id=repo_id,
|
||||
repo_type="dataset",
|
||||
commit_message=commit_message,
|
||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
|
||||
# README.md is excluded because when pushing to ``new_repo_id`` the
|
||||
# source card's links (e.g. the visualize badge) would keep pointing
|
||||
# at the source dataset; a fresh card is generated below instead.
|
||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
||||
)
|
||||
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
|
||||
|
||||
dataset_info = load_info(root)
|
||||
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
||||
card.push_to_hub(repo_id=repo_id, repo_type="dataset")
|
||||
|
||||
# Tag the upload with the codebase version. ``LeRobotDatasetMetadata``
|
||||
# resolves the dataset revision via ``get_safe_version`` which scans
|
||||
# for tags like ``v3.0``; without a tag it raises
|
||||
# ``RevisionNotFoundError``. Read the version straight from the
|
||||
# dataset's own ``meta/info.json`` so we tag whatever the writer
|
||||
# actually wrote (no accidental drift if the codebase floor moves).
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION # noqa: PLC0415
|
||||
|
||||
info_path = root / "meta" / "info.json"
|
||||
version_tag = CODEBASE_VERSION
|
||||
if info_path.exists():
|
||||
try:
|
||||
from lerobot.utils.io_utils import load_json # noqa: PLC0415
|
||||
|
||||
info = load_json(info_path)
|
||||
ds_version = info.get("codebase_version")
|
||||
if isinstance(ds_version, str) and ds_version.startswith("v"):
|
||||
version_tag = ds_version
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(
|
||||
f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
|
||||
flush=True,
|
||||
)
|
||||
version_tag = (
|
||||
dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
|
||||
)
|
||||
revision = getattr(commit_info, "oid", None)
|
||||
tag_kwargs = {
|
||||
"repo_id": repo_id,
|
||||
@@ -180,10 +182,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
tag_kwargs["revision"] = revision
|
||||
|
||||
try:
|
||||
from contextlib import suppress # noqa: PLC0415
|
||||
|
||||
from huggingface_hub.errors import RevisionNotFoundError # noqa: PLC0415
|
||||
|
||||
with suppress(RevisionNotFoundError):
|
||||
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
||||
api.create_tag(**tag_kwargs)
|
||||
|
||||
@@ -171,6 +171,9 @@ def update_policy(
|
||||
train_metrics.update_s = time.perf_counter() - start_time
|
||||
if torch.cuda.is_available():
|
||||
train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
# Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
|
||||
if output_dict:
|
||||
train_metrics.update_metrics(output_dict)
|
||||
return train_metrics, output_dict
|
||||
|
||||
|
||||
@@ -572,7 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
batch = preprocessor(batch)
|
||||
train_tracker.dataloading_s = time.perf_counter() - start_time
|
||||
|
||||
train_tracker, output_dict = update_policy(
|
||||
train_tracker, _ = update_policy(
|
||||
train_tracker,
|
||||
policy,
|
||||
batch,
|
||||
@@ -605,9 +608,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
train_tracker.samples_per_s = effective_batch_size / step_time
|
||||
logging.info(train_tracker)
|
||||
if wandb_logger:
|
||||
# Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
|
||||
# tracker by update_policy, so to_dict() already carries their windowed,
|
||||
# rank-reduced averages — no per-step output_dict passthrough needed.
|
||||
wandb_log_dict = train_tracker.to_dict()
|
||||
if output_dict:
|
||||
wandb_log_dict.update(output_dict)
|
||||
# Log sample weighting statistics if enabled
|
||||
if sample_weighter is not None:
|
||||
weighter_stats = sample_weighter.get_stats()
|
||||
|
||||
@@ -59,6 +59,20 @@ def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||
return device
|
||||
|
||||
|
||||
def resolve_safetensors_device(map_location: str | torch.device) -> str:
|
||||
"""Resolve a device string for a safetensors load, working around a device-mapping quirk.
|
||||
|
||||
safetensors' load maps the bare string "cuda" to cuda:0 regardless of the current device
|
||||
(unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under multi-GPU
|
||||
accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it before sharding.
|
||||
Resolve "cuda" to the concrete current-device index so each rank loads onto its own GPU.
|
||||
"""
|
||||
map_location = str(map_location)
|
||||
if map_location == "cuda" and torch.cuda.is_available():
|
||||
return f"cuda:{torch.cuda.current_device()}"
|
||||
return map_location
|
||||
|
||||
|
||||
def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
|
||||
"""
|
||||
mps is currently not compatible with float64
|
||||
|
||||
@@ -104,6 +104,7 @@ class MetricsTracker:
|
||||
"episodes",
|
||||
"epochs",
|
||||
"accelerator",
|
||||
"_caller_metrics",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
@@ -129,6 +130,9 @@ class MetricsTracker:
|
||||
self.episodes = self.samples / self._avg_samples_per_ep
|
||||
self.epochs = self.samples / self._num_frames
|
||||
self.accelerator = accelerator
|
||||
# Meter names the caller registered up front. update_metrics() leaves these untouched, so a
|
||||
# policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
|
||||
self._caller_metrics: set[str] = set(self.metrics)
|
||||
|
||||
def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
|
||||
if name in self.__dict__:
|
||||
@@ -156,6 +160,21 @@ class MetricsTracker:
|
||||
self.episodes = self.samples / self._avg_samples_per_ep
|
||||
self.epochs = self.samples / self._num_frames
|
||||
|
||||
def update_metrics(self, values: dict[str, Any]) -> None:
|
||||
"""Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
|
||||
|
||||
Non-numeric values and bools are ignored.
|
||||
Caller-registered metrics (those passed to the constructor) are never overridden.
|
||||
"""
|
||||
for name, value in values.items():
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
continue
|
||||
if name in self._caller_metrics:
|
||||
continue
|
||||
if name not in self.metrics:
|
||||
self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
|
||||
self.metrics[name].update(float(value))
|
||||
|
||||
def reduce_across_ranks(self) -> None:
|
||||
"""
|
||||
Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Behavior-pinning tests for the shared flow-matching sampling primitives.
|
||||
|
||||
``euler_integrate`` is compared against a verbatim copy of the historical pi0/pi05/
|
||||
smolvla sampling loop (including its RTC hook semantics): any divergence from that
|
||||
reference is a behavior change for released checkpoints.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.policies.common.flow_matching import (
|
||||
euler_integrate,
|
||||
sample_beta,
|
||||
sample_noise,
|
||||
sample_time_beta,
|
||||
)
|
||||
|
||||
|
||||
def test_sample_beta_range_dtype_and_reproducibility():
|
||||
torch.manual_seed(0)
|
||||
s1 = sample_beta(1.5, 1.0, 4096, "cpu")
|
||||
torch.manual_seed(0)
|
||||
s2 = sample_beta(1.5, 1.0, 4096, "cpu")
|
||||
assert torch.equal(s1, s2)
|
||||
assert s1.shape == (4096,) and s1.dtype == torch.float32
|
||||
assert s1.min() >= 0.0 and s1.max() <= 1.0
|
||||
# Beta(1.5, 1.0) mean is 1.5/2.5 = 0.6.
|
||||
assert abs(s1.mean().item() - 0.6) < 0.02
|
||||
|
||||
|
||||
def test_sample_time_beta_openpi_convention():
|
||||
torch.manual_seed(1)
|
||||
time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, offset=0.001)
|
||||
assert time.dtype == torch.float32
|
||||
assert time.min() >= 0.001 and time.max() <= 1.0
|
||||
# Exact composition: Beta sample * scale + offset, same RNG stream.
|
||||
torch.manual_seed(1)
|
||||
expected = sample_beta(1.5, 1.0, 4096, "cpu") * 0.999 + 0.001
|
||||
torch.testing.assert_close(time, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_sample_noise_seeded():
|
||||
torch.manual_seed(2)
|
||||
n1 = sample_noise((2, 8, 4), "cpu")
|
||||
torch.manual_seed(2)
|
||||
n2 = sample_noise((2, 8, 4), "cpu")
|
||||
assert torch.equal(n1, n2)
|
||||
assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4)
|
||||
|
||||
|
||||
def test_euler_integrate_constant_velocity_is_exact():
|
||||
# With v_t == c constant, x_0 = x_1 + sum(dt * c) = x_1 - c exactly (num_steps * dt = -1).
|
||||
noise = torch.randn(3, 5, 2)
|
||||
c = torch.randn(3, 5, 2)
|
||||
out = euler_integrate(lambda x_t, time: c, noise, num_steps=10)
|
||||
torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6)
|
||||
|
||||
|
||||
def test_euler_integrate_forward_constant_velocity_is_exact():
|
||||
# Forward convention: dt = +1/num_steps, so x_1 = x_0 + sum(dt * c) = x_0 + c exactly.
|
||||
noise = torch.randn(3, 5, 2)
|
||||
c = torch.randn(3, 5, 2)
|
||||
out = euler_integrate(lambda x_t, time: c, noise, num_steps=10, forward_euler=True)
|
||||
torch.testing.assert_close(out, noise + c, rtol=0, atol=1e-6)
|
||||
|
||||
|
||||
def _reference_forward_loop(denoise_fn, noise, num_steps):
|
||||
"""Verbatim structure of the groot/evo1/wall_x forward-Euler loop (t: 0 -> 1)."""
|
||||
bsize = noise.shape[0]
|
||||
device = noise.device
|
||||
dt = 1.0 / num_steps
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 0.0 + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
x_t = x_t + dt * denoise_fn(x_t, time_tensor)
|
||||
return x_t
|
||||
|
||||
|
||||
def test_euler_integrate_forward_matches_reference_loop():
|
||||
torch.manual_seed(7)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
ref = _reference_forward_loop(denoise_fn, noise, 10)
|
||||
out = euler_integrate(denoise_fn, noise, 10, forward_euler=True)
|
||||
assert torch.equal(out, ref)
|
||||
|
||||
|
||||
def _reference_pi0_loop(denoise_fn, noise, num_steps, rtc_enabled, rtc_processor, kw):
|
||||
"""Verbatim structure of the historical pi0/pi05/smolvla sample_actions loop."""
|
||||
bsize = noise.shape[0]
|
||||
device = noise.device
|
||||
dt = -1.0 / num_steps
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 1.0 + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return denoise_fn(input_x_t, current_timestep)
|
||||
|
||||
if rtc_enabled:
|
||||
v_t = rtc_processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=kw.get("prev_chunk_left_over"),
|
||||
inference_delay=kw.get("inference_delay"),
|
||||
time=time,
|
||||
original_denoise_step_partial=denoise_step_partial_call,
|
||||
execution_horizon=kw.get("execution_horizon"),
|
||||
)
|
||||
else:
|
||||
v_t = denoise_step_partial_call(x_t)
|
||||
x_t = x_t + dt * v_t
|
||||
if rtc_processor is not None and rtc_processor.is_debug_enabled():
|
||||
rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
|
||||
return x_t
|
||||
|
||||
|
||||
class _StubRTCProcessor:
|
||||
def __init__(self, debug_enabled: bool):
|
||||
self._debug = debug_enabled
|
||||
self.tracked = []
|
||||
self.guidance_calls = []
|
||||
|
||||
def is_debug_enabled(self):
|
||||
return self._debug
|
||||
|
||||
def denoise_step(
|
||||
self,
|
||||
x_t,
|
||||
prev_chunk_left_over,
|
||||
inference_delay,
|
||||
time,
|
||||
original_denoise_step_partial,
|
||||
execution_horizon,
|
||||
):
|
||||
self.guidance_calls.append(
|
||||
{
|
||||
"time": time,
|
||||
"inference_delay": inference_delay,
|
||||
"execution_horizon": execution_horizon,
|
||||
"x_t": x_t.clone(),
|
||||
}
|
||||
)
|
||||
return original_denoise_step_partial(x_t) * 0.5
|
||||
|
||||
def track(self, time, x_t, v_t):
|
||||
self.tracked.append({"time": time, "x_t": x_t.clone(), "v_t": v_t.clone()})
|
||||
|
||||
|
||||
def _make_denoise_fn():
|
||||
weight = torch.randn(4, 4) * 0.1
|
||||
|
||||
def denoise_fn(x_t, time_tensor):
|
||||
return x_t @ weight + time_tensor[:, None, None]
|
||||
|
||||
return denoise_fn
|
||||
|
||||
|
||||
def test_euler_integrate_matches_historical_loop():
|
||||
torch.manual_seed(3)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
ref = _reference_pi0_loop(denoise_fn, noise, 10, rtc_enabled=False, rtc_processor=None, kw={})
|
||||
out = euler_integrate(denoise_fn, noise, 10)
|
||||
assert torch.equal(out, ref)
|
||||
|
||||
|
||||
def test_euler_integrate_rtc_guidance_and_kwarg_forwarding():
|
||||
torch.manual_seed(4)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
leftover = torch.randn(2, 6, 4)
|
||||
kw = {"inference_delay": 3, "prev_chunk_left_over": leftover, "execution_horizon": 25}
|
||||
|
||||
ref_proc, new_proc = _StubRTCProcessor(False), _StubRTCProcessor(False)
|
||||
ref = _reference_pi0_loop(denoise_fn, noise, 6, rtc_enabled=True, rtc_processor=ref_proc, kw=kw)
|
||||
out = euler_integrate(
|
||||
denoise_fn,
|
||||
noise,
|
||||
6,
|
||||
rtc_processor=new_proc,
|
||||
rtc_enabled=True,
|
||||
inference_delay=3,
|
||||
prev_chunk_left_over=leftover,
|
||||
execution_horizon=25,
|
||||
)
|
||||
assert torch.equal(out, ref)
|
||||
assert len(new_proc.guidance_calls) == 6
|
||||
for ref_call, new_call in zip(ref_proc.guidance_calls, new_proc.guidance_calls, strict=True):
|
||||
assert ref_call["time"] == new_call["time"]
|
||||
assert new_call["inference_delay"] == 3 and new_call["execution_horizon"] == 25
|
||||
# Guidance sees the PRE-update x_t.
|
||||
assert torch.equal(ref_call["x_t"], new_call["x_t"])
|
||||
|
||||
|
||||
def test_euler_integrate_debug_tracking_fires_even_when_rtc_disabled():
|
||||
# Historical behavior: track() fires whenever the processor exists and has debugging
|
||||
# enabled, independent of whether RTC guidance is active.
|
||||
torch.manual_seed(5)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
proc = _StubRTCProcessor(True)
|
||||
out = euler_integrate(denoise_fn, noise, 4, rtc_processor=proc, rtc_enabled=False)
|
||||
assert len(proc.guidance_calls) == 0
|
||||
assert len(proc.tracked) == 4
|
||||
# track() receives the POST-update x_t; the last one is the returned sample.
|
||||
assert torch.equal(proc.tracked[-1]["x_t"], out)
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Behavior-pinning tests for the shared VLA helpers.
|
||||
|
||||
These helpers are the canonical versions of functions that used to be copy-pasted across
|
||||
the openpi-derived policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla). The expected
|
||||
values below encode the historical per-policy behavior exactly; a failure here means a
|
||||
behavior change that would silently affect released checkpoints.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.common.vla_utils import (
|
||||
create_sinusoidal_pos_embedding,
|
||||
make_att_2d_masks,
|
||||
pad_vector,
|
||||
prepare_attention_masks_4d,
|
||||
resize_with_pad,
|
||||
resize_with_pad_torch,
|
||||
)
|
||||
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
|
||||
|
||||
|
||||
def test_create_sinusoidal_pos_embedding_matches_openpi_formula():
|
||||
time = torch.tensor([0.0, 0.25, 1.0])
|
||||
dim, min_period, max_period = 8, 4e-3, 4.0
|
||||
emb = create_sinusoidal_pos_embedding(time, dim, min_period, max_period, device=torch.device("cpu"))
|
||||
|
||||
assert emb.shape == (3, dim)
|
||||
# Independent recomputation of the openpi formula in float64.
|
||||
fraction = torch.linspace(0.0, 1.0, dim // 2, dtype=torch.float64)
|
||||
period = min_period * (max_period / min_period) ** fraction
|
||||
scaling = 1.0 / period * 2 * math.pi
|
||||
sin_input = scaling[None, :] * time.to(torch.float64)[:, None]
|
||||
expected = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
|
||||
torch.testing.assert_close(emb, expected, rtol=1e-9, atol=1e-9)
|
||||
|
||||
|
||||
def test_create_sinusoidal_pos_embedding_validation():
|
||||
with pytest.raises(ValueError, match="divisible by 2"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2), 7, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
with pytest.raises(ValueError, match="batch_size"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2, 2), 8, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
|
||||
|
||||
def test_make_att_2d_masks_docstring_cases():
|
||||
# Pure causal attention: [[1 1 1]]
|
||||
pad = torch.ones(1, 3, dtype=torch.bool)
|
||||
att = torch.tensor([[1, 1, 1]], dtype=torch.int32)
|
||||
expected = torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.bool)
|
||||
assert torch.equal(make_att_2d_masks(pad, att), expected)
|
||||
|
||||
# Prefix-LM: [[0 0 1 1]] -> first two tokens attend bidirectionally, rest causal.
|
||||
att = torch.tensor([[0, 0, 1, 1]], dtype=torch.int32)
|
||||
pad = torch.ones(1, 4, dtype=torch.bool)
|
||||
expected = torch.tensor([[[1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]], dtype=torch.bool)
|
||||
assert torch.equal(make_att_2d_masks(pad, att), expected)
|
||||
|
||||
# Padding removes rows and columns.
|
||||
pad = torch.tensor([[True, True, False]])
|
||||
att = torch.tensor([[0, 1, 1]], dtype=torch.int32)
|
||||
out = make_att_2d_masks(pad, att)
|
||||
assert not out[0, :, 2].any() and not out[0, 2, :].any()
|
||||
|
||||
|
||||
def test_make_att_2d_masks_validation():
|
||||
with pytest.raises(ValueError):
|
||||
make_att_2d_masks(torch.ones(3, dtype=torch.bool), torch.ones(1, 3, dtype=torch.int32))
|
||||
with pytest.raises(ValueError):
|
||||
make_att_2d_masks(torch.ones(1, 3, dtype=torch.bool), torch.ones(3, dtype=torch.int32))
|
||||
|
||||
|
||||
def test_prepare_attention_masks_4d():
|
||||
masks = torch.tensor([[[True, False], [False, True]]])
|
||||
out = prepare_attention_masks_4d(masks)
|
||||
assert out.shape == (1, 1, 2, 2)
|
||||
expected = torch.tensor([[[[0.0, OPENPI_ATTENTION_MASK_VALUE], [OPENPI_ATTENTION_MASK_VALUE, 0.0]]]])
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
out_bf16 = prepare_attention_masks_4d(masks, dtype=torch.bfloat16)
|
||||
assert out_bf16.dtype == torch.bfloat16
|
||||
assert torch.equal(out_bf16, expected.to(torch.bfloat16))
|
||||
|
||||
|
||||
def test_pad_vector_openpi_semantics():
|
||||
v = torch.arange(6.0).reshape(2, 3)
|
||||
padded = pad_vector(v, 5)
|
||||
assert padded.shape == (2, 5)
|
||||
assert torch.equal(padded[:, :3], v) and not padded[:, 3:].any()
|
||||
# Already large enough (>=): returned unchanged, same object.
|
||||
assert pad_vector(v, 3) is v
|
||||
assert pad_vector(v, 2) is v
|
||||
# 3D input.
|
||||
v3 = torch.ones(2, 4, 3)
|
||||
assert pad_vector(v3, 7).shape == (2, 4, 7)
|
||||
|
||||
|
||||
def test_pad_vector_truncate_semantics():
|
||||
v = torch.arange(6.0).reshape(2, 3)
|
||||
out = pad_vector(v, 2, truncate=True)
|
||||
assert out.shape == (2, 2) and torch.equal(out, v[:, :2])
|
||||
out = pad_vector(v, 5, truncate=True)
|
||||
assert out.shape == (2, 5) and torch.equal(out[:, :3], v) and not out[:, 3:].any()
|
||||
assert pad_vector(v, 0, truncate=True).shape == (2, 0)
|
||||
assert pad_vector(v, 3, truncate=True) is v
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels_last", [True, False])
|
||||
def test_resize_with_pad_torch_centered(channels_last):
|
||||
img = torch.rand(2, 3, 30, 60) if not channels_last else torch.rand(2, 30, 60, 3)
|
||||
out = resize_with_pad_torch(img, 64, 64)
|
||||
if channels_last:
|
||||
assert out.shape == (2, 64, 64, 3)
|
||||
# Aspect ratio preserved: 30x60 -> 32x64, padded 16 top and 16 bottom (centered).
|
||||
assert not out[:, :16].any() and not out[:, -16:].any()
|
||||
assert out[:, 16:48].abs().sum() > 0
|
||||
else:
|
||||
assert out.shape == (2, 3, 64, 64)
|
||||
assert not out[:, :, :16].any() and not out[:, :, -16:].any()
|
||||
|
||||
|
||||
def test_resize_with_pad_torch_uint8_roundtrip():
|
||||
img = (torch.rand(1, 3, 20, 20) * 255).to(torch.uint8)
|
||||
out = resize_with_pad_torch(img, 40, 40)
|
||||
assert out.dtype == torch.uint8 and out.shape == (1, 3, 40, 40)
|
||||
with pytest.raises(ValueError, match="Unsupported image dtype"):
|
||||
resize_with_pad_torch(torch.rand(1, 3, 8, 8, dtype=torch.float64), 16, 16)
|
||||
|
||||
|
||||
def test_resize_with_pad_top_left():
|
||||
img = torch.rand(2, 3, 30, 60)
|
||||
out = resize_with_pad(img, 64, 64, pad_value=-1.0)
|
||||
assert out.shape == (2, 3, 64, 64)
|
||||
# 30x60 -> 32x64; this variant pads on the TOP only (32 rows of pad_value).
|
||||
assert torch.equal(out[:, :, :32], torch.full((2, 3, 32, 64), -1.0))
|
||||
assert out[:, :, 32:].min() >= 0
|
||||
# No-op fast path returns the same object.
|
||||
assert resize_with_pad(img, 30, 60, pad_value=0.0) is img
|
||||
with pytest.raises(ValueError, match="expected"):
|
||||
resize_with_pad(torch.rand(3, 8, 8), 16, 16, pad_value=0.0)
|
||||
|
||||
|
||||
def test_clone_past_key_values():
|
||||
pytest.importorskip("transformers")
|
||||
from transformers import DynamicCache
|
||||
|
||||
from lerobot.policies.common.vla_utils import clone_past_key_values
|
||||
|
||||
cache = DynamicCache()
|
||||
keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
|
||||
cache.update(keys, values, 0)
|
||||
cloned = clone_past_key_values(cache)
|
||||
(ck, cv, _), (ok, ov, _) = next(iter(cloned)), next(iter(cache))
|
||||
assert torch.equal(ck, ok) and torch.equal(cv, ov)
|
||||
# Deep copy: mutating the clone must not touch the original.
|
||||
ck.zero_()
|
||||
assert not torch.equal(ck, ok)
|
||||
|
||||
|
||||
def test_clone_past_key_values_is_fullgraph_compilable():
|
||||
pytest.importorskip("transformers")
|
||||
from transformers import DynamicCache
|
||||
|
||||
from lerobot.policies.common.vla_utils import clone_past_key_values
|
||||
|
||||
cache = DynamicCache()
|
||||
keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
|
||||
cache.update(keys, values, 0)
|
||||
|
||||
compiled_clone = torch.compile(clone_past_key_values, backend="eager", fullgraph=True)
|
||||
cloned = compiled_clone(cache)
|
||||
|
||||
(cloned_keys, cloned_values, _), (original_keys, original_values, _) = (
|
||||
next(iter(cloned)),
|
||||
next(iter(cache)),
|
||||
)
|
||||
assert torch.equal(cloned_keys, original_keys)
|
||||
assert torch.equal(cloned_values, original_values)
|
||||
@@ -25,13 +25,57 @@ pytest.importorskip("transformers")
|
||||
pytest.importorskip("torchdiffeq")
|
||||
|
||||
from lerobot.policies.factory import make_policy_config # noqa: E402
|
||||
from lerobot.policies.wall_x import WallXConfig # noqa: E402
|
||||
from lerobot.policies.wall_x import (
|
||||
WallXConfig, # noqa: E402
|
||||
)
|
||||
from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy # noqa: E402
|
||||
from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors # noqa: E402
|
||||
from lerobot.policies.wall_x.qwen_model import Qwen2_5_VLMoEModel, Qwen2_5_VLTextConfig # noqa: E402
|
||||
from lerobot.utils.random_utils import set_seed # noqa: E402
|
||||
from tests.utils import require_cuda, require_hf_token # noqa: E402
|
||||
|
||||
|
||||
def test_moe_model_captures_requested_hidden_states_and_attentions():
|
||||
hidden_size = 16
|
||||
expert_config = {
|
||||
"hidden_size": hidden_size,
|
||||
"intermediate_size": 32,
|
||||
"hidden_act": "silu",
|
||||
}
|
||||
config = Qwen2_5_VLTextConfig(
|
||||
vocab_size=32,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=32,
|
||||
num_hidden_layers=2,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=4,
|
||||
max_position_embeddings=32,
|
||||
layer_types=["full_attention", "full_attention"],
|
||||
rope_parameters={
|
||||
"rope_type": "default",
|
||||
"rope_theta": 1_000_000.0,
|
||||
"mrope_section": [1, 1, 0],
|
||||
},
|
||||
num_experts=2,
|
||||
experts=[expert_config, expert_config],
|
||||
dim_inputs=(hidden_size, hidden_size),
|
||||
mlp_moe=True,
|
||||
)
|
||||
config._attn_implementation = "eager"
|
||||
model = Qwen2_5_VLMoEModel(config)
|
||||
input_ids = torch.tensor([[1, 2, 3]])
|
||||
|
||||
output = model(
|
||||
input_ids=input_ids,
|
||||
moe_token_types=torch.zeros_like(input_ids),
|
||||
output_hidden_states=True,
|
||||
output_attentions=True,
|
||||
)
|
||||
|
||||
assert len(output.hidden_states) == config.num_hidden_layers + 1
|
||||
assert len(output.attentions) == config.num_hidden_layers
|
||||
|
||||
|
||||
@require_cuda
|
||||
@require_hf_token
|
||||
def test_policy_instantiation():
|
||||
|
||||
@@ -18,6 +18,8 @@ import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
# ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
|
||||
# exercises) imports ``lerobot.datasets``, which only ships under the
|
||||
@@ -26,11 +28,13 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
|
||||
|
||||
|
||||
def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
from lerobot.scripts.lerobot_annotate import _push_to_hub
|
||||
from lerobot.scripts import lerobot_annotate
|
||||
|
||||
root = tmp_path / "dataset"
|
||||
(root / "meta").mkdir(parents=True)
|
||||
(root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
|
||||
(root / "meta" / "info.json").write_text(
|
||||
json.dumps({"codebase_version": "v3.0", "fps": 30, "features": {}})
|
||||
)
|
||||
|
||||
calls = {}
|
||||
|
||||
@@ -43,9 +47,6 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
return SimpleNamespace(oid="abc123")
|
||||
|
||||
def delete_tag(self, repo_id, **kwargs):
|
||||
import requests
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
|
||||
# Simulate the common case: no stale tag to delete.
|
||||
raise RevisionNotFoundError("no such tag", response=requests.Response())
|
||||
@@ -53,7 +54,12 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
def create_tag(self, **kwargs):
|
||||
calls["create_tag"] = kwargs
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
|
||||
monkeypatch.setattr(lerobot_annotate, "HfApi", FakeHfApi)
|
||||
|
||||
def fake_card_push(self, **kwargs):
|
||||
calls["card_push"] = {"content": str(self), **kwargs}
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.DatasetCard.push_to_hub", fake_card_push)
|
||||
|
||||
cfg = SimpleNamespace(
|
||||
repo_id="source/dataset",
|
||||
@@ -62,7 +68,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
push_commit_message=None,
|
||||
)
|
||||
|
||||
_push_to_hub(root, cfg)
|
||||
lerobot_annotate._push_to_hub(root, cfg)
|
||||
|
||||
assert calls["create_repo"] == {
|
||||
"repo_id": "annotated/dataset",
|
||||
@@ -71,6 +77,13 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
"exist_ok": True,
|
||||
}
|
||||
assert calls["upload_folder"]["repo_id"] == "annotated/dataset"
|
||||
# The source README must not be copied over: its links (e.g. the
|
||||
# visualize badge) point at the source dataset. A card regenerated for
|
||||
# the target repo is pushed instead.
|
||||
assert "README.md" in calls["upload_folder"]["ignore_patterns"]
|
||||
assert calls["card_push"]["repo_id"] == "annotated/dataset"
|
||||
assert "visualize_dataset?path=annotated/dataset" in calls["card_push"]["content"]
|
||||
assert "source/dataset" not in calls["card_push"]["content"]
|
||||
# A stale tag (e.g. from a previous annotation run) is deleted first so
|
||||
# the new tag always points at the upload we just made.
|
||||
assert calls["delete_tag"] == {
|
||||
|
||||
@@ -233,3 +233,37 @@ def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
|
||||
# accumulate against the cluster view rather than the stale per-rank sum.
|
||||
meter = tracker.update_s
|
||||
assert meter.sum / meter.count == pytest.approx(meter.avg)
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_registers_and_averages():
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
|
||||
tracker.update_metrics({"latent_loss": 0.2, "action_loss": 0.4})
|
||||
tracker.update_metrics({"latent_loss": 0.4, "action_loss": 0.6})
|
||||
|
||||
# New keys are auto-registered as mean-reduced meters and averaged over the window.
|
||||
assert tracker.metrics["latent_loss"].reduction == "mean"
|
||||
assert tracker.metrics["latent_loss"].avg == pytest.approx(0.3)
|
||||
assert tracker.metrics["action_loss"].avg == pytest.approx(0.5)
|
||||
assert tracker.to_dict()["latent_loss"] == pytest.approx(0.3)
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_skips_non_numeric():
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
|
||||
tracker.update_metrics({"loss": 0.5, "head_mode": "sparse", "enabled": True})
|
||||
|
||||
# strings and bools ignored
|
||||
assert "loss" in tracker.metrics
|
||||
assert "head_mode" not in tracker.metrics
|
||||
assert "enabled" not in tracker.metrics
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_does_not_override_caller_meter():
|
||||
# A policy that echoes "loss" in its output dict must not overwrite the caller-owned,
|
||||
# already-aggregated loss meter.
|
||||
metrics = {"loss": AverageMeter("loss", ":.3f", reduction="mean")}
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
|
||||
tracker.loss = 1.0 # caller-set optimized loss
|
||||
tracker.update_metrics({"loss": 99.0, "latent_loss": 0.2})
|
||||
|
||||
assert tracker.metrics["loss"].avg == pytest.approx(1.0) # snapshot ignored
|
||||
assert tracker.metrics["latent_loss"].avg == pytest.approx(0.2)
|
||||
|
||||
Reference in New Issue
Block a user