Adds recorder manager in manager-based environments #1336

nvcyc · 2024-10-29T22:40:17Z

Description

This PR adds a recorder manager (RecorderManager) and relevant utility classes for recording data produced in various reset and step stages in manager-based environments.

Wither the built-in recorder manager, users can create custom recorder terms in their environment configurations with callback functions returning tensors to be recorded as environments advance. It is particularly useful for implementing an app that collects human-operated demos and for those who want to record robot actions for post-validation/replay in Isaac Lab environments.

The recorder manager works in both single- and multi-environment use cases. An episode for an environment instance is exported to a dataset file, via a dataset file handler, upon completion (a termination term is signaled a reset to the environment instance is triggered).

By default, the recorder manager is inactive (by assigning no recorder terms in the default configuration), which should have minimal performance impact for existing apps that do not require data recording.

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update -- to be updated in later PRs

Checklist

I have run the pre-commit checks with ./isaaclab.sh --format
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
I have updated the changelog and the corresponding version in the extension's config/extension.toml file
I have added my name to the CONTRIBUTORS.md or my name already exists there

kellyguo11 · 2024-11-10T10:40:40Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

@@ -303,8 +377,13 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
        if "interval" in self.event_manager.available_modes:
            self.event_manager.apply(mode="interval", dt=self.step_dt)

+        # -- compute observations
+        self.prev_obs_buf = self.obs_buf


would obs_buf need to be cloned?

Where does prev_obs_buf get used? I only see it set.

kellyguo11 · 2024-11-10T11:49:53Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

+        # -- update command
+        self.command_manager.compute(dt=self.step_dt)
+        # -- step interval events
+        if "interval" in self.event_manager.available_modes:


not 100% sure if it's ok to do randomization before reset. could there be cases where the events could trigger resets?

kellyguo11 · 2024-11-10T11:50:18Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

+        if "interval" in self.event_manager.available_modes:
+            self.event_manager.apply(mode="interval", dt=self.step_dt)
+
+        self.prev_obs_buf = self.obs_buf


this could introduce additional observation computation if recording is not needed

kellyguo11 · 2024-11-10T12:06:23Z

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py

+    * Pre-reset recording: This callback is invoked at the beginning of `env.reset()` before the reset is effective.
+    * Post-reset recording: This callback is invoked at the end of `env.reset()`.
+    * Pre-step recording: This callback is invoked at the beginning of `env.step()`, after the step action is processed
+          and before the action is applied the action manager.


applied "to" the action manager?

kellyguo11 · 2024-11-10T17:20:45Z

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py

+        for env_id in range(env.num_envs):
+            self._episodes[env_id] = EpisodeData()
+
+        if hasattr(env.cfg, "env_name"):


we can use getattr with default None

kellyguo11 · 2024-11-10T18:08:30Z

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

+            env_ids = slice(None)
+        # -- assets
+        for asset_name, articulation in self._articulations.items():
+            joint_position = torch.Tensor(state["articulation"][asset_name])


why does this need to be wrapped in torch.Tensor? and if so, we would also need to specify the device?

kellyguo11 · 2024-11-11T14:25:28Z

source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py

+
+        # set up the environment
+        self.num_envs = 20
+        self.device = "cpu"


can device only be cpu? should we test for GPU case as well?

agreed, ideally with subTest, but this will need to be set in each test_ method

Second this, ideally use subtest within each test_ method below.

kellyguo11 · 2024-11-11T14:26:07Z

source/extensions/omni.isaac.lab/test/utils/test_episode_data.py

+        episode = EpisodeData()
+        self.assertTrue(episode.is_empty())
+
+        episode.add("key", torch.Tensor([1, 2, 3], device="cpu"))


similar comment about device

kellyguo11 · 2024-11-17T13:33:14Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

+
+        # set the state
+        self.scene.reset_to(state, env_ids, is_relative=is_relative)
+        self.scene.write_data_to_sim()


reset_to seems to already call write_data_to_sim, is it necessary to call this again?

kellyguo11 · 2024-11-17T13:59:18Z

source/standalone/tools/record_demos.py

+        return torch.concat([delta_pose, gripper_vel], dim=1)
+
+
+def main():


this seems to be similar to the collect_demonstrations.py script, are there key differences for us to keep both scripts?

kellyguo11 · 2024-11-17T14:22:12Z

test_episode_data.py test seems to be failing

jtigue-bdai · 2024-11-18T14:14:30Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

@@ -303,8 +377,13 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
        if "interval" in self.event_manager.available_modes:
            self.event_manager.apply(mode="interval", dt=self.step_dt)

+        # -- compute observations
+        self.prev_obs_buf = self.obs_buf


Where does prev_obs_buf get used? I only see it set.

jtigue-bdai · 2024-11-18T14:35:07Z

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

+        state = dict()
+        # -- assets
+        state["articulation"] = dict()
+        for asset_name, articulation in self._articulations.items():


Also velocity, and applied torque

