[BugFix] LLMCollector trajectory collection methods only returns 1 trajectory when multiple are completed at the same step

[albertbou92]

Description

Fixes dones indexing to always return a list in trajectory collection methods forLLMCollector

The current behavior using dones.nonzero()[0].tolist() only works correctly when a single element in dones is True. When multiple are True, it returns only the first index, not the full list.

This PR replaces it with dones.nonzero(as_tuple=True)[0].tolist(), which consistently returns a list of all indices where dones is True.

import torch

# One True
dones = torch.zeros(5, dtype=torch.bool)
dones[0] = True
print("dones:", dones)
print("dones.nonzero()[0].tolist():", dones.nonzero()[0].tolist())
print("dones.nonzero(as_tuple=True)[0].tolist():", dones.nonzero(as_tuple=True)[0].tolist())

# Multiple True
dones[1] = True
print("\ndones:", dones)
print("dones.nonzero()[0].tolist():", dones.nonzero()[0].tolist())  # Only first index
print("dones.nonzero(as_tuple=True)[0].tolist():", dones.nonzero(as_tuple=True)[0].tolist())  # Full list

dones: tensor([ True, False, False, False, False])
dones.nonzero()[0].tolist(): [0]
dones.nonzero(as_tuple=True)[0].tolist(): [0]

dones: tensor([ True, True, False, False, False])
dones.nonzero()[0].tolist(): [0]
dones.nonzero(as_tuple=True)[0].tolist(): [0, 1]

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3018

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Good catch thanks!