mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 11:16:00 +00:00
Pepijn
5912af7f9c
- RoboMME env wrapper with image/wrist_image/state observations - Docker image with Vulkan, SAPIEN, mani-skill deps - CI workflow: 1-episode smoke eval with pepijn223/smolvla_robomme - preprocess_observation: handle image/wrist_image/state keys - pyproject.toml: robomme extra Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
122 lines
4.8 KiB
Python
122 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2025 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.
|
|
|
|
"""Extract natural-language task descriptions for a benchmark suite.
|
|
|
|
Runs inside the benchmark Docker container (where the env library is installed)
|
|
immediately after lerobot-eval, writing a JSON file that parse_eval_metrics.py
|
|
picks up and embeds in metrics.json.
|
|
|
|
Output format: {"<suite>_<task_idx>": "<nl instruction>", ...}
|
|
|
|
Usage:
|
|
python scripts/ci/extract_task_descriptions.py \\
|
|
--env libero --task libero_spatial \\
|
|
--output /tmp/eval-artifacts/task_descriptions.json
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def _libero_descriptions(task_suite: str) -> dict[str, str]:
|
|
from libero.libero import benchmark # type: ignore[import-untyped]
|
|
|
|
suite_dict = benchmark.get_benchmark_dict()
|
|
if task_suite not in suite_dict:
|
|
print(
|
|
f"[extract_task_descriptions] Unknown LIBERO suite '{task_suite}'. "
|
|
f"Available: {list(suite_dict.keys())}",
|
|
file=sys.stderr,
|
|
)
|
|
return {}
|
|
suite = suite_dict[task_suite]()
|
|
return {f"{task_suite}_{i}": suite.get_task(i).language for i in range(suite.n_tasks)}
|
|
|
|
|
|
def _metaworld_descriptions(task_name: str) -> dict[str, str]:
|
|
# MetaWorld tasks don't expose a separate NL description attribute;
|
|
# use a cleaned version of the task name as the description.
|
|
label = task_name.removeprefix("metaworld-").replace("-", " ").strip()
|
|
return {f"{task_name}_0": label}
|
|
|
|
|
|
_ROBOMME_DESCRIPTIONS = {
|
|
"BinFill": "Fill the target bin with the correct number of cubes",
|
|
"PickXtimes": "Pick the indicated cube the specified number of times",
|
|
"SwingXtimes": "Swing the object the specified number of times",
|
|
"StopCube": "Grasp and stop the moving cube",
|
|
"VideoUnmask": "Pick the cube shown in the reference video",
|
|
"VideoUnmaskSwap": "Pick the cube matching the reference video after a swap",
|
|
"ButtonUnmask": "Press the button indicated by the reference",
|
|
"ButtonUnmaskSwap": "Press the correct button after objects are swapped",
|
|
"PickHighlight": "Pick the highlighted cube",
|
|
"VideoRepick": "Repick the cube shown in the reference video",
|
|
"VideoPlaceButton": "Place the cube on the button shown in the video",
|
|
"VideoPlaceOrder": "Place cubes in the order shown in the video",
|
|
"MoveCube": "Move the cube to the target location",
|
|
"InsertPeg": "Insert the peg into the target hole",
|
|
"PatternLock": "Unlock the pattern by pressing buttons in sequence",
|
|
"RouteStick": "Route the stick through the required waypoints",
|
|
}
|
|
|
|
|
|
def _robomme_descriptions(task_names: str) -> dict[str, str]:
|
|
"""Return descriptions for each requested RoboMME task. Keys match the
|
|
video filename pattern `<task>_<task_id>` used by the eval script."""
|
|
out: dict[str, str] = {}
|
|
for name in (t.strip() for t in task_names.split(",") if t.strip()):
|
|
desc = _ROBOMME_DESCRIPTIONS.get(name, name)
|
|
out[f"{name}_0"] = desc
|
|
return out
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)")
|
|
parser.add_argument("--task", required=True, help="Task/suite name (e.g. libero_spatial)")
|
|
parser.add_argument("--output", required=True, help="Path to write task_descriptions.json")
|
|
args = parser.parse_args()
|
|
|
|
descriptions: dict[str, str] = {}
|
|
try:
|
|
if args.env == "libero":
|
|
descriptions = _libero_descriptions(args.task)
|
|
elif args.env == "metaworld":
|
|
descriptions = _metaworld_descriptions(args.task)
|
|
elif args.env == "robomme":
|
|
descriptions = _robomme_descriptions(args.task)
|
|
else:
|
|
print(
|
|
f"[extract_task_descriptions] No description extractor for env '{args.env}'.",
|
|
file=sys.stderr,
|
|
)
|
|
except Exception as exc:
|
|
print(f"[extract_task_descriptions] Warning: {exc}", file=sys.stderr)
|
|
|
|
out_path = Path(args.output)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
out_path.write_text(json.dumps(descriptions, indent=2))
|
|
print(f"[extract_task_descriptions] {len(descriptions)} descriptions → {out_path}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|