diff --git a/docs/source/api/lab/isaaclab.managers.rst b/docs/source/api/lab/isaaclab.managers.rst index b4cd0fc7dfcc..4cef8faf6226 100644 --- a/docs/source/api/lab/isaaclab.managers.rst +++ b/docs/source/api/lab/isaaclab.managers.rst @@ -228,9 +228,40 @@ Curriculum Manager :members: :exclude-members: __init__ +.. _recorder-capture: + Recorder Manager ---------------- +Tasks configure recorder terms to capture the state, controls and other fields their consumers need. +Terms run at pre/post step and reset callbacks; reset callbacks receive slices or indices directly. +Consumers own subsequent processing, replay and learning. + +Returning a field path and a batched tensor (or nested dictionary of tensors) retains samples. +For online use without retention, publish the batch on the term and return ``(None, None)``. +Consumers access it through ``env.recorder_manager.get_term(name)`` and copy data before the next capture. +``EXPORT_NONE`` disables file export only; disabling the term removes capture itself. + +Fields may have different sample rates. ``buffer_initial_capacity`` sets the initial samples per world +and field; choose it with payload size and episode length in mind. Storage grows without dropping +samples and is reused across resets. + +**Observation migration:** RL post-step callbacks run after rewards and terminations, before reset +and observation updates. ``env.obs_buf`` still contains the preceding policy input. Read physical +state from the scene or call ``env.observation_manager.compute_group("recording", update_history=False)`` +for current observations. Use a separate noiseless group without history or delay to avoid stale +measurements and additional policy noise draws. History is not advanced, but noise and modifiers still +run. Non-RL environments update their normal observations before post-step capture. + +**Storage migration:** ``get_episode(env_id)`` returns an owned tensor snapshot with a leading sample +dimension; reading copies data and may synchronize the device. Modify live recording through +``add_to_episodes`` and ``set_success_to_episodes``, not the snapshot. ``EpisodeData.add()`` and +``pre_export()`` are removed: collect through the manager, or assign completed tensors to +``episode.data``. HDF5 loading and playback keep the same tensor format. + +The unused ``RecorderTerm.close(file_path)`` hook is removed. Dataset handlers finalize files; +applications explicitly close any additional resources they own. + .. autoclass:: RecorderManager :members: :inherited-members: diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index e10f8a7b507b..4fa081b8229f 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -512,7 +512,6 @@ def annotate_episode_in_manual_mode( print(f'\tThe episode will be replayed again for re-marking subtask signals for the eef "{eef_name}".\n') - annotated_episode = env.recorder_manager.get_episode(0) for ( subtask_term_signal_name, subtask_term_signal_action_index, @@ -520,7 +519,9 @@ def annotate_episode_in_manual_mode( # subtask termination signal is false until subtask is complete, and true afterwards subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_term_signal_action_index] = False - annotated_episode.add(f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals) + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals.unsqueeze(0), [0] + ) if args_cli.annotate_subtask_start_signals: for ( @@ -529,8 +530,8 @@ def annotate_episode_in_manual_mode( ) in subtask_start_signal_action_indices.items(): subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_start_signal_action_index] = False - annotated_episode.add( - f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals.unsqueeze(0), [0] ) return True diff --git a/source/isaaclab/changelog.d/recorder-capture.changed.rst b/source/isaaclab/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..f44e5272497a --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,22 @@ +* Replaced per-world recorder episode containers with batched field allocations and + independent insertion indices, reusing storage across resets and growing without dropping samples. +* Added :meth:`~isaaclab.managers.RecorderManager.get_term` for direct access to captured batches. +* Preserve slices through recorder reset callbacks and normalize Python sequences to lists so tuple + environment indices work for capture, reset and success updates. +* **Breaking:** Removed the implicit observation computation before RL post-step recorder callbacks. + Physical capture no longer consumes extra observation noise. Custom terms needing observations + of the completed transition must call ``env.observation_manager.compute_group(name, update_history=False)`` + explicitly; ``env.obs_buf`` still contains the preceding policy input. Use a noiseless group without + history or delay for current measurements that do not consume policy noise draws. +* **Breaking:** Changed :meth:`~isaaclab.managers.RecorderManager.get_episode` to return an owned + snapshot with tensor leaves. To update live recording, use ``add_to_episodes`` and + ``set_success_to_episodes``; do not mutate the snapshot. Set ``buffer_initial_capacity`` + to control the initial allocation per field. +* **Breaking:** Removed ``EpisodeData.add()`` and ``pre_export()``. Collection now belongs to + ``RecorderManager.add_to_episodes()``; assign nested tensors with leading sample dimensions + to ``EpisodeData.data`` when constructing an episode directly. Playback and HDF5 loading remain supported. +* **Breaking:** Removed the unused ``RecorderTerm.close(file_path)`` hook. Dataset handlers own + file finalization; applications must explicitly close any additional resources they own. + +* Documented independent capture, retention and export choices, with task-selected schemas and + consumer-owned processing. Core recorder storage remains independent of task and learning packages. diff --git a/source/isaaclab/changelog.d/recorder-capture.major b/source/isaaclab/changelog.d/recorder-capture.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst new file mode 100644 index 000000000000..1302fe5c27d3 --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst @@ -0,0 +1,2 @@ +* Kept retained recorder allocations writable after leaving PyTorch inference mode, including + grown buffers and schema changes, so subsequent reset and export can finalize those episodes. diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 859cb837cd5b..a7bbe2de7e1e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -236,10 +236,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.reset_time_outs = self.termination_manager.time_outs 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() + 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).int() diff --git a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py index a9315bbca63d..55050f5f43a5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py +++ b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py @@ -14,7 +14,7 @@ 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): + def record_post_reset(self, env_ids: Sequence[int] | slice | None): def extract_env_ids_values(value): nonlocal env_ids if isinstance(value, dict): diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index f4894840afb9..e47e780f43eb 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -54,6 +54,12 @@ class RecorderManagerBaseCfg: export_in_close: bool = False """Whether to export episodes in the close call.""" + buffer_initial_capacity: int = 64 + """Initial samples per environment and field. Must be positive; storage grows as needed. + + Set to the expected episode length, including decimation callbacks, to avoid reallocations. + """ + dataset_compression: bool = True """Enable dataset compression.""" @@ -68,24 +74,15 @@ class RecorderTerm(ManagerTermBase): * 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 by the action manager. - * Post-step recording: This callback is invoked at the end of `env.step()` when all the managers are processed. + * Post-step recording: This callback captures the completed physical transition. In RL environments, + it runs after rewards and terminations, before automatic reset and the policy observation update. """ - def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv): - """Initialize the recorder term. - - Args: - cfg: The configuration object. - env: The environment instance. - """ - # call the base class constructor - super().__init__(cfg, env) - """ User-defined callbacks. """ - def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + def record_pre_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str | None, torch.Tensor | dict | None]: """Record data at the beginning of env.reset() before reset is effective. Args: @@ -100,7 +97,7 @@ def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, t """ return None, None - def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str | None, torch.Tensor | dict | None]: """Record data at the end of env.reset(). Args: @@ -122,7 +119,11 @@ def record_pre_step(self) -> tuple[str | None, torch.Tensor | dict | None]: return None, None def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]: - """Record data at the end of env.step() when all the managers are processed. + """Record the completed physical transition. + + In RL environments, this runs before automatic reset and observation updates; + ``env.obs_buf`` still holds the preceding policy input. See :ref:`recorder-capture` + for recording current observations. Non-RL environments update observations first. Returns: A tuple of key and value to be recorded. @@ -139,17 +140,6 @@ def record_post_physics_decimation_step(self) -> tuple[str | None, torch.Tensor """ 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.""" @@ -163,6 +153,8 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): """ self._term_names: list[str] = [] self._terms: dict[str, RecorderTerm] = {} + self._dataset_file_handler = None + self._failed_episode_dataset_file_handler = None # Do nothing if cfg is None or an empty dict if not cfg: @@ -177,30 +169,28 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): if not isinstance(cfg, RecorderManagerBaseCfg): raise TypeError("Configuration for the recorder manager is not of type RecorderManagerBaseCfg.") - # create episode data buffer indexed by environment id - self._episodes: dict[int, EpisodeData] = dict() - for env_id in range(env.num_envs): - self._episodes[env_id] = EpisodeData() + if cfg.buffer_initial_capacity < 1: + raise ValueError("buffer_initial_capacity must be positive.") + self._fields: list[_RecorderField] = [] + self._env_ids = torch.arange(env.num_envs, device=env.device) + # -1 denotes an episode whose success has not been set. + self._success = torch.full((env.num_envs,), -1, dtype=torch.int8, device=env.device) + self._exported_counts = torch.zeros(2, dtype=torch.long, device=env.device) env_name = getattr(env.cfg, "env_name", None) - self._dataset_file_handler = None if cfg.dataset_export_mode != DatasetExportMode.EXPORT_NONE: self._dataset_file_handler = cfg.dataset_file_handler_class_type() self._dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name ) - self._failed_episode_dataset_file_handler = None if cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: self._failed_episode_dataset_file_handler = cfg.dataset_file_handler_class_type() self._failed_episode_dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name ) - self._exported_successful_episode_count = {} - self._exported_failed_episode_count = {} - def __str__(self) -> str: """Returns: A string representation for recorder manager.""" msg = f" contains {len(self._term_names)} active terms.\n" @@ -232,36 +222,14 @@ def active_terms(self) -> list[str]: return self._term_names @property - def exported_successful_episode_count(self, env_id=None) -> int: - """Number of successful episodes. - - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. - - Returns: - The number of successful episodes. - """ - if not hasattr(self, "_exported_successful_episode_count"): - return 0 - if env_id is not None: - return self._exported_successful_episode_count.get(env_id, 0) - return sum(self._exported_successful_episode_count.values()) + def exported_successful_episode_count(self) -> int: + """Number of successful episodes.""" + return int(self._exported_counts[1].item()) if self.active_terms else 0 @property - def exported_failed_episode_count(self, env_id=None) -> int: - """Number of failed episodes. - - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. - - Returns: - The number of failed episodes. - """ - if not hasattr(self, "_exported_failed_episode_count"): - return 0 - if env_id is not None: - return self._exported_failed_episode_count.get(env_id, 0) - return sum(self._exported_failed_episode_count.values()) + def exported_failed_episode_count(self) -> int: + """Number of failed episodes.""" + return int(self._exported_counts[0].item()) if self.active_terms else 0 """ Operations. @@ -280,123 +248,118 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, torch # Do nothing if no active recorder terms are provided if len(self.active_terms) == 0: return {} - # resolve environment ids - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - + if isinstance(env_ids, Sequence): + env_ids = list(env_ids) for term in self._terms.values(): term.reset(env_ids=env_ids) - - for env_id in env_ids: - self._episodes[env_id] = EpisodeData() - - # nothing to log here + self._reset_buffers(env_ids) return {} - def get_episode(self, env_id: int) -> EpisodeData: - """Returns the episode data for the given environment id. + def get_term(self, name: str) -> RecorderTerm: + """Return a configured recorder term. Args: - env_id: The environment id. + name: Name of the term. Returns: - The episode data for the given environment id. + The configured term. """ - return self._episodes.get(env_id, EpisodeData()) + return self._terms[name] + + def get_episode(self, env_id: int) -> EpisodeData: + """Return an owned snapshot of the selected environment's current episode. - def add_to_episodes(self, key: str, value: torch.Tensor | dict, env_ids: Sequence[int] | slice | None = None): - """Adds the given key-value pair to the episodes for the given environment ids. + Reading copies data and may synchronize the device. For live writes, use + :meth:`add_to_episodes` and :meth:`set_success_to_episodes`. Args: - key: The key of the given value to be added to the episodes. The key can contain nested keys - separated by '/'. For example, "obs/joint_pos" would add the given value under ['obs']['policy'] - in the underlying dictionary in the episode data. - value: The value to be added to the episodes. The value can be a tensor or a nested dictionary of tensors. - The shape of a tensor in the value is (env_ids, ...). - env_ids: The environment ids. Defaults to None, in which case all environments are considered. + env_id: The environment index. + + Returns: + An episode with tensor leaves and a leading sample dimension. """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return + episode = EpisodeData() + if not self.active_terms or not 0 <= env_id < self._env.num_envs: + return episode + for field in self._fields: + length = int(field.lengths[env_id].item()) + if length: + data = episode.data + keys = field.name.split("/") + for key in keys[:-1]: + data = data.setdefault(key, {}) + data[keys[-1]] = field.data[:length, env_id].clone() + success = int(self._success[env_id].item()) + episode.success = None if success < 0 else bool(success) + return episode + + def add_to_episodes( + self, key: str, value: torch.Tensor | dict, env_ids: Sequence[int] | slice | None = None + ) -> None: + """Append a batch to the selected episodes. - # resolve environment ids - if key is None: + Args: + key: Field path, with nested keys separated by '/'. + value: Tensor of shape (selected environments, ...) or nested dictionary of tensors. + Warp arrays are also accepted. Storage copies the values and preserves dtype/device. + A field's shape, dtype, and device must remain constant while it contains samples. + env_ids: Unique environment indices in batch order. None selects all environments. + """ + if not self.active_terms or key is None: return - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - if isinstance(value, dict): for sub_key, sub_value in value.items(): - if isinstance(sub_value, wp.array): - sub_value = wp.to_torch(sub_value) self.add_to_episodes(f"{key}/{sub_key}", sub_value, env_ids) return - if isinstance(value, wp.array): value = wp.to_torch(value) - value = value.clone() # Clone once for all envs - for value_index, env_id in enumerate(env_ids): - if env_id not in self._episodes: - self._episodes[env_id] = EpisodeData() - self._episodes[env_id].env_id = env_id - self._episodes[env_id].add(key, value[value_index], clone=False) - - def set_success_to_episodes(self, env_ids: Sequence[int] | slice | None, success_values: torch.Tensor): - """Sets the task success values to the episodes for the given environment ids. + num_selected = ( + self._env.num_envs + if env_ids is None + else (len(range(self._env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)) + ) + if value.ndim < 1 or value.shape[0] != num_selected: + raise ValueError(f"Field '{key}' requires {num_selected} rows, got {tuple(value.shape)}.") + if not num_selected: + return + field = next((field for field in self._fields if field.name == key), None) + if field is None: + # Reject ambiguous leaf/subtree paths before allocating any storage. + if any(key.startswith(f"{f.name}/") or f.name.startswith(f"{key}/") for f in self._fields): + raise ValueError(f"Recorder field '{key}' conflicts with an existing field path.") + field = _RecorderField(key, value, self._env.num_envs, self.cfg.buffer_initial_capacity) + self._fields.append(field) + field.append(value, env_ids) + + def set_success_to_episodes(self, env_ids: Sequence[int] | slice | None, success_values: torch.Tensor) -> None: + """Set success flags for selected environments. Args: - env_ids: The environment ids. Defaults to None, in which case all environments are considered. - success_values: The task success values to be set to the episodes. The shape of the tensor is (env_ids, 1). + env_ids: The environment indices. None selects all environments. + success_values: Flags in selection order, with shape (num_selected,) or (num_selected, 1). """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: + if not self.active_terms: return - - # resolve environment ids if env_ids is None: env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - - for value_index, env_id in enumerate(env_ids): - self._episodes[env_id].success = success_values[value_index].item() + elif isinstance(env_ids, Sequence): + env_ids = list(env_ids) + self._success[env_ids] = success_values.reshape(-1).to(device=self._success.device, dtype=torch.int8) def record_pre_step(self) -> None: """Trigger recorder terms for pre-step functions.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_pre_step() self.add_to_episodes(key, value) def record_post_step(self) -> None: """Trigger recorder terms for post-step functions.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_post_step() self.add_to_episodes(key, value) def record_post_physics_decimation_step(self) -> None: """Trigger recorder terms for post-physics step functions in the decimation loop.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_post_physics_decimation_step() self.add_to_episodes(key, value) @@ -413,17 +376,15 @@ def record_pre_reset(self, env_ids: Sequence[int] | slice | None, force_export_o if env_ids is None: env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() + elif isinstance(env_ids, Sequence): + env_ids = list(env_ids) for term in self._terms.values(): key, value = term.record_pre_reset(env_ids) self.add_to_episodes(key, value, env_ids) # Set task success values for the relevant episodes - success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device) + success_results = torch.zeros_like(self._success[env_ids], dtype=torch.bool) # Check success indicator from termination terms if hasattr(self._env, "termination_manager"): if "success" in self._env.termination_manager.active_terms: @@ -443,8 +404,10 @@ def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> None: if len(self.active_terms) == 0: return - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) + if env_ids is None: + env_ids = slice(None) + elif isinstance(env_ids, Sequence): + env_ids = list(env_ids) for term in self._terms.values(): key, value = term.record_post_reset(env_ids) @@ -485,11 +448,11 @@ def export_episodes( return if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() + env_ids = self._env_ids + elif isinstance(env_ids, slice): + env_ids = self._env_ids[env_ids] + else: + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self._env_ids.device) # Handle demo_ids processing if demo_ids is not None: @@ -502,58 +465,46 @@ def export_episodes( duplicates = [x for i, x in enumerate(demo_ids) if demo_ids.index(x) != i] raise ValueError(f"demo_ids must be unique. Found duplicates: {list(set(duplicates))}") - # Export episode data through dataset exporter - need_to_flush = False + present = torch.zeros(self._env.num_envs, dtype=torch.bool, device=self._env_ids.device) + for field in self._fields: + present |= (field.lengths > 0).to(present.device) + succeeded = self._success == 1 - if any(env_id in self._episodes and not self._episodes[env_id].is_empty() for env_id in env_ids): + # Only file export materializes per-world snapshots. EXPORT_NONE stays batched. + need_to_flush = False + mode = self.cfg.dataset_export_mode + if mode != DatasetExportMode.EXPORT_NONE and present[env_ids].any().item(): ep_meta = self.get_ep_meta() - if self._dataset_file_handler is not None: - self._dataset_file_handler.add_env_args(ep_meta) - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.add_env_args(ep_meta) - - for i, env_id in enumerate(env_ids): - if env_id in self._episodes and not self._episodes[env_id].is_empty(): - self._episodes[env_id].pre_export() - - episode_succeeded = self._episodes[env_id].success - target_dataset_file_handler = None - if (self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_ALL) or ( - self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and episode_succeeded - ): - target_dataset_file_handler = self._dataset_file_handler - elif self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: - if episode_succeeded: - target_dataset_file_handler = self._dataset_file_handler - else: - target_dataset_file_handler = self._failed_episode_dataset_file_handler - if target_dataset_file_handler is not None: - # Use corresponding demo_id if provided, otherwise None - current_demo_id = demo_ids[i] if demo_ids is not None else None - target_dataset_file_handler.write_episode( - self._episodes[env_id], current_demo_id, self.cfg.dataset_compression - ) - need_to_flush = True - # Update episode count - if episode_succeeded: - self._exported_successful_episode_count[env_id] = ( - self._exported_successful_episode_count.get(env_id, 0) + 1 - ) - else: - self._exported_failed_episode_count[env_id] = self._exported_failed_episode_count.get(env_id, 0) + 1 - # Reset the episode buffer for the given environment after export - self._episodes[env_id] = EpisodeData() - + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.add_env_args(ep_meta) + selected = env_ids.tolist() + selected_present = present[env_ids].tolist() + selected_success = succeeded[env_ids].tolist() + for i, env_id in enumerate(selected): + if not selected_present[i]: + continue + handler = self._dataset_file_handler + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + if not selected_success[i]: + handler = self._failed_episode_dataset_file_handler + elif mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and not selected_success[i]: + continue + episode = self.get_episode(env_id) + demo_id = demo_ids[i] if demo_ids is not None else None + handler.write_episode(episode, demo_id, self.cfg.dataset_compression) + need_to_flush = True + + self._exported_counts[0] += (present & ~succeeded)[env_ids].sum() + self._exported_counts[1] += (present & succeeded)[env_ids].sum() + self._reset_buffers(env_ids) if need_to_flush: - if self._dataset_file_handler is not None: - self._dataset_file_handler.flush() - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.flush() + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.flush() def close(self): - """Closes the recorder manager by exporting any remaining data to file as well as properly - closes the recorder terms. - """ + """Export remaining episodes if configured and close dataset files.""" # Do nothing if no active recorder terms are provided if len(self.active_terms) == 0: return @@ -563,13 +514,17 @@ def close(self): self._dataset_file_handler.close() if self._failed_episode_dataset_file_handler is not None: self._failed_episode_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. """ + def _reset_buffers(self, env_ids: Sequence[int] | torch.Tensor | slice | None) -> None: + """Discard selected samples by resetting indices; retain allocations for the next episodes.""" + for field in self._fields: + field.reset(env_ids) + self._success[slice(None) if env_ids is None else env_ids] = -1 + def _prepare_terms(self): """Prepares a list of recorder terms.""" # check if config is dict already @@ -587,6 +542,7 @@ def _prepare_terms(self): "export_in_record_pre_reset", "export_in_close", "dataset_compression", + "buffer_initial_capacity", ]: continue # check if term config is None @@ -605,3 +561,65 @@ def _prepare_terms(self): # add term name and parameters self._term_names.append(term_name) self._terms[term_name] = term + + +class _RecorderField: + """One field's time-major allocation and independent per-environment insertion indices. + + No per-environment Python containers or tensor views are created during append. The host upper + bound avoids reading device indices in steady state; only capacity exhaustion checks synchronize. + """ + + def __init__(self, name: str, value: torch.Tensor, num_envs: int, capacity: int) -> None: + self.name = name + # Retained storage outlives the caller's inference context: reset/export may run outside it. + with torch.inference_mode(False): + self.data = torch.empty((capacity, num_envs, *value.shape[1:]), dtype=value.dtype, device=value.device) + self.lengths = torch.zeros(num_envs, dtype=torch.long, device=value.device) + self.worlds = torch.arange(num_envs, device=value.device) + self.upper_bound = 0 + + @torch.no_grad() + def append(self, value: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | slice | None) -> None: + if value.shape[1:] != self.data.shape[2:] or value.dtype != self.data.dtype or value.device != self.data.device: + if self.lengths.any().item(): + raise ValueError(f"Recorder field '{self.name}' changed shape, dtype, or device with active samples.") + with torch.inference_mode(False): + self.data = torch.empty( + (*self.data.shape[:2], *value.shape[1:]), dtype=value.dtype, device=value.device + ) + self.lengths = self.lengths.to(value.device) + self.worlds = self.worlds.to(value.device) + self.upper_bound = 0 + if env_ids is None: + worlds, indices = self.worlds, self.lengths + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + worlds = self.worlds[env_ids] + indices = self.lengths[worlds] + capacity = self.data.shape[0] + if self.upper_bound >= capacity: + self.upper_bound = int(self.lengths.max().item()) + if self.upper_bound >= capacity: + with torch.inference_mode(False): + data = self.data.new_empty((2 * capacity, *self.data.shape[1:])) + data[:capacity].copy_(self.data) + self.data = data + self.data[indices, worlds] = value + if env_ids is None: + self.lengths.add_(1) + else: + self.lengths[worlds] = indices + 1 + self.upper_bound += 1 + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None) -> None: + if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): + self.lengths.zero_() + self.upper_bound = 0 + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + self.lengths[env_ids] = 0 + if not isinstance(env_ids, slice) and len(env_ids) == len(self.lengths): + self.upper_bound = 0 diff --git a/source/isaaclab/isaaclab/utils/datasets/episode_data.py b/source/isaaclab/isaaclab/utils/datasets/episode_data.py index bb528d4c630a..c401906b3aa4 100644 --- a/source/isaaclab/isaaclab/utils/datasets/episode_data.py +++ b/source/isaaclab/isaaclab/utils/datasets/episode_data.py @@ -9,7 +9,11 @@ class EpisodeData: - """Class to store episode data.""" + """An episode's tensor data and playback cursors. + + Data leaves are tensors with a leading sample dimension. Collection belongs to the recorder + manager; this container holds snapshots or loaded datasets without accumulating samples. + """ def __init__(self) -> None: """Initializes episode data class.""" @@ -85,34 +89,6 @@ def is_empty(self): """Check if the episode data is empty.""" return not self._data - def add(self, key: str, value: torch.Tensor | dict, clone: bool = True): - """Add a key-value pair to the dataset. - - The key can be nested by using the "/" character. - For example: "obs/joint_pos". - - Args: - key: The key name. - value: The corresponding value of tensor type or of dict type. - clone: Whether to clone the tensor value before storing it in the episode data. - """ - if isinstance(value, dict): - for sub_key, sub_value in value.items(): - self.add(f"{key}/{sub_key}", sub_value, clone=clone) - return - - stored = value.clone() if (clone and isinstance(value, torch.Tensor)) else value - sub_keys = key.split("/") - data = self._data - for sub_key in sub_keys[:-1]: - if sub_key not in data: - data[sub_key] = {} - data = data[sub_key] - # Accumulate in lists to avoid copying tensors on each append. - if sub_keys[-1] not in data: - data[sub_keys[-1]] = [] - data[sub_keys[-1]].append(stored) - def get_initial_state(self) -> torch.Tensor | None: """Get the initial state from the dataset.""" return self._data.get("initial_state") @@ -158,18 +134,6 @@ def get_next_joint_target(self) -> dict | torch.Tensor | None: self._next_joint_target_index += 1 return joint_target - def pre_export(self): - """Prepare data for export by converting lists to tensors.""" - - def pre_export_helper(data): - for key, value in data.items(): - if isinstance(value, list): - data[key] = torch.stack(value) - elif isinstance(value, dict): - pre_export_helper(value) - - pre_export_helper(self._data) - def _index_nested(data: dict | torch.Tensor, index: int, keep_dim: bool) -> dict | torch.Tensor | None: """Select ``index`` along the first dimension of every tensor in a (nested) dict. diff --git a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py index e8042d488292..0a45428d7daf 100644 --- a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py +++ b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py @@ -12,8 +12,20 @@ import gymnasium as gym import numpy as np import pytest +import torch from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import ( + DatasetExportMode, + ObservationGroupCfg, + ObservationManager, + ObservationTermCfg, + RecorderManager, + RecorderManagerBaseCfg, + RecorderTerm, + RecorderTermCfg, +) +from isaaclab.utils.noise import UniformNoiseCfg pytestmark = pytest.mark.unit @@ -102,3 +114,99 @@ def test_obs_space_follows_clip_constraint(): assert term_space.shape == expected_shapes[term_name] assert np.all(term_space.low == low) assert np.all(term_space.high == high) + + +def _state_observation(env): + return env.scene.state + + +class _PhysicalRecorder(RecorderTerm): + def record_post_step(self): + self.latest = self._env.scene.state.clone() + return None, None + + +class _ObservationRecorder(RecorderTerm): + def record_post_step(self): + self.latest = self._env.observation_manager.compute_group("state", update_history=False) + return None, None + + +def _make_stepping_env(recorder_type=None): + """Exercise the real step/observation/recorder owners with a minimal deterministic physics leaf.""" + env = object.__new__(ManagerBasedRLEnv) + env._is_closed = True + env.cfg = SimpleNamespace(decimation=1, sim=SimpleNamespace(dt=0.02, render_interval=1), compute_final_obs=False) + env.scene = SimpleNamespace(num_envs=3, state=torch.zeros(3, 2)) + env.scene.write_data_to_sim = lambda: None + env.scene.update = lambda dt: None + env.sim = SimpleNamespace( + device="cpu", + is_playing=lambda: True, + is_rendering=False, + step=lambda render: env.scene.state.add_(1), + consume_reset_request=lambda: False, + ) + env._physics_handles_decimation = False + env._sim_step_counter = env.common_step_counter = 0 + env.episode_length_buf = torch.zeros(3, dtype=torch.long) + env.render_enabled = False + env.video_recorders = [] + env.extras = {} + env.action_manager = SimpleNamespace(process_action=lambda action: None, apply_action=lambda: None) + mask = torch.zeros(3, dtype=torch.bool) + + def terminate(): + mask[0] = env.common_step_counter == 4 + return mask + + env.termination_manager = SimpleNamespace( + active_terms=[], compute=terminate, terminated=mask, time_outs=torch.zeros_like(mask) + ) + env.reward_manager = SimpleNamespace(compute=lambda dt: torch.zeros(3)) + env.command_manager = SimpleNamespace(compute=lambda dt: None) + env.event_manager = SimpleNamespace(active_terms=[]) + group = ObservationGroupCfg(history_length=3, flatten_history_dim=False, enable_corruption=True) + group.state = ObservationTermCfg(func=_state_observation, noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1)) + state = ObservationGroupCfg(enable_corruption=False) + state.state = ObservationTermCfg(func=_state_observation) + env.observation_manager = ObservationManager({"policy": group, "state": state}, env) + cfg = RecorderManagerBaseCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE) + if recorder_type is not None: + cfg.capture = RecorderTermCfg(class_type=recorder_type) + env.recorder_manager = RecorderManager(cfg, env) + + def reset(env_ids): + env.scene.state[env_ids] = -100.0 + env.observation_manager.reset(env_ids) + env.recorder_manager.reset(env_ids) + + env._reset_idx = reset + env.obs_buf = env.observation_manager.compute(update_history=True) + return env + + +def test_raw_capture_preserves_observation_noise_history_and_rng(): + """Capture sees terminal state without implicit observation work or changes to policy history/RNG.""" + outputs, states = [], [] + for recorder_type in (None, _PhysicalRecorder, _ObservationRecorder): + torch.manual_seed(17) + env = _make_stepping_env(recorder_type) + history = [env.step(torch.empty(3, 0))[0]["policy"].clone() for _ in range(4)] + outputs.append(torch.stack(history)) + states.append(torch.get_rng_state()) + if recorder_type is not None: + term = env.recorder_manager.get_term("capture") + assert not env.recorder_manager._fields + # Both direct capture and explicit noiseless observations see the state before reset. + assert torch.equal(term.latest, torch.full((3, 2), 4.0)) + assert torch.equal(env.scene.state[0], torch.full((2,), -100.0)) + env.recorder_manager.close() + assert torch.equal(outputs[0], outputs[1]) + assert torch.equal(states[0], states[1]) + assert torch.equal(outputs[0], outputs[2]) + assert torch.equal(states[0], states[2]) + # Continuing worlds retain steps 2, 3 and 4; the reset world's history starts afresh. + expected = torch.tensor([2.0, 3.0, 4.0])[None, :, None] + assert torch.all((outputs[1][-1, 1:] - expected).abs() <= 0.1) + assert torch.all((outputs[1][-1, 0] + 100.0).abs() <= 0.1) diff --git a/source/isaaclab/test/managers/test_recorder_manager.py b/source/isaaclab/test/managers/test_recorder_manager.py index c07dfba99d10..f37214334707 100644 --- a/source/isaaclab/test/managers/test_recorder_manager.py +++ b/source/isaaclab/test/managers/test_recorder_manager.py @@ -4,18 +4,21 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -import os -import shutil -import tempfile -import uuid +import ast +import inspect +import sys from collections.abc import Sequence +from types import SimpleNamespace from typing import TYPE_CHECKING, cast +import h5py import pytest import torch +import warp as wp from isaaclab.managers import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets import EpisodeData pytestmark = pytest.mark.unit @@ -26,11 +29,13 @@ class DummyResetRecorderTerm(RecorderTerm): """Dummy recorder term that records dummy data.""" - def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: - return "record_pre_reset", torch.ones(self._env.num_envs, 2, device=self._env.device) + def record_pre_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str, torch.Tensor]: + self.env_ids = env_ids + return "record_pre_reset", torch.ones(self._env.num_envs, 2, device=self._env.device)[env_ids] - def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: - return "record_post_reset", torch.ones(self._env.num_envs, 3, device=self._env.device) + def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str, torch.Tensor]: + self.env_ids = env_ids + return "record_post_reset", torch.ones(self._env.num_envs, 3, device=self._env.device)[env_ids] class DummyStepRecorderTerm(RecorderTerm): @@ -47,137 +52,205 @@ def record_post_step(self) -> tuple[str | None, torch.Tensor | None]: class DummyRecorderManagerCfg(RecorderManagerBaseCfg): """Dummy recorder configurations.""" - @configclass - class DummyResetRecorderTermCfg(RecorderTermCfg): - """Configuration for the dummy reset recorder term.""" - - class_type: type[RecorderTerm] = DummyResetRecorderTerm - - @configclass - class DummyStepRecorderTermCfg(RecorderTermCfg): - """Configuration for the dummy step recorder term.""" - - class_type: type[RecorderTerm] = DummyStepRecorderTerm - - record_reset_term = DummyResetRecorderTermCfg() - record_step_term = DummyStepRecorderTermCfg() + record_reset_term = RecorderTermCfg(class_type=DummyResetRecorderTerm) + record_step_term = RecorderTermCfg(class_type=DummyStepRecorderTerm) dataset_export_mode = DatasetExportMode.EXPORT_ALL -@configclass -class DummyEnvCfg: - """Dummy environment configuration.""" - - @configclass - class DummySimCfg: - """Configuration for the dummy sim.""" - - dt: float = 0.01 - render_interval: int = 1 - - @configclass - class DummySceneCfg: - """Configuration for the dummy scene.""" - - num_envs: int = 1 - - decimation: int = 1 - sim: DummySimCfg = DummySimCfg() - scene: DummySceneCfg = DummySceneCfg() - - -class DummySimulation: - """Minimal simulation double used by :class:`RecorderManager`.""" - - def is_playing(self) -> bool: - """Return whether the simulated timeline is playing.""" - return True - - -class DummyTerminationManager: - """Minimal termination manager double without success terms.""" - - active_terms: list[str] = [] - - -class DummyEnv: - """Minimal environment double used by recorder terms and metadata export.""" - - def __init__(self, device: str = "cpu", num_envs: int = 20) -> None: - self.num_envs = num_envs - self.device = device - self.sim = DummySimulation() - self.cfg = DummyEnvCfg() - self.cfg.scene.num_envs = num_envs - self.termination_manager = DummyTerminationManager() - - -def create_dummy_env(device: str = "cpu", num_envs: int = 20) -> ManagerBasedEnv: +def create_dummy_env(device: str = "cpu", num_envs: int = 3) -> ManagerBasedEnv: """Create a minimal environment double.""" - - return cast("ManagerBasedEnv", DummyEnv(device=device, num_envs=num_envs)) - - -@pytest.fixture -def dataset_dir(): - """Create directory to dump results.""" - test_dir = tempfile.mkdtemp() - yield test_dir - # Cleanup - shutil.rmtree(test_dir) - - -def test_str(dataset_dir): - """Test the string representation of the recorder manager.""" - # create recorder manager - cfg = DummyRecorderManagerCfg() - cfg.dataset_export_dir_path = dataset_dir - cfg.dataset_filename = f"{uuid.uuid4()}.hdf5" - recorder_manager = RecorderManager(cfg, create_dummy_env()) - assert len(recorder_manager.active_terms) == 2 - manager_str = str(recorder_manager) - assert "contains 2 active terms" in manager_str - assert "record_reset_term" in manager_str - assert "record_step_term" in manager_str - # the dataset file is created on construction - assert os.path.exists(os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename)) - recorder_manager.close() - - -def test_record(dataset_dir): - """Test the recording of the data.""" + env = SimpleNamespace( + num_envs=num_envs, + device=device, + sim=SimpleNamespace(is_playing=lambda: True), + cfg=SimpleNamespace( + decimation=1, + sim=SimpleNamespace(dt=0.01, render_interval=1), + scene=SimpleNamespace(num_envs=num_envs), + ), + termination_manager=SimpleNamespace(active_terms=[]), + ) + return cast("ManagerBasedEnv", env) + + +def test_record(tmp_path): + """Capture callbacks and full/partial resets preserve episode boundaries and selectors.""" env = create_dummy_env() - # 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 torch.stack(episode.data["record_pre_step"]).shape == (2, 4) - assert torch.stack(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=slice(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=slice(None)) + for cfg in (None, {}, RecorderManagerBaseCfg()): + disabled = RecorderManager(cfg, env) + assert disabled.exported_successful_episode_count == disabled.exported_failed_episode_count == 0 + cfg = DummyRecorderManagerCfg(dataset_export_dir_path=str(tmp_path), dataset_filename="dataset.hdf5") + manager = RecorderManager(cfg, env) + assert (tmp_path / cfg.dataset_filename).is_file() + reset_term = manager.get_term("record_reset_term") + + for _ in range(2): + manager.record_pre_step() + manager.record_post_step() for env_id in range(env.num_envs): - episode = recorder_manager.get_episode(env_id) - assert torch.stack(episode.data["record_post_reset"]).shape == (1, 3) - recorder_manager.reset(slice(1, None, 2)) + episode = manager.get_episode(env_id) + assert episode.data["record_pre_step"].shape == (2, 4) + assert episode.data["record_post_step"].shape == (2, 5) + + manager.set_success_to_episodes((2, 0), torch.tensor([[False], [True]])) + assert [manager.get_episode(i).success for i in range(env.num_envs)] == [True, None, False] + manager.record_pre_reset(slice(None)) # Capture, export and clear the completed episodes. + assert reset_term.env_ids == slice(None) + assert all(manager.get_episode(i).is_empty() for i in range(env.num_envs)) + manager.record_post_reset(None) + assert reset_term.env_ids == slice(None) + manager.reset(slice(1, None, 2)) for env_id in range(env.num_envs): - assert recorder_manager.get_episode(env_id).is_empty() == (env_id % 2 == 1) - recorder_manager.close() + assert manager.get_episode(env_id).is_empty() == (env_id % 2 == 1) + manager.record_post_reset(slice(1, None, 2)) + assert reset_term.env_ids == slice(1, None, 2) + assert all(manager.get_episode(i).data["record_post_reset"].shape == (1, 3) for i in range(env.num_envs)) + manager.reset((2, 0)) + assert [manager.get_episode(i).is_empty() for i in range(env.num_envs)] == [True, False, True] + manager.record_post_reset((2, 0)) + manager.record_pre_reset((2, 0), force_export_or_skip=False) + assert [manager.get_episode(i).success for i in range(env.num_envs)] == [False, None, False] + manager.set_success_to_episodes((), torch.empty(0, 1, dtype=torch.bool)) + manager.reset(()) + assert all(manager.get_episode(i).data["record_post_reset"].shape == (1, 3) for i in range(env.num_envs)) + manager.close() + + +def test_batched_episode_indices_and_ownership(monkeypatch): + """Asynchronous resets, field cadence, growth and snapshots must preserve exact samples.""" + cfg = DummyRecorderManagerCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE, buffer_initial_capacity=2) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + manager.add_to_episodes("unused", torch.empty(0, 7), []) + assert not manager._fields + with pytest.raises(ValueError, match="requires 3 rows"): + manager.add_to_episodes("invalid", torch.empty(2, 7)) + assert not manager._fields + source = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + # Policies commonly collect under inference_mode, then reset/export after leaving it. + with torch.inference_mode(): + manager.add_to_episodes("obs", {"state": source}) + wp.init() + manager.add_to_episodes("initial_state", wp.from_torch(torch.tensor([[10], [20], [30]], dtype=torch.int32))) + source.fill_(99) + manager.add_to_episodes("obs/state", source) + snapshot = manager.get_episode(1) + assert snapshot.data["obs"]["state"][0, 0] == 2 + assert torch.isnan(snapshot.data["obs"]["state"][0, 1]) + assert torch.signbit(manager.get_episode(0).data["obs"]["state"][0, 1]) + assert snapshot.data["initial_state"].dtype == torch.int32 + assert snapshot.success is None + + manager.reset(torch.tensor([1])) + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.tensor([[42.0, 43.0], [52.0, 53.0]]), [2, 1]) + manager.add_to_episodes("obs/state", torch.tensor([[62.0, 63.0]]), slice(0, 1)) + assert torch.equal(manager.get_episode(1).data["obs"]["state"], torch.tensor([[52.0, 53.0]])) + assert manager.get_episode(1).data.keys() == {"obs"} + assert torch.equal(manager.get_episode(2).data["obs"]["state"][:, 0], torch.tensor([3.0, 99.0, 42.0])) + assert torch.equal(manager.get_episode(0).data["obs"]["state"][:, 0], torch.tensor([1.0, 99.0, 62.0])) + assert snapshot.data["obs"]["state"].shape == (2, 2) + snapshot = manager.get_episode(1) + snapshot.data["obs"]["state"].fill_(-100) + assert manager.get_episode(1).data["obs"]["state"][0, 0] == 52 + + with pytest.raises(ValueError, match="changed shape"): + manager.add_to_episodes("obs/state", torch.zeros(3, 5)) + with pytest.raises(ValueError, match="conflicts"): + manager.add_to_episodes("obs", torch.zeros(3, 2)) + manager.add_to_episodes("obs/state", torch.empty(0, 2), []) + assert manager.get_episode(0).data["obs"]["state"].shape == (3, 2) + + manager.reset() + # A new schema is allowed once every world has released this field, e.g. manual annotation. + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + assert manager.get_episode(0).data["obs"]["state"].dtype == torch.bool + manager.reset() + + def forbidden(*args, **kwargs): + raise AssertionError("Steady-state recording must reuse storage without episode objects or host index reads") + + # Architectural/performance gate: preallocated insertion and reset need neither episode objects, + # per-world container appends, nor device-to-host scalar reads. Export-none also stays batched. + with monkeypatch.context() as patch: + patch.setattr(EpisodeData, "__init__", forbidden) + patch.setattr(torch.Tensor, "item", forbidden) + patch.setattr(torch, "empty", forbidden) + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + manager.set_success_to_episodes([0, 2], torch.tensor([[True], [False]])) + manager.export_episodes([0, 2]) + manager.reset([1]) + assert manager.exported_successful_episode_count == 1 + assert manager.exported_failed_episode_count == 1 + assert all(manager.get_episode(i).is_empty() for i in range(3)) + manager.close() + + +@pytest.mark.parametrize("mode", list(DatasetExportMode)) +def test_export_batched_episodes(tmp_path, mode): + """Export modes, custom IDs, nested tensor schema and counters survive partial finalization.""" + cfg = DummyRecorderManagerCfg(dataset_export_dir_path=str(tmp_path), dataset_export_mode=mode, export_in_close=True) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + values = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + manager.add_to_episodes("actions", values) + manager.add_to_episodes("obs", {"flag": torch.tensor([True, False, True])}) + manager.set_success_to_episodes(None, torch.tensor([[True], [False], [True]])) + with pytest.raises(ValueError, match="must match"): + manager.export_episodes([2, 1], [30]) + with pytest.raises(ValueError, match="must be unique"): + manager.export_episodes([2, 1], [30, 30]) + manager.export_episodes(torch.tensor([2, 1]), [30, 10]) + assert manager.get_episode(1).is_empty() + assert manager.get_episode(2).is_empty() + assert manager.get_episode(0).success is True + manager.export_episodes(slice(0, 1), [20]) + manager.export_episodes() # Empty episodes do not affect counters. + assert manager.exported_successful_episode_count == 2 + assert manager.exported_failed_episode_count == 1 + manager.close() + if mode == DatasetExportMode.EXPORT_NONE: + assert list(tmp_path.iterdir()) == [] + return + with h5py.File(tmp_path / "dataset.hdf5") as file: + expected = {"demo_30", "demo_20"} + if mode == DatasetExportMode.EXPORT_ALL: + expected.add("demo_10") + assert set(file["data"]) == expected + assert file["data/demo_20/actions"][()].tobytes() == values[0:1].numpy().tobytes() + assert bool(file["data/demo_30/obs/flag"][0]) + assert bool(file["data/demo_20"].attrs["success"]) + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + with h5py.File(tmp_path / "dataset_failed.hdf5") as file: + assert set(file["data"]) == {"demo_10"} + assert file["data/demo_10/actions"][()].tobytes() == values[1:2].numpy().tobytes() + assert not bool(file["data/demo_10"].attrs["success"]) + + +def test_recorder_storage_boundary(): + """Keep storage independent of consumers, with batched ownership and no Python world loops.""" + import isaaclab.managers.recorder_manager as module + + source = ast.parse(inspect.getsource(module)) + # Batched buffers own collection; snapshots and capture terms have no second writer lifecycle. + assert not {"add", "pre_export"} & vars(EpisodeData).keys() + assert "__init__" not in vars(module.RecorderTerm) + assert not hasattr(module.RecorderTerm, "close") + # Observation evaluation belongs to explicit terms, never generic capture dispatch. + assert not hasattr(module.RecorderTerm, "requires_post_step_observations") + assert not hasattr(module.RecorderManager, "requires_post_step_observations") + assert not any( + isinstance(node, ast.Attribute) and node.attr in {"observation_manager", "_episodes"} + for node in ast.walk(source) + ) + allowed = sys.stdlib_module_names | {"isaaclab", "torch", "warp", "prettytable"} + for node in ast.walk(source): + if isinstance(node, ast.Import): + assert all(alias.name.split(".")[0] in allowed for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + assert (node.module or "").split(".")[0] in allowed + elif isinstance(node, ast.FunctionDef) and node.name.startswith("exported_"): + assert not any(isinstance(child, ast.Name) and child.id == "hasattr" for child in ast.walk(node)) + field = next(node for node in source.body if isinstance(node, ast.ClassDef) and node.name == "_RecorderField") + append = next(node for node in field.body if isinstance(node, ast.FunctionDef) and node.name == "append") + assert not any(isinstance(node, (ast.For, ast.While, ast.comprehension)) for node in ast.walk(append)) diff --git a/source/isaaclab/test/utils/test_episode_data.py b/source/isaaclab/test/utils/test_episode_data.py deleted file mode 100644 index 7aa071a42a54..000000000000 --- a/source/isaaclab/test/utils/test_episode_data.py +++ /dev/null @@ -1,147 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause -import pytest -import torch - -from isaaclab.test.utils import DeviceScope, test_devices -from isaaclab.utils.datasets import EpisodeData - -pytestmark = [pytest.mark.unit, pytest.mark.windows_ci] - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_is_empty(device): - """Test checking whether the episode is empty.""" - episode = EpisodeData() - assert episode.is_empty() - - episode.add("key", torch.tensor([1, 2, 3], device=device)) - assert not episode.is_empty() - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_add_tensors(device): - """Test appending tensor data to the episode.""" - dummy_data_0 = torch.tensor([0], device=device) - dummy_data_1 = torch.tensor([1], device=device) - expected_added_data = torch.cat((dummy_data_0.unsqueeze(0), dummy_data_1.unsqueeze(0))) - episode = EpisodeData() - - # test adding data to a key that does not exist - episode.add("key", dummy_data_0) - key_data = torch.stack(episode.data.get("key")) - assert key_data is not None - assert torch.equal(key_data, dummy_data_0.unsqueeze(0)) - - # test adding data to a key that exists - episode.add("key", dummy_data_1) - key_data = torch.stack(episode.data.get("key")) - assert key_data is not None - assert torch.equal(key_data, expected_added_data) - - # test adding data to a key with "/" in the name - episode.add("first/second", dummy_data_0) - first_data = episode.data.get("first") - assert first_data is not None - second_data = torch.stack(first_data.get("second")) - assert second_data is not None - assert torch.equal(second_data, dummy_data_0.unsqueeze(0)) - - # test adding data to a key with "/" in the name that already exists - episode.add("first/second", dummy_data_1) - first_data = episode.data.get("first") - assert first_data is not None - second_data = torch.stack(first_data.get("second")) - assert second_data is not None - assert torch.equal(second_data, expected_added_data) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_add_dict_tensors(device): - """Test appending dict data to the episode.""" - dummy_dict_data_0 = { - "key_0": torch.tensor([0], device=device), - "key_1": {"key_1_0": torch.tensor([1], device=device), "key_1_1": torch.tensor([2], device=device)}, - } - dummy_dict_data_1 = { - "key_0": torch.tensor([3], device=device), - "key_1": {"key_1_0": torch.tensor([4], device=device), "key_1_1": torch.tensor([5], device=device)}, - } - - episode = EpisodeData() - - # test adding dict data to a key that does not exist - episode.add("key", dummy_dict_data_0) - key_data = episode.data.get("key") - assert key_data is not None - key_0_data = torch.stack(key_data.get("key_0")) - assert key_0_data is not None - assert torch.equal(key_0_data, torch.tensor([[0]], device=device)) - key_1_data = key_data.get("key_1") - assert key_1_data is not None - key_1_0_data = torch.stack(key_1_data.get("key_1_0")) - assert key_1_0_data is not None - assert torch.equal(key_1_0_data, torch.tensor([[1]], device=device)) - key_1_1_data = torch.stack(key_1_data.get("key_1_1")) - assert key_1_1_data is not None - assert torch.equal(key_1_1_data, torch.tensor([[2]], device=device)) - - # test adding dict data to a key that exists - episode.add("key", dummy_dict_data_1) - key_data = episode.data.get("key") - assert key_data is not None - key_0_data = torch.stack(key_data.get("key_0")) - assert key_0_data is not None - assert torch.equal(key_0_data, torch.tensor([[0], [3]], device=device)) - key_1_data = key_data.get("key_1") - assert key_1_data is not None - key_1_0_data = torch.stack(key_1_data.get("key_1_0")) - assert key_1_0_data is not None - assert torch.equal(key_1_0_data, torch.tensor([[1], [4]], device=device)) - key_1_1_data = torch.stack(key_1_data.get("key_1_1")) - assert key_1_1_data is not None - assert torch.equal(key_1_1_data, torch.tensor([[2], [5]], device=device)) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_get_initial_state(device): - """Test getting the initial state of the episode.""" - dummy_initial_state = torch.tensor([1, 2, 3], device=device) - episode = EpisodeData() - - episode.add("initial_state", dummy_initial_state) - initial_state = torch.stack(episode.get_initial_state()) - assert initial_state is not None - assert torch.equal(initial_state, dummy_initial_state.unsqueeze(0)) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_get_next_action(device): - """Test getting next actions.""" - # dummy actions - action1 = torch.tensor([1, 2, 3], device=device) - action2 = torch.tensor([4, 5, 6], device=device) - action3 = torch.tensor([7, 8, 9], device=device) - - episode = EpisodeData() - assert episode.get_next_action() is None - - episode.add("actions", action1) - episode.add("actions", action2) - episode.add("actions", action3) - - # check if actions are returned in the correct order - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action1) - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action2) - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action3) - - # check if None is returned when all actions are exhausted - assert episode.get_next_action() is None diff --git a/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py b/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py index 8ada18509cb8..677de8d93327 100644 --- a/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py +++ b/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py @@ -3,10 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause import json -import os -import sys -import tempfile -import uuid import h5py import pytest @@ -18,68 +14,19 @@ pytestmark = [pytest.mark.unit, pytest.mark.windows_ci] -def create_test_episode(device): - """create a test episode with dummy data.""" - test_episode = EpisodeData() - - test_episode.seed = 0 - test_episode.success = True - - test_episode.add("initial_state", torch.tensor([1, 2, 3], device=device)) - - test_episode.add("actions", torch.tensor([1, 2, 3], device=device)) - test_episode.add("actions", torch.tensor([4, 5, 6], device=device)) - test_episode.add("actions", torch.tensor([7, 8, 9], device=device)) - - test_episode.add("obs/policy/term1", torch.tensor([1, 2, 3, 4, 5], device=device)) - test_episode.add("obs/policy/term1", torch.tensor([6, 7, 8, 9, 10], device=device)) - test_episode.add("obs/policy/term1", torch.tensor([11, 12, 13, 14, 15], device=device)) - - return test_episode - - -@pytest.fixture -def temp_dir(): - """Create a temporary directory for test datasets.""" - # On Windows, libhdf5 releases its file handle shortly after close(), so the cleanup can race it. - with tempfile.TemporaryDirectory(ignore_cleanup_errors=sys.platform == "win32") as d: - yield d - - -def test_create_dataset_file(temp_dir, monkeypatch): - """Test creating a new dataset file.""" - # create a dataset file given a file name with extension - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_path, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(dataset_file_path) - - # create a dataset file given a file name without extension - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_path, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(dataset_file_path + ".hdf5") - - # create a dataset file given a bare file name in the current working directory - monkeypatch.chdir(temp_dir) - dataset_file_name = f"{uuid.uuid4()}.hdf5" - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_name, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(os.path.join(temp_dir, dataset_file_name)) +def test_create_dataset_file(tmp_path, monkeypatch): + """Accept absolute paths, omitted extensions and bare filenames.""" + monkeypatch.chdir(tmp_path) + for path in (str(tmp_path / "absolute.hdf5"), str(tmp_path / "no_extension"), "relative.hdf5"): + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.create(path, "test_env_name") + dataset_file_handler.close() + assert {path.name for path in tmp_path.iterdir()} == {"absolute.hdf5", "no_extension.hdf5", "relative.hdf5"} -def test_add_env_args_preserves_existing_args_after_reopen(temp_dir): - """Test extending environment arguments after reopening a dataset.""" - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") +def test_env_args_follow_active_file(tmp_path): + """Preserve metadata on reopen, then discard it when reusing the handler for a new file.""" + dataset_file_path = str(tmp_path / "first.hdf5") dataset_file_handler = HDF5DatasetFileHandler() dataset_file_handler.create(dataset_file_path, "test_env_name") dataset_file_handler.close() @@ -94,16 +41,7 @@ def test_add_env_args_preserves_existing_args_after_reopen(temp_dir): assert env_args == {"env_name": "test_env_name", "type": 2, "custom_arg": "custom_value"} - -def test_create_resets_env_args_when_reusing_handler(temp_dir): - """Test environment arguments do not leak between datasets.""" - first_dataset_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - second_dataset_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(first_dataset_path, "first_env") - dataset_file_handler.add_env_args({"custom_arg": "custom_value"}) - dataset_file_handler.close() - + second_dataset_path = str(tmp_path / "second.hdf5") dataset_file_handler.create(second_dataset_path, "second_env") dataset_file_handler.close() @@ -113,10 +51,10 @@ def test_create_resets_env_args_when_reusing_handler(temp_dir): assert env_args == {"env_name": "second_env", "type": 2} -def test_convert_dataset_to_xyzw_updates_format_version(temp_dir): +def test_convert_dataset_to_xyzw_updates_format_version(tmp_path): """Converted datasets should load their XYZW quaternions without converting them again.""" - input_path = os.path.join(temp_dir, "legacy.hdf5") - output_path = os.path.join(temp_dir, "converted.hdf5") + input_path = str(tmp_path / "legacy.hdf5") + output_path = str(tmp_path / "converted.hdf5") with h5py.File(input_path, "w") as dataset: dataset.attrs["format_version"] = 0 data = dataset.create_group("data") @@ -139,16 +77,27 @@ def test_convert_dataset_to_xyzw_updates_format_version(temp_dir): @pytest.mark.parametrize("device", test_devices()) -def test_write_and_load_episode(temp_dir, device): - """Test writing and loading an episode to and from the dataset file.""" - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") +def test_write_and_load_episode(tmp_path, device): + """Round-trip episode data and metadata, then play every action through exhaustion.""" + dataset_file_path = str(tmp_path / "episodes.hdf5") dataset_file_handler = HDF5DatasetFileHandler() dataset_file_handler.create(dataset_file_path, "test_env_name") - test_episode = create_test_episode(device) + test_episode = EpisodeData() + assert test_episode.is_empty() + assert test_episode.get_next_action() is None + test_episode.seed = 0 + test_episode.success = True + test_episode.data = { + "initial_state": torch.tensor([[1, 2, 3]], device=device), + "actions": torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device), + "obs": {"policy": {"term1": torch.arange(1, 16, device=device).reshape(3, 5)}}, + } + assert not test_episode.is_empty() + # Reading an absent action must not advance the cursor before data becomes available. + assert torch.equal(test_episode.get_next_action(), test_episode.data["actions"][0]) # write the episode to the dataset - test_episode.pre_export() dataset_file_handler.write_episode(test_episode) dataset_file_handler.flush() @@ -178,9 +127,12 @@ def test_write_and_load_episode(temp_dir, device): assert loaded_episode.seed == test_episode.seed assert loaded_episode.success == test_episode.success - assert torch.equal(loaded_episode.get_initial_state(), test_episode.get_initial_state()) + assert torch.equal(loaded_episode.get_initial_state(), test_episode.data["initial_state"]) + + assert torch.equal(loaded_episode.data["obs"]["policy"]["term1"], test_episode.data["obs"]["policy"]["term1"]) for action in test_episode.data["actions"]: assert torch.equal(loaded_episode.get_next_action(), action) + assert loaded_episode.get_next_action() is None dataset_file_handler.close() diff --git a/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..29d7578a43f8 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,3 @@ +* **Breaking:** Removed the implicit observation computation before Warp RL post-step recorder callbacks. + Capture runs before automatic reset. Custom terms needing observations must request them explicitly + from the observation manager; ``env.obs_buf`` is no longer refreshed for recorder dispatch. diff --git a/source/isaaclab_experimental/changelog.d/recorder-capture.major b/source/isaaclab_experimental/changelog.d/recorder-capture.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index be522e073a2b..5ac88958d450 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -320,14 +320,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: timer=DEBUG_TIMER_STEP, ) - if len(self.recorder_manager.active_terms) > 0: - # update observations for recording if needed - self._manager_call_switch.call_stage( - stage="ObservationManager_compute_no_history", - warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, - timer=DEBUG_TIMER_STEP, - ) - self.recorder_manager.record_post_step() + self.recorder_manager.record_post_step() # -- reset envs that terminated/timed-out and log the episode information # NOTE: Interim path (intentional). diff --git a/source/isaaclab_teleop/changelog.d/recorder-replay-fixture.skip b/source/isaaclab_teleop/changelog.d/recorder-replay-fixture.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_teleop/test/test_replay_demos_loop_termination.py b/source/isaaclab_teleop/test/test_replay_demos_loop_termination.py index c3f57f4640e5..a985c18d1b7e 100644 --- a/source/isaaclab_teleop/test/test_replay_demos_loop_termination.py +++ b/source/isaaclab_teleop/test/test_replay_demos_loop_termination.py @@ -54,7 +54,7 @@ def __init__(self, actions: torch.Tensor): def load_episode(self, episode_name, device): episode = EpisodeData() - episode.data = {"initial_state": {}, "actions": list(self._actions)} + episode.data = {"initial_state": {}, "actions": self._actions} return episode