A question regarding ReinforceLoss and ProbabilisticActor behaviour

[ArturGoldman]

Hi, thank you for maintaining this project. I might be misunderstanding something, so I wanted to ask for clarification.

TL;DR

ProbabilisticActor resamples actions even when out_keys are present in the input TensorDict, in contrast to its description in documentation. As a result, when using ReinforceLoss together with ProbabilisticActor, the sample_log_prob appears to be computed on trajectories that are re-sampled inside the loss, rather than on the actions collected during rollout.

Setup

I was trying to setup a simple experiment using a combination of ProbabilisticActor and ReinforceLoss. For simplicity, I used data obtained from a single EnvBase.rollout().

While reading the docs and code, I noticed the following points:

• In the description of ProbabilisticActor for out_keys it says the following: "Importantly, if these keys are found in the input TensorDict, the sampling step will be skipped."
• In the rollout() method implementation a flag exclude_actions=False is hardcoded, see here: https://github.com/pytorch/rl/blob/main/torchrl/envs/common.py#L3645-L3646
• In ReinforceLoss.forward(), a forward pass of the actor is performed, see here: https://github.com/pytorch/rl/blob/main/torchrl/objectives/reinforce.py#L392

Expected behavior (my understanding)

Based on the above, I expected the following:

• During env.rollout() with a ProbabilisticActor, actions would only be sampled at the first step. On subsequent steps, since the input TensorDict already contains the action keys, sampling would be skipped. (This already seems somewhat unintuitive, but it is what I inferred from the documentation.)
• During the ReinforceLoss forward pass, actions would not be re-sampled, and the loss would instead compute sample_log_prob corresponding to the actions collected during rollout. This would align with the usual on-policy REINFORCE setting.

What I observed

In practice, I have observed the opposite behaviour:

• During env.rollout(), ProbabilisticActor samples a new action at every step, even when the action key is already present in the input TensorDict.
• During the ReinforceLoss forward pass, actions are also re-sampled, meaning that the loss appears to be computed using log-probabilities of actions that were not used to generate the trajectory.

This was surprising to me, as it seems to imply that ReinforceLoss may be operating on re-sampled (off-trajectory) actions.

Looking further into the ProbabilisticActor implementation, it seems that resampling is controlled not only by the presence of out_keys in the input TensorDict, but also by the skip_existing() context manager and the internal _requires_sample flag, see here: https://github.com/pytorch/tensordict/blob/main/tensordict/nn/probabilistic.py#L601-L608

Question

Is this behavior intended, or am I missing some assumption or configuration detail? I also noticed that other losses (e.g. PPOLoss) do not appear to have this issue, since they do not call the actor’s forward() method directly, but instead rely on get_dist() and dist.log_prob(), see here: https://github.com/pytorch/rl/blob/main/torchrl/objectives/ppo.py#L671-L701. I would really appreciate any clarification on the intended design here.

Code / references

I have also conducted a small test which demonstrates the behaviour I am describing:

import torch
import torch.nn as nn
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
from torchrl.envs import EnvBase
from torchrl.data import (
    CompositeSpec,
    UnboundedContinuousTensorSpec,
    DiscreteTensorSpec,
    BinaryDiscreteTensorSpec,
)
from torchrl.modules import ProbabilisticActor, ValueOperator
from torch.distributions import Categorical
from torchrl.envs.utils import step_mdp
from tensordict.nn import InteractionType, set_skip_existing
from torchrl.objectives.reinforce import ReinforceLoss
from torchrl.objectives import ClipPPOLoss

torch.manual_seed(0)


# --- Minimal env: state is a random vector; reward/done are trivial ---
class TinyEnv(EnvBase):
    def __init__(self, obs_dim=4, device="cpu", batch_size=()):
        super().__init__(device=device, batch_size=batch_size)
        self.obs_dim = obs_dim

        self.observation_spec = CompositeSpec(
            observation=UnboundedContinuousTensorSpec(shape=(*batch_size, obs_dim), device=device)
        )
        self.action_spec = DiscreteTensorSpec(2, shape=(*batch_size, 1), device=device)
        self.reward_spec = UnboundedContinuousTensorSpec(shape=(*batch_size, 1), device=device)
        self.done_spec = BinaryDiscreteTensorSpec(shape=(*batch_size, 1), device=device)

    def _reset(self, tensordict=None):
        obs = torch.randn(*self.batch_size, self.obs_dim, device=self.device)
        td = TensorDict(
            {
                "observation": obs,
                "done": torch.zeros(*self.batch_size, 1, dtype=torch.bool, device=self.device),
                "terminated": torch.zeros(*self.batch_size, 1, dtype=torch.bool, device=self.device),
            },
            batch_size=self.batch_size,
            device=self.device,
        )
        return td

    def _set_seed(self, seed) -> int:
        if seed is None:
            seed = torch.seed() % (2**31 - 1)
        return int(seed)

    def _step(self, tensordict):
        # ignore action; just advance observation
        obs = torch.randn(*self.batch_size, self.obs_dim, device=self.device)
        reward = torch.zeros(*self.batch_size, 1, device=self.device)

        # never ends (for this test)
        done = torch.zeros(*self.batch_size, 1, dtype=torch.bool, device=self.device)
        terminated = torch.zeros(*self.batch_size, 1, dtype=torch.bool, device=self.device)

        out = TensorDict(
                {
                    "observation": obs,
                    "reward": reward,
                    "done": done,
                    "terminated": terminated,
                },
                batch_size=self.batch_size,
                device=self.device,
            )
        return out


