normalizer

normalizer

Linear normalization module (min-max scaling, z-score, demean). Inspired by https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/model/common/normalizer.py

LinearNormalizer

LinearNormalizer(params_dict=None)

Bases: DictOfTensorMixin

Multi-key normalizer that stores per-key normalization parameters.

Supports dict input (per-key normalization) and tensor input (single-key '_default').

Source code in src/versatil/common/dict_of_tensor_mixin.py

def __init__(self, params_dict: nn.ParameterDict | None = None) -> None:
    super().__init__()
    self.params_dict = (
        params_dict if params_dict is not None else nn.ParameterDict()
    )

fit

fit(data, last_n_dims=1, dtype=float32, mode=value, output_max=1.0, output_min=-1.0, range_eps=0.0001, fit_offset=True, device=None, clamp_range=True, min_std=0.02, min_range=0.04, sample_size=0)

Fit normalization parameters to data.

Parameters:

Name	Type	Description	Default
data	dict | Tensor | ndarray	Dict of tensors (per-key fit) or a single tensor (stored under '_default').	required
last_n_dims	int	Number of trailing dimensions to treat as features.	1
dtype	dtype	Data type for the normalizer parameters.	float32
mode	str	Normalization mode ('min_max', 'gaussian', or 'demean').	value
output_max	float	Maximum output value for min_max mode.	1.0
output_min	float	Minimum output value for min_max mode.	-1.0
range_eps	float	Epsilon for detecting constant dimensions.	0.0001
fit_offset	bool	Whether to fit an offset (centering).	True
device	device | str | None	Device to place parameters on.	None
clamp_range	bool	Whether to clamp std/range to minimum values.	True
min_std	float	Minimum std value when clamp_range=True and mode='gaussian'.	0.02
min_range	float	Minimum range value when clamp_range=True and mode='min_max'.	0.04
sample_size	int | dict[str, int]	Number of data rows to store as an empirical sample under each key. Pass an int to use the same budget for every key, or a {key: size} dict to store samples only for selected keys (missing keys default to 0 = no sample stored).	0

Source code in src/versatil/data/normalization/normalizer.py

@torch.no_grad()
def fit(
    self,
    data: dict | torch.Tensor | np.ndarray,
    last_n_dims: int = 1,
    dtype: torch.dtype = torch.float32,
    mode: str = KinematicsNormalizationType.MIN_MAX.value,
    output_max: float = 1.0,
    output_min: float = -1.0,
    range_eps: float = 1e-4,
    fit_offset: bool = True,
    device: torch.device | str | None = None,
    clamp_range: bool = True,
    min_std: float = 2e-2,
    min_range: float = 4e-2,
    sample_size: int | dict[str, int] = 0,
):
    """Fit normalization parameters to data.

    Args:
        data: Dict of tensors (per-key fit) or a single tensor (stored under '_default').
        last_n_dims: Number of trailing dimensions to treat as features.
        dtype: Data type for the normalizer parameters.
        mode: Normalization mode ('min_max', 'gaussian', or 'demean').
        output_max: Maximum output value for min_max mode.
        output_min: Minimum output value for min_max mode.
        range_eps: Epsilon for detecting constant dimensions.
        fit_offset: Whether to fit an offset (centering).
        device: Device to place parameters on.
        clamp_range: Whether to clamp std/range to minimum values.
        min_std: Minimum std value when clamp_range=True and mode='gaussian'.
        min_range: Minimum range value when clamp_range=True and mode='min_max'.
        sample_size: Number of data rows to store as an empirical sample under
            each key. Pass an int to use the same budget for every key, or a
            ``{key: size}`` dict to store samples only for selected keys
            (missing keys default to 0 = no sample stored).
    """
    if isinstance(data, dict):
        for key, value in data.items():
            per_key_size = (
                sample_size.get(key, 0)
                if isinstance(sample_size, dict)
                else sample_size
            )
            self.params_dict[key] = _fit(
                value,
                last_n_dims=last_n_dims,
                dtype=dtype,
                mode=mode,
                output_max=output_max,
                output_min=output_min,
                range_eps=range_eps,
                fit_offset=fit_offset,
                device=device,
                clamp_range=clamp_range,
                min_std=min_std,
                min_range=min_range,
                sample_size=per_key_size,
            )
    else:
        single_size = (
            sample_size.get("_default", 0)
            if isinstance(sample_size, dict)
            else sample_size
        )
        self.params_dict["_default"] = _fit(
            data,
            last_n_dims=last_n_dims,
            dtype=dtype,
            mode=mode,
            output_max=output_max,
            output_min=output_min,
            range_eps=range_eps,
            fit_offset=fit_offset,
            device=device,
            clamp_range=clamp_range,
            min_std=min_std,
            min_range=min_range,
            sample_size=single_size,
        )

