Adds recorder manager in manager-based environments

.gitattributes

@@ -9,3 +9,4 @@
 *.mp4 filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.jit filter=lfs diff=lfs merge=lfs -text
+*.hdf5 filter=lfs diff=lfs merge=lfs -text

CONTRIBUTORS.md

@@ -41,6 +41,7 @@ Guidelines for modifications:
 * Brayden Zhang
 * Calvin Yu
 * Chenyu Yang
+* CY (Chien-Ying) Chen
 * David Yang
 * Dorsa Rohani
 * Felix Yu

source/extensions/omni.isaac.lab/config/extension.toml

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

source/extensions/omni.isaac.lab/docs/CHANGELOG.rst

@@ -1,6 +1,18 @@
 Changelog
 ---------
 
+0.27.17 (2024-12-02)
+
+Added
+
+* Added :class:`~omni.isaac.lab.managers.RecorderManager` and its utility classes to record data from the simulation.
+* Added :class:`~omni.isaac.lab.utils.datasets.EpisodeData` to store data for an episode.
+* Added :class:`~omni.isaac.lab.utils.datasets.DatasetFileHandlerBase` as a base class for handling dataset files.
+* Added :class:`~omni.isaac.lab.utils.datasets.HDF5DatasetFileHandler` as a dataset file handler implementation to
+  export and load episodes from HDF5 files.
+* Added ``record_demos.py`` script to record human-teleoperated demos for a specified task and export to an HDF5 file.
+* Added ``replay_demos.py`` script to replay demos loaded from an HDF5 file.
+
 0.27.16 (2024-11-21)
 ~~~~~~~~~~~~~~~~~~~~
 

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

@@ -11,7 +11,7 @@
 import omni.isaac.core.utils.torch as torch_utils
 import omni.log
 
-from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager
+from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager, RecorderManager
 from omni.isaac.lab.scene import InteractiveScene
 from omni.isaac.lab.sim import SimulationContext
 from omni.isaac.lab.utils.timer import Timer
@@ -45,6 +45,9 @@ class ManagerBasedEnv:
       This includes resetting the scene to a default state, applying random pushes to the robot at different intervals
       of time, or randomizing properties such as mass and friction coefficients. This is useful for training
       and evaluating the robot in a variety of scenarios.
+    * **Recorder Manager**: The recorder manager that handles recording data produced during different steps
+      in the simulation. This includes recording in the beginning and end of a reset and a step. The recorded data
+      is distinguished per episode, per environment and can be exported through a dataset file handler to a file.
 
     The environment provides a unified interface for interacting with the simulation. However, it does not
     include task-specific quantities such as the reward function, or the termination conditions. These
@@ -153,6 +156,9 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         # allocate dictionary to store metrics
         self.extras = {}
 
+        # initialize observation buffers
+        self.obs_buf = {}
+
     def __del__(self):
         """Cleanup for the environment."""
         self.close()
@@ -208,6 +214,9 @@ def load_managers(self):
 
         """
         # prepare the managers
+        # -- recorder manager
+        self.recorder_manager = RecorderManager(self.cfg.recorders, self)
+        print("[INFO] Recorder Manager: ", self.recorder_manager)
         # -- action manager
         self.action_manager = ActionManager(self.cfg.actions, self)
         print("[INFO] Action Manager: ", self.action_manager)
@@ -228,15 +237,18 @@ def load_managers(self):
     Operations - MDP.
     """
 
-    def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]:
-        """Resets all the environments and returns observations.
+    def reset(
+        self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[VecEnvObs, dict]:
+        """Resets the specified environments and returns observations.
 
-        This function calls the :meth:`_reset_idx` function to reset all the environments.
+        This function calls the :meth:`_reset_idx` function to reset the specified environments.
         However, certain operations, such as procedural terrain generation, that happened during initialization
         are not repeated.
 
         Args:
             seed: The seed to use for randomization. Defaults to None, in which case the seed is not set.
+            env_ids: The environment ids to reset. Defaults to None, in which case all environments are reset.
             options: Additional information to specify how the environment is reset. Defaults to None.
 
                 Note:
