From 900f1a42e97a8eaa4cef355023bd87f88419632e Mon Sep 17 00:00:00 2001 From: pepijn Date: Wed, 1 Apr 2026 19:12:14 +0000 Subject: [PATCH] fix relative stats computation with rename_map 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 --- src/lerobot/scripts/lerobot_train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index deb5a4681..a9c0ddb54 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -324,7 +324,9 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): temp_loader = torch.utils.data.DataLoader( 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 gc.collect() torch.save({"mean": mean, "std": std}, stats_path)