Repository navigation
[BugFix] Let TensorDictMaxValueWriter extend with batches of trajectories - #4713
rishi02102017 wants to merge 1 commit into
Conversation
…ries TensorDictMaxValueWriter ranks each element along the first dimension of the data passed to extend(), so a [batch, time] rollout is stored one trajectory per slot, as described in the class docstring. The returned index was allocated with the full data shape, though, so writing the per-trajectory storage indices into it failed with a shape mismatch as soon as more than one trajectory was kept. Allocate the index along the first dimension only, which matches what RoundRobinWriter returns for a one-dimensional storage.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4713
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @rishi02102017! 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! |
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! |
Description
TensorDictMaxValueWriter.extendranks each element along the first dimension of the data, so a[batch, time]rollout is stored one trajectory per slot, as described in the class docstring. The index it returns was allocated with the full data shape, so writing the per-trajectory storage indices into it raised a shape mismatch whenever more than one trajectory was kept:This PR allocates the returned index along the first dimension only, in both the insertion and the no-insertion branches, which matches what
RoundRobinWriterreturns for a one-dimensional storage. One-dimensional data behaves as before.A regression test in
test/rb/test_writers.pyextends a buffer of capacity 2 with four 5-step trajectories and checks the returned index, that the two trajectories with the highest summed key are kept, and that sampling returns whole trajectories.Motivation and Context
The trajectory example in the
TensorDictMaxValueWriterdocstring (rb.extend(rollout)followed bysample.shape == (5, 10)) currently fails on main and on torchrl 0.14.1.close #4712
Types of changes
Checklist