Unverified Commit f2d5b335 authored by chengronglai's avatar chengronglai Committed by GitHub

Adds trial counter to saved image filenames in blueprint cfg (#2073)

# Description

Uses recorder_manager's successful episode count as trial number and
adding it to saved image filenames.
Track successful episodes across environments by adding trial counter to
saved image filenames.

## Type of change

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

- New feature (non-breaking change which adds functionality)

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there
parent 6b794ac2
......@@ -83,9 +83,15 @@ def image(
os.makedirs(dir_path, exist_ok=True)
if images.dtype == torch.uint8:
images = images.float() / 255.0
# Get total successful episodes
total_successes = sum(env.recorder_manager._exported_successful_episode_count.values())
for tile in range(images.shape[0]):
tile_chw = torch.swapaxes(images[tile : tile + 1].unsqueeze(1), 1, -1).squeeze(-1)
save_image(tile_chw, f"{image_path}_{data_type}_{tile}_{env.common_step_counter}.png")
filename = (
f"{image_path}_{data_type}_trial_{total_successes}_tile_{tile}_step_{env.common_step_counter}.png"
)
save_image(tile_chw, filename)
return images.clone()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment