From 1cd7ca71a188eb35c8e8a6a1d20e7a6038b63d38 Mon Sep 17 00:00:00 2001 From: Thomas Wolf Date: Tue, 18 Jun 2024 13:44:28 +0200 Subject: [PATCH] update --- lerobot/common/policies/normalize.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/lerobot/common/policies/normalize.py b/lerobot/common/policies/normalize.py index a711dcaa..8f9096a6 100644 --- a/lerobot/common/policies/normalize.py +++ b/lerobot/common/policies/normalize.py @@ -162,6 +162,9 @@ class Normalize(nn.Module): output_batch[key] = output_batch[key] * 2 - 1 else: raise ValueError(mode) + for key in batch: + if key not in output_batch: + output_batch[key] = batch[key] return output_batch @@ -231,4 +234,7 @@ class Unnormalize(nn.Module): output_batch[key] = output_batch[key] * (max - min) + min else: raise ValueError(mode) + for key in batch: + if key not in output_batch: + output_batch[key] = batch[key] return output_batch