Adds record at close functionality to the RecorderManager

source/isaaclab/config/extension.toml

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

source/isaaclab/docs/CHANGELOG.rst

@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.40.12 (2025-06-30)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added write to file on close to :class:`~isaaclab.manager.RecorderManager`.
+
+
 0.40.11 (2025-06-27)
 ~~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/envs/manager_based_env.py

@@ -450,6 +450,7 @@ def close(self):
             del self.action_manager
             del self.observation_manager
             del self.event_manager
+            self.recorder_manager.close()
             del self.recorder_manager
             del self.scene
             # clear callbacks and instance

source/isaaclab/isaaclab/managers/recorder_manager.py

@@ -123,6 +123,17 @@ def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]:
         """
         return None, None
 
+    def close(self, file_path: str):
+        """Finalize and "clean up" the recorder term.
+
+        This can include tasks such as appending metadata (e.g. labels) to a file
+        and properly closing any associated file handles or resources.
+
+        Args:
+            file_path: the absolute path to the file
+        """
+        pass
+
 
 class RecorderManager(ManagerBase):
     """Manager for recording data from recorder terms."""
@@ -193,12 +204,7 @@ def __str__(self) -> str:
 
     def __del__(self):
         """Destructor for recorder."""
-        # Do nothing if no active recorder terms are provided
-        if len(self.active_terms) == 0:
-            return
-
-        if self._dataset_file_handler is not None:
-            self._dataset_file_handler.close()
+        self.close()
 
         if self._failed_episode_dataset_file_handler is not None:
             self._failed_episode_dataset_file_handler.close()
@@ -456,6 +462,17 @@ def export_episodes(self, env_ids: Sequence[int] | None = None) -> None:
             if self._failed_episode_dataset_file_handler is not None:
                 self._failed_episode_dataset_file_handler.flush()
 
+    def close(self):
+        """Closes the recorder manager by exporting any remaining data to file as well as properly closes the recorder terms."""
+        # Do nothing if no active recorder terms are provided
+        if len(self.active_terms) == 0:
+            return
+        if self._dataset_file_handler is not None:
+            self.export_episodes()
+            self._dataset_file_handler.close()
+        for term in self._terms.values():
+            term.close(os.path.join(self.cfg.dataset_export_dir_path, self.cfg.dataset_filename))
+
     """
     Helper functions.
     """

source/isaaclab/test/managers/test_recorder_manager.py

@@ -14,21 +14,28 @@
 
 """Rest everything follows."""
 
+import h5py
 import os
 import shutil
 import tempfile
 import torch
 import uuid
 from collections import namedtuple
 from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
+import omni.usd
 import pytest
 
-from isaaclab.envs import ManagerBasedEnv
+from isaaclab.envs import ManagerBasedEnv, ManagerBasedEnvCfg
 from isaaclab.managers import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg
+from isaaclab.scene import InteractiveSceneCfg
 from isaaclab.sim import SimulationContext
 from isaaclab.utils import configclass
 
+if TYPE_CHECKING:
+    import numpy as np
+
 
 class DummyResetRecorderTerm(RecorderTerm):
     """Dummy recorder term that records dummy data."""
@@ -78,6 +85,72 @@ class DummyStepRecorderTermCfg(RecorderTermCfg):
     dataset_export_mode = DatasetExportMode.EXPORT_ALL
 
 
+@configclass
+class EmptyManagerCfg:
+    """Empty manager specifications for the environment."""
+
+    pass
+
+
+@configclass
+class EmptySceneCfg(InteractiveSceneCfg):
+    """Configuration for an empty scene."""
+
+    pass
+
+
+def get_empty_base_env_cfg(device: str = "cuda", num_envs: int = 1, env_spacing: float = 1.0):
+    """Generate base environment config based on device"""
+
+    @configclass
+    class EmptyEnvCfg(ManagerBasedEnvCfg):
+        """Configuration for the empty test environment."""
+
+        # Scene settings
+        scene: EmptySceneCfg = EmptySceneCfg(num_envs=num_envs, env_spacing=env_spacing)
+        # Basic settings
+        actions: EmptyManagerCfg = EmptyManagerCfg()
+        observations: EmptyManagerCfg = EmptyManagerCfg()
+        recorders: EmptyManagerCfg = EmptyManagerCfg()
+
+        def __post_init__(self):
+            """Post initialization."""
+            # step settings
+            self.decimation = 4  # env step every 4 sim steps: 200Hz / 4 = 50Hz
+            # simulation settings
+            self.sim.dt = 0.005  # sim step every 5ms: 200Hz
+            self.sim.render_interval = self.decimation  # render every 4 sim steps
+            # pass device down from test
+            self.sim.device = device
+
+    return EmptyEnvCfg()
+
+
+def get_file_contents(file_name: str, num_steps: int) -> dict[str, np.ndarray]:
+    """Retrieves the contents of the hdf5 file
+    Args:
+        file_name: absolute path to the hdf5 file
+        num_steps: number of steps taken in the environment
+    Returns:
+        dict[str, np.ndarray]: dictionary where keys are HDF5 paths and values are the corresponding data arrays.
+    """
+    data = {}
+    with h5py.File(file_name, "r") as f:
+
+        def get_data(name, obj):
+            if isinstance(obj, h5py.Dataset):
+                if "record_post_step" in name:
+                    assert obj[()].shape == (num_steps, 5)
+                elif "record_pre_step" in name:
+                    assert obj[()].shape == (num_steps, 4)
+                else:
+                    raise Exception(f"The hdf5 file contains an unexpected data path, {name}")
+                data[name] = obj[()]
+
+        f.visititems(get_data)
+    return data
+
+
 def create_dummy_env(device: str = "cpu") -> ManagerBasedEnv:
     """Create a dummy environment."""
 
@@ -122,36 +195,61 @@ def test_initialize_dataset_file(dataset_dir):
     assert os.path.exists(os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename))
 
 
-def test_record(dataset_dir):
+@pytest.mark.parametrize("device", ("cpu", "cuda"))
+def test_record(device, dataset_dir):
     """Test the recording of the data."""
-    for device in ("cuda:0", "cpu"):
-        env = create_dummy_env(device)
-        # create recorder manager
-        cfg = DummyRecorderManagerCfg()
-        cfg.dataset_export_dir_path = dataset_dir
-        cfg.dataset_filename = f"{uuid.uuid4()}.hdf5"
-        recorder_manager = RecorderManager(cfg, env)
-
-        # record the step data
-        recorder_manager.record_pre_step()
-        recorder_manager.record_post_step()
-
-        recorder_manager.record_pre_step()
-        recorder_manager.record_post_step()
-
-        # check the recorded data
-        for env_id in range(env.num_envs):
-            episode = recorder_manager.get_episode(env_id)
-            assert episode.data["record_pre_step"].shape == (2, 4)
-            assert episode.data["record_post_step"].shape == (2, 5)
-
-        # Trigger pre-reset callbacks which then export and clean the episode data
-        recorder_manager.record_pre_reset(env_ids=None)
-        for env_id in range(env.num_envs):
-            episode = recorder_manager.get_episode(env_id)
-            assert episode.is_empty()
-
-        recorder_manager.record_post_reset(env_ids=None)
-        for env_id in range(env.num_envs):
-            episode = recorder_manager.get_episode(env_id)
-            assert episode.data["record_post_reset"].shape == (1, 3)
+    env = create_dummy_env(device)
+    # create recorder manager
+    cfg = DummyRecorderManagerCfg()
+    cfg.dataset_export_dir_path = dataset_dir
+    cfg.dataset_filename = f"{uuid.uuid4()}.hdf5"
+    recorder_manager = RecorderManager(cfg, env)
+
+    # record the step data
+    recorder_manager.record_pre_step()
+    recorder_manager.record_post_step()
+
+    recorder_manager.record_pre_step()
+    recorder_manager.record_post_step()
+
+    # check the recorded data
+    for env_id in range(env.num_envs):
+        episode = recorder_manager.get_episode(env_id)
+        assert episode.data["record_pre_step"].shape == (2, 4)
+        assert episode.data["record_post_step"].shape == (2, 5)
+
+    # Trigger pre-reset callbacks which then export and clean the episode data
+    recorder_manager.record_pre_reset(env_ids=None)
+    for env_id in range(env.num_envs):
+        episode = recorder_manager.get_episode(env_id)
+        assert episode.is_empty()
+
+    recorder_manager.record_post_reset(env_ids=None)
+    for env_id in range(env.num_envs):
+        episode = recorder_manager.get_episode(env_id)
+        assert episode.data["record_post_reset"].shape == (1, 3)
+
+
+@pytest.mark.parametrize("device", ("cpu", "cuda"))
+def test_close(device, dataset_dir):
+    """Test whether data is correctly exported in the close function fully integrated with ManagerBasedEnv."""
+    # create a new stage
+    omni.usd.get_context().new_stage()
+    # create environment
+    env_cfg = get_empty_base_env_cfg(device=device, num_envs=2)
+    cfg = DummyRecorderManagerCfg()
+    cfg.dataset_export_dir_path = dataset_dir
+    cfg.dataset_filename = f"{uuid.uuid4()}.hdf5"
+    env_cfg.recorders = cfg
+    env = ManagerBasedEnv(cfg=env_cfg)
+    num_steps = 3
+    for _ in range(num_steps):
+        act = torch.randn_like(env.action_manager.action)
+        obs, ext = env.step(act)
+    # check contents of hdf5 file
+    file_name = f"{env_cfg.recorders.dataset_export_dir_path}/{env_cfg.recorders.dataset_filename}"
+    data_pre_close = get_file_contents(file_name, num_steps)
+    assert len(data_pre_close) == 0
+    env.close()
+    data_post_close = get_file_contents(file_name, num_steps)
+    assert len(data_post_close.keys()) == 2 * env_cfg.scene.num_envs