ClipPPOLoss problem with MaskedCategorical dist

reinforcement-learning
1

Here is a loss module:

loss_module = ClipPPOLoss(
    actor_network = policy_module,
    critic_network = critic,
    clip_epsilon = clip_epsilon,
    entropy_bonus = bool(entropy_eps),
    entropy_coeff = entropy_eps,
    # these keys match by default but we set this for completeness
    critic_coeff = 1.0,
    loss_critic_type = "smooth_l1",
)

and in training loop:

        for _ in range(frames_per_batch // minibatch_size):
            subdata = replay_buffer.sample(minibatch_size)
            print(subdata)
            loss_vals = loss_module(subdata.to(device)) # <-- ERR unbatched !

Here is the result of print(subdata):

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        action_log_prob: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        action_mask: Tensor(shape=torch.Size([12, 1, 4]), device=cpu, dtype=torch.bool, is_shared=False),
        advantage: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        angle: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        bid: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        collector: TensorDict(
            fields={
                traj_ids: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([12]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        logits: Tensor(shape=torch.Size([12, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        next: TensorDict(
            fields={
                action_mask: Tensor(shape=torch.Size([12, 1, 4]), device=cpu, dtype=torch.bool, is_shared=False),
                angle: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                bid: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                done: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([12, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                params: TensorDict(
                    fields={
                        a: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                        b: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.int64, is_shared=False),
                        dt: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                        h: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                        k: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([12]),
                    device=cpu,
                    is_shared=False),
                posA: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                posV: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                solde: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                state_value: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                t: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([12]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([12, 6]), device=cpu, dtype=torch.float32, is_shared=False),
        params: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                b: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.int64, is_shared=False),
                dt: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                h: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False),
                k: Tensor(shape=torch.Size([12]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([12]),
            device=cpu,
            is_shared=False),
        posA: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        posV: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        solde: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        state_value: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        t: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        value_target: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([12]),
    device=cpu,
    is_shared=False)

and here is the traceback error:

Traceback (most recent call last):
  File "/home/fauche/Documents/RenforcementLearning/CosPPO2.py", line 220, in <module>
    loss_vals = loss_module(subdata.to(device)) # <-- ERR unbatched !
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1857, in _call_impl
    return inner()
           ^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1805, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torchrl/objectives/common.py", line 54, in new_forward
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/tensordict/nn/common.py", line 328, in wrapper
    return func(_self, tensordict, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torchrl/objectives/ppo.py", line 1174, in forward
    log_weight, dist, kl_approx = self._log_weight(
                                  ^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torchrl/objectives/ppo.py", line 681, in _log_weight
    log_prob, dist, is_composite = self._get_cur_log_prob(tensordict)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torchrl/objectives/ppo.py", line 647, in _get_cur_log_prob
    log_prob = dist.log_prob(action)
               ^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torchrl/modules/distributions/discrete.py", line 362, in log_prob
    result = -torch.nn.functional.cross_entropy(logits, value, reduce=False)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fauche/Documents/IAvenv/lib/python3.11/site-packages/torch/nn/functional.py", line 3494, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected target size [12, 12], got [12, 1]

why expected target size is [12, 12] ? minibatch_size = 12
For more info here is the policy_module:

policy_module = ProbabilisticActor(
    module = policy_module,
    spec = env.action_spec,
    in_keys={"logits": "logits", "mask": "action_mask"},
    out_keys=["action"],
    distribution_class = MaskedCategorical,
    return_log_prob = True
)
2

Hi Jean-Marc!

Without seeing the relevant code (nor knowing which version of pytorch you’re using), it’s

hard to know for sure what’s going on. But I think it’s pretty likely that you’re passing in a

target of the wrong type / shape to CrossEntropyLoss.

CrossEntropyLoss works in two modes: It can take a target consisting of integer class

labels (and no class dimension) or a target consisting of floating-point class probabilities

(with a class dimension). It determines which mode it thinks you are using from the type

and shape of the target passed in.

Try printing out the type and shape of both the prediction and target you pass into

CrossEntropyLoss (or whatever version of the loss criterion you are using) just before you

call it.

Here’s an illustration of how CrossEntropyLoss does (and doesn’t) work:

>>> import torch
>>> torch.__version__
'2.8.0+cu129'
>>> _ = torch.manual_seed (2025)
>>> pred = torch.randn (12, 4)
>>> trga = torch.randint (4, (12,))            # integer class labels -- no class dimension
>>> trgb = torch.rand (12, 4)                  # float class probabilities -- class dimension
>>> torch.nn.CrossEntropyLoss() (pred, trga)   # works
tensor(2.1806)
>>> torch.nn.CrossEntropyLoss() (pred, trgb)
tensor(3.6652)
>>> trgc = trga.float()                        # float class labels -- won't work
>>> trgd = trgc.unsqueeze (-1)                 # add singleton class dimension -- still won't work
>>> torch.nn.CrossEntropyLoss() (pred, trgc)   # fails
Traceback (most recent call last):
  File "<python-input-87>", line 1, in <module>
    torch.nn.CrossEntropyLoss() (pred, trgc)   # fails
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "<path_to_pytorch_install>/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "<path_to_pytorch_install>/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "<path_to_pytorch_install>/torch/nn/modules/loss.py", line 1310, in forward
    return F.cross_entropy(
           ~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...<4 lines>...
        label_smoothing=self.label_smoothing,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "<path_to_pytorch_install>/torch/nn/functional.py", line 3462, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...<4 lines>...
        label_smoothing,
        ^^^^^^^^^^^^^^^^
    )
    ^
RuntimeError: expected scalar type Long but found Float
>>> torch.nn.CrossEntropyLoss() (pred, trgd)   # fails
Traceback (most recent call last):
  File "<python-input-88>", line 1, in <module>
    torch.nn.CrossEntropyLoss() (pred, trgd)   # fails
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "<path_to_pytorch_install>/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "<path_to_pytorch_install>/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "<path_to_pytorch_install>/torch/nn/modules/loss.py", line 1310, in forward
    return F.cross_entropy(
           ~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...<4 lines>...
        label_smoothing=self.label_smoothing,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "<path_to_pytorch_install>/torch/nn/functional.py", line 3462, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...<4 lines>...
        label_smoothing,
        ^^^^^^^^^^^^^^^^
    )
    ^
RuntimeError: 0D or 1D target tensor expected, multi-target not supported

I don’t know where the second 12 in “expected target size’ is coming from. My guess would

be that CrossEntropyLoss has somehow deduced that you have twelve classes (and is

looking for twelve target floating-point probabilities, one for each class. (But I can’t reproduce

your error message in detail, although this maybe a pytorch version issue.)

Best.

K. Frank

3

Hi Frank! Thank you for your response. I found the problem that came from the difference in shapes between action [12,1] and action_mask [12,1,4] instead of [12] and [12,4] when batch_size = torch.Size([12]).
Best regards.
Jean-Marc