Integrates NoiseModel to manager-based workflows

source/isaaclab/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.40.6"
+version = "0.40.X"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/isaaclab/docs/CHANGELOG.rst

@@ -1,6 +1,20 @@
 Changelog
 ---------
 
+## [Unreleased]
+~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* :class:`NoiseModel` support for manager-based workflows.
+
+Changed
+^^^^^^^
+
+* Renamed :func:`~isaaclab.utils.noise.NoiseModel.apply` method to :func:`~isaaclab.utils.noise.NoiseModel.__call__`.
+
+
 0.40.6 (2025-06-12)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/envs/direct_marl_env.py

@@ -356,7 +356,7 @@ def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn:
         if self.cfg.action_noise_model:
             for agent, action in actions.items():
                 if agent in self._action_noise_model:
-                    actions[agent] = self._action_noise_model[agent].apply(action)
+                    actions[agent] = self._action_noise_model[agent](action)
         # process actions
         self._pre_physics_step(actions)
 
@@ -409,7 +409,7 @@ def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn:
         if self.cfg.observation_noise_model:
             for agent, obs in self.obs_dict.items():
                 if agent in self._observation_noise_model:
-                    self.obs_dict[agent] = self._observation_noise_model[agent].apply(obs)
+                    self.obs_dict[agent] = self._observation_noise_model[agent](obs)
 
         # return observations, rewards, resets and extras
         return self.obs_dict, self.reward_dict, self.terminated_dict, self.time_out_dict, self.extras

source/isaaclab/isaaclab/envs/direct_rl_env.py

@@ -329,7 +329,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         action = action.to(self.device)
         # add action noise
         if self.cfg.action_noise_model:
-            action = self._action_noise_model.apply(action)
+            action = self._action_noise_model(action)
 
         # process actions
         self._pre_physics_step(action)
@@ -386,7 +386,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         # add observation noise
         # note: we apply no noise to the state space (since it is used for critic networks)
         if self.cfg.observation_noise_model:
-            self.obs_buf["policy"] = self._observation_noise_model.apply(self.obs_buf["policy"])
+            self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"])
 
         # return observations, rewards, resets and extras
         return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras

source/isaaclab/isaaclab/managers/manager_term_cfg.py

@@ -14,7 +14,7 @@
 
 from isaaclab.utils import configclass
 from isaaclab.utils.modifiers import ModifierCfg
-from isaaclab.utils.noise import NoiseCfg
+from isaaclab.utils.noise import NoiseCfg, NoiseModelCfg
 
 from .scene_entity_cfg import SceneEntityCfg
 
