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.