diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index cf769e34cfb..b4719ba9424 100644 --- a/source/isaaclab/config/extension.toml +++ b/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" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index 7dfe075a3a4..ff16dfcd009 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/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) ~~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 1febf07d70a..226a93facd3 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/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 diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index 4b6ba98f1e1..40fa3bd175d 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/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. """ diff --git a/source/isaaclab/test/managers/test_recorder_manager.py b/source/isaaclab/test/managers/test_recorder_manager.py index 42c1b47e1a1..20e05321895 100644 --- a/source/isaaclab/test/managers/test_recorder_manager.py +++ b/source/isaaclab/test/managers/test_recorder_manager.py @@ -14,6 +14,7 @@ """Rest everything follows.""" +import h5py import os import shutil import tempfile @@ -21,14 +22,20 @@ 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