8000 Improvement: Update normalization in ACT + remove in-place modification + add test by thomwolf · Pull Request #281 · huggingface/lerobot · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

Improvement: Update normalization in ACT + remove in-place modification + add test #281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Improvement: Update normalization in ACT + remove in-place modification + add test #281

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cd9ace2
update
thomwolf Jun 18, 2024
c9150c3
update unnormalize
thomwolf Jun 18, 2024
1cd7ca7
update
thomwolf Jun 18, 2024
33166e1
update
thomwolf Jun 19, 2024
f9295e2
Merge branch 'main' into thomwolf_2024_06_18_fix_normalization
alexander-soare Jun 19, 2024
8efe947
Merge branch 'main' into thomwolf_2024_06_18_fix_normalization
alexander-soare Jun 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 33 additions & 8 deletions lerobot/common/policies/normalize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ def __init__(
shapes: dict[str, list[int]],
modes: dict[str, str],
stats: dict[str, dict[str, Tensor]] | None = None,
std_epsilon: float = 1e-5,
):
"""
Args:
Expand All @@ -120,18 +121,23 @@ def __init__(
not provided, as expected for finetuning or evaluation, the default buffers should to be
overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
dataset is not needed to get the stats, since they are already in the policy state_dict.
std_epsilon (float, optional): A small minimal value for the standard deviation to avoid division by
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Would it be more suitable to generalize the variable name (to epsilon maybe) and description? It pertains not just to the standard deviation, but also the min/max. In fact, I think it's very common place to use min/max normalization so it would be nice to make sure std_epsilon is not dismissed up front by the reader.

zero. Default is `1e-5`. We use `clamp_min` to make sure the standard deviation (or the difference
between min and max) is at least `std_epsilon`.
"""
super().__init__()
self.shapes = shapes
self.modes = modes
self.stats = stats
self.std_epsilon = std_epsilon
stats_buffers = create_stats_buffers(shapes, modes, stats)
for key, buffer in stats_buffers.items():
setattr(self, "buffer_" + key.replace(".", "_"), buffer)

# TODO(rcadene): should we remove torch.no_grad?
@torch.no_grad
def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
output_batch = {}
for key, mode in self.modes.items():
buffer = getattr(self, "buffer_" + key.replace(".", "_"))

Expand All @@ -140,19 +146,25 @@ def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
std = buffer["std"]
assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
assert not torch.isinf(std).any(), _no_stats_error_str("std")
batch[key] = (batch[key] - mean) / (std + 1e-8)
output_batch[key] = (batch[key] - mean) / std.clamp_min(self.std_epsilon)
elif mode == "min_max":
min = buffer["min"]
max = buffer["max"]
assert not torch.isinf(min).any(), _no_stats_error_str("min")
assert not torch.isinf(max).any(), _no_stats_error_str("max")
# To add to have at least std_epsilon between min and max
epsilon = (self.std_epsilon - (max - min).abs()).clamp_min(0)
max = max + epsilon
# normalize to [0,1]
batch[key] = (batch[key] - min) / (max - min + 1e-8)
output_batch[key] = (batch[key] - min) / (max - min)
# normalize to [-1, 1]
batch[key] = batch[key] * 2 - 1
output_batch[key] = output_batch[key] * 2 - 1
else:
raise ValueError(mode)
return batch
for key in batch:
if key not in output_batch:
output_batch[key] = batch[key]
return output_batch


class Unnormalize(nn.Module):
Expand All @@ -166,6 +178,7 @@ def __init__(
shapes: dict[str, list[int]],
modes: dict[str, str],
stats: dict[str, dict[str, Tensor]] | None = None,
std_epsilon: float = 1e-5,
):
"""
Args:
Expand All @@ -184,11 +197,16 @@ def __init__(
not provided, as expected for finetuning or evaluation, the default buffers should to be
overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
dataset is not needed to get the stats, since they are already in the policy state_dict.
std_epsilon (float, optional): A small minimal value for the standard deviation to avoid division by
zero in the Normalize step. We use the same value for unnormalization here to have a consistent
behavior. Default is `1e-5`. We use `clamp_min` to make sure the standard deviation (or the difference
between min and max) is at least `std_epsilon`.
"""
super().__init__()
self.shapes = shapes
self.modes = modes
self.stats = stats
self.std_epsilon = std_epsilon
# `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
stats_buffers = create_stats_buffers(shapes, modes, stats)
for key, buffer in stats_buffers.items():
Expand All @@ -197,6 +215,7 @@ def __init__(
# TODO(rcadene): should we remove torch.no_grad?
@torch.no_grad
def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
output_batch = {}
for key, mode in self.modes.items():
buffer = getattr(self, "buffer_" + key.replace(".", "_"))

Expand All @@ -205,14 +224,20 @@ def forward(self, batch: dict[str, Tensor]) -> dict[str, Tenso 83E8 r]:
std = buffer["std"]
assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
assert not torch.isinf(std).any(), _no_stats_error_str("std")
batch[key] = batch[key] * std + mean
output_batch[key] = batch[key] * std.clamp_min(self.std_epsilon) + mean
elif mode == "min_max":
min = buffer["min"]
max = buffer["max"]
assert not torch.isinf(min).any(), _no_stats_error_str("min")
assert not torch.isinf(max).any(), _no_stats_error_str("max")
batch[key] = (batch[key] + 1) / 2
batch[key] = batch[key] * (max - min) + min
# To add to have at least std_epsilon between min and max
epsilon = (self.std_epsilon - (max - min).abs()).clamp_min(0)
max = max + epsilon
output_batch[key] = (batch[key] + 1) / 2
output_batch[key] = output_batch[key] * (max - min) + min
else:
raise ValueError(mode)
return batch
for key in batch:
if key not in output_batch:
output_batch[key] = batch[key]
return output_batch
Loading
Loading
0