normalize

normalize(x)

Normalize input data. Dict keys without parameters pass through unchanged.

Source code in src/versatil/data/normalization/normalizer.py

def normalize(self, x: dict | torch.Tensor | np.ndarray) -> dict | torch.Tensor:
    """Normalize input data. Dict keys without parameters pass through unchanged."""
    return self._normalize_impl(x, forward=True)

unnormalize

unnormalize(x)

Unnormalize input data back to original scale.

Source code in src/versatil/data/normalization/normalizer.py

def unnormalize(self, x: dict | torch.Tensor | np.ndarray) -> dict | torch.Tensor:
    """Unnormalize input data back to original scale."""
    return self._normalize_impl(x, forward=False)

get_input_stats

get_input_stats()

Get input statistics (min, max, mean, std) for all keys.

Source code in src/versatil/data/normalization/normalizer.py

def get_input_stats(self) -> dict:
    """Get input statistics (min, max, mean, std) for all keys."""
    if len(self.params_dict) == 0:
        raise RuntimeError("Not initialized")
    if len(self.params_dict) == 1 and "_default" in self.params_dict:
        return self.params_dict["_default"]["input_stats"]

    result = {}
    for key, value in self.params_dict.items():
        if key != "_default":
            result[key] = value["input_stats"]
    return result

get_output_stats

get_output_stats(key='_default')

Get normalized output statistics for a specific key.

Parameters:

Name	Type	Description	Default
key	str	Parameter key. Use '_default' for tensor-fitted normalizers, or a specific key for dict-fitted normalizers.	'_default'

Returns:

Type	Description
dict	Dict with normalized min, max, mean, and scaled std.

Source code in src/versatil/data/normalization/normalizer.py

def get_output_stats(self, key: str = "_default") -> dict:
    """Get normalized output statistics for a specific key.

    Args:
        key: Parameter key. Use '_default' for tensor-fitted normalizers,
            or a specific key for dict-fitted normalizers.

    Returns:
        Dict with normalized min, max, mean, and scaled std.
    """
    params = self.params_dict[key]
    input_stats = params["input_stats"]
    scale = params["scale"]
    single_field = SingleFieldLinearNormalizer(params)
    result = {}
    if "min" in input_stats:
        result["min"] = single_field.normalize(input_stats["min"])
    if "max" in input_stats:
        result["max"] = single_field.normalize(input_stats["max"])
    if "mean" in input_stats:
        result["mean"] = single_field.normalize(input_stats["mean"])
    if "std" in input_stats:
        result["std"] = torch.abs(scale) * input_stats["std"]
    return result

SingleFieldLinearNormalizer

SingleFieldLinearNormalizer(params_dict=None)

Bases: DictOfTensorMixin

Single-field normalizer with scale and offset parameters.

Source code in src/versatil/common/dict_of_tensor_mixin.py

def __init__(self, params_dict: nn.ParameterDict | None = None) -> None:
    super().__init__()
    self.params_dict = (
        params_dict if params_dict is not None else nn.ParameterDict()
    )

fit

fit(data, last_n_dims=1, dtype=float32, mode=value, output_max=1.0, output_min=-1.0, range_eps=1e-10, fit_offset=True, device='cpu', clamp_range=True, min_std=0.02, min_range=0.04, sample_size=0)

Fit scale and offset for one field from data statistics.

Source code in src/versatil/data/normalization/normalizer.py

@torch.no_grad()
def fit(
    self,
    data: torch.Tensor | np.ndarray,
    last_n_dims: int = 1,
    dtype: torch.dtype = torch.float32,
    mode: str = KinematicsNormalizationType.MIN_MAX.value,
    output_max: float = 1.0,
    output_min: float = -1.0,
    range_eps: float = 1e-10,
    fit_offset: bool = True,
    device: torch.device | str | None = "cpu",
    clamp_range: bool = True,
    min_std: float = 2e-2,
    min_range: float = 4e-2,
    sample_size: int = 0,
):
    """Fit scale and offset for one field from data statistics."""
    self.params_dict = _fit(
        data,
        last_n_dims=last_n_dims,
        dtype=dtype,
        mode=mode,
        output_max=output_max,
        output_min=output_min,
        range_eps=range_eps,
        fit_offset=fit_offset,
        device=device,
        clamp_range=clamp_range,
        min_std=min_std,
        min_range=min_range,
        sample_size=sample_size,
    )

create_fit classmethod

create_fit(data, last_n_dims=1, dtype=float32, mode=value, output_max=1.0, output_min=-1.0, range_eps=1e-10, fit_offset=True, device='cpu', clamp_range=True, min_std=0.02, min_range=0.04)

Create a fitted normalizer in one step.

Source code in src/versatil/data/normalization/normalizer.py

