Repository navigation
[BugFix] Accumulate offline RewardSum along the named time dimension - #4539
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4539
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 8 PendingAs of commit da61c81 with merge base 264a3c2 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @aalquwayfili! Thank you for your pull request and welcome to our community. Action RequiredIn 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. ProcessIn 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 If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
4fe8739 to
da61c81
Compare
bsprenger
left a comment
bsprenger
left a comment
There was a problem hiding this comment.
LGTM thanks for contributing!
2 participants
Description
Use the actual index of the TensorDict batch dimension named
timewhen computing offlineRewardSumtotals. Add six numeric regression cases covering unbatched, batch-first and time-first trajectories, flat and nested reward keys, and independent reward components.Motivation and Context
Fixes #4538.
RewardSum.forwardsubtracts one from the named time dimension. For rewards shaped[time, 1], this accumulates over the reward feature; for[batch, time, 1], it accumulates across independent batch entries. For example, rewards[1, 2, 3]remain[1, 2, 3]instead of producing[1, 3, 6]. The existing one-dimensional reward test cannot distinguish axes0and-1.TensorDict batch dimensions are the leading dimensions of their tensors; looking up a reward leaf does not remove an axis. Using that index directly restores the documented time accumulation. Standard environment step/reset accumulation is unchanged; direct batch and replay calls observe the corrected totals.
Types of changes
Validation
All six added regressions fail numerically before the fix. The full
TestRewardSumclass passes on CPU: 40 passed, 1 skipped. The skip is the existing unsupported inverse-transform test. The two changed files pass the pinned upstream ufmt and flake8 checks.Additional local checks cover batch-axis permutations, nested agent dimensions, multiple reward keys, gradients, empty/singleton time axes, and two calls through
torch.compile(..., backend="eager", fullgraph=True). This compile check covers eager graph capture, not Inductor or GPU execution.A further public-API check used the same three-step CartPole rollout for online and offline accumulation, with one environment and two serial environments. The unchanged online accumulator and patched offline transform agree; the original offline transform does not. These rollouts contain no episode boundary.
Validation used Python 3.13, PyTorch 2.14.0 and TensorDict 0.14.3 with imports from this checkout. Full repository tests, GPU checks and the documentation build were not run.
Checklist