ClipPPOLoss problem with MaskedCategorical dist

Jean-Marc_Fauche: <p>Here is a loss module:</p>
<pre><code class="lang-auto">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",
)
</code></pre>
<p>and in training loop:</p>
<pre><code class="lang-auto">        for _ in range(frames_per_batch // minibatch_size):
            subdata = replay_buffer.sample(minibatch_size)
            print(subdata)
            loss_vals = loss_module(subdata.to(device)) # &lt;-- ERR unbatched !

</code></pre>
<p>Here is the result of print(subdata):</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>and here is the traceback error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/fauche/Documents/RenforcementLearning/CosPPO2.py", line 220, in &lt;module&gt;
    loss_vals = loss_module(subdata.to(device)) # &lt;-- 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]

</code></pre>
<p>why expected target size is [12, 12] ? minibatch_size = 12<br>
For more info here is the policy_module:</p>
<pre><code class="lang-auto">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
)
</code></pre>

KFrank: <p>Hi Jean-Marc!</p>
<aside class="quote no-group quote-modified" data-username="Jean-Marc_Fauche" data-post="1" data-topic="222394">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jean-marc_fauche/48/64545_2.png" class="avatar"> Jean-Marc_Fauche:</div>
<blockquote>
<p>and here is the traceback error:</p>
<pre><code class="lang-auto">RuntimeError: Expected target size [12, 12], got [12, 1]

</code></pre>
</blockquote>
</aside>
<p>Without seeing the relevant code (nor knowing which version of pytorch you’re using), it’s</p>
<p>hard to know for sure what’s going on.  But I think it’s pretty likely that you’re passing in a</p>
<p>target of the wrong type / shape to <code>CrossEntropyLoss</code>.</p>
<p><code>CrossEntropyLoss</code> works in two modes:  It can take a target consisting of integer class</p>
<p>labels (and no class dimension) or a target consisting of floating-point class probabilities</p>
<p>(with a class dimension).  It determines which mode it thinks you are using from the type</p>
<p>and shape of the target passed in.</p>
<p>Try printing out the type and shape of both the prediction and target you pass into</p>
<p><code>CrossEntropyLoss</code> (or whatever version of the loss criterion you are using) just before you</p>
<p>call it.</p>
<p>Here’s an illustration of how <code>CrossEntropyLoss</code> does (and doesn’t) work:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.8.0+cu129'
&gt;&gt;&gt; _ = torch.manual_seed (2025)
&gt;&gt;&gt; pred = torch.randn (12, 4)
&gt;&gt;&gt; trga = torch.randint (4, (12,))            # integer class labels -- no class dimension
&gt;&gt;&gt; trgb = torch.rand (12, 4)                  # float class probabilities -- class dimension
&gt;&gt;&gt; torch.nn.CrossEntropyLoss() (pred, trga)   # works
tensor(2.1806)
&gt;&gt;&gt; torch.nn.CrossEntropyLoss() (pred, trgb)
tensor(3.6652)
&gt;&gt;&gt; trgc = trga.float()                        # float class labels -- won't work
&gt;&gt;&gt; trgd = trgc.unsqueeze (-1)                 # add singleton class dimension -- still won't work
&gt;&gt;&gt; torch.nn.CrossEntropyLoss() (pred, trgc)   # fails
Traceback (most recent call last):
  File "&lt;python-input-87&gt;", line 1, in &lt;module&gt;
    torch.nn.CrossEntropyLoss() (pred, trgc)   # fails
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/loss.py", line 1310, in forward
    return F.cross_entropy(
           ~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...&lt;4 lines&gt;...
        label_smoothing=self.label_smoothing,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/functional.py", line 3462, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...&lt;4 lines&gt;...
        label_smoothing,
        ^^^^^^^^^^^^^^^^
    )
    ^
RuntimeError: expected scalar type Long but found Float
&gt;&gt;&gt; torch.nn.CrossEntropyLoss() (pred, trgd)   # fails
Traceback (most recent call last):
  File "&lt;python-input-88&gt;", line 1, in &lt;module&gt;
    torch.nn.CrossEntropyLoss() (pred, trgd)   # fails
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "&lt;path_to_pytorch_install&gt;/torch/nn/modules/loss.py", line 1310, in forward
    return F.cross_entropy(
           ~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...&lt;4 lines&gt;...
        label_smoothing=self.label_smoothing,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "&lt;path_to_pytorch_install&gt;/torch/nn/functional.py", line 3462, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        input,
        ^^^^^^
    ...&lt;4 lines&gt;...
        label_smoothing,
        ^^^^^^^^^^^^^^^^
    )
    ^
RuntimeError: 0D or 1D target tensor expected, multi-target not supported


</code></pre>
<aside class="quote no-group" data-username="Jean-Marc_Fauche" data-post="1" data-topic="222394">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jean-marc_fauche/48/64545_2.png" class="avatar"> Jean-Marc_Fauche:</div>
<blockquote>
<p>why expected target size is [12, 12] ? minibatch_size = 12</p>
</blockquote>
</aside>
<p>I don’t know where the second <code>12</code> in “expected target size’ is coming from.  My guess would</p>
<p>be that <code>CrossEntropyLoss</code> has somehow deduced that you have twelve classes (and is</p>
<p>looking for twelve target floating-point probabilities, one for each class.  (But I can’t reproduce</p>
<p>your error message in detail, although this maybe a pytorch version issue.)</p>
<p>Best.</p>
<p>K. Frank</p>

Jean-Marc_Fauche: <p>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]).<br>
Best regards.<br>
Jean-Marc</p>