Fixes Termination Manager logging to report aggregated percentage of environments done due to each term.

[ooctipus]

Description

Currently Termination Manager write current step's done count for each term if reset is detected. This leads to two problem.

1. User sees different counts just by varying num_envs
2. the count can be over-count or under-count depending on when reset happens, as pointed out in [Bug Report] Termination Overcounting Caused by Missing Log Buffer Reset in manager_based_rl_env.py #2977 (Thanks, @Kyu3224)

The cause of the bug is because we are reporting current step status into a buffer that suppose to record episodic done. So instead of write the entire buffer base on current value, we ask the update to respect the non-reseting environment's old value, and instead of reporting count, we report percentage of environment that was done due to the particular term.

Test on Isaac-Velocity-Rough-Anymal-C-v0

Before fix:

Red: num_envs = 4096, Orange: num_envs = 1024

After fix:

Red: num_envs = 4096, Orange: num_envs = 1024

Note that curve of the same color ran on same seed, and curves matched exactly, the only difference is the data gets reported in termination. The percentage version is a lot more clear in conveying how agent currently fails, and how much percentage of agent fails, and shows that increasing num_envs to 4096 helps improve agent avoiding termination by base_contact much quicker than num_envs=1024. Such message is a bit hard to tell in first image.

Checklist

• I have run the pre-commit checks 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

[Mayankm96]

This indexing here is not great. We use the value from here to give termination rewards in different environments. Doing the .index every time for reward computation may lead to slow downs (as it searches over the list in O(n) fashion).

[ooctipus]

Ahhh if you like the single tensor approach I can also store one more name -> idx dict, and make change here O(1), I wasn't aware of termination rewards, (my bad!!) and thought this function may be used infrequently. I should do a global search next time rather than assume!

[Mayankm96]

Is there a reason to change this as a dict of tensors to a single tensor?

[ooctipus]

I thought it last_episode_done_stats = self._term_dones.float().mean(dim=0) this operation is very nice and optimized, thats why I did it, but if you think dict is more clear I can revert back : ))

[Mayankm96]

If you want the ratio, isn't that simply?

self._term_dones[key][env_ids].sum() / len(env_ids)

[ooctipus]

Thanks for reveiwing!
I guess, doing with env_ids will be viewing ratio of resetting environments. I thought maybe report stats of all environment can be a bit more nicer as user can verify from the graph that all terms sum up to 1. Of course you can do self._term_dones[key].sum() / self.env.num_envs as well

But it seems like if this approach is what I after, do it in one tensor operation seems quite nice, both speed wise and memory utility wise.

[github-actions]

Test Results Summary

2 419 tests   2 011 ✅  2h 22m 39s ⏱️
   90 suites    408 💤
    1 files        0 ❌

Results for commit a0767c6.

♻️ This comment has been updated with latest results.

[ooctipus]

@Mayankm96
I did my best to optimize the operations in compute and reset,
index all changed from O(n) to O(1) -> this part not captured in benchmark below

benchmarked task: velocity rough anymal c, 4096 envs, 1000 steps

before change :
one step total: 79.038

| termination.compute     |       0.229 ms|
| termination.reset       |       0.007 ms|

after change:

one step total: 80.187 ms

| termination.compute     |       0.274 ms|  slower due to one extra operation
| termination.reset       |       0.004 ms|  twice as fast

the cost for this PR is about 0.04 ms / 80, 0.05%, mostly due to compute needs modify other terms's done to correctly update the done buffer. I think this is reasonable