@@ -245,20 +257,78 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
         Returns:
             A tuple containing the observations and extras.
         """
+        if env_ids is None:
+            env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
+
+        # trigger recorder terms for pre-reset calls
+        self.recorder_manager.record_pre_reset(env_ids)
+
         # set the seed
         if seed is not None:
             self.seed(seed)
 
         # reset state of scene
-        indices = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
-        self._reset_idx(indices)
+        self._reset_idx(env_ids)
+        self.scene.write_data_to_sim()
+
+        # trigger recorder terms for post-reset calls
+        self.recorder_manager.record_post_reset(env_ids)
 
         # if sensors are added to the scene, make sure we render to reflect changes in reset
         if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
             self.sim.render()
 
+        # compute observations
+        self.obs_buf = self.observation_manager.compute()
+
         # return observations
-        return self.observation_manager.compute(), self.extras
+        return self.obs_buf, self.extras
+
+    def reset_to(
+        self,
+        state: dict[str, dict[str, dict[str, torch.Tensor]]],
+        env_ids: Sequence[int] | None,
+        seed: int | None = None,
+        is_relative: bool = False,
+    ) -> None:
+        """Resets specified environments to known states.
+
+        Note that this is different from reset() function as it resets the environments to specific states
+
+        Args:
+            state: The state to reset the specified environments to.
+            env_ids: The environment ids to reset. Defaults to None, in which case all environments are reset.
+            seed: The seed to use for randomization. Defaults to None, in which case the seed is not set.
+            is_relative: If set to True, the state is considered relative to the environment origins. Defaults to False.
+        """
+        # reset all envs in the scene if env_ids is None
+        if env_ids is None:
+            env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
+
+        # trigger recorder terms for pre-reset calls
+        self.recorder_manager.record_pre_reset(env_ids)
+
+        # set the seed
+        if seed is not None:
+            self.seed(seed)
+
+        self._reset_idx(env_ids)
+
+        # set the state
+        self.scene.reset_to(state, env_ids, is_relative=is_relative)
+
+        # trigger recorder terms for post-reset calls
+        self.recorder_manager.record_post_reset(env_ids)
+
+        # if sensors are added to the scene, make sure we render to reflect changes in reset
+        if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
+            self.sim.render()
+
+        # compute observations
+        self.obs_buf = self.observation_manager.compute()
+
+        # return observations
+        return self.obs_buf, self.extras
 
     def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         """Execute one time-step of the environment's dynamics.
@@ -278,6 +348,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         # process actions
         self.action_manager.process_action(action.to(self.device))
 
+        self.recorder_manager.record_pre_step()
+
         # check if we need to do rendering within the physics loop
         # note: checked here once to avoid multiple checks within the loop
         is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors()
@@ -303,8 +375,12 @@ 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.obs_buf = self.observation_manager.compute()
+        self.recorder_manager.record_post_step()
+
         # return observations and extras
-        return self.observation_manager.compute(), self.extras
+        return self.obs_buf, self.extras
 
     @staticmethod
     def seed(seed: int = -1) -> int:
@@ -334,6 +410,7 @@ def close(self):
             del self.action_manager
             del self.observation_manager
             del self.event_manager
+            del self.recorder_manager
             del self.scene
             # clear callbacks and instance
             self.sim.clear_all_callbacks()
@@ -375,3 +452,6 @@ def _reset_idx(self, env_ids: Sequence[int]):
         # -- event manager
         info = self.event_manager.reset(env_ids)
         self.extras["log"].update(info)
+        # -- recorder manager
+        info = self.recorder_manager.reset(env_ids)
+        self.extras["log"].update(info)

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

@@ -13,6 +13,7 @@
 
 import omni.isaac.lab.envs.mdp as mdp
 from omni.isaac.lab.managers import EventTermCfg as EventTerm
+from omni.isaac.lab.managers import RecorderManagerBaseCfg as DefaultEmptyRecorderManagerCfg
 from omni.isaac.lab.scene import InteractiveSceneCfg
 from omni.isaac.lab.sim import SimulationCfg
 from omni.isaac.lab.utils import configclass