# --- Build env ---
env = TinyEnv(obs_dim=4, device="cpu", batch_size=())

# --- Build a simple probabilistic actor: obs -> logits -> action (Categorical) ---
obs_dim = env.observation_spec["observation"].shape[-1]
n_actions = 2

class LogNet(nn.Module):
    def __init__(self, obs_dim, n_actions):
        super().__init__()
        self.n_actions = n_actions

    def forward(self, obs):
        batch_shape = obs.shape[:-1]
        logits = torch.ones(*batch_shape, self.n_actions, device=obs.device)
        return logits / self.n_actions

logits_net = LogNet(obs_dim, n_actions)
logits_module = TensorDictModule(
    logits_net,
    in_keys=["observation"],
    out_keys=["logits"],
)


def run_reinforce_rollout_case(*, return_log_prob: bool, skip_existing: bool):
    print("\n--- reinforce + rollout case ---")
    print("return_log_prob:", return_log_prob, "skip_existing:", skip_existing)

    policy_case = ProbabilisticActor(
        module=logits_module,
        spec=env.action_spec,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=Categorical,
        return_log_prob=return_log_prob,
        default_interaction_type=InteractionType.RANDOM,
    )

    value_net = nn.Linear(obs_dim, 1)
    critic_case = ValueOperator(value_net, in_keys=["observation"])
    loss_mod = ReinforceLoss(policy_case, critic_case)

    with torch.no_grad():
        rollout_td = env.rollout(max_steps=5, policy=policy_case, auto_reset=True)

    action_before = rollout_td.get("action").clone()
    try:
        with set_skip_existing(True) if skip_existing else set_skip_existing(False):
            loss_td = loss_mod(rollout_td)
        action_after = rollout_td.get("action").clone()
        changed = (action_before != action_after).any().item()
        print("actions changed during loss:", bool(changed))
        print("loss keys:", sorted([k for k in loss_td.keys(True, True) if str(k).startswith("loss")]))
    except Exception as exc:
        print("reinforce_loss failed:", type(exc).__name__, str(exc))

run_reinforce_rollout_case(return_log_prob=True, skip_existing=False)
run_reinforce_rollout_case(return_log_prob=True, skip_existing=True)

def run_rollout_outkeys_present_resample():
    print("\n--- rollout out_keys present but resample ---")
    policy_case = ProbabilisticActor(
        module=logits_module,
        spec=env.action_spec,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=Categorical,
        return_log_prob=True,
        default_interaction_type=InteractionType.RANDOM,
    )

    class PrefillPolicy(nn.Module):
        def __init__(self, base):
            super().__init__()
            self.base = base
            self.i = 0

        def forward(self, td):
            for key in self.base.out_keys:
                if key not in td.keys(True, True):
                    td.set(key, torch.zeros(*td.batch_size, 1))
            td.set("action", torch.full((*td.batch_size, 1), 2, dtype=torch.long))
            print(f"[prefill call {self.i}] out_keys present: {all(k in td.keys(True, True) for k in self.base.out_keys)}")
            action_before = td.get("action").clone()
            self.i += 1
            td_out = self.base(td)
            action_after = td_out.get("action").clone()
            print("action before:", action_before.squeeze().tolist(), "action after:", action_after.squeeze().tolist())
            return td_out

    prefill_policy = PrefillPolicy(policy_case)
    with set_skip_existing(False):
        env.rollout(max_steps=3, policy=prefill_policy, auto_reset=True)


run_rollout_outkeys_present_resample()

Output of the script:

--- reinforce + rollout case ---
return_log_prob: True skip_existing: False
actions changed during loss: True
loss keys: ['loss_actor', 'loss_value']

--- reinforce + rollout case ---
return_log_prob: True skip_existing: True
actions changed during loss: False
loss keys: ['loss_actor', 'loss_value']

--- rollout out_keys present but resample ---
[prefill call 0] out_keys present: True
action before: 2 action after: 0
[prefill call 1] out_keys present: True
action before: 2 action after: 0
[prefill call 2] out_keys present: True
action before: 2 action after: 1

[theap06]

@ArturGoldman I just reproduced this scenario and it seems that there is bug where the rollouts get overwritten with new samples. I can draft a PR fixing this issue for the REINFORCE Loss. I'll also update the docstring for the ProbabilisticActor behavior in this pr.