jtigue-bdai · 2024-11-18T14:40:27Z

source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py

+
+        # set up the environment
+        self.num_envs = 20
+        self.device = "cpu"


agreed, ideally with subTest, but this will need to be set in each test_ method

jtigue-bdai · 2024-11-18T14:53:36Z

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py

+    def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv) -> None:
+        super().__init__(cfg, env)


NIT: This is unnecessary if you are not adding to the init of the parent class. Can be removed.

jtigue-bdai · 2024-11-18T14:55:56Z

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py

+    def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]:
+        """Record data at the beginning of env.reset() before reset is effective."""
+        return None, None
+
+    def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]:
+        """Record data at the end of env.reset()."""
+        return None, None
+
+    def record_pre_step(self) -> tuple[str | None, torch.Tensor | None]:
+        """Record data in the beginning of env.step() after action is cached/processed in the ActionManager."""
+        return None, None
+
+    def record_post_step(self) -> tuple[str | None, torch.Tensor | None]:
+        """Record data at the end of env.step() when all the managers are processed."""
+        return None, None


Can you add a description of the returns in the docstrings of these functions.

jtigue-bdai · 2024-11-18T14:59:11Z

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

+    def get_state(self, is_relative: bool = False) -> dict[str, dict[str, torch.Tensor]]:
+        """Returns the state of the scene entities.
+
+        Returns:
+            A dictionary of the state of the scene entities.
+        """


can you describe the is_relative in the docstring? It could be interpreted as relative joint position and it looks like its just relative to environment origins.

jtigue-bdai · 2024-11-18T15:09:33Z

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

+            articulation.write_joint_state_to_sim(joint_position, torch.zeros_like(joint_position), env_ids=env_ids)
+            articulation.set_joint_position_target(joint_position, env_ids=env_ids)
+            articulation.set_joint_velocity_target(torch.zeros_like(joint_position), env_ids=env_ids)


as with the comments in scene.get_state we will need to set velocity and effort here too.,

jtigue-bdai · 2024-11-18T15:10:37Z

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

+        state = dict()
+        # -- assets
+        state["articulation"] = dict()
+        for asset_name, articulation in self._articulations.items():


it would be good to get velocity and applied_torque too for more general use cases.

jtigue-bdai · 2024-11-18T15:12:33Z

source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py

+
+        # set up the environment
+        self.num_envs = 20
+        self.device = "cpu"


Second this, ideally use subtest within each test_ method below.

jtigue-bdai · 2024-11-18T15:17:08Z

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py

+        self.export_episodes(env_ids)
+
+    def record_post_reset(self, env_ids: Sequence[int] | None) -> None:
+        """Trigger recorder terms for post-reset functions."""


add env_ids to docstring

nvcyc requested review from Mayankm96, jsmith-bdai, Dhoeller19 and kellyguo11 as code owners October 29, 2024 22:40

nvcyc force-pushed the cy/recorder_manager branch 2 times, most recently from 37b6d9c to 62087f0 Compare November 2, 2024 19:08

nvcyc force-pushed the cy/recorder_manager branch 6 times, most recently from 0e16d99 to 18f8f71 Compare November 7, 2024 23:58

jsmith-bdai requested a review from jtigue-bdai November 14, 2024 14:32

kellyguo11 reviewed Nov 17, 2024

View reviewed changes

jtigue-bdai reviewed Nov 18, 2024

View reviewed changes

nvcyc added 5 commits November 19, 2024 16:20

Adds recorder manager in manager-based environments

2b2214b

Fixes add and get_next APIs in EpisodeData

e29f4b3

Removes unused prev_obs_buf in manager based env classes

ce1cda3

Simplifies getting env_name from cfg in RecorderManager

e9929f3

Updates docstrings in RecorderManager

d004d77

nvcyc force-pushed the cy/recorder_manager branch from acf1518 to d004d77 Compare November 20, 2024 00:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adds recorder manager in manager-based environments #1336

Adds recorder manager in manager-based environments #1336

nvcyc commented Oct 29, 2024

kellyguo11 Nov 10, 2024

jtigue-bdai Nov 18, 2024

kellyguo11 Nov 10, 2024

kellyguo11 Nov 10, 2024

kellyguo11 Nov 10, 2024

kellyguo11 Nov 10, 2024

kellyguo11 Nov 10, 2024

kellyguo11 Nov 11, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

kellyguo11 Nov 11, 2024

kellyguo11 Nov 17, 2024

kellyguo11 Nov 17, 2024

kellyguo11 commented Nov 17, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

jtigue-bdai Nov 18, 2024

		return torch.concat([delta_pose, gripper_vel], dim=1)


		def main():

		def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv) -> None:
		super().__init__(cfg, env)

Adds recorder manager in manager-based environments #1336

Are you sure you want to change the base?

Adds recorder manager in manager-based environments #1336

Conversation

nvcyc commented Oct 29, 2024

Description

Type of change

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kellyguo11 commented Nov 17, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment