omni.isaac.lab.utils.noise.noise_model 源代码
# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations
import torch
from collections.abc import Sequence
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from . import noise_cfg
##
# Noise as functions.
##
[文档]def constant_noise(data: torch.Tensor, cfg: noise_cfg.ConstantNoiseCfg) -> torch.Tensor:
"""Applies a constant noise bias to a given data set.
Args:
data: The unmodified data set to apply noise to.
cfg: The configuration parameters for constant noise.
Returns:
The data modified by the noise parameters provided.
"""
# fix tensor device for bias on first call and update config parameters
if isinstance(cfg.bias, torch.Tensor):
cfg.bias = cfg.bias.to(device=data.device)
if cfg.operation == "add":
return data + cfg.bias
elif cfg.operation == "scale":
return data * cfg.bias
elif cfg.operation == "abs":
return torch.zeros_like(data) + cfg.bias
else:
raise ValueError(f"Unknown operation in noise: {cfg.operation}")
[文档]def gaussian_noise(data: torch.Tensor, cfg: noise_cfg.GaussianNoiseCfg) -> torch.Tensor:
"""Applies a gaussian noise to a given data set.
Args:
data: The unmodified data set to apply noise to.
cfg: The configuration parameters for gaussian noise.
Returns:
The data modified by the noise parameters provided.
"""
# fix tensor device for mean on first call and update config parameters
if isinstance(cfg.mean, torch.Tensor):
cfg.mean = cfg.mean.to(data.device)
# fix tensor device for std on first call and update config parameters
if isinstance(cfg.std, torch.Tensor):
cfg.std = cfg.std.to(data.device)
if cfg.operation == "add":
return data + cfg.mean + cfg.std * torch.randn_like(data)
elif cfg.operation == "scale":
return data * (cfg.mean + cfg.std * torch.randn_like(data))
elif cfg.operation == "abs":
return cfg.mean + cfg.std * torch.randn_like(data)
else:
raise ValueError(f"Unknown operation in noise: {cfg.operation}")
##
# Noise models as classes
##
[文档]class NoiseModel:
"""Base class for noise models."""
def __init__(self, noise_model_cfg: noise_cfg.NoiseModelCfg, num_envs: int, device: str):
"""Initialize the noise model.
Args:
noise_model_cfg: The noise configuration to use.
num_envs: The number of environments.
device: The device to use for the noise model.
"""
self._noise_model_cfg = noise_model_cfg
self._num_envs = num_envs
self._device = device
[文档] def reset(self, env_ids: Sequence[int] | None = None):
"""Reset the noise model.
This method can be implemented by derived classes to reset the noise model.
This is useful when implementing temporal noise models such as random walk.
Args:
env_ids: The environment ids to reset the noise model for. Defaults to None,
in which case all environments are considered.
"""
pass
[文档] def apply(self, data: torch.Tensor) -> torch.Tensor:
"""Apply the noise to the data.
Args:
data: The data to apply the noise to. Shape is (num_envs, ...).
Returns:
The data with the noise applied. Shape is the same as the input data.
"""
return self._noise_model_cfg.noise_cfg.func(data, self._noise_model_cfg.noise_cfg)
[文档]class NoiseModelWithAdditiveBias(NoiseModel):
"""Noise model with an additive bias.
The bias term is sampled from a the specified distribution on reset.
"""
def __init__(self, noise_model_cfg: noise_cfg.NoiseModelWithAdditiveBiasCfg, num_envs: int, device: str):
# initialize parent class
super().__init__(noise_model_cfg, num_envs, device)
# store the bias noise configuration
self._bias_noise_cfg = noise_model_cfg.bias_noise_cfg
self._bias = torch.zeros((num_envs, 1), device=self._device)
[文档] def reset(self, env_ids: Sequence[int] | None = None):
"""Reset the noise model.
This method resets the bias term for the specified environments.
Args:
env_ids: The environment ids to reset the noise model for. Defaults to None,
in which case all environments are considered.
"""
# resolve the environment ids
if env_ids is None:
env_ids = slice(None)
# reset the bias term
self._bias[env_ids] = self._bias_noise_cfg.func(self._bias[env_ids], self._bias_noise_cfg)
[文档] def apply(self, data: torch.Tensor) -> torch.Tensor:
"""Apply bias noise to the data.
Args:
data: The data to apply the noise to. Shape is (num_envs, ...).
Returns:
The data with the noise applied. Shape is the same as the input data.
"""
return super().apply(data) + self._bias