@classmethod
def create_fit(
    cls,
    data: torch.Tensor | np.ndarray,
    last_n_dims: int = 1,
    dtype: torch.dtype = torch.float32,
    mode: str = KinematicsNormalizationType.MIN_MAX.value,
    output_max: float = 1.0,
    output_min: float = -1.0,
    range_eps: float = 1e-10,
    fit_offset: bool = True,
    device: torch.device | str | None = "cpu",
    clamp_range: bool = True,
    min_std: float = 2e-2,
    min_range: float = 4e-2,
):
    """Create a fitted normalizer in one step."""
    obj = cls()
    obj.fit(
        data,
        last_n_dims=last_n_dims,
        dtype=dtype,
        mode=mode,
        output_max=output_max,
        output_min=output_min,
        range_eps=range_eps,
        fit_offset=fit_offset,
        device=device,
        clamp_range=clamp_range,
        min_std=min_std,
        min_range=min_range,
    )
    return obj

create_manual classmethod

create_manual(scale, offset, input_stats_dict)

Create a normalizer from explicit scale, offset, and input stats.

Source code in src/versatil/data/normalization/normalizer.py

@classmethod
def create_manual(
    cls,
    scale: torch.Tensor | np.ndarray,
    offset: torch.Tensor | np.ndarray,
    input_stats_dict: dict[str, torch.Tensor | np.ndarray],
):
    """Create a normalizer from explicit scale, offset, and input stats."""

    def to_tensor(x):
        """Convert input to a flat float tensor."""
        if not isinstance(x, torch.Tensor):
            x = torch.from_numpy(x)
        x = x.flatten()
        return x

    for tensor in [offset] + list(input_stats_dict.values()):
        if tensor.shape != scale.shape:
            raise ValueError(
                f"All inputs must have same shape as scale {scale.shape}, got {tensor.shape}"
            )
        if tensor.dtype != scale.dtype:
            raise ValueError(
                f"All inputs must have same dtype as scale {scale.dtype}, got {tensor.dtype}"
            )

    params_dict = nn.ParameterDict(
        {
            "scale": to_tensor(scale),
            "offset": to_tensor(offset),
            "input_stats": nn.ParameterDict(
                dict_apply(input_stats_dict, to_tensor)
            ),
        }
    )
    for parameter in params_dict.parameters():
        parameter.requires_grad_(False)
    return cls(params_dict)

create_identity classmethod

create_identity(dtype=float32)

Create an identity normalizer (scale=1, offset=0).

Source code in src/versatil/data/normalization/normalizer.py

@classmethod
def create_identity(cls, dtype: torch.dtype = torch.float32):
    """Create an identity normalizer (scale=1, offset=0)."""
    scale = torch.tensor([1], dtype=dtype)
    offset = torch.tensor([0], dtype=dtype)
    input_stats_dict = {
        "min": torch.tensor([-1], dtype=dtype),
        "max": torch.tensor([1], dtype=dtype),
        "mean": torch.tensor([0], dtype=dtype),
        "std": torch.tensor([1], dtype=dtype),
    }
    return cls.create_manual(scale, offset, input_stats_dict)

normalize

normalize(x)

Apply normalization: x * scale + offset.

Source code in src/versatil/data/normalization/normalizer.py

def normalize(self, x: torch.Tensor | np.ndarray) -> torch.Tensor:
    """Apply normalization: x * scale + offset."""
    return _normalize(x, self.params_dict, forward=True)

unnormalize

unnormalize(x)

Apply inverse normalization: (x - offset) / scale.

Source code in src/versatil/data/normalization/normalizer.py

def unnormalize(self, x: torch.Tensor | np.ndarray) -> torch.Tensor:
    """Apply inverse normalization: (x - offset) / scale."""
    return _normalize(x, self.params_dict, forward=False)

get_input_stats

get_input_stats()

Get input statistics (min, max, mean, std).

Source code in src/versatil/data/normalization/normalizer.py

def get_input_stats(self) -> nn.ParameterDict:
    """Get input statistics (min, max, mean, std)."""
    return self.params_dict["input_stats"]

get_output_stats

get_output_stats()

Get normalized output statistics.

Source code in src/versatil/data/normalization/normalizer.py

def get_output_stats(self) -> dict:
    """Get normalized output statistics."""
    input_stats = self.params_dict["input_stats"]
    scale = self.params_dict["scale"]
    result = {}
    if "min" in input_stats:
        result["min"] = self.normalize(input_stats["min"])
    if "max" in input_stats:
        result["max"] = self.normalize(input_stats["max"])
    if "mean" in input_stats:
        result["mean"] = self.normalize(input_stats["mean"])
    if "std" in input_stats:
        result["std"] = torch.abs(scale) * input_stats["std"]
    return result

get_input_sample

get_input_sample()

Return the raw subsample stored at fit time, or None if absent.

