[BugFix] Keep StepCounter flag propagation within each done group

[mottopanikeiku]

Description

StepCounter propagates truncation to nested done entries (and done too, with update_done=True) from a parent key. The loops took the first parent key that had a value, even if it belonged to an independent sibling group. So with ("left", ...) and ("right", ...) groups that each have nested agents, left reaching its step limit also truncated the right agents, and a group's own limit could fail to reach its agents. I now skip tuple parent keys whose prefix differs from the nested key's prefix. This is the same grouping EnvBase uses for reset keys. Root string keys still apply to every nested entry, and the existing dimension expansion and logical OR are unchanged. I also added two lines to the class docstring.

With the environment from #4586 (two groups of 3 agents, StepCounter(3), left at step 2):

	right agents truncated
main	[[True], [True], [True]]
this PR	[[False], [False], [False]]

Tests:

• The new TestStepCounter.test_stepcounter_independent_nested_groups goes through the public reset() / rollout(). It covers either group truncating, with update_done off and on, and keeps an agent flag that was already set. All 4 cases fail on main.
• With the fix, test/transforms/test_env_transforms.py::TestStepCounter: 82 passed, 2 skipped.
• pre-commit passes on the changed files.
• I ran these with the Python sources on top of the released wheel, without building the optional C++ extension; this transform doesn't use it.

Motivation and Context

Multi-agent environments with independent groups (each with its own done/reset keys) got wrong truncation flags for agents in groups that had not hit their step limit.

close #4586

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)

Checklist

• 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.

I used AI assistance (Claude and GPT-6 Astra) to find, fix and test this bug.

[meta-cla]

Hi @mottopanikeiku!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 263ae5b with merge base e7a05a7 ():

NEW FAILURE - The following job has failed:

• Generate documentation / build-docs (3.12, 12.8) / linux-job (gh)
  Could not load credentials from any providers

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

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[mottopanikeiku]

@bsprenger could you approve the CI workflows for this PR? It is my first one here, so they are waiting for a maintainer. It fixes #4586: with independent done groups, StepCounter marked the agents of one group as truncated when another group reached its step limit.

[coder-jayp]

LGTM! Verified that prefix-matching in _propagate_to_nested_keys resolves cross-group leakage without breaking root-level propagation, and the new test thoroughly covers both groups, update_done states, and pre-existing flags.