Resolve the source root via metadata only instead of instantiating a full
LeRobotDataset (which memory-maps the entire frame index just to read a path),
and download videos only when the run actually recomputes image/video stats
(derived from --skip-image-video). This avoids the multi-TB video download for
numeric-only runs.
Add an opt-in update_episode_stats flag so recomputing stats can also rewrite
the per-episode stats/* columns in the episodes parquet, keeping them consistent
with meta/stats.json. compute_dataset_episode_stats now returns a {episode_index:
stats} mapping so stats can be written back to the right episode and shards merge
by key. Wired into recompute_stats, the lerobot-edit-dataset CLI, and the SLURM
example (per-episode shards + --update-episode-stats).
The pipeline steps are pickled and run on workers where this script's module
globals are unavailable, so referencing the module-level _load_dataset helper
raised NameError. Inline dataset loading and shard paths into each run() with
local imports, and let --venv-path and --env-command coexist.
Drops the per-worker hf-mount machinery in favor of a shared HF_LEROBOT_HOME
cache plus --venv-path and --env-command hooks, which is simpler and avoids
node-local mount setup. Removes the now-unused os import.
Adds HPC cluster support to the SLURM stats recomputation example: a --qos
passthrough, per-worker hf-mount of the read-only source via datatrove's
env_command hook, and --chain-aggregate to submit aggregate with an afterok
dependency on compute. Also switches to datatrove's native mem_per_cpu_gb field.
Depth video stats were computed on raw 12-bit codec values, leaving them in
codec space instead of the recorded depth unit. Dequantize decoded frames via
the feature's depth encoder config (matching DatasetReader) so recomputed stats
match record-time stats.
Also fix the SLURM example: the --skip-image-video flag was inverted (0 skipped
visual stats), and add a --video-backend option so pyav can be used when
torchcodec fails to load locally.
Expose the shardable unit of work behind recompute_stats: compute_dataset_episode_stats
computes per-episode stats for an episode subset, and aggregate_episode_stats merges the
concatenated shards (count-weighted) and writes stats.json. recompute_stats now composes
these, so single-process behavior is unchanged.
Add examples/dataset/slurm_recompute_stats.py, a datatrove compute/aggregate driver that
shards episodes across workers and is read-only safe (reference-copies the source when
--new-root is given). Most useful for the expensive image/video stats path.
CarolinePascal