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]
# Note: Semantic Versioning is used: https://semver.org/
version = "0.40.6"
version = "0.40.7"
# Description
title = "Isaac Lab framework for Robot Learning"
......
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)
~~~~~~~~~~~~~~~~~~~
......
......@@ -356,7 +356,7 @@ class DirectMARLEnv(gym.Env):
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 @@ class DirectMARLEnv(gym.Env):
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
......
......@@ -329,7 +329,7 @@ class DirectRLEnv(gym.Env):
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 @@ class DirectRLEnv(gym.Env):
# 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
......
......@@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any
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
......
......@@ -14,7 +14,7 @@ from collections.abc import Sequence
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 @@ class ObservationManager(ManagerBase):
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 @@ class ObservationManager(ManagerBase):
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 @@ class ObservationManager(ManagerBase):
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 @@ class ObservationManager(ManagerBase):
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 @@ class ObservationManager(ManagerBase):
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(
......
......@@ -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):
......
......@@ -130,7 +130,7 @@ class NoiseModel:
"""
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 @@ class NoiseModelWithAdditiveBias(NoiseModel):
# 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 @@ class NoiseModelWithAdditiveBias(NoiseModel):
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
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