@@ -165,7 +165,7 @@ class ObservationTermCfg(ManagerTermBaseCfg):
     For more information on modifiers, see the :class:`~isaaclab.utils.modifiers.ModifierCfg` class.
     """
 
-    noise: NoiseCfg | None = None
+    noise: NoiseCfg | NoiseModelCfg | None = None
     """The noise to add to the observation. Defaults to None, in which case no noise is added."""
 
     clip: tuple[float, float] | None = None

source/isaaclab/isaaclab/managers/observation_manager.py

@@ -14,7 +14,7 @@
 from prettytable import PrettyTable
 from typing import TYPE_CHECKING
 
-from isaaclab.utils import class_to_dict, modifiers
+from isaaclab.utils import class_to_dict, modifiers, noise
 from isaaclab.utils.buffers import CircularBuffer
 
 from .manager_base import ManagerBase, ManagerTermBase
@@ -239,7 +239,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
                 if term_name in self._group_obs_term_history_buffer[group_name]:
                     self._group_obs_term_history_buffer[group_name][term_name].reset(batch_ids=env_ids)
         # call all modifiers that are classes
-        for mod in self._group_obs_class_modifiers:
+        for mod in self._group_obs_class_instances:
             mod.reset(env_ids=env_ids)
 
         # nothing to log here
@@ -320,8 +320,10 @@ def compute_group(self, group_name: str) -> torch.Tensor | dict[str, torch.Tenso
             if term_cfg.modifiers is not None:
                 for modifier in term_cfg.modifiers:
                     obs = modifier.func(obs, **modifier.params)
-            if term_cfg.noise:
+            if isinstance(term_cfg.noise, noise.NoiseCfg):
                 obs = term_cfg.noise.func(obs, term_cfg.noise)
+            elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None:
+                obs = term_cfg.noise.func(obs)
             if term_cfg.clip:
                 obs = obs.clip_(min=term_cfg.clip[0], max=term_cfg.clip[1])
             if term_cfg.scale is not None:
@@ -384,9 +386,9 @@ def _prepare_terms(self):
         self._group_obs_concatenate_dim: dict[str, int] = dict()
 
         self._group_obs_term_history_buffer: dict[str, dict] = dict()
-        # create a list to store modifiers that are classes
+        # create a list to store classes instances, e.g., for modifiers and noise models
         # we store it as a separate list to only call reset on them and prevent unnecessary calls
-        self._group_obs_class_modifiers: list[modifiers.ModifierBase] = list()
+        self._group_obs_class_instances: list[modifiers.ModifierBase | noise.NoiseModel] = list()
 
         # make sure the simulation is playing since we compute obs dims which needs asset quantities
         if not self._env.sim.is_playing():
@@ -497,7 +499,7 @@ def _prepare_terms(self):
                                 mod_cfg.func = mod_cfg.func(cfg=mod_cfg, data_dim=obs_dims, device=self._env.device)
 
                                 # add to list of class modifiers
-                                self._group_obs_class_modifiers.append(mod_cfg.func)
+                                self._group_obs_class_instances.append(mod_cfg.func)
                         else:
                             raise TypeError(
                                 f"Modifier configuration '{mod_cfg}' of observation term '{term_name}' is not of"
@@ -527,6 +529,20 @@ def _prepare_terms(self):
                                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                                 )
 
+                # prepare noise model classes
+                if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg):
+                    noise_model_cls = term_cfg.noise.class_type
+                    if not issubclass(noise_model_cls, noise.NoiseModel):
+                        raise TypeError(
+                            f"Class type for observation term '{term_name}' NoiseModelCfg"
+                            f" is not a subclass of 'NoiseModel'. Received: '{type(noise_model_cls)}'."
+                        )
+                    # initialize func to be the noise model class instance
+                    term_cfg.noise.func = noise_model_cls(
+                        term_cfg.noise, num_envs=self._env.num_envs, device=self._env.device
+                    )
+                    self._group_obs_class_instances.append(term_cfg.noise.func)
+
                 # create history buffers and calculate history term dimensions
                 if term_cfg.history_length > 0:
                     group_entry_history_buffer[term_name] = CircularBuffer(

source/isaaclab/isaaclab/utils/noise/noise_cfg.py

@@ -78,6 +78,19 @@ class NoiseModelCfg:
     noise_cfg: NoiseCfg = MISSING
     """The noise configuration to use."""
 
+    func: Callable[[torch.Tensor], torch.Tensor] | None = None
+    """Function or callable class used by this noise model.
+
+    The function must take a single `torch.Tensor` (the batch of observations) as input
+    and return a `torch.Tensor` of the same shape with noise applied.
+
+    It also supports `callable classes <https://docs.python.org/3/reference/datamodel.html#object.__call__>`_,
+    i.e. classes that implement the ``__call__()`` method. In this case, the class should inherit from the
+    :class:`NoiseModel` class and implement the required methods.
+
+    This field is used internally by :class:ObservationManager and is not meant to be set directly.
+    """
+
 
 @configclass
 class NoiseModelWithAdditiveBiasCfg(NoiseModelCfg):

source/isaaclab/isaaclab/utils/noise/noise_model.py

@@ -130,7 +130,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         """
         pass
 
-    def apply(self, data: torch.Tensor) -> torch.Tensor:
+    def __call__(self, data: torch.Tensor) -> torch.Tensor:
         """Apply the noise to the data.
 
         Args:
@@ -170,7 +170,7 @@ def reset(self, env_ids: Sequence[int] | None = 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:
+    def __call__(self, data: torch.Tensor) -> torch.Tensor:
         """Apply bias noise to the data.
 
         Args:
@@ -179,4 +179,4 @@ def apply(self, data: torch.Tensor) -> torch.Tensor:
         Returns:
             The data with the noise applied. Shape is the same as the input data.
         """
-        return super().apply(data) + self._bias
+        return super().__call__(data) + self._bias