255 lines
12 KiB
Python
255 lines
12 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import numpy as np
|
|
import torch
|
|
from torch import Tensor, nn
|
|
|
|
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
|
|
|
|
|
def create_stats_buffers(
|
|
features: dict[str, PolicyFeature],
|
|
norm_map: dict[str, NormalizationMode],
|
|
stats: dict[str, dict[str, Tensor]] | None = None,
|
|
) -> dict[str, dict[str, nn.ParameterDict]]:
|
|
"""
|
|
Create buffers per modality (e.g. "observation.image", "action") containing their mean, std, min, max
|
|
statistics.
|
|
|
|
Args: (see Normalize and Unnormalize)
|
|
|
|
Returns:
|
|
dict: A dictionary where keys are modalities and values are `nn.ParameterDict` containing
|
|
`nn.Parameters` set to `requires_grad=False`, suitable to not be updated during backpropagation.
|
|
"""
|
|
stats_buffers = {}
|
|
|
|
for key, ft in features.items():
|
|
norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
|
|
if norm_mode is NormalizationMode.IDENTITY:
|
|
continue
|
|
|
|
assert isinstance(norm_mode, NormalizationMode)
|
|
|
|
shape = tuple(ft.shape)
|
|
|
|
if ft.type is FeatureType.VISUAL:
|
|
# sanity checks
|
|
assert len(shape) == 3, f"number of dimensions of {key} != 3 ({shape=}"
|
|
c, h, w = shape
|
|
assert c < h and c < w, f"{key} is not channel first ({shape=})"
|
|
# override image shape to be invariant to height and width
|
|
shape = (c, 1, 1)
|
|
|
|
# Note: we initialize mean, std, min, max to infinity. They should be overwritten
|
|
# downstream by `stats` or `policy.load_state_dict`, as expected. During forward,
|
|
# we assert they are not infinity anymore.
|
|
|
|
buffer = {}
|
|
if norm_mode is NormalizationMode.MEAN_STD:
|
|
mean = torch.ones(shape, dtype=torch.float32) * torch.inf
|
|
std = torch.ones(shape, dtype=torch.float32) * torch.inf
|
|
buffer = nn.ParameterDict(
|
|
{
|
|
"mean": nn.Parameter(mean, requires_grad=False),
|
|
"std": nn.Parameter(std, requires_grad=False),
|
|
}
|
|
)
|
|
elif norm_mode is NormalizationMode.MIN_MAX:
|
|
min = torch.ones(shape, dtype=torch.float32) * torch.inf
|
|
max = torch.ones(shape, dtype=torch.float32) * torch.inf
|
|
buffer = nn.ParameterDict(
|
|
{
|
|
"min": nn.Parameter(min, requires_grad=False),
|
|
"max": nn.Parameter(max, requires_grad=False),
|
|
}
|
|
)
|
|
|
|
# TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
|
|
if stats:
|
|
if isinstance(stats[key]["mean"], np.ndarray):
|
|
if norm_mode is NormalizationMode.MEAN_STD:
|
|
buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
|
|
buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
|
|
elif norm_mode is NormalizationMode.MIN_MAX:
|
|
buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
|
|
buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
|
|
elif isinstance(stats[key]["mean"], torch.Tensor):
|
|
# Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
|
|
# tensors anywhere (for example, when we use the same stats for normalization and
|
|
# unnormalization). See the logic here
|
|
# https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
|
|
if norm_mode is NormalizationMode.MEAN_STD:
|
|
buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
|
|
buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
|
|
elif norm_mode is NormalizationMode.MIN_MAX:
|
|
buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
|
|
buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
|
|
else:
|
|
type_ = type(stats[key]["mean"])
|
|
raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
|
|
|
|
stats_buffers[key] = buffer
|
|
return stats_buffers
|
|
|
|
|
|
def _no_stats_error_str(name: str) -> str:
|
|
return (
|
|
f"`{name}` is infinity. You should either initialize with `stats` as an argument, or use a "
|
|
"pretrained model."
|
|
)
|
|
|
|
|
|
class Normalize(nn.Module):
|
|
"""Normalizes data (e.g. "observation.image") for more stable and faster convergence during training."""
|
|
|
|
def __init__(
|
|
self,
|
|
features: dict[str, PolicyFeature],
|
|
norm_map: dict[str, NormalizationMode],
|
|
stats: dict[str, dict[str, Tensor]] | None = None,
|
|
):
|
|
"""
|
|
Args:
|
|
shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
|
|
are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
|
|
mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
|
|
is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
|
|
modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
|
|
are their normalization modes among:
|
|
- "mean_std": subtract the mean and divide by standard deviation.
|
|
- "min_max": map to [-1, 1] range.
|
|
stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
|
|
and values are dictionaries of statistic types and their values (e.g.
|
|
`{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
|
|
training the model for the first time, these statistics will overwrite the default buffers. If
|
|
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.
|
|
"""
|
|
super().__init__()
|
|
self.features = features
|
|
self.norm_map = norm_map
|
|
self.stats = stats
|
|
stats_buffers = create_stats_buffers(features, norm_map, 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]:
|
|
batch = dict(batch) # shallow copy avoids mutating the input batch
|
|
for key, ft in self.features.items():
|
|
if key not in batch:
|
|
# FIXME(aliberts, rcadene): This might lead to silent fail!
|
|
continue
|
|
|
|
norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
|
|
if norm_mode is NormalizationMode.IDENTITY:
|
|
continue
|
|
|
|
buffer = getattr(self, "buffer_" + key.replace(".", "_"))
|
|
|
|
if norm_mode is NormalizationMode.MEAN_STD:
|
|
mean = buffer["mean"]
|
|
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)
|
|
elif norm_mode is NormalizationMode.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")
|
|
# normalize to [0,1]
|
|
batch[key] = (batch[key] - min) / (max - min + 1e-8)
|
|
# normalize to [-1, 1]
|
|
batch[key] = batch[key] * 2 - 1
|
|
else:
|
|
raise ValueError(norm_mode)
|
|
return batch
|
|
|
|
|
|
class Unnormalize(nn.Module):
|
|
"""
|
|
Similar to `Normalize` but unnormalizes output data (e.g. `{"action": torch.randn(b,c)}`) in their
|
|
original range used by the environment.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
features: dict[str, PolicyFeature],
|
|
norm_map: dict[str, NormalizationMode],
|
|
stats: dict[str, dict[str, Tensor]] | None = None,
|
|
):
|
|
"""
|
|
Args:
|
|
shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
|
|
are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
|
|
mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
|
|
is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
|
|
modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
|
|
are their normalization modes among:
|
|
- "mean_std": subtract the mean and divide by standard deviation.
|
|
- "min_max": map to [-1, 1] range.
|
|
stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
|
|
and values are dictionaries of statistic types and their values (e.g.
|
|
`{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
|
|
training the model for the first time, these statistics will overwrite the default buffers. If
|
|
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.
|
|
"""
|
|
super().__init__()
|
|
self.features = features
|
|
self.norm_map = norm_map
|
|
self.stats = stats
|
|
# `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
|
|
stats_buffers = create_stats_buffers(features, norm_map, 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]:
|
|
batch = dict(batch) # shallow copy avoids mutating the input batch
|
|
for key, ft in self.features.items():
|
|
if key not in batch:
|
|
continue
|
|
|
|
norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
|
|
if norm_mode is NormalizationMode.IDENTITY:
|
|
continue
|
|
|
|
buffer = getattr(self, "buffer_" + key.replace(".", "_"))
|
|
|
|
if norm_mode is NormalizationMode.MEAN_STD:
|
|
mean = buffer["mean"]
|
|
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
|
|
elif norm_mode is NormalizationMode.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
|
|
else:
|
|
raise ValueError(norm_mode)
|
|
return batch
|