Present only when fit(..., sample_size>0) was used.

Source code in src/versatil/data/normalization/normalizer.py

def get_input_sample(self) -> torch.Tensor | None:
    """Return the raw subsample stored at fit time, or ``None`` if absent.

    Present only when ``fit(..., sample_size>0)`` was used.
    """
    if "sample" in self.params_dict:
        return self.params_dict["sample"]
    return None

get_output_sample

get_output_sample()

Return the stored subsample in normalized output space, or None.

Source code in src/versatil/data/normalization/normalizer.py

def get_output_sample(self) -> torch.Tensor | None:
    """Return the stored subsample in normalized output space, or ``None``."""
    raw_sample = self.get_input_sample()
    if raw_sample is None:
        return None
    return self.normalize(raw_sample)

SequentialNormalizer

SequentialNormalizer(normalizers)

Bases: SingleFieldLinearNormalizer

Applies a sequence of already-fitted normalizers.

Source code in src/versatil/data/normalization/normalizer.py

def __init__(self, normalizers: list[SingleFieldLinearNormalizer]):
    super().__init__()
    self.normalizers = nn.ModuleList(normalizers)
    if len(self.normalizers) == 0:
        raise ValueError("At least one normalizer required")
    effective_scale = self.normalizers[0].params_dict["scale"].clone().detach()
    effective_offset = self.normalizers[0].params_dict["offset"].clone().detach()
    for normalizer in self.normalizers[1:]:
        normalizer_scale = normalizer.params_dict["scale"].clone().detach()
        normalizer_offset = normalizer.params_dict["offset"].clone().detach()
        effective_offset = normalizer_offset + normalizer_scale * effective_offset
        effective_scale = normalizer_scale * effective_scale

    input_stats = self.normalizers[0].params_dict["input_stats"]
    self.params_dict = nn.ParameterDict(
        {
            "scale": nn.Parameter(effective_scale, requires_grad=False),
            "offset": nn.Parameter(effective_offset, requires_grad=False),
            "input_stats": nn.ParameterDict(
                {
                    k: nn.Parameter(v.clone().detach(), requires_grad=False)
                    for k, v in input_stats.items()
                }
            ),
        }
    )

fit

fit(data=None, last_n_dims=1, dtype=float32, mode=value, output_max=1.0, output_min=-1.0, range_eps=1e-10, fit_offset=True, device='cpu', clamp_range=True, min_std=0.02, min_range=0.04)

Fit per-key normalizers over a dict of arrays.

Source code in src/versatil/data/normalization/normalizer.py

def fit(
    self,
    data: torch.Tensor | np.ndarray = None,
    last_n_dims: int = 1,
    dtype: torch.dtype = torch.float32,
    mode: str = KinematicsNormalizationType.MIN_MAX.value,
    output_max: float = 1.0,
    output_min: float = -1.0,
    range_eps: float = 1e-10,
    fit_offset: bool = True,
    device: torch.device | str | None = "cpu",
    clamp_range: bool = True,
    min_std: float = 2e-2,
    min_range: float = 4e-2,
) -> None:
    """Fit per-key normalizers over a dict of arrays."""
    raise NotImplementedError("SequentialNormalizer does not support fitting")

normalize

normalize(x)

Apply all normalizers in sequence.

Source code in src/versatil/data/normalization/normalizer.py

def normalize(self, x: torch.Tensor | np.ndarray) -> torch.Tensor:
    """Apply all normalizers in sequence."""
    for normalizer in self.normalizers:
        x = normalizer.normalize(x)
    return x

unnormalize

unnormalize(x)

Apply all normalizers in reverse sequence.

Source code in src/versatil/data/normalization/normalizer.py

def unnormalize(self, x: torch.Tensor | np.ndarray) -> torch.Tensor:
    """Apply all normalizers in reverse sequence."""
    for normalizer in reversed(self.normalizers):
        x = normalizer.unnormalize(x)
    return x

get_input_stats

get_input_stats()

Get input stats of the first normalizer.

Source code in src/versatil/data/normalization/normalizer.py

def get_input_stats(self) -> nn.ParameterDict:
    """Get input stats of the first normalizer."""
    return self.normalizers[0].get_input_stats()

get_output_stats

get_output_stats()

Get output stats after propagating through all normalizers.

Source code in src/versatil/data/normalization/normalizer.py

def get_output_stats(self) -> dict:
    """Get output stats after propagating through all normalizers."""
    stats = self.get_input_stats()
    for normalizer in self.normalizers:
        scale = normalizer.params_dict["scale"]
        offset = normalizer.params_dict["offset"]
        stats = {
            "min": scale * stats["min"] + offset,
            "max": scale * stats["max"] + offset,
            "mean": scale * stats["mean"] + offset,
            "std": torch.abs(scale) * stats["std"],
        }
    return stats