Adds support to record additional steps in record_demos.py (#224)

Adds support to record additional steps after task success is detected
in record_demos.py.
This helps record demos with sufficient lengths to be used for training.

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)

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

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

scripts/tools/record_demos.py

@@ -19,6 +19,7 @@ optional arguments:
     --dataset_file            File path to export recorded demos. (default: "./datasets/dataset.hdf5")
     --step_hz                 Environment stepping rate in Hz. (default: 30)
     --num_demos               Number of demonstrations to record. (default: 0)
+    --num_success_steps       Number of continuous steps with task success for concluding a demo as successful. (default: 10)
 """
 
 """Launch Isaac Sim Simulator first."""
@@ -39,6 +40,12 @@ parser.add_argument("--step_hz", type=int, default=30, help="Environment steppin
 parser.add_argument(
     "--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite."
 )
+parser.add_argument(
+    "--num_success_steps",
+    type=int,
+    default=10,
+    help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.",
+)
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
 # parse the arguments
@@ -58,6 +65,8 @@ import gymnasium as gym
 import time
 import torch
 
+import omni.log
+
 from isaaclab.devices import Se3HandTracking, Se3Keyboard, Se3SpaceMouse
 from isaaclab.envs import ViewerCfg
 from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
@@ -127,8 +136,19 @@ def main():
     env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1)
     env_cfg.env_name = args_cli.task
 
-    # modify configuration such that the environment runs indefinitely
-    # until goal is reached
+    # extract success checking function to invoke in the main loop
+    success_term = None
+    if hasattr(env_cfg.terminations, "success"):
+        success_term = env_cfg.terminations.success
+        env_cfg.terminations.success = None
+    else:
+        omni.log.warn(
+            "No success termination term was found in the environment."
+            " Will not be able to mark recorded demos as successful."
+        )
+
+    # modify configuration such that the environment runs indefinitely until
+    # the goal is reached or other termination conditions are met
     env_cfg.terminations.time_out = None
 
     env_cfg.observations.policy.concatenate_terms = False
@@ -173,6 +193,7 @@ def main():
 
     # simulate environment -- run everything in inference mode
     current_recorded_demo_count = 0
+    success_step_count = 0
     with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
         while True:
             # get keyboard command
@@ -187,15 +208,29 @@ def main():
             # perform action on environment
             env.step(actions)
 
+            if success_term is not None:
+                if bool(success_term.func(env, **success_term.params)[0]):
+                    success_step_count += 1
+                    if success_step_count >= args_cli.num_success_steps:
+                        env.recorder_manager.record_pre_reset([0], force_export_or_skip=False)
+                        env.recorder_manager.set_success_to_episodes(
+                            [0], torch.tensor([[True]], dtype=torch.bool, device=env.unwrapped.device)
+                        )
+                        env.recorder_manager.export_episodes([0])
+                        should_reset_recording_instance = True
+                else:
+                    success_step_count = 0
+
             if should_reset_recording_instance:
                 env.unwrapped.recorder_manager.reset()
                 env.reset()
                 should_reset_recording_instance = False
+                success_step_count = 0
 
             # print out the current demo count if it has changed
             if env.unwrapped.recorder_manager.exported_successful_episode_count > current_recorded_demo_count:
                 current_recorded_demo_count = env.unwrapped.recorder_manager.exported_successful_episode_count
-                print(f"Recorded {current_recorded_demo_count} demonstrations.")
+                print(f"Recorded {current_recorded_demo_count} successful demonstrations.")
 
             if (
                 args_cli.num_demos > 0

source/isaaclab/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.33.6"
+version = "0.33.7"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/isaaclab/docs/CHANGELOG.rst

 Changelog
 ---------
 
-0.33.6 (2025-01-30)
+0.33.7 (2025-01-30)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -12,7 +12,7 @@ Fixed
   to the event being triggered at the wrong time after the reset.
 
 
-0.33.5 (2025-01-17)
+0.33.6 (2025-01-17)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -42,7 +42,7 @@ Fixed
   the :class:`omni.isaac.lab.assets.RigidObjectCollection` class.
 
 
-0.33.4 (2025-01-14)
+0.33.5 (2025-01-14)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -51,6 +51,16 @@ Fixed
 * Fixed the respawn of only wrong object samples in :func:`repeated_objects_terrain` of :mod:`omni.isaac.lab.terrains.trimesh` module. Previously, the function was respawning all objects in the scene instead of only the wrong object samples, which in worst case could lead to infinite respawn loop.
 
 
+0.33.4 (2025-01-10)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Added an optional parameter in the :meth:`record_pre_reset` method in
+  :class:`~isaaclab.managers.RecorderManager` to override the export config upon invoking.
+
+
 0.33.3 (2025-01-08)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/managers/recorder_manager.py

@@ -359,7 +359,7 @@ class RecorderManager(ManagerBase):
             key, value = term.record_post_step()
             self.add_to_episodes(key, value)
 
-    def record_pre_reset(self, env_ids: Sequence[int] | None) -> None:
+    def record_pre_reset(self, env_ids: Sequence[int] | None, force_export_or_skip=None) -> None:
         """Trigger recorder terms for pre-reset functions.
 
         Args:
@@ -385,7 +385,7 @@ class RecorderManager(ManagerBase):
             success_results |= self._env.termination_manager.get_term("success")[env_ids]
         self.set_success_to_episodes(env_ids, success_results)
 
-        if self.cfg.export_in_record_pre_reset:
+        if force_export_or_skip or (force_export_or_skip is None and self.cfg.export_in_record_pre_reset):
             self.export_episodes(env_ids)
 
     def record_post_reset(self, env_ids: Sequence[int] | None) -> None: