[Refactor] Extract discrete action-value selection into shared helper by coder-jayp · Pull Request #4367 · pytorch/rl

coder-jayp · 2026-09-13T07:27:58Z

Description

Refactored the discrete action-value selection logic that was previously duplicated across DQNLoss, DiscreteSACLoss, and QMixerLoss.

Specifically, this PR introduces a centralized private helper function _select_action_value in torchrl.objectives.utils that handles selecting the correct Q-value based on the action_space (e.g., categorical vs. one-hot) and properly manages the keepdim dimension parameters.

I've replaced the inline if action_space == "categorical": ... else: ... branching in the respective loss modules with this new helper. I also added rigorous behavioral regression tests to test_dqn.py and test_sac.py using a deterministic IdentityQNet to mathematically assert the correct selection indices and verify that Q-values are being properly gathered.

Refactor example

For example, DQNLoss.forward replaces its inline action selection with the shared helper:

- if self.action_space == "categorical":
-     if action.ndim != pred_val.ndim:
-         action = action.unsqueeze(-1)
-     pred_val_index = torch.gather(pred_val, -1, index=action).squeeze(-1)
- else:
-     action = action.to(torch.float)
-     pred_val_index = (pred_val * action).sum(-1)
+ pred_val_index = _select_action_value(self.action_space, action, pred_val)

DiscreteSACLoss uses the same helper. Double-DQN target evaluation and QMixerLoss pass keepdim=True to retain the trailing singleton dimension.

Motivation and Context

Multiple loss modules contained identical branching logic for extracting chosen Q-values from action spaces, marked with a TODO. Centralizing this removes code duplication, makes it easier to support new action encodings in the future, and helps prevent edge-case shape handling bugs (e.g., unsqueeze/squeeze issues on discrete categorical tensors) in individual modules.

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

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Refactor (non-breaking change which improves code quality and removes duplication)

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.

pytorch-bot · 2026-09-13T07:28:01Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

⚠️ 16 Awaiting Approval

As of commit 45fe387 with merge base 06d57a0 ():

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

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

coder-jayp · 2026-09-13T07:28:41Z

@torchrlbot reviewer @vmoens

github-actions · 2026-09-13T07:29:19Z

Requested review from @vmoens (requested by @coder-jayp).

vmoens

The helper extraction looks sound, and I found no production correctness regression. Please correct the SAC regression test before merging: it currently passes with incorrect action indexing and would fail with the identity critics described in its setup. The QMixer test setup can also be consolidated with the existing coverage.

Validation: 242 targeted CPU test cases passed, including all four new tests. I additionally confirmed that the SAC test passes when selection incorrectly uses (action + 1) % num_actions, and that restoring identity critic weights makes both losses correctly equal 8.

vmoens · 2026-09-13T16:51:24Z

+        # The two losses must differ because action=0 selects Q=2.0 and
+        # action=3 selects Q=-2.0; they produce different MSE errors.
+        # If gather returns the same element for both actions, the losses are equal.
+        assert not torch.isclose(loss_action0, loss_action3, atol=1e-3), (


[P2] Assert known correct SAC losses

Constructing DiscreteSACLoss expands and resamples the critic parameters, so the identity weights initialized above are replaced with random weights. This inequality therefore does not verify the claimed Q-values or correct indexing: I confirmed that it still passes if the helper selects (action + 1) % num_actions.

With the intended identity critics, Q-values +2 and -2 have the same squared error against the configured zero target. Both losses are correctly 8 after summing over the two critics, so restoring the intended weights makes this assertion fail. Please pin the functional critic parameters after constructing the loss, use asymmetric Q-values or a nonzero reward, and assert the exact expected losses.

I've updated the test to explicitly pin the functional parameters after the loss module is constructed. The test now uses asymmetric Q-values and asserts the exact mathematically expected losses (2.0 and 0.0) based on the identity weights.

vmoens · 2026-09-13T16:51:24Z

+            rtol=0.0,
+        )
+
+    def test_qmixer_categorical_loss_value_is_correct(self):


[P3] Consolidate the redundant QMixer test setup

This roughly 100-line case recreates the existing actor, mixer, and data factories, but only asserts that the loss is finite and nonnegative. TestQMixer.test_qmixer already exercises categorical selection, the required mixer input shape, and backward. If the single-agent case is useful, please add it to that test's parametrization; alternatively, give this case an independently calculated expected loss.

The PR adds 351 test lines for 56 implementation lines (about 6.3:1, even counting the helper docstring). Reusing the existing setup and consolidating the numerical cases would make the regression coverage easier to review and maintain.

Consolidating into TestQMixer allowed us to reuse the mock factories. The test now asserts that distinct asymmetric Q-values produce distinct scalar losses, verifying the exact gather logic.

- Pin SAC critic parameters explicitly to test exact identity-based Q-values and expected losses - Consolidate QMixer regression test into TestQMixer and assert distinct asymmetric Q-values produce distinct losses

vmoens

The SAC finding is addressed: the updated test pins the functional critic weights, checks the expected losses 2.0 and 0.0, and now fails with deliberately shifted action indexing. The production implementation is unchanged from my previous review.

The revised QMixer test still passes when the selector ignores the supplied action and always selects action zero, because the two calls use different global states. Please isolate the action change as described inline, and remove the duplicated direct-run entry point.

Validation: 242 targeted CPU test cases passed on this commit. The SAC incorrect-index mutation was detected; the QMixer constant-index mutation was not. The revised tests reuse the actor and mixer factories, but total test additions have grown to 392 lines for 56 implementation additions (about 7:1), so the consolidation concern remains.

I also added a before-and-after code example to the PR description.

vmoens · 2026-09-14T08:46:15Z

+                        },
+                        batch_size=[batch, n_agents],
+                    ),
+                    "state": torch.randn(batch, 64, 64, 3),


[P2] Reuse the same global state when comparing actions

Each make_td call samples a different global state, and QMixer uses that state to generate its mixing weights and biases. The losses can therefore differ even when action selection is broken. I replaced the selector with one that always gathers action zero: this test still passes, with losses approximately 17.14146 and 7.92686.

Build one input TensorDict, clone it for the second call, and change only the action. Prefer an independently computed expected loss so selecting two different wrong indices is also detected. If keeping the identity-Q explanation, zero the actor linear layer's bias as well: the existing actor factory creates a biased layer, and this test only resets its weight.

Switched to cloning a single base_td to guarantee identical hypernetwork states, and zeroed the actor bias. Verified locally that this now catches the constant-index mutation.

vmoens · 2026-09-14T08:46:16Z

+if __name__ == "__main__":
+    pytest.main([__file__])
+
+
+if __name__ == "__main__":


[P3] Keep a single direct-run entry point

There are now two consecutive if __name__ == "__main__" blocks calling pytest.main([__file__]). Running this test file directly executes the entire DQN/QMixer suite twice because the first call returns to the module. Please remove the duplicate block.

- Eliminate the standalone TestDiscreteActionValueSelectionRegression class. - Move DQN regression tests into TestDQN and reuse the existing _create_mock_actor factory. - Fix QMixer regression test state contamination by reusing a base TensorDict. - Zero actor linear bias in regression tests. - Remove duplicate __main__ block.

coder-jayp · 2026-09-14T09:22:46Z

Test bloat addressed. I completely deleted the 145-line standalone regression class and moved its tests directly into TestDQN, updating them to reuse self._create_mock_actor(...).

[Refactor] Extract discrete action-value selection into shared helper

7be88f1

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 13, 2026

github-actions Bot added Refactoring Refactoring of an existing feature Objectives and removed Refactoring Refactoring of an existing feature labels Sep 13, 2026

github-actions Bot requested a review from vmoens September 13, 2026 07:29

vmoens requested changes Sep 13, 2026

View reviewed changes

Fix SAC and QMixer discrete action-value regression tests

92b52c3

- Pin SAC critic parameters explicitly to test exact identity-based Q-values and expected losses - Consolidate QMixer regression test into TestQMixer and assert distinct asymmetric Q-values produce distinct losses

github-actions Bot added the Refactoring Refactoring of an existing feature label Sep 14, 2026

vmoens requested changes Sep 14, 2026

View reviewed changes

coder-jayp requested a review from vmoens September 14, 2026 09:25

Uh oh!

Conversation

coder-jayp commented Sep 13, 2026 • edited by vmoens Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Refactor example

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot Bot commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

⚠️ 16 Awaiting Approval

Uh oh!

coder-jayp commented Sep 13, 2026

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Sep 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coder-jayp commented Sep 13, 2026 •

edited by vmoens

Loading

pytorch-bot Bot commented Sep 13, 2026 •

edited

Loading