mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
Add community_dataset_v3 -> v3.0 SO-arm migration scripts
This commit is contained in:
CarolinePascal
@@ -0,0 +1,84 @@
|
||||
"""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
|
||||
|
||||
SO_PREFIXES = ("so100", "so101") # so100, so101, so100_follower, so101_follower, so100_bimanual, ...
|
||||
SO_EXACT = {"so_follower"}
|
||||
# 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 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 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}
|
||||
|
||||
is_so = (rt.startswith(SO_PREFIXES) or rt in SO_EXACT) and rt not in NEVER_FIX
|
||||
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"}
|
||||
|
||||
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 {**out, "is_so": True, "encoding": enc, "maxabs": round(maxabs, 2),
|
||||
"saturates": sat, "stats_key": key_used, "ambiguous": ambiguous}
|
||||
@@ -0,0 +1,80 @@
|
||||
"""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
|
||||
|
||||
from classify import classify
|
||||
from so_arm_frame import to_degrees
|
||||
|
||||
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 _rewrite_parquet(root: Path, encoding: str) -> None:
|
||||
for pq in sorted((root / "data").glob("*/*.parquet")):
|
||||
df = pd.read_parquet(pq)
|
||||
changed = False
|
||||
for col in VALUE_COLS:
|
||||
if col in df.columns:
|
||||
conv = to_degrees(_stack(df[col].values), encoding, n_joints_per_arm=6)
|
||||
df[col] = list(conv.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 fix_dataset_in_place(root) -> dict:
|
||||
"""Returns the classification dict augmented with the action taken."""
|
||||
root = Path(root)
|
||||
cls = classify(root)
|
||||
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"}
|
||||
# drop stray files that would otherwise be uploaded
|
||||
for junk in (root / "meta").glob("info.json.bak"):
|
||||
junk.unlink()
|
||||
_rewrite_parquet(root, enc)
|
||||
_regen_episode_stats(root)
|
||||
return {**cls, "converted": True,
|
||||
"action": f"structural v2.1->v3.0 + joint values converted ({enc} -> degrees)"}
|
||||
@@ -0,0 +1,250 @@
|
||||
"""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,
|
||||
--allow-uncalibrated (accept placeholder CANON ranges).
|
||||
"""
|
||||
import argparse, csv, json, shutil, sys, traceback
|
||||
from pathlib import Path
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
|
||||
import so_arm_frame
|
||||
from classify import classify, load_info
|
||||
from fix_dataset import fix_dataset_in_place
|
||||
|
||||
SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
|
||||
|
||||
|
||||
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) -> 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."""
|
||||
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 canonical joint ranges",
|
||||
"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: [`{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:** normalized (-100..100 / 0..100) joint values were un-normalized "
|
||||
"using *placeholder* canonical joint ranges (uncalibrated). These values are 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
|
||||
card = create_lerobot_dataset_card(
|
||||
tags=[rt] if rt else 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) -> dict:
|
||||
local = Path(work_dir) / sub
|
||||
if local.parent.exists():
|
||||
shutil.rmtree(local.parent, ignore_errors=True) # clean any partial
|
||||
snapshot_download(SRC_REPO, repo_type="dataset", revision="main",
|
||||
allow_patterns=[f"{sub}/*"], local_dir=work_dir)
|
||||
|
||||
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)"}
|
||||
|
||||
result = fix_dataset_in_place(local) # SO-arm value fix (or structural_only)
|
||||
|
||||
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) # document the conversion in the dataset 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.")
|
||||
ap.add_argument("--allow-uncalibrated", action="store_true",
|
||||
help="Permit converting 'normalized' (-100..100) datasets using the PLACEHOLDER "
|
||||
"canonical joint ranges in so_arm_frame.py. Omit this to force running "
|
||||
"calibrate_canonical_ranges.py first (recommended for a real run).")
|
||||
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.")
|
||||
|
||||
if args.allow_uncalibrated:
|
||||
so_arm_frame.CANON_IS_CALIBRATED = True
|
||||
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
|
||||
snapshot_download(SRC_REPO, repo_type="dataset",
|
||||
allow_patterns=[f"{sub}/meta/*"], local_dir=args.work_dir)
|
||||
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()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""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):
|
||||
already mid-range-zero; only the SCALE is missing (per-robot range_min/max
|
||||
is not stored) -> use assumed canonical per-joint spans below. APPROXIMATE.
|
||||
* 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 / 0..100
|
||||
# normalization when per-robot calibration is unavailable. joints 0..4 (RANGE_M100_100):
|
||||
# normalized +/-100 -> +/-HALF_RANGE. gripper (RANGE_0_100): 0..100 -> centered at 50,
|
||||
# span FULL_RANGE. 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_GRIPPER_FULL_RANGE_DEG = 90.0
|
||||
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.empty_like(x)
|
||||
new_deg[..., :5] = (x[..., :5] / 100.0) * CANON_HALF_RANGE_DEG
|
||||
new_deg[..., 5] = (x[..., 5] / 100.0 - 0.5) * CANON_GRIPPER_FULL_RANGE_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, or pass --allow-uncalibrated to accept them."
|
||||
)
|
||||
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)
|
||||
Reference in New Issue
Block a user