Skip to content

[BugFix] Keep MCAdvantage and TopKRewardSelector working with tensordict 0.17 - #4576

Open
peterdsharpe wants to merge 1 commit into
pytorch:mainfrom
peterdsharpe:fix-get-as-nested-tensor-dense
Open

peterdsharpe wants to merge 1 commit into
pytorch:mainfrom
peterdsharpe:fix-get-as-nested-tensor-dense

Conversation

@peterdsharpe

Copy link
Copy Markdown
Contributor

Description

MCAdvantage and TopKRewardSelector now pass as_nested_tensor=True to get only when the trajectories form a lazy stack. For a regular TensorDict, they read the rewards with a plain get.

The rewards are the same as before with every tensordict version from 0.14 on: the keyword does nothing on a regular TensorDict today. The FutureWarning from tensordict 0.15 goes away, and nothing breaks when tensordict 0.17 makes the keyword return a nested tensor for every kind of tensordict.

Tests: test_mc_advantage_same_length_trajectories and test_topk_regular_tensordict feed regular TensorDict trajectories, turn the FutureWarning into an error, and check the rewards and advantages. With tensordict main, both fail before this change and pass after. TestMCAdvantage and TestTopK pass locally (26 passed; the Ray test is skipped).

Motivation and Context

Fixes #4575. Details are in the issue.

  • 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)
  • I have updated the tests accordingly (required for a bug fix or a new feature).

…ict 0.17

Both read rewards with get(key, as_nested_tensor=True), also when the
trajectories concatenate into a regular TensorDict. The keyword does
nothing there today, but from tensordict 0.17 it returns a nested tensor
for every kind of tensordict, and tensordict 0.15 already warns about it.
Pass it only to lazy stacks, so both components keep today's rewards
with every tensordict version from 0.14 on.
@pytorch-bot

pytorch-bot Bot commented Oct 16, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⚠️ 14 Awaiting Approval

As of commit 3c75205 with merge base 603c269 (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

BugFix CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. llm/ LLM-related PR, triggers LLM CI tests Objectives

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] MCAdvantage and TopKRewardSelector will break with tensordict 0.17

1 participant