[BUG] GRPOLoss masking strategies produce incompatible shapes

Describe the bug

GRPOLoss is documented to support masking_strategy in {"sft", "rlhf", "generic"}, but the test that actually compares those strategies is skipped because of a shape mismatch. CI therefore will not catch a regression in the public loss when the mask strategy changes.

This issue asks for the intended output shape per strategy and for the test to be unskipped once the loss matches that contract. It is not a request to land a speculative fix.

To Reproduce

The skip sits on TestGRPOLoss.test_grpo_loss_with_real_models:

    @pytest.mark.skipif(not _has_vllm, reason="vllm not available")
    @pytest.mark.parametrize("masking_strategy", ["sft", "rlhf"])
    @pytest.mark.skip(
        reason="GRPOLoss shape mismatch between masking strategies - needs investigation"
    )
    def test_grpo_loss_with_real_models(
        self,
        vllm_instance,
        transformers_instance,
        masking_strategy,
    ):

The test builds tokens-mode input for "sft" and history-mode input for "rlhf" (lines 1413-1442), generates with vLLMWrapper(..., return_log_probs=True), then constructs

        loss_fn = GRPOLoss(actor_network=wrapper, masking_strategy=masking_strategy)

and expects a finite result.loss_objective.

On the loss side, masking_strategy selects different distribution helpers:

            if self.masking_strategy == "sft" and hasattr(
                self.actor_network, "_get_sft_dist"
            ):
                dist = self.actor_network._get_sft_dist(tensordict)
            elif self.masking_strategy == "rlhf" and hasattr(
                self.actor_network, "_get_rlhf_dist"
            ):
                dist = self.actor_network._get_rlhf_dist(tensordict)
            elif self.masking_strategy == "generic" and hasattr(
                self.actor_network, "_get_generic_dist"
            ):
                dist = self.actor_network._get_generic_dist(tensordict)

The class docstring already warns that a mismatch with the advantage mask produces shape errors (grpo.py:377-383):

The masking strategy must match the strategy used for advantage computation to avoid shape mismatches.

Because the only cross-strategy test is skipped, that warning is not enforced.

Expected behavior

Please state the intended output shape of GRPOLoss (and of the intermediate log-prob / mask tensors) for each of "sft", "rlhf", and "generic", including how it should interact with tokens vs history input and with the advantage tensor.

Once the loss matches that contract, unskip test_grpo_loss_with_real_models (or replace it with a smaller deterministic test that still fails if the shapes diverge again).

Screenshots

N/A

System info

Observed on a source checkout of pytorch/rl at 1d3de3db. The skipped test also needs vllm and transformers.

Additional context

I am not proposing a particular reduction or broadcast rule here. The skip reason is the whole report: the public API claims three strategies, and one of them currently cannot be compared to the other in CI.

Reason and Possible fixes

Unskip only after the intended shapes are specified and the loss (or the test inputs / advantage) is aligned to them. A drive-by squeeze / unsqueeze without that spec is likely to hide the next mismatch.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

Start with test/llm/test_llm_objectives.py:1398-1469 and torchrl/objectives/llm/grpo.py:377-383,609-620. Compare the tokens-mode and history-mode inputs and the distribution helpers for each masking strategy, then establish the intended shapes for loss, log-prob, masks, and advantages. Done means the loss matches that contract and the real-model test is unskipped, or a deterministic replacement catches future shape divergence.