Unverified Commit ea717fa5 authored by Özhan Özen's avatar Özhan Özen Committed by GitHub

Integrates `NoiseModel` to manager-based workflows (#2755)

# Description

This PR adds `NoiseModel` support for manager-based workflows. To
achieve this, I have:
- Added `NoiseModel` lifecycle management to `ObservationManager`.
- Added a `Callable` field, `func`, to `NoiseModelCfg`, which
`ObservationManager` uses to assign the class instance within, similar
to how it is done for `ModifierBase`.
- Renamed `apply()` to be `__call()__`, to be consistent with
function-based noises and `ModifierBase`.

Fixes #2715 and #1864.

Note: I left the changelog with the entry [Unreleased] until the PR is
given the green light.

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there
parent ad14a674
[package] [package]
# Note: Semantic Versioning is used: https://semver.org/ # Note: Semantic Versioning is used: https://semver.org/
version = "0.40.6" version = "0.40.7"
# Description # Description
title = "Isaac Lab framework for Robot Learning" title = "Isaac Lab framework for Robot Learning"
......
Changelog Changelog
--------- ---------
0.40.7 (2025-06-24)
~~~~~~~~~~~~~~~~~~~
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) 0.40.6 (2025-06-12)
~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~
......
...@@ -356,7 +356,7 @@ class DirectMARLEnv(gym.Env): ...@@ -356,7 +356,7 @@ class DirectMARLEnv(gym.Env):
if self.cfg.action_noise_model: if self.cfg.action_noise_model:
for agent, action in actions.items(): for agent, action in actions.items():
if agent in self._action_noise_model: 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 # process actions
self._pre_physics_step(actions) self._pre_physics_step(actions)
...@@ -409,7 +409,7 @@ class DirectMARLEnv(gym.Env): ...@@ -409,7 +409,7 @@ class DirectMARLEnv(gym.Env):
if self.cfg.observation_noise_model: if self.cfg.observation_noise_model:
for agent, obs in self.obs_dict.items(): for agent, obs in self.obs_dict.items():
if agent in self._observation_noise_model: 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 observations, rewards, resets and extras
return self.obs_dict, self.reward_dict, self.terminated_dict, self.time_out_dict, self.extras return self.obs_dict, self.reward_dict, self.terminated_dict, self.time_out_dict, self.extras
......
...@@ -329,7 +329,7 @@ class DirectRLEnv(gym.Env): ...@@ -329,7 +329,7 @@ class DirectRLEnv(gym.Env):
action = action.to(self.device) action = action.to(self.device)
# add action noise # add action noise
if self.cfg.action_noise_model: if self.cfg.action_noise_model:
action = self._action_noise_model.apply(action) action = self._action_noise_model(action)
# process actions # process actions
self._pre_physics_step(action) self._pre_physics_step(action)
...@@ -386,7 +386,7 @@ class DirectRLEnv(gym.Env): ...@@ -386,7 +386,7 @@ class DirectRLEnv(gym.Env):
# add observation noise # add observation noise
# note: we apply no noise to the state space (since it is used for critic networks) # note: we apply no noise to the state space (since it is used for critic networks)
if self.cfg.observation_noise_model: 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 observations, rewards, resets and extras
return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras
......
...@@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any ...@@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any
from isaaclab.utils import configclass from isaaclab.utils import configclass
from isaaclab.utils.modifiers import ModifierCfg 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 from .scene_entity_cfg import SceneEntityCfg
...@@ -165,7 +165,7 @@ class ObservationTermCfg(ManagerTermBaseCfg): ...@@ -165,7 +165,7 @@ class ObservationTermCfg(ManagerTermBaseCfg):
For more information on modifiers, see the :class:`~isaaclab.utils.modifiers.ModifierCfg` class. 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.""" """The noise to add to the observation. Defaults to None, in which case no noise is added."""
clip: tuple[float, float] | None = None clip: tuple[float, float] | None = None
......
...@@ -14,7 +14,7 @@ from collections.abc import Sequence ...@@ -14,7 +14,7 @@ from collections.abc import Sequence
from prettytable import PrettyTable from prettytable import PrettyTable
from typing import TYPE_CHECKING 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 isaaclab.utils.buffers import CircularBuffer
from .manager_base import ManagerBase, ManagerTermBase from .manager_base import ManagerBase, ManagerTermBase
...@@ -239,7 +239,7 @@ class ObservationManager(ManagerBase): ...@@ -239,7 +239,7 @@ class ObservationManager(ManagerBase):
if term_name in self._group_obs_term_history_buffer[group_name]: 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) self._group_obs_term_history_buffer[group_name][term_name].reset(batch_ids=env_ids)
# call all modifiers that are classes # 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) mod.reset(env_ids=env_ids)
# nothing to log here # nothing to log here
...@@ -320,8 +320,10 @@ class ObservationManager(ManagerBase): ...@@ -320,8 +320,10 @@ class ObservationManager(ManagerBase):
if term_cfg.modifiers is not None: if term_cfg.modifiers is not None:
for modifier in term_cfg.modifiers: for modifier in term_cfg.modifiers:
obs = modifier.func(obs, **modifier.params) 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) 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: if term_cfg.clip:
obs = obs.clip_(min=term_cfg.clip[0], max=term_cfg.clip[1]) obs = obs.clip_(min=term_cfg.clip[0], max=term_cfg.clip[1])
if term_cfg.scale is not None: if term_cfg.scale is not None:
...@@ -384,9 +386,9 @@ class ObservationManager(ManagerBase): ...@@ -384,9 +386,9 @@ class ObservationManager(ManagerBase):
self._group_obs_concatenate_dim: dict[str, int] = dict() self._group_obs_concatenate_dim: dict[str, int] = dict()
self._group_obs_term_history_buffer: dict[str, dict] = 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 # 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 # make sure the simulation is playing since we compute obs dims which needs asset quantities
if not self._env.sim.is_playing(): if not self._env.sim.is_playing():
...@@ -497,7 +499,7 @@ class ObservationManager(ManagerBase): ...@@ -497,7 +499,7 @@ class ObservationManager(ManagerBase):
mod_cfg.func = mod_cfg.func(cfg=mod_cfg, data_dim=obs_dims, device=self._env.device) mod_cfg.func = mod_cfg.func(cfg=mod_cfg, data_dim=obs_dims, device=self._env.device)
# add to list of class modifiers # 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: else:
raise TypeError( raise TypeError(
f"Modifier configuration '{mod_cfg}' of observation term '{term_name}' is not of" f"Modifier configuration '{mod_cfg}' of observation term '{term_name}' is not of"
...@@ -527,6 +529,20 @@ class ObservationManager(ManagerBase): ...@@ -527,6 +529,20 @@ class ObservationManager(ManagerBase):
f" and optional parameters: {args_with_defaults}, but received: {term_params}." 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 # create history buffers and calculate history term dimensions
if term_cfg.history_length > 0: if term_cfg.history_length > 0:
group_entry_history_buffer[term_name] = CircularBuffer( group_entry_history_buffer[term_name] = CircularBuffer(
......
...@@ -78,6 +78,19 @@ class NoiseModelCfg: ...@@ -78,6 +78,19 @@ class NoiseModelCfg:
noise_cfg: NoiseCfg = MISSING noise_cfg: NoiseCfg = MISSING
"""The noise configuration to use.""" """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 @configclass
class NoiseModelWithAdditiveBiasCfg(NoiseModelCfg): class NoiseModelWithAdditiveBiasCfg(NoiseModelCfg):
......
...@@ -130,7 +130,7 @@ class NoiseModel: ...@@ -130,7 +130,7 @@ class NoiseModel:
""" """
pass pass
def apply(self, data: torch.Tensor) -> torch.Tensor: def __call__(self, data: torch.Tensor) -> torch.Tensor:
"""Apply the noise to the data. """Apply the noise to the data.
Args: Args:
...@@ -170,7 +170,7 @@ class NoiseModelWithAdditiveBias(NoiseModel): ...@@ -170,7 +170,7 @@ class NoiseModelWithAdditiveBias(NoiseModel):
# reset the bias term # reset the bias term
self._bias[env_ids] = self._bias_noise_cfg.func(self._bias[env_ids], self._bias_noise_cfg) 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. """Apply bias noise to the data.
Args: Args:
...@@ -179,4 +179,4 @@ class NoiseModelWithAdditiveBias(NoiseModel): ...@@ -179,4 +179,4 @@ class NoiseModelWithAdditiveBias(NoiseModel):
Returns: Returns:
The data with the noise applied. Shape is the same as the input data. 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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment