mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-11 20:11:48 +00:00
Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Caroline Pascal
|
84b605d82c | ||
|
CarolinePascal
|
e36b0368d4 | ||
|
CarolinePascal
|
67b18d87b2 | ||
|
Mahbod
|
98052e5f6e | ||
|
Mahbod
|
f59260f4aa | ||
|
Mahbod
|
fc262fbc06 |
@@ -55,7 +55,7 @@ jobs:
|
||||
github.repository == 'huggingface/lerobot'
|
||||
permissions:
|
||||
contents: read
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
with:
|
||||
commit_sha: ${{ github.sha }}
|
||||
package: lerobot
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
with:
|
||||
commit_sha: ${{ github.event.pull_request.head.sha }}
|
||||
pr_number: ${{ github.event.number }}
|
||||
|
||||
@@ -162,11 +162,11 @@ Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50`
|
||||
|
||||
| Suite | Success rate | Checkpoint |
|
||||
| ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
|
||||
| LIBERO Spatial | 95% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 100% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 98% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 93% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **96.5%** | |
|
||||
| LIBERO Spatial | 91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **88.25%** | |
|
||||
|
||||
```bash
|
||||
export MODEL_ID=your_trained_model_on_huggingface
|
||||
|
||||
@@ -1,489 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# 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.
|
||||
|
||||
"""
|
||||
SLURM-distributed recomputation of a LeRobotDataset's ``meta/stats.json``.
|
||||
|
||||
Modified copy of lerobot's examples/dataset/slurm_recompute_stats.py
|
||||
(feat/recompute-stats-readonly-and-visual branch) with cluster-friendly additions:
|
||||
|
||||
1. --qos : pass a SLURM QoS through to every worker's sbatch.
|
||||
2. --venv-path : activate a venv on each worker before the python step.
|
||||
3. --env-command : raw shell snippet injected before the python step (e.g. to
|
||||
export HF_LEROBOT_HOME). Runs in addition to --venv-path.
|
||||
4. --chain-aggregate : submit ``aggregate`` with an afterok dependency on
|
||||
``compute`` so it only runs once all shards exist
|
||||
(no manual squeue-wait, no gap/overlap race).
|
||||
5. --update-episode-stats : in ``aggregate``, also rewrite the per-episode stats in the
|
||||
episodes parquet so they stay consistent with meta/stats.json
|
||||
(default: only stats.json is written).
|
||||
|
||||
Data access: no filesystem mount. Point HF_LEROBOT_HOME at a node-visible shared
|
||||
cache (e.g. /fsx/$USER/.cache) so the dataset downloads once and all workers read
|
||||
it. This is the download route; the source dataset is fetched from the Hub on the
|
||||
CPU workers.
|
||||
|
||||
IMPORTANT — how to run (do NOT sbatch this file):
|
||||
Run it as a normal python process on the LOGIN node. datatrove submits the
|
||||
workers for you. The reference copy (--new-root) is built on the login node and
|
||||
references the shared HF cache, so /fsx must be visible there (it is).
|
||||
|
||||
Requires: pip install 'lerobot[dataset]' datatrove
|
||||
|
||||
Example (single command, compute then dependent aggregate):
|
||||
|
||||
export HF_LEROBOT_HOME=/fsx/$USER/.cache
|
||||
|
||||
python slurm_recompute_stats_patched.py compute \
|
||||
--repo-id behavior-1k/2026-challenge-demos \
|
||||
--new-root /fsx/$USER/behavior-1k_recomputed \
|
||||
--shard-dir /fsx/$USER/behavior-1k_recomputed/stats_shards \
|
||||
--logs-dir /fsx/$USER/logs/recompute \
|
||||
--skip-image-video 0 \
|
||||
--workers 250 \
|
||||
--partition hopper-cpu \
|
||||
--qos normal \
|
||||
--cpus-per-task 8 --mem-per-cpu 4G \
|
||||
--venv-path /fsx/$USER/venvs/lerobot/bin/activate \
|
||||
--env-command 'export HF_LEROBOT_HOME=/fsx/'"$USER"'/.cache' \
|
||||
--chain-aggregate
|
||||
|
||||
REHEARSE FIRST with --workers 2 --skip-image-video 1 and inspect one worker's log
|
||||
under --logs-dir to confirm QoS was accepted and a numeric stats.json is written.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from datatrove.executor import LocalPipelineExecutor
|
||||
from datatrove.executor.slurm import SlurmPipelineExecutor
|
||||
from datatrove.pipeline.base import PipelineStep
|
||||
|
||||
class ComputeEpisodeStatsShards(PipelineStep):
|
||||
"""Each worker computes per-episode stats for its ``episodes[rank::world_size]`` shard."""
|
||||
|
||||
def __init__(self, repo_id, root, new_root, skip_image_video, shard_dir, video_backend=None):
|
||||
super().__init__()
|
||||
self.repo_id = repo_id
|
||||
self.root = root
|
||||
self.new_root = new_root
|
||||
self.skip_image_video = skip_image_video
|
||||
self.shard_dir = shard_dir
|
||||
self.video_backend = video_backend
|
||||
|
||||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
||||
# NOTE: this method is pickled and executed on a worker, where this script's module
|
||||
# globals are NOT available. Keep it self-contained: import locally and don't reference
|
||||
# module-level helpers/constants.
|
||||
import logging
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
from lerobot.datasets import LeRobotDataset, compute_dataset_episode_stats
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
init_logging()
|
||||
load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {}
|
||||
root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
|
||||
dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
|
||||
|
||||
my_episodes = list(range(dataset.meta.total_episodes))[rank::world_size]
|
||||
if not my_episodes:
|
||||
logging.info(f"Rank {rank}: no episodes assigned")
|
||||
return
|
||||
logging.info(f"Rank {rank}: {len(my_episodes)} / {dataset.meta.total_episodes} episodes")
|
||||
|
||||
episode_stats = compute_dataset_episode_stats(
|
||||
dataset,
|
||||
episode_indices=my_episodes,
|
||||
skip_image_video=self.skip_image_video,
|
||||
)
|
||||
|
||||
shard_dir = Path(self.shard_dir)
|
||||
shard_dir.mkdir(parents=True, exist_ok=True)
|
||||
out = shard_dir / f"episode_stats_{rank:05d}.pkl"
|
||||
with open(out, "wb") as f:
|
||||
pickle.dump(episode_stats, f)
|
||||
logging.info(f"Rank {rank}: saved {len(episode_stats)} episode stats to {out}")
|
||||
|
||||
|
||||
class AggregateEpisodeStats(PipelineStep):
|
||||
"""Merge all per-episode stat shards into meta/stats.json."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
repo_id,
|
||||
root,
|
||||
new_root,
|
||||
shard_dir,
|
||||
push_to_hub=False,
|
||||
video_backend=None,
|
||||
update_episode_stats=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.repo_id = repo_id
|
||||
self.root = root
|
||||
self.new_root = new_root
|
||||
self.shard_dir = shard_dir
|
||||
self.push_to_hub = push_to_hub
|
||||
self.video_backend = video_backend
|
||||
self.update_episode_stats = update_episode_stats
|
||||
|
||||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
||||
# NOTE: pickled and executed on a worker; keep self-contained (see ComputeEpisodeStatsShards.run).
|
||||
import logging
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
from lerobot.datasets import LeRobotDataset, aggregate_episode_stats
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
init_logging()
|
||||
if rank != 0:
|
||||
return
|
||||
|
||||
shard_dir = Path(self.shard_dir)
|
||||
shards = sorted(shard_dir.glob("episode_stats_*.pkl"))
|
||||
if not shards:
|
||||
raise FileNotFoundError(f"No episode stat shards found in {shard_dir}")
|
||||
|
||||
# Shards map episode_index -> stats; merging by key makes a dropped shard show up as a
|
||||
# missing episode and a re-run shard overwrite rather than double-count.
|
||||
all_episode_stats = {}
|
||||
for shard in shards:
|
||||
with open(shard, "rb") as f:
|
||||
all_episode_stats.update(pickle.load(f))
|
||||
logging.info(f"Aggregating {len(all_episode_stats)} episode stats from {len(shards)} shards")
|
||||
|
||||
load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {}
|
||||
root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
|
||||
dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
|
||||
|
||||
# Aggregation is order-independent, so the only way sharding changes the result is a
|
||||
# gap (dropped shard) or an overlap (episode counted twice). Verify the shards cover
|
||||
# every episode exactly once before writing stats.json.
|
||||
expected_episodes = dataset.meta.total_episodes
|
||||
if len(all_episode_stats) != expected_episodes:
|
||||
raise ValueError(
|
||||
f"Expected {expected_episodes} per-episode stats (one per episode) but got "
|
||||
f"{len(all_episode_stats)} across {len(shards)} shards. A compute shard is likely "
|
||||
"missing or was written more than once; re-run the failed shards before aggregating."
|
||||
)
|
||||
|
||||
# Frame-count check catches the case where a duplicate and a gap cancel out in the
|
||||
# episode count: summed per-episode frame counts must equal the dataset's total frames.
|
||||
stats_values = list(all_episode_stats.values())
|
||||
numeric_key = next(
|
||||
(
|
||||
k
|
||||
for k, v in dataset.meta.features.items()
|
||||
if v["dtype"] not in ("image", "video", "string") and stats_values and k in stats_values[0]
|
||||
),
|
||||
None,
|
||||
)
|
||||
if numeric_key is not None:
|
||||
total_frames = sum(int(s[numeric_key]["count"][0]) for s in stats_values)
|
||||
if total_frames != dataset.meta.total_frames:
|
||||
raise ValueError(
|
||||
f"Summed frame count from shards ({total_frames}) != dataset total_frames "
|
||||
f"({dataset.meta.total_frames}); episodes are double-counted or missing."
|
||||
)
|
||||
|
||||
new_stats = aggregate_episode_stats(
|
||||
dataset, all_episode_stats, update_episode_stats=self.update_episode_stats
|
||||
)
|
||||
if new_stats is None:
|
||||
raise RuntimeError("Aggregation produced no stats")
|
||||
logging.info(f"Wrote stats for features: {list(new_stats.keys())} to {dataset.root}")
|
||||
|
||||
if self.push_to_hub:
|
||||
logging.info(f"Pushing {self.repo_id} to hub")
|
||||
dataset.push_to_hub()
|
||||
|
||||
|
||||
def _mem_gb(mem: str) -> int:
|
||||
"""Parse '4G' / '4GB' / '4' into an int number of GB for datatrove's mem_per_cpu_gb."""
|
||||
s = str(mem).strip().lower().rstrip("b").rstrip("g")
|
||||
return int(float(s))
|
||||
|
||||
|
||||
def _make_executor(
|
||||
pipeline,
|
||||
logs_dir,
|
||||
job_name,
|
||||
slurm,
|
||||
workers,
|
||||
tasks,
|
||||
time,
|
||||
partition,
|
||||
cpus,
|
||||
mem,
|
||||
qos=None,
|
||||
env_command=None,
|
||||
venv_path=None,
|
||||
depends=None,
|
||||
):
|
||||
kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)}
|
||||
if slurm:
|
||||
kwargs.update(
|
||||
{
|
||||
"job_name": job_name,
|
||||
"tasks": tasks,
|
||||
"workers": workers,
|
||||
"time": time,
|
||||
"partition": partition,
|
||||
"cpus_per_task": cpus,
|
||||
"mem_per_cpu_gb": _mem_gb(mem), # datatrove's native field (int GB)
|
||||
"sbatch_args": {},
|
||||
}
|
||||
)
|
||||
if qos:
|
||||
kwargs["qos"] = qos # -> "#SBATCH --qos=<qos>" on every worker
|
||||
if venv_path:
|
||||
kwargs["venv_path"] = venv_path # datatrove sources this before the python step
|
||||
if env_command:
|
||||
kwargs["env_command"] = env_command # extra raw snippet before python (composes with venv_path)
|
||||
if depends is not None:
|
||||
kwargs["depends"] = depends # chains --dependency=afterok:<compute jobid>
|
||||
return SlurmPipelineExecutor(**kwargs)
|
||||
kwargs.update({"tasks": tasks, "workers": 1})
|
||||
return LocalPipelineExecutor(**kwargs)
|
||||
|
||||
|
||||
def _maybe_reference_copy(repo_id, root, new_root, download_videos):
|
||||
"""Create the read-only-safe reference copy once, before submitting workers.
|
||||
|
||||
Loads metadata only (to resolve the source root and revision) instead of a full
|
||||
``LeRobotDataset``, which would also memory-map the entire frame index just to read a
|
||||
path. Fetches the source into the shared cache so the copy's symlinks point at real
|
||||
files and workers don't each re-download, pulling videos only when the run needs them
|
||||
(i.e. when image/video stats are being recomputed).
|
||||
"""
|
||||
if not new_root:
|
||||
return
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.scripts.lerobot_edit_dataset import _reference_copy_dataset
|
||||
from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
|
||||
|
||||
new_root_path = Path(new_root)
|
||||
if new_root_path.exists():
|
||||
return
|
||||
|
||||
meta = LeRobotDatasetMetadata(repo_id, root=Path(root) if root else None)
|
||||
ignore_patterns = None if download_videos else "videos/"
|
||||
if root:
|
||||
snapshot_download(
|
||||
repo_id,
|
||||
repo_type="dataset",
|
||||
revision=meta.revision,
|
||||
local_dir=meta.root,
|
||||
ignore_patterns=ignore_patterns,
|
||||
)
|
||||
src_root = Path(meta.root)
|
||||
else:
|
||||
src_root = Path(
|
||||
snapshot_download(
|
||||
repo_id,
|
||||
repo_type="dataset",
|
||||
revision=meta.revision,
|
||||
cache_dir=HF_LEROBOT_HUB_CACHE,
|
||||
ignore_patterns=ignore_patterns,
|
||||
)
|
||||
)
|
||||
_reference_copy_dataset(src_root, new_root_path)
|
||||
|
||||
|
||||
def _add_shared_args(p):
|
||||
p.add_argument("--repo-id", type=str, required=True, help="Dataset identifier, e.g. 'user/dataset'.")
|
||||
p.add_argument("--root", type=str, default=None, help="Source dataset root (defaults to the Hub cache).")
|
||||
p.add_argument(
|
||||
"--new-root",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Writable output root; a read-only-safe reference copy of --root. If omitted, stats "
|
||||
"are written in place at --root.",
|
||||
)
|
||||
p.add_argument("--shard-dir", type=Path, default=Path("stats_shards"), help="Per-rank shard dir.")
|
||||
p.add_argument("--logs-dir", type=Path, default=Path("logs"), help="datatrove logs dir.")
|
||||
p.add_argument("--job-name", type=str, default=None, help="SLURM job name.")
|
||||
p.add_argument("--slurm", type=int, default=1, help="1 = submit via SLURM; 0 = run locally.")
|
||||
p.add_argument("--partition", type=str, default=None, help="SLURM partition, e.g. 'hopper-cpu'.")
|
||||
p.add_argument("--qos", type=str, default=None, help="SLURM QoS, e.g. 'normal'. Passed to every worker.")
|
||||
p.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per SLURM task.")
|
||||
p.add_argument("--mem-per-cpu", type=str, default="4G", help="Memory per CPU, e.g. '4G'.")
|
||||
p.add_argument(
|
||||
"--video-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Video decoding backend (e.g. 'pyav', 'torchcodec'). Defaults to the dataset's default; "
|
||||
"use 'pyav' if torchcodec fails to load locally.",
|
||||
)
|
||||
p.add_argument("--venv-path", type=str, default=None, help="venv activate script sourced on each worker.")
|
||||
p.add_argument(
|
||||
"--env-command",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Raw shell snippet injected into each worker's sbatch before the python step "
|
||||
"(e.g. to export HF_LEROBOT_HOME). Runs in addition to --venv-path.",
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="PATCHED SLURM-distributed LeRobotDataset stats recomputation",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
cp = sub.add_parser("compute", help="Distribute per-episode stats across SLURM workers.")
|
||||
_add_shared_args(cp)
|
||||
cp.add_argument("--workers", type=int, default=50, help="Number of parallel SLURM tasks.")
|
||||
cp.add_argument(
|
||||
"--skip-image-video",
|
||||
type=int,
|
||||
default=1,
|
||||
help="1 = numeric features only (fast); 0 = also recompute image/video stats (decodes frames).",
|
||||
)
|
||||
cp.add_argument(
|
||||
"--chain-aggregate",
|
||||
action="store_true",
|
||||
help="After building compute, submit aggregate with an afterok dependency (single command).",
|
||||
)
|
||||
cp.add_argument("--push-to-hub", action="store_true", help="For the chained aggregate: push after done.")
|
||||
cp.add_argument(
|
||||
"--update-episode-stats",
|
||||
action="store_true",
|
||||
help="For the chained aggregate: also rewrite per-episode stats in the episodes parquet.",
|
||||
)
|
||||
|
||||
ap = sub.add_parser("aggregate", help="Merge shards into meta/stats.json.")
|
||||
_add_shared_args(ap)
|
||||
ap.add_argument("--push-to-hub", action="store_true", help="Push the dataset after aggregation.")
|
||||
ap.add_argument(
|
||||
"--update-episode-stats",
|
||||
action="store_true",
|
||||
help="Also rewrite per-episode stats in the episodes parquet to match stats.json.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--depends-job-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional SLURM job id; aggregate waits for it (afterok) before running.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
slurm = args.slurm == 1
|
||||
|
||||
if args.command == "compute":
|
||||
# The reference copy (if any) is created once on the submitting node so workers
|
||||
# can all load --new-root without racing to build it. Videos are only fetched when
|
||||
# image/video stats are being recomputed.
|
||||
_maybe_reference_copy(
|
||||
args.repo_id, args.root, args.new_root, download_videos=not bool(args.skip_image_video)
|
||||
)
|
||||
|
||||
compute_exec = _make_executor(
|
||||
pipeline=[
|
||||
ComputeEpisodeStatsShards(
|
||||
args.repo_id,
|
||||
args.root,
|
||||
args.new_root,
|
||||
bool(args.skip_image_video),
|
||||
str(args.shard_dir),
|
||||
args.video_backend,
|
||||
)
|
||||
],
|
||||
logs_dir=args.logs_dir,
|
||||
job_name=args.job_name or "recompute_stats_compute",
|
||||
slurm=slurm,
|
||||
workers=args.workers,
|
||||
tasks=args.workers,
|
||||
time="24:00:00",
|
||||
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,
|
||||
)
|
||||
|
||||
if args.chain_aggregate and slurm:
|
||||
# Build aggregate depending on compute. datatrove launches the dependency
|
||||
# (compute) first, then submits aggregate with --dependency=afterok:<jobid>.
|
||||
aggregate_exec = _make_executor(
|
||||
pipeline=[
|
||||
AggregateEpisodeStats(
|
||||
args.repo_id,
|
||||
args.root,
|
||||
args.new_root,
|
||||
str(args.shard_dir),
|
||||
args.push_to_hub,
|
||||
args.video_backend,
|
||||
args.update_episode_stats,
|
||||
)
|
||||
],
|
||||
logs_dir=args.logs_dir,
|
||||
job_name="recompute_stats_aggregate",
|
||||
slurm=slurm,
|
||||
workers=1,
|
||||
tasks=1,
|
||||
time="02:00:00",
|
||||
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,
|
||||
depends=compute_exec,
|
||||
)
|
||||
aggregate_exec.run()
|
||||
else:
|
||||
compute_exec.run()
|
||||
else:
|
||||
aggregate_exec = _make_executor(
|
||||
pipeline=[
|
||||
AggregateEpisodeStats(
|
||||
args.repo_id,
|
||||
args.root,
|
||||
args.new_root,
|
||||
str(args.shard_dir),
|
||||
args.push_to_hub,
|
||||
args.video_backend,
|
||||
args.update_episode_stats,
|
||||
)
|
||||
],
|
||||
logs_dir=args.logs_dir,
|
||||
job_name=args.job_name or "recompute_stats_aggregate",
|
||||
slurm=slurm,
|
||||
workers=1,
|
||||
tasks=1,
|
||||
time="02:00:00",
|
||||
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,
|
||||
)
|
||||
if args.depends_job_id is not None:
|
||||
aggregate_exec.depends_job_id = args.depends_job_id
|
||||
aggregate_exec.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -25,8 +25,6 @@ from .compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats
|
||||
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
|
||||
from .dataset_tools import (
|
||||
add_features,
|
||||
aggregate_episode_stats,
|
||||
compute_dataset_episode_stats,
|
||||
convert_image_to_video_dataset,
|
||||
delete_episodes,
|
||||
merge_datasets,
|
||||
@@ -36,7 +34,6 @@ from .dataset_tools import (
|
||||
reencode_dataset,
|
||||
remove_feature,
|
||||
split_dataset,
|
||||
write_episode_stats,
|
||||
)
|
||||
from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps
|
||||
from .image_writer import safe_stop_image_writer
|
||||
@@ -81,10 +78,8 @@ __all__ = [
|
||||
"detect_available_encoders_pyav",
|
||||
"add_features",
|
||||
"aggregate_datasets",
|
||||
"aggregate_episode_stats",
|
||||
"aggregate_pipeline_dataset_features",
|
||||
"aggregate_stats",
|
||||
"compute_dataset_episode_stats",
|
||||
"convert_image_to_video_dataset",
|
||||
"create_initial_features",
|
||||
"compute_sampler_state",
|
||||
@@ -104,6 +99,5 @@ __all__ = [
|
||||
"resolve_delta_timestamps",
|
||||
"safe_stop_image_writer",
|
||||
"split_dataset",
|
||||
"write_episode_stats",
|
||||
"write_stats",
|
||||
]
|
||||
|
||||
@@ -519,6 +519,13 @@ def compute_episode_stats(
|
||||
if features[key]["dtype"] in {"string", "language"}:
|
||||
continue
|
||||
|
||||
# Features with zero-width shapes are skipped (no data to compute stats on)
|
||||
if any(d == 0 for d in features[key].get("shape", ())):
|
||||
logging.debug(
|
||||
f"Skipping statistics computation for feature '{key}' with a zero-width shape {features[key]['shape']}."
|
||||
)
|
||||
continue
|
||||
|
||||
if features[key]["dtype"] in ["image", "video"]:
|
||||
ep_ft_array = sample_images(data)
|
||||
axes_to_reduce = (0, 2, 3)
|
||||
|
||||
@@ -33,13 +33,11 @@ from pathlib import Path
|
||||
import datasets
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.configs import (
|
||||
DEFAULT_DEPTH_UNIT,
|
||||
DepthEncoderConfig,
|
||||
RGBEncoderConfig,
|
||||
VideoEncoderConfig,
|
||||
@@ -54,14 +52,10 @@ from lerobot.utils.utils import flatten_dict
|
||||
from .aggregate import aggregate_datasets
|
||||
from .compute_stats import (
|
||||
aggregate_stats,
|
||||
auto_downsample_height_width,
|
||||
compute_episode_stats,
|
||||
compute_relative_action_stats,
|
||||
get_feature_stats,
|
||||
sample_indices,
|
||||
)
|
||||
from .dataset_metadata import LeRobotDatasetMetadata
|
||||
from .depth_utils import dequantize_depth
|
||||
from .image_writer import write_image
|
||||
from .io_utils import (
|
||||
get_parquet_file_size_in_mb,
|
||||
@@ -83,7 +77,6 @@ from .utils import (
|
||||
update_chunk_file_indices,
|
||||
)
|
||||
from .video_utils import (
|
||||
decode_video_frames,
|
||||
encode_video_frames,
|
||||
reencode_video,
|
||||
)
|
||||
@@ -1566,173 +1559,6 @@ def modify_tasks(
|
||||
return dataset
|
||||
|
||||
|
||||
def _load_episode_image_frames(
|
||||
dataset: LeRobotDataset,
|
||||
key: str,
|
||||
ep_idx: int,
|
||||
frame_offsets: list[int],
|
||||
is_depth: bool,
|
||||
) -> np.ndarray:
|
||||
"""Load sampled frames of an image feature for one episode as a (N, C, H, W) array."""
|
||||
ep = dataset.meta.episodes[ep_idx]
|
||||
from_idx = ep["dataset_from_index"]
|
||||
column = dataset.hf_dataset.with_format(None).select_columns(key)
|
||||
|
||||
frames = []
|
||||
for offset in frame_offsets:
|
||||
img = column[from_idx + offset][key]
|
||||
if is_depth:
|
||||
arr = np.array(img)
|
||||
if arr.ndim == 2:
|
||||
arr = arr[np.newaxis, ...]
|
||||
else:
|
||||
arr = np.transpose(np.array(img.convert("RGB"), dtype=np.uint8), (2, 0, 1))
|
||||
frames.append(auto_downsample_height_width(arr))
|
||||
return np.stack(frames)
|
||||
|
||||
|
||||
def _load_episode_video_frames(
|
||||
dataset: LeRobotDataset,
|
||||
key: str,
|
||||
ep_idx: int,
|
||||
frame_offsets: list[int],
|
||||
is_depth: bool,
|
||||
) -> np.ndarray:
|
||||
"""Load sampled frames of a video feature for one episode as a (N, C, H, W) array."""
|
||||
ep = dataset.meta.episodes[ep_idx]
|
||||
video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, key)
|
||||
from_timestamp = ep[f"videos/{key}/from_timestamp"]
|
||||
timestamps = [from_timestamp + offset / dataset.meta.fps for offset in frame_offsets]
|
||||
|
||||
frames = decode_video_frames(
|
||||
video_path,
|
||||
timestamps,
|
||||
dataset.tolerance_s,
|
||||
backend=dataset._video_backend,
|
||||
return_uint8=not is_depth,
|
||||
is_depth=is_depth,
|
||||
)
|
||||
if is_depth:
|
||||
# ``decode_video_frames`` returns raw 12-bit codec values; dequantize back to
|
||||
# the recorded depth unit so stats match record-time stats (which are stored in
|
||||
# ``info.depth_unit`` and only rescaled to the output unit on read).
|
||||
info = dataset.meta.features[key].get("info") or {}
|
||||
depth_encoder = DepthEncoderConfig.from_video_info(info)
|
||||
frames = dequantize_depth(
|
||||
frames,
|
||||
depth_min=depth_encoder.depth_min,
|
||||
depth_max=depth_encoder.depth_max,
|
||||
shift=depth_encoder.shift,
|
||||
use_log=depth_encoder.use_log,
|
||||
output_unit=info.get("depth_unit") or DEFAULT_DEPTH_UNIT,
|
||||
)
|
||||
return np.stack([auto_downsample_height_width(frame) for frame in frames.numpy()])
|
||||
|
||||
|
||||
def _compute_visual_episode_stats(
|
||||
dataset: LeRobotDataset,
|
||||
ep_idx: int,
|
||||
visual_keys: list[str],
|
||||
) -> dict:
|
||||
"""Compute per-episode statistics for image/video features by sampling frames.
|
||||
|
||||
Mirrors the image/video branch of :func:`compute_episode_stats`: per-channel stats
|
||||
are computed on downsampled sampled frames, then RGB stats are rescaled to [0, 1]
|
||||
(depth maps keep their native units).
|
||||
"""
|
||||
ep_length = dataset.meta.episodes[ep_idx]["length"]
|
||||
frame_offsets = sample_indices(ep_length)
|
||||
|
||||
ep_stats = {}
|
||||
for key in visual_keys:
|
||||
is_depth = key in dataset.meta.depth_keys
|
||||
if dataset.meta.features[key]["dtype"] == "video":
|
||||
frames = _load_episode_video_frames(dataset, key, ep_idx, frame_offsets, is_depth)
|
||||
else:
|
||||
frames = _load_episode_image_frames(dataset, key, ep_idx, frame_offsets, is_depth)
|
||||
|
||||
stats = get_feature_stats(frames, axis=(0, 2, 3), keepdims=True)
|
||||
normalization_factor = 1.0 if is_depth else 255.0
|
||||
ep_stats[key] = {
|
||||
k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0)
|
||||
for k, v in stats.items()
|
||||
}
|
||||
|
||||
return ep_stats
|
||||
|
||||
|
||||
def compute_dataset_episode_stats(
|
||||
dataset: LeRobotDataset,
|
||||
episode_indices: list[int] | None = None,
|
||||
skip_image_video: bool = True,
|
||||
drop_keys: list[str] | None = None,
|
||||
) -> dict[int, dict]:
|
||||
"""Compute per-episode statistics for a subset of episodes.
|
||||
|
||||
This is the shardable unit of work behind :func:`recompute_stats`: distribute
|
||||
``episode_indices`` across workers (e.g. ``list(range(n))[rank::world_size]``),
|
||||
then combine the results with :func:`aggregate_episode_stats`.
|
||||
|
||||
Args:
|
||||
dataset: The LeRobotDataset to compute stats for.
|
||||
episode_indices: Episodes to process. When ``None``, all episodes are processed.
|
||||
skip_image_video: If True (default), only numeric features are computed. If False,
|
||||
image/video stats are also computed by sampling and decoding frames.
|
||||
drop_keys: Feature keys to exclude (e.g. ``action`` when it is computed separately
|
||||
in relative-action space).
|
||||
|
||||
Returns:
|
||||
A mapping of episode index to its per-episode stat dict. Keeping the episode index
|
||||
(rather than a bare list) lets callers write the stats back to the correct episode
|
||||
row, and survives sharding since shards can be merged by key.
|
||||
"""
|
||||
features = dataset.meta.features
|
||||
meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
|
||||
drop = set(drop_keys or [])
|
||||
features_to_compute = {
|
||||
k: v
|
||||
for k, v in features.items()
|
||||
if v["dtype"] != "string"
|
||||
and k not in meta_keys
|
||||
and k not in drop
|
||||
and (not skip_image_video or v["dtype"] not in ["image", "video"])
|
||||
}
|
||||
numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
|
||||
visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]]
|
||||
|
||||
if dataset.meta.episodes is None:
|
||||
dataset.meta.episodes = load_episodes(dataset.meta.root)
|
||||
|
||||
if episode_indices is None:
|
||||
episode_indices = list(range(dataset.meta.total_episodes))
|
||||
|
||||
# Group requested episodes by their data parquet file so each file is read once.
|
||||
file_to_episodes: dict[Path, list[int]] = {}
|
||||
for ep_idx in episode_indices:
|
||||
file_to_episodes.setdefault(dataset.meta.get_data_file_path(ep_idx), []).append(ep_idx)
|
||||
|
||||
all_episode_stats = {}
|
||||
for src_path, eps in tqdm(sorted(file_to_episodes.items()), desc="Computing stats from data files"):
|
||||
df = pd.read_parquet(dataset.root / src_path) if numeric_keys else None
|
||||
for ep_idx in sorted(eps):
|
||||
episode_data = {}
|
||||
if numeric_keys:
|
||||
ep_df = df[df["episode_index"] == ep_idx]
|
||||
for key in numeric_keys:
|
||||
if key in ep_df.columns:
|
||||
values = ep_df[key].values
|
||||
episode_data[key] = (
|
||||
np.stack(values) if hasattr(values[0], "__len__") else np.array(values)
|
||||
)
|
||||
|
||||
ep_stats = compute_episode_stats(episode_data, features_to_compute)
|
||||
if visual_keys:
|
||||
ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys))
|
||||
all_episode_stats[int(ep_idx)] = ep_stats
|
||||
|
||||
return all_episode_stats
|
||||
|
||||
|
||||
def recompute_stats(
|
||||
dataset: LeRobotDataset,
|
||||
skip_image_video: bool = True,
|
||||
@@ -1740,21 +1566,13 @@ def recompute_stats(
|
||||
relative_exclude_joints: list[str] | None = None,
|
||||
chunk_size: int = 50,
|
||||
num_workers: int = 0,
|
||||
update_episode_stats: bool = False,
|
||||
) -> LeRobotDataset:
|
||||
"""Recompute stats.json from scratch by iterating all episodes.
|
||||
|
||||
Args:
|
||||
dataset: The LeRobotDataset to recompute stats for.
|
||||
skip_image_video: If True (default), only recompute stats for numeric features
|
||||
(action, state, etc.) and keep existing image/video stats unchanged. If False,
|
||||
image/video stats are also recomputed by sampling and decoding frames from each
|
||||
episode (this reads the image/video files, unlike the numeric-only path).
|
||||
update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the
|
||||
episodes parquet files so they stay consistent with the aggregated ``stats.json``.
|
||||
Defaults to False (only ``stats.json`` is rewritten). Requires a writable
|
||||
``dataset.root``. Note that relative-action stats are aggregate-only and are not
|
||||
written per-episode.
|
||||
(action, state, etc.) and keep existing image/video stats unchanged.
|
||||
relative_action: If True, compute action stats in relative space by
|
||||
iterating all valid action chunks and subtracting the current state.
|
||||
This matches the normalization distribution the model sees during
|
||||
@@ -1770,12 +1588,24 @@ def recompute_stats(
|
||||
The same dataset with updated stats.
|
||||
"""
|
||||
features = dataset.meta.features
|
||||
meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
|
||||
numeric_features = {
|
||||
k: v
|
||||
for k, v in features.items()
|
||||
if v["dtype"] not in ["image", "video", "string"] and k not in meta_keys
|
||||
}
|
||||
|
||||
if skip_image_video:
|
||||
features_to_compute = numeric_features
|
||||
else:
|
||||
features_to_compute = {
|
||||
k: v for k, v in features.items() if v["dtype"] != "string" and k not in meta_keys
|
||||
}
|
||||
|
||||
# When relative_action is enabled, compute action stats via chunk-based sampling
|
||||
# (matching what the model sees during training) and skip action in the
|
||||
# per-episode pass below.
|
||||
relative_action_stats = None
|
||||
drop_keys = None
|
||||
if relative_action and ACTION in features and OBS_STATE in features:
|
||||
if relative_exclude_joints is None:
|
||||
relative_exclude_joints = ["gripper"]
|
||||
@@ -1786,105 +1616,56 @@ def recompute_stats(
|
||||
exclude_joints=relative_exclude_joints,
|
||||
num_workers=num_workers,
|
||||
)
|
||||
drop_keys = [ACTION]
|
||||
features_to_compute.pop(ACTION, None)
|
||||
|
||||
all_episode_stats = compute_dataset_episode_stats(
|
||||
dataset, skip_image_video=skip_image_video, drop_keys=drop_keys
|
||||
)
|
||||
logging.info(f"Recomputing stats for features: {list(features_to_compute.keys())}")
|
||||
|
||||
new_stats = aggregate_episode_stats(
|
||||
dataset,
|
||||
all_episode_stats,
|
||||
extra_stats={ACTION: relative_action_stats} if relative_action_stats else None,
|
||||
update_episode_stats=update_episode_stats,
|
||||
)
|
||||
if new_stats is None:
|
||||
data_dir = dataset.root / DATA_DIR
|
||||
parquet_files = sorted(data_dir.glob("*/*.parquet"))
|
||||
if not parquet_files:
|
||||
raise ValueError(f"No parquet files found in {data_dir}")
|
||||
|
||||
all_episode_stats = []
|
||||
# TODO: enable image and video stats re-computation
|
||||
numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
|
||||
|
||||
for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"):
|
||||
df = pd.read_parquet(parquet_path)
|
||||
|
||||
for ep_idx in sorted(df["episode_index"].unique()):
|
||||
ep_df = df[df["episode_index"] == ep_idx]
|
||||
episode_data = {}
|
||||
for key in numeric_keys:
|
||||
if key in ep_df.columns:
|
||||
values = ep_df[key].values
|
||||
if hasattr(values[0], "__len__"):
|
||||
episode_data[key] = np.stack(values)
|
||||
else:
|
||||
episode_data[key] = np.array(values)
|
||||
|
||||
ep_stats = compute_episode_stats(episode_data, features_to_compute)
|
||||
all_episode_stats.append(ep_stats)
|
||||
|
||||
if features_to_compute and not all_episode_stats:
|
||||
logging.warning("No episode stats computed")
|
||||
else:
|
||||
logging.info("Stats recomputed successfully")
|
||||
return dataset
|
||||
return dataset
|
||||
|
||||
new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {}
|
||||
|
||||
def write_episode_stats(dataset: LeRobotDataset, episode_stats: dict[int, dict]) -> None:
|
||||
"""Overwrite the per-episode ``stats/*`` columns in the episodes parquet files in place.
|
||||
if relative_action_stats is not None:
|
||||
new_stats[ACTION] = relative_action_stats
|
||||
|
||||
Only the features present in ``episode_stats[ep_idx]`` are rewritten; stats columns for
|
||||
features that were not recomputed are left untouched. Every other episode column (tasks,
|
||||
length, chunk/file indices, frame ranges, …) is preserved. ``dataset.root`` must be
|
||||
writable (e.g. the reference copy created for read-only sources).
|
||||
"""
|
||||
if not episode_stats:
|
||||
return
|
||||
|
||||
meta = dataset.meta
|
||||
if meta.episodes is None:
|
||||
meta.episodes = load_episodes(meta.root)
|
||||
|
||||
# Group episodes by the parquet file that holds them so each file is rewritten once.
|
||||
file_to_episodes: dict[tuple[int, int], list[int]] = {}
|
||||
for ep_idx in episode_stats:
|
||||
ep = meta.episodes[ep_idx]
|
||||
key = (ep["meta/episodes/chunk_index"], ep["meta/episodes/file_index"])
|
||||
file_to_episodes.setdefault(key, []).append(ep_idx)
|
||||
|
||||
for (chunk_idx, file_idx), eps in file_to_episodes.items():
|
||||
path = meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
||||
table = pq.read_table(path)
|
||||
rows = table.to_pylist()
|
||||
row_by_ep = {row["episode_index"]: row for row in rows}
|
||||
for ep_idx in eps:
|
||||
row = row_by_ep[ep_idx]
|
||||
for feature, feature_stats in episode_stats[ep_idx].items():
|
||||
for stat_name, value in feature_stats.items():
|
||||
col = f"stats/{feature}/{stat_name}"
|
||||
if col in row:
|
||||
row[col] = np.asarray(value).tolist()
|
||||
# Reuse the source schema so the rewritten stats keep the exact on-disk types.
|
||||
new_table = pa.Table.from_pylist(rows, schema=table.schema)
|
||||
pq.write_table(new_table, path, compression="snappy", use_dictionary=True)
|
||||
|
||||
|
||||
def aggregate_episode_stats(
|
||||
dataset: LeRobotDataset,
|
||||
episode_stats: dict[int, dict],
|
||||
extra_stats: dict | None = None,
|
||||
update_episode_stats: bool = False,
|
||||
) -> dict | None:
|
||||
"""Aggregate per-episode stats, merge with existing stats, and write ``stats.json``.
|
||||
|
||||
Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass the
|
||||
merged ``{episode_index: stats}`` mapping of every worker's per-episode stats. ``extra_stats``
|
||||
lets callers inject feature stats computed outside the per-episode pass (e.g. relative-action
|
||||
stats).
|
||||
|
||||
Args:
|
||||
dataset: The dataset whose ``meta/stats.json`` (and optionally episode stats) is updated.
|
||||
episode_stats: Mapping of episode index to its per-episode stat dict.
|
||||
extra_stats: Feature stats to inject into the aggregate (not written per-episode).
|
||||
update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the
|
||||
episodes parquet files via :func:`write_episode_stats`.
|
||||
|
||||
Returns the written stats dict, or ``None`` if there was nothing to aggregate.
|
||||
"""
|
||||
if not episode_stats and not extra_stats:
|
||||
return None
|
||||
|
||||
new_stats = aggregate_stats(list(episode_stats.values())) if episode_stats else {}
|
||||
if extra_stats:
|
||||
new_stats.update(extra_stats)
|
||||
|
||||
# Merge: keep existing stats for features we didn't recompute.
|
||||
# Merge: keep existing stats for features we didn't recompute
|
||||
if dataset.meta.stats:
|
||||
for key, value in dataset.meta.stats.items():
|
||||
new_stats.setdefault(key, value)
|
||||
if key not in new_stats:
|
||||
new_stats[key] = value
|
||||
|
||||
write_stats(new_stats, dataset.root)
|
||||
dataset.meta.stats = new_stats
|
||||
|
||||
if update_episode_stats:
|
||||
write_episode_stats(dataset, episode_stats)
|
||||
|
||||
return new_stats
|
||||
logging.info("Stats recomputed successfully")
|
||||
return dataset
|
||||
|
||||
|
||||
def convert_image_to_video_dataset(
|
||||
|
||||
@@ -67,9 +67,9 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
|
||||
elif ft["shape"] == (1,):
|
||||
hf_features[key] = datasets.Value(dtype=ft["dtype"])
|
||||
elif len(ft["shape"]) == 1:
|
||||
hf_features[key] = datasets.Sequence(
|
||||
length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
|
||||
)
|
||||
# pyarrow rejects fixed-size lists of length 0, so use a variable length list instead
|
||||
length = ft["shape"][0] if ft["shape"][0] > 0 else -1
|
||||
hf_features[key] = datasets.Sequence(length=length, feature=datasets.Value(dtype=ft["dtype"]))
|
||||
elif len(ft["shape"]) == 2:
|
||||
hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
|
||||
elif len(ft["shape"]) == 3:
|
||||
|
||||
@@ -167,9 +167,7 @@ Show dataset information without feature details:
|
||||
--operation.type info \
|
||||
--operation.show_features false
|
||||
|
||||
Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default). The source
|
||||
dataset is never modified: large files are symlinked and only meta/ is copied, so this also works
|
||||
on read-only source datasets:
|
||||
Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type recompute_stats
|
||||
@@ -180,19 +178,6 @@ Recompute stats and save to a specific new repo_id:
|
||||
--new_repo_id lerobot/pusht_new_stats \
|
||||
--operation.type recompute_stats
|
||||
|
||||
Recompute stats including image/video features (samples and decodes frames from each episode):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type recompute_stats \
|
||||
--operation.skip_image_video false
|
||||
|
||||
Recompute stats and also rewrite the per-episode stats in the episodes parquet (keeps
|
||||
meta/stats.json and the per-episode stats consistent):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type recompute_stats \
|
||||
--operation.update_episode_stats true
|
||||
|
||||
Recompute stats in-place (overwrites original dataset stats):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
@@ -340,7 +325,6 @@ class RecomputeStatsConfig(OperationConfig):
|
||||
relative_exclude_joints: list[str] | None = None
|
||||
chunk_size: int = 50
|
||||
num_workers: int = 0
|
||||
update_episode_stats: bool = False
|
||||
overwrite: bool = False
|
||||
|
||||
|
||||
@@ -393,30 +377,6 @@ def _resolve_io_paths(
|
||||
return output_repo_id, input_path, output_path
|
||||
|
||||
|
||||
def _reference_copy_dataset(input_root: Path, output_root: Path) -> None:
|
||||
"""Create a lightweight copy of a dataset that never modifies the source.
|
||||
|
||||
The directory tree is recreated with real directories, and every file is
|
||||
symlinked to its source counterpart so no data is duplicated and the source is
|
||||
only ever read. Files under ``meta/`` are instead copied as real, writable files
|
||||
so that stats/info can be rewritten without touching the original. Symlinking
|
||||
individual files (rather than whole directories) keeps ``push_to_hub`` working,
|
||||
since ``Path.glob`` follows file symlinks but does not descend into symlinked
|
||||
directories. This makes the operation safe on read-only source datasets.
|
||||
"""
|
||||
for src in input_root.rglob("*"):
|
||||
rel = src.relative_to(input_root)
|
||||
dst = output_root / rel
|
||||
if src.is_dir():
|
||||
dst.mkdir(parents=True, exist_ok=True)
|
||||
elif rel.parts[0] == "meta":
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copyfile(src, dst) # copyfile ignores source perms, so dst is writable
|
||||
else:
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
dst.symlink_to(src.resolve())
|
||||
|
||||
|
||||
def get_output_path(
|
||||
repo_id: str,
|
||||
new_repo_id: str | None,
|
||||
@@ -714,18 +674,14 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
|
||||
)
|
||||
dataset = LeRobotDataset(cfg.repo_id, root=input_root)
|
||||
else:
|
||||
logging.info(f"Referencing dataset from {input_root} into {output_root} (source is left untouched)")
|
||||
logging.info(f"Copying dataset from {input_root} to {output_root}")
|
||||
if output_root.exists():
|
||||
backup_path = output_root.with_name(output_root.name + "_old")
|
||||
logging.warning(f"Output directory {output_root} already exists. Moving to {backup_path}")
|
||||
if backup_path.exists():
|
||||
shutil.rmtree(backup_path)
|
||||
shutil.move(output_root, backup_path)
|
||||
# recompute_stats only reads data/ and rewrites files under meta/ (stats.json, and
|
||||
# the episodes parquet when update_episode_stats is set), so symlink the large
|
||||
# immutable files and copy only meta/. This avoids duplicating the dataset and works
|
||||
# even when the source dataset is read-only.
|
||||
_reference_copy_dataset(input_root, output_root)
|
||||
shutil.copytree(input_root, output_root)
|
||||
dataset = LeRobotDataset(output_repo_id, root=output_root)
|
||||
|
||||
logging.info(f"Recomputing stats for {cfg.repo_id}")
|
||||
@@ -742,7 +698,6 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
|
||||
relative_exclude_joints=cfg.operation.relative_exclude_joints,
|
||||
chunk_size=cfg.operation.chunk_size,
|
||||
num_workers=cfg.operation.num_workers,
|
||||
update_episode_stats=cfg.operation.update_episode_stats,
|
||||
)
|
||||
|
||||
logging.info(f"Stats written to {dataset.root}")
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
# 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.
|
||||
import logging
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
@@ -687,6 +688,28 @@ def test_compute_episode_stats_string_features_skipped():
|
||||
assert "q01" in stats["action"]
|
||||
|
||||
|
||||
def test_compute_episode_stats_zero_width_features_skipped(caplog):
|
||||
"""Test that features with a zero-width dim (e.g. shape=(0,)) are skipped with a debug log."""
|
||||
episode_data = {
|
||||
"empty": np.zeros((100, 0), dtype=np.float32), # Zero-width feature
|
||||
"action": np.random.normal(0, 1, (100, 5)),
|
||||
}
|
||||
features = {
|
||||
"empty": {"dtype": "float32", "shape": (0,)},
|
||||
"action": {"dtype": "float32", "shape": (5,)},
|
||||
}
|
||||
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
stats = compute_episode_stats(episode_data, features)
|
||||
|
||||
# Zero-width features should be skipped with a debug log, others computed as usual
|
||||
assert "empty" not in stats
|
||||
assert "empty" in caplog.text
|
||||
assert "action" in stats
|
||||
assert "q01" in stats["action"]
|
||||
assert stats["action"]["mean"].shape == (5,)
|
||||
|
||||
|
||||
def test_aggregate_feature_stats_with_quantiles():
|
||||
"""Test aggregating feature stats that include quantiles."""
|
||||
stats_ft_list = [
|
||||
|
||||
@@ -1804,3 +1804,11 @@ def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
|
||||
root=dataset.root,
|
||||
episode_filter=lambda ep: ep["not_a_real_field"] > 0,
|
||||
)
|
||||
|
||||
|
||||
def test_get_hf_features_zero_width_feature_does_not_raise_on_from_dict():
|
||||
import datasets
|
||||
|
||||
features = {"empty": {"dtype": "float32", "shape": (0,), "names": ["empty"]}}
|
||||
hf_features = get_hf_features_from_features(features)
|
||||
datasets.Dataset.from_dict({"empty": [[], []]}, features=hf_features)
|
||||
|
||||
Reference in New Issue
Block a user