@@ -78,6 +79,12 @@ class ManagerBasedEnvCfg:
     Please refer to the :class:`omni.isaac.lab.scene.InteractiveSceneCfg` class for more details.
     """
 
+    recorders: object = DefaultEmptyRecorderManagerCfg()
+    """Recorder settings. Defaults to recording nothing.
+
+    Please refer to the :class:`omni.isaac.lab.managers.RecorderManager` class for more details.
+    """
+
     observations: object = MISSING
     """Observation space settings.
 

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

@@ -158,6 +158,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         # process actions
         self.action_manager.process_action(action.to(self.device))
 
+        self.recorder_manager.record_pre_step()
+
         # check if we need to do rendering within the physics loop
         # note: checked here once to avoid multiple checks within the loop
         is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors()
@@ -190,14 +192,29 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         # -- reward computation
         self.reward_buf = self.reward_manager.compute(dt=self.step_dt)
 
+        if len(self.recorder_manager.active_terms) > 0:
+            # update observations for recording if needed
+            self.obs_buf = self.observation_manager.compute()
+            self.recorder_manager.record_post_step()
+
         # -- reset envs that terminated/timed-out and log the episode information
         reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
         if len(reset_env_ids) > 0:
+            # trigger recorder terms for pre-reset calls
+            self.recorder_manager.record_pre_reset(reset_env_ids)
+
             self._reset_idx(reset_env_ids)
+
+            # this is needed to make joint positions set from reset events effective
+            self.scene.write_data_to_sim()
+
             # if sensors are added to the scene, make sure we render to reflect changes in reset
             if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
                 self.sim.render()
 
+            # trigger recorder terms for post-reset calls
+            self.recorder_manager.record_post_reset(reset_env_ids)
+
         # -- update command
         self.command_manager.compute(dt=self.step_dt)
         # -- step interval events
@@ -353,6 +370,9 @@ def _reset_idx(self, env_ids: Sequence[int]):
         # -- termination manager
         info = self.termination_manager.reset(env_ids)
         self.extras["log"].update(info)
+        # -- recorder manager
+        info = self.recorder_manager.reset(env_ids)
+        self.extras["log"].update(info)
 
         # reset the episode length buffer
         self.episode_length_buf[env_ids] = 0

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

@@ -20,5 +20,6 @@
 from .curriculums import *  # noqa: F401, F403
 from .events import *  # noqa: F401, F403
 from .observations import *  # noqa: F401, F403
+from .recorders import *  # noqa: F401, F403
 from .rewards import *  # noqa: F401, F403
 from .terminations import *  # noqa: F401, F403

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

@@ -0,0 +1,9 @@
+# Copyright (c) 2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Various recorder terms that can be used in the environment."""
+
+from .recorders import *
+from .recorders_cfg import *

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

@@ -0,0 +1,45 @@
+# Copyright (c) 2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from omni.isaac.lab.managers.recorder_manager import RecorderTerm
+
+
+class InitialStateRecorder(RecorderTerm):
+    """Recorder term that records the initial state of the environment after reset."""
+
+    def record_post_reset(self, env_ids: Sequence[int] | None):
+        return "initial_state", self._env.scene.get_state(is_relative=True)
+
+
+class PostStepStatesRecorder(RecorderTerm):
+    """Recorder term that records the state of the environment at the end of each step."""
+
+    def record_post_step(self):
+        return "states", self._env.scene.get_state(is_relative=True)
+
+
+class PreStepActionsRecorder(RecorderTerm):
+    """Recorder term that records the actions in the beginning of each step."""
+
+    def record_pre_step(self):
+        return "actions", self._env.action_manager.action
+
+
+class PreStepFlatPolicyObservationsRecorder(RecorderTerm):
+    """Recorder term that records the policy group observations in each step."""
+
+    def record_pre_step(self):
+        return "obs", self._env.obs_buf["policy"]
+
+
+class PreStepSubtaskTermsObservationsRecorder(RecorderTerm):
+    """Recorder term that records the subtask completion observations in each step."""
+
+    def record_pre_step(self):
+        return "obs/subtask_term_signals", self._env.obs_buf["subtask_terms"]