Multi-Agent implementation only checks the first agent when computing the finished-episode info

[Toni-SM]

Discussed in #284

^{Originally posted by William-SKC February 27, 2025}
Hi, I have a question regarding the expected behavior of this line in record_transition in MultiAgent class:

finished_episodes = (next(iter(terminated.values())) + next(iter(truncated.values()))).nonzero(as_tuple=False)

From my understanding, terminated and truncated are dictionaries where the keys represent agent IDs and the values are tensors indicating whether an agent's episode has ended.

My Questions:
Does this line only check the first agent (since it uses next(iter(...))), rather than all agents?
If so, shouldn't we be stacking all values using torch.stack(list(terminated.values())) to ensure all agents are checked?
What is the expected output for finished_episodes? Should it track all finished agents or just the first one?
I appreciate any clarification on this! Thanks for your help. 🙏

[ibin-study]

Hello @Toni-SM ,

I'm a researcher currently using the skrl library for multi-agent reinforcement learning with the MAPPO algorithm in an Isaac Lab environment. First of all, thank you for providing such a great library for multi-agent training.

Regarding the recent discussion about how terminated and truncated dictionaries are structured, I understand and agree with the need for a proper fix. In my current setting, all agents share the same termination and truncation signals, so it's not causing functional issues during training. However, I encountered a strange logging behavior that I would like to report.

Specifically, I observed that Episode / Total timesteps (min) is frequently logged as 1, which raised a red flag. To investigate, I added print statements inside the following method from skrl/multi_agents/torch/base.py:

finished_episodes = (next(iter(terminated.values())) + next(iter(truncated.values()))).nonzero(as_tuple=False)

This returns a tensor of shape [N, 2]. However, the second column is always 0. This causes the following block:

self._cumulative_rewards[finished_episodes] = 0
self._cumulative_timesteps[finished_episodes] = 0

to always reset the timestep of environment 0, even if other environments terminated. As a result, when env 0 happens to terminate, its timestep is 1 and this gets recorded in _track_timesteps, which then incorrectly drives Total timesteps (min) to 1.

To fix this issue, I suggest modifying the reset lines as follows:

self._cumulative_rewards[finished_episodes[:, 0]] = 0
self._cumulative_timesteps[finished_episodes[:, 0]] = 0

This ensures that only the actually terminated environments are reset.

Here’s a snippet of what I observed in the terminal:

 0%|                                                                                                           | 0/80000 [00:00<?, ?it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([1], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 1/80000 [00:03<86:06:10,  3.87s/it]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([2], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([2], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 2/80000 [00:04<43:25:19,  1.95s/it]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([3], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([3], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 3/80000 [00:05<29:57:31,  1.35s/it]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([4], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([4], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 4/80000 [00:05<23:02:25,  1.04s/it]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([5], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([5], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 5/80000 [00:06<17:28:31,  1.27it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([6], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([6], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 6/80000 [00:06<13:46:27,  1.61it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([7], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([7], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 7/80000 [00:06<11:34:22,  1.92it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([8], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([8], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 8/80000 [00:06<10:01:56,  2.21it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([9], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([9], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                 | 9/80000 [00:07<9:03:06,  2.45it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([10], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([10], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 10/80000 [00:07<8:20:35,  2.66it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([11], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([11], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 11/80000 [00:07<7:47:15,  2.85it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([12], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([12], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 12/80000 [00:08<7:21:28,  3.02it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([13], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([13], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 13/80000 [00:08<7:30:48,  2.96it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([14], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([14], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 14/80000 [00:08<7:38:54,  2.90it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([15], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([15], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                                | 15/80000 [00:09<7:24:34,  3.00it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([16], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([16], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([0, 2])
  0%|                                                                                               | 16/80000 [00:10<10:53:40,  2.04it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([17], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([17], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([1, 2])
  0%|                                                                                               | 17/80000 [00:11<16:21:30,  1.36it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([18], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([1, 2])
  0%|                                                                                               | 18/80000 [00:11<15:52:18,  1.40it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([19], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([4, 2])
  0%|                                                                                               | 19/80000 [00:13<18:09:44,  1.22it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([20], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([5, 2])
  0%|                                                                                               | 20/80000 [00:13<18:20:26,  1.21it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([21], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([6, 2])
  0%|                                                                                               | 21/80000 [00:14<14:51:01,  1.50it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([22], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([8, 2])
  0%|                                                                                               | 22/80000 [00:14<12:19:49,  1.80it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([23], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([13, 2])
  0%|                                                                                               | 23/80000 [00:14<10:38:12,  2.09it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([24], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([22, 2])
  0%|                                                                                                | 24/80000 [00:15<9:23:17,  2.37it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([25], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([29, 2])
  0%|                                                                                                | 25/80000 [00:15<8:35:03,  2.59it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([26], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([29, 2])
  0%|                                                                                                | 26/80000 [00:15<8:02:43,  2.76it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([27], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([50, 2])
  0%|                                                                                                | 27/80000 [00:15<7:37:02,  2.92it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([28], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([73, 2])
  0%|                                                                                                | 28/80000 [00:16<7:20:26,  3.03it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([29], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([110, 2])
  0%|                                                                                                | 29/80000 [00:16<7:11:49,  3.09it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([30], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([131, 2])
  0%|                                                                                                | 30/80000 [00:16<7:06:31,  3.12it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([31], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([154, 2])
  0%|                                                                                                | 31/80000 [00:17<7:07:18,  3.12it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([32], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([183, 2])
  0%|                                                                                               | 32/80000 [00:18<10:29:54,  2.12it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([33], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([228, 2])
  0%|                                                                                                | 33/80000 [00:18<9:55:11,  2.24it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([34], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([246, 2])
  0%|                                                                                                | 34/80000 [00:18<9:31:58,  2.33it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([35], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([236, 2])
  0%|                                                                                                | 35/80000 [00:19<9:09:03,  2.43it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([36], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([227, 2])
  0%|                                                                                                | 36/80000 [00:19<8:42:11,  2.55it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([37], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([262, 2])
  0%|                                                                                                | 37/80000 [00:19<8:27:41,  2.63it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([38], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([237, 2])
  0%|                                                                                                | 38/80000 [00:20<8:17:55,  2.68it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([39], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([245, 2])
  0%|                                                                                                | 39/80000 [00:20<8:12:07,  2.71it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([40], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([181, 2])
  0%|                                                                                                | 40/80000 [00:20<8:03:27,  2.76it/s]cumulative timesteps shape: torch.Size([4096, 1])
cumulative timesteps [0]: tensor([1], device='cuda:3', dtype=torch.int32)
cumulative timesteps [1]: tensor([41], device='cuda:3', dtype=torch.int32)
finished episodes: torch.Size([180, 2])
  0%|                                                                                                | 41/80000 [00:21<8:04:02,  2.75it/s]
...

As seen above, the cumulative timestep of environment 0 remains at 1 repeatedly even as other environments accumulate properly.

So although this issue appears in the logging and visualization components, I believe it would be safer to review the training logic as well — to ensure that similar indexing issues don’t affect training stability or memory tracking under the hood.

Thank you again for your work on skrl, and I hope this report helps improve the framework further.