admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-20 08:21:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix relative stats computation with rename_map

Browse Source 
When rename_map maps a dataset key to observation.state, the raw
dataset used for stats computation still has the original key.
Reverse the rename_map to find the correct key.

Made-with: Cursor
This commit is contained in:
pepijn
2026-04-01 19:12:14 +00:00
parent 5a15a6a911
commit 900f1a42e9
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+3 -1
View File
@@ -324,7 +324,9 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
temp_loader = torch.utils.data.DataLoader( temp_loader = torch.utils.data.DataLoader(
stats_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0 stats_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0
) )
mean, std = compute_relative_action_stats(temp_loader, num_batches=1000) reverse_rename = {v: k for k, v in cfg.rename_map.items()} if cfg.rename_map else {}
raw_state_key = reverse_rename.get("observation.state", "observation.state")
mean, std = compute_relative_action_stats(temp_loader, state_key=raw_state_key, num_batches=1000)
del temp_loader, stats_dataset del temp_loader, stats_dataset
gc.collect() gc.collect()
torch.save({"mean": mean, "std": std}, stats_path) torch.save({"mean": mean, "std": std}, stats_path)